using BenchmarkTools

# see this section of docs https://docs.julialang.org/en/v1/base/arrays/#Broadcast-and-vectorization-1

function G_single(t::T) where T<:Real

   G = 1. + 2. *exp(-t*(2. /3.)) + 4. *exp(-t*(4. /5.))

end

function G_vec(t::Vector{T}) where T<:Real

   G = 1. .+ 2. .* exp.(-t .* (2. /3.)) .+ 4. .* exp.(-t .* (4. /5.))

end

function G_arraycomp(t::Vector{T}) where T<:Real
    [G_single(ti) for ti in t]
end

function G_loop(t::Vector{T}) where T<:Real

    out = Array{T}(undef, length(t))

    for (i,ti) in enumerate(t)
        out[i] = G_single(ti)
    end

    out
end

t = collect(0.0:0.1:10000)

println("test vectorised version of function")
@btime G_vec($t)
println("\n")

println("test broadcasting using . notation")
@btime G_single.($t)
println("\n")

println("test broadcast function explicitly (should be the same as . notation)")
@btime broadcast($G_single, $t)
println("\n")

println("test using array comprehension")
@btime G_arraycomp($t)
println("\n")

println("test using for loop")
@btime G_loop($t)
println("\n")

println("Benchmarks complete")


import FastGaussQuadrature
import FastTransforms

using DecFP: Dec128

using ..CommonFunctions

struct Quadrature{T,N}
    order::Int
    nodes::Vector{T}
    weights::Vector{T}
end

function Quadrature(order, nodes::Vector{T}, weights::Vector{T}) where {T}
    @assert length(nodes) == length(weights)
    Quadrature{T, length(nodes)}(order, nodes, weights)
end

CommonFunctions.nnodes(Q::Quadrature{T,N}) where {T,N} = N
CommonFunctions.nodes(Q::Quadrature) = Q.nodes
CommonFunctions.order(Q::Quadrature) = Q.order
weights(Q::Quadrature) = Q.weights
Base.eachindex(Q::Quadrature) = eachindex(Q.nodes, Q.weights)

"""
Integrate a function `f(x)` over the interval [0,1] using the quadrature `quad`.
"""
function quadrature(quad::Quadrature, f::Function)
    sum(quad.weights[i] * f(quad.nodes[i]) for i in eachindex(quad))
end

"Scale nodes and weights from the interval [-1,+1] to the interval [0,1]."
function shift!(b,c)
    b .= b ./ 2
    c .= (c .+ 1) ./ 2
end

"Scale nodes and weights from the interval [0,1] to the interval [-1,+1]."
function unshift!(b,c)
    b .= b .* 2
    c .= 2 .* c .- 1
end

function RiemannQuadratureLeft(T=Float64)
    Quadrature(1, [0.0,], [1.0,])
end

function RiemannQuadratureRight(T=Float64)
    Quadrature(1, [1.0,], [1.0,])
end

function GaussLegendreQuadrature(s, T=Float64; shift=true)
    c, b = FastGaussQuadrature.gausslegendre(s)
    if shift shift!(b,c) end
    Quadrature(2s, c, b)
end

function LobattoLegendreQuadrature(s, T=Float64; shift=true)
    c, b = FastGaussQuadrature.gausslobatto(s)
    if shift shift!(b,c) end
    Quadrature(2s-2, c, b)
end

function get_nodes_gauss_chebyshev(s, k=1, T=Float64)
    x = zeros(T, s)
    if k == 1
        for i in eachindex(x)
            x[i] = @dec128 ( sin( π*(s-2i+1) / (2s) ) + 1 ) / 2
        end
        return x[end:-1:1]
    elseif k == 2
        for i in eachindex(x)
            x[i] = @dec128 ( cos( π*i / (s-1) ) + 1 ) / 2
        end
        return x
    end
end

function GaussChebyshevQuadrature(s, k=1, T=Float64; shift=true)
    c = get_nodes_gauss_chebyshev(s,k,T)
    b = zeros(c)

    tj::Dec128 = 0
    th::Dec128 = 0

    if k == 1
        for i in eachindex(b)
            tj = 0
            th = @dec128 π * (2i-1) / (2s)
            for j in 1:div(s,2)
                tj += @dec128 cos(2j*th) / (4j^2 - 1)
            end
            b[i] = @dec128 (1 - 2tj) / s
        end
    elseif k == 2
        for i in eachindex(b)
            tj = 0
            th = @dec128 π * i / (s+1)
            for j in 1:div(s+1,2)
                tj += @dec128 sin((2j-1) * th) / (2j - 1)
            end
            b[i] = @dec128 (2tj * sin(th)) / (s+1)
        end
    else
        error("GaussChebyshevQuadrature: k=", k, " not supported.")
    end

    if !shift unshift!(b,c) end

    Quadrature(2s-2,c,b)
end

function get_nodes_lobatto_chebyshev(s, T=Float64)
    x = zeros(T, s)
    for i in eachindex(x)
        x[i] = @dec128 ( cos( (i-1)*π / (s-1) ) + 1 ) / 2
    end
    return x[end:-1:1]
end

function LobattoChebyshevQuadrature(s, T=Float64)
    c = get_nodes_lobatto_chebyshev(s)
    b = zeros(c)

    # TODO compute weights

    error("LobattoChebyshevQuadrature not implemented, yet!")
end

function ClenshawCurtisQuadrature(s, T=Float64; shift=true)
    c, b = FastTransforms.clenshawcurtis(s, zero(T), zero(T))
    c = flipdim(c,1)
    if shift shift!(b,c) end
    Quadrature(s^2, c, b)
end


#! /usr/bin/env julia

using QuadGK

function Ei(x, minfloat=1e-3, maxfloat=100)
    f = t -> exp(-t) / t  # inline function
    if x > 0
        return quadgk(f, -x, -minfloat)[1]
             + quadgk(f, minfloat, maxfloat)[1]
    else
        return quadgk(f, -x, maxfloat)[1]
    end
end

X = linspace(-1, 1, 1000)   # 1000 points
Y = [ Ei(x) for x in X ]

using Winston
plot(X, Y)
title("The function Ei(x)")
xlabel("x"); ylabel("y")
savefig("figures/Ei_integral.png")


function HMMforward(v,phghm,ph1,pvgh;UseLogArray=true)
    #HMMFORWARD HMM Forward Pass
    # [alpha,loglik]=HMMforward(v,phghm,ph1,pvgh)
    #
    # Inputs:
    # v : visible (observation) sequence being a vector v=[2 1 3 3 1 ...]
    # phghm : homogeneous transition distribution phghm(i,j)=p(h(t)=i|h(t-1)=j)
    # ph1 : initial distribution
    # pvgh : homogeneous emission distribution pvgh(i,j)=p(v(t)=i|h(t)=j)
    #
    # Outputs:
    # alpha : p(h(t)|v(1:t))
    #           p(h(t)|v(1:t)) = p(v(t)|h(t)) * sum_{h(t-1)} {p(h(t)|h(t-1)) * p(h(t-1)|v(1:t-1))}
    #           alpha(t) = p(h(t)|v(1:t))
    #           alpha(t-1) = p(h(t-1)|v(1:t-1))
    # loglik : sequence log likelihood log p(v(1:T))

    T=length(v); H=length(ph1);

    pvghtrans=pvgh'
    z=zeros(1,T) # local normalisation factors
    alpha=zeros(H,T)
    alpha[:,1] = pvghtrans[:,v[1]].*ph1
    z[1]=sum(alpha[:,1])
    alpha[:,1]=alpha[:,1]./z[1]
    for t=2:T
        alpha[:,t]=pvghtrans[:,v[t]].*(phghm*alpha[:,t-1])
        z[t]=sum(alpha[:,t])
        alpha[:,t]=alpha[:,t]./z[t]
    end
    loglik = sum(log.(z)) # log likelihood

    return alpha, loglik

end


# # [Lotka-Volterra Model](@id predation_example)
#
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/examples/predation/lotka-volterra.ipynb)

using AlgebraicPetri

using Petri
using OrdinaryDiffEq
using Plots

using Catlab
using Catlab.Theories
using Catlab.Programs
using Catlab.CategoricalAlgebra.ShapeDiagrams
using Catlab.WiringDiagrams
using Catlab.Graphics

display_wd(ex) = to_graphviz(ex, orientation=LeftToRight, labels=true);

# #### Step 1: Define the building block Petri nets needed to construct the model

birth_petri = PetriCospan([1], Petri.Model([1], [(Dict(1=>1), Dict(1=>2))]), [1]);
Graph(decoration(birth_petri))
#-
predation_petri = PetriCospan([1,2], Petri.Model(1:2, [(Dict(1=>1,2=>1), Dict(2=>1.15))]), [2]);
Graph(decoration(predation_petri))
#-
death_petri = PetriCospan([1], Petri.Model([1], [(Dict(1=>1), Dict{Int,Int}())]), [1]);
Graph(decoration(death_petri))

# #### Step 2: Define a presentation of the free biproduct category
# that encodes the domain specific information

@present Predation(FreeBiproductCategory) begin
    prey::Ob
    predator::Ob
    birth::Hom(prey,prey)
    predation::Hom(prey⊗predator,predator)
    death::Hom(predator,predator)
end;

rabbits,wolves,birth,predation,death = generators(Predation);

F(ex) = functor((PetriCospanOb, PetriCospan), ex, generators=Dict(
                 rabbits=>PetriCospanOb(1),wolves=>PetriCospanOb(1),
                 birth=>birth_petri, predation=>predation_petri, death=>death_petri));

# #### Step 3: Generate models using the hom expression or program notations

lotka_volterra = (birth ⊗ id(wolves)) ⋅ predation ⋅ death
lotka_petri = decoration(F(lotka_volterra))
display_wd(lotka_volterra)
#-
Graph(lotka_petri)

# Generate appropriate vector fields, define parameters, and visualize solution

u0 = [40, 2];
p = [4.0, 1.0, 2.0];
prob = ODEProblem(lotka_petri,u0,(0.0,8.0),p);
sol = solve(prob,Tsit5(),abstol=1e-6);
plot(sol)

# There is also a second syntax that is easier to write for programmers
# than the hom expression syntax. Here is an example of the same model
# as before along with a test of equivalency

lotka_volterra2 = @program Predation (r::prey, w::predator) begin
  r_2 = birth(r)
  w_2 = predation(r_2, w)
  return death(w_2)
end
lotka_petri2 = decoration(F(to_hom_expr(FreeBiproductCategory, lotka_volterra2)))
lotka_petri == lotka_petri2

# #### Step 4: Extend your presentation to handle more complex phenomena
# such as a small food chain

@present DualPredation <: Predation begin
    Predator::Ob
    Predation::Hom(predator⊗Predator,Predator)
    Death::Hom(Predator,Predator)
end;

fish,Fish,Shark,birth,predation,death,Predation,Death = generators(DualPredation);

F(ex) = functor((PetriCospanOb, PetriCospan), ex, generators=Dict(
                 fish=>PetriCospanOb(1),Fish=>PetriCospanOb(1),
                 birth=>birth_petri, predation=>predation_petri, death=>death_petri,
                 Shark=>PetriCospanOb(1),Predation=>predation_petri, Death=>death_petri));

# Define a new model where fish are eaten by Fish which are then eaten by Sharks

dual_lv = @program DualPredation (fish::prey, Fish::predator, Shark::Predator) begin
  f_2 = birth(fish)
  F_2 = predation(f_2, Fish)
  F_3 = death(F_2)
  S_2 = Predation(F_3, Shark)
  S_3 = Death(S_2)
end
display_wd(dual_lv)
#-
dual_lv_petri = decoration(F(to_hom_expr(FreeBiproductCategory, dual_lv)))
Graph(dual_lv_petri)

# Generate a new solver, provide parameters, and analyze results

u0 = [40, 2, 1];
p = [4.0, 1.0, 2.0, 2.5, 1.5];
prob = ODEProblem(dual_lv_petri,u0,(0.0,8.0),p);
sol = solve(prob,Tsit5(),abstol=1e-6);
plot(sol)


"""
    Ormsby(; <keyword arguments>)

Create a Ormsby wavelet sampled every dt seconds with corner frequencies
defined by the vector f = [f1, f2, f3, f4]. The final wavelet is multiplied by
a Hamming window.

# Arguments

**Keyword arguments**

* `dt::Real=0.002`: sampling interval in secs.
* `f::Vector{Real}=[2.0, 10.0, 40.0, 60.0]`: corner frequencies in Hz.

      ^
    1 |     ***************
      |    *               *
      |   *                 *
      |  *                   *
      | *                     *
      -----------------------------> f
        f1  f2           f3  f4

# Example
```julia
julia> w = Ormsby(); plot(w);
```
"""

function Ormsby{T<:Real}(; dt::Real=0.002, f::Vector{T}=[2.0, 10.0, 40.0, 60.0])

    f1 = f[1]
    f2 = f[2]
    f3 = f[3]
    f4 = f[4]

    fc = (f2+f3)/2.0
    nw = 2.2/(fc*dt)
    nc = floor(Int, nw/2)
    t = dt*collect(-nc:1:nc)
    nw = 2*nc + 1
    a4 = (pi*f4)^2/(pi*(f4-f3))
    a3 = (pi*f3)^2/(pi*(f4-f3))
    a2 = (pi*f2)^2/(pi*(f2-f1))
    a1 = (pi*f1)^2/(pi*(f2-f1))

    u = a4*(sinc(f4*t)).^2 - a3*(sinc(f3*t)).^2
    v = a2*(sinc(f2*t)).^2 - a1*(sinc(f1*t)).^2

    w = u - v
    w = w.*Hamming(nw)/maximum(w)

end


include("utils.jl")

using Plots
using LaTeXStrings
using Statistics

n = 100
maxk = 50
repeats = 20

A = hilbert(n)
ks = 1:maxk
λ = 1e-4

errors = zeros(maxk, repeats, 3)

methods = [
    ridge_weights(A, λ),
    uniform_weights(A),
    squared_norms_weights(A)
]

for k in ks
    for rep = 1:repeats
        for (i, w) in enumerate(methods)
            C = samplemat(A, w, k)
            C = unique(C, dims=2)
            errors[k, rep, i] = approxerror(A, C)
        end
    end
end

best_approx_errors = zeros(maxk)

for k in ks
    C = best_approx(A, k)
    best_approx_errors[k] = approxerror(A, C)
end

means = mean(errors, dims=2)[:, 1, :]
stds = std(errors, dims=2)[:, 1, :]

plot(xlabel=L"k", ylabel=L"\textrm{error}")
plot!(ks, means[:, 2], yaxis=:log, label="uniform", w=2)
plot!(ks, means[:, 3], yaxis=:log, label="squared norm", w=2)
plot!(ks, means[:, 1], yaxis=:log, label="ridge", w=2)
p = plot!(ks, best_approx_errors, yaxis=:log, label="best", c=:gray, w=2, ls=:dashdot)
savefig(p, "./plot.pdf")
display(p)  # sometimes the display doesn't work


