@testset "Interpolate" begin
    @testset "LineSegment" begin
        for T in Floats
            P₁ = Point{3,T}(1, 1, 1)
            P₂ = Point{3,T}(3, 3, 3)
            l = LineSegment(P₁, P₂)
            @test l(0) ≈ P₁
            @test l(1) ≈ P₂
            @test l(1//2) ≈ Point{3,T}(2, 2, 2)
        end
    end

    @testset "QuadraticSegment" begin
        for T in Floats
            P₁ = Point{3,T}(0, 0, 0)
            P₂ = Point{3,T}(2, 0, 0)
            P₃ = Point{3,T}(1, 1, 0)
            q = QuadraticSegment(P₁, P₂, P₃)
            for r = LinRange{T}(0, 1, 5)
                @test q(r) ≈ Point{3,T}(2r, -(2r)^2 + 4r, 0)
            end
        end
    end

    @testset "Triangle" begin
        for T in Floats
            P₁ = Point{3,T}(0, 0, 0)
            P₂ = Point{3,T}(0, 1, 0)
            P₃ = Point{3,T}(0, 0, 1)
            tri = Triangle(P₁, P₂, P₃)
            @test tri(0, 0) ≈ P₁
            @test tri(1, 0) ≈ P₂
            @test tri(0, 1) ≈ P₃
            @test tri(1//2, 1//2) ≈ Point{3,T}(0, 1//2, 1//2)
        end
    end

    @testset "Quadrilateral" begin
        for T in Floats
            P₁ = Point{3,T}(0, 0, 0)
            P₂ = Point{3,T}(0, 1, 0)
            P₃ = Point{3,T}(0, 1, 1)
            P₄ = Point{3,T}(0, 0, 1)
            quad = Quadrilateral(P₁, P₂, P₃, P₄)
            @test quad(0, 0) ≈ P₁
            @test quad(1, 0) ≈ P₂
            @test quad(1, 1) ≈ P₃
            @test quad(0, 1) ≈ P₄
            @test quad(1//2, 1//2) ≈ Point{3,T}(0, 1//2, 1//2)
        end
    end

    @testset "QuadraticTriangle" begin
        for T in Floats
            P₁ = Point{3,T}(0, 0, 0)
            P₂ = Point{3,T}(0, 1, 0)
            P₃ = Point{3,T}(0, 1, 1)
            P₄ = Point{3,T}(0, 1//2, 0)
            P₅ = Point{3,T}(0, 1, 1//2)
            P₆ = Point{3,T}(0, 1//2, 1//2)
            tri6 = QuadraticTriangle(P₁, P₂, P₃, P₄, P₅, P₆)
            @test tri6(0, 0) ≈ P₁
            @test tri6(1, 0) ≈ P₂
            @test tri6(0, 1) ≈ P₃
            @test tri6(1//2, 0) ≈ P₄
            @test tri6(1//2, 1//2) ≈ P₅
            @test tri6(0, 1//2) ≈ P₆
        end
    end

    @testset "QuadraticQuadrilateral" begin
        for T in Floats
            P₁ = Point{3,T}(0, 0, 0)
            P₂ = Point{3,T}(0, 1, 0)
            P₃ = Point{3,T}(0, 1, 1)
            P₄ = Point{3,T}(0, 0, 1)
            P₅ = Point{3,T}(0, 1//2,    0)
            P₆ = Point{3,T}(0,    1, 1//2)
            P₇ = Point{3,T}(0, 1//2,    1)
            P₈ = Point{3,T}(0,    0, 1//2)
            quad8 = QuadraticQuadrilateral(P₁, P₂, P₃, P₄, P₅, P₆, P₇, P₈)
            @test quad8(0, 0) ≈ P₁
            @test quad8(1, 0) ≈ P₂
            @test quad8(1, 1) ≈ P₃
            @test quad8(0, 1) ≈ P₄
            @test quad8(1//2,    0) ≈ P₅
            @test quad8(   1, 1//2) ≈ P₆
            @test quad8(1//2,    1) ≈ P₇
            @test quad8(   0, 1//2) ≈ P₈
            @test quad8(1//2, 1//2) ≈ Point{3,T}(0, 1//2, 1//2)
        end
    end

    @testset "Tetrahedron" begin
        for T in Floats
            P₁ = Point{3,T}(0, 0, 0)
            P₂ = Point{3,T}(1, 0, 0)
            P₃ = Point{3,T}(0, 1, 0)
            P₄ = Point{3,T}(0, 0, 1)
            tet = Tetrahedron(P₁, P₂, P₃, P₄)
            @test tet(0, 0, 0) ≈ P₁
            @test tet(1, 0, 0) ≈ P₂
            @test tet(0, 1, 0) ≈ P₃
            @test tet(0, 0, 1) ≈ P₄
        end
    end

    @testset "Hexahedron" begin
        for T in Floats
            P₁ = Point{3,T}(0, 0, 0)
            P₂ = Point{3,T}(1, 0, 0)
            P₃ = Point{3,T}(0, 1, 0)
            P₄ = Point{3,T}(1, 1, 0)
            P₅ = Point{3,T}(0, 0, 1)
            P₆ = Point{3,T}(1, 0, 1)
            P₇ = Point{3,T}(0, 1, 1)
            P₈ = Point{3,T}(1, 1, 1)
            hex = Hexahedron(P₁, P₂, P₃, P₄, P₅, P₆, P₇, P₈)
            @test hex(0, 0, 0) ≈ P₁
            @test hex(1, 0, 0) ≈ P₂
            @test hex(1, 1, 0) ≈ P₃
            @test hex(0, 1, 0) ≈ P₄
            @test hex(0, 0, 1) ≈ P₅
            @test hex(1, 0, 1) ≈ P₆
            @test hex(1, 1, 1) ≈ P₇
            @test hex(0, 1, 1) ≈ P₈
        end
    end

    @testset "QuadraticTetrahedron" begin
        for T in Floats
            P₁  = Point{3,T}(0, 0, 0)
            P₂  = Point{3,T}(1, 0, 0)
            P₃  = Point{3,T}(0, 1, 0)
            P₄  = Point{3,T}(0, 0, 1)
            P₅  = Point{3,T}(1//2,    0,    0)
            P₆  = Point{3,T}(1//2, 1//2,    0)
            P₇  = Point{3,T}(0,    1//2,    0)
            P₈  = Point{3,T}(0,       0, 1//2)
            P₉  = Point{3,T}(1//2,    0, 1//2)
            P₁₀ = Point{3,T}(0,    1//2, 1//2)
            tet10 = QuadraticTetrahedron(P₁, P₂, P₃, P₄, P₅, P₆, P₇, P₈, P₉, P₁₀)
            @test tet10(0, 0, 0) ≈ P₁
            @test tet10(1, 0, 0) ≈ P₂
            @test tet10(0, 1, 0) ≈ P₃
            @test tet10(0, 0, 1) ≈ P₄
            @test tet10(1//2,    0,    0) ≈ P₅
            @test tet10(1//2, 1//2,    0) ≈ P₆
            @test tet10(0,    1//2,    0) ≈ P₇
            @test tet10(0,       0, 1//2) ≈ P₈
            @test tet10(1//2,    0, 1//2) ≈ P₉
            @test tet10(0,    1//2, 1//2) ≈ P₁₀
        end
    end

    @testset "QuadraticHexahedon" begin
        for T in Floats
            P₁  = Point{3,T}(0, 0, 0)
            P₂  = Point{3,T}(1, 0, 0)
            P₃  = Point{3,T}(1, 1, 0)
            P₄  = Point{3,T}(0, 1, 0)
            P₅  = Point{3,T}(0, 0, 1)
            P₆  = Point{3,T}(1, 0, 1)
            P₇  = Point{3,T}(1, 1, 1)
            P₈  = Point{3,T}(0, 1, 1)
            P₉  = Point{3,T}(1//2,    0,    0)
            P₁₀ = Point{3,T}(   1, 1//2,    0)
            P₁₁ = Point{3,T}(1//2,    1,    0)
            P₁₂ = Point{3,T}(   0, 1//2,    0)
            P₁₃ = Point{3,T}(1//2,    0,    1)
            P₁₄ = Point{3,T}(   1, 1//2,    1)
            P₁₅ = Point{3,T}(1//2,    1,    1)
            P₁₆ = Point{3,T}(   0, 1//2,    1)
            P₁₇ = Point{3,T}(0, 0, 1//2)
            P₁₈ = Point{3,T}(1, 0, 1//2)
            P₁₉ = Point{3,T}(1, 1, 1//2)
            P₂₀ = Point{3,T}(0, 0, 1//2)
            hex20 = QuadraticHexahedron(P₁,  P₂,  P₃,  P₄,  P₅,  P₆,  P₇,
                                        P₈,  P₉,  P₁₀, P₁₁, P₁₂, P₁₃, P₁₄,
                                        P₁₅, P₁₆, P₁₇, P₁₈, P₁₉, P₂₀)
            @test hex20(0, 0, 0) ≈ P₁
            @test hex20(1, 0, 0) ≈ P₂
            @test hex20(1, 1, 0) ≈ P₃
            @test hex20(0, 1, 0) ≈ P₄
            @test hex20(0, 0, 1) ≈ P₅
            @test hex20(1, 0, 1) ≈ P₆
            @test hex20(1, 1, 1) ≈ P₇
            @test hex20(0, 1, 1) ≈ P₈
            @test hex20(1//2,    0,    0)≈ P₉
            @test hex20(   1, 1//2,    0)≈ P₁₀
            @test hex20(1//2,    1,    0)≈ P₁₁
            @test hex20(   0, 1//2,    0)≈ P₁₂
            @test hex20(1//2,    0,    1)≈ P₁₃
            @test hex20(   1, 1//2,    1)≈ P₁₄
            @test hex20(1//2,    1,    1)≈ P₁₅
            @test hex20(   0, 1//2,    1)≈ P₁₆
            @test hex20(0, 0, 1//2)≈ P₁₇
            @test hex20(1, 0, 1//2)≈ P₁₈
            @test hex20(1, 1, 1//2)≈ P₁₉
            @test hex20(0, 0, 1//2)≈ P₂₀
        end
    end
end


function niche_network(S::Int, C::Float64)

    cond = false
    adj = fill(0, S, S)
    n = fill(0.0, S)
    r = fill(0.0, S)
    c = fill(0.0, S)

    while !cond
        n = sort(rand(S))
        β = 1 / (2 * C) - 1
        r = n .* rand(Beta(1, β), S)

        # set r[position of min[η]] = 0, so that every web has at least one basal species
        r[1] = 0.0

        for pred in 1:S

            c[pred] = rand(Uniform(r[pred]/2, n[pred]))

            for prey in 1:S
                if ((c[pred] - r[pred] / 2) <= n[prey] <= (c[pred] + r[pred] / 2))
                    adj[prey, pred] = 1
                else
                    adj[prey, pred] = 0
                end
            end
        end

        cond = is_connected(DiGraph(adj))
    end

    return (fw = adj, n = n, r = r, c = c)
end

function prob_network(S, C)

    cond = false
    adj = fill(0, S, S)

    n = fill(0.0, S)
    r = fill(0.0, S)
    c = fill(0.0, S)

    while !cond

        n = sort(rand(S))
        β = 1 / (2 * C) - 1
        r = n .* rand(Beta(1, β), S)

        # set r[position of min[η]] = 0, so that every web has at least one basal species
        r[1] = 0.0

        @inbounds for pred in 1:S

            c[pred] = rand(Uniform(r[pred] / 2, n[pred]))

            for prey in 1:S
                ϕ = exp(-((n[prey] - c[pred]) / (r[pred] / 2)) ^ 2)
                adj[prey, pred] = rand(Bernoulli(ϕ), 1)[1]
            end
        end

        cond = is_connected(DiGraph(adj))
    end

    return (fw = adj, n = n, r = r, c = c)
end

function ϕ_mat(θ, m)

    ϕ = fill(0.0, m.S, m.S)

    @inbounds for pred in 1:m.nonbasal
        for prey in 1:m.S
            ϕ[prey, m.basal + pred] = 0.99 * exp(-((θ.n[prey] - θ.c[pred] * θ.n[m.basal + pred]) / (θ.r[pred] / 2)) ^ 2)
        end
    end

    return ϕ

end


#=
A Pythagorean triplet is a set of three natural numbers, a < b < c, for which,

a2 + b2 = c2
For example, 32 + 42 = 9 + 16 = 25 = 52.

There exists exactly one Pythagorean triplet for which a + b + c = 1000.
Find the product abc.
=#

# Kinda ugly brute force approach
function findABC()
  max = 1000
  for a in 1:max, b in 1:max, c in 1:max
    if a >= b || b >= c || a + b + c != max
      continue
    end
    if a^2 + b^2 == c^2
      return a*b*c
    end
  end
end

@time println(findABC())


"""
    lognormal(x, μ, σ)

```math
\\frac{(x-μ)^2}{2σ^2}
```
"""
function lognormal(x, μ, σ)
  if σ == Inf
    return 0
  end
  if σ ≈ 0
    return Inf
  end
  (x-μ)^2/2/σ^2
end

"""
    logpoisson(n, x...)

```math
  λ > 0 \\
  λ = sum([x...])
  λ - n*log(λ) + n*log(n) - n
```math
"""
function logpoisson(n, x...)
  λ = sum([x...])
  if λ <= 0
    return Inf
  end
  λ - n*log(λ) + n*log(n) - n
end

"""
    batlog(x)

Compute the natural log, but create bounds to ignore very large
numbers / infinities.
```jldoctest
julia> batlog.([12.0, 1e6, 0])
3-element Array{Float64, 1}:
  2.4849
  13.8155
  -100000.0
```
"""
function batlog(x::Float64; up::Float64=1000.0, down::Float64=-1000.0)
  if x <= 0
    return -1e8
  end
  y = log(x)
  if y < down
    @debug "Batlog hit floor" y
  end
  if y > up
    @debug "Batlog hit ceiling" y
  end
  minimum([up, maximum([down, y])])
end


#Array{Float64,1}
a=[168.3,179.2,165.8,170.5,188.2,174.6,162.8,175.5,178.1,177.1]

#get arithmetic mean
@show mean(a)


# Utility functions for common algebraic operations.

"""
    rotateTensorToVectorZ(tensor, vector)

Rotates `tensor` with a rotation matrix that aligns z-axis with `vector`.
"""
function rotateTensorToVectorZ(tensor::AbstractMatrix{T}, v::AbstractVector{T}) where T
   unitz = SVector{3, T}(0.0, 0.0, 1.0)
   vz = v × unitz::SVector{3, T}
   vz ./= hypot(vz[1], vz[2], vz[3])
   angle = acos(v ⋅ unitz / hypot(v[1], v[2], v[3]))
   R = getRotationMatrix(vz, angle)
   # Rotate tensor
   R * tensor * R'
end

"""
    getRotationMatrix(vector, angle)

Creates a rotation matrix that rotates around a unit `vector` by an `angle` in radians.
References: https://en.wikipedia.org/wiki/Rodrigues'_rotation_formula
https://en.wikipedia.org/wiki/Rotation_matrix#Rotation_matrix_from_axis_and_angle
"""
function getRotationMatrix(v::AbstractVector, θ)
   sinθ, cosθ = sincos(eltype(v)(θ))
   tmp = 1 - cosθ
   m =  @SMatrix [
        cosθ+v[1]^2*tmp         v[1]*v[2]*tmp-v[3]*sinθ v[1]*v[3]*tmp+v[2]*sinθ;
        v[1]*v[2]*tmp+v[3]*sinθ cosθ+v[2]^2*tmp         v[2]*v[3]*tmp-v[1]*sinθ;
        v[1]*v[3]*tmp-v[2]*sinθ v[3]*v[2]*tmp+v[1]*sinθ cosθ+v[3]^2*tmp]
end

"""
    getRotationB(B) -> SMatrix

Obtain a rotation matrix with each column being a unit vector which is parallel (`v3`) and
perpendicular (`v1,v2`) to the magnetic field `B`. The two perpendicular directions are
chosen based on the reference vector of z-axis in the Cartesian coordinates.
"""
function getRotationB(B::AbstractVector{T}) where T
   b = hypot(B[1], B[2], B[3])
   if B[3] / b ≈ -1.0 # B aligned with reference vector
      R = SMatrix{3,3,T}([0.0 -1.0 0.0; 1.0 0.0 0.0; 0.0 0.0 1.0])
   elseif B[3] / b ≈ 1.0 # B aligned with reference vector
      R = SMatrix{3,3,T}([0.0 -1.0 0.0; 1.0 0.0 0.0; 0.0 0.0 1.0])
   else
      v0 = SVector{3,T}(0.0, 0.0, 1.0) # reference vector
      v3 = SVector{3,T}(B[1]/b, B[2]/b, B[3]/b) # unit vector along B
      v1 = v0 × v3::SVector{3, T}
      v2 = v3 × v1::SVector{3, T}

      R = @SMatrix [v1[1] v2[1] v3[1]; v1[2] v2[2] v3[2]; v1[3] v2[3] v3[3]]
   end
   R
end

"""
    rotateWithB(T, B) -> Matrix

Rotate the tensor `T` with the 3rd direction aligned with `B`.
See also: [`rotateWithB!`](@ref)
"""
function rotateWithB(T::AbstractMatrix, B::AbstractVector)
   R = getRotationB(B)
   R * T * R'
end

"""
    rotateWithB!(T, B)

Rotate the tensor `T` with the 3rd direction aligned with `B`.
See also: [`rotateWithB`](@ref)
"""
function rotateWithB!(T::AbstractMatrix, B::AbstractVector)
   R = getRotationB(B)
   T[:] = R * T * R'
end


using AmbientForcing
using OrdinaryDiffEq, Distributions

# Adopting the Robertson Example from the DifferentialEquation Docs
# https://diffeq.sciml.ai/stable/tutorials/advanced_ode_example/#Handling-Mass-Matrices
function rober(du,u,p,t)
  y₁,y₂,y₃ = u
  du[1] = -0.04 * y₁ + 1e4 * y₂ * y₃
  du[2] =  0.04 * y₁ - 1e4 * y₂ * y₃ - 3e7 * y₂^2
  du[3] =  y₁ + y₂ + y₃ - 1
  nothing
end

# Creating the mass matrix M
# The last row depicts our constraint
M = [1. 0  0
     0  1. 0
     0  0  0];

# Setting up the DAE as an ODE in mass matrix form
ode_rober = ODEFunction(rober, mass_matrix = M)

# Choosing an inital condition which fulfills the constraint
u0 = [1.0,0.0,0.0];

# Acsessing the constraint equations g
g_rober = constraint_equations(ode_rober)
g_rober(u0)

# Randomly perturbing all variables
Frand = random_force(ode_rober, [0.0,1], Uniform)
z_new = ambient_forcing(ode_rober, u0, 2.0, Frand)

# g ≈ 0 means the constraint is fulfilled
g_rober(z_new)

# Only perturbing the second variable y₂
h = [0, 1 ,0]
z_new = ambient_forcing(ode_rober, u0, 2.0, h)

g_rober(z_new)


using Printf

@testset "Hermite 1D" begin

    for nx in 5:5:35

        ϵ, γ = 0.1, 3.0
        xmin, xmax = -1, 1

        interp = Hermite(Chebyshev(xmin, xmax, nx), 0.1, 3.0)

        xk = interp.nodes
        fk = cos.(xk.^2)

	nxe = 128
        xe  = collect(range(xmin, stop=xmax, length=nxe))
        dx  = xe[2]-xe[1]
        fe  = cos.(xe.^2)

        s   = interpolate( interp, fk, xe )

        l2_error = sqrt(trapz((s - fe).^2, dx))
        l1_error = maximum(abs.(s .- fe))
	@printf( "%02d points - L2 : %.3e - L∞ : %.3e \n",
	        nx, l2_error, l1_error )

	@test l2_error < max(1.0e-14, 10.0^(-nx÷2+1))
    end

end

@testset "Radial 1D" begin

    for nx in 5:5:35

        ϵ, γ = 0.1, 3.0
        xmin, xmax = -1, 1

        interp = Radial(Chebyshev(xmin, xmax, nx), 0.1)

        xk = interp.nodes
        fk = cos.(xk.^2)

	nxe = 128
        xe  = collect(range(xmin, stop=xmax, length=nxe))
        dx  = xe[2]-xe[1]
        fe  = cos.(xe.^2)

        @time s   = interpolate( interp, fk, xe )

        l2_error = sqrt(trapz((s - fe).^2, dx))
        l1_error = maximum(abs.(s .- fe))
	@printf( "%02d points - L2 : %.3e - L∞ : %.3e \n",
	        nx, l2_error, l1_error )

	@test true
    end

end


@doc raw"""
The linear trend model is defined by:
```math
\begin{gather*}
    \begin{aligned}
        y_{t} &=  \mu_{t} + \gamma_{t} + \varepsilon_{t} \quad &\varepsilon_{t} \sim \mathcal{N}(0, \sigma^2_{\varepsilon})\\
        \mu_{t+1} &= \mu_{t} + \nu_{t} + \xi_{t} \quad &\xi_{t} \sim \mathcal{N}(0, \sigma^2_{\xi})\\
        \nu_{t+1} &= \nu_{t} + \zeta_{t} \quad &\zeta_{t} \sim \mathcal{N}(0, \sigma^2_{\zeta})\\
    \end{aligned}
\end{gather*}
```

# Example
```jldoctest
julia> model = LocalLinearTrend(rand(100))
LocalLinearTrend model
```

See more on [Finland road traffic fatalities](@ref)

# References
 * Durbin, James, & Siem Jan Koopman. (2012). "Time Series Analysis by State Space Methods:
    Second Edition." Oxford University Press. pp. 44
"""
mutable struct LocalLinearTrend <: StateSpaceModel
    hyperparameters::HyperParameters
    system::LinearUnivariateTimeInvariant
    results::Results

    function LocalLinearTrend(y::Vector{Fl}) where Fl
        Z = Fl.([1.0; 0.0])
        T = Fl.([1 1; 0 1])
        R = Fl.([1 0; 0 1])
        d = zero(Fl)
        c = zeros(Fl, 2)
        H = one(Fl)
        Q = Fl.([1 0; 0 1])

        system = LinearUnivariateTimeInvariant{Fl}(y, Z, T, R, d, c, H, Q)

        names = ["sigma2_ε", "sigma2_ξ", "sigma2_ζ"]
        hyperparameters = HyperParameters{Fl}(names)

        return new(hyperparameters, system, Results{Fl}())
    end
end

function default_filter(model::LocalLinearTrend)
    Fl = typeof_model_elements(model)
    steadystate_tol = Fl(1e-5)
    a1 = zeros(Fl, 2)
    P1 = Fl(1e6) .* Matrix{Fl}(I, 2, 2)
    return UnivariateKalmanFilter(a1, P1, 2, steadystate_tol)
end

function initial_hyperparameters!(model::LocalLinearTrend)
    Fl = typeof_model_elements(model)
    initial_hyperparameters = Dict{String,Fl}(
        "sigma2_ε" => var(model.system.y),
        "sigma2_ξ" => var(model.system.y),
        "sigma2_ζ" => one(Fl),
    )
    set_initial_hyperparameters!(model, initial_hyperparameters)
    return model
end

function constrain_hyperparameters!(model::LocalLinearTrend)
    constrain_variance!(model, "sigma2_ε")
    constrain_variance!(model, "sigma2_ξ")
    constrain_variance!(model, "sigma2_ζ")
    return model
end

function unconstrain_hyperparameters!(model::LocalLinearTrend)
    unconstrain_variance!(model, "sigma2_ε")
    unconstrain_variance!(model, "sigma2_ξ")
    unconstrain_variance!(model, "sigma2_ζ")
    return model
end

function fill_model_system!(model::LocalLinearTrend)
    model.system.H = get_constrained_value(model, "sigma2_ε")
    model.system.Q[1] = get_constrained_value(model, "sigma2_ξ")
    model.system.Q[end] = get_constrained_value(model, "sigma2_ζ")
    return model
end

function reinstantiate(::LocalLinearTrend, y::Vector{Fl}) where Fl
    return LocalLinearTrend(y)
end

has_exogenous(::LocalLinearTrend) = false


#=
Aodwt.jl
2019-02-23 Jeff Fessler, University of Michigan
=#

export Aodwt

using LinearMapsAA: LinearMapAA, LinearMapAM, LinearMapAO
using Wavelets: dwt!, idwt!, wavelet, WT

"""
    A, levels, mfun = Aodwt(dims ; level::Int=3, wt=wavelet(WT.haar))

Create orthogonal discrete wavelet transform (ODWT) `LinearMapAA`

in
- `dims::Dims` tuple of dimensions

option
- `level::Int` # of levels; default 3
- `wt` wavelet transform type (see `Wavelets` package); default Haar
- `operator::Bool=true` default to `LinearMapAO`

out
- `A` a `LinearMapAX` object
- `scales` array of size `dims` showing the scale of each coefficient
which is useful when imposing scale-dependent regularization
- `mfun` convenience function for A*X when X is a Matrix or Array (not vector)

2019-02-23 Jeff Fessler, University of Michigan
"""
function Aodwt(
	dims::Dims ;
	T::DataType = Float32,
	level::Int = 3,
	wt = wavelet(WT.haar),
	operator::Bool = true, # !
)

	function mfunA(lev)
		# todo: avoiding convert involves `dwt!` limitation; see:
		# https://github.com/JuliaDSP/Wavelets.jl/issues/53
		# https://github.com/JuliaDSP/Wavelets.jl/pull/54
	#	forw!(y,x) = dwt!(y, convert(AbstractArray{T},x), wt, level)
	#	back!(x,y) = idwt!(x, convert(AbstractArray{T},y), wt, level)
		forw!(y,x) = dwt!(y, x, wt, lev)
		back!(x,y) = idwt!(x, y, wt, lev)

		if operator
			mfun = (A, x) -> A * x
			return mfun, LinearMapAA(forw!, back!,
				(prod(dims), prod(dims)) ;
				prop = (wt=wt, level=lev), T=T,
				operator = true, idim=dims, odim=dims,
			)
		else
			mfun = (A, x) -> reshape(A * vec(x), dims)
			return mfun, LinearMapAA(
				(y,x) -> vec(forw!(reshape(y, dims), reshape(x, dims))),
				(x,y) -> vec(back!(reshape(x, dims), reshape(y, dims))),
			#	x -> vec(dwt(reshape(x, dims), wt, level)),
			#	y -> vec(idwt(reshape(y, dims), wt, level)),
				(prod(dims), prod(dims)) ;
				prop = (wt=wt, level=lev), T=T,
			)
		end
	end

	mfun, A = mfunA(level)

	scales = zeros(dims)
	for il=1:level
		_,Al = mfunA(il)
		tmp = mfun(Al, ones(dims)) .== 0
		scales += il * (tmp .& (scales .== 0))
	end

	return A, scales, mfun
end


# # clip-02-05.jl

# Load Julia packages (libraries) needed  for the snippets in chapter 0

using StatisticalRethinking, Optim
#gr(size=(600,600));

# ### snippet 3.2

p_grid = range(0, step=0.001, stop=1)
prior = ones(length(p_grid))
likelihood = [pdf(Binomial(9, p), 6) for p in p_grid]
posterior = likelihood .* prior
posterior = posterior / sum(posterior)
samples = sample(p_grid, Weights(posterior), length(p_grid));
samples[1:5]

# ### snippet 3.3
# Draw 10000 samples from this posterior distribution

N = 10000
samples = sample(p_grid, Weights(posterior), N);

# In StatisticalRethinkingJulia samples will always be stored
# in an MCMCChains.Chains object.

chn = MCMCChains.Chains(reshape(samples, N, 1, 1), ["toss"]);

# Describe the chain

MCMCChains.describe(chn)

# Plot the chain

plot(chn)

# ### snippet 3.4

# Create a vector to hold the plots so we can later combine them

p = Vector{Plots.Plot{Plots.GRBackend}}(undef, 2)
p[1] = scatter(1:N, samples, markersize = 2, ylim=(0.0, 1.3), lab="Draws")

# ### snippet 3.5

# Analytical calculation

w = 6
n = 9
x = 0:0.01:1
p[2] = density(samples, ylim=(0.0, 5.0), lab="Sample density")
p[2] = plot!( x, pdf.(Beta( w+1 , n-w+1 ) , x ), lab="Conjugate solution")

# Add quadratic approximation

plot(p..., layout=(1, 2))

# End of `03/clip-02-05.jl`


using Sobol
using Base.Test

# test integrand from Joe and Kuo paper (integrates to 1)
function testfunc(x::Vector{Float64})
    f = 1.0
    for j = 1:length(x)
        cj = j^0.3333333333333333333
        f *= (abs(4*x[j] - 2) + cj) / (1 + cj)
    end
    return f
end

function testint(N, n)
    s = SobolSeq(N)
    x = Array(Float64, N)
    skip(s, n)
    sum = 0.0
    for j = 1:n
        sum += testfunc(next!(s, x))
    end
    sum / n
end

# test results from Joe and Kuo (2003)
N = [50, 100, 200, 300, 400, 500, 600, 700, 800, 900, 1000, 1111]
n = [1009, 1997, 4001, 8009, 16001, 32003, 64007, 128021]
JoeKuo = [0.9743 0.9544 0.9838 0.9982 0.9995 0.9976 0.9969 0.9975;
          1.0157 0.9288 0.9649 0.9866 0.9985 0.9817 0.9955 0.9995;
          1.1109 0.8877 0.9070 1.0312 1.0216 0.9772 0.9744 0.9925;
          0.9770 0.8715 0.9305 0.9546 1.0060 0.9712 0.9770 0.9843;
          1.1780 0.8908 0.9710 0.9527 1.0203 0.9754 0.9779 0.9725;
          1.0567 0.8548 0.9700 0.9704 0.9768 0.9634 0.9745 0.9670;
          0.8735 0.7674 1.0130 0.9874 0.9677 0.9416 0.9580 0.9564;
          0.8752 0.8241 1.0067 1.0078 0.9987 0.9236 0.9517 0.9475;
          0.9139 0.8822 1.0926 1.0560 1.0210 0.9454 0.9330 0.9359;
          0.7779 1.0225 1.0593 1.0991 1.0470 0.9115 0.9688 0.9615;
          0.8348 1.0443 1.0710 1.0862 1.0483 0.9188 0.9735 0.9683;
          0.7524 0.9139 0.9592 0.9703 1.0116 0.9124 1.0411 0.9904]

for i = 1:length(N)
    println("testing integrals of dimension $(N[i])")
    for j = 1:length(n)
        println("  ... with $(n[j]) samples")
        t = testint(N[i], n[j])
        @test_approx_eq_eps t JoeKuo[i,j] 1e-4
    end
end


"""
2D Predator-Prey Model

Taken from:
[14.3 pp2: A 2D Predator-Prey Model][auto07p/14.3]
of AUTO-07p manual.

[auto07p/14.3]: http://www.dam.brown.edu/people/sandsted/auto/auto07p.pdf#page=138
"""
module PredatorPrey

using DiffEqBase: ODEProblem
using StaticArrays: SVector
using Setfield: @lens
import Setfield

using ...Bifurcations: BifurcationProblem
using ...Codim2: cast_container, DiffEqCodim2Problem

f1(u, p) =
    p[2] * u[1] * (1 - u[1]) - u[1] * u[2] - p[1] * (1 - exp(- p[3] * u[1]))
f2(u, p) =
    - u[2] + p[4] * u[1] * u[2]

f(u::SVector, p, t) = SVector(f1(u, p), f2(u, p))

function f(du, u, p, t)
    du[1] = f1(u, p)
    du[2] = f2(u, p)
    nothing
end

const _f = f

function make_prob(
        p = (0.0, 3, 5, 3);
        f = _f,
        u0 = SVector(0.0, 0.0),
        tspan = (0.0, 30.0),
        ode = ODEProblem{!(u0 isa SVector)}(f, u0, tspan, p),
        param_axis = (@lens _[1]),
        t_domain = (0.0, 1.0),
        kwargs...)
    x_min = SVector(-0.2, -Inf)
    x_max = SVector(1.0, Inf)
    x_min = cast_container(typeof(ode.u0), x_min)
    x_max = cast_container(typeof(ode.u0), x_max)
    return BifurcationProblem(
        ode, param_axis, t_domain;
        phase_space = (x_min, x_max),
        kwargs...)
end

prob = make_prob()
ode = prob.p.de_prob
p = ode.p
u0 = ode.u0
param_axis = prob.p.param_axis

function make_codim2_prob(
        p = ode.p;
        u0 = SVector(0.0, 0.0),
        tspan = (0.0, 30.0),
        ode = ODEProblem{!(u0 isa SVector)}(f, u0, tspan, p),
        param_axis1 = (@lens _[1]),
        param_axis2 = (@lens _[3]),
        t_domain = ([-Inf, -Inf], [Inf, Inf]),
        kwargs...)
    return DiffEqCodim2Problem(
        ode,
        param_axis1,
        param_axis2,
        t_domain;
        kwargs...)
end

end  # module


using Plots
using Random
using Pkg; Pkg.add("LaTeXStrings")
using LaTeXStrings

# Set working directory
cd("/Users/elvis/Documents/DSI/courses/content/stochastic-approximations/julia-code/")

# Set seed
Random.seed!(123)

# Generate MC points
n = 4000
points_x = 3 .* rand(n)
points_y = -3 .+ 7 .* rand(n)

# Define functions
f1(x) = 3*abs(cos(x)) .+ 2*sin(x)
f2(x) = -3*abs(cos(x)) .+ 2*sin(x)

# Label points in circle
labels = [f2(x) <= y <= f1(x) for (x,y) in zip(points_x, points_y)]

# Create curves
x_vals = 0:.01:3
f1_x = 3*abs.(cos.(x_vals)) .+ 2*sin.(x_vals)
f2_x = -3*abs.(cos.(x_vals)) .+ 2*sin.(x_vals)

# Create plot of MC points
plot(points_x, points_y, seriestype = :scatter,
     group = .!labels,
     legend = false,
     ms = 3,
     bg = RGB(247/255, 236/255, 226/255))

# Add curves to plot
plot!(x_vals, f1_x, seriestype= :line, label = L"f_1(x)", lc = :blue,
    legend = false)
plot!(x_vals, f2_x, seriestype= :line, label = L"f_2(x)", lc = :violet)
plot!(x_vals, f1_x, fillrange = [f2_x, f1_x], fillcolor = :blue, alpha = .12)

# Area estimate
area = 21*sum(labels)/n

savefig("../images/area-ex1-monte-carlo.svg")


using LinearAlgebra

@testset "core_nonmutating" begin

    # Basic Linear Algebra: cross3, dot3, norm3, norm3!
    a = [1. 2. 9.; 2. 5. 6]
    b = [4. 5. 7.; 3. 2. 8.]
    xab = [-31. 29. -3.; 28. 2. -11.]
    dab = [77., 64.]
    na = [sqrt(86), sqrt(65)]
    @test xab ≈ DER.cross3(a,b)
    @test dab ≈ DER.dot3(a,b)
    @test na ≈ DER.norm3(a)

    # Edges and Tangents
    ea = [1. 3. -3.]
    ta = ea./sqrt(19)
    @test ea ≈ DER.edges(a)
    @test ta ≈ DER.tangents(a)

    # Parallel transport
    θ = [1.0, 0.01, 0.0]
    ptic_θ = [(1 - cos(1.0))/sin(1.0)^2, (1 - cos(0.01))/sin(0.01)^2, 0.5]

    s2 = sin.(θ).^2
    c = cos.(θ)
    @test ptic_θ ≈ DER.ptransport_inner_coefficient.(s2, c)

    t1 = [1. 0. 0.]
    t2 = [1. 1. 1.]/sqrt(3)
    v1 = [0. 1. 1.]
    v2 = [-2. 1. 1.]/sqrt(3)
    @test DER.ptransport(t1, t1, t2) ≈ t2
    @test DER.ptransport(DER.cross3(t1, t2), t1, t2) ≈ DER.cross3(t1, t2)
    @test DER.ptransport([v1; v1], [t1; t1], [t2; t2]) ≈ [v2; v2]

    # Rotations of orthogonal unit vectors
    v1 = [0. 1. 0.; 0. 0. 1.]
    v2 = [0. 1. 1.; 0. -1. 1.]/sqrt(2)
    θ = π/4
    t = [1. 0. 0.]
    @test DER.rotate_orthogonal_unit(v1, [t; t], [θ; θ]) ≈ v2

    # Basic Rod update
    x0 = [0. 0. 0.; 1. 0. 0.; 2. 0. 0.]
    d0 = [0. 0. 1.; 0. 0. 1.]
    Δx = [0. 0. 0.; 0. 0. 1.; 0. 0. 0.]
    Δθ = [0., π/2]
    x, d = DER.rod_update(x0, d0, Δx, Δθ)
    @test x ≈ x0 + Δx
    @test d ≈ [-1/sqrt(2) 0. 1/sqrt(2); 0. -1. 0.]

    # Full kinematics
    l, κ, τ = DER.full_kinematics(x, d)
    @test l ≈ [sqrt(2), sqrt(2)]
    @test κ ≈ [-1. 1.]
    @test τ ≈ [π/2]
end

@testset "core_mutating" begin

    # Basic Linear Algebra: cross3, dot3, norm3, norm3!
    a = [1. 2. 9.; 2. 5. 6]
    b = [4. 5. 7.; 3. 2. 8.]
    xab = [-31. 29. -3.; 28. 2. -11.]
    dab = [77., 64.]
    na = [sqrt(86), sqrt(65)]
    ta = a./na

    c = similar(a)
    d = zeros(2)
    cache = similar(a)

    DER.cross3!(c,a,b)
    @test xab ≈ c

    DER.dot3!(d,a,b)
    @test dab ≈ d

    DER.norm3!(d,a)
    @test na ≈ d

    c = copy(a)
    DER.norm3!(c)
    @test ta ≈ c

    # Edges and tangents
    eb = [-1.0 -3.0 1.0]
    c = similar(eb)
    DER.edges!(c,b)
    @test eb ≈ c

    tb = eb./sqrt(11)
    c = similar(tb)
    DER.tangents!(c,b)
    @test tb ≈ c

    # Parallel transport inner coeff: TBD!
    θ = [1.0, 0.01, 0.0]
    ptic_θ = [(1 - cos(1.0))/sin(1.0)^2, (1 - cos(0.01))/sin(0.01)^2, 0.5]

    s2 = sin.(θ).^2
    c = cos.(θ)
    out = similar(θ)
    DER.ptransport_inner_coefficient!(out, s2, c)
    @test ptic_θ ≈ out

    # Parallel transport
    t1 = [1. 0. 0.]
    t2 = [1. 1. 1.]/sqrt(3)
    v1 = [0. 1. 1.]
    v2 = [-2. 1. 1.]/sqrt(3)
    cache = zeros(1, 9)

    val = similar(t1)
    DER.ptransport!(val, t1, t1, t2, cache)
    @test val ≈ t2

    DER.ptransport!(val, DER.cross3(t1, t2), t1, t2, cache)
    @test val ≈ DER.cross3(t1, t2)

    cache = zeros(2, 9)
    val = [v1; v1]
    DER.ptransport!(val, [v1; v1], [t1; t1], [t2; t2], cache)
    @test val ≈ [v2; v2]

    # Rotations of orthogonal unit vectors
    v1 = [0. 1. 0.; 0. 0. 1.]
    v2 = [0. 1. 1.; 0. -1. 1.]/sqrt(2)
    θ = π/4
    t = [1. 0. 0.]
    cache = zeros(2, 5)
    DER.rotate_orthogonal_unit!(v1, [t; t], [θ; θ], cache)
    @test  v1 ≈ v2

    # Basic Rod update
    x1 = [0. 0. 0.; 1. 0. 0.; 2. 0. 0.]
    d1 = [0. 0. 1.; 0. 0. 1.]
    Δx = [0. 0. 0.; 0. 0. 1.; 0. 0. 0.]
    Δθ = [0., π/2]
    x2 = x1 + Δx
    d2 = [-1/sqrt(2) 0. 1/sqrt(2); 0. -1. 0.]

    xout = similar(x1)
    dout = similar(d1)
    cache = DER.allocate_cache(DER.rod_update!, Float64, 1)
    DER.rod_update!(xout, dout, x1, d1, Δx, Δθ, cache)
    @test xout ≈ x2
    @test dout ≈ d2

    # Full Kinematics
    l = zeros(2)
    κ = zeros(1,2)
    τ = zeros(1)
    caches = DER.allocate_cache(DER.full_kinematics!, Float64, 1)
    DER.full_kinematics!(l, κ, τ, x2, d2, caches)
    @test l ≈ [sqrt(2), sqrt(2)]
    @test κ ≈ [-1. 1.]
    @test τ ≈ [π/2]

end

# Next: Elastic and Gravitational energy


using Turing, ReverseDiff, Memoization
using DifferentialEquations
using Plots
using Statistics
using StatsPlots
Turing.setadbackend(:reversediff)
Turing.setrdcache(true)
using Distributions
using LinearAlgebra

# here we are simulating the signal of two coupled oscillators
# with amplitudes A1 and A2 of kBT/k and kBT/(k+c) where c is the coupling constant
# D = kBT/gamma
# in terms of the standard parameters of a OU process:
# sigma = sqrt(2*D)
# Theta = D/A

function corr_osc(c)
    # function that returns two time series that are correlated through a coupling coefficient c
    μ = 0.0 # mean is zero
    σ = sqrt(2) # D=1
    Θ1 = 1.0
    Θ2 = 1.0+abs(c)

    W1 = OrnsteinUhlenbeckProcess(Θ1,μ,σ,0.0,1.0)
    W2 = OrnsteinUhlenbeckProcess(Θ2,μ,σ,0.0,1.0)
    prob1 = NoiseProblem(W1,(0.0,100.0))
    prob2 = NoiseProblem(W2,(0.0,100.0))
    sol1 = solve(prob1;dt=0.1)
    sol2 = solve(prob2;dt=0.1)

    # creating the two correlated
    x1 = (sol1.u .+ sol2.u)/2
    if c>0
        x2 = (sol1.u .- sol2.u)/2
    else
        x2 = (sol2.u .- sol1.u)/2
    end
    return x1,x2
end

# Ornstein-Uhlenbeck process
@model ou(rn,T,delta_t) = begin
    ampl ~ Uniform(0.0,5.0)
    b ~ beta(5.0,1.0)

    rn[1] ~ Normal(0,sqrt(ampl))

    for i=2:T
        rn[i] ~ Normal(rn[i-1]*b,sqrt(ampl*(1-b^2)))
    end
end

# Ornstein-Uhlenbeck process of two coupled oscillators
@model ou_corr(rn1,rn2,T,delta_t) = begin
    ampl1 ~ Uniform(0.0,5.0)
    ampl2 ~ Uniform(0.0,5.0)
    d ~ Uniform(0.0,5.0)
    b1 = exp(-delta_t*d/ampl1)
    b2 = exp(-delta_t*d/ampl2)

    rn1[1] ~ Normal(0,sqrt(ampl1))
    rn2[1] ~ Normal(0,sqrt(ampl2))

    for i=2:T
        rn1[i] ~ Normal(rn1[i-1]*b1,sqrt(ampl1*(1-b1^2)))
        rn2[i] ~ Normal(rn2[i-1]*b2,sqrt(ampl2*(1-b2^2)))
    end
end

# Ornstein-Uhlenbeck process with added Gaussian noise
@model oupn(rn,T,delta_t,::Type{R}=Vector{Float64}) where {R} = begin
    ampl ~ Uniform(0.0,5.0)
    b ~ Beta(5.0,1.0)
    noise_ampl ~ Uniform(0.0,1)

    b = exp(-delta_t/tau)
    r = R(undef, T)

    r[1] ~ Normal(0,sqrt(ampl))

    for i=2:T
        r[i] ~ Normal(r[i-1]*b,sqrt(ampl*(1-b^2)))
    end
    rn ~ MvNormal(r,sqrt(noise_ampl))
end

x1, x2 = corr_osc(-0.8)

p1 = Plots.plot(x1)
p1 = Plots.plot!(x2)

pearson = Statistics.cor(x1,x2)
println(pearson)

# lets see whether we can estimate c from the data
y1 = x1 .+ x2
y2 = x1 .- x2
chn = sample(ou_corr(y1,y2,length(y1),0.1), NUTS(0.65), 10000)

print(describe(chn))
p2 = plot(chn)

ampl1 = Array(chn[:ampl1])
ampl2 = Array(chn[:ampl2])

if ampl1>ampl2
    c = (ampl1 .- ampl2)./ampl2
else
    c = (ampl1 .- ampl2)./ampl1
end

println("A1,A2: ",mean(ampl1),",",mean(ampl2))
println("c estimate: ",mean(c),"std: ",std(c))


T = 3
n = 2
m = 1

A = [1.0 1.0; 0.0 1.0]
B = [0.0; 1.0]

Q = [1.0 0.0; 0.0 1.0]
R = [1.0]

x_init = [1.0; 0.0]

nz = T * n + (T - 1) * m + T * n

function unpack(z)
	x3 = z[1:n]
	u2 = z[n .+ (1:m)]
	x2 = z[n + m .+ (1:n)]
	u1 = z[n + m + n .+ (1:m)]
	x1 = z[n + m + n + m .+ (1:n)]
	y3 = z[n + m + n + m + n .+ (1:n)]
	y2 = z[n + m + n + m + n + n .+ (1:n)]
	y1 = z[n + m + n + m + n + n + n .+ (1:n)]

	return x3, u2, x2, u1, x1, y3, y2, y1
end

z0 = rand(nz)

function lagrangian(z, x_init)
	x3, u2, x2, u1, x1, y3, y2, y1 = unpack(z)

	L = 0.0

	# objective
	L += transpose(x3) * Q * x3
	L += (transpose(u2) * R * u2)[1]
	L += transpose(x2) * Q * x2
	L += (transpose(u1) * R * u1)[1]
	L += transpose(x1) * Q * x1

	# constraints
	L += transpose(y3) * (A * x2 + B * u2[1] - x3)
	L += transpose(y2) * (A * x1 + B * u1[1] - x2)
	L += transpose(y1) * (x_init - x1)

	return L
end

@variables z_sym[1:nz]
@variables x_sym[1:n]

L = lagrangian(z_sym, x_sym)
L = simplify.(L)

dL = Symbolics.gradient(L, z_sym)
ddL = Symbolics.hessian(L, z_sym)
ddLx = Symbolics.jacobian(dL, x_sym)

L_grad = eval(Symbolics.build_function(dL, z_sym, x_sym)[1])
L_hess = eval(Symbolics.build_function(ddL, z_sym, x_sym)[1])
Lx_hess = eval(Symbolics.build_function(ddLx, z_sym, x_sym)[1])

function solve(z0, x_init)
	z = copy(z0)

	∇L = L_grad(z, x_init)

	println()
	for i = 1:10
		r_norm = norm(∇L)
		println("iter ($i)")
		println("	grad norm: $(r_norm)")
		r_norm < 1.0e-8 && break

		∇²L = L_hess(z, x_init)

		Δ = ∇²L \ ∇L

		α = 1.0
		ẑ = z - α * Δ

		∇L̂ = L_grad(ẑ, x_init)

		iter = 1
		while norm(∇L̂) > norm(∇L)
			α *= 0.5
			ẑ = z - α * Δ
			∇L̂ = L_grad(ẑ, x_init)
			iter += 1

			if iter > 100
				break
			end
		end

		z = ẑ
		∇L = ∇L̂
	end

	return z, ∇L, L_hess(z, x_init)
end

z_sol, ∇L_sol, ∇²L_sol = solve(z0, x_init)
K_direct = (∇²L_sol \ Lx_hess(z_sol, x_init))[n + m + n .+ (1:m), :]

# LQR solution
K, P = tvlqr([A for t = 1:T-1], [B for t = 1:T-1], [Q for t = 1:T], [R[1,1] for t = 1:T-1])

@assert norm(K[1] - K_direct) < 1.0e-8


using Pkg; Pkg.activate(pwd())
using DEMC
using Random
using LinearAlgebra
using Distributions
Random.seed!(319531501)

# generate data
nobs = 1000
npar = 25
Σ = ones(npar, npar)
for i = 1:npar
    for iprime = 1:i
        if i != iprime
            Σ[i, iprime] = 0.25
            Σ[iprime, i] = 0.25
        end
    end
end

DataDistribution = MvNormal(zeros(npar), Σ)
X = zeros(nobs, npar+1)
X[:,1] .= 1.
for iobs = 1:nobs
    X[iobs, 2:end] .= rand(DataDistribution)
end

β = 1 .+ rand(npar+1)*3
y = X*β + randn(nobs)

# log objective function - for correct std errors need to use optimal weights
log_obj(b) = -0.5 * sum((y.-X*b).^2)

# set up of DEMCz chain
ndim = length(β)
opts = DEMC.demcopt(ndim)
opts.blockindex = [1:ndim] # parameter blocks: here choose all parameters to be updated simultaenously
opts.Nblocks = length(opts.blockindex)
opts.eps_scale = 1e-5*ones(ndim) # scale of random error around DE update
opts.γ = 2. # scale of DE update, 2.38 is the "optimal" number for a normal distribution
opts.N = 5 # number of chains
opts.K = 10 # every K steps add current N draws to Z
# Number of iterations in Chain
opts.Ngeneration = 100000
opts.autostop = :Rhat
opts.autostop_every = 2000
opts.autostop_Rhat = 1.1
opts.print_step = 1000
opts.verbose = false

Z = randn((10*ndim, ndim))

# run chain
mc, Z = DEMC.demcz_sample(log_obj, Z, opts)

# drop first half of chain
N, Npar, Ntot = size(mc.chain)
keep = Int(Ntot-opts.autostop_every)+1:Ntot
Ngen_burned = length(keep)
chain_burned = mc.chain[:,:,keep]
chainflat = DEMC.flatten_chain(chain_burned, N, Ngen_burned, Npar)'

# covariance of estimates
b, Σb = DEMC.mean_cov_chain(chain_burned)
# correct std errors for incorrect weighting matrix
# OLS
bols = (X'*X)\(X'*y)
ui = y .- X*bols
s2hat = (1/(nobs-npar)) * sum(ui.^2)
Σols = inv(X'X)*s2hat

# difference simulated to OLS
println("\n MCMC estimates: ", b, "\n dist to OLS: ", b .- bols)

accept_ratio, Rhat = DEMC.convergence_check(mc.chain, mc.log_obj, "none"; verbose = false, parnames = [])


# indicator of the zero cone

export IndZero

"""
**Indicator of the zero cone**

    IndZero()

Returns the indicator function of the set containing the origin, the "zero cone".
"""
struct IndZero <: ProximableFunction end

is_separable(f::IndZero) = true
is_convex(f::IndZero) = true
is_singleton(f::IndZero) = true
is_cone(f::IndZero) = true
is_affine(f::IndZero) = true

function (f::IndZero)(x::AbstractArray{C}) where {R <: Real, C <: Union{R, Complex{R}}}
    for k in eachindex(x)
        if x[k] != C(0)
            return R(Inf)
        end
    end
    return R(0)
end

function prox!(y::AbstractArray{C}, f::IndZero, x::AbstractArray{C}, gamma=R(1)) where {R <: Real, C <: Union{R, Complex{R}}}
    for k in eachindex(x)
        y[k] = C(0)
    end
    return R(0)
end

fun_name(f::IndZero) = "indicator of the zero cone"
fun_dom(f::IndZero) = "AbstractArray{Real}, AbstractArray{Complex}"
fun_expr(f::IndZero) = "x ↦ 0 if all(x = 0), +∞ otherwise"
fun_params(f::IndZero) = "none"

function prox_naive(f::IndZero, x::AbstractArray{C}, gamma=R(1)) where {R <: Real, C <: Union{R, Complex{R}}}
    return zero(x), R(0)
end


# Test Benders Algorithm for Mixed Integer Linear Programming
# Edward J. Xu, edxu96@outlook.com
# April 5th, 2019
# ---------------------------------------------------------------------------
push!(LOAD_PATH, "$(homedir())/Desktop/Benders Decomposition, MILP")
cd("$(homedir())/Desktop/Benders Decomposition, MILP")
using Benders
# ---------------------------------------------------------------------------
# Test Example 1: mat_a and mat_b being column vector
n_x = 1
n_y = 1
vec_min_y = hcat([0])
vec_max_y = hcat([10])
vec_c = hcat([5])
vec_f = hcat([-3])
vec_b = hcat([4; 0; -13])
mat_a = hcat([1; 2; 1])
mat_b = hcat([2; -1; -3])
Benders.milp(n_x = n_x,
             n_y = n_y,
             vec_min_y = vec_min_y,
             vec_max_y = vec_max_y,
             vec_c = vec_c,
             vec_f = vec_f,
             vec_b = vec_b,
             mat_a = mat_a,
             mat_b = mat_b,
             epsilon = 0,
             timesIterationMax = 5)
# ---------------------------------------------------------------------------
# Test Example 2: mat_a and mat_b being matrix
n_x_2 = 2
n_y_2 = 2
vec_min_y_2 = hcat([0; 0])
vec_max_y_2 = hcat([10; 10])
vec_c_2 = hcat([5; 3])
vec_f_2 = hcat([-3; 1])
vec_b_2 = hcat([4; 0; -13])
mat_a_2 = [1 3; 2 1; 1 -5]
mat_b_2 = [2 -4; -1 2; -3 1]
Benders.milp(n_x = n_x_2,
             n_y = n_y_2,
             vec_min_y = vec_min_y_2,
             vec_max_y = vec_max_y_2,
             vec_c = vec_c_2,
             vec_f = vec_f_2,
             vec_b = vec_b_2,
             mat_a = mat_a_2,
             mat_b = mat_b_2,
             epsilon = 0.000001,
             timesIterationMax = 10)
# ---------------------------------------------------------------------------
# Test Example 3: sub and ray problem
n_x_3 = 2
n_y_3 = 2
vec_min_y_3 = hcat([0; 0])
vec_max_y_3 = hcat([2; 2])
vec_c_3 = hcat([2; 6])
vec_f_3 = hcat([2; 3])
vec_b_3 = hcat([5; 4])
mat_a_3 = [-1 2; 1 -3]
mat_b_3 = [3 -1; 2 2]
Benders.milp(n_x = n_x_3,
             n_y = n_y_3,
             vec_min_y = vec_min_y_3,
             vec_max_y = vec_max_y_3,
             vec_c = vec_c_3,
             vec_f = vec_f_3,
             vec_b = vec_b_3,
             mat_a = mat_a_3,
             mat_b = mat_b_3,
             epsilon = 0.000001,
             timesIterationMax = 10)


#
# Denoise an SPD Example with Douglas Rachford
#
# Denoise an SPD Example with Parallel Douglas-Rachford to minimize the
#
# L2-TV functional with anisotropic TV
#
# where the example is the same data and cost as for SPD_Image_CP_vs_CPPA
#
# This example and its data is used for comparison in Section 6.2 of
#
# > R. Bergmann, R. Herzog, M. Silva Louzeiro, D. Tenbrinck, J. Vidal Núñez:
# > _Fenchel Duality Theory and a Primal-Dual Algorithm on Riemannian Manifolds_,
# > arXiv: [1908.02022](https://arxiv.org/abs/1908.02022)
#
using Manopt, Manifolds
using Images, CSV, DataFrames, LinearAlgebra, JLD2, Dates
#
# Settings
experiment_name = "SPD_Image_DR"
export_orig = true
export_result = true
export_table = true
use_debug = true
asy_render_detail = 2
results_folder = joinpath(@__DIR__, "Image_TV")
comparison_data = joinpath(
    @__DIR__, "..", "CyclicProximalPoint", "Image_TV", "SPD_Image_CPPA-cost.jld2"
)
!isdir(results_folder) && mkdir(results_folder)
#
# Parameters
η = 0.58
λ = 0.93
α = 6.0
#
# Manifold & Data
f = artificial_SPD_image2(32)
if export_orig
    fn = joinpath(results_folder, experiment_name * "-orig.asy")
    asymptote_export_SPD(fn; data=f, scale_axes=(7.5, 7.5, 7.5))
    render_asymptote(fn; render=asy_render_detail)
end
sC = StopAfterIteration(400)
try
    cppa_data = load(comparison_data)
    cost_threshold = cppa_data["cost_function_value"]
    cppa_iter = cppa_data["iterations"]
    global sC = StopWhenAny(StopAfterIteration(400), StopWhenCostLess(cost_threshold))
    @info "Comparison to CPPA (`SPDImage_CPPA.jl`) and its cost after $(cppa_iter) iterations of $cost_threshold."
catch e
    msg = sprint(showerror, e)
    @info "Comparison to CPPA (`CyclicProximalPoint/SPDImage_CPPA.jl`) not possible, data file $(comparison_data) either missing or corrupted.\n Error: $(msg).\n\n Starting with a default stopping criterion."
end
#
# Build Problem for L2-TV
pixelM = SymmetricPositiveDefinite(3);
M = PowerManifold(pixelM, NestedPowerRepresentation(), size(f)...)
d = length(size(f))
rep(d) = (d > 1) ? [ones(Int, d)..., d] : d
fidelity(x) = 1 / 2 * distance(M, x, f)^2
Λ(x) = forward_logs(M, x) # on T_xN
prior(x) = norm(norm.(Ref(pixelM), repeat(x, rep(d)...), Λ(x)), 1)
#
# Setup & Optimize
print("--- Douglas–Rachford with η: $(η) and λ: $(λ) ---\n")
cost(x) = 1 / α * fidelity(x[1]) + prior(x[1])
N = PowerManifold(M, NestedPowerRepresentation(), 5)
prox1 = (η, x) -> [prox_distance(M, η, f, x[1]), prox_parallel_TV(M, α * η, x[2:5])...]
prox2 = (η, x) -> fill(mean(M, x, GradientDescentEstimation(); stop_iter=4), 5)
x0 = fill(f, 5)
@time o = DouglasRachford(
    N,
    cost,
    [prox1, prox2],
    x0;
    λ=i -> η,
    α=i -> λ, # map from Paper notation of BPS16 to toolbox notation
    debug=use_debug ? [:Iteration, " | ", :Cost, "\n", 10, :Stop] : missing,
    record=export_table ? [:Iteration, :Cost] : missing,
    stopping_criterion=sC,
    return_options=true,
)
y = get_solver_result(o)[1]
export_table && (r = get_record(o))
#
# Result
numIter = length(r)
if export_result
    fn = joinpath(
        results_folder,
        experiment_name * "img-result-$(numIter)-α$(replace(string(α), "." => "-")).asy",
    )
    asymptote_export_SPD(fn; data=y, scale_axes=(7.5, 7.5, 7.5))
    render_asymptote(fn; render=asy_render_detail)
end
if export_table
    # scale cost back for saving such that its comparable with the other two results
    A = cat([ri[1] for ri in r], [ri[2] / α for ri in r]; dims=2)
    CSV.write(
        joinpath(results_folder, experiment_name * "-Cost.csv"),
        DataFrame(A);
        writeheader=false,
    )
end


# # [Linear Time Continuous System with Controls](@id linear_continuous_controls)
#
# Now we will extend the [`previous example`](@ref linear_continuous) by adding some exegeneous control signals.
# As always, we will generate some data via `OrdinaryDiffEq.jl`

using DataDrivenDiffEq
using ModelingToolkit
using LinearAlgebra
using OrdinaryDiffEq
#md using Plots

A = [-0.9 0.2; 0.0 -0.2]
B = [0.0; 1.0]
u0 = [10.0; -10.0]
tspan = (0.0, 10.0)

f(u,p,t) = A*u .+ B .* sin(0.5*t)

sys = ODEProblem(f, u0, tspan)
sol = solve(sys, Tsit5(), saveat = 0.05);

# We will use the data provided by our problem, but add the control signal `U = sin(0.5*t)` to it.
X = Array(sol)
t = sol.t
control(u,p,t) = [sin(0.5*t)]
prob = ContinuousDataDrivenProblem(X, t, U = control)

# And plot the problems data.

#md plot(prob)

# Again, we will use `gDMD` to estimate the systems dynamics. Since we have a control signal
# defined in the problem, the algorithm will detect it automatically and use `gDMDc`:

res = solve(prob, DMDSVD(), digits = 1)
#md println(res)

# We see that the system has been recovered correctly, indicated by the small error and high AIC score of the result. We can confirm this by looking at the resulting [`Basis`](@ref)

#md
system = result(res)
#md println(system)

# And also plot the prediction of the recovered dynamics

#md plot(res)

# Again, we can have a look at the generator of the system, which is independent from the inputs.

generator(system)

# Sticking to the same procedure as earlier, we now use a linear sparse regression to solve the problem

@parameters t
@variables x[1:2](t) u[1:1](t)

basis = Basis([x; u], x, controls = u, independent_variable = t, name = :LinearBasis)
#md print(basis) #hide

# Note that we added a new variable `u[1](t)` as a control to both the equations and the basis constructor.
# Afterwards, we simply `solve` the already defined problem with our `Basis` and a `SparseOptimizer`

sparse_res = solve(prob, basis, STLSQ(1e-1))
#md println(sparse_res)

# Which holds the same equations
sparse_system = result(sparse_res)
#md println(sparse_system)

# Again, we can have a look at the result

#md plot(
#md     plot(prob), plot(sparse_res), layout = (1,2)
#md )

# Both results can be converted into an `ODESystem`. To include the control signal, we simply
# substitute the control variables in the corresponding equations.

subs_control = (u[1] => sin(0.5*t))

eqs = map(equations(sparse_system)) do eq
    eq.lhs ~ substitute(eq.rhs, subs_control)
end

@named sys = ODESystem(
    eqs,
    get_iv(sparse_system),
    states(sparse_system),
    parameters(sparse_system)
    );

# And simulated using `OrdinaryDiffEq.jl` using the (known) initial conditions and the parameter mapping of the estimation.

x0 = [x[1] => u0[1], x[2] => u0[2]]
ps = parameter_map(sparse_res)

ode_prob = ODEProblem(sys, x0, tspan, ps)
estimate = solve(ode_prob, Tsit5(), saveat = prob.t);

# And look at the result
#md plot(sol, color = :black)
#md plot!(estimate, color = :red, linestyle = :dash)

#md # ## [Copy-Pasteable Code](@id linear_continuous_copy_paste)
#md #
#md # ```julia
#md # @__CODE__
#md # ```

@test all(aic(sparse_res) .>= 1e3) #src
@test all(l2error(sparse_res) .<= 5e-1) #src
@test all(determination(sparse_res) .>= 0.97) #src
@test Array(sol) ≈ Array(estimate) rtol = 5e-2 #src


#*************************************************************************#
#
# AbstractType
#
#*************************************************************************#
abstract type ElementwiseOperator <: AparametricModel end

#*************************************************************************#
#
# Generate Buffers
#
#*************************************************************************#
function dx_buffer(f::ElementwiseOperator)
    return Diagonal{Float64}(zeros(f.noutputs))
end

#*************************************************************************#
#
# LogisticFunction
#
#*************************************************************************#
"""
    LogisticFunction(n)

Initialize `ParametricModel` representing logistic function. It applies
`f(x) = 1 / (1+exp(-x))` elementwise to a given input vector.
`n` is the input vector dimension.
"""
struct LogisticFunction <: ElementwiseOperator
    ninputs::Int
    noutputs::Int
end

function LogisticFunction(n::Int)
    LogisticFunction(n, n)
end

function _evaluate!(f::LogisticFunction,
                    x::AbstractVector{<:Real},
                    z::AbstractVector{Float64},
                    dx::AbstractMatrix{Float64})

    z .= 1 ./ (1+exp.(-x))

    if !isempty(dx)
        dx.diag .=  z .* (1-z)
    end
end

#*************************************************************************#
#
# Hyperbolic Tangent
#
#*************************************************************************#
"""
    HyperbolicTangent(n)

Initialize `ParametricModel` representing hyperbolic tangent. It
applies `f(x) = tanh(x)` elementwise to a given input vector.
 `n` is the input vector dimension.
"""
struct HyperbolicTangent <: ElementwiseOperator
    ninputs::Int
    noutputs::Int
end

function HyperbolicTangent(n::Int)
    HyperbolicTangent(n, n)
end

function _evaluate!(f::HyperbolicTangent,
                    x::AbstractVector{<:Real},
                    z::AbstractVector{Float64},
                    dx::AbstractMatrix{Float64})
    z .= tanh.(x)

    if !isempty(dx)
        dx.diag .=  1 - z.^2
    end
end

#*************************************************************************#
#
# AffineMap
#
#*************************************************************************#
"""
    AffineMap(a[, b])

Initialize `ParametricModel` representing affine function. It applies
`f(x) = a*x + b` elementwise to a given input vector. `a` and
`b` are vectors.
"""
struct AffineMap <: ElementwiseOperator
    ninputs::Int
    noutputs::Int
    a::Vector{Float64}
    b::Vector{Float64}
end

function AffineMap{R<:Real, S<:Real}(a::Vector{R}, b::Vector{S}=zeros(length(a)))
    if length(a) != length(b)
        throw(ArgumentError("Vector `a` and `b` should have the same size."))
    end

    n = length(a)
    AffineMap(n, n, a, b)
end

function _evaluate!(f::AffineMap,
                    x::AbstractVector{<:Real},
                    z::AbstractVector{Float64},
                    dx::AbstractMatrix{Float64})
    z .= f.a.*x + f.b

    if !isempty(dx)
        dx.diag .=  f.a
    end
end

#*************************************************************************#
#
# Identity
#
#*************************************************************************#
"""
    Identity(n)

Initialize `ParametricModel` representing identity function.  It applies
`f(x) = x` elementwise to a given input vector. `n` is the input vector
dimension.
"""
struct Identity <: ElementwiseOperator
    ninputs::Int
    noutputs::Int
end

function Identity(n::Int)
    Identity(n, n)
end

function _evaluate!(f::Identity,
                    x::AbstractVector{<:Real},
                    z::AbstractVector{Float64},
                    dx::AbstractMatrix{Float64})
    z .= x

    if !isempty(dx)
        dx.diag .=  1
    end
end


using LinearAlgebra;
using Printf;

function linspace(start::Number, stop::Number, n::Int64)
	linx = zeros(Float64, n);
	delta = (stop - start) / (n - 1);
	for i = 1:n
		linx[i] = start + delta * (i - 1)
	end
	return linx;
end

lx = 1.0;
t = 0.0
nx = 11;
dx = lx / (nx - 1);
dt = 1e-2;
tlim = 0.1;
kappa = 1.0
beta = dt * kappa / 2.0dx^2;
t_save = linspace(0.0, tlim, nx);
save_count = 1;

function solve_linear(n::Int64, a::Float64, b::Float64, c::Float64, r::Array)
	A = zeros(n, n);
	for i = 2:n+1
		for j = 2:n+1
			if(i-1 == j-1-1)
				A[i-1, j-1] = c;
			elseif (i-1==j-1)
				A[i-1,j-1] = b;
			elseif (i-1-1 == j-1)
				A[i-1,j-1] = a;
			end
		end
	end
	return A \ r;
end

function crank_nicolson(u::Array)
	a = Float64(0.0);
	b = Float64(0.0);
	c = Float64(0.0);
	r = zeros(Float64, nx - 2);
	for i = 2:nx-1
		a = -beta;
		b = 1.0 + 2.0beta;
		c = -beta;
		r[i-1] = u[i] + beta*(u[i+1] - 2.0u[i] + u[i-1]);
	end
	r[begin] = r[begin] - a*u[begin]
	r[end] = r[end] - c * u[end];
	u[2:end-1] = solve_linear(nx - 2, a, b, c, r)
end

function output(x::Array, u::Array, t::Float64)
	filename = @sprintf("cn_t=%4.3f.dat", t);
	open(filename, "w") do io
		for i = 1:nx
			write(io, @sprintf("%.5f\t%.5f\n", x[i], u[i]));
		end
	end;
end

function main()
	@printf("beta = %.5f, dt = %.5f, dx = %.5f\n", beta, dt, dx);
	u = ones(Float64, nx);
	u[begin] = u[end] = 0.0;
	x = linspace(0, lx, nx);
	while t < tlim
		output(x, u, t);
		crank_nicolson(u);
		global t = t + dt;
	end
end

main()


using RunTests
using Base.Test

@testmodule ErrorAnalysisTest begin
  using MLLibrary

  function test_confusionMatrix()
    testset = [1 0 1 1 1 1 1 0 0 0 0]
    predset = [1 1 0 1 1 1 0 1 0 0 0]

    @test MLLibrary.Utils.ErrorAnalysis.generateConfusionMatrix(testset,predset)==[0 0 1;
                    0 3 2;
                    1 2 4;]

    testset = [1 2 3 4 1 2 4 1 3 2 1 2]
    predset = [1 1 4 1 1 2 3 2 3 2 1 4]

    @test MLLibrary.Utils.ErrorAnalysis.generateConfusionMatrix(testset,predset)==[0 1 2 3 4;
                    1 3 1 0 0;
                    2 1 2 0 1;
                    3 0 0 1 1;
                    4 1 0 1 0;]
  end

  function test_generateWeightedErrorRate()
    testset = [1 0 1 1 1 1 1 0 0 0 0]
    predset = [1 1 0 1 1 1 0 1 0 0 0]

    confusionmatrix =  MLLibrary.Utils.ErrorAnalysis.generateConfusionMatrix(testset,predset)

    @test round(MLLibrary.Utils.ErrorAnalysis.calculateWeightedSuccessRate(confusionmatrix),2) == .63

    testset = [1 2 3 4 1 2 4 1 3 2 1 2]
    predset = [1 1 4 1 1 2 3 2 3 2 1 4]

    confusionmatrix= MLLibrary.Utils.ErrorAnalysis.generateConfusionMatrix(testset,predset)

    @test round(MLLibrary.Utils.ErrorAnalysis.calculateWeightedSuccessRate(confusionmatrix),2) == .44

  end

  function test_generateErrorRate()
    testset = [1 0 1 1 1 1 1 0 0 0 0]
    predset = [1 1 0 1 1 1 0 1 0 0 0]

    @test round(MLLibrary.Utils.ErrorAnalysis.calculateErrorRate(testset,predset),2)==.36
    #
    testset = [1 2 3 4 1 2 4 1 3 2 1 2]
    predset = [1 1 4 1 1 2 3 2 3 2 1 4]

    @test MLLibrary.Utils.ErrorAnalysis.calculateErrorRate(testset,predset)==.5

  end

  function test_calculateRootMeanStandardError()
    y=[1;2;3;4;5;6;1;2;3]
    observed = [0;4;2;1;5;2;1;6;1]

    @test round(MLLibrary.Utils.ErrorAnalysis.calculateRootMeanStandardError(observed,y),2)==2.38

  end

  function test_calculateWCSS()
    centroids=[1 2 3 4;5 6 7 8]
    data=[1 2 1 1;4 1 2 2; 5 2 1 4;3 3 2 1]
    y=[1;1;2;2]
    wcss = MLLibrary.Utils.ErrorAnalysis.calculateWCSS(centroids,data,y)
    println(wcss)

    @test wcss==183
  end

end


@testset "Poisson" begin

    xmin = 0.0
    xmax = 2π
    ymin = 0.0
    ymax = 2π

    nx = 128
    ny = 128

    ux = zeros(Float64, nx, ny)
    uy = zeros(Float64, nx, ny)

    omega = zeros(Float64, nx, ny)
    psi = zeros(Float64, nx, ny)

    ux_exact = zeros(nx, ny)
    uy_exact = zeros(nx, ny)
    psi_exact = zeros(nx, ny)

    grid = Grid(xmax-xmin, nx, ymax - ymin, ny)

    poisson = Poisson(grid)

    mode = 2
    for i = 1:nx, j = 1:ny
        x = xmin + (i - 1) * (xmax - xmin) / nx
        y = ymin + (j - 1) * (ymax - ymin) / ny
        psi_exact[i, j] = mode * sin(mode * x) * cos(mode * y)
        ux_exact[i, j] = mode^2 * sin(mode * x) * sin(mode * y)
        uy_exact[i, j] = mode^2 * cos(mode * x) * cos(mode * y)
        omega[i, j] = 2 * mode^3 * sin(mode * x) * cos(mode * y)
    end

    rhs = copy(omega)
    compute_psi!(psi, poisson, rhs)
    @test psi_exact ≈ psi

    rhs = copy(omega)
    compute_velocities!(ux, uy, poisson, rhs)
    @test ux_exact ≈ ux
    @test uy_exact ≈ uy

end


data = split(strip(read(joinpath(@__DIR__, "./2.txt"), String)), '\n')

#part 1
position = Dict(:horizontal => 0, :depth => 0)

command = Dict("forward" => (:horizontal,+), "down" => (:depth,+), "up" => (:depth,-))

function executeCommands(data)
  global position,command
  direction,distance = split(data,' ')
  position[command[direction][1]] = command[direction][2](position[command[direction][1]],parse(Int,distance))
end

executeCommands.(data)
position[:horizontal]*position[:depth]

# part 2
position = Dict(:horizontal => 0, :depth => 0, :aim => 0)

command = Dict("forward" => (:horizontal,+), "down" => (:aim,+), "up" => (:aim,-))

function executeCommands2(data)
  global position
  direction,distance = split(data,' ')
  d = parse(Int,distance)
  position[command[direction][1]] = command[direction][2](position[command[direction][1]],d)
  if direction == "forward"
    position[:depth] = position[:depth] + d * position[:aim]
  end
end

executeCommands2.(data)
position[:horizontal]*position[:depth]


using LinearAlgebra
include("poisson_linear_system.jl")
include("simpson_integrate.jl")

function poisson_error(numelements)
    xL = -1.0
    xR = +1.0
    wavenumber = pi / 2
    bcleft = bcright = 0.0

    numpts = numelements + 1
    stepsize = (xR - xL) / numelements
    xrange = range(xL, stop = xR, length = numpts)

    rhs = cos.(wavenumber * xrange)
    op = -poisson_linear_system(numpts)
    rhs[1] = bcleft / stepsize^2
    rhs[end] = bcright / stepsize^2
    sol = stepsize^2 * (op \ rhs)
    exactsolution =
        1.0 / wavenumber^2 * (cos.(wavenumber * xrange) .- cos(wavenumber))
    err = simpson_integrate(sol - exactsolution, stepsize)
    return err
end

function convergence_rate(err,dx)
    return diff(log.(err)) ./ diff(log.(dx))
end

powers = 1:10
numelements = 2 .^ powers
stepsize = 2.0 ./ numelements
err = poisson_error.(numelements)
rate = convergence_rate(err,stepsize)

# using Plots
# plot(stepsize,err,xscale = :log, yscale = :log)
# scatter!(stepsize,err)


#=============================================

Definitions of various similarity functions

=============================================#

using QuadGK
using LinearAlgebra: dot, norm

#====================
Definitions of built-in similarity functions
====================#

#====================
Cosine similarity
====================#

@doc raw"""
    cossim(x,y)

Computes the cosine similarity between two inputs ``x`` and ``y``. Cosine similarity is defined as

``\text{cossim}(x,y) = \frac{\left\langle x,y\right\rangle}{\|x\|\cdot\|y\|}``

where ``\left\langle\cdot,\cdot\right\rangle`` is an inner product (e.g. dot product) and ``\|\cdot\|`` is its derived norm. This is roughly interpreted as being related to the angle between the inputs ``x`` and ``y``: when ``x`` and ``y`` have low angle between them, `cossim(x,y)` is high (close to ``1``). When ``x`` and ``y`` have large angle between them, `cossim(x,y)` is low (close to ``-1``).

# Arguments
- `x` and `y`: two inputs for which `dot(x,y)`, `norm(x)`, and `norm(y)` are defined.

# Examples
```jldoctest; setup = :(using LSHFunctions)
julia> using LinearAlgebra: dot, norm;

julia> x, y = rand(4), rand(4);

julia> cossim(x,y) == dot(x,y) / (norm(x) * norm(y))
true

julia> z = rand(5);

julia> cossim(x,z)
ERROR: DimensionMismatch("dot product arguments have lengths 4 and 5")
```

See also: [`SimHash`](@ref)
"""
function cossim(x::AbstractVector, y::AbstractVector)
    norm_x = norm(x)
    norm_y = norm(y)

    if norm_x == 0 || norm_y == 0
        "x and y must be nonzero" |> ErrorException |> throw
    end

    dot(x,y) / (norm_x * norm_y)
end

function cossim(f, g, interval::LSHFunctions.RealInterval)
    norm_f = L2_norm(f, interval)
    norm_g = L2_norm(g, interval)

    if norm_f == 0 || norm_g == 0
        "f and g must be nonzero" |> ErrorException |> throw
    end

    inner_prod(f, g, interval) / (norm_f * norm_g)
end

#====================
L^p distance
====================#

@doc raw"""
    ℓp(x::AbstractVector, y::AbstractVector, p::Real=2)
    ℓ1(x::AbstractVector, y::AbstractVector)
    ℓ2(x::AbstractVector, y::AbstractVector)

Computes the ``\ell^p`` distance between a pair of vectors, given by

``\ell^p(x,y) \coloneqq \|x - y\|_p = \left(\sum_i \left|x_i - y_i\right|^p\right)^{1/p}``

`ℓ1(x,y)` is the same as `ℓp(x,y,1)`, and `ℓ2(x,y)` is the same as `ℓp(x,y,2)`.

# Examples
```jldoctest; setup = :(using LSHFunctions)
julia> x = [1, 2, 3];

julia> y = [4, 5, 6];

julia> ℓp(x,y,2) == (abs(1-4)^2 + abs(2-5)^2 + abs(3-6)^2)^(1/2)
true

julia> ℓp(x,y,3) == (abs(1-4)^3 + abs(2-5)^3 + abs(3-6)^3)^(1/3)
true
```

See also: [`ℓp_norm`](@ref), [`L1Hash`](@ref), [`L2Hash`](@ref)
"""
ℓp(x::AbstractVector, y::AbstractVector, p::Real=2) = Lp(x, y, p)

@doc (@doc ℓp)
ℓ1(x::AbstractVector, y::AbstractVector) = L1(x, y)

@doc (@doc ℓp)
ℓ2(x::AbstractVector, y::AbstractVector) = L2(x, y)

@doc raw"""
    Lp(x::AbstractVector, y::AbstractVector, p::Real=2)
    L1(x::AbstractVector, y::AbstractVector)
    L2(x::AbstractVector, y::AbstractVector)

Computes the ``ℓ^p`` distance between a pair of vectors ``x`` and ``y``. Identical to `ℓp(x,y,p)`, `ℓ1(x,y)`, and `ℓ2(x,y)`, respectively.

See also: [`ℓp`](@ref)
"""
function Lp(x::AbstractVector{T}, y::AbstractVector, p::Real=2) where T
    if p ≤ 0
        "p must be positive" |> ErrorException |> throw
    elseif length(x) != length(y)
        "length(x) != length(y)" |> DimensionMismatch |> throw
    end

    result = T(0)
    @inbounds @simd for ii = 1:length(x)
        result += abs(x[ii] - y[ii])^p
    end

    return result^(1/p)
end

@doc (@doc Lp)
function L1(x::AbstractVector{T}, y::AbstractVector) where T
    if length(x) != length(y)
        "length(x) != length(y)" |> DimensionMismatch |> throw
    end

    result = T(0)
    @inbounds @simd for ii = 1:length(x)
        result += abs(x[ii] - y[ii])
    end

    return result
end

@doc (@doc Lp)
function L2(x::AbstractVector{T}, y::AbstractVector) where T
    if length(x) != length(y)
        "length(x) != length(y)" |> DimensionMismatch |> throw
    end

    result = T(0)
    @inbounds @simd for ii = 1:length(x)
        result += abs2(x[ii] - y[ii])
    end

    return √result
end

# Function space L^p distances

@doc raw"""
    Lp(f, g, interval::LSHFunctions.RealInterval, p)
    L1(f, g, interval::LSHFunctions.RealInterval)
    L2(f, g, interval::LSHFunctions.RealInterval)

Computes the ``L^p`` distance between two functions, given by

``L^p(f,g) \coloneqq \|f - g\|_p = \left(\int_a^b \left|f(x) - g(x)\right|^p \hspace{0.15cm} dx\right)^{1/p}``

# Examples
Below we compute the ``L^1``, ``L^2``, and ``L^3`` distances between ``f(x) = x^2 + 1`` and ``g(x) = 2x`` over the interval ``[0,1]``. The distances are computed by evaluating the integral

``\left(\int_0^1 \left|f(x) - g(x)\right|^p \hspace{0.15cm}dx\right)^{1/p} = \left(\int_0^1 \left|x^2 - 2x + 1\right|^p \hspace{0.15cm}dx\right)^{1/p} = \left(\int_0^1 (x - 1)^{2p} \hspace{0.15cm}dx\right)^{1/p}``

for ``p = 1``, ``p = 2``, and ``p = 3``.

```jldoctest; setup = :(using LSHFunctions)
julia> f(x) = x^2 + 1; g(x) = 2x;

julia> interval = @interval(0 ≤ x ≤ 1);

julia> Lp(f, g, interval, 1) ≈ L1(f, g, interval) ≈ 3^(-1)
true

julia> Lp(f, g, interval, 2) ≈ L2(f, g, interval) ≈ 5^(-1/2)
true

julia> Lp(f, g, interval, 3) ≈ 7^(-1/3)
true
```

See also: [`Lp_norm`](@ref), [`ℓp`](@ref)
"""
Lp(f, g, interval::LSHFunctions.RealInterval, p::Real=2) =
    Lp_norm(x -> f(x) - g(x), interval, p)

@doc (@doc Lp)
L1(f, g, interval::LSHFunctions.RealInterval) = L1_norm(x -> f(x) - g(x), interval)

@doc (@doc Lp)
L2(f, g, interval::LSHFunctions.RealInterval) = L2_norm(x -> f(x) - g(x), interval)

#====================
Jaccard similarity
====================#

@doc raw"""
    jaccard(A::Set, B::Set) :: Float64

Computes the Jaccard similarity between sets ``A`` and ``B``, which is defined as

``\text{Jaccard}(A,B) = \frac{\left|A \cap B\right|}{\left|A \cup B\right|}``

# Arguments
- `A::Set`, `B::Set`: two sets whose Jaccard similarity we would like to compute.

# Examples
```jldoctest; setup = :(using LSHFunctions)
julia> A, B = Set([1, 2, 3]), Set([2, 3, 4]);

julia> jaccard(A,B)
0.5

julia> jaccard(A,B) == length(A ∩ B) / length(A ∪ B)
true
```

See also: [`MinHash`](@ref)
"""
function jaccard(A::Set, B::Set) :: Float64
    if isempty(A)
        # Use the convention that if A = B = ∅, their Jaccard
        # similarity is zero.
        Float64(0)
    else
        length(A ∩ B) / length(A ∪ B)
    end
end

@doc raw"""
    function jaccard(x::BitArray{1}, y::BitArray{1})

Computes the Jaccard similarity between a pair of binary vectors:

``J(x, y) = \frac{\sum_{i} \min{(x_i,y_i)}}{\sum_{i} \max{(x_i,y_i)}}``

# Arguments
- `x::BitArray{1}`, `y::BitArray{1}`: two binary vectors, in the form of `BitArray`s.

# Examples
```jldoctest; setup = :(using LSHFunctions)
julia> x = BitArray([true, false, true, true, false]);

julia> y = BitArray([false, false, true, true, true]);

julia> jaccard(x,y)
0.5
```
"""
function jaccard(x::BitArray{1}, y::BitArray{1}) :: Float64
    union = sum(x .| y)
    if union == 0
        # Use the convention that if x and y are full of zeros, their Jaccard
        # similarity is zero.
        Float64(0)
    else
        intersection = sum(x .& y)
        intersection / union
    end
end

@doc raw"""
    function jaccard(x::AbstractVector{<:Real}, y::AbstractVector{<:Real})

Computes the Jaccard similarity between a pair of vectors of real numbers:

``J(x, y) = \frac{\sum_{i} \min{(x_i,y_i)}}{\sum_{i} \max{(x_i,y_i)}}``

# Arguments
- `x::AbstractVector{<:Real}`, `y::AbstractVector{<:Real}`: a pair of vectors containing real numbers (subtypes of `Real`).

# Examples
```jldoctest; setup = :(using LSHFunctions)
julia> x = [0.8, 0.1, 0.3, 0.4, 0.1];

julia> y = [1.0, 0.6, 0.0, 0.4, 0.5];

julia> jaccard(x,y)
0.5
```
"""
function jaccard(x::AbstractVector{T},
                 y::AbstractVector) :: Float64 where {T <: Real}
    if length(x) != length(y)
        DimensionMismatch("dimensions must match") |> throw
    end

    intersection = T(0)
    union = T(0)

    @inbounds @simd for ii = 1:length(x)
        if 0 ≤ x[ii] ≤ y[ii]
            intersection += x[ii]
            union += y[ii]
        elseif 0 ≤ y[ii] < x[ii]
            intersection += y[ii]
            union += x[ii]
        else
            ErrorException("vectors must have non-negative elements") |> throw
        end
    end

    if union == T(0)
        # Use the convention that if x and y are full of zeros, their Jaccard
        # similarity is zero.
        Float64(0)
    else
        Float64(intersection / union)
    end
end

jaccard(x::AbstractVector{<:Integer}, y::AbstractVector{<:AbstractFloat}) =
    jaccard(y, x)

@doc raw"""
    function jaccard(A::Set{<:K},
                     B::Set{<:K},
                     weights::Dict{K,V}) where {K,V<:Number}

Computes the weighted Jaccard similarity between two sets:

``J(x, y) = \frac{\sum_{x\in A\cap B} w_x}{\sum_{y\in A\cup B} w_y}``

# Arguments
- `A::Set`, `B::Set`: two sets whose Jaccard similarity we would like to compute.
- `weights::Dict`: a dictionary mapping symbols in the sets `A` and `B` to numerical weights. These weights must be positive.

# Examples
```jldoctest; setup = :(using LSHFunctions)
julia> A = Set(["a", "b", "c"]);

julia> B = Set(["b", "c", "d"]);

julia> W = Dict("a" => 0.2, "b" => 2.4, "c" => 0.6, "d" => 1.8);

julia> jaccard(A,B,W)
0.6
```
"""
function jaccard(A::Set{<:K},
                 B::Set{<:K},
                 weights::Dict{K,V}) :: Float64 where {K,V<:Real}

    union_weight = V(0)

    for el in A ∪ B
        w = weights[el]
        if w < 0
            ErrorException("weights must be non-negative") |> throw
        end
        union_weight += w
    end

    intersection_weight = sum(weights[el] for el in A ∩ B)

    # By convention, if A = B = ∅, their Jaccard similarity is zero
    if union_weight == V(0)
        Float64(0)
    else
        Float64(intersection_weight / union_weight)
    end
end

#====================
Inner product and norms
====================#

### Inner products

@doc raw"""
    inner_prod(x::AbstractVector, y::AbstractVector)

Computes the ``\ell^2`` inner product (dot product)

``\left\langle x, y\right\rangle = \sum_i x_iy_i``

# Examples
```jldoctest; setup = :(using LSHFunctions)
julia> using LinearAlgebra: dot;

julia> x, y = randn(4), randn(4);

julia> inner_prod(x,y) ≈ dot(x,y)
true
```
"""
inner_prod(x::AbstractVector, y::AbstractVector) = dot(x,y)

# 1-dimensional inner product between L^2 functions
@doc raw"""
    inner_prod(f, g, interval::LSHFunctions.RealInterval)

Computes the ``L^2`` inner product

``\left\langle f, g\right\rangle = \int_a^b f(x)g(x) \hspace{0.15cm} dx``

where the interval we're integrating over is specified by the `interval` argument.

# Examples
```jldoctest; setup = :(using LSHFunctions)
julia> f(x) = cos(x); g(x) = sin(x);

julia> inner_prod(f, g, @interval(0 ≤ x ≤ π/2)) ≈ 1/2
true
```
"""
inner_prod(f, g, interval::LSHFunctions.RealInterval) =
    quadgk(x -> f(x)g(x), interval.lower, interval.upper)[1]

### L^p norms
@doc raw"""
    Lp_norm(x::AbstractVector, p::Real = 2)
    L1_norm(x::AbstractVector)
    L2_norm(x::AbstractVector)

Compute the ``\ell^p`` norm of a vector ``x``. Identical to `ℓp_norm(x,p)`, `ℓ1_norm(x)`, and `ℓ2_norm(x)`, respectively.

See also: [`ℓp_norm`](@ref)
"""
Lp_norm(x::AbstractVector, p::Real = 2) = norm(x,p)

@doc (@doc Lp_norm)
L1_norm(x::AbstractVector) = norm(x,1)

@doc (@doc Lp_norm)
L2_norm(x::AbstractVector) = norm(x)

@doc raw"""
    ℓp_norm(x::AbstractVector, p::Real = 2)
    ℓ1_norm(x::AbstractVector)
    ℓ2_norm(x::AbstractVector)

Compute the ``\ell^p`` norm of a point ``x``, defined as

``\|x\|_p = \left(\sum_i \left|x_i\right|^p\right)^{1/p}``

# Examples

```jldoctest; setup = :(using LSHFunctions)
julia> x = randn(4);

julia> ℓp_norm(x, 1) ≈ ℓ1_norm(x) ≈ (map(u -> abs(u)^1, x) |> sum)^(1/1)
true

julia> ℓp_norm(x, 2) ≈ ℓ2_norm(x) ≈ (map(u -> abs(u)^2, x) |> sum)^(1/2)
true

julia> ℓp_norm(x, 3) ≈ (map(u -> abs(u)^3, x) |> sum)^(1/3)
true
```

See also: [`ℓp`](@ref), [`Lp_norm`](@ref)
"""
ℓp_norm(x::AbstractVector, p::Real = 2) = Lp_norm(x, p)

@doc (@doc ℓp_norm)
ℓ1_norm(x::AbstractVector) = L1_norm(x)

@doc (@doc ℓp_norm)
ℓ2_norm(x::AbstractVector) = L2_norm(x)

# 1-dimensional L^p norms

@doc raw"""
    Lp_norm(f, interval::LSHFunctions.RealInterval, p::Real=2)
    L1_norm(f, interval::LSHFunctions.RealInterval)
    L2_norm(f, interval::LSHFunctions.RealInterval)

Computes the ``L^p`` function-space norm of a function ``f``, which is given by the equation

``\|f\|_p = \left(\int_a^b \left|f(x)\right|^p \hspace{0.15cm} dx\right)^{1/p}``

`L1_norm(f, interval)` is the same as `Lp_norm(f, interval, 1)`, and `L2_norm(f, interval)` is the same as `Lp_norm(f, interval, 2)`.

# Examples

```jldoctest; setup = :(using LSHFunctions)
julia> f(x) = x;

julia> interval = @interval(0 ≤ x ≤ 1);

julia> Lp_norm(f, interval, 1) ≈ L1_norm(f, interval) ≈ 2^(-1/1)
true

julia> Lp_norm(f, interval, 2) ≈ L2_norm(f, interval) ≈ 3^(-1/2)
true

julia> Lp_norm(f, interval, 3) ≈ 4^(-1/3)
true
```
"""
Lp_norm(f, interval::LSHFunctions.RealInterval, p::Real=2) = (quadgk(x -> abs(f(x)).^p, interval.lower, interval.upper)[1])^(1/p)

@doc (@doc Lp_norm)
L1_norm(f, interval::LSHFunctions.RealInterval) = quadgk(x -> abs(f(x)), interval.lower, interval.upper)[1]

@doc (@doc Lp_norm)
L2_norm(f, interval::LSHFunctions.RealInterval) = √quadgk(x -> abs2(f(x)), interval.lower, interval.upper)[1]


module JMcDM

# Dependencies
import DataFrames: DataFrame, DataFrameRow
import JuMP: @variable, @objective, @constraint
import JuMP: Model, MOI, optimize!, JuMP, objective_value
import GLPK

# for Pretty printing
# of MCDM results
import Base.show

# includes
include("types.jl")
include("utilities.jl")
include("topsis.jl")
include("vikor.jl")
include("electre.jl")
include("moora.jl")
include("dematel.jl")
include("ahp.jl")
include("nds.jl")
include("singlecriterion.jl")
include("game.jl")
include("dataenvelop.jl")
include("grey.jl")
include("saw.jl")
include("aras.jl")
include("wpm.jl")
include("waspas.jl")
include("edas.jl")
include("marcos.jl")
include("mabac.jl")
include("mairca.jl")
include("copras.jl")
include("promethee.jl")
include("cocoso.jl")
include("critic.jl")
include("entropy.jl")
include("codas.jl")

include("summary.jl")

include("mcdm.jl")
include("print.jl")

include("copeland.jl")

# export imported functions
export DataFrame

# export MCDM methods
export MCDMMethod
export ArasMethod
export CocosoMethod
export CodasMethod
export CoprasMethod
export CriticMethod
export EdasMethod
export ElectreMethod
export GreyMethod
export MabacMethod
export MaircaMethod
export MooraMethod
export PrometheeMethod
export SawMethod
export TopsisMethod
export VikorMethod
export WPMMethod
export WaspasMethod
export MarcosMethod

export MCDMSetting

# export MCDM types
export MCDMResult
export TopsisResult
export VikorResult
export ElectreResult
export MooraResult
export AHPConsistencyResult
export AHPResult
export NDSResult
export GreyResult
export SawResult
export ARASResult
export WPMResult
export WASPASResult
export EDASResult
export MARCOSResult
export MABACResult
export MAIRCAResult
export COPRASResult
export PrometheeResult
export CoCoSoResult
export CRITICResult
export EntropyResult
export CODASResult

# export game type
export GameResult

#  export SCDM types
export SCDMResult
export LaplaceResult
export MaximinResult
export MaximaxResult
export MinimaxResult
export MiniminResult
export SavageResult
export HurwiczResult
export MLEResult
export ExpectedRegretResult
export DataEnvelopResult

# export utility functions
export euclidean
export normalize
export colmaxs
export colmins
export unitize
export makeDecisionMatrix
export reverseminmax
export makeminmax

#  export MCDM tools
export topsis
export vikor
export electre
export moora
export dematel
export ahp_RI, ahp_consistency, ahp
export nds
export grey
export saw
export aras
export wpm
export waspas
export edas
export marcos
export mabac
export mairca
export copras
export promethee, prometLinear, prometVShape, prometUsual, prometQuasi, prometLevel
export cocoso
export critic
export entropy
export codas

#  export SCDM tools
export laplace
export maximin
export maximax
export minimax
export minimin
export savage
export hurwicz
export mle
export expectedregret

# export game solver
export game

# export data envelop
export dataenvelop

#  export summary function
export summary
export mcdm

# export Copeland module
import JMcDM.Copeland: copeland
export copeland

end # module


using DiffEqSensitivity, Flux, Optimization, OptimizationFlux, OptimizationOptimJL, OrdinaryDiffEq, Test

x = Float32[0.8; 0.8]
tspan = (0.0f0,10.0f0)

ann = Chain(Dense(2,10,tanh), Dense(10,1))
p = Float32[-2.0,1.1]
p2,re = Flux.destructure(ann)
_p = [p;p2]
θ = [x;_p]

function dudt2_(u,p,t)
    x, y = u
    [(re(p[3:end])(u)[1]),p[1]*y + p[2]*x]
end

prob = ODEProblem(dudt2_,x,tspan,_p)
solve(prob,Tsit5())

function predict_rd(θ)
  Array(solve(prob,Tsit5(),u0=θ[1:2],p=θ[3:end],abstol=1e-7,reltol=1e-5,sensealg=TrackerAdjoint()))
end
loss_rd(p) = sum(abs2,x-1 for x in predict_rd(p))
l = loss_rd(θ)

cb = function (θ,l)
  @show l
  # display(plot(solve(remake(prob,u0=Flux.data(_x),p=Flux.data(p)),Tsit5(),saveat=0.1),ylim=(0,6)))
  false
end

# Display the ODE with the current parameter values.
cb(θ,l)

loss1 = loss_rd(θ)
optfunc = Optimization.OptimizationFunction((x, p) -> loss_rd(x), Optimization.AutoZygote())
optprob = Optimization.OptimizationProblem(optfunc, θ)
res = Optimization.solve(optprob, BFGS(initial_stepnorm = 0.01), callback = cb)
loss2 = res.minimum
@test 3loss2 < loss1

## Partial Neural Adjoint

u0 = Float32[0.8; 0.8]
tspan = (0.0f0,25.0f0)

ann = Chain(Dense(2,10,tanh), Dense(10,1))

p1,re = Flux.destructure(ann)
p2 = Float32[-2.0,1.1]
p3 = [p1;p2]
θ = [u0;p3]

function dudt_(du,u,p,t)
    x, y = u
    du[1] = re(p[1:41])(u)[1]
    du[2] = p[end-1]*y + p[end]*x
end
prob = ODEProblem(dudt_,u0,tspan,p3)
solve(prob,Tsit5(),abstol=1e-8,reltol=1e-6)

function predict_adjoint(θ)
  Array(solve(prob,Tsit5(),u0=θ[1:2],p=θ[3:end],saveat=0.0:1:25.0))
end
loss_adjoint(θ) = sum(abs2,x-1 for x in predict_adjoint(θ))
l = loss_adjoint(θ)

cb = function (θ,l)
  @show l
  # display(plot(solve(remake(prob,p=Flux.data(p3),u0=Flux.data(u0)),Tsit5(),saveat=0.1),ylim=(0,6)))
  false
end

# Display the ODE with the current parameter values.
cb(θ,l)

loss1 = loss_adjoint(θ)
optfunc = Optimization.OptimizationFunction((x, p) -> loss_adjoint(x), Optimization.AutoZygote())
optprob = Optimization.OptimizationProblem(optfunc, θ)
res1 = Optimization.solve(optprob, ADAM(0.01), callback = cb, maxiters = 100)

optprob = Optimization.OptimizationProblem(optfunc, res1.minimizer)
res = Optimization.solve(optprob, BFGS(initial_stepnorm = 0.01), callback = cb)
loss2 = res.minimum
@test 3loss2 < loss1


@doc raw"""
    gaussian([T=Float64], point; x, y, fwhm, amp=1, theta=0, bkg=0)
    gaussian([T=Float64], px, py; x, y, fwhm, amp=1, theta=0, bkg=0)

An unnormalized bivariate Gaussian distribution. The position can be specified in `(x, y)` coordinates as a `Tuple`, `AbstractVector`, or as separate arguments. If `theta` is given, the PSF will be rotated by `theta` degrees counter-clockwise from the x-axis. If `bkg` is given it will be added as a scalar to the PSF.

The `fwhm` can be a scalar (isotropic) or a vector/tuple (diagonal). Keep in mind that `theta` has no effect for isotropic distributions and is degenerate with the `fwhm` parameters (i.e., theta=90 is the same as reversing the `fwhm` tuple)

# Functional form
```
f(x | x̂, FWHM) = exp[-4ln(2) * ||x - x̂|| / FWHM^2]
```
where `x̂` and `x` are position vectors (indices) `||⋅||` represents the square-distance, and `FWHM` is the full width at half-maximum. If `FWHM` is a scalar, the Gaussian distribution will be isotropic. If `FWHM` is a vector or tuple, the weighting is applied along each axis (diagonal).
"""
gaussian(T, px, py; x, y, fwhm, amp=one(T), theta=0, bkg=0) = convert(T, _gaussian(px, py, x, y, fwhm, amp, theta, bkg))

"""
    normal

An alias for [`gaussian`](@ref)
"""
const normal = gaussian

# this is the factor to convert 1/(2σ²) to 1/(2fwhm²)
const GAUSS_PRE = -4 * log(2)

# isotropic
function _gaussian(px, py, x, y, fwhm, amp, theta, background)
    # find offset from center
    dx = px - x
    dy = py - y
    # rotate
    !iszero(theta) && @warn "isotropic gaussian is not affected by non-zero rotation angle $theta"
    # unnormalized gaussian likelihood
    sqmahab = dx^2 + dy^2
    sqmahab /= fwhm^2
    return amp * exp(GAUSS_PRE * sqmahab) + background
end

# bivariate
function _gaussian(px, py, x, y, fwhm::BivariateLike, amp, theta, background)
    # find offset from center
    dx = px - x
    dy = py - y
    # rotate
    if !iszero(theta)
        dx, dy = rotate_point(dx, dy, theta)
    end
    # unnormalized gaussian likelihood
    fwhmx, fwhmy = fwhm
    sqmahab = (dx / fwhmx)^2 + (dy / fwhmy)^2
    return amp * exp(GAUSS_PRE * sqmahab) + background
end


module YaWorkingData

using LinearAlgebra

export de_μ, directional_σ, ∇directional_σ,
  first_principal_component, pca, transform

const DT{T} = Vector{T} where {T <: Real}
const VDT{T} = Vector{DT{T}} where {T <: Real}

"""
Re-center to have mean 0. across all dim

```julia-repl
julia> de_μ([[1., 3., 2., 4.], [1., 3., 2., 4.]])
2-element Vector{Vector{Float64}}:
 [-1.5, 0.5, -0.5, 1.5]
 [-1.5, 0.5, -0.5, 1.5]
```
"""
function de_μ(data::VDT)::VDT
  """
  Re-center to have mean 0. across all dim
  """
  μ₀ = μ.(data)
  map(((d, μ)=t) -> d .- μ, zip(data, μ₀))
end

"""
Returns the variance of data in the direction of v
"""
function directional_σ(data::VDT, v::DT)::Float64
  dir_v = direction(v)
  sum(dot(u, dir_v) .^ 2 for u ∈ data)
end

function ∇directional_σ(data::VDT, v::DT)::DT
  """
  The gradient of directional variance with respect to v
  """
  v_dir = direction(v)
  sum(map(v -> 2 * dot(v, v_dir), data) .* data)
end

function first_principal_component(data::VDT, n::Integer=100, η=0.1)::DT
  # Start with a random guess
  guess = ones(eltype(data[1]), length(data[1]))
  for _ ∈ 1:n
    dσ = directional_σ(data, guess)
    ∇ = ∇directional_σ(data, guess)
    guess .+= ∇ * η
  end
  direction(guess)
end

function pca(data::VDT, ncomp::Integer)::VDT
  comps = VDT{eltype(data[1])}()
  for _ ∈ 1:ncomp
    comp = first_principal_component(data)
    push!(comps, comp)
    data = rm_project(data, comp)
  end
  comps
end

function transform(data::VDT, comps::VDT)::VDT
  transform_vect.(data, Ref(comps))
end

## Internals
"""
```julia-repl
julia> μ([1., 3., 2., 4.])
2.5
```
"""
function μ(data::DT)::Float64
  sum(data) / length(data)
end

function direction(v::DT)::DT
  v ./ norm(v)
end

function project(u::DT, v::DT)::DT
  """return the projection of u onto the direction v"""
  proj_len = dot(u, v)
  proj_len * v
end

function rm_project_from_vect(u::DT, v::DT)::DT
  """projects u onto v and subtracts the result from u"""
  u - project(u, v)
end

function rm_project(data::VDT, v::DT)::VDT
  rm_project_from_vect.(data, Ref(v))
end

function transform_vect(v::DT, comps::VDT)::DT
  dot.(Ref(v), comps)
end

end  # module

Data Science From Scratch/src/YaWorkingData.jl


using Test, Flux
using DiffEqDevTools
using DiffEqBase, NeuralNetDiffEq

# Homogeneous on scalar
f = (du, u, p, t) -> -du+6*u
tspan = (0.0f0, 1.0f0)
u0 = 1.0f0
du0 = 0.0f0
dt = 1/5f0
f_analytic = t -> [1/5 * exp(-3*t) *(3*exp(5*t)+2.0)]
prob = SecondOrderODEProblem(ODEFunction{false}(f, analytic=f_analytic), u0, du0, tspan)
chain = Chain(Dense(1,5,σ),Dense(5,1))
opt = ADAM(1, (0.9, 0.95))
sol = solve(prob, NeuralNetDiffEq.NNODE2(chain,opt), dt=dt, verbose = true, maxiters = 200)
err = sqrt(sum((sol.u - [f_analytic(t)[1] for t in tspan[1]:dt:tspan[2]]).^2)[1])
@test err< 0.5

#Example 2: Test on vector
f = (du, u, p, t) -> -2*du-5*u+5*t^2+12
tspan = (0.0f0, 1.0f0)
u0 = [1.0f0]
du0 = [0.0f0]
dt = 1/5f0
f_analytic = t -> [1/50 * exp(-t) * (2*exp(t) * (25*t^2 - 20*t + 58) - 13*sin(2*t) - 66*cos(2*t))]
prob = SecondOrderODEProblem(ODEFunction{false}(f, analytic=f_analytic), u0, du0, tspan)
chain = Chain(Dense(1,5,σ),Dense(5,1))
opt = ADAM(0.1, (0.9, 0.95))
sol = solve(prob, NeuralNetDiffEq.NNODE2(chain,opt), dt=dt, verbose = true, maxiters = 300)
err = sqrt(sum((sol.u - [f_analytic(t)[1] for t in tspan[1]:dt:tspan[2]]).^2)[1])
@test err < 0.5

#Example 3: Test on system of 2nd order ODEs
function func(ddu, du, u, p, t)
	ddu[1] = (-p[3]*u[1]-p[4]*(u[1]-u[2]))/p[1]
	ddu[2] = (-p[4]*(u[2]-u[1]))/p[2]
end
function func(du, u, p, t)
	[(-p[3]*u[1]-p[4]*(u[1]-u[2]))/p[1], (-p[4]*(u[2]-u[1]))/p[2]]
end
p = [1.0f0, 1.0f0, 6.0f0, 4.0f0]
u0 = [1.0f0, 0.0f0]
du0 = [2.0f0, 0.0f0]
prob = SecondOrderODEProblem{false}(func, [u0], [du0], (0.0f0, 3.0f0), p)
dt = 1/20f0
opt = ADAM(1e-01)
chain = Chain(Dense(1,5,σ),Dense(5,2))
sol = solve(prob, NeuralNetDiffEq.NNODE2(chain, opt), dt=dt, verbose=true, abstol=1e-10, maxiters=50)


# Cubic interpolation
a = 4.3
b = 1.7
c = 6.2
d = 5.9
cfunc = x -> a*x.^3 .+ b*x.^2 .+ c*x .+ d
cgrad = x -> 3a*x.^2 .+ 2b*x .+ c
chess = x -> 6a*x .+ 2b
xg = Float64[1:8;]
y = cfunc(xg)

# Low-level interface
yc = copy(y)
interp_invert!(yc,BCnan,InterpCubic)
ic = InterpGridCoefs(yc,InterpCubic)

x = 3.7
set_position(ic, BCnil, true, true, [x])
yi = interp(ic,yc)
@test_approx_eq_eps cfunc(x) yi 0.01*(abs(cfunc(x))+abs(yi))
set_gradient_coordinate(ic, 1)
gi = interp(ic, yc)
@test_approx_eq_eps cgrad(x) gi 0.1*(abs(cgrad(x))+abs(yi))
set_hessian_coordinate(ic, 1, 1)
hi = interp(ic, yc)
@test_approx_eq_eps chess(x) hi 0.1*(abs(chess(x))+abs(yi))

# High-level interface
yi = InterpGrid(y, BCnil, InterpCubic)
@test_approx_eq_eps yi[x] cfunc(x) 0.01*(abs(cfunc(x))+abs(yi[x]))
v,g = valgrad(yi,x)
@test_approx_eq_eps v cfunc(x) 0.01*(abs(cfunc(x))+abs(yi[x]))
@test_approx_eq_eps g cgrad(x) 0.01*(abs(cgrad(x))+abs(yi[x]))
v,g,h = valgradhess(yi,x)
@test_approx_eq_eps v cfunc(x) 0.01*(abs(cfunc(x))+abs(yi[x]))
@test_approx_eq_eps g cgrad(x) 0.01*(abs(cgrad(x))+abs(yi[x]))
@test_approx_eq_eps h chess(x) 0.01*(abs(chess(x))+abs(yi[x]))

# Test derivatives with other boundary conditions
func = x -> yi[x]
gnum = derivative_numer(func,x)
hnum = derivative2_numer(func,x)

for BC in (BCnan, BCna)
    yi = InterpGrid(y, BC, InterpCubic)
    @test_throws BoundsError yi[1.1,2.8]  # wrong dimensionality
    func = x -> yi[x]
    v,g = valgrad(yi,x)
    gnum = derivative_numer(func,x)
    @test_approx_eq_eps g gnum 10*cbrt(eps())*(abs(g)+abs(gnum))
    g = zeros(2)
    @test_throws ErrorException valgrad(g, yi, x)
    v,g,h = valgradhess(yi,x)
    gnum = derivative_numer(func,x)
    hnum = derivative2_numer(func,x)
    @test_approx_eq_eps g gnum 10*cbrt(eps())*(abs(g)+abs(gnum))
    @test_approx_eq_eps h hnum 10*cbrt(eps())*(abs(h)+abs(hnum))
    g = zeros(1)
    h = zeros(2,2)
    @test_throws ErrorException valgradhess(g,h,yi,x)
    g = zeros(2)
    h = zeros(1,1)
    @test_throws ErrorException valgradhess(g,h,yi,x)
end

# Test derivatives of 2D interpolation

y = rand(7,8)
yi = InterpGrid(y,BCnan,InterpCubic)
x = [3.4,6.1]
v,g = valgrad(yi,x...)
v2,g2,h2 = valgradhess(yi,x...)
func = (x1,x2) -> yi[x1,x2]

gnum = derivative_numer(func, tuple(x...), 1)
@test_approx_eq_eps g[1] gnum cbrt(eps())*(abs(g[1])+abs(gnum))
@test_approx_eq_eps g2[1] gnum cbrt(eps())*(abs(g2[1])+abs(gnum))

gnum = derivative_numer(func, tuple(x...), 2)
@test_approx_eq_eps g[2] gnum cbrt(eps())*(abs(g[2])+abs(gnum))
@test_approx_eq_eps g2[2] gnum cbrt(eps())*(abs(g2[2])+abs(gnum))

hnum = derivative2_numer(func, tuple(x...), (1, 1))
@test_approx_eq_eps h2[1,1] hnum 10*sqrt(sqrt(eps()))*(abs(h2[1,1])+abs(hnum))
hnum = derivative2_numer(func, tuple(x...), (1, 2))
@test_approx_eq_eps h2[1,2] hnum 10*sqrt(sqrt(eps()))*(abs(h2[1,2])+abs(hnum))
@test_approx_eq_eps h2[2,1] hnum 10*sqrt(sqrt(eps()))*(abs(h2[2,1])+abs(hnum))
hnum = derivative2_numer(func, tuple(x...), (2, 2))
@test_approx_eq_eps h2[2,2] hnum 10*sqrt(sqrt(eps()))*(abs(h2[2,2])+abs(hnum))


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

using Markdown
using InteractiveUtils

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

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

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

# ╔═╡ 1aa41898-3afa-11eb-2f49-570853ce6435
begin
	mu1, sig1, n1 = 0, 2, 8
	mu2, sig2, n2 = 0, 30, 15
	dist1 = Normal(mu1, sig1)
	dist2 = Normal(mu2, sig2)

	N = 10^4
	tdArray = Array{Tuple{Float64,Float64}}(undef,N)

	df(s1,s2,n1,n2) =
		(s1^2/n1 + s2^2/n2)^2 / ( (s1^2/n1)^2/(n1-1) + (s2^2/n2)^2/(n2-1) )
end;

# ╔═╡ fd06acdc-3b1d-11eb-1d5f-03f5cf0f595a
begin
	for i in 1:N
		x1Data = rand(dist1, n1)
		x2Data = rand(dist2, n2)
		x1Bar,x2Bar = mean(x1Data),mean(x2Data)
		s1,s2 = std(x1Data),std(x2Data)
		tStat = (x1Bar - x2Bar) / sqrt(s1^2/n1 + s2^2/n2)
		tdArray[i] = (tStat , df(s1,s2,n1,n2))
	end
	sort!(tdArray, by = first)

	invVal(v,i) = quantile(TDist(v),i/(N+1))

	xCoords  = Array{Float64}(undef,N)
	yCoords1 = Array{Float64}(undef,N)
	yCoords2 = Array{Float64}(undef,N)

	for i in 1:N
		xCoords[i] = first(tdArray[i])
		yCoords1[i] = invVal(last(tdArray[i]), i)
		yCoords2[i] = invVal(n1+n2-2, i)
	end
end;

# ╔═╡ fd06e9ae-3b1d-11eb-1646-d9c440ed4b98
begin
	scatter(xCoords, yCoords1, c=:blue, label="Calculated v", msw=0)
	scatter!(xCoords, yCoords2, c=:red, label="Fixed v", msw=0)
	plot!([-10,10], [-10,10],
		 c=:black, lw=0.3, xlims=(-8,8), ylims=(-8,8), ratio=:equal, label="",
		 xlabel="Theoretical t-distribution quantiles",
		 ylabel="Simulated t-distribution quantiles", legend=:topleft)
end

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

# ╔═╡ Cell order:
# ╟─ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
# ╠═e077ab16-22e5-11eb-15ed-5f0654bf7328
# ╠═e75a9a06-22e5-11eb-2e36-8d4f62b830ed
# ╠═1aa41898-3afa-11eb-2f49-570853ce6435
# ╠═fd06acdc-3b1d-11eb-1d5f-03f5cf0f595a
# ╠═fd06e9ae-3b1d-11eb-1646-d9c440ed4b98
# ╟─475ff888-22e6-11eb-2354-09f8f40a8e12


"""
    Univariate optimization and root finding examples with a few different packages
"""

using Plots
#=
 You should load all packages at beginning of file!
 I just load the optimizers later to make it clear which package is being demonstrated.
=#

# Function definition
# f(x) = x[1]^2 - 0.5
f(x) = x[1]^3 - x[1]
xlim_low = -1.5
xlim_high = 1.5
plot1 = plot(f, xlim_low, xlim_high, label = "f(x)")
plot1 = plot!(x -> 0, xlim_low, xlim_high, label = "")
plot(plot1)

# Optimization
initial_value = -0.01
scatter!([initial_value],
         [0.],
         label = "Initial Value",
         markershape = :x)

## Optim
using Optim

optim_result = Optim.optimize(f, [initial_value], BFGS())
print("Optim minimum: ")
if optim_result.ls_success
    println(optim_result.minimizer)
    scatter!([optim_result.minimizer], [optim_result.minimum], label = "Optim minimum")
else
    println("Failed to converge")
end

# Root finding

## Roots
using Roots
roots_zero = find_zero(f, initial_value)
println("Roots 'find_zero' zero: $(roots_zero)")
scatter!([roots_zero],[f(roots_zero)], label = "Roots 'find_zero' zero")

## Alternative syntax
roots_fzero_zero = fzero(f, initial_value)
println("Roots 'fzero' zero: $(roots_fzero_zero)")
scatter!([roots_fzero_zero],[f(roots_zero)], label = "Roots 'fzero' zero")

## Alternative syntax
fx = ZeroProblem(f, initial_value)
roots_solve_zero = solve(fx)
println("Roots 'solve' zero: $(roots_solve_zero)")
scatter!([roots_solve_zero],[f(roots_zero)], label = "Roots 'solve' zero")

## NLsolve
using NLsolve
f!(F, x) = F[1] = f(x)
nls_zero = nlsolve(f!, [initial_value])
println("NLsolve zero: $(nls_zero.zero)")
scatter!([nls_zero.zero], [f(nls_zero.zero)], label = "NLsolve zero")

## LeastSquaresOptim
### Nonlinear least squares
### minₓ F(x)ᵀF(x)
using LeastSquaresOptim

ls_optim_result = LeastSquaresOptim.optimize(f, [initial_value], Dogleg())
print("LeastSquaresOptim minimum: ")
if ls_optim_result.converged
    println(ls_optim_result.minimizer)
    scatter!([ls_optim_result.minimizer], [f(ls_optim_result.minimizer)], label = "LeastSquaresOptim minimum")
else
    println("Failed to converge")
end


@testset "Interval arithmetic" begin
    a = -1.5 ± 0.5
    b = -1..1
    @test a + b == -3.0..0.0
    @test a * b == -2.0..2.0
    @test a * b + a == -4.0..1.0
    @test a * (b + 1) == -4.0..0.0
end

@testset "Interval matrix arithmetic" begin
    a = 1.0..1.3; b = 2.0..3.5; c = -0.5 ± 0.5; d = 0.0 ± 0.1
    a₊ = 2.0..2.6; b₊ = 4.0..7.0; c₊ = -2.0..0.0; d₊ = -0.2..0.2
    a₋ = -0.3..0.3; b₋ = -1.5..1.5; c₋ = -1.0..1.0; d₋ = -0.2..0.2
    A = IntervalMatrix([a b; c d])

    B = A + A
    @test B isa IntervalMatrix && B == IntervalMatrix([a₊ b₊; c₊ d₊])

    B = A - A
    @test B isa IntervalMatrix && B == IntervalMatrix([a₋ b₋; c₋ d₋])

    # multiply interval and interval matrix
    x = 0.0..2.0
    @test x * A == A * x ==
          IntervalMatrix([0.0..2.6 0.0..7.0; -2.0..0.0 -0.2..0.2])
    # multiply scalar and interval matrix
    x = 1.0
    for A2 in [x * A, A * x, A / x]
        @test A2 == A && typeof(A2) == typeof(A)
    end

    # arithmetic closure using interval matrices and non-interval matrices
    Ainf = inf(A)
    @test A + Ainf isa IntervalMatrix
    @test Ainf + A isa IntervalMatrix
    @test A - Ainf isa IntervalMatrix
    @test Ainf - A isa IntervalMatrix
    @test A * Ainf isa IntervalMatrix
    @test Ainf * A isa IntervalMatrix
end

@testset "Matrix multiplication" begin

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

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

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

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


# # Custom halftoning with SDFs
# Ordered dithering & halftoning algorithms typically use predefined threshold matrices
# with unique values ranging from `1//n` to `(n-1)/n`, where `n` is the number of elements
# in the matrix.
#
# In this example, we will be generating these threshold matrices
# from signed distance functions (SDFs). A good resource for those is Inigo Quilez'
# blogpost ["2D distance functions"](https://www.iquilezles.org/www/articles/distfunctions2d/distfunctions2d.htm).
#
# Here we reimplement his function `sdStar5`:
using Images
using DitherPunk
using ColorSchemes
using TestImages

function star_sdf(x, y; r=0.5, rf=2.0)
    k1 = [0.809016994375, -0.587785252292]
    k2 = [-k1[1], k1[2]]
    p = [abs(x), y]
    p -= 2.0 * maximum([k1 ⋅ p, 0.0]) * k1
    p -= 2.0 * maximum([k2 ⋅ p, 0.0]) * k2
    p = [abs(p[1]), p[2] - r]
    ba = rf * [-k1[2], k1[1]] - [0, 1]
    h = clamp((p ⋅ ba) / (ba ⋅ ba), 0.0, r)
    return norm(p - ba * h) * sign(p[2] * ba[1] - p[1] * ba[2])
end;

# This SDF can be visualized by its contour:
#
# ![](https://user-images.githubusercontent.com/20258504/144755939-19b13e29-e50b-4744-bdeb-234f33790ca8.png)
#
# To convert an SDF to an `OrderedDither` algorithm, use this function:
function sdf2halftone(sdf, n)
    rg = range(-1, 1; length=n)
    A = [sdf(x, y) for y in rg, x in rg]
    p = sortperm(reshape(-A, :))
    B = Vector{Int}(undef, n^2)
    B[p] .= 1:(n^2)
    return OrderedDither(reshape(B, size(A)...)//(n^2 + 1))
end;

# Let's try it on a test image:
img = testimage("fabio_gray_512")
alg = sdf2halftone(star_sdf, 7)
dither(img, alg)

# We can also change the size of the matrix.
# Large halftoning matrices look great on high resolution images.
alg = sdf2halftone(star_sdf, 15)
dither(img, alg)


@defcomp co2forcing begin
    c_CO2concentration = Parameter(index=[time], unit="ppbv")
    f0_CO2baseforcing = Parameter(unit="W/m2", default=1.68)
    fslope_CO2forcingslope = Parameter(unit="W/m2", default=5.5)
    c0_baseCO2conc = Parameter(unit="ppbv", default=400859.5833)
    f_CO2forcing = Variable(index=[time], unit="W/m2")

    function run_timestep(p, v, d, t)

        # eq.13 in Hope 2006
        # the max() condition was added to prevent numerical issues due to  non-positive CO2 concentration levels which occurs in 2300 for low-emission scenarios
        # if the growth effects are paired with GDP-emission feedback
        if p.c_CO2concentration[t] > 0
            v.f_CO2forcing[t] = max(0, p.f0_CO2baseforcing + p.fslope_CO2forcingslope * log(p.c_CO2concentration[t] / p.c0_baseCO2conc) )
        else
            v.f_CO2forcing[t] = 0
        end
    end
end


# K=funLtoK(L)
# The function funLtoK(L) converts a kernel L matrix into a (normalized)
# kernel K matrix. The K matrix has to be semi-positive definite.
#
# Code available here:
# Keeler, 2018, https://github.com/hpaulkeeler/DetPoisson_Julia
#
# #TEST
# B=[3, 2, 1; 4, 5,6; 9, 8,7];
# L=B'*B
# L =
#
#   106    98    90
#    98    93    88
#    90    88    86
# K=funLtoK(L)
# K =
#
#     0.7602    0.3348   -0.0906
#     0.3348    0.3320    0.3293
#    -0.0906    0.3293    0.7492

function funLtoK(L)
    eigenVectLK=eigvecs(L); #eigen decomposition -- vectors
    eigenValL=(eigvals(L)); #eigen decomposition -- values
    eigenValK = eigenValL./(1 .+eigenValL); #eigenvalues of K
    eigenValK=Diagonal(eigenValK); #eigenvalues of L as diagonal matrix
    K=eigenVectLK*eigenValK*(eigenVectLK'); #recombine from eigen components
    K=real(K); #make sure all values are real
    return K;
end