using Primes, Printf
function strfactor(n::Integer)
    n > -2 || return "-1 × " * strfactor(-n)
    isprime(n) || n < 2 && return dec(n)
    f = factor(Vector{typeof(n)}, n)
    return join(f, " × ")
end

lo, hi = -4, 40
println("Factor print $lo to $hi:")
for n in lo:hi
    @printf("%5d = %s\n", n, strfactor(n))
end


# indicator of the L1 norm ball with given radius

export IndBallL1

"""
**Indicator of a ``L_1`` norm ball**

    IndBallL1(r=1.0)

Returns the indicator function of the set
```math
S = \\left\\{ x : \\sum_i |x_i| \\leq r \\right\\}.
```
Parameter `r` must be positive.
"""
struct IndBallL1{R} <: ProximableFunction
    r::R
    function IndBallL1{R}(r::R) where R
        if r <= 0
            error("parameter r must be positive")
        else
            new(r)
        end
    end
end

is_convex(f::IndBallL1) = true
is_set(f::IndBallL1) = true
is_prox_accurate(f::IndBallL1) = false

IndBallL1(r::R=1.0) where R = IndBallL1{R}(r)

function (f::IndBallL1)(x::AbstractArray{T}) where {R <: Real, T <: RealOrComplex{R}}
    if norm(x, 1) - f.r > f.r*eps(R)
        return R(Inf)
    end
    return R(0)
end

function prox!(y::AbstractArray{R}, f::IndBallL1, x::AbstractArray{R}, _::R=R(1)) where {R <: Real}
    if norm(x, 1) <= f.r
        y .= x
        return R(0)
    else # do a projection of abs(x) onto simplex then recover signs
        abs_x = abs.(x)
        simplex_proj_condat!(y, f.r, abs_x)
        y .*= sign.(x)
        return R(0)
    end
end

function prox!(y::AbstractArray{T}, f::IndBallL1, x::AbstractArray{T}, _::R=R(1)) where {R <: Real, T <: Complex{R}}
    if norm(x, 1) <= f.r
        y .= x
        return R(0)
    else # do a projection of abs(x) onto simplex then recover signs
        abs_x = real.(abs.(x))
        y_temp = similar(abs_x)
        simplex_proj_condat!(y_temp, f.r, abs_x)
        y .= y_temp .* sign.(x)
        return R(0)
    end
end

fun_name(f::IndBallL1) = "indicator of an L1 norm ball"
fun_dom(f::IndBallL1) = "AbstractArray{Real}, AbstractArray{Complex}"
fun_expr(f::IndBallL1) = "x ↦ 0 if ‖x‖_1 ⩽ r, +∞ otherwise"
fun_params(f::IndBallL1) = "r = $(f.r)"

function prox_naive(f::IndBallL1, x::AbstractArray{T}, _::R=R(1)) where {R <: Real, T <: RealOrComplex{R}}
    # do a simple bisection (aka binary search) on λ
    L = R(0)
    U = maximum(abs, x)
    λ = L
    v = R(0)
    maxit = 120
    for _ in 1:maxit
        λ = (L + U) / 2
        v = sum(max.(abs.(x) .- λ, R(0)))
        # modify lower or upper bound
        (v < f.r) ? U = λ : L = λ
        # exit condition
        if abs(L - U) < (1 + abs(U))*eps(R)
            break
        end
    end
    return sign.(x) .* max.(R(0), abs.(x) .- λ), R(0)
end


vertices(array :: AbstractArray{Bool}) = sum(array)

function ortvec(N, i)
    vec = zeros(Int, N)
    vec[i] = 1

    return CartesianIndex(vec...)
end

function segments(array :: AbstractArray{Bool, N}) where N
    δ = [ortvec(N, i) for i in 1:N]

    indices    = CartesianIndices(array)
    fidx, lidx = first(indices), last(indices)
    uidx       = oneunit(fidx)

    mapreduce(+, indices) do idx
        slice = array[idx:min(idx + uidx, lidx)]

        mapreduce(+, 1:N) do i
            idx2 = uidx + δ[i]
            (checkbounds(Bool, slice, idx2) && slice[uidx] && slice[idx2]) ? 1 : 0
        end
    end
end

function checkface(array :: AbstractArray{Bool, N},
                   δ1    :: CartesianIndex{N},
                   δ2    :: CartesianIndex{N}) where N
    uidx = array |> CartesianIndices |> first |> oneunit

    idx1 = uidx
    idx2 = uidx + δ1
    idx3 = uidx + δ2
    idx4 = uidx + δ1 + δ2

    return checkbounds(Bool, array, idx2) &&
        checkbounds(Bool, array, idx3) &&
        checkbounds(Bool, array, idx4) &&
        array[idx1] && array[idx2] && array[idx3] && array[idx4]
end

function faces(array :: AbstractArray{Bool, 3})
    indices    = CartesianIndices(array)
    fidx, lidx = first(indices), last(indices)
    uidx       = oneunit(fidx)

    mapreduce(+, indices) do idx
        slice = array[idx:min(idx + uidx, lidx)]

        checkface(slice, CartesianIndex(1, 0, 0), CartesianIndex(0, 1, 0)) +
            checkface(slice, CartesianIndex(1, 0, 0), CartesianIndex(0, 0, 1)) +
            checkface(slice, CartesianIndex(0, 1, 0), CartesianIndex(0, 0, 1))
    end
end

function volumes(array :: AbstractArray{Bool})
    indices    = CartesianIndices(array)
    fidx, lidx = first(indices), last(indices)
    uidx       = oneunit(fidx)

    mapreduce(+, fidx:(lidx - uidx)) do idx
        slice = array[idx:(idx + uidx)]
        all(isone, slice)
    end
end

"""
    euler_characteristic(array :: AbstractArray{Bool})

Calculate Euler characteristic for a set of points `S` on a regular
cubic or square grid. The points are defined by a binary array `array`
in such a manner that if `array[idx] == true` then `idx ∈ S`.
"""
function euler_characteristic end

euler_characteristic(array :: AbstractArray{Bool, 2}) =
    vertices(array) - segments(array) + volumes(array)

euler_characteristic(array :: AbstractArray{Bool, 3}) =
    vertices(array) - segments(array) + faces(array) - volumes(array)


function nthroot(n::Integer, r::Real)
    r < 0 || n == 0 && throw(DomainError())
    n < 0 && return 1 / nthroot(-n, r)
    r > 0 || return 0
    x = r / n
    prevdx = r
    while true
        y = x ^ (n - 1)
        dx = (r - y * x) / (n * y)
        abs(dx) ≥ abs(prevdx) && return x
        x += dx
        prevdx = dx
    end
end

@show nthroot.(-5:2:5, 5.0)
@show nthroot.(-5:2:5, 5.0) - 5.0 .^ (1 ./ (-5:2:5))


# Test using the Boston housing and neighborhood dataset
@testset "Bostonhsg" begin

    boston = sdataset("Bostonhsg")

    # ----------------
    #  Distance  with 3
    # ----------------
    Wdist3 = dnearneigh(boston, threshold = 3)

    @test nobs(Wdist3) == 506
    @test mean(Wdist3) ≈ 44.711462 atol = 1e-5
    @test median(Wdist3) ≈ 29.0
    @test minimum(Wdist3) == 0
    @test maximum(Wdist3) == 122
    @test nislands(Wdist3) == 4
    @test islands(Wdist3) == [55, 65, 286, 287]

    # Moran's I
    mboston = moran(boston.MEDV, Wdist3, permutations = 0)

    # ----------------
    #  Distance  with 4
    # ----------------
    Wdist4 = dnearneigh(boston, threshold = 4)

    @test nobs(Wdist4) == 506
    @test mean(Wdist4) ≈ 72.225296 atol = 1e-5
    @test median(Wdist4) ≈ 56.0
    @test minimum(Wdist4) == 1
    @test maximum(Wdist4) == 178
    @test nislands(Wdist4) == 0
    @test islands(Wdist4) == []

    # Moran's I with 999 permutations
    mboston = moran(boston.MEDV, Wdist4, rng = StableRNG(1234567), permutations = 999)

    @test score(mboston) ≈ 0.267418 atol = 1e-5
    @test pvalue(mboston) ≈ 0.001 atol = 1e-5
    @test std(mboston) ≈ 0.013811 atol = 1e-5
    @test zscore(mboston) ≈ 19.493602 atol = 1e-5

    # ----------------
    #  KNN with k = 10
    # ----------------
    Wknn10 = knearneigh(boston, k = 10)

    @test nobs(Wknn10) == 506
    @test mean(Wknn10) == 10
    @test median(Wknn10) ≈ 10.0
    @test minimum(Wknn10) == 10
    @test maximum(Wknn10) == 10
    @test nislands(Wknn10) == 0

    # Moran I with 999 permutations
    mboston = moran(boston.MEDV, Wknn10, rng = StableRNG(1234567), permutations = 999)

    @test score(mboston) ≈ 0.530743 atol = 1e-5
    @test std(mboston) ≈ 0.018956 atol = 1e-5
    @test zscore(mboston) ≈ 28.080397 atol = 1e-5
    @test pvalue(mboston) ≈ 0.001 atol = 1e-5

    # Moran's I with 9999 permutations
    mboston = moran(boston.MEDV, Wknn10, rng = StableRNG(1234567), permutations = 9999)

    @test score(mboston) ≈ 0.530743 atol = 1e-5
    @test std(mboston) ≈ 0.018541 atol = 1e-5
    @test zscore(mboston) ≈ 28.727399 atol = 1e-5
    @test pvalue(mboston) ≈ 0.0001 atol = 1e-5

    # Geary's c I with 9999 permutations
    cboston = geary(boston.MEDV, Wknn10, rng = StableRNG(1234567), permutations = 9999)

    @test score(cboston) ≈ 0.432582 atol = 1e-5
    @test std(cboston) ≈ 0.022467 atol = 1e-5
    @test zscore(cboston) ≈ -25.249835 atol = 1e-5
    @test pvalue(cboston) ≈ 0.0001 atol = 1e-5

    # Local Moran
    lmboston = localmoran(boston.MEDV, Wknn10, rng = StableRNG(1234567), permutations = 9999)

    # ----------------
    #  Plot Recipes
    # ----------------

    # Test Moran's I distribution
    @test_nowarn RecipesBase.apply_recipe(Dict{Symbol, Any}(), mboston)

    # Test Moran's Plot
    @test_nowarn RecipesBase.apply_recipe(Dict{Symbol, Any}(), boston.MEDV, Wknn10)
    @test_nowarn RecipesBase.apply_recipe(Dict{Symbol, Any}([(:standardize, false)]), boston.MEDV, Wknn10)

    # Test LISA Cluster Map with points
    @test_nowarn RecipesBase.apply_recipe(Dict{Symbol, Any}(), boston, lmboston)
    @test_nowarn RecipesBase.apply_recipe(Dict{Symbol, Any}(), boston.geometry, lmboston)

end


export chebygrid, ischebygrid, chebycoef, cheby, chebyderiv
export ChebyshevInterpolator, BichebyshevInterpolator

ξ2x(ξ, a, b) = (ξ + 1)*((b - a)/2) + a

x2ξ(x, a, b) = 2*(x - a)/(b - a) - 1

x2θ(x, a, b) = acos(x2ξ(x, a, b))

"""
    chebygrid(n)

Create an array of `n` chebyshev nodes in [-1,1]
"""
@memoize chebygrid(n::Int) = cos.(π*(n-1:-1:0)/(n-1))

"""
    chebygrid(xa, xb, n)

Create an array of `n` chebyshev nodes in [`xa`,`xb`]
"""
chebygrid(xa, xb, n::Int) = ξ2x.(chebygrid(n), xa, xb)

"""
    chebygrid(xa, xb, nx, ya, yb, ny)

Create a two-dimensional grid of chebyshev nodes using `nx` points along the first axis, in [`xa`,`xb`], and `ny` points along the second axis, in [`ya`,`yb`].
"""
function chebygrid(xa, xb, nx::Int, ya, yb, ny::Int)
    X = chebygrid(xa, xb, nx) .* ones(ny)'
    Y = chebygrid(ya, yb, ny)' .* ones(nx)
    return X, Y
end

function ischebygrid(x)::Bool
    n = length(x)
    c = chebygrid(n)
    xa, xb = minimum(x), maximum(x)
    @inbounds for i ∈ eachindex(x)
        ξ = x2ξ(x[i], xa, xb)
        #shift from [-1,1] to [1,3] to avoid issues comparing to zero
        !isapprox(ξ+2, c[i]+2) && return false
    end
    return true
end

@memoize function chebymatrix(n::Int)
    @assert n > 1 "can't construct cheby matrix smaller than 2 x 2"
    A = zeros(n,n)
    ξ = chebygrid(n)
    for j ∈ 1:n, k ∈ 1:n
        @inbounds A[j,k] = cos((k-1)*acos(ξ[j]))
    end
    return A
end

function chebyrecurrance!(T, ξ, L::Int)::Nothing
    Tₖ₋₂ = one(ξ)
    Tₖ₋₁ = ξ
    @inbounds T[2] = Tₖ₋₁
    for k = 3:L
        #compute next value
        Tₖ = 2ξ*Tₖ₋₁ - Tₖ₋₂
        #set array value
        @inbounds T[k] = Tₖ
        #swaps
        Tₖ₋₂ = Tₖ₋₁
        Tₖ₋₁ = Tₖ
    end
    nothing
end

function chebyrecurrance(ξ::U, L::Int) where {U}
    T = Vector{U}(undef, L)
    @inbounds T[1] = one(ξ)
    chebyrecurrance!(T, ξ, L)
    return T
end

"""
    chebycoef(y)

Compute the Chebyshev expansion coefficients for a set of points `y`, which are assumed to be located on the Chebyshev points for some interval.
"""
function chebycoef(y)
    n = length(y)
    @assert n > 1 "must have at least 2 points to form chebyshev coefficients"
    Aᴵ = invertedchebymatrix(n)
    return Aᴵ*y
end

"""
    cheby(coef, x, xa, xb)

Evaluates the Chebyshev expansion represented by the coefficients in `coef` and defined on the interval [`xa`,`xb`] at the point `x`.
"""
function cheby(coef, x, xa, xb)
    N = length(coef)
    @assert xa <= x <= xb "x must be inside [xa,xb]"
    #get coordinate in ξ space
    ξ = x2ξ(x, xa, xb)
    #first two elements of cheby recursion
    Tₖ₋₂ = one(ξ)
    Tₖ₋₁ = ξ
    #first two terms of dot product
    @inbounds y = Tₖ₋₂*coef[1] + Tₖ₋₁*coef[2]
    #cheby recursion and rest of terms in dot product, all at once
    for k = 3:N
        #next value in recursion
        Tₖ = 2*ξ*Tₖ₋₁ - Tₖ₋₂
        #next term in dot product
        @inbounds y += Tₖ*coef[k]
        #swaps
        Tₖ₋₂ = Tₖ₋₁
        Tₖ₋₁ = Tₖ
    end
    return y