## quadrature.jl : nystrom integration methods

abstract type QuadratureRule end

struct GaussLegendre <: QuadratureRule end
struct EOLE <: QuadratureRule end
struct Simpson <: QuadratureRule end
struct Midpoint <: QuadratureRule end
struct Trapezoidal <: QuadratureRule end

# get nodes and weights of Gauss-Legendre quadrature on [a,b]
function get_nodes_and_weights(n,a,b,q::GaussLegendre)
    nodes, weights = gausslegendre(n)
    weights = (b-a)/2*weights
    nodes = (b-a)/2*nodes .+ (a+b)/2
    return nodes, weights
end

# get nodes and weights of structured grid on [a,b]
function get_nodes_and_weights(n,a,b,q::EOLE)
    nodes = range(a,stop = b,length = n)
    weights = (b-a)/n*fill(1.0, size(nodes))
    return nodes, weights
end

# get nodes and weights of Simpson's rule on [a,b]
function get_nodes_and_weights(n,a,b,q::Simpson)
    n % 2 == 0 || begin
        warn("to use Simpson's rule, n must be even (received $(n)). I will continue with n = $(n+1)")
        n += 1
    end
    Δx = (b-a)/n
    nodes = a:Δx:b
    weights = repeat(2:2:4, outer=Int(n/2))
    weights[1] = 1
    push!(weights,1)
    weights *= Δx/3
    return nodes, weights
end

# get nodes and weights of Midpoint rule on [a,b]
function get_nodes_and_weights(n,a,b,q::Midpoint)
    Δx = (b-a)/n
    nodes = a+Δx/2:Δx:b-Δx/2
    weights = Δx*fill(1.0, size(nodes))
    return nodes, weights
end

# get nodes and weights of Trapezoidal rule on [a,b]
function get_nodes_and_weights(n,a,b,q::Trapezoidal)
    Δx = (b-a)/n
    nodes = a:Δx:b
    weights = Δx*fill(1.0, size(nodes))
    weights[1] /= 2
    weights[end] /= 2
    return nodes, weights
end


"""
    logit10(x::Real)
Function transforming interval [0,1] to [-Inf, Inf] using logit transformation.
"""
function logit10(x::Real)
    log10(x / (1.0 - x))
end

"""
    logistic10(x::Real)
Function transforming interval [-Inf, Inf] to [0,1] using logistic transformation.
Inverse function for [`logit10`](@ref).
"""
function logistic10(x::Real)
    exp10(x) / (exp10(x) + 1.0)
end

"""
    scaling(x::Real, scale::Symbol = :direct)
Transforms values from specific scale to range [-Inf, Inf] based on option.

## Return
Transformed value.

## Arguments
* `x`: input value.
* `scale`: transformation type: `:direct` (`:lin`), `:log`, `:logit`.
"""
function scaling(x::Real, scale::Symbol = :direct)
    if scale == :direct || scale == :lin
        return x
    elseif scale == :log
        return log10(x)
    elseif scale == :logit
        return logit10(x)
    else
        throw(DomainError(scale, "scale type is not supported"))
    end
end

"""
    unscaling(x::Real, scale::Symbol = :direct)
Transforms values from [-Inf, Inf] to specific scale based on option. Inverse function
for [`scaling`](@ref).

## Return
Transformed value.

## Arguments
* `x`: input value.
* `scale`: transformation type: `:direct` (`:lin`), `:log`, `:logit`.
"""
function unscaling(x::Real, scale::Symbol = :direct)
    if scale == :direct || scale == :lin
        return x
    elseif scale == :log
        return exp10(x)
    elseif scale == :logit
        return logistic10(x)
    else
        throw(DomainError(scale, "scale type is not supported"))
    end
end

scaling(x_tup::Tuple{Float64,Float64}, scale::Symbol = :direct) = scaling.(x_tup, scale)

unscaling(x_tup::Tuple{Float64,Float64}, scale::Symbol = :direct) = unscaling.(x_tup, scale)

scaling(::Nothing, ::Symbol) = nothing

unscaling(::Nothing, ::Symbol) = nothing

"""
    function get_endpoint(
        theta_init::Vector{Float64},
        theta_num::Int,
        loss_func::Function,
        method::Symbol,
        direction::Symbol = :right;

        loss_crit::Float64 = 0.0,
        scale::Vector{Symbol} = fill(:direct, length(theta_init)),
        theta_bounds::Vector{Tuple{Float64,Float64}} = unscaling.(
            fill((-Inf, Inf), length(theta_init)),
            scale
            ),
        scan_bound::Float64 = unscaling(
            (direction==:left) ? -9.0 : 9.0,
            scale[theta_num]
            ),
        scan_tol::Float64 = 1e-3,
        loss_tol::Float64 = 0.,
        local_alg::Symbol = :LN_NELDERMEAD,
        max_iter::Int = 10^5,
        loss_grad::Union{Function, Symbol} = :EMPTY,
        silent::Bool = false
    )

Calculates confidence interval's right or left endpoints for a given parameter `theta_num`.

## Return
[`EndPoint`](@ref) object storing confidence interval's endpoint and intermediate profile points.

## Arguments
- `theta_init`: starting values of parameter vector ``\\theta``. The starting values should not necessary be the optimum values of `loss_func` but `loss_func(theta_init)` should be lower than `loss_crit`.
- `theta_num`: index of vector component for identification: `theta_init(theta_num)`.
- `loss_func`: loss function ``\\Lambda\\left(\\theta\\right)`` for profile likelihood-based (PL) identification. Usually we use log-likelihood for PL analysis: ``\\Lambda( \\theta ) = - 2 ln\\left( L(\\theta) \\right)``.
- `method`: computational method to estimate confidence interval's endpoint. Currently the following methods are implemented: `:CICO_ONE_PASS`, `:LIN_EXTRAPOL`, `:QUADR_EXTRAPOL`.
- `direction`: `:right` or `:left` endpoint to estimate.

## Keyword arguments
- `loss_crit` : critical level of loss function. Confidence interval's endpoint value is the intersection point of profile likelihood and `loss_crit` level.
- `scale` : vector of scale transformations for each parameters' component. Possible values: `:direct` (`:lin`), `:log`, `:logit`. This option can speed up the optimization, especially for wide `theta_bounds`. The default value is `:direct` (no transformation) for all parameters.
- `theta_bounds` : vector of tuple `(lower_bound, upper_bound)` for each parameter. Bounds define the ranges for possible parameter values. Default bounds are `(-Inf,Inf)`.
- `scan_bound` : value which states the area of confidence point analysis.
- `scan_tol` : Absolute tolerance for `theta_num` parameter used as termination criterion.
- `loss_tol` : Absolute tolerance controlling `loss_func` closenes to `loss_crit` (termination criterion). Currently doesn't work for `:CICO_ONE_PASS` method because of limitation in `LN_AUGLAG` interface.
- `local_alg` : algorithm of optimization. Derivative-free and gradient-based algorithms form NLopt package.
- `max_iter` : maximal number of fitter iterations. If reaches the result status will be `:MAX_ITER_STOP`.
- `loss_grad` : For gradient optimization methods it is necessary to set how the gradient of `loss_func` should be calculated.
    There are options:
    - `:EMPTY` (default) no gradient is set. It works only for gradient-free methods.
    - `:AUTODIFF` means autodifferentiation from `ForwardDiff` package is used.
    - `:FINITE` means finite difference method from `Calculus` is used.
    - It is also possible to set gradient function here `function(x::Vector{Float64})` which returns gradient vector.
- `silent` : Boolean argument declaring whether we display the optimization progress. Default is `false`

See also [`get_interval`](@ref)
"""
function get_endpoint(
    theta_init::Vector{Float64},
    theta_num::Int,
    loss_func::Function,
    method::Symbol,
    direction::Symbol = :right;

    loss_crit::Float64 = 0.0,
    scale::Vector{Symbol} = fill(:direct, length(theta_init)), # :direct, :lin, :log, :logit
    theta_bounds::Vector{Tuple{Float64,Float64}} = unscaling.(
        fill((-Inf, Inf), length(theta_init)),
        scale
        ),
    scan_bound::Float64 = unscaling(
        (direction==:left) ? -9.0 : 9.0,
        scale[theta_num]
        ),
    scan_tol::Float64 = 1e-3,
    loss_tol::Float64 = 0.,
    local_alg::Symbol = :LN_NELDERMEAD,
    max_iter::Int = 10^5,
    loss_grad::Union{Function, Symbol} = :EMPTY,
    silent::Bool = false,
    #kwargs... # other options for get_right_endpoint
)
    isLeft = direction == :left
    n_theta = length(theta_init)

    # checking arguments
    # theta_bound[1] < theta_init < theta_bound[2]
    theta_init_outside_theta_bounds = .! [theta_bounds[i][1] < theta_init[i] < theta_bounds[i][2] for i in 1:length(theta_init)]
    if any(theta_init_outside_theta_bounds)
        throw(ArgumentError("theta_init is outside theta_bound: $(findall(theta_init_outside_theta_bounds))"))
    end
    # scan_bound should be within theta_bounds
    !(theta_bounds[theta_num][1] < scan_bound < theta_bounds[theta_num][2]) &&
        throw(ArgumentError("scan_bound are outside of the theta_bounds $(theta_bounds[theta_num])"))
    # theta_init should be within scan_bound
    if (theta_init[theta_num] >= scan_bound && !isLeft) || (theta_init[theta_num] <= scan_bound && isLeft)
        throw(ArgumentError("init values are outside of the scan_bound $scan_bound"))
    end
    # 0 <= theta_bound[1] for :log
    less_than_zero_theta_bounds = (scale .== :log) .& [theta_bounds[i][1] < 0 for i in 1:length(theta_init)]
    if any(less_than_zero_theta_bounds)
        throw(ArgumentError(":log scaled theta_bound min is negative: $(findall(less_than_zero_theta_bounds))"))
    end
    # 0 <= theta_bounds <= 1 for :logit
    less_than_zero_theta_bounds = (scale .== :logit) .& [theta_bounds[i][1] < 0 || theta_bounds[i][2] > 1 for i in 1:length(theta_init)]
    if any(less_than_zero_theta_bounds)
        throw(ArgumentError(":logit scaled theta_bound min is outside range [0,1]: $(findall(less_than_zero_theta_bounds))"))
    end
    # loss_func(theta_init) < loss_crit
    !(loss_func(theta_init) < loss_crit) &&
        throw(ArgumentError("Check theta_init and loss_crit: loss_func(theta_init) should be < loss_crit"))

    # set counter in the scope
    counter::Int = 0
    # set supreme, maximal or minimal value of scanned parameter inside critical
    supreme_gd = nothing

    # progress info
    prog = ProgressUnknown("$direction CP counter:"; spinner=false, enabled=!silent, showspeed=true)

    # transforming loss
    theta_init_gd = scaling.(theta_init, scale)
    if isLeft theta_init_gd[theta_num] *= -1 end # change direction

    function loss_func_gd(theta_gd)
        theta_g = copy(theta_gd)
        if isLeft theta_g[theta_num] *= -1 end # change direction
        theta = unscaling.(theta_g, scale)
        # calculate function
        loss_norm = loss_func(theta) - loss_crit

        # update counter
        counter += 1
        # update supreme ?
        update_supreme = (loss_norm < 0.) &&
            (typeof(supreme_gd)==Nothing || (theta_gd[theta_num] > supreme_gd))
        if update_supreme
            supreme_gd = theta_gd[theta_num]
        end
        # display current
        supreme = if (isLeft && typeof(supreme_gd)!==Nothing)
            unscaling(-supreme_gd, scale[theta_num])
        else
            unscaling(supreme_gd, scale[theta_num])
        end
        ProgressMeter.update!(prog, counter, spinner="⠋⠙⠹⠸⠼⠴⠦⠧⠇⠏"; showvalues = [(:supreme,round(supreme; sigdigits=4))])

        return loss_norm
    end
    theta_bounds_gd = scaling.(theta_bounds, scale)
    if isLeft theta_bounds_gd[theta_num] = (-1*theta_bounds_gd[theta_num][2], -1*theta_bounds_gd[theta_num][1]) end # change direction

    scan_bound_gd = scaling(scan_bound, scale[theta_num])
    if isLeft scan_bound_gd *= -1 end # change direction

    # transform gradient
    loss_grad_gd = loss_grad
    loss_grad_gd = if isa(loss_grad, Function)
        function(theta_gd)
            theta_g = copy(theta_gd)
            if isLeft theta_g[theta_num] *= -1 end # change direction
            theta = unscaling.(theta_g, scale)
            loss_grad_value = loss_grad(theta)
            loss_grad_value_gd = zeros(n_theta)

            for i in 1:n_theta
                if scale[i] == :log
                    loss_grad_value_gd[i] = loss_grad_value[i] * theta[i] * log(10.)
                elseif scale[i] == :logit
                    loss_grad_value_gd[i] = loss_grad_value[i] * theta[i] * (1. - theta[i]) * log(10.)
                else # for :direct and :lin
                    loss_grad_value_gd[i] = loss_grad_value[i]
                end
            end

            if isLeft loss_grad_value_gd[theta_num] *= -1 end # change direction

            return loss_grad_value_gd
        end
    else
        loss_grad
    end

    # calculate endpoint using base method
    (optf_gd, pp_gd, status) = get_right_endpoint(
        theta_init_gd,
        theta_num,
        loss_func_gd,
        Val(method);
        theta_bounds = theta_bounds_gd,
        scan_bound = scan_bound_gd,
        scan_tol,
        loss_tol,
        local_alg,
        max_iter,
        loss_grad = loss_grad_gd
        #kwargs...
    )

    # transforming back
    if (isLeft && typeof(optf_gd)!==Nothing) optf_gd *= -1 end # change direction
    optf = unscaling(optf_gd, scale[theta_num])
    temp_fun = (pp::ProfilePoint) -> begin
        if isLeft pp.params[theta_num] *= -1 end # change direction
        ProfilePoint(
            unscaling(pp.params[theta_num], scale[theta_num]),
            pp.loss + loss_crit,
            unscaling.(pp.params, scale),
            pp.ret,
            pp.counter
        )
    end
    pps = [ temp_fun(pp_gd[i]) for i in 1:length(pp_gd) ]
    # transforming supreme back
    supreme = if (isLeft && typeof(supreme_gd)!==Nothing)
        unscaling(-supreme_gd, scale[theta_num])
    else
        unscaling(supreme_gd, scale[theta_num])
    end
    ProgressMeter.finish!(prog)

    EndPoint(optf, pps, status, direction, counter, supreme)
