## ---- Fit
"""
    fit(::Type{RationalFunction}, r::Polynomial, m, n; var=:x)

Fit a Pade approximant ([`pade_fit`](@ref)) to `r`.

Examples:

```jldoctext
julia> using Polynomials, PolynomialRatios

julia> x = variable()
Polynomial(x)

julia> ex = 1 + x + x^2/2 + x^3/6 + x^4/24 + x^5/120 # Taylor polynomial for e^x
Polynomial(1.0 + 1.0*x + 0.5*x^2 + 0.16666666666666666*x^3 + 0.041666666666666664*x^4 + 0.008333333333333333*x^5)

julia> maximum(abs, exp(x) - fit(RationalFunction, ex, 1,1)(x) for x ∈ 0:.05:0.5)
0.017945395966538547

julia> maximum(abs, exp(x) - fit(RationalFunction, ex, 1,2)(x) for x ∈ 0:.05:0.5)
0.0016624471707165078

julia> maximum(abs, exp(x) - fit(RationalFunction, ex, 2,1)(x) for x ∈ 0:.05:0.5)
0.001278729299871717

julia> maximum(abs, exp(x) - fit(RationalFunction, ex, 2,2)(x) for x ∈ 0:.05:0.5)
7.262205147950951e-5
```
"""
function Polynomials.fit(::Type{RationalFunction},r::Polynomial, m::Integer, n::Integer;var=:x)
    p,q = pade_fit(r, m,n, var=var)
    p // q
end

## https://mathworld.wolfram.com/PadeApproximant.html
"""
    pade_fit(r::Polynomial, m,n)

For a polynomial `r` of degree `d ≥ m + n`, find a rational function `p/q` with
`degree(p) ≤ m`, `degree(q) ≤ n` and `q*r - p = x^{m+n+1}*s(x)` for some polynomial `s`.

This implementation sets up a system of equations to identify `p` and `q`.
"""
function pade_fit(p::Polynomial{T}, m::Integer, n::Integer; var=:x) where {T}
    d = degree(p)
    @assert (0 <= m) && (1 <= n) && (m + n <= d)

    # could be much more perfomant
    c = convert(LaurentPolynomial, p) # for better indexing
    cs = [c[m+j-i] for j ∈ 1:n, i ∈ 0:n]

    qs′ = cs[:, 2:end] \ cs[:,1]
    qs = vcat(1, -qs′)

    cs = [c[0 + j - i] for j ∈ 0:m, i∈0:n]
    ps = cs * qs

    Polynomial(ps, var), Polynomial(qs,var)
end
export pade_fit


function sample_interpolation(x,y,xmin,xmax,step)

  nx = length(x)
  n = round(Int64,(xmax-xmin)/step+1)
  xnew = Vector{Float64}(undef,n)
  ynew = Vector{Float64}(undef,n)

  for i in 1:n
   xnew[i] = xmin + (i-1)*step
   ix = findfirst( x -> x > xnew[i], x )
   if ix == 1
     ynew[i] = y[1]
   elseif ix == nothing
     ynew[i] = y[nx]
   else
     ynew[i] = y[ix-1] + (xnew[i]-x[ix-1])*((y[ix]-y[ix-1])/(x[ix]-x[ix-1]))
   end
  end

  return xnew, ynew

end


@doc raw"""
```
Z,p,q = retrieve_surface(algorithm::DiscreteShapeBound, img::AbstractArray, iterations::Int=2000)
```

Same as [`DiscreteShape`](@ref retrieve_surface(::DiscreteShape, ::AbstractArray, ::iterations))
except it has its initial conditions bound by the image as per bellow where
``E`` is the brightness of the image:
```math
\begin{gathered}
Z_{i,j}=\begin{cases}-100E_{i,j} &\text{if } E_{i,j}>0.75\\0 &\text{otherwise}\\
\end{cases}\\\\p,q=\nabla E
\end{gathered}
```
# Example
Compute the heightmap for a synthetic image generated by `generate_surface`.
```julia
using Images, Makie, ShapeFromShading

#generate synthetic image
img = generate_surface(SynthSphere(), 1, [0.2,0,0.9], radius = 5)

#calculate the heightmap (using 500 iterations)
Z,p,q = retrieve_surface(DiscreteShapeBound(), img, 500)

#normalize to maximum of 1 (not necessary but makes displaying easier)
Z = Z./maximum(Z)

#display using Makie (Note: Makie can often take several minutes first time)
r = 0.0:0.1:2
surface(r, r, Z)
```
# Reference
1. S. Elhabian, "Hands on Shape from Shading", Computer Vision and Image Processing, 2008.
"""
function retrieve_surface(algorithm::DiscreteShapeBound, img::AbstractArray, iterations::Int=2000; smoothness::Int=1000)
    ρ,I,σ,τ = estimate_img_properties(img)
    λ = smoothness
    return retrieve_surface(DiscreteShapeBound(), img, ρ, I, iterations, smoothness=λ)
end

function retrieve_surface(algorithm::DiscreteShapeBound, img::AbstractArray, albedo::Real, illumination_direction::Vector{T} where T <: Real, iterations::Int=2000; smoothness::Int=1000)
    ρ = albedo
    I = illumination_direction
    E = Array{Float64}(img)

    #initialize variables
    p = zeros(Complex{Float64},axes(E))
    q = zeros(Complex{Float64},axes(E))
    R = zeros(Complex{Float64},axes(E))
    Z = zeros(axes(E))

    #apply bounding conditions
    for i in CartesianIndices(E)
        if E[i] > 0.75
            Z[i] = -100*E[i]
        else
            Z[i] = 0.0
        end
    end
    q, p = Array{Complex{Float64}}.(imgradients(Z, KernelFactors.sobel, "replicate"))
    return solve_EulerLagrange(ρ, I, iterations, p, q, R, smoothness, E, Z)
end


### A Pluto.jl notebook ###
# v0.12.7

using Markdown
using InteractiveUtils

# ╔═╡ 6fd616f6-1ea1-11eb-3814-8bfb4a096c49
using Pkg, DrWatson

# ╔═╡ 6ffe2628-1ea1-11eb-24ea-57f985146a72
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using Distributions, Random, StatsBase, DataFrames, Plots
	Random.seed!(1)
end;

# ╔═╡ ed174bc4-1ea0-11eb-1e2f-a32874cec549
md"## Listing 3.16"

# ╔═╡ 70179b62-1ea1-11eb-13ce-7fb4d54bad4b
function rouletteSpins(p)
    x = 0
    while true
        x += 1
        if rand() < p
            return x
        end
    end
end

# ╔═╡ beee83e0-2162-11eb-10f1-7f63ce3708d2
begin
	p, xGrid, N = 18/37, 1:7, 10^6
	mcEstimate = counts([rouletteSpins(p) for _ in 1:N],xGrid)/N

	gDist = Geometric(p)
	gPmf = [pdf(gDist,x-1) for x in xGrid]
end;

# ╔═╡ beeeb36a-2162-11eb-0698-5fc29070f3ee
begin
	plot(xGrid, mcEstimate, line=:stem, marker=:circle,
		c=:blue, ms=10, msw=0, lw=4, label="MC estimate")
	plot!( xGrid, gPmf, line=:stem, marker=:xcross,
		c=:red, ms=6, msw=0, lw=2, label="PMF",
		ylims=(0,0.5), xlabel="x", ylabel="Probability")
end

# ╔═╡ 70182cd8-1ea1-11eb-094d-8d0d49cf15f3
md"## End of listing 3.16"

# ╔═╡ Cell order:
# ╠═ed174bc4-1ea0-11eb-1e2f-a32874cec549
# ╠═6fd616f6-1ea1-11eb-3814-8bfb4a096c49
# ╠═6ffe2628-1ea1-11eb-24ea-57f985146a72
# ╠═70179b62-1ea1-11eb-13ce-7fb4d54bad4b
# ╠═beee83e0-2162-11eb-10f1-7f63ce3708d2
# ╠═beeeb36a-2162-11eb-0698-5fc29070f3ee
# ╟─70182cd8-1ea1-11eb-094d-8d0d49cf15f3


using DataStructures: MutableBinaryMaxHeap, top_with_handle, delete!

function proj_simplex0(v::Array{T, 1}, s::T=one(T)) where {T<:Real}
    n = length(v)
    if sum(v) == s && all(v .≥ 0)
        w = v
    else
        u = sort(v, rev=true)
        cssv = cumsum(u)
        ρ = findall(>(0), u .* Array((1:n)) .> (cssv .- s))[end]
        θ = (cssv[ρ] - s) / (ρ)
        w = max.(v .- θ, zero(T))
    end
    return w
end

""" Algorithm from Held, M., Wolfe, P., Crowder, H.: 'Validation of
subgradient optimization'. The description is from the Condat L.:
'Fast Projection onto the Simplex and the l1 Ball'. (Algorithm 1)"""

function proj_simplex1(y::Array{T, 1}, a::T=one(T)) where {T<:Real}
    n = length(y)
    if sum(y) == a && all(y .≥ 0)
        x = y
    else
        let τ
            u = sort(y, rev=true)
            cumsum_u = zero(T)
            for k in 1:n
                if cumsum_u + u[k] < k * u[k] + a
                    cumsum_u += u[k]
                    τ = (cumsum_u - a) / k
                else
                    break
                end
            end
            x = max.(y .- τ, zero(T))
        end
    end
    return x
end

function proj_simplex12(y::Array{T, 1}, a::T=one(T)) where {T<:Real}
    n = length(y)
    if sum(y) == a && all(y .≥ 0)
        return y
    else
        τ = 0.0
        u = sort(y, rev=true)
        cumsum_u = zero(T)
        for k in 1:n
            if cumsum_u + u[k] < k * u[k] + a
                cumsum_u += u[k]
                τ = (cumsum_u - a) / k
            else
                break
            end
        end
        y = max.(y .- τ, zero(T))
    end
    return y
end

""" Algorithm from van den Berg, E., Friedlander, M.P.: 'Probing the
Pareto frontier for basis pursuit solution'. The description is from
Condat L: 'Fast Projection onto the Simplex and the l1
Ball'. (Algorithm 2)"""

function proj_simplex2(y::Array{T, 1}, a::T=one(T)) where {T<:Real}
    N = length(y)
    if sum(y) == a && all(y .≥ 0)
        x = y
    else
        τ = zero(T)
        v = MutableBinaryMaxHeap(y)
        cumsum_u = zero(T)
        for k in 1:N
            u = first(v)
            if cumsum_u + u < k * u + a
                cumsum_u += u
                i = top_with_handle(v)[2]
                delete!(v, i)
                τ = (cumsum_u - a) / k
            else
                break
            end
        end
        x = max.(y .- τ, zero(T))
    end
    return x
end

function proj_simplex22(y::Array{Float64, 1}, a=1.0)
    N = length(y)
    if sum(y) == a && all(y .≥ 0)
        x = y
    else
        τ = 0.
        v = MutableBinaryMaxHeap(y)
        cumsum_u = 0.
        for k in 1:N
            u = first(v)
            if cumsum_u + u < k * u + a
                cumsum_u += u
                i = top_with_handle(v)[2]
                delete!(v, i)
                τ = (cumsum_u - a) / k
            else
                break
            end
        end
        x = max.(y .- τ, 0.)
    end
    return x
end

"""
From Condat “Fast projection onto the simplex and the l_1 ball”. In: Mathematical Programming 158.1
(2016), pp. 575–585
"""
function proj_simplex_condat(y::Array{Float64, 1}, a=Float64(1.0))
    N = length(y)
    v = [y[1]]
    v_tilde = Float64[]
    ρ = y[1] - a
    for n in 2:N
        yn = y[n]
        if yn > ρ
            ρ += (yn - ρ) / (length(v) + 1)
            if ρ > yn - a
                append!(v, y[n])
            else
                append!(v_tilde, v)
                v = [yn]
                ρ = yn - a
            end
        end
    end

    if !isempty(v_tilde)
        for yi in v_tilde
            if yi > ρ
                append!(v, yi)
                ρ += (yi - ρ) / length(v)
            end
        end
    end
    # if during the loop rho is increases at least once, then the flag is true. Otherwise we stop
    flag = true
    while flag
        flag = false
        for (i, yi) in enumerate(v)
            if yi ≤ ρ
                deleteat!(v, i)
                ρ += (ρ - yi) / length(v)
                flag = true
            end
        end
    end
    τ = ρ
    x = max.(y .- τ, zero(y))
    return x
end

function proj_simplex3(v, z=1.)
    n = length(v)
    U = Array((1:n))
    s = 0
    ρ = 0
    while length(U) > 0
        G = []
        L = []
        k = U[rand(1:length(U))]
        ds = v[k]
        for j in U
            if v[j] >= v[k]
                if j != k
                    ds += v[j]
                    append!(G, j)
                end
            elseif v[j] < v[k]
                append!(L, j)
            end
        end
        drho = length(G) + 1
        if s + ds - (ρ + drho) * v[k] < z
            s += ds
            ρ += drho
            U = L
        else
            U = G
        end
    end
    theta = (s - z) / ρ
    return max.(v .- theta, 0.)

end

function proj_simplex4(v, z=1., τ=1e-7, max_iter=1000)
    lower = 0
    upper = maximum(v)
    current = Inf
    w = zeros(n)
    for it in 1:max_iter
        if abs(current) / z < τ && current < 0.
            break
        end
        theta = (upper + lower) / 2.0
        w = max.(v .- theta, 0.)
        current = sum(w) - z
        if current <= 0.
            upper = theta
        else
            lower = theta
        end
    end
    return w
end

function softmax(x::Array{T, 1}) where {T<:Real}
    # compare with NNlib implementation
    res = exp.(x .- maximum(x))
    res ./= sum(res)
    return res
end


import Polynomials
using Polynomials

#
# given the coefficients of a polynomial p(x) = p(n)*x^n + + p(1)*x + p(0)
#
# determne a scaling factor for x and a linear factor t that the transformed
# polynomial q(z) = t * p(s*z) has coefficient p(n) = 1 and all |p(j)| <= 1
#
# Invariant:
# if q, shift, smax, u = polyscale(p) then q = polytransform(p, shift, smax) * u
# up to numerical precision

FloatOrComplex = Union{AbstractFloat, Complex, Complex{Float32}, Complex{BigFloat}}

function polyscale(P::Poly{T}) where {T<:FloatOrComplex}
	  n = length(P.a) -1
    shift = - P.a[n] / n
    P = polytransform(P, shift)
	  p = P.a
    p[n] = 0
    pni = 1 / p[n+1]
    smaxlog = maximum(log(abs(p[1:n] * pni)) ./ (n+1 - (1:n)))
    smax = exp(smaxlog)
    u = pni / exp(smaxlog * n)
    P = polytransform(Poly(p), 0.0, smax) * u
    P.a[n+1] = T(1)
    P, shift, smax, u
end

"""
given a monic polynomial (x^n + ...) determine a scaling factor s that the transformed
polynomial q(z) = p(s*z) is half of the time < 1, rest >= 1. Ignore zero coefficients.
"""
function polyscale2(P::Poly{T}) where T<:FloatOrComplex

  p = abs.(P.a)
  n = length(p) - 1
  pni = T(1) / p[n+1]
  scale!(p, pni)
  p = view(p,1:n)
  bigg = log(realmax(T))
  tiny = log(realmin(T))
  clip(x::T) = x == -Inf ? -x : x

  llim = exp(maximum(((log.(p) - bigg) ./ (n:-1:1))))
  ulim = exp(minimum(clip, ((log.(p) - tiny) ./ (n:-1:1))))
  med = exp(sort(log.(p) ./ (n:-1:1))[n÷2+1])

  med, llim, ulim

end

"""

Find scaling factor for a polynomial p_1 + p_2*x^1 + ... + p_(n+1)*x^n,
which has the minimal possible quotient of coefficients.
Coefficients which are zero ignored.
"""
function polyscale_mini_quotient(P::Poly{T}) where T<:FloatOrComplex
  p = abs.(P.a)
  n = length(p) - 1
  pni = log(abs(p[n+1]))

  plog = log.(p[1:n]) - pni

  f = Array{real(T)}(2n)
  copy!(f, plog)
  copy!(f, n+1, -f, 1, n)
  g = Array{real(T)}(2n)
  g =
  copy!(g, -(n:-1:1))
  copy!(g, n+1, -g, 1, n)

  xa, xb, val = minimize(f, g)
  exp(xa), exp(xb), exp(val)