end

"""
    chebyderiv(coef, xa, xb)

Generates the expansion coefficents for the derivative of a preexisting Chebyshev expansion defined on the interval [`xa`,`xb`].
"""
function chebyderiv(coef, xa, xb)
    n = length(coef)
    T = eltype(coef)
    @assert xa < xb "xa must be less than xb"
    #length 1 case is a trivial zero slope
    n == 1 && return zeros(T, 1)
    #start the array and check for length 2 case
    d = zeros(T, n-1)
    d[n-1] = 2*(n-2)*coef[n]
    n == 2 && return d
    #recurrance for all the rest
    d[n-2] = 2*(n-3)*coef[n-1]
    @inbounds for k ∈ n-2:-1:2
        d[k-1] = d[k+1] + 2*(k-1)*coef[k]
    end
    #interval width factor
    d .*= 2/(xb - xa)
    #extra DC scaling
    d[1] /= 2
    return d
end

#-------------------------------------------------------------------------------
# caching function for inverted cheby matrices, needed for interpolator setup

@memoize invertedchebymatrix(n::Int64)::Matrix{Float64} = inv(chebymatrix(n))

#-------------------------------------------------------------------------------
# one-dimensional interpolation

struct ChebyshevInterpolator{N,T}
    #lowest value in range
    xa::T
    #highest value in range
    xb::T
    #interpolation coefficents
    a::NTuple{N,T}
    #must always have strict boundaries
    boundaries::StrictBoundaries
end

function Base.show(io::IO, ϕ::ChebyshevInterpolator)
    print(io, "$(typeof(ϕ)) ∈ [$(ϕ.xa), $(ϕ.xb)]")
end

"""
    ChebyshevInterpolator(x, y)

Construct a `ChebyshevInterpolator` for the points defined by coordinates `x` and values `y`. The `x` coordinates *must* be arranged on a chebyshev grid, which can be generated using the [`chebygrid`](@ref) function.
"""
function ChebyshevInterpolator(x, y)
    #same types
    T = promote_type(eltype(x), eltype(y))
    x = collect(T, x)
    y = collect(T, y)
    #check for basic issues
    rangecheck(x, y, 3)
    #demand that the input points have chebyshev spacing
    @assert ischebygrid(x) "points must be on a chebyshev grid"
    #generate expansion coefficients
    a = Tuple(chebycoef(y))
    #construct
    ChebyshevInterpolator(minimum(x), maximum(x), a, StrictBoundaries())
end

"""
    ChebyshevInterpolator(f, xa, xb, n)

Construct a `ChebyshevInterpolator` for the function `f` using `n` function evaluations in the range [`xa`,`xb`]. The function evaluations will occur on the chebyshev nodes.
"""
function ChebyshevInterpolator(f::F, xa, xb, n::Int) where {F}
    #set up the range coordinates
    x = chebygrid(xa, xb, n)
    #evaluate the function at those coordinates
    y = f.(x)
    #call the other constructor
    ChebyshevInterpolator(x, y)
end

function (ϕ::ChebyshevInterpolator{N,U})(x) where {N,U}
    #always enforce boundaries
    ϕ.boundaries(x, ϕ.xa, ϕ.xb)
    #evaluate the expansion
    return cheby(ϕ.a, x, ϕ.xa, ϕ.xb)
end

"""
    chebyderiv(ϕ::ChebyshevInterpolator)

Construct a ChebyshevInterpolator representing the derivative of a preexisting interpolator.
"""
function chebyderiv(ϕ::ChebyshevInterpolator)
    #derivative expansion's coefficients
    coef = Tuple(chebyderiv(ϕ.a, ϕ.xa, ϕ.xb))
    #new interpolator
    ChebyshevInterpolator(ϕ.xa, ϕ.xb, coef, StrictBoundaries())
end

#-------------------------------------------------------------------------------
# bichebyshev interpolation, a little trickier now!

struct BichebyshevInterpolator{M,N,U}
    #lowest value on axis 1
    xa::U
    #highest value on axis 1
    xb::U
    #lowest value on axis 2
    ya::U
    #highest value on axis 2
    yb::U
    #matrix and vectors for doing the interpolation
    A::Matrix{U} # size (ny by nx) or (M by N)
    a::Vector{U} # length ny for cosine expansion in θy
    b::Vector{U} # length nx for cosine expansion in θx
    c::Vector{U} # length ny for doing M*b in place
    #must always use strict boundaries
    boundaries::StrictBoundaries
end

function Base.show(io::IO, ϕ::BichebyshevInterpolator)
    print(io, "$(typeof(ϕ)) ∈ [$(ϕ.xa), $(ϕ.xb)], [$(ϕ.ya), $(ϕ.yb)]")
end

"""
    BichebyshevInterpolator(x, y, Z)

Construct a `BichebyshevInterpolator` for the grid of points defined by coordinates (`x`,`y`) and values `Z`. The given points must lie on a chebyshev grid in each direction. These can be generated with the [`chebygrid`](@ref) function or the interpolator can be constructed directly from a function using the method below.

!!! warning

    The Bichebyshev interpolator is *not thread-safe*. It computes a cosine expansion and does some linear algebra in-place using arrays stored with the object. A single `BichebyshevInterpolator` should never be called by multiple threads at once.
"""
function BichebyshevInterpolator(x, y, Z)
    #same types
    T = promote_type(eltype(x), eltype(y), eltype(Z))
    x = collect(T, x)
    y = collect(T, y)
    Z = collect(T, Z)
    #check for basic grid problems
    gridcheck(x, y, Z, 3)
    #grid properties
    nx, ny = length(x), length(y)
    xa, xb = T(minimum(x)), T(maximum(x))
    ya, yb = T(minimum(y)), T(maximum(y))
    #reject any non-cheby grid spacing
    @assert ischebygrid(x) "axis 1 coordinates must be on a chebyshev grid"
    @assert ischebygrid(y) "axis 2 coordinates must be on a chebyshev grid"
    #get inverted matrices from cache or generate them
    B = invertedchebymatrix(nx)
    #generate interpolation coefficients along axis 1 for each value of axis 2
    α = zeros(T, nx, ny)
    for j = 1:ny
        mul!(view(α,:,j), B, view(Z,:,j))
    end
    #then combine α and A
    A = invertedchebymatrix(ny)*α'
    #other vectors we need for doing the actual interpolation
    a = ones(T, ny)
    b = ones(T, nx)
    c = zeros(T, ny)
    #done
    BichebyshevInterpolator{ny,nx,T}(xa, xb, ya, yb, A, a, b, c, StrictBoundaries())
end

"""
    BichebyshevInterpolator(f, xa, xb, nx, ya, yb, ny)

Construct a `BichebyshevInterpolator` for the function `f` using a grid of `nx` points on the first axis in [`xa`,`xb`] and `ny` points on the second axis in [`ya`,`yb`].
"""
function BichebyshevInterpolator(f::F, xa, xb, nx::Int, ya, yb, ny::Int) where {F}
    #set up the grid
    X, Y = chebygrid(xa, xb, nx, ya, yb, ny)
    #evaluate the function at chebyshev grid points
    Z = f.(X, Y)
    #call the other constructor
    BichebyshevInterpolator(X[:,1], Y[1,:], Z)
end

#=====
This is the fast implementation. It's executed when the types of the
input coordinates match the type of the stored coefficients in the
interpolator. When the types match, the Chebyshev expansions can be
evaluated in-place, using vectors pre-allocated in the interpolator.
See the a, b, and c fields of the struct. This method also guarantees
that the interpolator type and the coordinate types are <: AbstractFloat,
making the low-level linear algebra functions safe. Without such a
guarantee, there can be issues with, for example, the dual numbers in
FowardDiff routines.
=====#
function (Φ::BichebyshevInterpolator{M,N,U})(x::U, y::U) where {M,N,U<:AbstractFloat}
    #always enforce boundaries for Chebyshev
    Φ.boundaries(x, Φ.xa, Φ.xb, y, Φ.ya, Φ.yb)
    #evaluate Chebyshev polys at the coordinates recursively and in-place
    ξy = x2ξ(y, Φ.ya, Φ.yb)
    chebyrecurrance!(Φ.a, ξy, M)
    ξx = x2ξ(x, Φ.xa, Φ.xb)
    chebyrecurrance!(Φ.b, ξx, N)
    #perform M*b, which interpolates along the first axis, also in-place
    mul!(Φ.c, Φ.A, Φ.b)
    #then a'*c interpolates along the second axis
    return dot(Φ.a, Φ.c)
end

#=====
This is the slow, but type-flexible, implementation. It's
executed whenever the type of the interpolator's coefficients or
the coordinate types don't match OR are not AbstractFloats. The
price of this flexibility is allocations for the expansions and
loss of the low-level linear algebra routines 😭.
=====#
function (Φ::BichebyshevInterpolator{M,N,U})(x, y) where {M,N,U}
    #always enforce boundaries for Chebyshev
    Φ.boundaries(x, Φ.xa, Φ.xb, y, Φ.ya, Φ.yb)
    #coordinates in ξ space
    ξx, ξy = promote(x2ξ(x, Φ.xa, Φ.xb), x2ξ(y, Φ.ya, Φ.yb))
    #allocating expansion in the x direction
    b = chebyrecurrance(ξx, N)
    c = Φ.A * b
    #then perform the y-axis expansion and dot product simulataneously
    Tₖ₋₂ = one(ξy)
    Tₖ₋₁ = ξy
    @inbounds z = c[1] + ξy*c[2]
    for i = 3:M
        #compute next value
        Tₖ = 2ξy*Tₖ₋₁ - Tₖ₋₂
        #update running dot product
        @inbounds z += c[i]*Tₖ
        #swaps
        Tₖ₋₂ = Tₖ₋₁
        Tₖ₋₁ = Tₖ
    end
    return z
end


@testset "softmax" begin
    X = rand(100,5)
    max_ = maximum(X, dims=2)
    exp_ = exp.(X .- max_)
    @test isapprox(MLJMultivariateStatsInterface.softmax(X), exp_ ./ sum(exp_, dims=2))
    X_ = MLJMultivariateStatsInterface.softmax!(X)
    @test X_ === X
    @test isapprox(X_, exp_ ./ sum(exp_, dims=2))
end


using MCObservables

function test_simpleObservable(numtest, MCS)
  inform("Start testing SimpleObservable")

  nsigma = 3

  counts = zeros(Int, nsigma)

  for i in 1:numtest
    obs = SimpleObservable()
    for j in 1:MCS
      obs << randn()
    end
    m = abs(mean(obs))
    er = stderror(obs)
    for isigma in 1:nsigma
      if m < isigma*er
        counts[isigma] += 1
      end
    end
  end
  sigma = counts ./ numtest
  @show sigma

  @test 0.68 < sigma[1] < 0.69
  @test 0.95 < sigma[2] < 0.96
  @test 0.99 < sigma[3]
  inform("Success testing SimpleObservable")
end

result_simpleObservable = test_simpleObservable(100_000, 8192)


using Test
using Mods

p = 23
@test Mod(2,p) == Mod(25,p)
@test Mod(2,p)+Mod(21,p) == Mod(0,p)
@test Mod(10,p)*Mod(5,p) == Mod(4,p)
@test Mod(10,p)/Mod(5,p) == Mod(2,p)
@test Mod(17,p) * inv(Mod(17,p)) == Mod(1,p)
@test Mod(17^6,p)==Mod(17,p)^6
@test Mod(17,p)^(-2) == inv(Mod(17,p))^2
@test Mod(17,p) == 17

q = 91
a = Mod(17,p)
b = Mod(32,q)

x = CRT(a,b)
@test Mod(x.val,p)==a
@test Mod(x.val,q)==b
@test inv(a)*a == 1

p = 9223372036854775783   # This is a big prime
x = Mod(-2,p)
@test x*x == 4
@test x+x == -4
@test x/x == 1
@test x^4 == 16
@test x^(p-1) == 1   # Fermat Little Theorem test
@test 2x == x+x
@test x-x == 0
y = inv(x)
@test x*y == 1
@test x+p == x
@test x*p == 0
@test p-x == x-2x


@testset "Matrix multiplication" begin

    # test default settings
    @test get_multiplication_mode() == Dict(:multiplication => :fast)

    A = [2..4 -2..1; -1..2 2..4]
    set_multiplication_mode(:slow)
    @test A * A == [0..18 -16..8; -8..16 0..18]

    set_multiplication_mode(:rank1)
    @test A * A == [0..18 -16..8; -8..16 0..18]

    set_multiplication_mode(:fast)
    @test A * A == [-2..19.5 -16..10; -10..16 -2..19.5]
    @test A * mid.(A) == [5..12.5 -8..2; -2..8 5..12.5]
    @test mid.(A) * A == [5..12.5 -8..2; -2..8 5..12.5]
end


function lgd20()
#
#  test polynomial suggested by Goedecker
#   Legendre polynomial of degree 20
#
    p = lgd(20);
    z = [-.9931285991850949; -.9639719272779138; -.9122344282513259; -.8391169718222188;
         -.7463319064601508; -.6360536807265150; -.5108670019508271; -.3737060887154196;
         -.2277858511416451; -.7652652113349733e-1; .7652652113349733e-1; .2277858511416451;
         .3737060887154196; .5108670019508271; .6360536807265150; .7463319064601508;
         .8391169718222188; .9122344282513259; .9639719272779138; .9931285991850949]
    z = [z ones(20)]

    p, PolyZeros(z)
end


dt = 0.005
function lorenz63(u0, s, n)
    sigma, rho, beta = s
    d, m = size(u0)
    n = n+1
    u_trj = zeros((m,d,n))
    u_trj[:,:,1] = u0'
    for i = 2:n
    	x = u_trj[:,1,i-1]
    	y = u_trj[:,2,i-1]
    	z = u_trj[:,3,i-1]

    	u_trj[:,1,i] = x + dt*(sigma.*(y - x))
    	u_trj[:,2,i] = y + dt*(x.*(rho .- z) - y)
    	u_trj[:,3,i] = z + dt*(x.*y - beta.*z)
    end
    return permutedims(u_trj,[3,2,1])