end

"""
    function get_endpoint(
        theta_init::Vector{Float64},
        scan_func::Function,
        loss_func::Function,
        method::Symbol,
        direction::Symbol = :right;

        loss_crit::Float64 = 0.0,
        scale::Vector{Symbol} = fill(:direct, length(theta_init)),
        theta_bounds::Vector{Tuple{Float64,Float64}} = unscaling.(
            fill((-Inf, Inf), length(theta_init)),
            scale
            ),
        scan_bound::Float64 = unscaling(
            (direction==:left) ? -9.0 : 9.0,
            scale[theta_num]
            ),
        scan_tol::Float64 = 1e-3,
        loss_tol::Float64 = 0.,
        local_alg::Symbol = :LN_NELDERMEAD,
        kwargs...
        )

Calculates confidence interval's right or left endpoints for a function of parameters `scan_func`.

## Return
[`EndPoint`](@ref) object storing confidence interval's endpoint and intermediate profile points.

## Arguments
- `theta_init`: starting values of parameter vector ``\\theta``. The starting values should not necessary be the optimum values of `loss_func` but `loss_func(theta_init)` should be lower than `loss_crit`.
- `scan_func`: function of parameters.
- `loss_func`: loss function ``\\Lambda\\left(\\theta\\right)`` for profile likelihood-based (PL) identification. Usually we use log-likelihood for PL analysis: ``\\Lambda( \\theta ) = - 2 ln\\left( L(\\theta) \\right)``.
- `method`: computational method to estimate confidence interval's endpoint. Currently the only supported method is: `:CICO_ONE_PASS`.
- `direction`: `:right` or `:left` endpoint to estimate.

## Keyword arguments
- `loss_crit`: critical level of loss function. Confidence interval's endpoint value is the intersection point of profile likelihood and `loss_crit` level.
- `scale`: vector of scale transformations for each parameters' component. Possible values: `:direct` (`:lin`), `:log`, `:logit`. This option can speed up the optimization, especially for wide `theta_bounds`. The default value is `:direct` (no transformation) for all parameters.
- `theta_bounds`: vector of tuple `(lower_bound, upper_bound)` for each parameter. Bounds define the ranges for possible parameter values. Default bounds are `(-Inf,Inf)`.
- `scan_bound`: value which states the area of confidence point analysis.
- `scan_tol`: Absolute tolerance for `theta_num` parameter used as termination criterion.
- `loss_tol`: Absolute tolerance controlling `loss_func` closenes to `loss_crit` (termination criterion). Currently doesn't work for `:CICO_ONE_PASS` method because of limitation in `LN_AUGLAG` interface.
- `local_alg`: algorithm of optimization. Derivative-free and gradient-based algorithms form NLopt package.
- `max_iter` : maximal number of fitter iterations. If reaches the result status will be `:MAX_ITER_STOP`.
- `scan_grad` : For gradient optimization methods it is necessary to set how the gradient of `scan_func` should be calculated.
    - `:EMPTY` (default) no gradient is set. It works only for gradient-free methods.
    - `:AUTODIFF` means autodifferentiation from `ForwardDiff` package is used.
    - `:FINITE` means finite difference method from `Calculus` is used.
    - `function(x::Vector{Float64})` which returns gradient vector.
- `loss_grad` : For gradient optimization methods it is necessary to set how the gradient of `loss_func` should be calculated.
    - `:EMPTY` (default) no gradient is set. It works only for gradient-free methods.
    - `:AUTODIFF` means autodifferentiation from `ForwardDiff` package is used.
    - `:FINITE` means finite difference method from `Calculus` is used.
    - `function(x::Vector{Float64})` which returns gradient vector.
- `silent` : Boolean argument declaring whether we display the optimization progress. Default is `false`

See also [`get_interval`](@ref)
"""
function get_endpoint(
    theta_init::Vector{Float64},
    scan_func::Function,
    loss_func::Function,
    method::Symbol,
    direction::Symbol = :right;

    loss_crit::Float64 = 0.0,
    # :direct, :lin, :log, :logit
    scale::Vector{Symbol} = fill(:direct, length(theta_init)),
    theta_bounds::Vector{Tuple{Float64,Float64}} = unscaling.(
        fill((-Inf, Inf), length(theta_init)),
        scale
        ),
    scan_bound::Float64 = (direction==:left) ? -1e9 : 1e9, # log scan bound is not implemented
    scan_tol::Float64 = 1e-3,
    loss_tol::Float64 = 0.,
    local_alg::Symbol = :LN_NELDERMEAD,
    max_iter::Int = 10^5,
    scan_grad::Union{Function, Symbol} = :EMPTY,
    loss_grad::Union{Function, Symbol} = :EMPTY,
    silent::Bool = false,
    #kwargs... # other options for get_right_endpoint
)
    isLeft = direction == :left
    n_theta = length(theta_init)
    loss_init = loss_func(theta_init)

    # checking arguments
    # theta_bound[1] < theta_init < theta_bound[2]
    theta_init_outside_theta_bounds = .! [theta_bounds[i][1] < theta_init[i] < theta_bounds[i][2] for i in 1:length(theta_init)]
    if any(theta_init_outside_theta_bounds)
        throw(ArgumentError("theta_init is outside theta_bound: $(findall(theta_init_outside_theta_bounds))"))
    end

    # 0 <= theta_bound[1] for :log
    less_than_zero_theta_bounds = (scale .== :log) .& [theta_bounds[i][1] < 0 for i in 1:length(theta_init)]
    if any(less_than_zero_theta_bounds)
        throw(ArgumentError(":log scaled theta_bound min is negative: $(findall(less_than_zero_theta_bounds))"))
    end
    # 0 <= theta_bounds <= 1 for :logit
    less_than_zero_theta_bounds = (scale .== :logit) .& [theta_bounds[i][1] < 0 || theta_bounds[i][2] > 1 for i in 1:length(theta_init)]
    if any(less_than_zero_theta_bounds)
        throw(ArgumentError(":logit scaled theta_bound min is outside range [0,1]: $(findall(less_than_zero_theta_bounds))"))
    end
    # loss_func(theta_init) < loss_crit
    !(loss_init < loss_crit) &&
        throw(ArgumentError("Check theta_init and loss_crit: loss_func(theta_init) should be < loss_crit"))

    # set counter in the scope
    prog = ProgressUnknown("Fitter counter:"; spinner=false, enabled=!silent, showspeed=true)
    counter::Int = 0
    # set supreme, maximal or minimal value of loss_func inside critical
    supreme_gd = nothing
    scan_val_gd = nothing

    # transforming
    theta_init_gd = scaling.(theta_init, scale)

    function scan_func_gd(theta_gd)
        theta = unscaling.(theta_gd, scale)
        scan_val = scan_func(theta)
        scan_val_gd = isLeft ? (-1)*scan_val : scan_val

        return scan_val_gd
    end

    function loss_func_gd(theta_gd)
        #theta_g = copy(theta_gd) # why copy?
        theta = unscaling.(theta_gd, scale)
        loss_value = loss_func(theta) - loss_crit

        counter += 1
        #println("$loss_value => $supreme_gd")
        should_update = (loss_value < 0.) &&
            !isa(scan_val_gd, Nothing) &&
            (isa(supreme_gd, Nothing) || (scan_val_gd > supreme_gd)) &&
            !isa(scan_val_gd, ForwardDiff.Dual)
        if should_update
            supreme_gd = scan_val_gd
        end
        supreme = if isa(supreme_gd, Nothing)
            "-"
        elseif isLeft
            round(-supreme_gd; sigdigits=4)
        else
            round(supreme_gd; sigdigits=4)
        end
        ProgressMeter.update!(prog, counter; showvalues = [(:supreme,supreme)])

        return loss_value
    end
    theta_bounds_gd = scaling.(theta_bounds, scale)

    # TODO: transformed by scan_scale: scaling(scan_bound, scan_scale)
    scan_bound_gd = scan_bound
    if isLeft scan_bound_gd *= -1 end # change direction

    # transform gradients
    scan_grad_gd = if isa(scan_grad, Function)
        function(theta_gd)
            theta = unscaling.(theta_gd, scale)
            scan_grad_value = scan_grad(theta)
            scan_grad_value_gd = zeros(n_theta)

            for i in 1:n_theta
                if scale[i] == :log
                    scan_grad_value_gd[i] = scan_grad_value[i] * theta[i] * log(10.)
                elseif scale[i] == :logit
                    scan_grad_value_gd[i] = scan_grad_value[i] * theta[i] * (1. - theta[i]) * log(10.)
                else # for :direct and :lin
                    scan_grad_value_gd[i] = scan_grad_value[i]
                end
            end

            return scan_grad_value_gd
        end
    else
        scan_grad
    end

    loss_grad_gd = if isa(loss_grad, Function)
        function(theta_gd)
            theta = unscaling.(theta_gd, scale)
            loss_grad_value = loss_grad(theta)
            loss_grad_value_gd = zeros(n_theta)

            for i in 1:n_theta
                if scale[i] == :log
                    loss_grad_value_gd[i] = loss_grad_value[i] * theta[i] * log(10.)
                elseif scale[i] == :logit
                    loss_grad_value_gd[i] = loss_grad_value[i] * theta[i] * (1. - theta[i]) * log(10.)
                else # for :direct and :lin
                    loss_grad_value_gd[i] = loss_grad_value[i]
                end
            end

            return loss_grad_value_gd
        end
    else
        loss_grad
    end

    # calculate endpoint using base method
    (optf_gd, pp_gd, status) = get_right_endpoint(
        theta_init_gd,
        scan_func_gd,
        loss_func_gd,
        Val(method);

        theta_bounds = theta_bounds_gd,
        scan_bound = scan_bound_gd,
        scan_tol,
        loss_tol,
        local_alg,
        max_iter,
        scan_grad = scan_grad_gd,
        loss_grad = loss_grad_gd,
        #kwargs...
    )

    # transforming back
    temp_fun = (pp::ProfilePoint) -> begin
        value = isLeft ? (-1)*pp.value : pp.value # TODO: scan_scale
        ProfilePoint(
            value,
            pp.loss + loss_crit,
            unscaling.(pp.params, scale),
            pp.ret,
            pp.counter
        )
    end
    pps = [ temp_fun(pp_gd[i]) for i in 1:length(pp_gd) ]
    # transforming supreme back
    if (isLeft && typeof(supreme_gd)!==Nothing) supreme_gd *= -1 end # change direction

    if (isLeft && typeof(optf_gd)!==Nothing) optf_gd *= -1 end # change direction
    # optf = unscaling(optf_gd, scan_scale)
    optf = optf_gd

    EndPoint(optf, pps, status, direction, counter, supreme_gd)
end


module IntervalAnalysis
using Plots
using Base
using Random
using Symbolics
using Documenter

export Interval,*,+,-,subdivide

"""
    Interval

Defines a single interval between two real numbers.

# Examples
```jldoctest
julia> Interval(-1,1)
-1.0 ←-→ 1.0
```
"""
struct Interval
	"Lower interval value"
	l::Real
	"Upper interval value"
	u::Real
	Interval(l,u) = l > u ? new(u,l) : new(l,u)

end

Base.show(io::IO, x::Interval) = print(io,round(x.l,digits=6)," ←-→ ",round(x.u,digits=6))

"""

	+(x₁::Interval, x₂::Interval)

Adds two intervals

# Examples
```jldoctest
julia> Interval(-1,1) + Interval(-2,4)
-3.0 ←-→ 5.0
```
"""
function Base.:+(x₁::Interval, x₂::Interval)
	Interval(x₁.l + x₂.l,x₁.u + x₂.u)
    end

function Base.:+(x₁::Real, x₂::Interval)
	Interval(x₁ + x₂.l,x₁ + x₂.u)
    end

function Base.:+(x₁::Interval, x₂::Real)
	Interval(x₁.l + x₂,x₁.u + x₂)
    end

 """

	-(x₁::Interval, x₂::Interval)

Subtracts one interval from another

# Examples
```jldoctest
julia> Interval(-1,1) - Interval(-2,4)
-5.0 ←-→ 3.0
```
"""
function Base.:-(x₁::Interval, x₂::Interval)
	Interval(x₁.l - x₂.u,x₁.u - x₂.l)
    end

function Base.:-(x::Interval)
	Interval(-x.l,-x.u)
    end

function Base.:-(x₁::Interval, x₂::Real)
	Interval(x₁.l - x₂,x₁.u - x₂)
    end

    function Base.:-(x₁::Real, x₂::Interval)
	Interval(x₁ - x₂.l,x₁ - x₂.u)
    end

"""

	*(x₁::Interval, x₂::Interval)

Multiplies two intervals

# Examples
```jldoctest
julia> Interval(-1,1) * Interval(-2,4)
-4.0 ←-→ 4.0
```
"""
function Base.:*(x₁::Interval, x₂::Interval)
	if 0 ≤ x₁.l
		if 0 ≤ x₂.l
			return Interval(x₁.l*x₂.l,x₁.u*x₂.u)
		elseif x₂.l < 0 < x₂.u
			return Interval(x₁.u*x₂.l,x₁.u*x₂.u)
		elseif x₂.u ≤ 0
			return Interval(x₁.u*x₂.l,x₁.l*x₂.u)
		end
	elseif x₁.l < 0 < x₁.u
		if 0 ≤ x₂.l
			return Interval(x₁.l*x₂.u,x₁.u*x₂.u)
		elseif x₂.u ≤ 0
			return Interval(x₁.u*x₂.l,x₁.l*x₂.l)
		elseif x₂.l < 0 < x₂.u
			return Interval(min(x₁.l*x₂.u,x₁.u*x₂.l),max(x₁.l*x₂.l,x₁.u*x₂.u))
		end
	elseif x₁.u ≤ 0
		if 0 ≤ x₂.l
			return Interval(x₁.l*x₂.u,x₁.u*x₂.l)
		elseif x₂.l < 0 < x₂.u
			return Interval(x₁.l*x₂.u,x₁.l*x₂.l)
		elseif x₂.u ≤ 0
			return Interval(x₁.u*x₂.u,x₁.l*x₂.l)
		end
	end
end

function Base.:*(x₁::Interval,x₂::Real)
	Interval(x₁.l*x₂,x₁.u*x₂)
end

function Base.:*(x₁::Real,x₂::Interval)
	Interval(x₁*x₂.l,x₁*x₂.u)
end

"""

	/(x₁::Interval, x₂::Interval)

Divides two intervals

# Examples
```jldoctest
julia> Interval(-1,1) / Interval(-2,4)
-0.5 ←-→ 0.5
```
"""
function Base.:/(x₁::Interval, x₂::Interval)
	x₁ * Interval(1/x₂.u,1/x₂.l)
    end

function Base.:/(x₁::Real, x₂::Interval)
	Interval(x₁/x₂.u,x₁/x₂.l)
    end