end

"""

  `minimize(f::Array, g::Array of same size) -> xa, xb, vmin`

Solve the following special linear optimization problem:
Define `f(x) = max{ f1[k] + g1[k] * x / k = 1..n1}`.
Minimize `f(x)` for real x.
The function f is convex as maximum of linear (convex) functions is convex.
If f is bounded below, it has minmum value `vmin` and a solution interval `[xa, xb]`
with `f(x) = vmin for all x in [xa, xb]`.
If f is not bounded below, return `vmin = -Inf` and `xa = xb = ±Inf`.
"""
function minimize(f::AbstractVector{T}, g::AbstractVector{T}) where T<:Real

  n::Int = length(f)
  n == length(g) || error(ArgumentError("lengths of f1 and g1 are different"))

  Z = zero(T)
  inf = typemax(T)

  # remove unwanted infinities
  f, g = unzip(Iterators.filter(x -> x[1] !=  inf, zip(f, g)))::NTuple{2,AbstractArray{T}}
  n = length(f)

  if n == 0
    return -inf, inf, -inf
  elseif n == 1
    if g[1] == Z
      return -inf, inf, f[1]
    elseif g[1] < Z
      return inf, inf, -inf
    else
      return -inf, -inf, -inf
    end
  end

  eva(x::T)::Tuple{T,T,Int} = maximum(1:n) do k; ((f[k] + g[k] * x), g[k], k) end

  function nextto(x::T)::Tuple{T, T, T}
    v, gk, k = eva(x)
    cross(j::Int) = -(f[j] - f[k]) / (g[j] - gk)

    if isnan(v)
      xa = xb = x
    elseif gk > Z
      xa = xb = maximum(filter(y->y<x, cross.(filter(j->j!=k, 1:n))))
    elseif gk < Z
      xa = xb = minimum(filter(y->y>x, cross.(filter(j->j!=k, 1:n))))
    elseif gk == Z
      xa = maximum(filter(y->y<x, cross.(filter(j->j!=k, 1:n))))
      xb = minimum(filter(y->y>x, cross.(filter(j->j!=k, 1:n))))
    else
      xa, xb = inf, -inf
    end
    xa, xb, gk
  end

  # find extreme slopes
  gmax::T, fmax::T, kmax::Int = maximum(1:n) do k; (g[k], f[k], k) end
  gmin::T, fmin::T, kmin::Int = minimum(1:n) do k; (g[k], f[k], k) end

  xmax::T = maximum(filter(x->!isnan(x), -(f-f[kmax]) ./ (g - gmax)))
  xmin::T = minimum(filter(x->!isnan(x), -(f-f[kmin]) ./ (g - gmin)))
  xa, xb = xmax, xmin

  while xmin < xa || xb < xmax
    xa, xb, gab = nextto((xmin + xmax) / 2)
    gs::Int = cmp(gab, Z)
    if gs >= 0
      xmax = xb
    end
    if gs <= 0
      xmin = xa
    end
  end
  xa, xb, eva((xa + xb)/ 2)...
end

maximum1(f::Function, v0, itr) = mapreduce(f, Base.scalarmax, v0, itr)
maximum1(v0, itr) = mapreduce(identity, Base.scalarmax, v0, itr)
maximum1(f::Function, itr) = mapreduce(f, Base.scalarmax, lowerbound(f, itr), itr)
maximum1(itr) = mapreduce(identity, Base.scalarmax, lowerbound(itr), itr)
lowerbound(itr) = lowerbound(eltype(itr))
lowerbound(f::Function, itr) = lowerbound(promote_type(Base.return_types(f, (eltype(itr),))))
lowerbound(T::Type) = error("no lower bound for type $T")
lowerbound(::Type{T}) where T<:Union{AbstractFloat,Unsigned} = typemin(T)
lowerbound(S::Type{Rational{T}}) where T<:Integer = typemin(S)
lowerbound(::Type{String}) = ""

#Base.mr_empty(::typeof(identity),::typeof(Base.scalarmax),::Type{T}) where T<:Real = typemin(T)
#Base.mr_empty(::typeof(identity),::typeof(Base.scalarmin),::Type{T}) where T<:Real = typemax(T)

"""
  `unzip(itr)`

The iterable object itr must deliver N-tuples.
Convert an array of N-tuples into N-tuple of arrays.
"""
function unzip(aot)
  tuple(map(x->Array(collect(x)), zip(aot...))...)
end


# * Linear interpolation

# https://github.com/JuliaLang/julia/pull/18777
lerp(a::T,b::T,t) where T = T(fma(t, b, fma(-t, a, a)))
lerp(a::R,b::R,t::C) where {R<:Real,C<:Complex} = lerp(a,b,real(t)) + im*lerp(a,b,imag(t))
lerp(a::C,b::C,t::R) where {R<:Real,C<:Complex} = lerp(real(a),real(b),t) + im*lerp(imag(a),imag(b),(t))

# * Gaußian quadrature

"""
    change_interval!(xs, ws, x, w[, a=0, b=1, γ=1])

Transform the Gaußian quadrature roots `x` and weights `w` on the
elementary interval `[-1,1]` to the interval `[γ*a,γ*b]` and store the
result in `xs` and `ws`, respectively. `γ` is an optional root of
unity, used to complex-rotate the roots (but not the weights).
"""
function change_interval!(xs::AbstractVector{T}, ws, x, w,
                          a=zero(T), b=one(T), γ=one(T)) where T
    xs .= lerp.(γ*a, γ*b, (x .+ 1)/2)
    ws .= (b-a)*w/2
    xs,ws
end

change_interval(x::AbstractVector{T}, w::AbstractVector{T},
                a=zero(T), b=one(T), γ::U=one(T)) where {T,U} =
    change_interval!(similar(x,U), similar(w), x, w, a, b)

# * Gauß–Legendre quadrature

"""
    num_quadrature_points(k, k′)

The number of quadrature points needed to exactly compute the matrix
elements of an operator of polynomial order `k′` with respect to a
basis of order `k`.
"""
function num_quadrature_points(k, k′)
    N2 = 2*(k-1) + k′
    N2>>1 + N2&1
end

"""
    lgwt(t, N) -> (x,w)

Generate the `N` Gauß–Legendre quadrature roots `x` and associated
weights `w`, with respect to the B-spline basis generated by the knot
set `t`.

# Examples

```jldoctest
julia> CompactBases.lgwt(LinearKnotSet(2, 0, 1, 3), 2)
([0.0704416, 0.262892, 0.403775, 0.596225, 0.737108, 0.929558], [0.166667, 0.166667, 0.166667, 0.166667, 0.166667, 0.166667])

julia> CompactBases.lgwt(ExpKnotSet(2, -4, 2, 7), 2)
([2.11325e-5, 7.88675e-5, 0.000290192, 0.000809808, 0.00290192, 0.00809808, 0.0290192, 0.0809808, 0.290192, 0.809808, 2.90192, 8.09808, 29.0192, 80.9808], [5.0e-5, 5.0e-5, 0.00045, 0.00045, 0.0045, 0.0045, 0.045, 0.045, 0.45, 0.45, 4.5, 4.5, 45.0, 45.0])
```
"""
function lgwt(t::AbstractKnotSet{k,ml,mr,T}, N) where {k,ml,mr,T}
    2N-1 ≥ 2(k-1) || @warn "N = $N quadrature point$(N > 1 ? "s" : "") not enough to calculate overlaps between polynomials of order k = $k"
    x, w = gausslegendre(N)

    nei = nonempty_intervals(t)
    ni = length(nei)
    xo = zeros(T, ni*length(x))
    wo = zeros(T, ni*length(x))

    for (i,j) in enumerate(nei)
        sel = (i-1)*N+1 : i*N
        change_interval!(view(xo, sel), view(wo, sel),
                         x, w,
                         t[j], t[j+1])
    end

    xo,wo
end

# * Gauß–Lobatto

function element_grid(order, a::T, b::T, c::T=zero(T), eiϕ=one(T)) where T
    x,w = gausslobatto(order)
    xs,ws = change_interval(x, w, a-c, b-c, eiϕ)
    c .+ xs, ws
end


"""
```julia
towel(u0 = [0.085, -0.121, 0.075])
```
```math
\\begin{aligned}
x_{n+1} &= a x_n (1-x_n) -0.05 (y_n +0.35) (1-2z_n) \\\\
y_{n+1} &= 0.1 \\left( \\left( y_n +0.35 \\right)\\left( 1+2z_n\\right) -1 \\right)
\\left( 1 -1.9 x_n \\right) \\\\
z_{n+1} &= 3.78 z_n (1-z_n) + b y_n
\\end{aligned}
```
The folded-towel map is a hyperchaotic mapping due to Rössler [1]. It is famous
for being a mapping that has the smallest possible dimensions necessary for hyperchaos,
having two positive and one negative Lyapunov exponent.
The name comes from the fact that when plotted looks like a folded towel, in every
projection.

Default values are the ones used in the original paper.

[1] : O. E. Rössler, Phys. Lett. **71A**, pp 155 (1979)
"""
function towel(u0=[0.085, -0.121, 0.075])
    return DDS(eom_towel, u0, nothing, jacob_towel)
end# should result in lyapunovs: [0.432207,0.378834,-3.74638]
function eom_towel(x, p, n)
    @inbounds x1, x2, x3 = x[1], x[2], x[3]
    SVector( 3.8*x1*(1-x1) - 0.05*(x2+0.35)*(1-2*x3),
    0.1*( (x2+0.35)*(1-2*x3) - 1 )*(1 - 1.9*x1),
    3.78*x3*(1-x3)+0.2*x2 )
end
function jacob_towel(x, p, n)
    @SMatrix [3.8*(1 - 2x[1]) -0.05*(1-2x[3]) 0.1*(x[2] + 0.35);
    -0.19((x[2] + 0.35)*(1-2x[3]) - 1)  0.1*(1-2x[3])*(1-1.9x[1])  -0.2*(x[2] + 0.35)*(1-1.9x[1]);
    0.0  0.2  3.78(1-2x[3]) ]
end

function eom_towel_iip(dx, x, p, n)
    @inbounds begin
        x1, x2, x3 = x[1], x[2], x[3]
        dx[1] = 3.8*x1*(1-x1) - 0.05*(x2+0.35)*(1-2*x3)
        dx[2] = 0.1*( (x2+0.35)*(1-2*x3) - 1 )*(1 - 1.9*x1)
        dx[3] = 3.78*x3*(1-x3)+0.2*x2
    end
end
function jacob_towel_iip(J, x, p, n)
    @inbounds begin
        J[1,1] = 3.8*(1 - 2x[1])
        J[2,1] = -0.19((x[2] + 0.35)*(1-2x[3]) - 1)
        J[3,1] = 0.0
        J[1,2] = -0.05*(1-2x[3])
        J[2,2] =  0.1*(1-2x[3])*(1-1.9x[1])
        J[3,2] = 0.2
        J[1,3] = 0.1*(x[2] + 0.35)
        J[2,3] = -0.2*(x[2] + 0.35)*(1-1.9x[1])
        J[3,3] = 3.78(1-2x[3])
    end
end

"""
```julia
standardmap(u0=0.001rand(2); k = 0.971635)
```
```math
\\begin{aligned}
\\theta_{n+1} &= \\theta_n + p_{n+1} \\\\
p_{n+1} &= p_n + k\\sin(\\theta_n)
\\end{aligned}
```
The standard map (also known as Chirikov standard map) is a two dimensional,
area-preserving chaotic mapping due to Chirikov [1]. It is one of the most studied
chaotic systems and by far the most studied Hamiltonian (area-preserving) mapping.

The map corresponds to the  Poincaré's surface of section of the kicked rotor system.
Changing the non-linearity parameter `k` transitions the system from completely periodic
motion, to quasi-periodic, to local chaos (mixed phase-space) and finally to global
chaos.

The default parameter `k` is the critical parameter where the golden-ratio torus is
destroyed, as was calculated by Greene [2]. The e.o.m. considers the angle variable
`θ` to be the first, and the angular momentum `p` to be the second, while
both variables
are always taken modulo 2π (the mapping is on the [0,2π)² torus).

The parameter container has the parameters in the same order as stated in this
function's documentation string.

[1] : B. V. Chirikov, Preprint N. **267**, Institute of
Nuclear Physics, Novosibirsk (1969)

[2] : J. M. Greene, J. Math. Phys. **20**, pp 1183 (1979)
"""
function standardmap(u0=0.001rand(2); k = 0.971635)
    return DDS(standardmap_eom, u0, [k], standardmap_jacob)
end
@inbounds function standardmap_eom(x, par, n)
    theta = x[1]; p = x[2]
    p += par[1]*sin(theta)
    theta += p
    while theta >= twopi; theta -= twopi; end
    while theta < 0; theta += twopi; end
    while p >= twopi; p -= twopi; end
    while p < 0; p += twopi; end
    return SVector(theta, p)
end
@inbounds standardmap_jacob(x, p, n) =
@SMatrix [1 + p[1]*cos(x[1])    1;
          p[1]*cos(x[1])        1]

"""
```julia
coupledstandardmaps(M::Int, u0 = 0.001rand(2M); ks = ones(M), Γ = 1.0)
```
```math
\\begin{aligned}
\\theta_{i}' &= \\theta_i + p_{i}' \\\\
p_{i}' &= p_i + k_i\\sin(\\theta_i) - \\Gamma \\left[
\\sin(\\theta_{i+1} - \\theta_{i}) + \\sin(\\theta_{i-1} - \\theta_{i})
\\right]
\\end{aligned}
```
A discrete system of `M` nonlinearly coupled standard maps, first
introduced in [1] to study diffusion and chaos thresholds.
The *total* dimension of the system
is `2M`. The maps are coupled through `Γ`
and the `i`-th map has a nonlinear parameter `ks[i]`.

The first `M` entries of the state are the angles, the last `M` are the momenta.

[1] : H. Kantz & P. Grassberger, J. Phys. A **21**, pp 127–133 (1988)
"""
function coupledstandardmaps end
using SparseArrays
function coupledstandardmaps(M::Int, u0 = 0.001rand(2M);
    ks = ones(M), Γ = 1.0)

    SV = SVector{M, Int}
    idxs = SV(1:M...) # indexes of thetas
    idxsm1 = SV(circshift(idxs, +1)...)  #indexes of thetas - 1
    idxsp1 = SV(circshift(idxs, -1)...)  #indexes of thetas + 1

    csm = CoupledStandardMaps{M}(idxs, idxsm1, idxsp1)
    J = zeros(eltype(u0), 2M, 2M)
    # Set ∂/∂p entries (they are eye(M,M))
    # And they dont change they are constants
    for i in idxs
        J[i, i+M] = 1
        J[i+M, i+M] = 1
    end
    sparseJ = sparse(J)
    p = (ks, Γ)
    csm(sparseJ, u0, p, 0)
    return DDS(csm, u0, p, csm, sparseJ)
end
struct CoupledStandardMaps{N}
    idxs::SVector{N, Int}
    idxsm1::SVector{N, Int}
    idxsp1::SVector{N, Int}
end
function (f::CoupledStandardMaps{N})(xnew::AbstractVector, x, p, n) where {N}
    ks, Γ = p
    @inbounds for i in f.idxs

        xnew[i+N] = mod2pi(
            x[i+N] + ks[i]*sin(x[i]) -
            Γ*(sin(x[f.idxsp1[i]] - x[i]) + sin(x[f.idxsm1[i]] - x[i]))
        )

        xnew[i] = mod2pi(x[i] + xnew[i+N])
    end
    return nothing