end
function dlorenz63(u, s)
    sigma, rho, beta = s
    n, d = size(u)
    x = view(u,:,1)
    y = view(u,:,2)
    z = view(u,:,3)
    du = zeros(n, d, d)
    @. du[:,1,1] = 1.0 - dt*sigma
    @. du[:,1,2] = dt*sigma
    @. du[:,2,1] = dt*(rho - z)
    @. du[:,2,2] = 1.0 - dt
    @. du[:,2,3] = -dt*x
    @. du[:,3,1] = dt*y
    @. du[:,3,2] = dt*x
    @. du[:,3,3] = 1.0 - dt*beta
    return reshape([du[:,:,1]'; du[:,:,2]';
    				du[:,:,3]'], d, d, n)
end
function perturbation(u,s)
    n, d = size(u)
    # the perturbation in row i in T_{u_(i+1)} M
    return [zeros(1,n); dt*u[:,1]'; zeros(1,n)]
end
function vectorField(u,s)
    n, d = size(u)
    sigma, rho, beta = s
    u = u'
    x, y, z = u[1,:], u[2,:], u[3,:]
    return [sigma.*(y - x)  x.*(rho .- z) - y  x.*y - beta.*z]'
end
function lorenz63_rhs_ad(du, u, s, t)
    du[1] = s[1]*(u[2] - u[1])
	du[2] = u[1]*(s[2] - u[3]) - u[2]
	du[3] = u[1]*u[2] - s[3]*u[3]
end
function lorenz63_ad(u0, s, n)
	t = n*dt
    prob = ODEProblem(lorenz63_rhs_ad, u0, (0.,t), s)
	sol = Array(solve(prob, Tsit5(), saveat=dt))
	return sol[:,end]
end
function obj_fun(u0, s)
    prob = ODEProblem(lorenz63_rhs_ad, u0, (0.,1.1), s)
	#_prob = remake(prob,u0=u0,p=s)
	sol = solve(prob, Tsit5(), saveat=0.005)
	sum(sol[3,:])/size(sol,2)
end


# Parts of this code were taken / derived from Graphs.jl. See LICENSE for
# licensing details.
"""
    struct FloydWarshallState{T, U}

An [`AbstractPathState`](@ref) designed for Floyd-Warshall shortest-paths calculations.
"""
struct FloydWarshallState{T,U<:Integer} <: AbstractPathState
    dists::Matrix{T}
    parents::Matrix{U}
end

@doc """
    floyd_warshall_shortest_paths(g, distmx=weights(g))

Use the [Floyd-Warshall algorithm](http://en.wikipedia.org/wiki/Floyd–Warshall_algorithm)
to compute the shortest paths between all pairs of vertices in graph `g` using an
optional distance matrix `distmx`. Return a [`LightGraphs.FloydWarshallState`](@ref) with relevant
traversal information.

### Performance
Space complexity is on the order of ``\\mathcal{O}(|V|^2)``.
"""
function floyd_warshall_shortest_paths(
    g::AbstractGraph{U},
    distmx::AbstractMatrix{T}=weights(g)
) where T<:Real where U<:Integer
    nvg = nv(g)
    # if we do checkbounds here, we can use @inbounds later
    checkbounds(distmx, Base.OneTo(nvg), Base.OneTo(nvg))

    dists = fill(typemax(T), (Int(nvg), Int(nvg)))
    parents = zeros(U, (Int(nvg), Int(nvg)))

    @inbounds for v in vertices(g)
        dists[v, v] = zero(T)
    end
    undirected = !is_directed(g)
    @inbounds for e in edges(g)
        u = src(e)
        v = dst(e)

        d = distmx[u, v]

        dists[u, v] = min(d, dists[u, v])
        parents[u, v] = u
        if undirected
            dists[v, u] = min(d, dists[v, u])
            parents[v, u] = v
        end
    end

    @inbounds for pivot in vertices(g)
        # Relax dists[u, v] = min(dists[u, v], dists[u, pivot]+dists[pivot, v]) for all u, v
        for v in vertices(g)
            d = dists[pivot, v]
            d == typemax(T) && continue
            p = parents[pivot, v]
            for u in vertices(g)
                ans = (dists[u, pivot] == typemax(T) ? typemax(T) : dists[u, pivot] + d)
                if dists[u, v] > ans
                    dists[u, v] = ans
                    parents[u, v] = p
                end
            end
        end
    end
    fws = FloydWarshallState(dists, parents)
    return fws
end

function enumerate_paths(s::FloydWarshallState{T,U}, v::Integer) where T where U<:Integer
    pathinfo = s.parents[v, :]
    paths = Vector{Vector{U}}()
    for i in 1:length(pathinfo)
        if (i == v) || (s.dists[v, i] == typemax(T))
            push!(paths, Vector{U}())
        else
            path = Vector{U}()
            currpathindex = i
            while currpathindex != 0
                push!(path, currpathindex)
                if pathinfo[currpathindex] == currpathindex
                    currpathindex = 0
                else
                    currpathindex = pathinfo[currpathindex]
                end
            end
            push!(paths, reverse(path))
        end
    end
    return paths
end

enumerate_paths(s::FloydWarshallState) = [enumerate_paths(s, v) for v in 1:size(s.parents, 1)]
enumerate_paths(st::FloydWarshallState, s::Integer, d::Integer) = enumerate_paths(st, s)[d]


function run_gradient_descent_iterations(func, grad, x0, lr, numiter; rtol = 1e5)
    xprev = copy(x0)
    xnext = copy(x0)

    g = grad(xprev)
    gnorm = g'*g
    iter = 1
    tol = rtol*eps()

    while gnorm > tol && iter < numiter
        f = func(xprev)
        g = grad(xprev)
        gnorm = g'*g

        xnext .= xprev - lr * (f / gnorm) * g
        xprev .= xnext

        iter += 1
    end
    return xnext
end


using Plots
pyplot()
using QuantEcon

P =[0.971 0.029 0.000
    0.145 0.778 0.077
    0.000 0.508 0.492]

psi = [0.0 0.2 0.8]

t = 20
x_vals = Array(Float64, t+1)
y_vals = Array(Float64, t+1)
z_vals = Array(Float64, t+1)
colors = []

for i=1:t
    x_vals[i] = psi[1]
    y_vals[i] = psi[2]
    z_vals[i] = psi[3]
    psi = psi*P
    push!(colors, :red)
end
push!(colors, :black)

mc = MarkovChain(P)
psi_star = stationary_distributions(mc)[1]
x_vals[t+1] = psi_star[1]
y_vals[t+1] = psi_star[2]
z_vals[t+1] = psi_star[3]
scatter(x_vals, y_vals, z_vals, color=colors)
plot!(lims=(0, 1), ticks=[0.25 0.5 0.75]', legend=:none)


# TODO: combine into one function to avoid code repetition
# pass integration 'method' parameter, enum?

function left_riemann(f::Function, a::Real, b::Real, n::Int)

    delta = (b - a) / n # Calculate delta between samples
    X = a .+(0:n-1) * delta # Get range of samples
    fX = map(f, X) # Calculate f(x) for each sample in range
    return sum(fX) * delta # Sum and divide by sample width

end

function right_riemann(f::Function, a::Real, b::Real, n::Int)

    delta = (b - a) / n
    X = a .+(1:n) * delta
    fX = map(f, X)
    return sum(fX) * delta

end

function trapezoid(f::Function, a::Real, b::Real, n::Int)

    delta = (b - a) / n
    X = a .+(0:n) * delta
    fX = map(f, X)

    # must calculate mean height of left and right
    return sum(fX[2:end] + diff(fX) * 0.5) * delta

end


#=

# 2D Lid-driven cavity (AD-Newton)
([source code](SOURCE_URL))

This example solves the lid-driven cavity problem where one seeks
a velocity ``\mathbf{u}`` and pressure ``\mathbf{p}`` of the incompressible Navier--Stokes problem
```math
\begin{aligned}
- \mu \Delta \mathbf{u} + (\mathbf{u} \cdot \nabla) \mathbf{u} + \nabla p & = 0\\
\mathrm{div}(u) & = 0
\end{aligned}
```
where ``\mathbf{u} = (1,0)`` along the top boundary of a square domain.

This examples highlights the use of automatic differentation to obtain Newton derivatives of nonlinear PDEOperators. The user can
switch between a Newton obtained by automatic differentation (ADnewton = true) and a 'manual' Newton scheme (ADnewton = false). The runtime
of the manual newton is slightly faster, but requires much more code input from the user that is more prone to errors.

=#

using PDELib
using Printf

function fe_liddrivencavity_autonewton(; verbosity = 0, Plotter = nothing, ADnewton = true, nref=2)

    function boundary_data_top!(result)
        result[1] = 1;
        result[2] = 0;
    end

    ## grid
    xgrid = uniform_refine(grid_unitsquare(Triangle2D), nref);

    ## problem parameters
    viscosity = 1e-2
    maxIterations = 50  # termination criterion 1 for nonlinear mode
    maxResidual = 1e-12 # termination criterion 2 for nonlinear mode
    broken_p = false

    ## choose one of these (inf-sup stable) finite element type pairs
    #FETypes = [H1P2{2,2}, H1P1{1}] # Taylor--Hood
    #FETypes = [H1P2B{2,2}, H1P1{1}]; broken_p = true # P2-bubble
    #FETypes = [H1CR{2}, H1P0{1}] # Crouzeix--Raviart
    #FETypes = [H1MINI{2,2}, H1P1{1}] # MINI element on triangles only
    FETypes = [H1BR{2}, H1P0{1}]; broken_p = true # Bernardi--Raugel

    #####################################################################################
    #####################################################################################

    ## negotiate data functions to the package
    user_function_bnd = DataFunction(boundary_data_top!, [2,2]; name = "u_bnd", dependencies = "", quadorder = 0)

    ## load linear Stokes problem prototype and assign data
    ## we are adding the nonlinar convection term ourself below
    ## to discuss the details
    StokesProblem = IncompressibleNavierStokesProblem(2; viscosity = viscosity, nonlinear = false)
    add_boundarydata!(StokesProblem, 1, [1,2,4], HomogeneousDirichletBoundary)
    add_boundarydata!(StokesProblem, 1, [3], BestapproxDirichletBoundary; data = user_function_bnd)

    ## store matrix of Laplace operator for nonlinear solver
    StokesProblem.LHSOperators[1,1][1].store_operator = true

    ## add Newton for convection term
    if ADnewton
        ## AUTOMATIC DIFFERENTATION
        ## requries kernel function for NonlinearForm action (u * grad) u
        function ugradu_kernel_AD(result, input)
            ## input = [u, grad(u)]
            ## compute (u * grad) u = grad(u)*u
            for j = 1 : 2
                result[j] = 0.0
                for k = 1 : 2
                    result[j] += input[k]*input[2 + (j-1)*2+k]
                end
            end
            return nothing
        end
        action_kernel = ActionKernel(ugradu_kernel_AD, [2,6]; dependencies = "", quadorder = 1)

        ## generate and add nonlinear PDEOperator (modifications to RHS are taken care of automatically)
        NLConvectionOperator = GenerateNonlinearForm("(u * grad) u  * v", [Identity, Gradient], [1,1], Identity, action_kernel; ADnewton = true)
        add_operator!(StokesProblem, [1,1], NLConvectionOperator)
    else
        ## MANUAL DIFFERENTATION
        ## uses the following kernel function for the linearised convection operator
        function ugradu_kernel_nonAD(result, input_current, input_ansatz)
            ## input_current = [current, grad(current)]
            ## input_ansatz = [ansatz, grad(ansatz)]
            ## compute (current * grad) ansatz + (current * grad) ansatz
            for j = 1 : 2
                result[j] = 0.0
                for k = 1 : 2
                    result[j] += input_current[k]*input_ansatz[2 + (j-1)*2+k]
                    result[j] += input_ansatz[k]*input_current[2 + (j-1)*2+k]
                end
            end
            return nothing
        end
        ## and a similar kernel function for the right-hand side
        function ugradu_kernel_rhs(result, input_current)
            ## input_current = [current, grad(current)]
            ## input_ansatz = [ansatz, grad(ansatz)]
            ## compute (current * grad) current
            for j = 1 : 2
                result[j] = 0.0
                for k = 1 : 2
                    result[j] += input_current[k]*input_current[2 + (j-1)*2+k]
                end
            end
            return nothing
        end
        newton_action_kernel = NLActionKernel(ugradu_kernel_nonAD, [2,6]; dependencies = "", quadorder = 1)
        action_kernel_rhs = ActionKernel(ugradu_kernel_rhs, [2,6]; dependencies = "", quadorder = 1)

        ## generate and add nonlinear PDEOperator (modifications to RHS are taken care of by optional arguments)
        NLConvectionOperator = GenerateNonlinearForm("(u * grad) u  * v", [Identity, Gradient], [1,1], Identity, newton_action_kernel; ADnewton = false, action_kernel_rhs = action_kernel_rhs)
        add_operator!(StokesProblem, [1,1], NLConvectionOperator)
    end

    ## generate FESpaces
    FESpaceVelocity = FESpace{FETypes[1]}(xgrid)
    FESpacePressure = FESpace{FETypes[2]}(xgrid; broken = broken_p)
    Solution = FEVector{Float64}("Stokes velocity",FESpaceVelocity)
    append!(Solution,"Stokes pressure",FESpacePressure)

    ## show configuration and solve Stokes problem
    Base.show(StokesProblem)
    GradientRobustMultiPhysics.solve!(Solution, StokesProblem; verbosity = verbosity, maxIterations = maxIterations, maxResidual = maxResidual)

    ## plot
    GradientRobustMultiPhysics.plot(Solution, [1,2], [Identity, Identity]; Plotter = Plotter, verbosity = verbosity, use_subplots = true)
    true
end


#=
Created on 09/01/2021 12:54:49
Last update: -

@author: Michiel Stock
michielfmstock@gmail.com

Estimate pi through Monte Carlo sampling.
Do this by simuting to throw `n` pebbles in the [-1, 1] x [-1, 1] square
and track the fraction that land in the unit circle.
=#

function estimatepi(n)
    hits = 0
    for i in 1:n
        x = rand()
        y = rand()
        if x^2 + y^2 ≤ 1.0
            hits += 1
        end
    end
    return 4hits/n
end

estimatepi2(n) = 4count(x-> x ≤ 1.0, sum(rand(n, 2).^2, dims=2)) / n


using PlasmaDispersionFunctions
using ArbNumerics, QuadGK, Random, SpecialFunctions, Test

import SpecialFunctions.erfcx
erfcx(x::Complex{BigFloat}) = erfcx(ArbComplex(x))

@testset "Tests" begin
Random.seed!(0)

@testset "Plasma dispersion function" begin
  @test plasma_dispersion_function(0.0) ≈ im*sqrt(pi) rtol=1e-5
  @test plasma_dispersion_function(im) ≈ im*0.757872156141312 rtol=1e-5
  @test plasma_dispersion_function(ComplexF64(-1.52, 0.47), 0) ≈
    ComplexF64(0.6088888957234254, 0.33494583882874024) rtol=1e-5
  @test plasma_dispersion_function(big(0.0+0im)) ≈
    ArbComplex(im*sqrt(pi)) rtol=1e-5
end

@testset "using Z₋₁" begin
  Z0 = plasma_dispersion_function(0.0, 0)
  Z1 = plasma_dispersion_function(0.0, 1)
  @test Z1 == plasma_dispersion_function(0.0, 1, Z0)
  Z2 = plasma_dispersion_function(0.0, 2)
  @test Z2 == plasma_dispersion_function(0.0, 2, Z1)
end

@testset "vs quadrature" begin
  function test_quad(Z, pow)
    for z in Z
      principal(x) = x.^pow .* exp(-x.^2)/sqrt(pi)
      integrand(x) = principal(x) / (x-z)
      principalpart = if iszero(imag(z))
        folded(v) = (principal(v + real(z)) - principal(-v + real(z))) / v
        QuadGK.quadgk(folded, nextfloat(0.0), Inf, rtol=eps())[1]
      else
        QuadGK.quadgk(integrand, -Inf, Inf, rtol=eps())[1]
      end
      σ = imag(z) < 0 ? 2 : imag(z) == 0 ? 1 : 0
      cauchyresidue = im * (σ * π * principal(z))
      b = plasma_dispersion_function(z, pow)
      a = principalpart + cauchyresidue
      @test real(a) ≈ real(b) rtol=sqrt(eps()) atol=0.0
      @test imag(a) ≈ imag(b) rtol=sqrt(eps()) atol=0.0
    end
  end
  Zs = []
  push!(Zs,  1.0)
  push!(Zs, -1.0)
  push!(Zs,  1.0 + im)
  push!(Zs, -1.0 + im)
  push!(Zs,  1.0 - im)
  push!(Zs, -1.0 - im)
  for j in 1:10
    push!(Zs, 4 * (rand() - 0.5) + im * 4 * (rand() - 0.5))
  end
  for i in Unsigned.(0:10)
    @testset "QuadGK of maxwellian with $(i)th moment" begin
      test_quad(Zs, i)
    end
  end
end

@testset "check errors are caught" begin
  @test_throws ArgumentError plasma_dispersion_function(1.0, -1)
end

end


# function Mpz(pz_pml::SparseMatrixCSC{Float64,Int64}, lambda::Array{Float64,2}, dz::Float64, dt::Float64, ext::Int64, iflag::Int64)
#     lambda = modExpand(lambda, ext, iflag)
#     (m, n) = size(lambda)
#     a1 = 9/8;   a2 = -1/24;
#     c1 = a1/dz; c2 = a2/dz;
#     C3 = zeros(m*n)
#     denum = zeros(m*n)
#     for ix = 1 : n
#         for iz = 1 : m
#             a = 1/dt + pz_pml[iz,ix]/2
#             b = 1/dt - pz_pml[iz,ix]/2
#             denum[(ix-1)*m+iz] = lambda[iz,ix] / a
#             C3[(ix-1)*m+iz]= b / a
#         end
#     end
#     tmp = spzeros(m, m)
#     tmp[1,1] =  c1; tmp[1,2] = c2;
#     tmp[2,1] = -c1; tmp[2,2] = c1; tmp[2,3] = c2;
#     for iz = 3: m-1
#         tmp[iz,iz-1] = -c1; tmp[iz, iz  ] = c1;
#         tmp[iz,iz-2] = -c2; tmp[iz, iz+1] = c2;
#     end
#     tmp[m,m-1] = -c1; tmp[m,m] = c1; tmp[m, m-2] = -c2
#     MpzBpz = spdiagm(C3)
#     MpzBvz = spdiagm(denum) * kron(speye(n), tmp)
#     return MpzBpz, MpzBvz
# end

function Mpz(pz_pml::SparseMatrixCSC{Float64,Int64}, v::Array{Float64,2}, dz::Float64, dt::Float64, ext::Int64, iflag::Int64)
    v = modExpand(v, ext, iflag)
    (m, n) = size(pz_pml)
    a1 = 9/8;   a2 = -1/24;
    c1 = a1/dz; c2 = a2/dz;
    C3 = zeros(m*n)
    denum = zeros(m*n)
    for ix = 1 : n
        for iz = 1 : m
            a = 1/dt + pz_pml[iz,ix]/2
            b = 1/dt - pz_pml[iz,ix]/2
            denum[(ix-1)*m+iz] = (v[iz,ix])^2 / a
            C3[(ix-1)*m+iz]= b / a
        end
    end
    tmp = spzeros(m, m)
    tmp[1,1] =  c1; tmp[1,2] = c2;
    tmp[2,1] = -c1; tmp[2,2] = c1; tmp[2,3] = c2;
    for iz = 3: m-1
        tmp[iz,iz-1] = -c1; tmp[iz, iz  ] = c1;
        tmp[iz,iz-2] = -c2; tmp[iz, iz+1] = c2;
    end
    tmp[m,m-1] = -c1; tmp[m,m] = c1; tmp[m, m-2] = -c2
    MpzBpz = spdiagm(C3)
    MpzBvz = spdiagm(denum) * kron(speye(n), tmp)
    return MpzBpz, MpzBvz
end

function Mpz_back(pz_pml::SparseMatrixCSC{Float64,Int64}, v::Array{Float64,2}, dz::Float64, dt::Float64, ext::Int64, iflag::Int64)
    v = modExpand(v, ext, iflag)
    (m, n) = size(pz_pml)
    a1 = 9/8;   a2 = -1/24;
    c1 = a1/dz; c2 = a2/dz;
    C3 = zeros(m*n)
    denum = zeros(m*n)
    for ix = 1 : n
        for iz = 1 : m
            a = 1/dt - pz_pml[iz,ix]/2
            b = 1/dt + pz_pml[iz,ix]/2
            denum[(ix-1)*m+iz] = (v[iz,ix])^2 / a
            C3[(ix-1)*m+iz]= b / a
        end
    end
    tmp = spzeros(m, m)
    tmp[1,1] =  c1; tmp[1,2] = c2;
    tmp[2,1] = -c1; tmp[2,2] = c1; tmp[2,3] = c2;
    for iz = 3: m-1
        tmp[iz,iz-1] = -c1; tmp[iz, iz  ] = c1;
        tmp[iz,iz-2] = -c2; tmp[iz, iz+1] = c2;
    end
    tmp[m,m-1] = -c1; tmp[m,m] = c1; tmp[m, m-2] = -c2
    MpzBpz = spdiagm(C3)
    MpzBvz = -spdiagm(denum) * kron(speye(n), tmp)
    return MpzBpz, MpzBvz
end


using FrankWolfe
using ProgressMeter
using Arpack
using Plots
using DoubleFloats
using ReverseDiff

using LinearAlgebra

n = Int(1e2);
k = n

f(x) = dot(x, x)
function grad!(storage, x)
    @. storage = 2 * x
end

# pick feasible region
lmo = FrankWolfe.ProbabilitySimplexOracle{Rational{BigInt}}(1); # radius needs to be integer or rational

# compute some initial vertex
x0 = FrankWolfe.compute_extreme_point(lmo, zeros(n));

# benchmarking Oracles
FrankWolfe.benchmark_oracles(f, grad!, () -> rand(n), lmo; k=100)

# the algorithm runs in rational arithmetic even if the gradients and the function itself are not rational
# this is because we replace the descent direction by the directions of the LMO are rational

@time x, v, primal, dual_gap, trajectory = FrankWolfe.frank_wolfe(
    f,
    grad!,
    lmo,
    x0,
    max_iteration=k,
    line_search=FrankWolfe.Agnostic(),
    print_iter=k / 10,
    verbose=true,
    emphasis=FrankWolfe.blas,
);

println("\nOutput type of solution: ", eltype(x))

# you can even run everything in rational arithmetic using the shortstep rule
# NOTE: in this case the gradient computation has to be rational as well

@time x, v, primal, dual_gap, trajectory = FrankWolfe.frank_wolfe(
    f,
    grad!,
    lmo,
    x0,
    max_iteration=k,
    line_search=FrankWolfe.RationalShortstep(),
    L=2,
    print_iter=k / 10,
    verbose=true,
    emphasis=FrankWolfe.blas,
);

println("\nOutput type of solution: ", eltype(x))

println("\nNote: the last step where we exactly close the gap. This is not an error. ")
fract = 1 // n
println(
    "We have *exactly* computed the optimal solution with with the $fract * (1, ..., 1) vector.\n",
)
println("x = $x")

####################################################################################################################
### APPRROXIMATE CARATHEODORY WITH PLANTED SOLUTION
####################################################################################################################

rhs = 1
n = 40
k = 1e5

xpi = rand(big(1):big(100), n)
total = sum(xpi)
xp = xpi .// total

f(x) = norm(x - xp)^2
function grad!(storage, x)
    @. storage = 2 * (x - xp)
end

lmo = FrankWolfe.ProbabilitySimplexOracle{Rational{BigInt}}(rhs)
direction = rand(n)
x0 = FrankWolfe.compute_extreme_point(lmo, direction)

@time x, v, primal, dual_gap, trajectory = FrankWolfe.frank_wolfe(
    f,
    grad!,
    lmo,
    x0,
    max_iteration=k,
    line_search=FrankWolfe.Agnostic(),
    print_iter=k / 10,
    emphasis=FrankWolfe.memory,
    verbose=true,
)
println("\nOutput type of solution: ", eltype(x))

println("Computed solution: x = $x")


# The goal of this file is to provide tools that implements the MFK algorithm from these two references:
# @article{furman2017,
#   title = {On {{Log}}-{{Normal Convolutions}}: {{An Analytical}}-{{Numerical Method With Applications}} to {{Economic Capital Determination}}},
#   shorttitle = {On {{Log}}-{{Normal Convolutions}}},
#   author = {Furman, Edward and Hackmann, Daniel and Kuznetsov, Alexey},
#   year = {2017},
#   issn = {1556-5068},
#   abstract = {We put forward an efficient algorithm for approximating the sums of independent and lognormally distributed random variables. Namely, by merging tools from probability theory and numerical analysis, we are able to compute the cumulative distribution functions of the just-mentioned sums with any desired precision. Importantly, our algorithm is fast and can tackle equally well sums with just a few or thousands of summands. We illustrate the effectiveness of the new method in the contexts of the individual and collective risk models, aggregate economic capital determination, and economic capital allocation.},
#   journal = {SSRN Electronic Journal},
#   language = {en}
# }

# @article{miles2019,
#   ids = {mile,milesa},
#   title = {Risk Aggregation: {{A}} General Approach via the Class of {{Generalized Gamma Convolutions}}},
#   author = {Miles, Justin and Furman, Edward and Kuznetsov, Alexey},
#   year = {2019},
#   journal = {Variance},
#   keywords = {gamma}
# }

"""

    compute_g(dist,n, integrator)

This function will compute the n first theretical (-1)-exponentially shifted moments of a distribution,
using takashi-mori tanh-sinh exponentials formulas from the DoubleExponentialFormulas package.

This might take a long time. As an integrator, you should pass the result of `QuadDE(BigFloat)` after setting enough precision.

"""
function compute_g(dist,n, integrator)
    g = Array{BigFloat}(undef, 1,2n+1)
    residuals_g = deepcopy(g)
    for i in 0:(2n)
        g[i+1],residuals_g[i+1] = integrator(x -> (-x)^(i) * Distributions.pdf(dist,x) * exp(-x), 0, +Inf)
        print("g_{",i,"} = ",Float64(g[i+1]),", rez = ",Float64(residuals_g[i+1]),"\n")
    end
    return g
end

"""

    E_from_g(g)

Simply computes the empirical laguerre coefficients from the exponentially shifted moments,
if you computed those moments from a theoretical distribution.

"""
function E_from_g(g)
    # this function should compute laguerre coefficients from g.
    a = deepcopy(g)
    for i in 1:length(a)
        a[i] = sqrt(2) * sum(binomial(big(i-1),big(k))*big(2)^(k)/factorial(big(k)) * g[k+1] for k in 0:(i-1))
    end
    return a
end

function build_s!(s,g,facts)
    s[1] = -g[2]/g[1]
    for k in 1:(length(s)-1)
        s[k+1] = g[k+2] / facts[k+1]
        for i in 0:(k-1)
            s[k+1] += s[i+1] * g[k-i+1] / facts[k-i+1]
        end
        s[k+1] = - s[k+1]/g[1]
    end
end

function s_from_k(k)
    if isodd(length(k))
        n = Int((length(k)-1)/2)
    else
        n = Int(length(k)/2) - 1
    end
    # n = Int(round(length(k)/2))
    s = k[2:2n+1] .* (-1) .^(2:(2n+1))
    return s
end

function MFK_end(s::Vector{T}) where T
    n = Int(length(s)//2)
    S = zeros(T,(n,n))
    for i in 0:(n-1)
        for j in 0:(n-1)
            S[i+1,j+1] = s[i+j+1]
        end
    end

    sol_b = LinearAlgebra.Symmetric(S) \ (-s[(n+1):end])
    b = deepcopy(sol_b)
    append!(b,1)
    b = reverse(b)
    b_deriv = reverse(sol_b) .* (1:n)

    a = zeros(T,n)
    a[1] = s[1]
    for k in 1:(n-1)
        a[k+1] = s[k+1]
        for i in 0:(k-1)
            a[k+1] = a[k+1] + b[i+2] * s[k-i]
        end
    end

    z = real.(PolynomialRoots.roots(b, polish=true))

    beta = -z .-1
    alpha = deepcopy(beta)

    for i in 1:length(alpha)
        rez_num = 0
        rez_denom = 0
        for k in 1:length(a)
            rez_num += a[k] * z[i]^(k-1)
        end
        for k in 1:length(b_deriv)
            rez_denom += b_deriv[k] * z[i]^(k-1)
        end
        alpha[i] = rez_num/rez_denom
    end
    theta = 1 ./beta
    return alpha,theta
end

"""

    MFK_Projection(g_integrals,n_gammas)

From a set of g_integrals computed by the compute_g function, this function runs the algorithm from Miles, Furman & Kuznetsov
to produce a univariate gamma convolution.

This algorithm works only if the distribution you computed the g_integrals from is, indeed, a generalized gamma convolution.
otherwise, it might fail and return garbage.

"""
function MFK_Projection(g_integrals::Vector{T},n) where T
    facts = factorial.(T.(0:2n))
    s = Array{T}(undef, 2n)
    build_s!(s,g_integrals,facts)
    alpha,beta = MFK_end(s,n)
    return ThorinDistributions.UnivariateGammaConvolution(Float64.(alpha), Float64.(1 ./ beta))
end


using BenchmarkTools
using InteractiveUtils

function trapz(a::Float64,b::Float64,n::Int64,f::Function)
    h::Float64=(b-a)/n
    sum=0.5*(f(a)+f(b))
    for i=1:n-1
        sum+=f(i*h)
    end
    sum*=h
end

@inline f(x::Float64)=+exp(-x)*x^2

function g(x::Float64)
    h=0.0
    if x<0.5
        h=-exp(-x)
    else
        h= exp(x)
    end
    return h*x^2
end

# Note: replacing f by f1 and g by g1 do not change the computing time:
# f1(x)=x < 1.5 ? +exp(-x)*x^2 :  +exp(-x)*x^2
# g1(x::Float64)=x < 0.5 ? -exp(-x)*x^2 :  exp(x)*x^2

function implicit(t::Float64)
    # implicit = root of  4*sin(x)-exp(x)+t
    # Newton iterations, starting from zero:
    x= 0.0
    F= 4*sin(x)-exp(x)+t
    while abs(F)> 1.e-15
        #
        x-= F/(4*cos(x) - exp(x))
        F=  4*sin(x)-exp(x)+t
        #
    end
    x
end

io = IOContext(stdout, :compact => false)

fw=open("RunningOn"*gethostname(),"w")
# Note: it seems that a loop like:
#
# for F in [f,g,implicit]
#   bench_res = @benchmark trapz(0.,1.,1000,F)
# end
#
# is incompatible with @benchmark... who knows why?

f(0.5)
println("f:")
bench_res = @benchmark trapz(0.,1.,1000,f);
show(io, bench_res)
write(fw,"\nf: "*string(bench_res)*"\n")

g(0.5)
println("\ng:")
bench_res = @benchmark trapz(0.,1.,1000,g);
show(io, bench_res)
write(fw,"g: "*string(bench_res)*"\n")

implicit(0.5)
println("\nimplicit:")
bench_res = @benchmark trapz(0.,1.,1000,implicit)
show(io, bench_res)
write(fw,"implicit: "*string(bench_res)*"\n")

println("\nend.")
close(fw)


using Statistics
using StatsBase

"""
Bounds given metric between -1.0 and 1.0, where 1.0 is perfect fit.

Suitable for use with any metric that ranges from 1 to -∞.

# References
1. Mathevet, T., Michel, C., Andréassian, V., Perrin, C., 2006.
    A bounded version of the Nash-Sutcliffe criterion for better model
    assessment on large sets of basins.
    IAHS-AISH Publication 307, 211–219.
    https://iahs.info/uploads/dms/13614.21--211-219-41-MATHEVET.pdf

# Example
```julia
julia> import Streamfall: @bound, KGE
julia> @bound KGE([1,2], [3,2])
-0.35653767993482094
```
"""
macro bound(metric)
    tmp = :($metric)
    return :($tmp / (2.0 - $tmp))
end

"""
Normalizes given metric between 0.0 and 1.0, where 1.0 is perfect fit.

Suitable for use with any metric that ranges from 1 to -∞.

# References
1. Nossent, J., Bauwens, W., 2012.
    Application of a normalized Nash-Sutcliffe efficiency to improve the
    accuracy of the Sobol’ sensitivity analysis of a hydrological model.
    EGU General Assembly Conference Abstracts 237.

# Example
```julia
julia> import Streamfall: @normalize, KGE
julia> @normalize KGE([1,2], [3,2])
0.1111111111111111
```
"""
macro normalize(metric)
    return :(1.0 / (2.0 - $metric))
end

"""
Applies mean inverse approach to a metric.

Suitable for use with any metric that ranges from 1 to -∞.

If using with other macros such as `@normalize` or `@bound`,
these must come first.

# References
1. Garcia, F., Folton, N., Oudin, L., 2017.
    Which objective function to calibrate rainfall–runoff
        models for low-flow index simulations?
    Hydrological Sciences Journal 62, 1149–1166.
    https://doi.org/10.1080/02626667.2017.1308511

# Example
```julia
julia> import Streamfall: @normalize, @mean_inverse, KGE
julia> @normalize @mean_inverse KGE [1,2] [3,2] 1e-6
0.3193506006429825
```
"""
macro mean_inverse(metric, obs, sim, ϵ=1e-2)
    obj, o, s, ϵ = eval(metric), eval(obs), eval(sim), eval(ϵ)
    q = obj(o, s)
    q2 = obj(1.0 ./ (o .+ ϵ), 1.0 ./ (s .+ ϵ))
    return mean([q, q2])
end

"""
Applies split meta metric approach

If using with other macros such as `@normalize` or `@bound`,
these must come first.

# References
1. Fowler, K., Peel, M., Western, A., Zhang, L., 2018.
    Improved Rainfall-Runoff Calibration for Drying Climate:
    Choice of Objective Function.
    Water Resources Research 54, 3392–3408.
    https://doi.org/10.1029/2017WR022466

# Example
```julia
julia> using Statistics
julia> import Streamfall: @normalize, @split, KGE
julia> @normalize @split KGE repeat([1,2], 365) repeat([3,2], 365) 365 mean
0.3217309561946589
```
"""
macro split(metric, obs, sim, n, agg_func=mean)
    obj, o, s, t, func = eval(metric), eval(obs), eval(sim), eval(n), eval(agg_func)
    return naive_split_metric(o, s; n_members=t, metric=obj, comb_method=func)
end

"""The Nash-Sutcliffe Efficiency score"""
NSE(obs, sim) = 1.0 - sum((obs .- sim).^2) / sum((obs .- mean(obs)).^2)

"""Normalized Nash-Sutcliffe Efficiency score (bounded between 0 and 1).

# References
1. Nossent, J., Bauwens, W., 2012.
    Application of a normalized Nash-Sutcliffe efficiency to improve the accuracy of the Sobol’ sensitivity analysis of a hydrological model.
    EGU General Assembly Conference Abstracts 237.
"""
NNSE(obs, sim) = 1.0 / (2.0 - NSE(obs, sim))

"""Root Mean Square Error"""
RMSE(obs, sim) = (sum((sim .- obs).^2)/length(sim))^0.5

"""Coefficient of determination (R^2)

Aliases `NSE()`
"""
function R2(obs, sim)::Float64
    return NSE(obs, sim)
end

"""Adjusted R^2

# Arguments
- `obs::Vector` : observations
- `sim::Vector` : modeled results
- `p::Int` : number of explanatory variables
"""
function ADJ_R2(obs, sim, p::Int64)::Float64
    n = length(obs)
    adj_r2 = 1 - (1 - R2(obs, sim)) * ((n - 1) / (n - p - 1))

    return adj_r2
end

"""
Mean Absolute Error
"""
MAE(obs, sim) = mean(abs.(sim .- obs))

"""
Mean Error
"""
ME(obs, sim) = mean(sim .- obs)

"""
    PBIAS(obs::Vector, sim::Vector)::Float64

Percent bias between `sim` and `obs`

Model performance for streamflow can be determined to be
satisfactory if the Nash-Sutcliffe Efficiency (NSE)
score > 0.5, the RMSE standard deviation ratio (RSR) < 0.7
and percent bias (PBIAS) is +/- 25% (see [1]).

# References
1. Moriasi, D.N., Arnold, J.G., Liew, M.W.V., Bingner, R.L.,
    Harmel, R.D., Veith, T.L., 2007.
    Model Evaluation Guidelines for Systematic Quantification
    of Accuracy in Watershed Simulations.
    Transactions of the ASABE 50, 885–900.
    https://doi.org/10.13031/2013.23153
"""
PBIAS(obs, sim) = (sum(obs .- sim) * 100) / sum(obs)

"""
    RSR(obs::Vector, sim::Vector)::Float64

The RMSE-observations standard deviation ratio (RSR).

Varies between 0 and a large positive value, where 0
indicates an RMSE value of 0.

# References
1. Moriasi, D.N., Arnold, J.G., Liew, M.W.V., Bingner, R.L.,
    Harmel, R.D., Veith, T.L., 2007.
    Model Evaluation Guidelines for Systematic Quantification
    of Accuracy in Watershed Simulations.
    Transactions of the ASABE 50, 885–900.
    https://doi.org/10.13031/2013.23153
"""
function RSR(obs, sim)::Float64
    rmse = RMSE(obs, sim)
    σ_obs = std(obs)
    rsr = rmse / σ_obs
    return rsr
end

"""
    EV(obs, sim)

Explained Variance.

Indicates the amount of variation in the observations which the predictions
are able to explain.
"""
function EV(obs, sim)
    return 1.0 - (var(obs - sim) / obs)
end

"""
Relative Skill Score.

Provides an indication of model performance relative to a known benchmark score.

Suitable for use with least-squares approaches that provide skill scores ranging from 1 to -∞.

# Arguments
- `Sb` : Benchmark score
- `Sm` : Model score

# References
1. Knoben, W.J.M., Freer, J.E., Woods, R.A., 2019.
    Technical note: Inherent benchmark or not?
        Comparing Nash-Sutcliffe and Kling-Gupta efficiency scores (preprint).
    Catchment hydrology/Modelling approaches.
    https://doi.org/10.5194/hess-2019-327
"""
function relative_skill_score(Sb::Float64, Sm::Float64)
    return (Sm - Sb) / (1.0 - Sb)
end

"""
    NSE_logbias(obs, sim; metric::Function=NSE, bias_threshold::Float64=5.0, shape::Float64=2.5)

The NSE_logbias meta-metric provides a weighted combination of a least-squares approach and a
logarithmic function of bias. The metric penalizes predictions with an overall bias above a
threshold (defined as 5% in [1]).

It is also referred to as the Viney F score.

# Extended help
The penalty applied is non-symmetrical (or multiplicatively symmetrical) in that
predictions that are _double_ the observed are penalized identically to predictions that are
_half_ the observed volume.

# Arguments
- `obs::Vector` : Historic observations to compare against
- `sim::Vector` : Modeled time series
- `metric::Function` : least-squares method to use, defaults to NSE
- `bias_threshold::Float64` : Bias threshold after which the score given by `metric` is penalized, defaults to 5 (%)
- `shape::Float64` : Exponent value controlling shape of penalization (see Figure 2 in [1]).

# References
1. Viney, N. R., Perraud, J., Vaze, J., Chiew, F.H.S., Post, D.A., Yang, A. 2009
    The usefulness of bias constraints in model calibration for regionalisation
        to ungauged catchments
    18th World IMACS / MODSIM Congress, Cairns, Australia, 13 - 17 July 2009
    Available at:
    https://www.researchgate.net/publication/294697092_The_usefulness_of_bias_constraints_in_model_calibration_for_regionalisation_to_ungauged_catchments

2. Teng, J., Potter, N.J., Chiew, F.H.S., Zhang, L., Wang, B., Vaze, J., Evans, J.P., 2015.
    How does bias correction of regional climate model precipitation affect modelled runoff?
    Hydrology and Earth System Sciences 19, 711–728.
    https://doi.org/10.5194/hess-19-711-2015
"""
function NSE_logbias(obs, sim; metric=NSE, bias_threshold=5.0, shape=2.5)
    E_ns = metric(obs, sim)
    B = PBIAS(obs, sim) / 100.0
    return E_ns - bias_threshold * abs(log(1 + B))^shape
end

"""
    KGE(obs::Vector, sim::Vector; scaling::Tuple=nothing)::Float64

Calculate the 2009 Kling-Gupta Efficiency (KGE) metric.

Decomposes NSE into correlation (`r`), relative variability (`α`), and bias (`β`) terms.

A KGE score of 1 means perfect fit.
A score < -0.41 indicates that the mean of observations
provides better estimates (see Knoben et al., [2]).

The `scaling` argument expects a three-valued tuple
which scales `r`, `α` and `β` factors respectively.
If not specified, defaults to `1`.

Note: Although similar, NSE and KGE cannot be directly compared.

# References
1. Gupta, H.V., Kling, H., Yilmaz, K.K., Martinez, G.F., 2009.
    Decomposition of the mean squared error and NSE performance criteria:
    Implications for improving hydrological modelling.
    Journal of Hydrology 377, 80–91.
    https://doi.org/10.1016/j.jhydrol.2009.08.003

2. Knoben, W.J.M., Freer, J.E., Woods, R.A., 2019.
    Technical note: Inherent benchmark or not? Comparing Nash-Sutcliffe and Kling-Gupta efficiency scores (preprint).
    Catchment hydrology/Modelling approaches.
    https://doi.org/10.5194/hess-2019-327

3. Mizukami, N., Rakovec, O., Newman, A.J., Clark, M.P., Wood, A.W.,
    Gupta, H.V., Kumar, R., 2019.
    On the choice of calibration metrics for “high-flow”
        estimation using hydrologic models.
    Hydrology and Earth System Sciences 23, 2601–2614.
    https://doi.org/10.5194/hess-23-2601-2019
"""
function KGE(obs, sim; scaling=nothing)::Float64
    if isnothing(scaling)
        scaling = (1, 1, 1)
    end

    # Correlation
    r = Statistics.cor(obs, sim)
    if isnan(r)
        r = 0.0
    end

    # relative variance
    α = std(sim) / std(obs)

    # bias
    β = mean(sim) / mean(obs)

    rs = scaling[1]
    as = scaling[2]
    bs = scaling[3]

    kge = 1 - sqrt(rs*(r - 1)^2 + as*(α - 1)^2 + bs*(β - 1)^2)

    return kge
end

"""Bounded KGE, bounded between -1 and 1.

# Arguments
- `obs::Vector` : observations
- `sim::Vector` : modeled results
"""
function BKGE(obs, sim)::Float64
    kge = KGE(obs, sim)
    return kge / (2 - kge)
end

"""Normalized KGE between 0 and 1.

# Arguments
- `obs::Vector` : observations
- `sim::Vector` : modeled results
"""
function NKGE(obs, sim; scaling=nothing)::Float64
    return 1 / (2 - KGE(obs, sim; scaling=scaling))
end

"""Calculate the modified KGE metric (2012).

Also known as KGE prime (KGE').

# Extended help

It is not recommended to apply KGE' with log-transformed flows (see [2]).
Numerical instabilities arise as flow approaches values close to zero.
This is possible under extreme dry conditions or by chance when sub-sampling.

In cases where observations are constant or otherwise displays zero variance or zero
mean flow, this implementation applies a simple logistic function (ℯ⁻ˣ) to gain an
indication of simulated data's distance to zero.

This is to:
- avoid NaNs influencing subsequent calculations
- allow use with split methods which may partition streamflows into periods of 0 flows.

# Arguments
- `obs::Vector`: observations
- `sim::Vector` : modeled results
- `scaling::Tuple` : scaling factors in order of timing (r), magnitude (β), variability (γ).
                     Defaults to (1,1,1).

# References
1. Kling, H., Fuchs, M., Paulin, M., 2012.
    Runoff conditions in the upper Danube basin under an ensemble of climate change scenarios.
    Journal of Hydrology 424–425, 264–277.
    https://doi.org/10.1016/j.jhydrol.2012.01.011

2. Santos, L., Thirel, G., Perrin, C., 2018.
    Technical note: Pitfalls in using log-transformed flows within the KGE criterion.
    Hydrology and Earth System Sciences 22, 4583–4591.
    https://doi.org/10.5194/hess-22-4583-2018
"""
function mKGE(obs, sim; scaling=nothing)::Float64
    if isnothing(scaling)
        scaling = (1,1,1)
    end

    # Timing (Pearson's correlation)
    r = Statistics.cor(obs, sim)
    if isnan(r)
        # can happen if obs or sim is a constant (std of 0)
        r = all(obs .== sim) ? 1.0 : 0.0
    end

    # Variability
    cv_s = StatsBase.variation(sim)
    if isnan(cv_s)
        cv_s = 0.0
    end

    cv_o = StatsBase.variation(obs)
    if isnan(cv_o)
        cv_o = 0.0
    end

    if cv_o == 0.0 && cv_s == 0.0
        γ = 1.0
    elseif cv_o == 0.0
        # Use logistic function to indicate distance from 0
        γ = 1.0 / exp(cv_s)
    else
        γ = cv_s / cv_o
    end

    # Magnitude
    μ_o = mean(obs)
    μ_s = mean(sim)
    if μ_o == 0.0
        # use logistic function to indicate distance from 0 (μ_o)
        β = 1.0 / exp(μ_s)
    else
        β = μ_s / μ_o
    end

    rs = scaling[1]
    βs = scaling[2]
    γs = scaling[3]

    mod_kge = 1.0 - sqrt(rs*(r - 1)^2 + βs*(β - 1)^2 + γs*(γ - 1)^2)

    return mod_kge
end

"""Bounded modified KGE between -1 and 1.

# Arguments
- `obs::Vector` : observations
- `sim::Vector` : modeled results
"""
function BmKGE(obs, sim; scaling=nothing)::Float64
    mkge = mKGE(obs, sim; scaling=scaling)
    return mkge / (2 - mkge)
end

"""Normalized modified KGE between 0 and 1.

# Arguments
- `obs::Vector` : observations
- `sim::Vector` : modeled results
"""
function NmKGE(obs, sim; scaling=nothing)::Float64
    return 1 / (2 - mKGE(obs, sim; scaling=scaling))
end

"""
Mean Inverse NmKGE

Said to produce better fits for low-flow indices
compared to mKGE (see [1]).

# Arguments
- `obs::Vector` : observations
- `sim::Vector` : modeled results
- `scaling::Tuple` : scaling factors for r, α, and β (defaults to 1.0)
- `ϵ::Float64` : small constant to use with inverse flow to allow consideration of periods with no flow. Defaults to 1e-2.

# References
1. Garcia, F., Folton, N., Oudin, L., 2017.
    Which objective function to calibrate rainfall–runoff
        models for low-flow index simulations?
    Hydrological Sciences Journal 62, 1149–1166.
    https://doi.org/10.1080/02626667.2017.1308511
"""
mean_NmKGE(obs, sim; scaling=nothing, ϵ=1e-2) = mean([Streamfall.NmKGE(obs, sim; scaling=scaling), Streamfall.NmKGE(1.0 ./ (obs .+ ϵ), 1.0 ./ (sim .+ ϵ); scaling=scaling)])

"""Calculate the non-parametric Kling-Gupta Efficiency (KGE) metric.

# Arguments
- `obs::Vector` : observations
- `sim::Vector` : modeled
- `scaling::Tuple` : scaling factors for timing (s), variability (α), magnitude (β)

# References
1. Pool, S., Vis, M., Seibert, J., 2018.
    Evaluating model performance: towards a non-parametric variant of the Kling-Gupta efficiency.
    Hydrological Sciences Journal 63, 1941–1953.
    https://doi.org/10.1080/02626667.2018.1552002

"""
function npKGE(obs, sim; scaling=nothing)::Float64
    if isnothing(scaling)
        scaling = (1,1,1)
    end

    # flow duration curves
    μ_s = mean(sim)
    if μ_s == 0.0
        fdc_sim = repeat([0.0], length(sim))
    else
        x = length(sim) * μ_s
        fdc_sim = sort(sim / x)
    end

    μ_o = mean(obs)
    if μ_o == 0.0
        fdc_obs = repeat([0.0], length(obs))
    else
        x = length(obs) * μ_o
        fdc_obs = sort(obs / x)
    end

    α = 1 - 0.5 * sum(abs.(fdc_sim - fdc_obs))

    # Magnitude
    if μ_o == 0.0
        β = μ_s == 0.0 ? 1.0 : μ_s
    else
        β = μ_s / μ_o
    end

    # Timing and flow dynamics
    r = StatsBase.corspearman(fdc_obs, fdc_sim)
    if isnan(r)
        r = 1.0  # can occur if identical sequences are used (e.g., 0 flows)
    end

    rs = scaling[1]
    αs = scaling[2]
    βs = scaling[3]

    kge = 1 - sqrt(rs*(r - 1)^2 + αs*(α - 1)^2 + βs*(β - 1)^2)

    return kge
end

"""Bounded non-parametric KGE between -1 and 1.

# Arguments
- `obs::Vector` : observations
- `sim::Vector` : modeled results
"""
function BnpKGE(obs, sim; scaling=nothing)::Float64
    npkge = npKGE(obs, sim; scaling=scaling)
    return npkge / (2 - npkge)
end

"""Normalized non-parametric KGE between 0 and 1.

# Arguments
- `obs::Vector` : observations
- `sim::Vector` : modeled results
"""
function NnpKGE(obs, sim; scaling=nothing)::Float64
    return 1 / (2 - npKGE(obs, sim; scaling=scaling))
end

"""Liu Mean Efficiency metric (LME).

Reformulation of the KGE metric said to be advantageous for capturing extreme
flow events.

# Arguments
- `obs::Vector` : observations
- `sim::Vector` : modeled results

# References
1. Liu, D., 2020.
    A rational performance criterion for hydrological model.
    Journal of Hydrology 590, 125488.
    https://doi.org/10.1016/j.jhydrol.2020.125488
"""
function LME(obs, sim)::Float64
    μ_o = mean(obs)
    μ_s = mean(sim)
    β = (μ_s / μ_o)

    r = Statistics.cor(obs, sim)
    σ_s = std(sim)
    σ_o = std(obs)
    α = (σ_s / σ_o)
    k_1 = r * α

    LME = 1 - sqrt((k_1 - 1)^2 + (β - 1)^2)

    return LME
end

function naive_split_metric(obs::Vector, sim::Vector, n_members::Int, metric::Function=NNSE)
    obs_chunks = collect(Iterators.partition(obs, n_members))
    sim_chunks = collect(Iterators.partition(sim, n_members))
    scores = Array{Float64, 1}(undef, length(obs_chunks))

    for (idx, h_chunk) in enumerate(obs_chunks)
        scores[idx] = metric(h_chunk, sim_chunks[idx])
    end

    return scores
end

"""Naive approach to split metrics.

Split metrics are a meta-objective optimization approach which "splits" data
into subperiods. The objective function is calculated for each subperiod and
then recombined. The approach addresses the lack of consideration of dry years
with least-squares.

In Fowler et al., [1] the subperiod is one year. The implementation offered here
is "naive" in that the data is partitioned into `N` chunks of `n_members` and
does not consider date/time.

# Arguments
- `obs::Vector` : Historic observations to compare against
- `sim::Vector` : Modeled time series
- `n_members::Int` : number of members per chunk (i.e., sub-samples), defaults to 365
- `metric::Function` : Objective function to apply, defaults to NNSE
- `comb_method::Function` : Recombination method, defaults to `mean`

# References
1. Fowler, K., Peel, M., Western, A., Zhang, L., 2018.
    Improved Rainfall-Runoff Calibration for Drying Climate: Choice of Objective Function.
    Water Resources Research 54, 3392–3408.
    https://doi.org/10.1029/2017WR022466
"""
function naive_split_metric(obs, sim; n_members::Int=365, metric::Function=NNSE, comb_method::Function=mean)
    scores = naive_split_metric(obs, sim, n_members, metric)
    return comb_method(scores)
end

"""
    inverse_metric(obs, sim; metric, comb_method::Function=mean)

A meta-objective function which combines the performance of the
given metric as applied to the discharge and the inverse of the
discharge.

By default, the combination method is to take the mean.

# Arguments
- obs : observed
- sim : modeled results
- `metric::Function` : objective function
- `comb_method::Function` : mean
- ϵ : offset value to use (enables use with zero-flow time steps), defaults to 1e-2

# References
1. Garcia, F., Folton, N., Oudin, L., 2017.
    Which objective function to calibrate rainfall–runoff models
        for low-flow index simulations?
    Hydrological Sciences Journal 62, 1149–1166.
    https://doi.org/10.1080/02626667.2017.1308511
"""
function inverse_metric(obs, sim, metric::Function; comb_method::Function=mean, ϵ=1e-2)
    return comb_method([metric(obs, sim), metric(1.0 ./ (obs .+ ϵ), 1.0 ./ (sim .+ ϵ))])
end

"""
Allows comparison of any model compared against a pre-defined benchmark, assuming both scores were obtained with
the same objective function.

Positive values indicate a model is better than the benchmark, and negative values indicate a model performs worse.

# Extended help
It is noted in Knoben et al., [1] that the skill score should always be contextualized with the original benchmark
value. Interpreting skill scores by themselves may become difficult if the benchmark score is already quite high.
A small improvement of no real practical value could be misconstrued as a large improvement.
As an example, if the benchmark has an KGE score of 0.999 and its counterpart 0.9995, then a skill score of 0.5 will
be reported.

# References
1. Knoben, W.J.M., Freer, J.E., Woods, R.A., 2019.
    Technical note: Inherent benchmark or not? Comparing Nash-Sutcliffe and Kling-Gupta efficiency scores (preprint).
    Catchment hydrology/Modelling approaches.
    https://doi.org/10.5194/hess-2019-327

2. Towner, J., Cloke, H.L., Zsoter, E., Flamig, Z., Hoch, J.M., Bazo, J., Coughlan de Perez, E., Stephens, E.M., 2019.
    Assessing the performance of global hydrological models for capturing peak river flows in the Amazon basin.
    Hydrology and Earth System Sciences 23, 3057–3080.
    https://doi.org/10.5194/hess-23-3057-2019
"""
function skill_score(model_score, benchmark_score)
    return (model_score - benchmark_score) / (1.0 - benchmark_score)
end


mutable struct OptimizationParameters
    d::Int # dimension of embedding
    # main optimization parameters
    max_iter::Int
    lr::Real
    beta_1::Real
    beta_2::Real
    weight_decay::Real
    eps::Real # adding to denominator of ADAM rescaling
    # early optimization parameters
    early_iter::Int
    early_lr::Real
    early_beta_1::Real
    early_beta_2::Real
    early_compression::Real
    early_exaggeration::Real
    # boolean flags
    dofast::Bool
    debias::Bool
    verbose::Bool
end

function OptimizationParameters(; d::Int = 2, max_iter::Int = 1024, lr::Real = 1,
            beta_1::Real = .9, beta_2::Real = .999, weight_decay::Real = 0,
            early_iter::Int = 64, early_lr::Real = lr, early_beta_1::Real = 0.5,
            early_beta_2::Real = 0.9, early_compression::Real = 1,
            early_exaggeration::Real = 12, eps::Real = 1e-8, verbose::Bool = false,
            debias::Bool = true, dofast::Bool = size(X, 2) ≥ fast_algorithm_min_size)
    early_exaggeration ≥ 1 || throw(DomainMismatch("early_exaggeration < 1"))
    OptimizationParameters(d, max_iter, lr, beta_1, beta_2, weight_decay, eps,
                early_iter, early_lr, early_beta_1, early_beta_2, early_compression,
                early_exaggeration, dofast, debias, verbose)
end


# Based upon alg 1 in
# A unified sparse optimization framework to learn parsimonious physics-informed models from data
# by K Champion et. al.

mutable struct SR3{U,T} <: AbstractOptimiser
    λ::U
    ν::U
    R::T
end

function SR3(λ = 1e-1, ν = 1.0)
    R = NormL1
    return SR3(λ, ν, R)
end

function set_threshold!(opt::SR3, threshold)
    opt.λ = threshold^2*opt.ν /2
    return
end

get_threshold(opt::SR3) = sqrt(2*opt.λ/opt.ν)

init(o::SR3, A::AbstractArray, Y::AbstractArray) =  A \ Y
init!(X::AbstractArray, o::SR3, A::AbstractArray, Y::AbstractArray) =  ldiv!(X, qr(A, Val(true)), Y)

function fit!(X::AbstractArray, A::AbstractArray, Y::AbstractArray, opt::SR3; maxiter::Int64 = 10)
    f = opt.R(get_threshold(opt))

    n, m = size(A)
    W = copy(X)

    # Init matrices
    P = inv(A'*A+I(m)/(opt.ν))
    X̂ = A'*Y
    for i in 1:maxiter
        # Solve rigde regression
        X .= P*(X̂+W/(opt.ν))
        # Add proximal iteration
        prox!(W, f, X, opt.ν*opt.λ)
    end

    # This is the effective threshold of the SR3 algorithm
    # See Unified Framework paper supplementary material S1
    #η = sqrt(2*opt.λ*opt.ν)
    X[abs.(X) .< get_threshold(opt)] .= zero(eltype(X))
    return
end


#' ---
#' title: Nonlinear Modelling
#' ---

#' **Originally Contributed by**: Arpit Bhatia

#' This tutorial provides a breif introduction to nonlinear modelling in JuMP.
#' For more details and specifics, visit the [JuMP docs](http://www.juliaopt.org/JuMP.jl/stable/nlp/).

#' ## Nonlinear Programs
#' While we have already seen examples of linear, quadratic and conic programs,
#' JuMP also supports other general smooth nonlinear (convex and nonconvex) optimization problems.

#' A JuMP model object can contain a mix of linear, quadratic, and nonlinear contraints or objective functions.
#' Thus, a model object for a nonlinear program is constructed in the same way as before.

using JuMP, Ipopt
model = Model(optimizer_with_attributes(Ipopt.Optimizer, "print_level" => 0));

#' ### Variables
#' Variables are modelled using the `@variable` macro as usual and
#' a starting point may be provided by using the `start` keyword argument

@variable(model, x, start = 4)
@variable(model, y, start = -9.66);

#' ### Parameters
#' Only in the case of nonlinear models, JuMP offers a syntax for "parameter" objects
#' which can refer to a numerical value.

@NLparameter(model, p == 0.003); # Providing a starting value is necessary for parameters
@NLparameter(model, l[i = 1:10] == 4 - i); # A collection of parameters

#' The `value` and `set_value` functions are used to query and update the value of a parameter respectively.

value(l[1])

#+

set_value(l[1], -4)
value(l[1])

#' Parameters are useful since it's faster to modify a model in-place by changing the value of the parameter
#' compared to creating an entirely new model object.

#' ### Expressions
#' JuMP also supports the creation of arithmetic expressions which can then be inserted into
#' constraints, the objective and other expressions.

@NLexpression(model, expr_1, sin(x))
@NLexpression(model, expr_2, asin(expr_1)); # Inserting one expression into another

#' There are some [syntax rules](https://pkg.julialang.org/docs/JuMP/DmXqY/0.19.2/nlp/#Syntax-notes-1)
#' which must be followed while writing a nonlinear expression.

#' Note that JuMP also supports linear and quadratic expression.
#' You can find out more about this functionality in the [docs](https://pkg.julialang.org/docs/JuMP/DmXqY/0.19.2/expressions/).

#' ### Nonlinear Objectives and Constraints
#' Nonlinear objectives and constraints are specified by using the `@NLobjective` and `@NLconstraint` macros.

@NLconstraint(model, exp(x) + y^4 <= 0)
@NLobjective(model, Min, tan(x) + log(y))

#' ### User-defined Functions
#' In addition to supporting a library of built-in functions,
#' JuMP supports the creation of user-defined nonlinear functions to use within nonlinear expressions.
#' The `register` function is used to enable this functionality.

my_function(a,b) = (a * b)^-6 + (b / a)^3
register(model, :my_function, 2, my_function, autodiff = true)

#' The arguements to the function are:
#' - model for which the function is being registered
#' - Julia symbol object corresponding to the name of the function
#' - Number of arguments the function takes
#' - name of the Julia method
#' - instruction for JuMP to compute exact gradients automatically

#' ## MLE using JuMP

#' Since we already have a bit of JuMP experience at this point,
#' let's try a modelling example and apply what we have learnt.
#' In this example, we compute the maximum likelihood estimate (MLE) of
#' the parameters of a Gaussian distribution i.e. the sample mean and variance.

#' If $X_{1}, \ldots, X_{n}$ are an id sample from a population with pdf or pmf
#' $f\left(x | \theta_{1}, \ldots, \theta_{k}\right),$ the likelihood function is defined by

#' $$
#' L(\theta | \mathbf{x})=L\left(\theta_{1}, \ldots, \theta_{k} | x_{1}, \ldots, x_{n}\right)=\prod_{i=1}^{n} f\left(x_{i} | \theta_{1}, \ldots, \theta_{k}\right)
#' $$

#' For each sample point $\mathbf{x}$, let $\hat{\theta}(\mathbf{x})$ be a parameter value
#' at which $L(\theta | \mathbf{x})$ attains its maximum as a function of $\theta,$ with $\mathbf{x}$ held fixed.
#' A maximum likelihood estimator (MLE) of the parameter $\theta$ based on a sample $\mathbf{X}$ is
#' $\hat{\theta}(\mathbf{X})$.

#'The Gaussian likelihood is -

#' $$
#' L(\theta | \mathbf{x})=\prod_{i=1}^{n} \frac{1}{(2 \pi)^{1 / 2}} e^{-(1 / 2)\left(x_{i}-\theta\right)^{2}}=\frac{1}{(2 \pi)^{n / 2}} e^{(-1 / 2) \Sigma_{i=1}^{n}\left(x_{i}-\theta\right)^{2}}
#' $$

#' In most cases, the natural logarithm of
#' $L(\theta | \mathbf{x}), \log L(\theta | \mathbf{x})$ (known as the log likelihood),
#' is used rather than $L(\theta | \mathbf{x})$ directly.
#' The reason is that the log likelihood is easier to differentiate.
#' This substituion is possible because the log function is strictly increasing on $(0, \infty)$,
#' which implies that the extrema of $L(\theta | \mathbf{x})$ and $\log L(\theta | \mathbf{x})$ coincide.

using Random, Statistics

Random.seed!(1234)

n = 1_000
data = randn(n)

mle = Model(optimizer_with_attributes(Ipopt.Optimizer, "print_level" => 0))
@NLparameter(mle, problem_data[i = 1:n] == data[i])
μ0 = randn()
σ0 = rand() + 1
@info "Starting guess, mean: $μ0, std: $σ0"
@variable(mle, μ, start = μ0)
@variable(mle, σ >= 0.0, start = σ0)
@NLexpression(mle, loglikelihood,
    -(n / 2) * (log(2π) + 2 * log(σ)) - inv(2 * σ^2) * sum((xi - μ)^2 for xi in problem_data)
)

@NLobjective(mle, Max, loglikelihood)

optimize!(mle)

println("μ = ", value(μ))
println("mean(data) = ", mean(data))
println("σ^2 = ", value(σ)^2)
println("var(data) = ", var(data))
println("MLE value: ", exp(objective_value(mle)))

#+

# Changing the data

data = randn(n)
optimize!(mle)

println("μ = ", value(μ))
println("mean(data) = ", mean(data))
println("σ^2 = ", value(σ)^2)
println("var(data) = ", var(data))
println("MLE objective: ", objective_value(mle))

#' ## Writing Convex Models

#' Nonlinear solvers like Ipopt are usually local solvers.
#' For convex problems, the local optima is also the global optima,
#' and thus Ipopt is able to provide us with the correct solution.
#' However, in case a problem is not written in the convex form, then we may be unable to solve it.

#' A tool that helps us in dealing with this issue is [Disciplined Convex Programming](https://dcp.stanford.edu) (DCP).
#' DCP is a system for constructing mathematical expressions with known curvature from a given library of base functions.
#' Specifically, it helps us to construct convex optimization models when possible,
#' i.e. minimize convex function or maximize concave function and use constraints that are convex $f$ <= concave $g$


"""Calculates the 2nd order centred gradient in x-direction on any grid (u,v,T or q).
The size of dudx must be m-1,n compared to m,n = size(u)"""
function ∂x!(dudx::Matrix{T},u::Matrix{T}) where {T<:AbstractFloat}
    m,n = size(dudx)
    @boundscheck (m+1,n) == size(u) || throw(BoundsError())

    @inbounds for j ∈ 1:n, i ∈ 1:m
        dudx[i,j] = u[i+1,j] - u[i,j]
    end
end

"""Calculates the 2nd order centred gradient in y-direction on any grid (u,v,T or q).
The size of dudy must be m,n-1 compared to m,n = size(u)."""
function ∂y!(dudy::Array{T,2},u::Array{T,2}) where {T<:AbstractFloat}
    m,n = size(dudy)
    @boundscheck (m,n+1) == size(u) || throw(BoundsError())

    @inbounds for j ∈ 1:n, i ∈ 1:m
            dudy[i,j] = u[i,j+1] - u[i,j]
    end
end

""" ∇² is the 2nd order centred Laplace-operator ∂/∂x^2 + ∂/∂y^2.
The 1/Δ²-factor is omitted and moved into the viscosity coefficient."""
function ∇²!(du::Matrix{T},u::Matrix{T}) where {T<:AbstractFloat}
    m, n = size(du)
    @boundscheck (m+2,n+2) == size(u) || throw(BoundsError())

    @inbounds for j ∈ 1:n
        for i ∈ 1:m
            #        1
            # the 1 -4 1  -stencil in low-precision resilient form
            #        1
            ui1j1 = u[i+1,j+1]
            du[i,j] = ((u[i,j+1] - ui1j1) + (u[i+2,j+1] - ui1j1)) + ((u[i+1,j] - ui1j1) + (u[i+1,j+2] - ui1j1))
        end
    end
end

"""∂x is the 2nd order centred Gradient-operator ∂/∂x with grid spacing Δ (default 1)."""
function ∂x(u::Array{T,2},Δx::Real) where {T<:AbstractFloat}

    m,n = size(u)

    dudx = Array{T,2}(undef,m-1,n)
    one_over_dx = T(1.0/Δx)

    @inbounds for j ∈ 1:n
        for i ∈ 1:m-1
            dudx[i,j] = one_over_dx*(u[i+1,j] - u[i,j])
        end
    end

    return dudx
end

"""∂y is the 2nd order centred Gradient-operator ∂/∂y with grid spacing Δ (default 1)."""
function ∂y(u::Array{T,2},Δy::Real=1) where {T<:AbstractFloat}

    m,n = size(u)

    dudy = Array{T,2}(undef,m,n-1)
    one_over_dy = T(1.0/Δy)

    @inbounds for j ∈ 1:n-1
        for i ∈ 1:m
            dudy[i,j] = one_over_dy*(u[i,j+1] - u[i,j])
        end
    end

    return dudy
end

""" ∇² is the 2nd order centred Laplace-operator ∂/∂x^2 + ∂/∂y^2 with grid spacing Δ (default 1)."""
function ∇²(u::Array{T,2},Δ::Real=1) where {T<:AbstractFloat}

    m, n = size(u)
    du = Array{T,2}(undef,m-2,n-2)

    minus_4 = T(-4.0)
    one_over_dx² = T(1/Δ^2)

    @inbounds for j ∈ 2:n-1
        for i ∈ 2:m-1
            du[i-1,j-1] = one_over_dx²*(minus_4*u[i,j] + u[i,j-1] + u[i,j+1] + u[i-1,j] + u[i+1,j])
        end
    end
    return du
end


### A Pluto.jl notebook ###
# v0.14.3

using Markdown
using InteractiveUtils

# ╔═╡ d1f4afa8-2501-4447-96a6-7fb93f463e2c
using IterativeSolvers, DataFrames

# ╔═╡ 57c64cca-55ee-4c2b-ad27-2942a262df7b
md" ### Ex-02-08.jl"

# ╔═╡ c7b3e4fb-8f17-4e0e-82a5-18d30f8b198a
begin
	A = Float64[
	  16 4 8;
	  4 5 -4;
	  8 -4 22;
	]
	a = Float64[4; 2; 5]
	r = vcat(collect(1:6), collect(10:10:100))
end;

# ╔═╡ 21a0af0a-7873-45f5-bd59-31c43cbb0f0c
md" ##### Jacobi method:"

# ╔═╡ 137e76a3-ea70-49e0-a1f7-a2eadaac9963
x_j = [(i=i, j=jacobi(A, a; maxiter=i)) for i in r]

# ╔═╡ 31d9aa57-e96b-409b-b1ab-9e5966af6bf5
x_j[1]

# ╔═╡ 70d35d30-1798-4982-9b9e-3812a575d179
A * x_j[11].j

# ╔═╡ 954c3dec-86ad-4a29-9528-a790a9ebf2c0
md" ##### Gauss-Seidel method:"

# ╔═╡ 53491c26-f5d1-48d0-a795-7c20a31b00df
x_gs = [(i=i, gs=gauss_seidel(A, a; maxiter=i)) for i in r]

# ╔═╡ b2a3e4de-cd59-466a-9872-f72e78bb91ba
A * x_gs[11].gs

# ╔═╡ df81e70f-ec63-4be6-8517-d962b315dc0f
md" ##### Succesive Over-Relaxation method:"

# ╔═╡ cebd114d-8a7c-44c0-bbca-e39b124a263e
x_sor = [(i=i, sor=sor(A, a, 1.0; maxiter=i)) for i in r]

# ╔═╡ 9b86fca1-6867-4c70-b7a2-6e648bbfc26e
A * x_sor[11].sor

# ╔═╡ c1dd5a8b-5256-410f-9abc-58afd36cb1af
md" ##### Sequential Succesive Over-Relaxation method:"

# ╔═╡ 09fd7541-fe59-435b-8f6d-73b72b0c01f0
x_ssor = [(i=i, ssor=ssor(A, a, 1.0; maxiter=i)) for i in r]

# ╔═╡ 96d4ef12-7a35-47c5-a255-5056e582e618
A * x_ssor[11].ssor

# ╔═╡ Cell order:
# ╟─57c64cca-55ee-4c2b-ad27-2942a262df7b
# ╠═d1f4afa8-2501-4447-96a6-7fb93f463e2c
# ╠═c7b3e4fb-8f17-4e0e-82a5-18d30f8b198a
# ╟─21a0af0a-7873-45f5-bd59-31c43cbb0f0c
# ╠═137e76a3-ea70-49e0-a1f7-a2eadaac9963
# ╠═31d9aa57-e96b-409b-b1ab-9e5966af6bf5
# ╠═70d35d30-1798-4982-9b9e-3812a575d179
# ╟─954c3dec-86ad-4a29-9528-a790a9ebf2c0
# ╠═53491c26-f5d1-48d0-a795-7c20a31b00df
# ╠═b2a3e4de-cd59-466a-9872-f72e78bb91ba
# ╟─df81e70f-ec63-4be6-8517-d962b315dc0f
# ╠═cebd114d-8a7c-44c0-bbca-e39b124a263e
# ╠═9b86fca1-6867-4c70-b7a2-6e648bbfc26e
# ╟─c1dd5a8b-5256-410f-9abc-58afd36cb1af
# ╠═09fd7541-fe59-435b-8f6d-73b72b0c01f0
# ╠═96d4ef12-7a35-47c5-a255-5056e582e618


# Load X and y variable
using JLD
data = load("nonLinear.jld")
(X,y,Xtest,ytest) = (data["X"],data["y"],data["Xtest"],data["ytest"])

# Compute number of training examples and number of features
(n,d) = size(X)

display(size(X))

ntrain = Int(floor(n/2))

Xtrain = X[1:ntrain, :]
ytrain = y[1:ntrain, :]
Xvalidate = X[ntrain+1:n, :]
yvalidate = y[ntrain+1:n, :]

display(size(Xtrain))
display(size(Xvalidate))

# Fit least squares model
#include("leastSquares.jl")
#model = leastSquares(X,y)

bestError = 10000000000000000
bestl = 0
bestsigma = 0

include("leastSquaresRBFL2.jl")

# Find best l, sigma values
for l = 0:0.1:3
  for sigma = 0.1:0.1:3
		model = leastSquaresRBFL2(Xtrain,ytrain,l, sigma)

		# Report the error on the validation set
		t = size(Xvalidate,1)
		yhat = model.predict(Xvalidate)
		validationError = sum((yhat - yvalidate).^2)/t
		@printf("l = %f, sigma = %f, ValidationError = %.2f\n", l, sigma, validationError)
    if validationError < bestError
      bestError = validationError
      bestl = l
      bestsigma = sigma
    end
  end
end

@printf("Best validation l = %f, sigma = %f, error = %f", bestl, bestsigma, bestError)

model = leastSquaresRBFL2(X,y,bestl,bestsigma)

# Report the error on the test set
t = size(Xtest,1)
yhat = model.predict(Xtest)
testError = sum((yhat - ytest).^2)/t
@printf("l = %f, sigma = %f, TestError = %.2f\n", bestl, bestsigma, testError)

# Plot model
using PyPlot
figure()
plot(X,y,"b.")
plot(Xtest,ytest,"g.")
Xhat = minimum(X):.1:maximum(X)
Xhat = reshape(Xhat,length(Xhat),1) # Make into an n by 1 matrix
yhat = model.predict(Xhat)
plot(Xhat,yhat,"r")
ylim((-300,400))
show()


# This script provides implementations of
# multi-step methods of different orders
# for assessing their performance on a test
# equation y' = λ y

println('\n', " "^4, "> Loading the packages...")

using Printf

# Define the value of λ
const λ = -1.5

# Integrate the test equation using the
# Euler's method, return the value of the
# solution at the last step
function euler(
    y₀::F,
    h::F,
    n::I,
)::F where
{F <: AbstractFloat,I <: Unsigned}
    # Prepare a buffer for the solution
    y = y₀
    # Compute the solution
    for _ in 1:n
        y += h * λ * y
    end
    return y
end

# Integrate the test equation using the
# two-step Adams–Bashforth's method, return
# the value of the solution at the last step
function ab2(
    y₀::F,
    h::F,
    n::I,
)::F where
{F <: AbstractFloat,I <: Unsigned}
    # Prepare a buffer for the solution
    y = y₀
    # Prepare a buffer for previous values
    yₖ₋₁ = y₀
    # Compute the second value of the solution
    # using the one-step Euler's method
    y = euler(y₀, h, UInt(1))
    # Compute the rest in two steps
    for _ in 2:n
        k = y
        y += h * λ * (3 / 2 * y - 1 / 2 * yₖ₋₁)
        yₖ₋₁ = k
    end
    return y
end

# Integrate the test equation using the
# 4th order Runge-Kutta's method, return
# the value of the solution at the last step
function rk4(
    y₀::F,
    h::F,
    n::I,
)::F where
{F <: AbstractFloat,I <: Unsigned}
    # Prepare a buffer for the solution
    y = y₀
    # Compute the rest in two steps
    for _ in 1:n
        k₁ = λ * y
        k₂ = λ * (y + h * k₁ / 2)
        k₃ = λ * (y + h * k₂ / 2)
        k₄ = λ * (y + h * k₃)
        y += h / 6 * (k₁ + 2 * k₂ + 2 * k₃ + k₄)
    end
    return y
end

# Integrate the test equation using the
# three-step Adams–Bashforth's method
# (with Euler's method as the starter),
# return the value of the solution at
# the last step
function ab3_euler(
    y₀::F,
    h::F,
    n::I,
)::F where
{F <: AbstractFloat,I <: Unsigned}
    # Prepare a buffer for the solution
    y = y₀
    # Prepare buffers for previous values
    yₖ₋₁ = y₀
    yₖ₋₂ = y₀
    # Compute the second value of the solution
    # using the one-step Euler's method
    y = euler(y₀, h, UInt(1))
    # Compute the third value of the solution
    # using the two-step Adams–Bashforth's method
    yₖ₋₁ = y
    y = ab2(y₀, h, UInt(2))
    # Compute the rest in two steps
    for _ in 3:n
        k = y
        y += h * λ * (23 / 12 * y - 4 / 3 * yₖ₋₁ + 5 / 12 * yₖ₋₂)
        yₖ₋₂ = yₖ₋₁
        yₖ₋₁ = k
    end
    return y
end

# Integrate the test equation using the
# three-step Adams–Bashforth's method
# (with 4th order Runge-Kutta's method as the
# starter), return the value of the solution
# at the last step
function ab3_rk4(
    y₀::F,
    h::F,
    n::I,
)::F where
{F <: AbstractFloat,I <: Unsigned}
    # Prepare a buffer for the solution
    y = y₀
    # Prepare buffers for previous values
    yₖ₋₁ = y₀
    yₖ₋₂ = y₀
    # Compute the second and third values of the solution
    # using the 4th-order Runge-Kutta's method
    yₖ₋₁ = rk4(y₀, h, UInt(1))
    y = rk4(y₀, h, UInt(2))
    # Compute the rest in two steps
    for _ in 3:n
        k = y
        y += h * λ * (23 / 12 * y - 4 / 3 * yₖ₋₁ + 5 / 12 * yₖ₋₂)
        yₖ₋₂ = yₖ₋₁
        yₖ₋₁ = k
    end
    return y
end

println(" "^4, "> Integrating the test equation...")

# Define the initial values of the position and velocity
y₀ = 1.0

# Calculate the analytical solution
yₐ = ℯ^λ

# Print the initial values
println(
    '\n',
    " "^6, "y₀: ", y₀,
)

# Define the output directories
const data_dir = "$(@__DIR__)/../data/test"
const files = joinpath.(
    data_dir,
    [
        "euler.dat",
        "ab2.dat",
        "rk4.dat",
        "ab3_euler.dat",
        "ab3_rk4.dat",
    ]
)

# Prepare a list of methods
const methods = (euler, ab2, rk4, ab3_euler, ab3_rk4)

# Integrate the test equation using all available
# methods and write the values of the solutions
# on the last step
function integrate(h::Vector{F}, n::Vector{I}) where {F <: AbstractFloat,I <: Unsigned}
    # Open the data files
    io = open.(files, "a")

    # Prepare results vectors
    y = [ Vector{F}() for _ in 1:length(io) ]
    Δy = [ Vector{F}() for _ in 1:length(io) ]

    # For each pair of parameters
    for i in eachindex(h)
        # Integrate and get the results of the last steps
        for j in eachindex(io)
            yᵢ = methods[j](y₀, h[i], n[i])
            Δyᵢ = abs(yₐ - yᵢ)
            push!(y[j], yᵢ)
            push!(Δy[j], Δyᵢ)
        end
    end

    # Get the format string based on the index of the iteration
    function get_format(i)
        Printf.Format(
            "\$ 10^{-$(length("$(n[i])") - 1)} \$ " *
            "& \$ $(n[i] % 10 == 0 ? "10^$(length("$(n[i])") - 1)" : n[i]) \$ " *
            "& \$ %.14f \$ " *
            "& \$ %.14f \$ " *
            "\\\\\n"
        )
    end

    # Print the data in the files in the LaTeX format

    for j in eachindex(io)
        println(io[j], "# Solution")
        for i in eachindex(h)
            Printf.format(io[j], get_format(i), y[j][i], Δy[j][i])
        end
    end
end

# Truncate the previous results
open.(files; truncate=true)

# Integrate n = 2^m iterations
integrate(
    [1e-2, 1e-3, 1e-4, 1e-5, 1e-6, 1e-7],
    UInt.([ 10^i for i in 2:7 ]),
)

println()


@testset "Gauss–Lobatto" begin
    # Check error
    @test_throws DomainError gausslobatto(-1)
    @test_throws DomainError gausslobatto(0)
    @test_throws DomainError gausslobatto(1)

    n = 2
    x,w = gausslobatto(n)
    @test x[1] ≈ -1.
    @test x[2] ≈ 1.
    @test w[1] ≈ 1.
    @test w[2] ≈ 1.

    n = 3
    x,w = gausslobatto(n)
    @test x[1] ≈ -1.
    @test abs(x[2])<1e-15
    @test x[3] ≈ 1.
    @test w[1] ≈ 1/3
    @test w[2] ≈ 4/3
    @test w[3] ≈ 1/3

    tol = 1e-14
    n = 42
    x,w = gausslobatto(n)
    @test ( length(x) == n) && ( length(w) == n )
    @test abs(x[37] - 0.922259214258616) < tol
    @test abs(w[37] - 0.029306411216166) < tol
    @test ( x[1] == -1 && x[n] == 1 )

    @test dot( w,(x.^2)) ≈ 2/3
    @test dot( w,exp.(x)) ≈ exp(1)-exp(-1)
end


module kMeansClustering

function squared_euclidean_distance(v::Array{Float64},w::Array{Float64})
    dist = 0
    for d in 1:length(v)
        dist+= (v[d]-w[d])^2
    end
    return dist
end

function agroup(matrix::Matrix{Float64})
    l = [Float64[]]
    pop!(l)
    dim,n_points = size(matrix)
    for p in 1:n_points
        v = Float64[]
        for d in 1:dim
            push!(v,matrix[dim*(p-1)+d])
        end
        push!(l,v)
    end
    return l,dim
end

function compare_distance(p::Array{Float64,1},old_ctr::Array{Float64,1},new_ctr::Array{Float64,1})
    old_dist = squared_euclidean_distance(p,old_ctr)
    new_dist = squared_euclidean_distance(p,new_ctr)
    if new_dist<=old_dist
        return new_ctr
    end
    return old_ctr
end

function dic(l::Array{Array{Float64,1},1},centers::Array{Array{Float64,1},1})
    dic_centers = Dict()
    for ctr in centers
        dic_centers[ctr] = [Float64[]]
        pop!(dic_centers[ctr])
    end
    for p in l
        center = centers[1]
        for ctr in centers
            center = compare_distance(p,center,ctr)
        end
        c = dic_centers[center]
        push!(c,p)
        dic_centers[center] = c
    end
    return dic_centers
end

function reorganize(clusters::Array{Array{Array{Float64,1},1},1},dim::Int64)
    centers = [Float64[]]
    pop!(centers)
    for cluster in clusters
        s = zeros(dim)
        for p in cluster
            for n in 1:dim
                s[n]+= p[n]
            end
        end
        means = Float64[]
        for d in s
            push!(means,d/length(cluster))
        end
        push!(centers,means)
    end
    return centers
end

function cost(dic_centers::Dict{Any,Any})
    cst = 0
    for (centr,clus) in dic_centers
        for p in clus
            cst+=squared_euclidean_distance(p,centr)
        end
    end
    return cst
end

function assign(clusters::Array{Array{Array{Float64,1},1},1},l::Array{Array{Float64,1},1})
    dic_assign = Dict()
    for n in 1:length(clusters)
        for el in clusters[n]
            dic_assign[el] = n
        end
    end
    l_assign = Int64[]
    for p in l
        push!(l_assign,dic_assign[p])
    end
    return l_assign
end

function new_format(centers::Array{Array{Float64,1},1})
    new = hcat(centers[1],centers[2])
    for p in centers[3:length(centers)]
        new = hcat(new,p)
    end
    return new
end

function build_clusters(dic_centers::Dict{Any,Any})
    clusters = [[Float64[]]]
    pop!(clusters)
    for (center,cluster) in dic_centers
        push!(clusters,cluster)
    end
    return clusters
end

function mykmeansclustering(matrix::Matrix{Float64},k::Int64)
    l,dim = agroup(matrix)
    centers = [Float64[]]
    pop!(centers)
    while length(unique(centers))<k
        centers = rand(l,k) # Forgy Method
    end
    initials = centers
    cop = [Float64[]]
    pop!(cop)
    while centers!=cop
        dic_centers = dic(l,centers)
        cop = centers[1:k]
        clusters = build_clusters(dic_centers)
        centers = reorganize(clusters,dim)
    end
    dic_centers = dic(l,centers)
    clusters = build_clusters(dic_centers)
    cst = cost(dic_centers)
    assignment = assign(clusters,l)
    centers = new_format(centers)
    initials = new_format(initials)
    return centers,cst,assignment,initials
end

export mykmeansclustering
end


using Test, JLD2, LinearAlgebra
#
# gradient computation by parial automatic differentiation and adjoint merthods
# of
#           E(v) = ∫ |U(v)|² dx
# where
# A(v) U(v) = F(v) the discrete version of - div(M ∇ Ua) = Fa
# and v are the vertices of the voronoi cells
#

# load polygonal mesh
filename, nc = "Lpolmesh_coarse", 100
mesh_filename = "$(@__DIR__)/data/$(filename)_$(nc).jld2"
println("read file $(mesh_filename)")
JLD2.@load(mesh_filename, pv, cellsb, cellsbt, t, pb, tb)
# boundary points / rhs
rhs, boundary_condition, meshboundary = if occursin("square", mesh_filename)
    PolygonalFem.rhs_sqr, PolygonalFem.boundary_condition_sqr, []
else
    EPS = 1e-4
    PolygonalFem.rhs_L_elas, PolygonalFem.boundary_condition_L_elas,
    findall(pv[:, 2] .> (1 - EPS))
end
np, ε = size(pv, 1), 1e-4
x, dp = pv, 2 * rand(size(pv)...) .- 1
# IMPORTANT Do not move boundary nodes
dp[meshboundary, :] .= 0.
xp, xm = x .+ ε * dp, x .- ε * dp
# finite differences
f = x -> PolygonalFem.costnormU_elas2D(x, cellsb, t, rhs, boundary_condition,
                                       meshboundary)
fp, fm = f(xp), f(xm)
@time g = PolygonalFem.∇costnormU_elas2D(x, cellsb, t, rhs, boundary_condition,
                                         meshboundary)
@time g = PolygonalFem.∇costnormU_elas2D(x, cellsb, t, rhs, boundary_condition,
                                         meshboundary)
println("")
println(fp)
println(fm)
println(" ")
println((fp - fm) / (2 * ε))
println(dot(g, dp))
println("")
@testset "test eval gradient dot function for polygonal costnormU_elas2D" begin
    @test abs(dot(g, dp) - (fp - fm) / (2 * ε)) < 1e-4
end


using ControlSystems, RobustAndOptimalControl
using Plots

"""
This is a simple SISO example which was used for debugging the implementation,
as this exact example was use in the lecture notes of the "Principles of Optimal
Control" cours of the MIT OpenCourseWare [1], however, this example is in discrete time

[1] https://ocw.mit.edu/courses/aeronautics-and-astronautics/16-323-principles-of-optimal-control-spring-2008/lecture-notes/lec15.pdf

The example can be set to visualize and save plots using the two variables
  makeplots - true/false (true if plots are to be generated, false for testing)
"""
makeplots = true

# Define the process
ts = 0.005
Gd = bilinearc2d(ss(tf([200], [0.025,1.0025,10.1,1])),ts)

# Sensitivity weight function
M, wB, A = 1.5, 10, 1e-4
WS = tf([1/M, wB],[1, wB*A])

# Output sensitivity weight function
WU = ss(0.1)

# Complementary sensitivity weight function
WT = []

# Create continuous time approximation of the process
Gc = bilineard2c(ss(Gd))

# Form augmented P dynamics in state-space
Pc = hinfpartition(Gc, WS, WU, WT)

# Check that the assumptions are satisfied
flag = hinfassumptions(Pc)

# Synthesize the H-infinity optimal controller
flag, Cc, γ = hinfsynthesize(Pc)

# Extract the transfer functions defining some signals of interest, but do so
# using discrete equivalent of the continuous time objects Pc, Cc and Gc
Pcl, S, CS, T = hinfsignals(
  bilinearc2d(Pc, ts),
  bilinearc2d(Gc, ts),
  bilinearc2d(Cc, ts)
)

# Visualize results
if makeplots
  specificationplot([S, CS, T], [WS, WU, WT], γ)
  specificationplot(Pcl, γ; s_labels=["\$\\sigma(P_{cl}(j\\omega))\$"], w_labels=["\$\\gamma\$"])
end