function Base.:/(x₁::Interval, x₂::Real)
	Interval(x₁.l/x₂,x₁.u/x₂)
    end

function w(x::Interval)
	x.u-x.l
end

function w(X::Vector{Interval})
	widths = w.(X)
	d,i = findmax(widths)
end

function m(x::Interval)
	(x.u+x.l)/2
end

function m(X::Vector{Interval})
	m.(X)
end

function Base.sqrt(x::Interval)
	Interval(sqrt(x.l),sqrt(x.u))
end

function Base.:^(x::Interval,p::Integer)
	if x.l > 0 || p%2 != 0
		return Interval(x.l^p,x.u^p)
	elseif x.u < 0 && p%2 == 0
		return Interval(x.u^p,x.l^p)
	else
		return Interval(0,abs(x)^p)
	end
end

function Base.exp(x::Interval)
	Interval(exp(x.l),exp(x.u))
end

function Base.log(x::Interval)
	Interval(log(x.l),log(x.u))
end

function Base.log10(x::Interval)
	Interval(log10(x.l),log10(x.u))
end

function Base.log2(x::Interval)
	Interval(log2(x.l),log2(x.u))
end

function Base.sin(x::Interval)
	if x.l ≥ -(π/2) && x.u ≤ π/2
			return Interval(sin(x.l),sin(x.u))
	end
	return Interval(-1,1)
end

function Base.cos(x::Interval)
	if x.l ≥ 0 && x.u ≤ π
			return Interval(cos(x.l),cos(x.u))
	end
	return Interval(-1,1)
end

function Base.abs(x::Interval)
	max(abs(x.l),abs(x.u))
end

function qdist(x₁::Interval,x₂::Interval)
	max(abs(x₁.l-x₂.l),abs(x₁.u-x₂.u))
end

"""

	subdivide(x::Interval,N::Integer)

Subdivides an interval into N intervals.

# Examples
```jldoctest
julia> subdivide(Interval(0,1),5)
[0.0 ←-→ 0.2, 0.2 ←-→ 0.4, 0.4 ←-→ 0.6, 0.6 ←-→ 0.8, 0.8 ←-→ 1.0]
```
"""
function subdivide(x::Interval,N::Integer)
	A = Array{Interval}(undef,N)
	w_x = w(x)/N
	for j ∈ 1:N
		A[j] = Interval(x.l+(j-1)*w_x,x.l+j*w_x)
	end
	return A
end

function subdivide_vector(X::Vector{Interval},N::Integer)
	A = Array{Vector{Interval}}(undef,length(X))
	for i ∈ 1:length(X)
		A[i] = subdivide(X[i],N)
	end
	return A
end

function bisect_box(X::Vector{Interval})
	# bisects a box over the largest axis
	d,i = w(X)
	X_s = subdivide(X[i],2)
	X₁ = copy(X)
	X₁[i] = X_s[1]
	X₂ = copy(X)
	X₂[i] = X_s[2]
	return X₁,X₂
end

function hull(X::Vector{Interval})
	l = [x.l for x ∈ X]
	u = [x.u for x ∈ X]
	Interval(minimum(l),maximum(u))
end

function refinement(x,f,N)
	X = subdivide(x,N)
	Y = f.(X)
	hull(Y)
end

function plot_refinement(x,f,∇f,N)
	plot(f,x.l,x.u,c="black",linewidth=3,label="")
	X = subdivide(x,N)
	Y = f.(X)
	Yₘᵥ = mean_value_form.(f,∇f,X)
	for i ∈ 1:length(X)
		if i == 1
			plot!([X[i].l,X[i].u],[Y[i].l,Y[i].l],fillrange=[Y[i].u,Y[i].u],linewidth=0,c="black",label="Interval Extension",fillalpha=0.15)
			plot!([X[i].l,X[i].u],[Yₘᵥ[i].l,Yₘᵥ[i].l],fillrange=[Yₘᵥ[i].u,Yₘᵥ[i].u],linewidth=0,c="red",label="Mean Value Form",fillalpha=0.25)
		else
			plot!([X[i].l,X[i].u],[Y[i].l,Y[i].l],fillrange=[Y[i].u,Y[i].u],linewidth=0,c="black",label="",fillalpha=0.15)
			plot!([X[i].l,X[i].u],[Yₘᵥ[i].l,Yₘᵥ[i].l],fillrange=[Yₘᵥ[i].u,Yₘᵥ[i].u],linewidth=0,c="red",label="",fillalpha=0.25)
		end

	end
	display(plot!(dpi=600))
end

function skelboe_moore(f,X)

	fᵤ = f(X).u # smallest upper bound (AKA the only upper bound at this stage)

	# creating a store of intervals, mean values, and interval values
	L = [Dict("X"=>X,"f_m"=>f(m(X)),"F"=>f(X))]
	Lᵣ = []
	it_count = 0

	# until tolerance...
	for it ∈ 1:30

		# initialise list of intervals to sack off
		del_index =[]
		for i in 1:length(L)
			# if the lowest interval value is above the smallest upper bound..
			if L[i]["F"].l > fᵤ
				# remove interval, we know global minimum is not here...
				push!(del_index,i)
			end
		end

		# delete all the ones we said we would
		for d_i ∈ del_index
			push!(Lᵣ,L[d_i])
		end
		deleteat!(L,del_index)

		# pick the remaining interval with smallest mean value
		s = argmin([l["f_m"] for l ∈ L])

		# divide the box into two
		X_s = bisect_box(L[s]["X"])

		# get rid of the containing interval
		deleteat!(L,s)

		for x ∈ X_s
			# evaluate interval values
			F = f(x)

			# check if we find a smallest upper bound
			if F.u < fᵤ
				fᵤ = F.u
			end
			# add new interval information to store
			push!(L,Dict("X"=>x,"f_m"=>f(m(x)),"F"=>F))

		end

	end
	return L,Lᵣ
end

function plot_boxes(L,Lᵣ)
	for l ∈ L
		x1 = l["X"][1]
		x2 = l["X"][2]
		plot!(Shape([x1.l,x1.l,x1.u,x1.u] , [x2.u,x2.l,x2.l,x2.u]),label="",linewidth=1,alpha=0,linealpha=1)
	end
	for l ∈ Lᵣ
		x1 = l["X"][1]
		x2 = l["X"][2]
		plot!(Shape([x1.l,x1.l,x1.u,x1.u] , [x2.u,x2.l,x2.l,x2.u]),c="black",alpha=0.55,label="",linewidth=1)
	end
end

function mean_value_form(f,∇f,X)
	if typeof(X) == Interval
		D = ∇f(X)[1]
	else
		D = ∇f(X)
	end
	mean = m(X)
	fₘ = f(mean)
	if typeof(X) == Interval
		return fₘ + D*(X-mean)
	else
		return fₘ + sum(D[i]*(X[i]-mean[i]) for i ∈ 1:length(X))
	end

end

function build_gradient(f,len_x)
	if len_x > 1
		@variables x[1:len_x]
		∇f = Symbolics.gradient(f(x),[x[i] for i ∈ 1:len_x])
		∇f_expr = build_function(∇f, x)
		∇f = eval(∇f_expr[1])
		return ∇f
	else
		@variables x
		∇f = Symbolics.gradient(f(x),[x])
		∇f_expr = build_function(∇f, x)
		∇f = eval(∇f_expr[1])
		return ∇f
	end

end

function g(x1,x2)
	return x1*exp(x1+x2^2)-x2^2
end

function g_vec(X)
	return g(X[1],X[2])
end

# f(x) = x^3 - x^2
# ∇f = build_gradient(f,1)
# plot_refinement(Interval(0,1),f,∇f,20)
# print(Interval(-1,1) / Interval(-2,4))

end


import Base: exp, log, sqrt, ^
import Base: cos, sin, tan, cosh, sinh, tanh, acos, asin, atan, acosh, asinh, atanh

function exp(q::Quaternion)
    z = sqrt(q.q1*q.q1 + q.q2*q.q2 + q.q3*q.q3)
    return exp(q.q0)*(cos(z) + imag(q)*(z != zero(q.q0) ? sin(z)/z : one(q.q0)))
end

function log(q::Quaternion)
    z = abs(imag(q))
    return log(abs(q)) + ((imag(q)!=zero(q)) ? imag(q)/z*atan(z,q.q0) :
    q.q0 < zero(q.q0) ? quaternion(pi*im) : zero(q))
end

function cos(q::Quaternion)
    z = abs(imag(q))
    return cos(q.q0)*cosh(z)-imag(q)*sin(q.q0)*(z != zero(q.q0) ? sinh(z)/z : one(q.q0))
end

function sin(q::Quaternion)
    z = abs(imag(q))
    return sin(q.q0)*cosh(z)+imag(q)*cos(q.q0)*(z != zero(q.q0) ? sinh(z)/z : one(q.q0))
end

function tan(q::Quaternion)
    return sin(q)/cos(q)
end

function cosh(q::Quaternion)
    return (exp(q)+exp(-q))/2
end

function sinh(q::Quaternion)
    return (exp(q)-exp(-q))/2
end

function tanh(q::Quaternion)
    return sinh(q)/cosh(q)
end

function sqrt(q::Quaternion)
    return exp(log(q)/2)
end

function (^)(q::Quaternion, p::AbstractFloat)
    return exp(p*log(q))
end

function (^)(q::Quaternion, p::Quaternion)
    return exp(p*log(q))
end

function acos(q::Quaternion)
    x = imag(q)
    y = abs(x)
    z = y!=zero(q.q0) ? x/y : quaternion(one(q.q0)*im)

    return -z*log(q + z*sqrt(one(q.q0) - q*q))
    # return -z*log(q - z*sqrt(one(q.q0) - q*q))
end

function asin(q::Quaternion)
    x = imag(q)
    y = abs(x)
    z = y!=zero(q.q0) ? x/y : quaternion(one(q.q0)*im)

    return -z*log(z*q + sqrt(one(q.q0) - q*q))
    #= return -z*log(z*q - sqrt(one(q.q0) - q*q)) =#
end

function atan(q::Quaternion)
    x = imag(q)
    y = abs(x)
    z = y!=zero(q.q0) ? x/y : quaternion(one(q.q0)*im)

    return -z*log((one(q.q0)+z*q)/(one(q.q0)-z*q))/2
end


"""
    erlang_rng(k, λ, shape=1; seed=nothing)

Generate a `shape` element array of random variables from a Erlang_{`k`}(`λ`) distribution. Optionally you can set a specific seed.

# Notes

The pdf of an Erlang_{k}(λ) distribution is given as:

``f(x, k, λ) = \\frac{λ^k e^{-λx} x^{k-1}}{(k-1)!} \\quad x ≥ 0``

# Examples

```julia-repl
julia> erlang_rng(5, .5)
1-element Vector{Float64}:
 10.803989701023117

julia> erlang_rng(3, 1, (2,2))
2×2×1 Array{Float64, 3}:
[:, :, 1] =
 2.19956  4.18505
 5.46892  2.5633
```

# References

D. Goldsman, P. Goldsman. A first course in probability and statistics. 2021.

L. Martino, D. Luengo. Extremely efficient generation of Gamma random variables for α ≥ 1. 2013.
"""
function erlang_rng(k::Int, λ::Real, shape::Union{Int, Tuple{Vararg{Int}}}=1; seed::Union{Int, Nothing}=nothing)
    U = get_std_uniform((shape..., k), seed=seed)
    k_dim = ndims(U)  # k is final dimension
    X = (-1/λ) .* log.(prod(U, dims=k_dim))  # Here (-1/λ) represents mean
    return X
end


## Spectral clustering on artists
using DelimitedFiles
using SparseArrays
names = readlines("../2-unit-1-demos/artistsim.names")
data = readdlm("../2-unit-1-demos/artistsim.smat")
xy = readdlm("artistsim.xy")
A = sparse(Int.(data[2:end,1]).+1,
           Int.(data[2:end,2]).+1,
           Int.(data[2:end,3]),
           Int(data[1,1]),Int(data[1,1]))