end
function (f::CoupledStandardMaps{M})(
    J::AbstractMatrix, x, p, n) where {M}

    ks, Γ = p
    # x[i] ≡ θᵢ
    # x[[idxsp1[i]]] ≡ θᵢ+₁
    # x[[idxsm1[i]]] ≡ θᵢ-₁
    @inbounds for i in f.idxs
        cosθ = cos(x[i])
        cosθp= cos(x[f.idxsp1[i]] - x[i])
        cosθm= cos(x[f.idxsm1[i]] - x[i])
        J[i+M, i] = ks[i]*cosθ + Γ*(cosθp + cosθm)
        J[i+M, f.idxsm1[i]] = - Γ*cosθm
        J[i+M, f.idxsp1[i]] = - Γ*cosθp
        J[i, i] = 1 + J[i+M, i]
        J[i, f.idxsm1[i]] = J[i+M, f.idxsm1[i]]
        J[i, f.idxsp1[i]] = J[i+M, f.idxsp1[i]]
    end
    return J
end

"""
```julia
henon(u0=zeros(2); a = 1.4, b = 0.3)
```
```math
\\begin{aligned}
x_{n+1} &= 1 - ax^2_n+y_n \\\\
y_{n+1} & = bx_n
\\end{aligned}
```
The Hénon map is a two-dimensional mapping due to Hénon [1] that can display a strange
attractor (at the default parameters). In addition, it also displays many other aspects
of chaos, like period doubling or intermittency, for other parameters.

According to the author, it is a system displaying all the properties of the
Lorentz system (1963) while being as simple as possible.
Default values are the ones used in the original paper.

The parameter container has the parameters in the same order as stated in this
function's documentation string.

[1] : M. Hénon, Commun.Math. Phys. **50**, pp 69 (1976)
"""
function henon(u0=zeros(2); a = 1.4, b = 0.3)
    return DDS(hoop, u0, [a,b], hoop_jac)
end # should give lyapunov exponents [0.4189, -1.6229]
hoop(x, p, n) = SVector{2}(1.0 - p[1]*x[1]^2 + x[2], p[2]*x[1])
hoop_jac(x, p, n) = @SMatrix [-2*p[1]*x[1] 1.0; p[2] 0.0]

function henon_iip(u0=zeros(2); a = 1.4, b = 0.3)
    return DDS(hiip, u0, [a, b], hiip_jac)
end
function hiip(dx, x, p, n)
    dx[1] = 1.0 - p[1]*x[1]^2 + x[2]
    dx[2] = p[2]*x[1]
    return
end
function hiip_jac(J, x, p, n)
    J[1,1] = -2*p[1]*x[1]
    J[1,2] = 1.0
    J[2,1] = p[2]
    J[2,2] = 0.0
    return
end

"""
```julia
logistic(x0 = rand(); r = 4.0)
```
```math
x_{n+1} = rx_n(1-x_n)
```
The logistic map is an one dimensional unimodal mapping due to May [1] and is used by
many as the archetypal example of how chaos can arise from very simple equations.

Originally intentend to be a discretized model of polulation dynamics, it is now famous
for its bifurcation diagram, an immensly complex graph that that was shown
be universal by Feigenbaum [2].

The parameter container has the parameters in the same order as stated in this
function's documentation string.

[1] : R. M. May, Nature **261**, pp 459 (1976)

[2] : M. J. Feigenbaum, J. Stat. Phys. **19**, pp 25 (1978)
"""
function logistic(x0=rand(); r = 4.0)
    return DDS(logistic_eom, x0, [r], logistic_jacob)
end
logistic_eom(x, p, n) = p[1]*x*(1-x)
logistic_jacob(x, p, n) = p[1]*(1-2x)

"""
    pomaeu_manneville(u0 = 0.2; z = 2.5)
The Pomeau-Manneville map is a one dimensional discrete map which is
characteristic for displaying intermittency [1]. Specifically, for
z > 2 the average time between chaotic bursts diverges, while
for z > 2.5, the map iterates are long range correlated [2].

Notice that here we are providing the "symmetric" version:
```math
x_{n+1} = \\begin{cases}
-4x_n + 3, & \\quad x_n \\in (0.5, 1] \\\\
x_n(1 + |2x_n|^{z-1}), & \\quad |x_n| \\le 0.5 \\\\
-4x_n - 3, & \\quad x_n \\in [-1, 0.5)
\\end{cases}
```

[1] : Manneville & Pomeau, Comm. Math. Phys. **74** (1980)

[2] : Meyer et al., New. J. Phys **20** (2019)
"""
function pomeau_manneville(u0 = 0.2, z = 2.5)
    return DDS(pm_eom, u0, [z], pm_jac)
end
function pm_eom(x, p, n)
    if x < -0.5
        -4x - 3
    elseif -0.5 ≤ x ≤ 0.5
        @inbounds x*(1 + abs(2x)^(p[1]-1))
    else
        -4x + 3
    end
end
function pm_jac(x, p, n)
    if x < -0.5
        -4.0
    elseif -0.5 ≤ x ≤ 0.5
        @inbounds z = p[1]
        0.5(x^2 * 2^z * (z-1)*abs(x)^(z-3) + 2^z * abs(x)^(z-1) + 2)
    else
        -4.0
    end
end


immutable Uniform <: ContinuousUnivariateDistribution
    a::Float64
    b::Float64
    function Uniform(a::Real, b::Real)
	    if a < b
	    	new(float64(a), float64(b))
	    else
	    	error("a < b required for range [a, b]")
	    end
	end
end

Uniform() = Uniform(0.0, 1.0)

@_jl_dist_2p Uniform unif

entropy(d::Uniform) = log(d.b - d.a + 1.0)

insupport(d::Uniform, x::Number) = isreal(x) && d.a <= x <= d.b

kurtosis(d::Uniform) = -6.0 / 5.0

mean(d::Uniform) = (d.a + d.b) / 2.0

median(d::Uniform) = (d.a + d.b) / 2.0

function mgf(d::Uniform, t::Real)
	a, b = d.a, d.b
	return (exp(t * b) - exp(t * a)) / (t * (b - a))
end

function cf(d::Uniform, t::Real)
	a, b = d.a, d.b
	return (exp(im * t * b) - exp(im * t * a)) / (im * t * (b - a))
end

modes(d::Uniform) = error("The uniform distribution has no modes")

rand(d::Uniform) = d.a + (d.b - d.a) * rand()

skewness(d::Uniform) = 0.0

function var(d::Uniform)
	w = d.b - d.a
	return w * w / 12.0
end


function sieve(x::Int)
    sieve = fill(true, x)
    for i in 2:isqrt(x)
        if sieve[i]
            j = 2i
            while j <= x
                sieve[j] = false
                j += i
            end
        end
    end
    primes = findall(sieve)
    popfirst!(primes)
    return primes
end


using Plots, Statistics, JLD, LaTeXStrings, StatsBase
function movement(first_pos, step_num)
    directions=[[0,1],[0,-1],[1,0],[-1,0]]
    walking=rand(directions, step_num)
    return first_pos .+ cumsum(hcat(walking),dims=1)[step_num]
end
function Rg(data)
    COM=(mean(data, dims=1))
    return sqrt(mean((data[:,1] .- COM[1]) .^2 + (data[:,2] .- COM[2]) .^2))
end

step_num_list=[i for i in 100:100:1000]
first_pos=[0,0]
run_num=100000
AllData=[]
MeanSquaredDistance=[]
RgData=[]
for step_num in step_num_list
    final_destinations=zeros(Int, run_num, 2)
    squared_distance=[]
    for i in 1:run_num
        final_destinations[i,:]=movement(first_pos, step_num)
        push!(squared_distance, (final_destinations[i,1]^2 + final_destinations[i,2]^2))
    end
    push!(SavedData, final_destinations)
    push!(MeanSquaredDistance, mean(squared_distance))
    push!(RgData, Rg(final_destinations))
end
#Radius of Gyration plot:
scatter(log.(RgData), log.(step_num_list), legend=false, xlabel="log(Time)", ylabel="log(Rg)",
    title="Logarithm of Radious of Gyration over the Logarithm of Time", dpi=400, titlefontsize=9)
savefig("C:\\Users\\Narges\\Documents\\GitHub\\computational_physics\\chapter5\\4.5\\Figs\\Rg_Time.png")
# mean(r^2) plot:
scatter(MeanSquaredDistance, step_num_list, ylabel=L"<r^2>", xlabel="Time",
    title="Mean Squared Distance over Time", legend=false, dpi=400)
savefig("C:\\Users\\Narges\\Documents\\GitHub\\computational_physics\\chapter5\\4.5\\Figs\\MeanSquare(r)_Time.png")

save("C:\\Users\\Narges\\Documents\\GitHub\\computational_physics\\chapter5\\4.5\\2DWalker.jld",
    "data", AllData, "MeanSquaredDistanceData",MeanSquaredDistance, "RgData",RgData )


function harmonic_oscillator(; X=Universe(2))
    A = [0.0 1.0; -1.0 0.0]
    prob = @ivp(x' = Ax, x(0) ∈ Singleton([1.0, 0.0]), x ∈ X)
    tspan = (0.0, 20.0)
    return prob, tspan
end

function harmonic_oscillator_homog_RFEM()
    A = [0 1; -(4π)^2 0]
    X0 = Hyperrectangle([1.0, 0.0], [0.1, 0.1])
    prob = @ivp(x' = A*x, x(0) ∈ X0)
    tspan = (0.0, 2.0)
    return prob, tspan
end

function harmonic_oscillator_forced_RFEM()
    A = [0 1; -(4π)^2 0]
    X0 = Hyperrectangle([1.0, 0.0], [0.1, 0.1])
    B = Matrix(1.0I, 2, 2)
    X = Universe(2)
    U = Interval(0.0) × Interval(0.8, 1.2)
    prob = @ivp(x' = A*x + B*u, x(0) ∈ X0, x ∈ X, u ∈ U)
    tspan = (0.0, 2.0)
    return prob, tspan
end


@doc raw"""
    bollingerbands(ta::TimeArray, ma=20, width=2.0)

Bollinger Bands

# Formula

```math
\begin{align*}
  \text{Up}   & = \text{SMA} + \text{width} \times \sigma(P) \\
  \text{Mean} & = \text{SMA} \\
  \text{Down} & = \text{SMA} - \text{width} \times \sigma(P)
\end{align*}
```
"""
function bollingerbands(ta::TimeArray, ma::Integer = 20, width::AbstractFloat = 2.0)
  tama   = sma(ta, ma)
  upband = tama .+ moving(nanstd, ta, ma) .* width .* sqrt((ma - 1) / ma) # take out Bessel correction, per algorithm
  dnband = tama .- moving(nanstd, ta, ma) .* width .* sqrt((ma - 1) / ma)
  bands  =  merge(upband, dnband)
  merge(bands, tama, colnames = [:up, :down, :mean])
end

@doc raw"""
    donchianchannels(ohlc::TimeArray, n = 20; h = :High, l = :Low)

Donchian Channels

# Formula

```math
\begin{align*}
  \text{Up}   & = \max(P_1^\text{High}, \dots, P_t^\text{High}) \\
  \text{Mid}  & = \frac{\text{Up} + \text{Down}}{2} \\
  \text{Down} & = \min(P_1^\text{Low}, \dots, P_t^\text{Low})
\end{align*}
```

# References

- [TradingView Wiki]
  (https://www.tradingview.com/wiki/Donchian_Channels_(DC))
"""
function donchianchannels(ohlc::TimeArray, n::Integer = 20; h = :High, l = :Low)
  up   = rename(moving(nanmax, ohlc[h], n), :up)
  down = rename(moving(nanmin, ohlc[l], n), :down)
  mid  = rename((up .+ down) ./ 2, :mid)
  merge(up, merge(mid, down))
end

@doc raw"""
    truerange(ohlc::TimeArray; h = :High, l = :Low, c = :Close)

True Range

# Formula

```math
\text{TR} = \max (P_t^\text{High}, P_{t-1}^\text{Close}) -
            \min (P_t^\text{Low},  P_{t-1}^\text{Close})
```
"""
function truerange(ohlc::TimeArray; h = :High, l = :Low, c = :Close)
  highs    = merge(ohlc[h], lag(ohlc[c]))
  lows     = merge(ohlc[l], lag(ohlc[c]))
  truehigh = TimeArray(timestamp(highs), maximum(values(highs), dims=2), [:hi], meta(highs))
  truelow  = TimeArray(timestamp(lows),  minimum(values(lows), dims=2),  [:lo], meta(lows))
  rename(truehigh .- truelow, :tr)
end

@doc raw"""
    atr(ohlc::TimeArray, n = 14; h= :High, l= :Low, c = :Close)

Average True Range

It's the exponential moving average of [`truerange`](@ref)

# Formula

```math
\text{ATR} = \text{EMA}(\text{TR}, n)
```
"""
function atr(ohlc::TimeArray, n::Integer = 14; h = :High, l = :Low, c = :Close)
  # atr was invented by Wilder, so only his ema is currently supported
  res = ema(truerange(ohlc; h = h, l = l, c = c), n, wilder = true)
  TimeArray(timestamp(res), values(res), [:atr], meta(ohlc))
end

@doc raw"""
    keltnerbands(ohlc, n = 20, w = 2; h = :High, l = :Low, c = :Close)

Keltner Channels

Linda Bradford Raschke introduced the newer version of Keltner Channels
in the 1980s. We implement the newer version.

# Formula

```math
\begin{align*}
  \text{Up}   & = \text{Mid} + w \times \text{ATR}(n) \\
  \text{Mid}  & = \text{EMA}(P^{typical}, n) \\
  \text{Down} & = \text{Mid} - w \times \text{ATR}(n)
\end{align*}
```

# References

- [StockCharts]
  (http://stockcharts.com/school/doku.php?id=chart_school:technical_indicators:keltner_channels)

- [Wikipedia]
  (https://en.wikipedia.org/wiki/Keltner_channel)
"""
function keltnerbands(ohlc::TimeArray, n::Integer = 20, w::AbstractFloat = 2.0;
                      h = :High, l = :Low, c = :Close)
  kma = rename(ema(typical(ohlc, h=h, l=l, c=c), n), :kma)
  rng = atr(ohlc, n, h=h, l=l, c=c)

  kup = rename(kma .+ (2 .* rng), :kup)
  kdn = rename(kma .- (2 .* rng), :kdn)

  merge(kup, merge(kma, kdn))
end

@doc raw"""
    chaikinvolatility(ta, n = 10, p = 10; h = High, l = :Low)

Chaikin Volatility

# Arguments

- `n` is the smooth period
- `p` is the previous period

# Formula

```math
\text{Chaikin Vola} =
  \frac{\text{EMA}(P^\text{High}_t - P^\text{Low}_t, n) - \text{EMA}(P^\text{High}_{t-p} - P^\text{Low}_{t-p}, n)}
       {\text{EMA}(P^\text{High}_{t-p} - P^\text{Low}_{t-p}, n)}
  \times 100
```

# References

- [IncredibleCharts]
  (https://www.incrediblecharts.com/indicators/chaikin_volatility.php)
"""
function chaikinvolatility(ta::TimeArray, n::Integer = 10, p::Integer = 10;
                           h = :High, l = :Low)
  rng = ema(ta[h] .- ta[l], n)
  prev = lag(rng, p)
  rename(@.((rng - prev) / prev * 100), :chaikinvolatility)
end


export LogZero
export LogSum2Exp
export LogSum3Exp
export LogSumExp

LogZero(T::DataType) = - floatmax(T)

"""
    LogSum2Exp(a::Real, b::Real) -> max(a,b) + log(1.0 + exp(-abs(a-b)))
```julia
julia> LogSum2Exp(Float32(1.2),Float64(3.3))
3.4155195283818967
```
"""
function LogSum2Exp(a::Real, b::Real)
    isinf(a) && return b
    isinf(b) && return a
    if a < b
        a, b = b, a
    end
    return (a + log(1.0 + exp(b-a)))
end

function LogSum3Exp(a::Real, b::Real, c::Real)
    return LogSum2Exp(LogSum2Exp(a,b),c)
end

function LogSumExp(a)
    tmp = LogZero(eltype(a))
    for i = 1:length(a)
        tmp = LogSum2Exp(tmp, a[i])
    end
    return tmp
end


print("Enter latitude       => ")
lat = parse(Float64, readline(STDIN))
print("Enter longitude      => ")
lng = parse(Float64, readline(STDIN))
print("Enter legal meridian => ")
ref = parse(Float64, readline(STDIN))
println()

slat = sin(deg2rad(lat))
@printf "    sine of latitude:   %.3f\n" slat
@printf "    diff longitude:     %.3f\n" (lng - ref)

println("\nHour, sun hour angle, dial hour line angle from 6am to 6pm\n")

for h in -6:6
  hra = 15 * h
  hra -= lng - ref
  hla = rad2deg(atan(slat * tan(deg2rad(hra))))
  @printf "HR = %3d; HRA = %7.3f; HLA = %7.3f\n" h hra hla
end


### A Pluto.jl notebook ###
# v0.12.7

using Markdown
using InteractiveUtils

# ╔═╡ 6fd616f6-1ea1-11eb-3814-8bfb4a096c49
using Pkg, DrWatson

# ╔═╡ 6ffe2628-1ea1-11eb-24ea-57f985146a72
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using Distributions, Random, StatsBase, DataFrames, Plots
	Random.seed!(1)
end;

# ╔═╡ ed174bc4-1ea0-11eb-1e2f-a32874cec549
md"## Listing 3.29"

# ╔═╡ 70179b62-1ea1-11eb-13ce-7fb4d54bad4b
Z() = sqrt(-2*log(rand()))*cos(2*pi*rand())

# ╔═╡ 0000b90c-21c2-11eb-2426-fdbaf4727f0c
begin
	xGrid = -4:0.01:4

	histogram([Z() for _ in 1:10^6], bins=50,
			normed=true, label="MC estimate")
	plot!(xGrid, pdf.(Normal(),xGrid),
		 c=:red, lw=4, label="PDF",
		 xlims=(-4,4), ylims=(0,0.5), xlabel="x", ylabel="f(x)")
end

# ╔═╡ 70182cd8-1ea1-11eb-094d-8d0d49cf15f3
md"## End of listing 3.29"

# ╔═╡ Cell order:
# ╟─ed174bc4-1ea0-11eb-1e2f-a32874cec549
# ╠═6fd616f6-1ea1-11eb-3814-8bfb4a096c49
# ╠═6ffe2628-1ea1-11eb-24ea-57f985146a72
# ╠═70179b62-1ea1-11eb-13ce-7fb4d54bad4b
# ╠═0000b90c-21c2-11eb-2426-fdbaf4727f0c
# ╟─70182cd8-1ea1-11eb-094d-8d0d49cf15f3


### A Pluto.jl notebook ###
# v0.12.4

using Markdown
using InteractiveUtils

# ╔═╡ 81915028-18e7-11eb-0aed-57a75839d4a7
begin
	using CSV
	using DataFrames
	using Plots
	using StatsPlots
	using Pipe: @pipe
	using GLM
	using StatsBase
	using Statistics
	using Dates
end

# ╔═╡ 02eb9d36-18e8-11eb-24a9-db6f58080189
md"""
# Seoul Bike Sharing Demand

This notebook demonstrates how to use GLM to analyze the Seoul Bike Sharing Demand
data set.

## Data definitions

The following definitions are captured from the [UCI Datasets Archive page](https://archive.ics.uci.edu/ml/datasets/Seoul+Bike+Sharing+Demand).

- Date: year-month-day
- Rented Bike count: Count of bikes rented at each hour
- Hour: Hour of he day
- Temperature: Temperature in Celsius
- Humidity: %
- Windspeed: m/s
- Visibility: 10m
- Dew point temperature: Celsius
- Solar radiation: MJ/m2
- Rainfall: mm
- Snowfall: cm
- Seasons: Winter, Spring, Summer, Autumn
- Holiday: Holiday/No holiday
- Functional Day: NoFunc(Non Functional Hours), Fun(Functional hours)

As the dataset is targeted for bike sharing demand, the key metrics should be the Rented Bike Count field. The rest of the columns are variables that may be correlated to bike sharing demand.

Taking a quick peek at the variables, it makes sense that they may directly or indirectly affect people's decision of renting a bike. For example, when it's very hot or very cold, then I may take on some other transportation. Likewise, humidity, solar radiation, visibility, rainfall, and snowfall probably play a part as well.

Most variables are continuous variables. The last three columns seasons, holiday, and functional day are discrete variables. As for the hour of day variable, we may want to consider that as discrete because increasing the value (later in the day) does not always affect the response variable (rented bike count) in the same direction.
"""

# ╔═╡ eedb4bf8-19a9-11eb-106b-69b047d4118e
md"""
## Reading data
"""

# ╔═╡ dfa5fbf0-18e7-11eb-391b-1d878b9c96f0
file = "https://archive.ics.uci.edu/ml/machine-learning-databases/00560/SeoulBikeData.csv"

# ╔═╡ aed964bc-18e7-11eb-11f1-a3d9957a45f5
df = DataFrame(CSV.File(download(file); dateformats = Dict(:Date => "dd/mm/yyyy")));

# ╔═╡ 4d4a9158-18ec-11eb-2fb0-c178637f903a
md"""
## Data wrangling
"""

# ╔═╡ 09577b1e-19aa-11eb-3e6c-175705b4b951
md"""
Let's quickly examine the data frame and its columns.
"""

# ╔═╡ c5b4e26a-18e7-11eb-2bef-2d937faa252d
describe(df, :eltype, :min, :mean, :max, :nmissing)

# ╔═╡ 6ddeb944-18ec-11eb-05c7-b51e7b01578e
md"It looks pretty clean in general but let's rename some columns so they can be referenced more easily in the code below."

# ╔═╡ ebc9c09a-18e9-11eb-1f56-ef92258af11f
rename!(df,
	2 => :RentedBikeCount,
	4 => :Temperature, 5 => :Humidity, 6 => :WindSpeed, 7 => :Visibility,
	8 => :DewPointTemperature, 9 => :SolarRadiation, 10 => :Rainfall, 11 => :Snowfall,
	14 => :FunctioningDay
);

# ╔═╡ 16c31048-19b6-11eb-2ad8-1f984acfb06f
md"""
## Correlation analysis
"""

# ╔═╡ 9564ab04-19b4-11eb-3c4c-815fe8a6031a
names(df)

# ╔═╡ 3084d7a6-19b5-11eb-1fa5-9b9f7cd482f4
let
	cols = 3:11
	colnames = names(df)[cols]
	Σ = cor(Matrix(df[:, cols]))
	heatmap(colnames, colnames, Σ; xrotation = 45, seriescolor = :PiYG_4)
end

# ╔═╡ 2ef9d8a6-19b6-11eb-0cc5-e1831ce6c40a
md"""
It appears that temperature and dew point temperature columns are highly correlated. For that reason, we may want to exclude dew point temperature in our model below.
"""

# ╔═╡ 7deee71e-18ec-11eb-113d-fb0cdc6e07ef
md"""
## Seasonal analysis
Since this is a time series, let's check if there's a trend.
"""

# ╔═╡ 32cd8006-18ec-11eb-221a-d38281aadffb
plot(df.Date, df.RentedBikeCount;
	legend = :none,
	ylabel = "Rented Bike Count")

# ╔═╡ a9f3aeda-18ec-11eb-2cb9-9f12fd001d88
md"""
As we only have 1 year of data, we cannot really see any seasonal trends. However, my intuition is that people don't tend to ride bikes during winter due to the low temperature. When spring comes around March/April, the demand picked up.

Now, December 2018 still look like a higher demand when compared with December 2017. A possibility is that biking may have gotten popular? Not sure...
"""

# ╔═╡ 69199458-18ef-11eb-1a46-7d71379a4c04
md"""
## Time of day

My initial guess is that people probably do not rent bikes early morning or late night. Let's verify that.
"""

# ╔═╡ 5de23772-18ed-11eb-1105-85661397bcd3
let
	data = @pipe df |>
		select(_, :Date, :Hour, :RentedBikeCount) |>
		unstack(_, :Hour, :RentedBikeCount) |>
		sort(_, :Date)

	matrix = Matrix(data[:, 2:end])'

	heatmap(data[:, 1], 0:23, matrix;
		title = "Rented Bike Count by Date / Hour",
		ylabel = "Hour",
		xrotation = 45.0,
		xticks = Date(2017,12,1):Month(1):Date(2019,1,1),
		yticks = 0:4:24)
end

# ╔═╡ 8789d48e-18ef-11eb-2b87-2337472759c7
md"""
**The two horizontal stripes show that 8 AM and 6 PM are popular hours for bike rentals.** That makes sense because people might want to do some exercise right before work or after work. Or, they may use it for transportation.

**Korean seems to stay up late.** Bike rental demand continues after 6 PM until almost midnight. By 3 AM, it's total silence.

Now, let's take a look at the same data from a different angle using a boxplot.
"""

# ╔═╡ 5ef37b5a-18eb-11eb-3f75-cfb5947950fb
begin
	hline([1800], color = :green)
	@df df StatsPlots.boxplot!(:Hour, :RentedBikeCount;
		legend = :none,
		color = palette(:default)[1],  # reset to first color due to `hline` above
		xlabel = "Hour",
		ylabel = "Rented Bike Count")
end

# ╔═╡ 3dc1127e-19a4-11eb-2c09-4f2624f479ac
md"""
The boxplot shows the quartiles for each hour across all dates in the year. It confirms our findings that 8 AM and 6 AM are popular times.
"""

# ╔═╡ 716cb212-19b1-11eb-190a-7702595f6bab
md"""
## Categorical Variables

Some variables should be converted to categorical such that GLM.jl can encode dummy variables as such. Instead of mutating the existing `Hour` column, I will create a new `Hour2` column and make it categorical.
"""

# ╔═╡ 79056460-19b6-11eb-1c5c-bf9224a0b18e
df.Hour2 = df.Hour

# ╔═╡ 9506548a-19b6-11eb-02dc-ad2832999a5b
categorical!(df, :Hour2);

# ╔═╡ 7553988e-19a5-11eb-3549-bfd4ea98b2fd
md"""
## Linear Regression

Using linear regression, we can fit the data to a linear equation. Here, the response variable is `RentedBikeCount`. We can choose any of the other fields as explanatory variables. Let's start with something simple - an ordinary linear model.
"""

# ╔═╡ c18fa04e-19a5-11eb-2f5c-0bdba322a26f
hour_model = @formula(RentedBikeCount ~ 1 + Hour2)

# ╔═╡ 558893b4-19a6-11eb-3e84-ef27b2749b80
md"""
To fit the model, we can use the `lm` function.
"""

# ╔═╡ 66b6f680-19a6-11eb-2d2f-afe8d206ee5f
ols = lm(hour_model, df)

# ╔═╡ a46a544a-19a6-11eb-3de3-87dcf7bf431f
md"""
We can now use the fitted model to predict bike demand.
"""

# ╔═╡ bfe39e20-19a6-11eb-1b9c-b95be62dadb5
let
	y = df.RentedBikeCount
	yhat = round.(Int, predict(ols))
	(r_squared = r2(ols), rmsd = rmsd(y, yhat))
end

# ╔═╡ 079fe824-19a7-11eb-3f1d-b3aa9c9b0f33
md"""
The $R^2$ is very low, meaning that the current model cannot make very accurate prediction. That's understandable because we have only used a single explanatory variable.

The root mean squared deviation (`rmsd`) value shows how much the predicted values deviates from the actual values.

Now, let's design a more complex model but we will continue to use an ordinary linear model.
"""

# ╔═╡ 41fb7ff4-19a7-11eb-0945-017db54aa411
let
	model = @formula(RentedBikeCount ~
						1 + Hour2 + Temperature + Humidity + WindSpeed +
						Visibility + SolarRadiation + Rainfall + Snowfall +
						Seasons + Holiday + FunctioningDay)
	fitted = lm(model, df)
	y = df.RentedBikeCount
	yhat = round.(Int, predict(fitted))
	(r_squared = r2(fitted), rmsd = rmsd(y, yhat))
end

# ╔═╡ e3615da0-19a7-11eb-3dc7-d734a6da4420
md"""
That's great result. The $R^2$ has jumped to 0.66 now! The RMSD is also reduced quite significantly from 543 to 375.
"""

# ╔═╡ ab20e2f6-19a9-11eb-09cd-eff9e75608ee
md"""
## Generalized Linear Model (GLM)

I wonder if we can do better. The GLM.jl package supports Generalized Linear Model (GLM) which is more flexible that linear regression. We will demonstrate how it works below.

Given that the response variable is a count, the general wisdom (not mine) is to design the model with Poisson distribution.
"""

# ╔═╡ 1accea7a-19a8-11eb-2e37-b943977ad801
let
	model = @formula(RentedBikeCount ~
						1 + Hour2 + Temperature + Humidity + WindSpeed +
						Visibility + SolarRadiation + Rainfall + Snowfall +
						Seasons + Holiday + FunctioningDay)
	fitted = glm(model, df, Poisson(), LogLink())
	y = df.RentedBikeCount
	yhat = round.(Int, predict(fitted))
	(rmsd = rmsd(y, yhat), )
end

# ╔═╡ 61eb2598-19a8-11eb-055d-b166da58fa50
md"""
There is no `r2` function defined for GLM models, so we just show RMSD here. As you can see, RMSD is further reduced using this model. That's an improvement.

What if we build an even more complex model? So far, all variables are independent. If we introduce interaction terms (multiple variables interacting with each other) then we may create a more powerful predictor.

The question is how to choose the right variables for the interaction terms. My gut feeling is that it would be appropriate to choose variables that are "orthogonal" to each other. For example, humidity and rainfall should be highly correlated and the interaction between them would be somewhat uninteresting. Hence, I have chosen to mix hour of day, temperature, and humidity in the following experiment.
"""

# ╔═╡ 66967c54-19b3-11eb-1b13-856ef40132fa
let
	model = @formula(RentedBikeCount ~
						1 + Hour2 + Temperature + Humidity + WindSpeed +
						Visibility + SolarRadiation + Rainfall + Snowfall +
						Seasons + Holiday + FunctioningDay +
						Hour2 * Temperature * Humidity
	)
	fitted = glm(model, df, Poisson(), LogLink())
	y = df.RentedBikeCount
	yhat = round.(Int, predict(fitted))
	(rmsd = rmsd(y, yhat), )
end

# ╔═╡ 1de7c9b2-19b4-11eb-2085-7b7383899819
md"""
That's great! The RMSD is now further reduced although not by a whole lot.
"""

# ╔═╡ 34be34aa-19b4-11eb-1a17-e1253c48b050
md"""
## Todo's

So far, I have been fitting the model with the complete data set. In order to test the predictive power of the model, I should test it against unseen data. Of course, I don't have any more data than what I have downloaded. What I should do is to split the data set and do cross validation.
"""

# ╔═╡ 544e7ac0-19b2-11eb-3931-8f7f10c23212
md"""
## Resources

I don't know much about GLM before working on this. I found the following resources useful as I learn about the subject:

- Foundations of Linear and Generalized Linear Models by Alan Agresti (Wiley 2015)
- [Introduction to Generalized Linear Models](https://online.stat.psu.edu/stat504/node/216/)
- [MIT 18.650 Statistics for Applications, Phillipe Rigolle, Lecture 21-22](https://www.youtube.com/watch?v=X-ix97pw0xY)
"""

# ╔═╡ 7c1251d4-19a9-11eb-0450-8dcae5fdc135
md"""
Thanks you for reading. I hope you enjoy this notebook.

_Tom Kwong_,
_October 2020_
"""

# ╔═╡ Cell order:
# ╟─02eb9d36-18e8-11eb-24a9-db6f58080189
# ╟─eedb4bf8-19a9-11eb-106b-69b047d4118e
# ╠═81915028-18e7-11eb-0aed-57a75839d4a7
# ╠═dfa5fbf0-18e7-11eb-391b-1d878b9c96f0
# ╠═aed964bc-18e7-11eb-11f1-a3d9957a45f5
# ╟─4d4a9158-18ec-11eb-2fb0-c178637f903a
# ╟─09577b1e-19aa-11eb-3e6c-175705b4b951
# ╠═c5b4e26a-18e7-11eb-2bef-2d937faa252d
# ╟─6ddeb944-18ec-11eb-05c7-b51e7b01578e
# ╠═ebc9c09a-18e9-11eb-1f56-ef92258af11f
# ╟─16c31048-19b6-11eb-2ad8-1f984acfb06f
# ╠═9564ab04-19b4-11eb-3c4c-815fe8a6031a
# ╠═3084d7a6-19b5-11eb-1fa5-9b9f7cd482f4
# ╟─2ef9d8a6-19b6-11eb-0cc5-e1831ce6c40a
# ╟─7deee71e-18ec-11eb-113d-fb0cdc6e07ef
# ╠═32cd8006-18ec-11eb-221a-d38281aadffb
# ╟─a9f3aeda-18ec-11eb-2cb9-9f12fd001d88
# ╟─69199458-18ef-11eb-1a46-7d71379a4c04
# ╠═5de23772-18ed-11eb-1105-85661397bcd3
# ╟─8789d48e-18ef-11eb-2b87-2337472759c7
# ╠═5ef37b5a-18eb-11eb-3f75-cfb5947950fb
# ╟─3dc1127e-19a4-11eb-2c09-4f2624f479ac
# ╟─716cb212-19b1-11eb-190a-7702595f6bab
# ╠═79056460-19b6-11eb-1c5c-bf9224a0b18e
# ╠═9506548a-19b6-11eb-02dc-ad2832999a5b
# ╟─7553988e-19a5-11eb-3549-bfd4ea98b2fd
# ╠═c18fa04e-19a5-11eb-2f5c-0bdba322a26f
# ╟─558893b4-19a6-11eb-3e84-ef27b2749b80
# ╠═66b6f680-19a6-11eb-2d2f-afe8d206ee5f
# ╟─a46a544a-19a6-11eb-3de3-87dcf7bf431f
# ╠═bfe39e20-19a6-11eb-1b9c-b95be62dadb5
# ╟─079fe824-19a7-11eb-3f1d-b3aa9c9b0f33
# ╠═41fb7ff4-19a7-11eb-0945-017db54aa411
# ╟─e3615da0-19a7-11eb-3dc7-d734a6da4420
# ╟─ab20e2f6-19a9-11eb-09cd-eff9e75608ee
# ╠═1accea7a-19a8-11eb-2e37-b943977ad801
# ╟─61eb2598-19a8-11eb-055d-b166da58fa50
# ╠═66967c54-19b3-11eb-1b13-856ef40132fa
# ╟─1de7c9b2-19b4-11eb-2085-7b7383899819
# ╟─34be34aa-19b4-11eb-1a17-e1253c48b050
# ╟─544e7ac0-19b2-11eb-3931-8f7f10c23212
# ╟─7c1251d4-19a9-11eb-0450-8dcae5fdc135


# fbp2_example.jl

#using Revise
using Sinograms
using MIRT: image_geom, sino_geom
# todo: ImagePhantoms
using MIRT: ellipse_im_params, ellipse_im, ellipse_sino #, :shepplogan
using MIRTjim: jim
using Plots

# Geometry
down = 2 # down-sampling factor for quicker tests
ig = image_geom(nx=512, ny=504, fov=500)
ig = ig.down(down)

sg = sino_geom(:par, nb=1088, na=984, orbit=360, orbit_start=0, d=541/949, offset=0.25,)
sg = sg.down(down)

# Phantom object
ell = ellipse_im_params(ig, :shepplogan)
xtrue = ellipse_im(ig, ell, oversample = 4, hu_scale=1000)
sino = ellipse_sino(sg, ell, oversample = 4) * 1000 # hu_scale

clim = (1 .+ (-1, 1) .* 0.05) .* 1000
p1 = jim(ig.x, ig.y, xtrue, "Phantom"; clim)
p2 = jim(sg.r, sg.ad, sino, "Sinogram", aspect_ratio=:none)

plan = fbp2(sg, ig)
@info "plan complete"
result,sino_filtered = fbp2(plan, sino)
@info "fbp complete"

p3 = jim(ig.x, ig.y, result, "FBP"; clim = clim)
#p3 = jim(ig.x, ig.y, result)
#p4 = jim(sr.r, sg.ad, sino_filtered, "Filtered Sino")
#p4 = plot(sino_filtered, label="")
p4 = jim(ig.x, ig.y, result-xtrue, "error"; clim=(-1,1).*100)

jim(p1,p2,p3,p4)

#=
#original example:

down = 2
ig = image_geom(nx=512, ny=504, fov=500)
ig = ig.down(down)
sg = sino_geom(:par, nb=888, na=984, orbit=180 , d=541/949, offset=0.25)
sg=sg.down(down)

ell = ellipse_im_params(ig, :shepplogan)
xtrue = ellipse_im(ig, ell, oversample = 4, hu_scale=1000)
sino = ellipse_sino(sg, ell, oversample = 4)

clim = (1 .+ (-1, 1) .* 0.05) .* 1000
p1 = jim(ig.x, ig.y, xtrue; clim)
p2 = jim(sg.r, sg.ad, sino, aspect_ratio=:none)

plan=fbp2(sg, ig)
println("plan complete")
result,sino_filtered=fbp2(plan,sino)
println("fbp complete")

p3 = jim(ig.x, ig.y, result; clim)
p4 = jim(real(sino_filtered))
jim(p1,p2,p3,p4)

=#


using Plots, LaTeXStrings, JLD, Statistics
TotalTime=1000
probability=[0.25,0.5,0.75,0.95]
function RandomWalk(TotalTime, probability)
    arr=zeros((TotalTime,length(probability)))
    for j in 1:length(probability)
        for i in 1:TotalTime
            if probability[j]>rand()
                arr[i,j]=1
            else
                arr[i,j]=-1
            end
        end
    end
    return cumsum(arr,dims=1)
end
#TheMatrix=RandomWalk(TotalTime, probability)
run_num=10000
# a list cinsisting of lists to the number of probabilities,
#each list having lists of each position in each run number per time:
RandomWalkers=[[[] for j in 1:TotalTime] for i in 1:length(probability)]
for n in 1:run_num
    TheMatrix=RandomWalk(TotalTime, probability)
    for col in 1:length(probability)
        for i in 1:TotalTime
            push!(RandomWalkers[col][i],TheMatrix[i,col])
        end
    end
end
SavedData_mean=[[] for i in 1:length(probability)]
SavedData_var=[[] for i in 1:length(probability)]
for i in 1:length(RandomWalkers)
    for j in RandomWalkers[i]
        push!(SavedData_var[i],var(j))
        push!(SavedData_mean[i],mean(j))
    end
end
save("C:\\Users\\Narges\\Documents\\GitHub\\computational_physics\\chapter5\\4.2\\chapter5_4.2_MeanAndStdData.jld",
 "MeanData", SavedData_mean, "VarData", SavedData_var)
 #plots:
load("C:\\Users\\Narges\\Documents\\GitHub\\computational_physics\\chapter5\\4.2\\chapter5_4.2_MeanAndStdData.jld")
scatter(dpi=400)
for i in 1:length(probability)
    p=probability[i]
    scatter!(1:TotalTime, SavedData_mean[i],label=L"%$p",markersize=1,markerstrokewidth=0, alpha=0.7)
end
scatter!(title="<x> per time for different probabilities", xlabel="t", ylabel="<x>")
savefig("C:\\Users\\Narges\\Documents\\GitHub\\computational_physics\\chapter5\\4.2\\Figs\\MeanPlot.png")
scatter(dpi=400)
for i in 1:length(probability)
    p=probability[i]
    scatter!(1:TotalTime, SavedData_var[i],label=L"%$p",markersize=1,markerstrokewidth=0, alpha=0.7)
end
scatter!(title="Var[x] per time  for different probabilities", xlabel="t", ylabel="Var[x]")
savefig("C:\\Users\\Narges\\Documents\\GitHub\\computational_physics\\chapter5\\4.2\\Figs\\VarPlot.png")


push!(LOAD_PATH, expanduser(".")) # Assumed to be ran from the repl folder.

#include("RadioPropagation.jl")
using RadioPropagation
using Plots

frequency_ghz = collect(1:1000);

attenuation_db_km = RadioPropagation.atmospheric_attenuation_db_per_km.( frequency_ghz );

plot1 = plot(	frequency_ghz, attenuation_db_km,
	 			xlabel 	= "Frequency [GHz]",
				ylabel 	= "One-way loss [dB/km]",
				title  	= "Atmospheric Attenuation.",
				#ylims 	= (),
				label 	= "Standard",
				yaxis=:log,
				dpi=300
				)

attenuation_db_km = RadioPropagation.atmospheric_attenuation_db_per_km.( frequency_ghz, 288.15, 0 );

plot!( 	frequency_ghz, attenuation_db_km,
	 	label = "Dry air",
	 	yaxis=:log,
	 	)

savefig("figures/attenuation_example")


using LinearAlgebraicRepresentation, ViewerGL, SparseArrays
Lar = LinearAlgebraicRepresentation; GL = ViewerGL
using QHull

V,CV = Lar.rod()([4,1])
ch = QHull.chull(convert(Lar.Points,V'))
FV = ch.simplices
EV = Lar.simplexFacets(FV)
GL.VIEW([ GL.GLGrid(V,EV) ]);

cyl = Lar.Struct([ Lar.t(0,0,-1.5), (V,FV,EV) ])
tris = Lar.Struct([ cyl,
                    Lar.Struct([Lar.r(pi/2,0,0), cyl]),
                    Lar.Struct([Lar.r(0,pi/2,0), cyl])
                    ])
V,FV,EV = Lar.struct2lar(tris)
GL.VIEW([ GL.GLGrid(V,FV) ]);

VV = [[k] for k=1:size(V,2)]
GL.VIEW( GL.numbering(.2)((V,[VV,EV,FV]), GL.COLORS[4]) );

cop_EV = convert(Lar.ChainOp, Lar.coboundary_0(EV::Lar.Cells));
cop_FE = Lar.coboundary_1(V, FV::Lar.Cells, EV::Lar.Cells);
W = convert(Lar.Points, V');
V, copEV, copFE, copCF = Lar.Arrangement.spatial_arrangement( W, cop_EV, cop_FE)

W = convert(Lar.Points, V')
V,CVs,FVs,EVs = Lar.pols2tria(W, copEV, copFE, copCF)

GL.VIEW(GL.GLExplode(V,FVs,1.5,1.5,1.5,99,1));
GL.VIEW(GL.GLExplode(V,EVs,1.5,1.5,1.5,99,1));
meshes = GL.GLExplode(V,CVs[1:end],8,4,6,99,1);
GL.VIEW( push!( meshes, GL.GLFrame) );


# Random forests regression

"""
Finds the threshold to split `features` with that minimizes the
mean-squared-error loss over `labels`.
Returns (best_val, best_thresh), where `best_val` is -MSE

Variable name changes from original code
  labels   -> y
  features -> x
  nl       -> n_left
  nr       -> n_right
  su       -> sum_y
  su2      -> sum_y2
  s_l      -> sum_y_left
  s2_l     -> sum_y2_left
  s_r      -> sum_y_right
  s2_r     -> sum_y2_right

"""
function _best_mse_loss{T<:Float64, U<:Real}(y::Vector{T}, x::DataArray{U, 1}, domain)
    best_val = -Inf
    best_thresh = 0.0
    n = length(y)

    sum_y_left = sum_y2_left = zero(T)            # scalar values of 0

    sum_y = sum(y)                             # scalar sum of all y_i
    sum_y2 = zero(T);                             # scalar value of 0

    # get sum of squares
    for i = 1:n
        sum_y2 += y[i]^2
    end

    n_left = 0
    i = 1

    # Since`x` is sorted, below is an O(n) algorithm for finding the optimal
    # threshold in `domain`. We iterate through the array and update sum_y_left
    # and sum_y_right (= sum(y) - sum_y_left) as we go. - @cstjean
    for thresh in domain

        # this loop checks which side of the split this x_i is on
        while i <= n && x[i] < thresh
            sum_y_left += y[i]
            sum_y2_left += y[i]^2
            n_left += 1
            i += 1
        end
        sum_y_right = sum_y - sum_y_left
        sum_y2_right = sum_y2 - sum_y2_left
        n_right = n - n_left

        # This check is necessary I think because in theory all y could
        # be the same, then either n_left or n_right would be 0. - @cstjean
        if n_right > 0 && n_left > 0

            # This isn't really squared-error loss. We're computing the sum of variance
            # in the daughter nodes times the N for that daughter node. This last point
            # is why the two paranthetical terms here look a bit diffent than mere variance.
            # We have canceled a constant term N for each term.
            loss = (sum_y2_left - sum_y_left^2/n_left) + (sum_y2_right - sum_y_right^2/n_right)

            # update best value and threshold
            if -loss > best_val
                best_val = -loss
                best_thresh = thresh
            end
        end
    end
    return best_val, best_thresh
end

"""
This function returns a tuple with the column index and the threshold
value of the predictor in the matrix X that minimizes MSE. The function
assumes complete data in both y and X.

Variable name changes from original code
  labels     -> y
  features   -> x
  nr         -> n
  nf         -> p
  i          -> j
  features_i -> x_j
  labels_i   -> y_ord
  domain_i   -> domain_j
  inds       -> col_indcs

"""

function _split_mse_df{T<:Float64}(y::Vector{T}, X::DataFrame, mtry::Int)
    n, p = size(X)
    best = (0, 0.0)
    best_val = -Inf
    if mtry > 0
        r = randperm(p)
        col_indcs = r[1:mtry]
    else
        col_indcs = collect(1:p)
    end
    for j in col_indcs
        keep_row = !isna(X[:, j])

        # x_obs = convert(Vector, X[keep_row, j])

        if !any(keep_row)
            continue
        end

        # y_obs = y[keep_row]
        # ord::Array{Int64, 1} = sortperm(x_obs)
        # x_j = convert(Array{typeof(x_obs[1])}, x_obs[ord])
        # y_ord = y_obs[ord]

        x_j, y_ord = prepare_xy_obs(X[keep_row, j], y[keep_row])

        if n > 100
            if VERSION >= v"0.4.0-dev"
                domain_j = quantile(x_j, linspace(0.01, 0.99, 99); sorted=true)
            else  # sorted=true isn't supported on StatsBase's Julia 0.3 version
                domain_j = quantile(x_j, linspace(0.01, 0.99, 99))
            end
        else
            domain_j = x_j
        end
        value, thresh = _best_mse_loss(y_ord, x_j, domain_j)
        if value > best_val
            best_val = value
            best = (j, thresh)
        end
    end
    return best
end

function prepare_xy_obs{T}(x_obs::DataArray{T, 1}, y_obs)
    ord = sortperm(x_obs)
    x_j::DataArray{T, 1} = x_obs[ord]
    y_ord = y_obs[ord]
    return (x_j, y_ord)
end

function find_na_cols(dat)
    cols_with_na = Array{Int, 1}(0)

    for j in 1:ncol(dat)
        if sum(isna(dat[:, j])) ≠ 0
            push!(cols_with_na, j)
        end
    end
    return cols_with_na
end

# function find_surrogates{T::BitArray}(left_node::T, y::Vector, X::Matrix, mtry::Int)

function build_stump{T <: Float64, U<:Real}(y::Vector{T}, X::Matrix{U})
    S = _split_mse(y, X, 0)

    if S == NO_BEST
        return Leaf(mean(y), y)
    end

    col_idx, thresh = S
    split = X[:, col_idx] .< thresh

    return Node(col_idx,
                thresh,
                Leaf(mean(y[split]), y[split]),
                Leaf(mean(y[!split]), y[!split]))
end

function build_tree_df{T <: Float64}(y::Vector{T}, X::DataFrame, maxlabels = 5, mtry = 0, maxdepth = -1, n_surrogates = 5)
    n, p = size(X)
    if mtry == 0
        mtry = p
    end
    if maxdepth < -1
        error("Unexpected value for maxdepth: $(maxdepth) (expected: maxdepth >= 0, or maxdepth = -1 for infinite depth)")
    end
    if length(y) <= maxlabels || maxdepth == 0          # stopping rules
        return Leaf(mean(y), y)
    end

    S = _split_mse_df(y, X, mtry)                       # get [complete data] optimal split point

    if S == NO_BEST
        return Leaf(mean(y), y)
    end

    col_idx, thresh = S

    if n_surrogates ≠ 0
        cols_with_na = find_na_cols(X)                  # needed for all recursive steps since rows of X change
        surrogate_vars = Vector{Tuple{Int, Real}}(n_surrogates)
    end

    # if col_idx in cols_with_na
    if col_idx in 1:p                             # find surrogates even when no missing data
        na_rows = isna(X[:, col_idx])
        split_with_na = Array{Any, 1}(n)                # vector of Bools with some NA values

        for i = 1:n
            split_with_na[i] = isna(X[i, col_idx]) ? NA : X[i, col_idx] < thresh
        end
        row_indcs = collect(1:n)[!na_rows]
        col_indcs = deleteat!(collect(1:p), col_idx)

        # Here we need a function that splits so as to optimize agreement
        # with the `split_with_na` result for each observed `row_indcs`.
        surrogate_vars = surrogate_splits(split_with_na, X, row_indcs, col_indcs, 5)
        split = apply_surrogates(split_with_na, X, surrogate_vars)
    else
        split = X[:, col_idx] .< thresh
    end

    return Node(col_idx,
                thresh,
                surrogate_vars,
                build_tree_df(y[split], X[split,:], maxlabels, mtry, max(maxdepth-1, -1)),
                build_tree_df(y[!split], X[!split,:], maxlabels, mtry, max(maxdepth-1, -1)))
end

function build_forest_df{T <: Real}(y::Vector{T}, X::DataFrame, mtry, ntrees, maxlabels = 5, maxdepth = -1; nthreads = 1, n_surrogates = 5, oob_measure = "rsq")

    n = length(y)
    tree_arr = Array{Node, 1}(ntrees)

    if nthreads ≥ 2
        yhat_mat = fill(-Inf, (n, ntrees))              # use to compute oob-score
        mse_inc_dicts = Array{Dict{Int, Float64}, 1}(ntrees)

        @threads for t in 1:ntrees
            inds = sample(1:n, n)
            tree_arr[t] = build_tree_df(y[inds], X[inds, :], maxlabels, mtry, maxdepth, n_surrogates)

            # get OOB indices to calculate OOB score
            oob_indcs = setdiff(collect(1:n), inds)
            yhat, mse_inc = apply_tree(tree_arr[t], X[oob_indcs, :], y[oob_indcs])

            yhat_mat[oob_indcs, t] = yhat
            mse_inc_dicts[t] = mse_inc              # dict with MSE increases for OOB permutation importance
        end

        oob_score = par_oob_score(yhat_mat, y, oob_measure)
        mse_inc_dict = par_mse_inc(mse_inc_dicts)
    else
        yhat_mat = zeros(n, 2)                      # col 1 is numerator, col 2 is denominator
        yhat_mat[:, 1] = fill(-Inf, n)
        mse_inc_dict = Dict{Int, Float64}()

        for t in 1:ntrees
            inds = sample(1:n, n)
            tree_arr[t] = build_tree_df(y[inds], X[inds, :], maxlabels, mtry, maxdepth, n_surrogates)

            # get OOB indices to calculate OOB score
            oob_indcs = setdiff(collect(1:n), inds)
            yhat, mse_inc = apply_tree(tree_arr[t], X[oob_indcs, :], y[oob_indcs])

            if !all(isfinite(yhat_mat[oob_indcs, 1]))
                for idx in oob_indcs
                    if !isfinite(yhat_mat[idx, 1])
                        yhat_mat[oob_indcs, 1] = 0.0
                    end
                end
            end

            yhat_mat[oob_indcs, 1] += yhat
            yhat_mat[oob_indcs, 2] += 1
            update_mse_inc!(mse_inc_dict, mse_inc)
        end

        if all(isfinite(yhat_mat[:, 1]))

            if oob_measure == "rsq"
                oob_score = rsq(y, yhat_mat[:, 1] ./ yhat_mat[:, 2])
            elseif oob_measure == "mse"
                oob_score = mean_squared_error(y, yhat_mat[:, 1] ./ yhat_mat[:, 2])
            end
        else
            warn("Not enough trees built to compute OOB score. Some rows were never out-of-bag")
            return nothing
        end
    end
    return (Ensemble(tree_arr), oob_score, mse_inc_dict)
end

function build_forest_df_novarimp{T <: Real}(y::Vector{T}, X::DataFrame, mtry, ntrees, maxlabels = 5, maxdepth = -1; nthreads = 1, n_surrogates = 5, oob_measure = "rsq")

    n = length(y)
    tree_arr = Array{Node, 1}(ntrees)

    if nthreads ≥ 2
        yhat_mat = fill(-Inf, (n, ntrees))       # use to compute oob-score

        @threads for t in 1:ntrees
            inds = sample(1:n, n)
            tree_arr[t] = build_tree_df(y[inds], X[inds, :], maxlabels, mtry, maxdepth, n_surrogates)

            # get OOB indices to calculate OOB score
            oob_indcs = setdiff(collect(1:n), inds)
            yhat = apply_tree(tree_arr[t], X[oob_indcs, :])

            yhat_mat[oob_indcs, t] = yhat
        end
        oob_score = par_oob_score(yhat_mat, y, oob_measure)
    else
        yhat_mat = zeros(n, 2)                  # col 1 is numerator, col 2 is denominator
        yhat_mat[:, 1] = fill(-Inf, n)

        for t in 1:ntrees
            inds = sample(1:n, n)
            tree_arr[t] = build_tree_df(y[inds], X[inds, :], maxlabels, mtry, maxdepth, n_surrogates)

            # get OOB indices to calculate OOB score
            oob_indcs = setdiff(collect(1:n), inds)
            yhat = apply_tree(tree_arr[t], X[oob_indcs, :])

            if !all(isfinite(yhat_mat[oob_indcs, 1]))
                for idx in oob_indcs
                    if !isfinite(yhat_mat[idx, 1])
                        yhat_mat[oob_indcs, 1] = 0.0
                    end
                end
            end

            yhat_mat[oob_indcs, 1] += yhat
            yhat_mat[oob_indcs, 2] += 1
        end

        if all(isfinite(yhat_mat[:, 1]))

            if oob_measure == "rsq"
                oob_score = rsq(y, yhat_mat[:, 1] ./ yhat_mat[:, 2])
            elseif oob_measure == "mse"
                oob_score = mean_squared_error(y, yhat_mat[:, 1] ./ yhat_mat[:, 2])
            end
        else
            warn("Not enough trees built to compute OOB score. Some rows were never out-of-bag")
            return nothing
        end
    end
    return (Ensemble(tree_arr), oob_score)
end

# This function is a helper used for getting oob-score for
# parallel version of random forest regression. The parallel case
# needs to be treated a bit differently to avoid threads overwriting
# each other. NOTE: There might be a more elegant solution.
function par_oob_score(yhat_mat, y, measure = "rsq")
    n, ntrees = size(yhat_mat)
    yhat = fill(-Inf, n)
    num_preds = zeros(Int, n)       # this will be our denomiator

    for t = 1:ntrees
        for i = 1:n
            if isfinite(yhat_mat[i, t])
                if !isfinite(yhat[i])
                    yhat[i] = 0.0
                end
                yhat[i] += yhat_mat[i, t]
                num_preds[i] += 1
            end
        end
    end
    if all(isfinite(yhat))
        yhat ./= num_preds
    else
        warn("Not enough trees built to compute OOB score. Some rows never out-of-bag")
        return nothing
    end
    if measure == "rsq"
        oob_score = rsq(y, yhat)
    elseif measure == "mse"
        oob_score = mean_squared_error(y, yhat)
    end
    return oob_score
end

function update_mse_inc!(aggregate_dict, tree_dict)
    for k in keys(tree_dict)
        aggregate_dict[k] = get(aggregate_dict, k, 0.0) + tree_dict[k]
    end
end

function par_mse_inc(mse_inc_dicts)
    out = Dict{Int, Float64}()
    ntrees = length(mse_inc_dicts)
    for i = 1:ntrees
        update_mse_inc!(out, mse_inc_dicts[i])
    end
    return out
end


#=

  Euler #28 in Julia.

  Problem 28:
  """
  Starting with the number 1 and moving to the right in a clockwise
  direction a 5 by 5 spiral is formed as follows:

     21 22 23 24 25
     20  7  8  9 10
     19  6  1  2 11
     18  5  4  3 12
     17 16 15 14 13

  It can be verified that the sum of the numbers on the diagonals is 101.

  What is the sum of the numbers on the diagonals in a 1001 by 1001 spiral formed in the same way?
  """

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

=#

include("Euler.jl")

#
function euler28a()
    s = 1
    n = 3
    while n <= 1001
        s += 4*n^2 - 6*n+6
        n+=2
    end
    return s
end

run_euler(euler28a)


# ### m10.3t.jl

using Pkg, DrWatson

@quickactivate "StatisticalRethinkingTuring"
using Turing
using StatisticalRethinking
Turing.setprogress!(false)

delim=';'
df = CSV.read(sr_datadir("chimpanzees.csv"), DataFrame; delim);

# pulled_left, condition, prosoc_left
@model function ppl10_3(y, x₁, x₂)
    α ~ Normal(0, 10)
    βp ~ Normal(0, 10)
    βpC ~ Normal(0, 10)

    logits = α .+ (βp .+ βpC * x₁) .* x₂
    y .~ BinomialLogit.(1, logits)
end;

m10_3t = ppl10_3(df.pulled_left, df.condition, df.prosoc_left)
nchains = 4; sampler = NUTS(0.65); nsamples=2000
chns10_3t = mapreduce(c -> sample(m10_3t, sampler, nsamples), chainscat, 1:nchains)

# Rethinking result

m_10_03t_result = "
      Mean StdDev lower 0.89 upper 0.89 n_eff Rhat
 a    0.05   0.13      -0.15       0.25  3284    1
 bp   0.62   0.22       0.28       0.98  3032    1
 bpC -0.11   0.26      -0.53       0.29  3184    1
";

# End of m10.03t.jl


#
#   SVD decomposition of a matrix truncated to a rank
#
using Manopt, Manifolds
import LinearAlgebra: norm, svd, Diagonal

"""
    truncated_svd(A, p)

return a singular value decomposition of a real valued matrix A truncated to
rank p.

# Input
* `A` – a real-valued matrix A of size mxn
* `p` – an integer p in min { m, n }

# Output
* `U` – an orthonormal matrix of size mxp
* `V` – an orthonormal matrix of size nxp
* `S` – a diagonal matrix of size pxp with nonnegative and decreasing diagonal
        entries
"""
function truncated_svd(A::Array{Float64,2} = randn(42, 60), p::Int64 = 5)
    (m, n) = size(A)

    if p > min(m,n)
        throw( ErrorException("The Rank p=$p must be smaller than the smallest dimension of A = $min(m, n).") )
    end

    M = ProductManifold(Grassmann(m,p), Grassmann(n,p))

    function cost(X::ProductRepr)
        return cost([submanifold_components(X)...])
    end
    function cost(X::Array{Matrix{Float64},1})
        return -0.5 * norm(transpose(X[1]) * A * X[2])^2
    end

    function egrad(X::Array{Matrix{Float64},1})
        U = X[1]
        V = X[2]
        AV = A*V
        AtU = transpose(A)*U
        return [ -AV*(transpose(AV)*U), -AtU*(transpose(AtU)*V) ];
    end

    function rgrad(M::ProductManifold, X::ProductRepr)
        eG = egrad([submanifold_components(M,X)...])
        x = [submanifold_components(M,X)...]
        return Manifolds.ProductRepr(project.(M.manifolds, x, eG)...)
    end

    function e2rHess(M::Grassmann, x, ξ, eGrad::Matrix{T},Hess::Matrix{T}) where T<:Union{U, Complex{U}} where U<:AbstractFloat
        pxHess = project(M,x,Hess)
        xtGrad = x'*eGrad
        ξxtGrad = ξ*xtGrad
        return pxHess - ξxtGrad
    end

    function eHess(X::Array{Matrix{Float64},1}, H::Array{Matrix{Float64},1})
        U = X[1]
        V = X[2]
        Udot = H[1]
        Vdot = H[2]
        AV = A*V
        AtU = transpose(A)*U
        AVdot = A*Vdot
        AtUdot = transpose(A)*Udot
        return [ -(AVdot*transpose(AV)*U + AV*transpose(AVdot)*U + AV*transpose(AV)*Udot),
                 -(AtUdot*transpose(AtU)*V + AtU*transpose(AtUdot)*V + AtU*transpose(AtU)*Vdot)
            ]
    end

    function rhess(M::ProductManifold, X::ProductRepr, H::ProductRepr)
        x = [submanifold_components(M,X)...]
        h = [submanifold_components(M,H)...]
        eG = egrad(x)
        eH = eHess(x,h)
        return Manifolds.ProductRepr(e2rHess.(M.manifolds, x, h, eG, eH)...)
    end

    x = random_point(M)
    X = trust_regions(M, cost, rgrad, x, rhess;
        Δ_bar=4*sqrt(2*p),
        debug = [:Iteration, " ", :Cost, " | ", DebugEntry(:Δ), "\n", 1, :Stop]
    )

    U = get_component(M,X,1)
    V = get_component(M,X,2)

    Spp = transpose(U)*A*V
    SVD = svd(Spp)
    U = U*SVD.U
    S = SVD.S
    V = V*SVD.V

    return [U, S, V]
end

A=[1. 2. 3. 4.; 5. 6. 7. 8.; 9. 10. 11. 12.; 13. 14. 15. 16.]

truncated_svd(A,2)


"""
    LogNormal

* Parametrization
parametrized in \\mu and \\sigma^2

* Score

* Fisher Information

* `time_varying_params` map.

* Default link
"""
LogNormal

function score!(score_til::Matrix{T}, y::T, ::Type{LogNormal}, param::Matrix{T}, t::Int) where T
    score_til[t, 1] = (log(y) - param[t, 1])/param[t, 2]
    score_til[t, 2] = -(0.5/param[t, 2]) * (1 - ((log(y) - param[t, 1])^2)/param[t, 2])
    return
end

function fisher_information!(aux::AuxiliaryLinAlg{T}, ::Type{LogNormal}, param::Matrix{T}, t::Int) where T
    aux.fisher[1, 1] = 1/(param[t, 2])
    aux.fisher[2, 2] = 1/(2*(param[t, 2]^2))
    aux.fisher[2, 1] = 0
    aux.fisher[1, 2] = 0
    return
end

function log_likelihood(::Type{LogNormal}, y::Vector{T}, param::Matrix{T}, n::Int) where T
    loglik = -0.5*log(2*pi)*n
    for t in 1:n
        loglik -= log(y[t] * sqrt(param[t, 2])) + 0.5*(log(y[t]) - param[t, 1])^2/param[t, 2]
    end
    return -loglik
end

# Links
function link!(param_tilde::Matrix{T}, ::Type{LogNormal}, param::Matrix{T}, t::Int) where T
    param_tilde[t, 1] = link(IdentityLink, param[t, 1])
    param_tilde[t, 2] = link(LogLink, param[t, 2], zero(T))
    return
end
function unlink!(param::Matrix{T}, ::Type{LogNormal}, param_tilde::Matrix{T}, t::Int) where T
    param[t, 1] = unlink(IdentityLink, param_tilde[t, 1])
    param[t, 2] = unlink(LogLink, param_tilde[t, 2], zero(T))
    return
end
function jacobian_link!(aux::AuxiliaryLinAlg{T}, ::Type{LogNormal}, param::Matrix{T}, t::Int) where T
    aux.jac[1] = jacobian_link(IdentityLink, param[t, 1])
    aux.jac[2] = jacobian_link(LogLink, param[t, 2], zero(T))
    return
end

# utils
function update_dist(::Type{LogNormal}, param::Matrix{T}, t::Int) where T
    # lognormal here is parametrized as sigma^2
    return LogNormal(param[t, 1], sqrt(param[t, 2]))
end

function params_sdm(d::LogNormal)
    pars = Distributions.params(d)
    return (pars[1], pars[2]^2)
end

function num_params(::Type{LogNormal})
    return 2
end


simulate_bm(t::Tree, k::Int; σ²::Float64=1.0, x₀::Float64=0.0) =
    simulate_bm(t; σ²=fill(σ², k), x₀=fill(x₀, k))

"""
    simulate_bm(t::Tree, k::Int, σ::Vector{Float64} = ones(k), x₀::Vector{Float64} = zeros(k))

Simulate the evolution of continuous traits under univariate Brownian motion on a phylogenetic tree.

Creates a `SpeciesDataMatrix` with `k` independent traits, with `σ²` as the vector of diffusion coeffients of the traits (all 1.0 by default), and `x₀` as the vector of trait values at the root (all 0.0 by default).
"""
function simulate_bm(
    t::Tree;
    σ²::Union{Float64, Vector{Float64}} = 1.0,
    x₀::Union{Float64, Vector{Float64}} = 0.0
    )
    ! t.rooted && @warn "The tree is unrooted. The simulation was done using the origin node as the root."

    if σ² isa Float64 && x₀ isa Float64
        σ² = [σ²]
        x₀ = [x₀]
    elseif σ² isa Float64 && x₀ isa Vector{Float64}
        σ² = fill(σ², length(x₀))
    elseif σ² isa Vector{Float64} && x₀ isa Float64
        x₀ = fill(x₀, length(σ²))
    elseif length(σ²) ≠ length(x₀)
        msg = "the vectors `σ²` and `x₀` must both have `k` number of elements."
        throw(ArgumentError(msg))
    end

    k = length(x₀)
    σ = .√(σ²)

    return _simulate_bm_species(t, k, σ, x₀)
end

function _simulate_bm_species(t::Tree, k::Int, σ::Vector{Float64}, x₀::Vector{Float64})::SpeciesDataMatrix{Float64}
    x = SpeciesDataMatrix{Float64}(t.dir, k)
    for p ∈ neighbours(t.origin)
        _simulate_bm!(x, p, t.origin, k, σ, x₀)
    end

    return x
end

function _simulate_bm!(
    x::SpeciesDataMatrix{Float64},
    p::Node,
    q::Node,
    k::Int,
    σ::Vector{Float64},
    x₀::Vector{Float64}
    )
    xₚ = x₀ .+ σ .* √(brlength(q, p)) .* randn(k)
    if p.species > 0 && p.species ∈ x.dir
        x[p.species,:] = xₚ
    end
    for r in neighbours(p)
        r == q && continue
        _simulate_bm!(x, r, p, k, σ, xₚ)
    end
end


using Distributions
using QuadGK

function kde_kth_neighbor_1d(xout::Vector{T}, xin::Vector{T}; k::Integer = 10, sigma::Real = 2.0) where T <: Real
   issorted(xin) ? kde_kth_neighbor_1d_sorted(xout,xin,k=k,sigma=sigma)  :
                   kde_kth_neighbor_1d_unsorted(xout,xin,k=k,sigma=sigma)
end

function kde_kth_neighbor_1d_unsorted(xout::Vector{T}, xin::Vector{T}; k::Integer = 10, sigma::Real = 2) where T <: Real
   @assert(length(xin)>=k+1)
   @assert issorted(xout)
   yout = similar(xout)
   ilo = 1
   ihi = k
   perm = sortperm(xin)
   #print("# ilo,ihi = ",ilo,", ", ihi, " length(xin)=", length(xin), " len(xout)=",length(xout) )
   #println( " perm[ilo,ihi,ihi+1] = ", perm[ilo],", ", perm[ihi], ", ", perm[ihi+1])
   dlo = abs(xout[1]-xin[perm[ilo]])
   dhi = abs(xin[perm[ihi]]-xout[1])
   dstar = abs(xin[perm[ihi+1]]-xout[1])
   for (i,x) in enumerate(xout)
       while dstar<dlo && ihi<length(xin)
          ilo += 1
          ihi += 1
          dlo = abs(x-xin[perm[ilo]])
          dhi = dstar
          if ihi+1<=length(perm)
            dstar = abs(xin[perm[ihi+1]]-x)
          end
       end
       #width = (xin[perm[ihi]]-xin[perm[ilo]])*sigma/k
       width = max(dlo,dhi)*sigma/k
       distrib = Distributions.Normal(x,width)
       yout[i] = sum(pdf(distrib,xin[perm[ilo:ihi]]))/k
   end
   return yout
end

function kde_kth_neighbor_1d_sorted(xout::Vector{T}, xin::Vector{T}; k::Integer = 10, sigma::Real = 2) where T <: Real
   @assert(length(xin)>=k+1)
   #@assert issorted(xin)
   yout = similar(xout)
   ilo = 1
   ihi = k
   dlo = abs(xout[1]-xin[ilo])
   dhi = abs(xin[ihi]-xout[1])
   dstar = abs(xin[ihi+1]-xout[1])
   for (i,x) in enumerate(xout)
       while dstar<dlo && ihi<length(xin)
          ilo += 1
          ihi += 1
          dlo = abs(x-xin[ilo])
          dhi = dstar
          if ihi+1<=length(perm)
             dstar = abs(xin[ihi+1]-x)
          end
       end
       #width = (xin[ihi]-xin[ilo])*sigma/k
       width = max(dlo,dhi)*sigma/k
       distrib = Distributions.Normal(x,width)
       yout[i] = sum(pdf(distrib,xin[ilo:ihi]))/k
   end
   return yout
end

function kl_integral_arg(q::T, p::T; epsilon::T = 1e-10) where T <: Real
  if q > epsilon
     return p*log(p/q)
  else
     return p
  end
end

function hellinger_integral_arg(p::T, q::T) where T <: Real
  sqrt(p*q)
end

function calc_kl_distance_ab(x1::Vector{T}, x2::Vector{T}, a::T, b::T; n::Integer = 100, k::Integer = max(min((length(x1) - 1), (length(x2) - 1), 20), 1)) where T <: Real
  @assert b>a
  @assert n>=2
  if(length(x1)<3 || length(x2)<3) return 0.0  end
  xgrid = collect(linspace(a,b,n+1))
  y1 = kde_kth_neighbor_1d(xgrid,x1,k=k)
  y2 = kde_kth_neighbor_1d(xgrid,x2,k=k)
  integral = 0.5*(kl_integral_arg(y1[1],y2[1])+kl_integral_arg(y1[end],y2[end]))
  integral += sum([kl_integral_arg(y1[i],y2[i]) for i in 2:(length(y1)-1) ])
  integral /= n
end

function calc_kl_distance(x1::Vector{T}, x2::Vector{T}; n::Integer = 100, k::Integer = max(min((length(x1) - 1), (length(x2) - 1), 20), 1)) where T <: Real
  @assert b>a
  @assert n>=2
  n1 = length(x1)
  n2 = length(x2)
  xmin = minimum(x1,x2)
  xmax = maximum(x1,x2)
  padding = k*(xmax-xmin)/(n1+n2)
  calc_kl_distance_ab(x1,x2,xmin,xmax,n=n,k=k)
end


struct ImbensKalyanaraman <: BandwidthSelector end

_string(::ImbensKalyanaraman) = "Imbens Kalyanaraman bandwidth"

function kernel_constant(::ImbensKalyanaraman, kernel)
    kernel = EquivalentKernel(kernel)
    νs = [kernel_moment(kernel, Val(j)) for j = 0:3]
    πs = [squared_kernel_moment(kernel, Val(j)) for j = 0:3]
    C1 = 1 / 4 * abs2((abs2(νs[3]) - νs[2] * νs[4]) / (νs[3] * νs[1] - abs2(νs[2])))
    C2_num = abs2(νs[3]) * πs[1] - 2 * νs[2] * νs[3] * πs[2] + abs2(νs[2]) * πs[3]
    C2_denom = abs2(νs[3] * νs[1] - abs2(νs[2]))
    C2 = C2_num / C2_denom
    (C2 / (4 * C1))^(1 / 5)
end

function bandwidth(ik::ImbensKalyanaraman, kernel::SupportedKernels, ZsR::RDData)
    ZsR_untreated = ZsR[Untreated()]
    ZsR_treated = ZsR[Treated()]

    cutoff = ZsR.cutoff

    N = nobs(ZsR)

    N_untreated = nobs(ZsR_untreated)
    N_treated = nobs(ZsR_treated)

    # Step 1: Density and conditional variance at 0
    h₁ = 1.84 * Statistics.std(ZsR.Zs) * N^(-1 / 5)

    interval_h₁ = Interval{:closed,:closed}(cutoff-h₁, cutoff+h₁)

    ZsR_untreated_h₁ = ZsR_untreated[interval_h₁]
    ZsR_treated_h₁ = ZsR_treated[interval_h₁]

    N_h₁_untreated = nobs(ZsR_untreated_h₁)
    N_h₁_treated = nobs(ZsR_treated_h₁)

    Ȳ_h₁_untreated = mean(ZsR_untreated_h₁.Ys)
    sd_Y_h₁_untreated = Statistics.std(ZsR_untreated_h₁.Ys)
    Ȳ_h₁_treated = mean(ZsR_treated_h₁.Ys)
    sd_Y_h₁_treated = Statistics.std(ZsR_treated_h₁.Ys)

    f̂₀ = (N_h₁_untreated + N_h₁_treated) / 2 / N / h₁

    # Step 2: Estimation of second derivatives
    global_cubic_lm = lm(@formula(Ys ~ 1 + (ZsC >= 0) + ZsC + ZsC^2 + ZsC^3), ZsR)

    m̂₀_triple_prime = 6 * coef(global_cubic_lm)[5]

    h₂_untreated = 3.56 * (sd_Y_h₁_untreated^2 / f̂₀ / m̂₀_triple_prime^2)^(1 / 7) * N_untreated^(-1 / 7)
    h₂_treated =
        3.56 * (sd_Y_h₁_treated^2 / f̂₀ / m̂₀_triple_prime^2)^(1 / 7) * N_treated^(-1 / 7)

    ZsR_untreated_h₂ = ZsR_untreated[Interval{:closed,:closed}(cutoff-h₂_untreated, cutoff+h₂_untreated)]
    ZsR_treated_h₂ = ZsR_treated[Interval{:closed,:closed}(cutoff-h₂_treated, cutoff+h₂_treated)]

    N_h₂_untreated = nobs(ZsR_untreated_h₂)
    N_h₂_treated = nobs(ZsR_treated_h₂)

    quadratic_fit_untreated = lm(@formula(Ys ~ 1 + ZsC + ZsC^2), ZsR_untreated_h₂)
    m̂₀_double_prime_untreated = 2 * last(coef(quadratic_fit_untreated))

    quadratic_fit_treated = lm(@formula(Ys~ 1 + ZsC + ZsC^2), ZsR_treated_h₂)
    m̂₀_double_prime_treated  = 2 * last(coef(quadratic_fit_treated ))

    # Step 3: Calculation of regularization terms

    r̂_untreated = 2160 * sd_Y_h₁_untreated^2 / N_h₂_untreated/ h₂_untreated^4
    r̂_treated = 2160 * sd_Y_h₁_treated^2 / N_h₂_treated / h₂_treated^4

    Ckernel = kernel_constant(ik, kernel)# 3.4375

    regularized_squared_diff =
        abs2(m̂₀_double_prime_treated - m̂₀_double_prime_untreated) + r̂_untreated + r̂_treated
    ĥ_IK =
        Ckernel *
        ((sd_Y_h₁_untreated^2 + sd_Y_h₁_treated^2) / f̂₀ / regularized_squared_diff)^(1 / 5) *
        N^(-1 / 5)
    ĥ_IK
end


using Jecco
using LinearAlgebra
using SparseArrays
using Plots

source(x,y) = exp(-x^2 - y^2) * (-4 + 3 * (x^2 + y^2) + 4 * x^2 * y^2)

# returns axx Dxx + ayy Dyy + axy Dxy + bx Dx + by Dy + cc. note that this
# function overwrites the input matrices to save memory
function build_operator(Dxx::SparseMatrixCSC, Dyy::SparseMatrixCSC, Dxy::SparseMatrixCSC,
                        Dx::SparseMatrixCSC, Dy::SparseMatrixCSC,
                        axx::Diagonal, ayy::Diagonal, axy::Diagonal,
                        bx::Diagonal, by::Diagonal, cc::Diagonal)
    Dxx_ = similar(Dxx)
    Dyy_ = similar(Dyy)
    Dxy_ = similar(Dxy)
    Dx_  = similar(Dx)
    Dy_  = similar(Dy)

    mul!(Dxx_, axx, Dxx)
    mul!(Dyy_, ayy, Dyy)
    mul!(Dxy_, axy, Dxy)
    mul!(Dx_, bx, Dx)
    mul!(Dy_, by, Dy)

    Dxx_ + Dyy_ + Dxy_ + Dx_ + Dy_ + cc
end

#=
use the Kronecker product (kron) to build the 2-dimensional derivation matrices
from the 1-dimensional ones. see for instance:

  https://en.wikipedia.org/wiki/Kronecker_product

  https://arxiv.org/pdf/1801.01483.pdf (section 5)
=#
function deriv_operators(hx, hy, Nx::Int, Ny::Int, ord::Int)
    Dx_op  = CenteredDiff{1}(1, ord, hx, Nx)
    Dxx_op = CenteredDiff{1}(2, ord, hx, Nx)

    Dy_op  = CenteredDiff{2}(1, ord, hy, Ny)
    Dyy_op = CenteredDiff{2}(2, ord, hy, Ny)

    Dx  = kron(I(Ny), SparseMatrixCSC(Dx_op))
    Dxx = kron(I(Ny), SparseMatrixCSC(Dxx_op))
    Dy  = kron(SparseMatrixCSC(Dy_op), I(Nx))
    Dyy = kron(SparseMatrixCSC(Dyy_op), I(Nx))

    Dx, Dy, Dxx, Dyy, Dx * Dy
end

x_min    = -5.0
x_max    =  5.0
x_nodes  =  128
y_min    = -5.0
y_max    =  5.0
y_nodes  =  64

ord = 4

xcoord  = Cartesian{1}("x", x_min, x_max, x_nodes, endpoint=false)
ycoord  = Cartesian{2}("y", y_min, y_max, y_nodes, endpoint=false)

hx = Jecco.delta(xcoord)
hy = Jecco.delta(ycoord)

Nx = xcoord.nodes
Ny = ycoord.nodes

f_exact = [exp(-xcoord[i]^2 - ycoord[j]^2) for i in 1:Nx, j in 1:Ny]

Dx, Dy, Dxx, Dyy, Dxy = deriv_operators(hx, hy, Nx, Ny, ord)

f0    = zeros(Nx,Ny)
ind2D = LinearIndices(f0)

M = Nx * Ny

b_vec = zeros(M)

axx     = ones(M)
ayy     = ones(M)
axy     = zeros(M)
bx      = zeros(M)
by      = zeros(M)
cc      = zeros(M)

for j in 1:Ny, i in 1:Nx
    idx = ind2D[i,j]

    xi  = xcoord[i]
    yi  = ycoord[j]

    axy[idx] = xi * yi
    bx[idx]  = xi
    by[idx]  = yi
    cc[idx]  = xi^2 + yi^2

    b_vec[idx] = source(xi, yi)
end

# build operator A = Dxx + Dyy + x y Dxy + x Dx + y Dy + (x^2 + y^2)
A_mat = build_operator(Dxx, Dyy, Dxy, Dx, Dy, Diagonal(axx), Diagonal(ayy),
                       Diagonal(axy), Diagonal(bx), Diagonal(by), Diagonal(cc))

A_fact = factorize(A_mat)
sol    = A_fact \ b_vec

@inbounds for idx in eachindex(f0)
    f0[idx] = sol[idx]
end

j_slice = div(Ny,2) + 1
x = xcoord[:]

plot(x, f_exact[:,j_slice])
scatter!(x, f0[:,j_slice])


using DiffEqBayes
using Distributions
using OrdinaryDiffEq, RecursiveArrayTools, ParameterizedFunctions
using Plots

gr(fmt=:png)

g1 = @ode_def LorenzExample begin
  dx = σ*(y-x)
  dy = x*(ρ-z) - y
  dz = x*y - β*z
end σ ρ β

r0 = [1.0; 0.0; 0.0]
tspan = (0.0, 30.0)
p = [10.0,28.0,2.66]

prob = ODEProblem(g1,r0,tspan,p)
sol = solve(prob,Tsit5())

t = collect(range(1,stop=30,length=30))
sig = 0.49
data = convert(Array, VectorOfArray([(sol(t[i]) + sig*randn(3)) for i in 1:length(t)]))

Plots.scatter(t, data[1,:],markersize=4,color=:purple)
Plots.scatter!(t, data[2,:],markersize=4,color=:yellow)
Plots.scatter!(t, data[3,:],markersize=4,color=:black)
plot!(sol)

cb = AdaptiveProbIntsUncertainty(5)
monte_prob = MonteCarloProblem(prob)
sim = solve(monte_prob,Tsit5(),num_monte=100,callback=cb,reltol=1e-5,abstol=1e-5)
plot(sim,vars=(0,1),linealpha=0.4)

cb = AdaptiveProbIntsUncertainty(5)
monte_prob = MonteCarloProblem(prob)
sim = solve(monte_prob,Tsit5(),num_monte=100,callback=cb,reltol=1e-6,abstol=1e-6)
plot(sim,vars=(0,1),linealpha=0.4)

cb = AdaptiveProbIntsUncertainty(5)
monte_prob = MonteCarloProblem(prob)
sim = solve(monte_prob,Tsit5(),num_monte=100,callback=cb,reltol=1e-8,abstol=1e-8)
plot(sim,vars=(0,1),linealpha=0.4)

priors = [Truncated(Normal(10,2),1,15),Truncated(Normal(30,5),1,45),Truncated(Normal(2.5,0.5),1,4)]

@time bayesian_result = stan_inference(prob,t,data,priors;reltol=1e-8,abstol=1e-8,vars=(StanODEData(),InverseGamma(3,2)))

plot_chain(bayesian_result)

@time bayesian_result_turing = turing_inference(prob,Tsit5(),t,data,priors)

plot_chain(bayesian_result_turing)


### A Pluto.jl notebook ###
# v0.11.2

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 9d996162-d41c-11ea-09b4-994ed99dd014

# ╔═╡ 9d3661fc-d41c-11ea-1997-175b6d9a4095

# ╔═╡ d5eaa180-d41b-11ea-32f6-6ff398c3e4c2
begin
	using LinearAlgebra
	using PlutoUI
	using StatsKit
	import Plots
	import StatsPlots

	md"""
	# Julia 102: Estatística com Pluto.jl

	Vamos explorar um pouco as distribuições e funções estatísticas no Julia. A primeira coisa a notar é que o Julia básico contém o [pacote](https://docs.julialang.org/en/v1/stdlib/Statistics/) `Statistics`, com operações bem simples como `mean()` e `cov()`, e o [pacote](https://docs.julialang.org/en/v1/stdlib/Random/) `Random`, para a geração de números (pseudo)aleatórios. No entanto, há todo um ambiente de computação estatística em desenvolvimento para a linguagem. Vamos utilizar o [meta-pacote](https://github.com/JuliaStats/StatsKit.jl) `StatsKit` que importa alguns dos principais pacotes estatísticos.

	Inicialmente, vamos explorar algumas funções e visualizações de interesse para a Estatística por meio das distribuicões univariadas do pacote [Distributions](https://juliastats.org/Distributions.jl/stable/univariate/).
	"""
end

# ╔═╡ 61be6a20-d41c-11ea-0eee-654e9ec43f39
begin
	md"""
	Vamos brincar um pouco com uma distribuição Normal na qual escolheremos os parâmetros μ e σ², comparando o seu gráfico com o da Normal Padrão.

	**Selecione o parâmetro μ da Normal**:

	$(@bind μ Slider(-50.0:0.1:50.0, 0))

	**Selecione o parâmetro σ² da Normal**:

	$(@bind σ² Slider(1:0.01:50.0))

	Acima podemos conferir o resultado, em que a Normal Padrão está em azul escuro e a nossa Normal está em vermelho. Note que o valor inicial de μ é `0` e o σ² é `1`, logo as curvas estão sobrepostas.
	"""
end

# ╔═╡ 517a652c-d41c-11ea-2118-fdbf45828452
md"""### Distribuição Normal
**Valor de μ**: $μ

**Valor de σ²**: $σ²
"""

# ╔═╡ 57bc858e-d41c-11ea-392b-f1c840845820
#Visualização das distribuições
begin
	Plots.plot(Normal(μ, σ²), color = :red, lw = 2, legend = false, fill = (0, :red), fillalpha = 0.5)
	Plots.plot!(Normal(0,1), color = :royalblue, lw = 2, legend = false, fill = (0, :royalblue), fillalpha = 0.5)
end

# ╔═╡ 664c20be-d41c-11ea-1467-e1c670479af6
begin

	x = rand(Normal(μ, σ²), 100)

	md"""
	Vamos agora gerar uma amostra aleatória de tamanho 100 da distribuição da nossa Normal(μ, σ²) usando o pacote `Random`. O comando está abaixo.

	"""
end

# ╔═╡ acc667f8-d59c-11ea-0805-41b21640540b
begin
	md"""
	Agora vejamos a distribuição LogNormal.

	**Selecione o parâmetro μₗ da LogNormal**:

	$(@bind μₗ Slider(0.0:0.1:25.0))

	**Selecione o parâmetro σ²ₗ da LogNormal**:

	$(@bind σ²ₗ Slider(0.01:0.01:5.0))

	Note que o valor inicial de μₗ é `1.0` e o σ²ₗ é `0.01`.
	"""
end

# ╔═╡ a3dc9324-d59c-11ea-3787-39b6dc720a4b
md"""### Distribuição LogNormal
**Valor de μₗ**: $μₗ

**Valor de σ²ₗ**: $σ²ₗ
"""

# ╔═╡ a651c700-d59c-11ea-029c-1d55b204a41c
#Visualização da distribuição
Plots.plot(LogNormal(μₗ, sqrt(σ²ₗ)), color = :red, lw = 2, legend = false, fill = (0, :red), fillalpha = 0.5)

# ╔═╡ 6ebf2016-d41c-11ea-1a8d-6566e98d1f34
begin
	md"""
	Agora vejamos a distribuição Gamma e o seu formato de acordo com os parâmetros k e θ, onde θ é o parâmetro de escala. Note que nesta parametrização, temos que a função de densidade é dada por

	\begin{equation}
	f(x| k, \theta) = \frac{x^{k-1} e^{-x/θ}}{\Gamma(k) \theta^{k}} \text{ ,}
	\end{equation}

	para $x > 0$, $k > 0$ e $\theta > 0$.

	**Selecione o parâmetro k da Gamma**:

	$(@bind k Slider(1.0:0.1:25.0))

	**Selecione o parâmetro θ da Gamma**:

	$(@bind θ Slider(0.1:0.1:25.0))

	Novamente podemos conferir o resultado. Note que o valor inicial de k é `1.0` e o θ é `0.1`.
	"""
end

# ╔═╡ 6a7b3bf2-d41c-11ea-014e-8b94632190ff
md"""### Distribuição Gamma
**Valor de k**: $k

**Valor de θ**: $θ
"""

# ╔═╡ 6e2923fe-d41c-11ea-03d2-0956b986677c
#Visualização da distribuição
Plots.plot(Gamma(k, θ), color = :red, lw = 2, legend = false, fill = (0, :red), fillalpha = 0.5)

# ╔═╡ 801c622e-d41c-11ea-0441-e300c24904ae
begin
	md"""
	Agora a distribuição mais versátil de todas: a distribuição Beta e o seu formato de acordo com os parâmetros α e β.

	**Selecione o parâmetro α da Beta**:

	$(@bind α Slider(0.1:0.1:10.0))

	**Selecione o parâmetro β da Beta**:

	$(@bind β Slider(0.1:0.1:10.0))

	Os valores iniciais de α e de Β são ambos `0.1`.
	"""
end

# ╔═╡ 6f79fbca-d41c-11ea-0c49-3ffd9d1d5ab2
md"""### Distribuição Beta
**Valor de α**: $α

**Valor de β**: $β
"""

# ╔═╡ 80571608-d41c-11ea-2e72-b7f20192a991
#Visualização da distribuição
Plots.plot(Beta(α, β), color = :red, lw = 2, legend = false, fill = (0, :red), fillalpha = 0.5)

# ╔═╡ 70188128-d41c-11ea-1e8f-5371e257a3bf
begin
	md"""
	Vemos agora a distribuição que ninguém se lembra como é a função de densidade: a distribuição Χ² e o seu formato de acordo com o parâmetro ν de graus de liberdade.

	**Selecione o parâmetro ν da Chi-Quadrado**:

	$(@bind ν Slider(1:50))

	O valor inicial de ν é `1` e só assume valores inteiros positivos.

	"""
end

# ╔═╡ 7fd8c4ce-d41c-11ea-0790-f7afc7a1a84d
#Verificando o valor do parâmetro
md"""
### Distribuição Chi-Quadrado
**Valor de ν**: $ν
"""

# ╔═╡ 7f661988-d41c-11ea-3167-753a433c34e3
#Visualização da distribuição
Plots.plot(Chisq(ν), color = :red, lw = 2, legend = false, fill = (0, :red), fillalpha = 0.5)

# ╔═╡ 9dcde59a-d41c-11ea-36ec-db54e9680172
begin
	md"""
	Como **com certeza** alguém disse na anterior que sabe sim a forma da função de densidade da Χ², então agora sim a distribuição cuja função de densidade que **ninguém** se lembra. Algum desafiante?

	**Selecione o parâmetro d₁ da F-Snedecor**:

	$(@bind d₁ Slider(1:50))

	**Selecione o parâmetro d₂ da F-Snedecor**:

	$(@bind d₂ Slider(1:50))

	Os valores iniciais de d₁ e d₂ são ambos `1`, e só assumem valores inteiros positivos.
	"""
end

# ╔═╡ 9bcc9a70-d41c-11ea-3235-0ffa6ccc0bd5
#Verificando o valor dos parâmetros
md"""
### Distribuição F-Snedecor
**Valor de d₁**: $d₁

**Valor de d₂**: $d₂
"""

# ╔═╡ 9dfe153a-d41c-11ea-2a80-79631819f1ac
#Visualização da distribuição
Plots.plot(FDist(d₁, d₂), color = :red, lw = 2, legend = false, fill = (0, :red), fillalpha = 0.5)

# ╔═╡ Cell order:
# ╠═d5eaa180-d41b-11ea-32f6-6ff398c3e4c2
# ╟─517a652c-d41c-11ea-2118-fdbf45828452
# ╟─57bc858e-d41c-11ea-392b-f1c840845820
# ╟─61be6a20-d41c-11ea-0eee-654e9ec43f39
# ╟─664c20be-d41c-11ea-1467-e1c670479af6
# ╟─a3dc9324-d59c-11ea-3787-39b6dc720a4b
# ╟─a651c700-d59c-11ea-029c-1d55b204a41c
# ╟─acc667f8-d59c-11ea-0805-41b21640540b
# ╟─6a7b3bf2-d41c-11ea-014e-8b94632190ff
# ╟─6e2923fe-d41c-11ea-03d2-0956b986677c
# ╟─6ebf2016-d41c-11ea-1a8d-6566e98d1f34
# ╟─6f79fbca-d41c-11ea-0c49-3ffd9d1d5ab2
# ╟─80571608-d41c-11ea-2e72-b7f20192a991
# ╟─801c622e-d41c-11ea-0441-e300c24904ae
# ╟─7fd8c4ce-d41c-11ea-0790-f7afc7a1a84d
# ╠═7f661988-d41c-11ea-3167-753a433c34e3
# ╟─70188128-d41c-11ea-1e8f-5371e257a3bf
# ╟─9bcc9a70-d41c-11ea-3235-0ffa6ccc0bd5
# ╟─9dfe153a-d41c-11ea-2a80-79631819f1ac
# ╟─9dcde59a-d41c-11ea-36ec-db54e9680172
# ╠═9d996162-d41c-11ea-09b4-994ed99dd014
# ╠═9d3661fc-d41c-11ea-1997-175b6d9a4095


# indicator of a generic box

"""
  IndBox(lb, ub)

Returns the function `g = ind{x : lb ⩽ x ⩽ ub}`. Parameters `lb` and `ub` can be
either scalars or arrays of the same dimension as `x`, and must satisfy `lb <= ub`.
Bounds are allowed to take values `-Inf` and `+Inf`.
"""

immutable IndBox{T <: Union{Real, AbstractArray}, S <: Union{Real, AbstractArray}} <: IndicatorConvex
  lb::T
  ub::S
  function IndBox(lb::T, ub::S)
    if !(eltype(lb) <: Real && eltype(ub) <: Real)
      error("lb and ub must be real")
    end
    if any(lb .> ub)
      error("lb and ub must satisfy lb <= ub")
    else
      new(lb, ub)
    end
  end
end

IndBox{T <: Real}(lb::T, ub::T) = IndBox{T, T}(lb, ub)

IndBox{T <: AbstractArray, S <: Real}(lb::T, ub::S) = IndBox{T, S}(lb, ub)

IndBox{T <: Real, S <: AbstractArray}(lb::T, ub::S) = IndBox{T, S}(lb, ub)

IndBox{T <: AbstractArray, S <: AbstractArray}(lb::T, ub::S) =
  size(lb) != size(ub) ? error("bounds must have the same dimensions, or at least one of them be scalar") :
  IndBox{T, S}(lb, ub)

IndBox_lb{T <: Real, S}(f::IndBox{T, S}, i) = f.lb
IndBox_lb{T <: AbstractArray, S}(f::IndBox{T, S}, i) = f.lb[i]
IndBox_ub{T, S <: Real}(f::IndBox{T, S}, i) = f.ub
IndBox_ub{T, S <: AbstractArray}(f::IndBox{T, S}, i) = f.ub[i]

function (f::IndBox){R <: Real}(x::AbstractArray{R})
  for k in eachindex(x)
    if x[k] < IndBox_lb(f,k) || x[k] > IndBox_ub(f,k)
      return +Inf
    end
  end
  return 0.0
end

function prox!{R <: Real}(f::IndBox, x::AbstractArray{R}, y::AbstractArray{R}, gamma::Real=1.0)
  for k in eachindex(x)
    if x[k] < IndBox_lb(f,k)
      y[k] = IndBox_lb(f,k)
    elseif x[k] > IndBox_ub(f,k)
      y[k] = IndBox_ub(f,k)
    else
      y[k] = x[k]
    end
  end
  return 0.0
end

"""
  IndBallLinf(r::Real=1.0)

Returns the indicator function of an infinity-norm ball, that is function
`g(x) = ind{maximum(abs(x)) ⩽ r}` for `r ⩾ 0`.
"""

IndBallLinf{R <: Real}(r::R=1.0) = IndBox(-r, r)

fun_name(f::IndBox) = "indicator of a box"
fun_dom(f::IndBox) = "AbstractArray{Real}"
fun_expr(f::IndBox) = "x ↦ 0 if all(lb ⩽ x ⩽ ub), +∞ otherwise"
fun_params(f::IndBox) =
  string( "lb = ", typeof(f.lb) <: AbstractArray ? string(typeof(f.lb), " of size ", size(f.lb)) : f.lb, ", ",
          "ub = ", typeof(f.ub) <: AbstractArray ? string(typeof(f.ub), " of size ", size(f.ub)) : f.ub)

function prox_naive{R <: Real}(f::IndBox, x::AbstractArray{R}, gamma::Real=1.0)
  y = min.(f.ub, max.(f.lb, x))
  return y, 0.0
end