n = size(A,1)
##
function simplepagerank(A,α,v)
  @assert(0 ≤ α < 1, "needs probably α")
  @assert(all(vi -> vi ≥ 0, v), "needs non-negative v")
  v = v ./ sum(v) # we can normalize for them.
  d = vec(sum(A,dims=2)) # compute the degrees
  x = copy(v) # start of with v
  nsteps = 2*ceil(Int,log(eps(1.0))/log(α)) # upper bound on steps
  for i=1:nsteps
    x = α*(A'*(x./d)) .+ (1-α).*v
  end
  return x./sum(x) # renormalize to probability
end
##
using MatrixNetworks
pr = MatrixNetworks.pagerank(A,0.85)
##
x = simplepagerank(A,0.85,ones(n)./n)
##
using LinearAlgebra
@show norm(x-pr)

##
@time pr = MatrixNetworks.pagerank(A,0.85)
@time x = simplepagerank(A,0.85,ones(n)./n)
@show norm(x-pr)
## Let's watch PageRank converge
using Plots
function pagerank_converence_animation(A,xy,α,nsteps)
  n = size(A,1)
  v = ones(n)/n
  @assert(0 ≤ α < 1, "needs probably α")
  @assert(all(vi -> vi ≥ 0, v), "needs non-negative v")
  v = v ./ sum(v) # we can normalize for them.
  d = vec(sum(A,dims=2)) # compute the degrees
  x = copy(v) # start of with v

  @assert(nsteps <= size(A,1))
  anim = @animate for i=1:nsteps
    x = α*(A'*(x./d)) .+ (1-α).*v
    f2 = copy(x)
    f2 = f2/sum(f2)
    p = sortperm(f2) # this gets things to draw in the right z-order
    scatter(xy[p,1],xy[p,2], marker_z=-abs.(log10.(f2[p])).^0.5,colorbar=false,alpha=0.5,
      color=:magma,
      size=(1200,1200),
      markerstrokecolor=nothing,markerstrokewidth=0,label="",framestyle=:none,background=:black)
  end
  return anim
end
anim = pagerank_converence_animation(A,xy,0.85,25)
gif(anim,"pagerank-convergence-artistsim.gif",fps=5)

## Change between iterates
function pagerank_converence_plot(A,xy,α,nsteps)
  n = size(A,1)
  v = ones(n)/n
  @assert(0 ≤ α < 1, "needs probably α")
  @assert(all(vi -> vi ≥ 0, v), "needs non-negative v")
  v = v ./ sum(v) # we can normalize for them.
  d = vec(sum(A,dims=2)) # compute the degrees
  x = copy(v) # start of with v
  y = copy(v)
  hist = zeros(0)
  for i=1:nsteps
    x = α*(A'*(x./d)) .+ (1-α).*v
    push!(hist, norm(y - x/sum(x),1))
    y = x/sum(x)
  end
  return hist
end
hist = pagerank_converence_plot(A,xy,0.85,100)
##
pyplot()
theme(:dark)
plot(hist, yscale=:log10, legend=false, size=(300,300))
xlabel!("Iteration")
ylabel!("1-norm change in normalized vector")
savefig("pagerank-convergence-artistsim-norm.pdf")
##
x85 = simplepagerank(A,0.85,ones(n)./n)
x99 = simplepagerank(A,0.99,ones(n)./n)
x5 = simplepagerank(A,0.5,ones(n)./n)
##

## Show a few different vectors
gr()
p=sortperm(x99)
scatter(xy[p,1],xy[p,2],marker_z=-abs.(log10.(x99[p])).^0.5,markerstrokewidth=0,
color=:magma)
##
gr()
p=sortperm(x5)
scatter(xy[p,1],xy[p,2],marker_z=-abs.(log10.(x5[p])).^0.5,markerstrokewidth=0,
color=:magma)

## Now show the entire graph
using GraphRecipes
using LinearAlgebra
using Plots

ei, ej = findnz(A)
graphplot(ei, ej, x =xy[:,1], y=xy[:,2],
  markercolor=:black, markerstrokecolor=:white,
  size=(1200,1200),dpi=300,
  markersize=0, linecolor=1, linealpha=0.01, linewidth=0.5,
  markeralpha=0.2,colorbar=false,
  axis_buffer=0.02, background=nothing)
p=sortperm(x5)
scatter!(xy[p,1],xy[p,2],marker_z=-abs.(log10.(x5[p])).^0.5,markerstrokewidth=0)
savefig("artistsim-pagerank-50.png")

##
ei, ej = findnz(A)
graphplot(ei, ej, x =xy[:,1], y=xy[:,2],
  markercolor=:black, markerstrokecolor=:white,
  size=(1200,1200),dpi=300,
  markersize=0, linecolor=1, linealpha=0.01, linewidth=0.5,
  markeralpha=0.2,colorbar=false,
  axis_buffer=0.02, background=nothing)
p=sortperm(x99)
scatter!(xy[p,1],xy[p,2],marker_z=-abs.(log10.(x99[p])).^0.5,markerstrokewidth=0)
savefig("artistsim-pagerank-99.png")


### A Pluto.jl notebook ###
# v0.12.16

using Markdown
using InteractiveUtils

# ╔═╡ e077ab16-22e5-11eb-15ed-5f0654bf7328
using Pkg, DrWatson

# ╔═╡ e75a9a06-22e5-11eb-2e36-8d4f62b830ed
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using Random, StatsBase, Distributions, Plots, LaTeXStrings
	Random.seed!(0)
end

# ╔═╡ ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
md"## Listing5.16"

# ╔═╡ fa68607e-22e5-11eb-0558-c9a4d9f77426
begin
	mu0, mu1a, mu1b, mu1c, sd = 15, 16, 18, 20, 2
	tauGrid = 5:0.1:25

	dist0 = Normal(mu0,sd)
	dist1a, dist1b, dist1c  = Normal(mu1a,sd), Normal(mu1b,sd), Normal(mu1c,sd)

	falsePositive = ccdf.(dist0,tauGrid)
	truePositiveA, truePositiveB, truePositiveC =
	    ccdf.(dist1a,tauGrid), ccdf.(dist1b,tauGrid), ccdf.(dist1c,tauGrid)
end

# ╔═╡ a80f7e58-3af9-11eb-04b3-57fcb9ee857c
begin
	plot(falsePositive, [truePositiveA truePositiveB truePositiveC],
		c=[:blue :red :green],
		label=[L"H1a: \mu_1 = 16" L"H1b: \mu_1 = 18" L"H1c: \mu_1 = 20"])
	plot!([0,1], [0,1], c=:black, ls=:dash, label="H0 = H1 = 15",
		xlims=(0,1), ylims=(0,1), xlabel=L"\alpha", ylabel="Power",
		ratio=:equal, legend=:bottomright)
end

# ╔═╡ 475ff888-22e6-11eb-2354-09f8f40a8e12
md"## End of listing5.16"

# ╔═╡ Cell order:
# ╟─ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
# ╠═e077ab16-22e5-11eb-15ed-5f0654bf7328
# ╠═e75a9a06-22e5-11eb-2e36-8d4f62b830ed
# ╠═fa68607e-22e5-11eb-0558-c9a4d9f77426
# ╠═a80f7e58-3af9-11eb-04b3-57fcb9ee857c
# ╟─475ff888-22e6-11eb-2354-09f8f40a8e12


#"Read" data for parametric sinusoidal "simulation" & parameter extraction
#-------------------------------------------------------------------------------

using CMDimData
using CMDimData.MDDatasets

const π = MathConstants.π

#==Emulate reading in simulated data file
===============================================================================#

#=COMMENT
The code below emulates a parametric "simulation" of a sinusoidal response where
the 𝜙, A, and 𝑓 parameters of `signal = A * sin(𝜔*t + 𝜙); 𝜔 = 2π*𝑓` are varied.

The parametric signal can therefore be fully represented as:
	signal(𝜙, A, 𝑓, t)
=#

#But really construct multidimensional DataRS dataset manually:
signal = fill(DataRS, PSweep("phi", [0, 0.5, 1] .* (π/4))) do 𝜙
	fill(DataRS, PSweep("A", [1, 2, 4] .* 1e-3)) do A
	#Inner-most sweep: need to specify element type (DataF1):
	#(Other (scalar) element types: DataInt/DataFloat/DataComplex)
	fill(DataRS{DataF1}, PSweep("freq", [1, 4, 16] .* 1e3)) do 𝑓
		𝜔 = 2π*𝑓
		T = 1/𝑓
		Δt = T/100 #Define resolution from # of samples per period
		Tsim = 4T #Simulated time

		t = DataF1(0:Δt:Tsim) #DataF1 creates a t:{y, x} container with y == x
		sig = A * sin(𝜔*t + 𝜙) #Still a DataF1 sig:{y, x=t} container
		return sig
end; end; end

return signal


#######################################################################################################################################################################################################
#
# Changes made to this function
# General
#     2020-Mar-30: migrate the function from SCOPE
#     2021-Oct-21: rename the function to average_transmittance
#     2022-Feb-02: fix documentation
# To do
#     TODO: figure out where does this equation comes from
#
#######################################################################################################################################################################################################
"""

    average_transmittance(α::FT, nr::FT) where {FT<:AbstractFloat}

Return the average transmittance of isotropic radiation across an interface between two dielectrics, given
- `α` angle of incidence
- `nr` Index of refraction

# References
- Stern (1964) Transmission of isotropic radiation across an interface between two dielectrics. Applied Optics 3(1): 111-113.
- Allen (1973) Transmission of isotropic light across a dielectric surface in two and three dimensions. Journal of the Optical Society of America 63(6): 664-666.
"""
function average_transmittance(α::FT, nr::FT) where {FT<:AbstractFloat}
    @assert 0 < α <= 90;

    # some shortcuts to avoid overly comlicated equation
    _a     = (nr + 1) ^ 2 / 2;
    _a³    = _a ^ 3;
    _n²    = nr ^ 2;
    _n⁴    = nr ^ 4;
    _n⁶    = nr ^ 6;
    _n²p   = _n² + 1;
    _n²p²  = _n²p ^ 2;
    _n²p³  = _n²p ^ 3;
    _n²m²  = (_n² - 1) ^ 2;
    _k     = -1 * _n²m² / 4;
    _k²    = _k ^ 2;
    _sin²α = sind(α) ^ 2;
    _b₂    = _sin²α - _n²p/2;
    _b₁    = (α==90 ? 0 : sqrt(_b₂ ^ 2 + _k));
    _b     = _b₁ - _b₂;
    _b³    = _b ^ 3;
    _npanm = 2 * _n²p * _a - _n²m²;
    _npbnm = 2 * _n²p * _b - _n²m²;

    # S polarization
    _ts  = ( _k² / (6*_b³) + _k/_b - _b/2 ) - ( _k² / (6*_a³) + _k/_a - _a/2 );

    # P polarization
    _tp₁ = -2 * _n² * (_b - _a) / _n²p²;
    _tp₂ = -2 * _n² * _n²p * log(_b / _a) / _n²m²;
    _tp₃ = _n² * (1/_b - 1/_a) / 2;
    _tp₄ = 16 * _n⁴ * (_n⁴ + 1) * log(_npbnm / _npanm) / (_n²p³ * _n²m²);
    _tp₅ = 16 * _n⁶ * (1/_npbnm - 1/_npanm) / _n²p³;
    _tp  = _tp₁ + _tp₂ + _tp₃ + _tp₄ + _tp₅;

    return  (_ts + _tp) / (2 * _sin²α)
end


export logadd, logsub, log1mexp, log1pexp, inf

function inf(::Type{T})::T where T
    one(T) / zero(T)
end

function inf(x::T)::T where T
    inf(T)
end

"log(x+y) from log(x) and log(y)"
function logadd(lx, ly)
    if lx > ly
        return logadd(ly, lx)
    elseif -Inf < lx ≤ ly < Inf
        return ly + log1p(exp(lx - ly))
    else
        return ly
    end
end

logadd(lx, ly, lz, ls...) = foldl(logadd, (logadd(logadd(lx,ly),lz), ls...))

"log(|x-y|) from log(x) and log(y)"
function logsub(lx, ly)
    if min(lx, ly) > -Inf
        return max(lx, ly) + log1mexp(abs(lx - ly))
    else
        return max(lx, ly)
    end
end

"log(1 - exp(-x)), x ≥ 0"
function log1mexp(x)
    # https://cran.r-project.org/web/packages/Rmpfr/vignettes/log1mexp-note.pdf, Eq. 7
    x ≥ 0 || throw(DomainError())
    if x ≤ log(2)
        log(-expm1(-x))
    else
        log1p(-exp(-x))
    end
end

"log(1 + exp(x))"
function log1pexp(x)
    # https://cran.r-project.org/web/packages/Rmpfr/vignettes/log1mexp-note.pdf, Eq. 10
    if x ≤ -37
        exp(x)
    elseif -37 < x ≤ 18
        log1p(exp(x))
    elseif 18 < x ≤ 33.3
        x + exp(-x)
    else
        x
    end
end


module SDM

using LinearAlgebra
using Printf

include("utility.jl")
include("graphics.jl")

export sdm

function sdm(fun, fun_derivative, starting_point, step_size_params, epsilon; display_g=true)
    type = get(step_size_params, "type", "constant")
    method = get(SDM_step_size_methods, type, constant_step_size_sdm)
    return method(fun, fun_derivative, starting_point, step_size_params, epsilon, display_g)
end

function constant_step_size_sdm(fun, fun_derivative, starting_point, step_size_params, epsilon, display_g)
    @printf("Press enter to continue with iterations\n")
    @printf("step\t\tposition\t\tgradient_norm\n")

    gradient = fun_derivative(starting_point[1], starting_point[2])
    x = [starting_point[1], starting_point[2]]

    step_counter = 1
    while(true)

        #we move in the opposite direction of the gradient
        d = -gradient
        #step size is constant
        alpha = get(step_size_params, "step_size", 0.1)
        #move to the next point
        movement = alpha*d
        if display_g
            display(add_vector(x, movement))
        end
        x = x + movement

        gradient_norm = norm(gradient, 2)
        @printf("%d\t\t[%g,%g]\t\t%g", step_counter, x[1], x[2], gradient_norm)

        #we stop whenever the norm of the gradient is lower than a certain epsilon
        if (gradient_norm < epsilon)
            break
        end

        wait_for_key("")

        step_counter += 1

        #calculate the gradient for the new iteration
        gradient = fun_derivative(x[1], x[2])
    end

    @printf("\n")
    return x
end

function quadratic_step_size_sdm(fun, fun_derivative, starting_point, step_size_params, epsilon, display_g)
    @printf("Press enter to continue with iterations\n")
    @printf("step\t\tposition\t\tgradient_norm\n")

    Q = get(step_size_params, "Q", false)
    if Q==false error("Q not set in step_size_params"); return; end;

    gradient = fun_derivative(starting_point[1], starting_point[2])
    x = [starting_point[1], starting_point[2]]

    step_counter = 1
    while(true)

        #we move in the opposite direction of the gradient
        d = -gradient
        #best step size, if the function is a quadratic function
        alpha = (norm(d,2)^2)/(d'*Q*d)
        #move to the next point
        movement = alpha*d
        if display_g
            display(add_vector(x, movement))
        end
        x = x + movement

        gradient_norm = norm(gradient, 2)
        @printf("%d\t\t[%f,%f]\t\t%g", step_counter, x[1], x[2], gradient_norm)

        #we stop whenever the norm of the gradient is lower than a certain epsilon
        if (gradient_norm < epsilon)
            break
        end

        wait_for_key("")

        step_counter += 1

        #calculate the gradient for the new iteration
        gradient = fun_derivative(x[1], x[2])
    end

    @printf("\n")
    return x
end

SDM_step_size_methods = Dict([
    ("constant", constant_step_size_sdm),
    ("quadratic", quadratic_step_size_sdm)
])

end


""" Predator prey model with logistic function from http://www.math.lsa.umich.edu/~rauch/256/F2Lab5.pdf """
function foxes_and_rabbits(oct::Bool = true)
    m = Model(SOLVER_SILENT)
    t = 20
    r = 0.2
    x1 = 0.6
    y1 = 0.5
    @variable(m, x[1:t] >= 0.01) # Note: Ipopt solution does not converge with an upper bound!!
    @variable(m, dx[1:t-1])
    @variable(m, y[1:t] >= 0.01)
    @variable(m, dy[1:t-1])
    @constraint(m, x[1] == x1)
    @constraint(m, y[1] == y1)
    @constraint(m, [i=2:t], x[i] == x[i-1] + dx[i-1])
    @constraint(m, [i=2:t], y[i] == y[i-1] + dy[i-1])

    # GlobalModel representation
    set_upper_bound.(x, 1)
    set_upper_bound.(dx, 1)
    set_lower_bound.(dx, -1)
    set_upper_bound.(y, 1)
    set_upper_bound.(dy, 1)
    set_lower_bound.(dy, -1)

    if !oct
        @NLconstraint(m, [i=1:t-1], dx[i] == x[i]*(1-x[i]) - x[i]*y[i]/(x[i]+1/5))
        @NLconstraint(m, [i=1:t-1], dy[i] == r*y[i]*(1-y[i]/x[i]))
        return m
    else
        gm = GlobalModel(model = m, name = "foxes_rabbits")
        add_nonlinear_constraint(gm, :((x, y, dx) -> dx[1] - (x[1]*(1-x[1]) -x[1]*y[1]/(x[1]+$(r)))),
                                vars = [x[1], y[1], dx[1]], equality=true)
        add_nonlinear_constraint(gm, :((x, y, dy) -> dy[1] - $(r)*y[1]*(1-y[1]/x[1])),
                                vars = [x[1], y[1], dy[1]], equality=true)
        for i = 2:t-1
            add_linked_constraint(gm, gm.bbls[1], [x[i], y[i], dx[i]])
            add_linked_constraint(gm, gm.bbls[2], [x[i], y[i], dy[i]])
        end
        return gm
    end
end


#' ## What is learning?
#'
#' Computers read data, as we saw in notebooks 1 and 2. We can then build
#' functions that model that data to make decisions, as we saw in notebooks
#' 3 and 5.
#'
#' But how do you make sure that the model actually fits the data well?
#' In the last notebook, we saw that we can fiddle with the parameters of our
#' function defining the model to reduce the loss function. However,
#' we don't want to have to pick the model parameters ourselves.
#' Choosing parameters ourselves works *well enough* when we have a
#' simple model and only a few data points, but can quickly become
#' extremely complex for more detailed models and larger data sets.
#'
#' Instead, we want our machine to *learn* the parameters that fit the model
#' to our data, without needing us to fiddle with the parameters ourselves.
#' In this notebook, we'll talk about the "learning" in machine learning.

#' ### Motivation: Fitting parameters by hand
#'
#' Let's go back to our example of fitting parameters from notebook 3.
#' Recall that we looked at whether the amount of green in the pictures
#' could distinguish between an apple and a banana, and used a sigmoid
#' function to model our choice of "apple or banana" using the amount of
#' green in an image.

using Plots
gr()
using Images, Statistics

sigma(x,w,b) = 1/(1 + exp(-w*x + b))

apple = load("data/10_100.jpg")
banana = load("data/104_100.jpg")
apple_g = mean(Float64.(green.(apple)))
banana_g = mean(Float64.(green.(banana)))

w = 10.0 # Try manipulating w between 0 and 30 to see how the plot changes
b = 15.0 # Try manipulating b bewteen 0 and 30

pl = plot(x->sigma(x,w,b), xlims=(0, 1))
scatter!([apple_g], [0.0], label="Apple")
scatter!([banana_g], [1.0], label="Banana")

#' Intuitively, how did you tweak the sliders so that way the model sends
#' apples to 0 and bananas to 1? Most likely, you did the following:

#' #### Move the sliders a bit, see whether the curve moves in the right
#' direction, and if it did, keep doing it.

#' For a machine, "learning" is that same process, translated into math!

#' ## "Learning by nudging": The process of descent
#'
#' Let's start to formalize this idea. In order to push the curve in the
#' "right direction", we need some measurement of "how right" and "how wrong"
#' the model is. When we translate the idea of a "right direction" into math,
#' we end up with a **loss function**, `L(w, b)`, as we saw in notebook 5.
#' We say that the loss function is lowest when the model `σ(x, w, b)`
#' performs the best.
#'
#' Now we want to create a loss function that is the lowest when the apple is
#' at `0` and the banana is at `1`. If the data (the amount of green) for
#' our apple is $x_1$, then our model will output $σ(x_1,w, b)$ for our apple.
#' So, we want the difference $0 - σ(x_1, w, b)$ to be small. Similarly, if our
#' data for our banana (the banana's amount of green) is $x_2$, we want the
#' difference $1 - σ(x_2, w, b)$ to be small.
#'
#' To create our loss function, let's add together the squares of the
#' difference of the model's output from the desired output for the apple
#' and the banana. We get
#'
#' $$ L(w,b) = (0 - σ(x_1, w, b))^2 + (1 - σ(x_2, w, b))^2. $$
#'
#' $L(w, b)$ is lowest when it outputs `0` for the apple and `1` for the
#' banana, and thus the cost is lowest when the model "is correct".
#'
#' We can visualize this function by plotting it in 3D with the
# `surface` function or in 2D with contour lines

# plotly() # The plotly backend is nice for 3d surface plots
# gr() # The GR backend is good for faster interactive plots

L(w,b) = (0 - sigma(apple_g,w,b))^2 + (1 - sigma(banana_g,w,b))^2

w_range = 10:0.1:13
b_range = 0:1:20
L_values = [L(w,b) for b in b_range, w in w_range]

w = 11.5 # Try manipulating w with values from w_range (between 10 and 13)
b = 7 # Try manipulating b with values from b_range (between 0 and 20)
##     p1 = surface(w_range, b_range, L_values, xlabel="w", ylabel="b", cam=(70,40), cbar=false, leg=false)
##     scatter!(p1, [w], [b], [L(w,b)+1e-2], markersize=5, color = :blue)
p1 = contour(w_range, b_range, L_values, levels = 0.05:0.1:1,
    xlabel="w", ylabel="b", cam=(70,40), cbar=false, legend=false)
scatter!(p1, [w], [b], markersize=5, color=RGB(0.05, 0.05, 0.90))

p2 = plot(x->sigma(x,w,b), 0, 1, label="Model", legend=(0.1,0.9), lw=3)
scatter!(p2, [apple_g], [0.0], label="Apple", markersize=8)
scatter!(p2, [banana_g], [1.0], label="Banana", markersize=8, xlim=(0,1),
    ylim=(-0.1,1.1))
plot(p1, p2, layout=(2,1))

#' The blue ball on the 3D plot shows the current parameter choices,
#' plotted as `(w,b)`. Shown below the 3D plot is a 2D plot of the
#' corresponding model with those parameters. Notice that as the blue ball
#' rolls down the hill, the model becomes a better fit. Our loss function
#' gives us a mathematical notion of a "hill", and the process of
#' "learning by nudging" is simply rolling the ball down that hill.

#' To do this mathematically, we need to know which direction is "downhill".
#' Recall from calculus that the derivative of `L` with respect to `b` tells
#' you how `L` changes when `b` changes. Thus to roll downhill, we should go
#' in the direction where the derivative is negative (the function goes down)
#' for each parameter. This direction is the negative of what's called the
#' **gradient**, $\nabla L$. This means that the "learn by nudging method"
#' can be rephrased in mathematical terms as:
#'
#' 1. Calculate the gradient
#' 2. Move a little bit in the direction of the negative gradient
#' 3. Repeat

#' This process of rolling the ball in the direction of the negative gradient
#' is called **gradient descent**; written mathematically, it is
#'
#' $$p_{n+1} = p_n - \eta \nabla L(p_n).$$
#'
#' Here, $p_n$ represents the vector of current parameters
#' $(w, b)$; $\nabla L(p_n)$ is the gradient of the loss function, given
#' those parameters. We start from $p_n$ and change it by $\eta \nabla L(p_n)$,
#' where $\eta$ is a small step size that determines how far we move the
#' parameters in the direction of the negative gradient; notice that if you
#' step too far, you'll overshoot the minimum!. The result is
#' $p_{n+1}$, the new vector of parameters.
#'
#' [Picture of Gradient Descent Vectors]
#'
#' If we repeat this process, then we will end up at parameters where the
#' model correctly labels apples as `0` and bananas as `1`. When this happens,
#' the model has learned from the data and can then read pictures and tell
#' you whether they are apples or bananas!

#' #### Exercise 1
#'
#' Use the following terms to fill in the sentences below.
#' Terms may be used more than once or not at all:
#' > gradient, loss function, derivative, gradient descent, learning.
#'
#' * We can think of a _(A)_ as a 1D version of a _(B)_.
#' * We can visualize a _(C)_ as a hill.
#' * In the explanation above, rolling downhill is called _(D)_ and means traveling along the _(E)_.
#' * To quantify the correctness of a model we use a _(F)_.
#' * When our program can minimize a _(G)_ on its own, we say it is _(H)_.
#'
#'

#'
#' A)

#' B)

#' C)

#' D)

#' E)

#' F)

#' G)

#' H)


# Linear ODE
linear = (u,p,t) -> (p*u)
linear_analytic = (u0,p,t) -> u0*exp(p*t)
"""
Linear ODE

```math
\\frac{du}{dt} = αu
```

with initial condition ``u0=1/2``, ``α=1.01``, and solution

```math
u(t) = u0e^{αt}
```

with Float64s. The parameter is ``α``
"""
prob_ode_linear = ODEProblem(
                  ODEFunction(linear,analytic=linear_analytic),
                  1/2,(0.0,1.0),1.01)

"""
Linear ODE

```math
\\frac{du}{dt} = αu
```

with initial condition ``u0=1/2``, ``α=1.01``, and solution

```math
u(t) = u0e^{αt}
```

with BigFloats
"""
prob_ode_bigfloatlinear = ODEProblem(
                          ODEFunction(linear,analytic=linear_analytic),
                          big(0.5),(0.0,1.0),big(1.01))

f_2dlinear = (du,u,p,t) -> (@. du = p*u)
f_2dlinear_analytic = (u0,p,t) -> @. u0*exp(p*t)
"""
4x2 version of the Linear ODE

```math
\\frac{du}{dt} = αu
```

with initial condition ``u0=1/2``, ``α=1.01``, and solution

```math
u(t) = u0e^{αt}
```

with Float64s
"""
prob_ode_2Dlinear = ODEProblem(
                    ODEFunction(f_2dlinear,analytic=f_2dlinear_analytic),
                    rand(4,2),(0.0,1.0),1.01)

"""
100x100 version of the Linear ODE

```math
\\frac{du}{dt} = αu
```

with initial condition ``u0=1/2``, ``α=1.01``, and solution

```math
u(t) = u0e^{αt}
```

with Float64s
"""
prob_ode_large2Dlinear = ODEProblem(
                    ODEFunction(f_2dlinear,analytic=f_2dlinear_analytic),
                    rand(100,100),(0.0,1.0),1.01)

"""
4x2 version of the Linear ODE

```math
\\frac{du}{dt} = αu
```

with initial condition ``u0=1/2``, ``α=1.01``, and solution

```math
u(t) = u0e^{αt}
```

with BigFloats
"""
prob_ode_bigfloat2Dlinear = ODEProblem(
                    ODEFunction(f_2dlinear,analytic=f_2dlinear_analytic),
                    BigFloat.(rand(4,2)).*ones(4,2)/2,(0.0,1.0),big(1.01))

f_2dlinear_notinplace = (u,p,t) -> p*u
"""
4x2 version of the Linear ODE

```math
\\frac{du}{dt} = αu
```

with initial condition ``u0=1/2``, ``α=1.01``, and solution

```math
u(t) = u0e^{αt}
```

on Float64. Purposefully not in-place as a test.
"""
prob_ode_2Dlinear_notinplace = ODEProblem(
                    ODEFunction(f_2dlinear_notinplace,analytic=f_2dlinear_analytic),
                    rand(4,2),(0.0,1.0),1.01)


######### Stan program example  ###########

using Mamba, Stan

ProjDir = dirname(@__FILE__)
cd(ProjDir) do

  bernoullimodel = "
  data {
    int<lower=1> N;
    int<lower=0,upper=1> y[N];
  }
  parameters {
    real<lower=0,upper=1> theta;
  }
  model {
    theta ~ beta(1,1);
    y ~ bernoulli(theta);
  }
  "

  observeddata = [
    Dict("N" => 10, "y" => [0, 1, 0, 1, 0, 0, 0, 0, 0, 1]),
    Dict("N" => 10, "y" => [0, 1, 0, 0, 0, 0, 1, 0, 0, 1]),
    Dict("N" => 10, "y" => [0, 0, 0, 0, 0, 0, 1, 0, 1, 1]),
    Dict("N" => 10, "y" => [0, 0, 0, 1, 0, 0, 0, 1, 0, 1])
  ]

  initparms = Dict{String, Any}[
    Dict("theta" => 0.1),
    Dict("theta" => 0.4),
    Dict("theta" => 0.5),
    Dict("theta" => 0.9),
  ]

  monitor = ["theta", "lp__", "accept_stat__"]

  global stanmodel, rc, sim1
  stanmodel = Stanmodel(name="bernoulli", model=bernoullimodel, num_warmup=1);
  rc, sim1 = stan(stanmodel, observeddata, init=initparms, CmdStanDir=CMDSTAN_HOME)

  rc == 0 && sim1.value[1, end, :] |> display
end # cd


# sample from a harmonic potential in 1D using GJF
using Plots
using Printf
using Random
using LinearAlgebra
using ForwardDiff
using BasicMD

include("potentials.jl")

β = 5.0;
γ = 1.5;
# M = Float64[2.1];
M = 2.1;
q₀ = [0.0];
p₀ = [0.0];
Δt = 1e-1;
n_iters = 10^4;

V = x->Harmonic(x);
cfg = ForwardDiff.GradientConfig(V, q₀);
gradV! = (gradV, x)-> ForwardDiff.gradient!(gradV, V, x, cfg);

sampler = GJF(gradV!, β, γ, M, Δt)

Random.seed!(100);
Q₀ = copy(q₀);
P₀ = copy(p₀);
sample_trajectory!([Q₀, P₀], sampler, options=MDOptions(n_iters=n_iters));
@printf("In Place (Q,P) after %d iterations: (%g,%g)\n",n_iters, Q₀[1], P₀[1]);
#
Random.seed!(100);
Xvals = sample_trajectory([q₀, p₀], sampler, options=MDOptions(n_iters=n_iters,n_save_iters=n_iters));
Q = Xvals[end][1][1]
P = Xvals[end][2][1]
@printf("(Q,P) after %d iterations: (%g,%g)\n",n_iters, Q[1], P[1]);
#
Random.seed!(100);
Xvals = sample_trajectory([q₀, p₀], sampler, options=MDOptions(n_iters=n_iters));
histogram([X[1][1] for X in Xvals],label="Samples",normalize=true)
qq=LinRange(-2,2,200)
plot!(qq, sqrt((β)/(2*π))*exp.(-0.5 * β * qq.^2),label="Density")
xlabel!("q")
ylabel!("Frequency")


# basisfunctions.jl

"""
    fourier(k::Int)

fourier implements the Fourier basis of functions ϕ_k, defined by ϕ_0≡1, and if k≥1 and k is odd, by
ϕ_k(x)=sqrt(2)sin((k+1)*π*x) and if k≥1 and even, by ϕ_k(x)=sqrt(2)cos(π*k*x).
This function is defined for k≥0.

# Examples

```julia
x = -1.0:0.01:2.0
y = fourier(3).(x)
```
"""
function fourier(k::Int)
  k ≥ 0 || throw(AssertionError("k should be non-negative"))
  if k == 0
    return x -> 1.0
  elseif isodd(k)
    kplusonefloat = float(k+1)
    return x -> 1.4142135623730951 * sinpi(kplusonefloat*x)
  else
    kfloat = float(k)
    return x -> 1.4142135623730951 * cospi(kfloat*x)
  end
end

"""
    faberschauderone(x)

Implements the first Faber-Schauder function defined by 1-2x for 0≤x≤1/2 and
-1+2x for 0.5≤x≤1, and is 1-periodically extended to all x∈R.

#See also: faberschauder

#Warning

Note the difference between faberschauderone and faberschauder. The first is a
function that takes a Float64 and returns a Float64, the second takes (j,k) and
returns an anonymous function that takes a Float64 and returns a Float64.

#Examples

```julia
x=-2.0:0.001:2.0
y=faberschauderone.(x)
```
"""
function faberschauderone(x)
    y = mod(x, 1.0)
    if 0≤y≤0.5
        return 1.0 - 2*y
    else
        return -1.0 + 2*y
    end
end

"""
    faberschauder(j::Int, k::Int)

faberschauder implements the k-th Faber-Schauder function of level j. Here, j≥0 and 1≤k≤2^j.
It is a one-periodic function and defined on [0,1] by 2^(j+1)(x-(k-1)2^(-j)) on
(k-1)2^(-j)≤x≤(k-1/2)2^(-j) and 1 - 2^(j+1)(x-(k-1/2)2^(-j)) on
[(k-1/2)2^(-j), k2^(-j)] and zero outside these intervals.

# See also: faberschauderone.

#Warning

Note the difference between faberschauderone and faberschauder. The first is a
function that takes a Float64 and returns a Float64, the second takes (j,k) and
returns an anonymous function that takes a Float64 and returns a Float64.

#Example with Plots
```julia
using Plots
J=2
x=0.0:0.001:1.0
p=plot()
for j in 0:J
  for k in 1:2^j
      y = faberschauder(j,k).(x)
      plot!(p,x,y)
  end
end
display(p)
```
"""
function faberschauder(j::Int,k::Int)
  j ≥ 0 || throw(AssertionError("j should be a nonnegative integer."))
  1≤k≤2^j || throw(AssertionError("k should be an integer between 1 and 2^j."))
  return function(x::Float64)
    y = mod(x,1.0)
    if y ≤ ldexp(float(k-1), -j) || y ≥ ldexp(float(k), -j)
      return 0.0
    elseif ldexp(float(k-1), -j) < y ≤ ldexp(k-0.5, -j)
      return ldexp(y, j+1)-2*(k-1)
    else
      return 2*k - ldexp(y, j+1)
    end
  end
end


export cellvolume

"""
    cellvolume(a, b, c, α, β, γ)

Calculates the cell volume from 6 cell parameters.
"""
cellvolume(a, b, c, α, β, γ) =
    a * b * c * sqrt(sind(α)^2 - cosd(β)^2 - cosd(γ)^2 + 2 * cosd(α) * cosd(β) * cosd(γ))
"""
    cellvolume(l::Lattice)
    cellvolume(c::Cell)

Calculates the cell volume from a `Lattice` or a `Cell`.
"""
cellvolume(lattice::AbstractLattice) = abs(det(lattice.data))
cellvolume(cell::Cell) = cellvolume(cell.lattice)
"""
    cellvolume(g::MetricTensor)

Calculates the cell volume from a `MetricTensor`.
"""
cellvolume(g::MetricTensor) = sqrt(det(g.data))  # `sqrt` is always positive!


import MathOptInterface
const MOI = MathOptInterface

"""
    strengthen(a, u0, v0, ::MOI.Nonnegatives)

Compute strenghtened coefficient for non-negative variable.

The strengthened coefficient is given by
```
α = min(-u0 * ⌊π⌋, a + v0 * ⌊π⌋)
```
where ``π = -a / (u0 + v0)``
"""
function strengthen(a::Float64, u0::Float64, v0::Float64, k::MOI.Nonnegatives)

    # Sanity checks
    MOI.dimension(k) == 1 || error("Trying to strengthen a $(typeof(k)) of dimension $(MOI.dimension(k)) (should be 1)")
    @assert u0 + v0 > 1e-4  "u0 + v0 = $(u0 + v0) (should be > 1e-4)"  # Otherwise cut is essentially a K* cut

    π = -a / (u0 + v0)
    return min(
        -u0 * floor(π),
        a + v0 * ceil(π)
    )
end

"""
    strengthen(a, u0, v0, ::MOI.Nonpositive)

Compute strenghtened coefficient for non-positive variable.

The strengthened coefficient is given by
```
α = min(-u0 * ⌊π⌋, a + v0 * ⌊π⌋)
```
where ``π = -a / (u0 + v0)``
"""
function strengthen(a::Float64, u0::Float64, v0::Float64, k::MOI.Nonpositives)

    # Sanity checks
    MOI.dimension(k) == 1 || error("Trying to strengthen a $(typeof(k)) of dimension $(MOI.dimension(k)) (should be 1)")
    @assert u0 + v0 > 1e-4  # Otherwise cut is essentially a K* cut

    π = -a / (u0 + v0)
    return max(
        -u0 * ceil(π),
        a + v0 * floor(π)
    )
end

# TODO: strengthening for conic components


function gausschebyshev( n::Int32, kind::Int32=1 )
# GAUSS-CHEBYSHEV NODES AND WEIGHTS.

x = (Array(Float64,n), Array(Float64,n))
# Use known explicit formulas. Complexity O(n).
if kind == 1
    # Gauss-ChebyshevT quadrature, i.e., w(x) = 1/sqrt(1-x^2)
    x = (cos((2*[n:-1:1]-1)*pi/2n), pi./n*ones(n))
elseif kind == 2
    # Gauss-ChebyshevU quadrature, i.e., w(x) = sqrt(1-x^2)
    x = (cos([n:-1:1]*pi./(n+1)), pi/(n+1)*sin([n:-1:1]./(n+1)*pi).^2 )
elseif kind == 3
    # Gauss-ChebyshevV quadrature, i.e., w(x) = sqrt((1+x)/(1-x))
    x = (cos(([n:-1:1]-.5)*pi/(n+.5)), 2*pi/(n+.5)*cos(([n:-1:1]-.5)*pi/(2(n+.5))).^2)
elseif kind == 4
    # Gauss-ChebyshevW quadrature, i.e., w(x) = sqrt((1-x)/(1+x))
    x = (cos([n:-1:1]*pi/(n+.5)), 2*pi/(n+.5)*sin([n:-1:1]*pi/(2(n+.5))).^2)
else
   throw(ArgumentError("Chebyshev kind should be 1, 2, 3, or 4"))
end
return x
end


using DimensionalData

@inline gaussian(t, t₀, w) = exp(-((t - t₀) / w)^2)

const KDM = 4.148808e3 # MHz^2 pc^-1 cm^3 s

"""
Δt(DM, ν₁, ν₂)
Calculates the time delay corresponding to a dispersed pulse at dispersion measure `DM` in pc/cc between
frequencies `ν₁` and `ν₂`
"""
Δt(DM, ν₁, ν₂) = KDM * DM * (ν₁^-2 - ν₂^-2)

"""
fake_pulse(DM,f_low,f_high)

Generates a `Filterbank` file corresponding to a fake pulse of dispersion measure `DM`
from frequencies `f_low` to `f_high` (MHz).

# Optional Arguments
- `channels`: Number of frequency channels
- `samples`: Number of time samples
- 'δt`: The time step (s) represented by one sample
- `w`: Width of the pulse in time samples
- `A`: Amplitude of pulse in SNR
- `α`: Spectral index, i.e. attenuation of pulse over frequency
- `start`: Index of start of pulse
- `noise_floor`: Value of the top of the noise floor. This is 1 for `dtype` of Floats and one tenth `typemax` for integers by default.
- `dtype`: Data type of data
"""
function fake_pulse(DM, f_low, f_high;
                    channels=1024,
                    samples=1024,
                    δt=1e-3,
                    w=8,
                    A=2,
                    α=4,
                    start=1,
                    noise_floor::Union{T,Nothing}=nothing,
                    dtype::Type{T}=Float32) where {T}
    @assert start < samples "Starting sample must be less than the number of samples"

    if isnothing(noise_floor)
        if dtype <: Integer
            # Default to one tenth the range of the integer type
            noise_floor = typemax(T) ÷ 0xA
        else
            noise_floor = one(T)
        end
    end

    freqs = range(; start=f_high, stop=f_low, length=channels)
    time = range(; start=0, step=δt, length=samples)

    t_start = start * δt

    # Caclulate shifts
    shifts = @. Δt(DM, freqs', f_high)

    # Generate the raw pulse, these are floats at this point
    raw_pulse = @. gaussian(time, t_start + shifts, w * δt) * A * noise_floor * (freqs' / f_high)^α

    # Prep the dynamic spectrum with the background noise
    dyn_spec = rand(zero(T):noise_floor, samples, channels)

    # Add in the pulse
    if dtype <: Integer
        dyn_spec += round.(dtype,raw_pulse)
    else
        dyn_spec += raw_pulse
    end

    # Build Filterbank
    return Filterbank(DimArray(dyn_spec,
                               (Ti(time),
                                Freq(freqs))),
                      Dict("tsamp" => δt,
                           "nbits" => sizeof(dtype) * 8,
                           "nsamples" => samples,
                           "nchans" => channels,
                           "fch1" => f_high,
                           "foff" => step(freqs)))
end

export fake_pulse


"""
    boys_quadgk(x, n::Int; rtol=1e-12)

Calculates Boys function with Gaussian Quadrature using QuadGK.jl.
Accuracy can be tuned with `rtol` argument.
This is accurate, but slow.
"""
function boys_quadgk(x, n::Int; rtol=1e-12)
    tmp = quadgk( 0, 1; rtol=rtol) do t
        t^(2n) * exp( -x*t^2 )
    end
    return tmp[1]
end

"""
    boys_gamma(x, n::Int)

Evaluate Boys function with gamma function.
This is accurate.
"""
function boys_gamma(x, n::Int)
    upper = gamma( n + 1//2 ) - gamma( n + 1//2, x )
    lower = 2x^( n + 1//2 )
    return upper / lower
end

"""
    boys_rec_largex(x, n::Int)

For large `x` > 30. Calculate using recursion.
Uses approximations.
This is fast, but loses accuracy on high `n` values.

Based on:
A. K. H. Weiss, C. Ochsenfeld J. Comput.
Chem. 2015, 36, 1390–1398. DOI: 10.1002/jcc.23935
"""
function boys_rec_largex(x, n::Int)
    if n == 0
        return √π / (2√x)
    else
        return ( n - 1//2 ) * boys_rec_largex(x, n-1) / x
    end
end


"""
    expand_poly(x, [degree = 5], [obsdim]) -> Matrix

Perform a simple polynomial basis expansion of the given `degree`
for the vector `x`.

The optional parameter `obsdim` can be used to specify which
dimension of the resulting `Matrix` should correspond to the
observations. By default each column will denote an observation,
which means that the resulting `Matrix` would have the size
`(degree, length(x))`.

```@jldoctest
julia> expand_poly(1:5, degree = 3)
3×5 Array{Int64,2}:
 1  2   3   4    5
 1  4   9  16   25
 1  8  27  64  125

julia> expand_poly(1f0:5f0, 3) # positional arguments are type-stable
3×5 Array{Float32,2}:
 1.0  2.0   3.0   4.0    5.0
 1.0  4.0   9.0  16.0   25.0
 1.0  8.0  27.0  64.0  125.0
```

Alternatively it is also possible to specify `obsdim = 1`, which
will cause each row to denote an observation.

```@jldoctest
julia> expand_poly(1:5, degree = 3, obsdim = 1)
5×3 Array{Int64,2}:
 1   1    1
 2   4    8
 3   9   27
 4  16   64
 5  25  125

julia> expand_poly(1:5, 3, ObsDim.First()); # same but type-stable
```
"""
function expand_poly(x::AbstractVector; degree::Integer = 5, obsdim = ObsDim.Last())
    expand_poly(x, degree, convert(LearnBase.ObsDimension, obsdim))
end

function expand_poly(x::AbstractVector, degree::Integer, ::ObsDim.Last)
    expand_poly(x, degree, ObsDim.Constant{2}())
end

function expand_poly(
    x::AbstractVector{T},
    degree::Integer,
    ::ObsDim.Constant{2} = ObsDim.Constant{2}()
) where T

    n = length(x)
    X = zeros(floattype(T), (degree, n))
    for i in 1:n
        for d in 1:degree
            @inbounds X[d, i] += float(x[i])^d
        end
    end
    return X
end

function expand_poly(
    x::AbstractVector{T},
    degree::Integer,
    ::ObsDim.Constant{1}
) where T

    n = length(x)
    X = zeros(floattype(T), (n, degree))
    for d in 1:degree
        for i in 1:n
            @inbounds X[i, d] += float(x[i])^d
        end
    end
    return X
end

floattype(::Type{T}) where T = float(T)
floattype(::Type{Missing}) = Missing
floattype(U::Union) = Union{floattype(U.a), floattype(U.b)}


using SparseArrays
using LinearAlgebra

function euler_fluxes(r, ru, rv, rE)
    gamma = 7.0/5
    u = @. ru/r
    v = @. rv/r
    p = @. (gamma - 1.0)*(rE-(ru^2/r+rv^2/r)/2.0)

    Frx = ru
    Fry = rv
    Frux= @. r*u^2+p
    Fruy= @. ru*v
    Frvx= @. ru*v
    Frvy= @. r*v^2+p
    FrEx= @. u*(rE+p)
    FrEy= @. v*(rE+p)
    return Frx, Fry, Frux, Fruy, Frvx, Frvy, FrEx, FrEy
end

function compact_div(Fx, Fy, h)
    global m
    m=ceil(Int,1/h)
    #x = h * (1:m)
    #m = 4  # testing purpose
    #println("m=",m)
    LHS = SymTridiagonal(ones(m), ones(m-1)/4);
    #println(LHS)
    RHS = (3 / 4h) * Tridiagonal(-ones(m-1), zeros(m), ones(m-1));
    #println(RHS)
    RHS = sparse(RHS)
    #println(RHS)
    dFx = LHS \ (RHS * Fx)
    dFy = LHS \ (RHS * Fy)
    divF = dFx+dFy
    return divF
end

function compact_filter(u, alpha)
    a = 5/8+3*alpha/4
    b = alpha+1/2
    c = alpha/4-1/8
    global m
    LHS = SymTridiagonal(ones(m), alpha*ones(m-1));
    RHS=sparse([1:m;2:m;1:m-1;3:m;1:m-2], [1:m;1:m-1;2:m;1:m-2;3:m],
           [a*ones(m); b/2.0*ones(2*(m-1)); c/2.0*ones(2*(m-2))])
    u_filt = LHS \ (RHS * u)
    return u_filt
end

function euler_rhs(r, ru, rv, rE, h)
    Frx, Fry, Frux, Fruy, Frvx, Frvy, FrEx, FrEy = euler_fluxes(r, ru, rv, rE)
    #println("Frx=",Frx)
    #println("Fry=",Fry)
    fr  =-compact_div(Frx,  Fry,  h)
    fru =-compact_div(Frux, Fruy, h)
    frv =-compact_div(Frvx, Frvy, h)
    frE =-compact_div(FrEx, FrEy, h)
    return fr, fru, frv, frE
end

function ff(u,h)
    #println("size of u in ff=",size(u))
    #println("u in ff=",u)
    fr, fru, frv, frE = euler_rhs(u[:,1], u[:,2], u[:,3], u[:,4], h)
    return [fr fru frv frE]
end

function euler_rk4step(r, ru, rv, rE, h, k, alpha)
    u = [r ru rv rE]
    #println("u=",u)
    #println("size of u=",size(u))
    k1 = k * ff(u, h)
    k2 = k * ff(u + 0.5k1, h)
    k3 = k * ff(u + 0.5k2, h)
    k4 = k * ff(u + k3, h)
    u = u + (k1 + 2k2 + 2k3 + k4) / 6
    #println(u)
    u = compact_filter(u,alpha)
    return u[:,1], u[:,2], u[:,3], u[:,4]
end

if true
  ru = [1;2;3;4]
  rv = [5;6;7;8]
  r  = [2;3;4;2]
  rE = [1;2;2;2]
  println(euler_rk4step(r, ru, rv, rE, 0.25, 0.2, 0.3))
end

#f(x, y) = x * sqrt(y)
function rk4(f::Function, x0::Float64, y0::Float64, x1::Float64, n)
    vx = Array{Float64}(undef, n+1)
    vy = Array{Float64}(undef, n+1)
    vx[1] = x = x0
    vy[1] = y = y0
    h = (x1 - x0) / n
    for i in 1:n
        k1 = h * f(x, y)
        k2 = h * f(x + 0.5h, y + 0.5k1)
        k3 = h * f(x + 0.5h, y + 0.5k2)
        k4 = h * f(x + h, y + k3)
        vx[i + 1] = x = x0 + i * h
        vy[i + 1] = y = y + (k1 + 2k2 + 2k3 + k4) / 6
    end
    return vx, vy
end


# Pulse Parameters
vmax = 5  # volts
fc = 5e9  # GHz
bw = 50  # percent
v0 = 0.00001  # Maximum voltage allowed at t = 0
ptype = "normal"

# Time vector
t = linspace(0, 5e-9, 2001);

# Compute sigma (control's pulse width)
sigma = 1/(2*pi*(fc*bw/100)/2.354820045030949);

# Compute mu (time offset)
# Note that this is only approximate for the Differentiated Pulse.
mu = sigma*sqrt(-2*log(v0/vmax));

if ptype == "normal"
    pulse = vmax*exp.(-(t - mu).^2/(2*sigma^2));
elseif ptype == "diff"
    pulse = -vmax/sigma*(t - mu).*exp.(0.5-(t - mu).^2/(2*sigma^2));
elseif ptype == "sinmod"
    pulse = vmax*exp.(-(t - mu).^2/(2*sigma^2)).*cos.(2*pi*fc*(t - mu));
else
    error("Unsupported pulse type")
end

using PyPlot
plot(t*1e9, pulse, linewidth=2.0)
xlabel("Time (ns)")
ylabel("Amplitude (Volts)")
grid(true)


function score(x, y)
    dist = sqrt(x^2 + y^2)
    if dist > 10
        return 0
    elseif 5 < dist <= 10
        return 1
    elseif 1 < dist <= 5
        return 5
    else
        return 10
    end
end


## anisotropic_exponential.jl : implementation of anisotropic exponential covariance function

## AnisotropicExponential ##
struct AnisotropicExponential{T,M} <: AnisotropicCovarianceStructure{T}
    A::M
    σ::T
end

"""
    AnisotropicExponential(A, σ=1)

Create an anisotropic exponential covariance structure with anisotropy matrix A and (optional) marginal standard deviation `σ`.

# Examples
```jldoctest
julia> A = [1 0.5; 0.5 1]
2×2 Array{Float64,2}:
 1.0  0.5
 1.0  0.5

julia> a1 = AnisotropicExponential(A)
anisotropic exponential (A=[1.0 0.5; 0.5 1.0], σ=1.0)

```
"""
function AnisotropicExponential(A::Matrix{T} where {T<:Real}; σ=1.0::T where {T<:Real})
	isposdef(A) || throw(ArgumentError("anisotropy matrix A must be positive definite"))
    σ > 0 || throw(ArgumentError("marginal standard deviation σ of exponential covariance cannot be negative or zero"))
	T = promote_type(eltype(A),typeof(σ))
	AnisotropicExponential{T,typeof(A)}(A,convert(T,σ))
end

# evaluate exponential covariance
function apply(a::AnisotropicExponential,x::Vector{T}) where {T<:Real}
    exp(-x'*a.A*x)
end

show(io::IO, a::AnisotropicExponential) = print(io, "anisotropic exponential (A=$(a.A), σ=$(a.σ))")


## MNIST data using Non-negative matrix Factorization
using MLDatasets
train_x, train_y = MNIST.traindata()
test_x,  test_y  = MNIST.testdata()

## Create the input for NMF
m = size(train_x,1)*size(train_x,2)
n = size(train_x,3)
X = Float64.(reshape(train_x, m, n))
A = copy(X') # input

## Show the average image
using Images
heatmap(Gray.(reshape(sum(A;dims=1)/size(A,1),28,28)'),
  framestyle=:none,yflip=true,colorbar=true)

## Show a random set of images
imgs = Gray.(reshape(A[1:25,:]',28,28,25))
p = plot([heatmap(imgs[:,:,i]') for i=1:25]...,framestyle=:none,yflip=true)
## Turn this into a function
myheatmap(X) = heatmap(Gray.(X'./maximum(X)),framestyle=:none,yflip=true)
plotimgs(X) = plot([myheatmap(X[:,:,i]) for i=1:size(X,3)]...)
plotimgs(reshape(A[1:25,:]',28,28,25))
## Use NMF to find structure
#using ImageMagick
using NMF
## Make a rank 5 factorization
k = 5
W, H = NMF.nndsvd(A, k; variant= :ar)
# alginst = NMF.MultUpdate{Float64}(obj=:mse, maxiter=1000, verbose=false)
alginst = NMF.ALSPGrad{Float64}(maxiter=100)
r = NMF.solve!(alginst, A, W, H)
## show the images
plotimgs(reshape(r.H',28,28,k))
## Turn this into a method
function nmf(A,k)
  W, H = NMF.nndsvd(A, k; variant= :ar)
  alginst = NMF.ALSPGrad{Float64}(maxiter=100)
  r = NMF.solve!(alginst, A, W, H)
  p = plotimgs(reshape(r.H',28,28,k))
  savefig("nmf-mnist-$k.pdf")
  return p, r
end
p5, r5 = nmf(A,5)
p5

## Run a set of these
p8, r8 = nmf(A,8)
p10, r10 = nmf(A,10)
p15, r15 = nmf(A,15)
p25, r25 = nmf(A,25)
p35, r35 = nmf(A,35)
p50, r50 = nmf(A,50)
p100, r100 = nmf(A,100)

## Use these to re-assemble digits
ndig = 60
Xs = reshape((r10.W*r10.H)[1:ndig,:]',28,28,ndig)
As = reshape(A[1:ndig,:]',28,28,ndig)
anim = @animate for i=1:60
  plot(myheatmap(Xs[:,:,i]),myheatmap(As[:,:,i]))
end
gif(anim, "nmf-mnist-10-approx.gif",fps=3)

##
ndig = 60
Xs = reshape((r50.W*r50.H)[1:ndig,:]',28,28,ndig)
As = reshape(A[1:ndig,:]',28,28,ndig)
anim = @animate for i=1:60
  plot(myheatmap(Xs[:,:,i]),myheatmap(As[:,:,i]))
end
gif(anim, "nmf-mnist-50-approx.gif",fps=3)


"""
$(TYPEDSIGNATURES)

Resurrect linspace despite https://github.com/JuliaLang/julia/pull/25896#issuecomment-363769368
"""
linspace(a,b,n)=collect(range(a,b,length=n))

"""
$(TYPEDSIGNATURES)

(Try to) create a subdivision of interval (a,b) stored in the
returned array X such that
  - `X[1]==a, X[end]==b`
  - `(X[2]-X[1])<=ha+tol*(b-a)`
  - `(X[end]-X[end-1])<=hb+tol*(b-a)`
  - There is a number q such that  `X[i+1]-X[i] == q*(X[i]-X[i-1])`
  - X is the array with the minimal possible number of points with the above property

Caveat: the algorithm behind this is  tested for many cases but unproven.

Returns an Array containing the points of the subdivision.
"""
function geomspace(a, b, ha, hb ; tol=1.0e-10, maxiterations=100)
    a=Float64(a)
    b=Float64(b)
    ha=Float64(ha)
    hb=Float64(hb)

    function _geomspace0(l,h0, hl, tol=1.0e-10)
        @assert (l>0.0)
        @assert (h0>0.0)
        @assert (hl>=h0)
        @assert((hl+h0)<l)

        #  We need to  adjust two things:

        # The sum of the geometric progression must
        # match the length of the interval, so lmismatch
        # should be zero:
        function lmismatch(q,k)
            return l - h0*(1-q^k)/(1-q)
        end

        # The claim from experimenral evidence (Wolfram)
        # is that, if written as a polynomial,
        # it has two real zeros: one and the value searched for which
        # is slightly larger than one. All other zeros are on one circle.

        # The size of the last interval should be close to
        # to hl, so hmismatch should be close to one and not larger than one
        function  hmismatch(q,k)
            return  h0*q^(k-1)/hl
        end

        # define initial number of intervals from
        # average of minmal and maximal h
        n=Int(ceil((2.0*l/(h0+hl))))

        if n==1
            n=2
        end

        # define initial q such that hmismatch is one.
        q=(hl/h0)^(1.0/(n-1.0))

        # Iteration until both mismatches are satisfactory
        # Outer loop runs until hmismatch is less than 1
        hmiss=10.0 # some initial value >1 just to run the loop at least once
        if abs(q-1.0)<tol
            hmiss=1.0
        end

        while  hmiss>1.0
            # increase number of intervals until
            # lmismatch becomes less than zero
            while  lmismatch(q,n)>tol
                n+=1
            end

            # find initial interval for q containing
            # value with zero lmismatch
            ns=0

            while lmismatch(q,n)<tol &&  ns<maxiterations
                q*=0.99
                ns+=1
            end

            ql=q*0.9
            qr=q*1.1
            @assert ns<maxiterations "Unable to determine geomspace data after $(maxiterations) iterations"

            # bisection to define q with zero lmismatch
            ns=0
            qm=0.5*(ql+qr)
            while abs(qr-ql)>tol && ns<maxiterations
                ns+=1
                mmm=lmismatch(qm,n)
                if abs(mmm)<0
                    break
                elseif lmismatch(ql,n)*mmm<0
                    qr=qm
                else
                    ql=qm
                end
                qm=0.5*(ql+qr)
            end
            # increase q slightly to increase probability
            # for last interval to be <=hl
            q=qm*(1.0+tol)
            @assert ns<maxiterations "Unable to determine geomspace data after $(maxiterations) iterations"
            hmiss=hmismatch(q,n)
            if hmiss>1.0+tol
                n=n+1
            end
        end
        #  printf("%d %g %g %g\n",n,q,lmismatch(q,n),hmismatch(q,n))

        X = Array{Float64,1}(undef,n+1)
        X[1]=0
        h=h0
        for i=1:n
            X[i+1]=X[i]+h
            h*=q
        end
        X[n+1]=l

        # Fix last interval and remove "overshoots"
        if X[end]>l
            X[end]=l
        end

        while X[end-1]+tol>X[end]
            pop!(X)
            X[end]=l
        end

        return X
    end

    @assert (ha>0.0) "Start step size $(ha) should be positive"
    @assert (hb>0.0) "End step size $(hb) should be positive"
    @assert (a<b)    "Interval ends $(a), $(b) should be increasing"
    @assert ((ha+hb)<b-a) "Sum of step sizes $(ha)+$(hb) should not exceed interval size $(b)-$(a)"

    # Map things to [0,b-a]
    tol=tol*(b-a)
    if ha<hb-tol
        X=_geomspace0(b-a,ha,hb,tol)
        X.+=a
    elseif ha>hb+tol
        X=-reverse(_geomspace0(b-a,hb,ha,tol))
        X.+=b
    else
        n=Int(ceil((b-a)/ha))
        X=collect(range(a,b,length=n+1))
    end
#    @show X[2]-X[1],ha, X[end]-X[end-1],hb

    @assert (X[2]-X[1])<=ha+tol  "First interval turned out to be larger than $(ha)"
    @assert (X[2]-X[1])>ha/2  "First interval turned out to be less  than $(ha)/10"
    @assert (X[end]-X[end-1])<=hb+tol  "Last interval turned out to be larger than $(hb)"
    @assert (X[end]-X[end-1])>=hb/2  "Last interval turned out to be less than $(hb)/10"
    @assert abs(X[1]-a)<tol  "Range start $(X[1]) differs from $(a)"
    @assert abs(X[end]-b)<tol "Range end $(X[end])  differs from $(b)"

    return X
end

collect_or_assign(X)=X
collect_or_assign(X::AbstractRange)=collect(X)
is_monotone(X)=all(X[2:end]-X[1:end-1].>0)

"""
    c=glue(a,b)

Glue together two vectors `a` and b resulting in a vector c. They last element
of `a` shall be equal (up to tol) to the first element of b.
The result fulfills `length(c)=length(a)+length(b)-1`
"""
function glue(_a::AbstractVector, _b::AbstractVector; tol=1.0e-10)
    a=collect_or_assign(_a)
    b=collect_or_assign(_b)

    is_monotone(a) || error("non-monotonous first argument of glue")
    is_monotone(b) || error("non-monotonous second argument of glue")
    Tv=promote_type(eltype(a),eltype(b))
    na=length(a)
    nb=length(b)

    d=b[1]-a[na-1]
    @assert(d>0)
    d=b[1]-a[na]
    @assert(d>-tol)
    @assert(d<tol)

    c=Vector{Tv}(undef,na+nb-1)
    ic=0
    for ia=1:na
        ic+=1
        c[ic]=a[ia]
    end
    for ib=2:nb
        ic+=1
        c[ic]=b[ib]
    end
    return c
end


function ctr(f::Function, box=[0;1], N = 10)
    h = (box[2]-box[1])/N
    val = 0
    val += f(box[1])/2
    for j = 1:N-1
        x = box[1] + j*h
        val += f(x)
    end
    val += f(box[2])/2
    val *= h
    return val
end