# Jonas August, 2/2022

#=
This problem was asked by Two Sigma.

Using a function rand7() that returns an integer from 1 to 7 (inclusive) with uniform probability, implement a function rand5() that returns an integer from 1 to 5 (inclusive).
=#

rand7() = rand(1:7)

function rand5()
    x = 6
    while x > 5  # Keep sampling uniformly on 1:7 until in range 1:5
        x = rand7()
    end
    x
end

sample = [rand5() for _ in 1:50000]

" Calculate 5 bin histogram. "
function hist5(sample::Vector{Int})
    h = zeros(Int, 5)
    for x in sample
        h[x] += 1
    end
    h
end

# Test uniformity of probability: All bins should have similar counts.
println(hist5(sample))


### A Pluto.jl notebook ###
# v0.11.14

using Markdown
using InteractiveUtils

# ╔═╡ 9e38d21a-f785-11ea-1ce0-5d313a4e9a7e
using Pkg, DrWatson

# ╔═╡ 9e393624-f785-11ea-077b-cb4f64059bb3
begin
	@quickactivate "StatisticalRethinkingTuring"
	using StatisticalRethinking
end

# ╔═╡ 4e27a698-f785-11ea-1ded-a548d1953a14
md"## Fig3.1t.jl"

# ╔═╡ 9e3b92a2-f785-11ea-090f-a78cbded0a8d
begin
	N = 10000
	p_grid = range(0, stop=1, length=N)
	prior = ones(length(p_grid))
	likelihood = pdf.(Binomial.(9, p_grid), 6)
	posterior = likelihood .* prior
	posterior = posterior / sum(posterior)
	samples = sample(p_grid, Weights(posterior), length(p_grid));
end;

# ╔═╡ 9e43adc0-f785-11ea-2b6a-f11004009094
begin
	p1 = scatter(samples, ylim=(0, 1), xlab="Sample number",
	  ylab="Proportion water(p)", leg=false)
	p2 = density(samples, xlim=(0.0, 1.0), ylim=(0.0, 3.0),
	  xlab="Proportion water (p)",
	  ylab="Density", leg=false)
end;

# ╔═╡ d13a98ea-f785-11ea-0cb5-e51b20ddcdac
plot(p1, p2, layout=(1,2))

# ╔═╡ 9e520366-f785-11ea-1d50-8f7ab0c29c9c
md"## End of Fig3.1t.jl"

# ╔═╡ Cell order:
# ╟─4e27a698-f785-11ea-1ded-a548d1953a14
# ╠═9e38d21a-f785-11ea-1ce0-5d313a4e9a7e
# ╠═9e393624-f785-11ea-077b-cb4f64059bb3
# ╠═9e3b92a2-f785-11ea-090f-a78cbded0a8d
# ╠═9e43adc0-f785-11ea-2b6a-f11004009094
# ╠═d13a98ea-f785-11ea-0cb5-e51b20ddcdac
# ╟─9e520366-f785-11ea-1d50-8f7ab0c29c9c


using DataStructures
# author: Younis Shah

# Gates: H, CNOT, S, Pauli gates

b₁ = [0, 1]
b₀ = [1, 0]

# Hadamard gate
# H = [1\√2  1\√2]
#     [1\√2 -1\√2]
H = round(1/sqrt(2) * [1 1; 1 -1], 3)

# Pauli X gate
σₓ = [0 1;
      1 0]

# Pauli Y gate
Y = [0 -im;
     im 0]

# Pauli Z gate
Z = [1 0;
     0 -1]

# Phase gate
# S= [1 0]
#    [0 ι]
S = [1 0;
     0 im]

# Controlled NOT gate
CNOT = [1 0 0 0;
        0 1 0 0;
        0 0 0 1;
        0 0 1 0]

# 2x2 Identity matrix
Ι = eye(Int64, 2)

gate_symbol_list = (:X, :Y, :Z, :H, :S)

gate_list = ("I", "H", "X", "Y", "Z", "CNOT", "S")

# |0⟩ = [1]
#       [0]
pure_state_0 = [1, 0]
#|1⟩ = [0]
#      [1]
pure_state_1 = [0, 1]

# Define a new Qubit type
# Rest -TODO
type Qubit
  index::Int64
  ket::Array
end

# Define a QuantumRegister type
# Rest - TODO
type QuantumRegister
  register::Array{Qubit}
end

# Define a new Stage output type
type StageOutput
  index::Int64
  output::Union(Array, SparseMatrixCSC)
end

# Define a new type to store all the Stage outputs
type StageOutputDeque
  stages::Deque{StageOutput}
end

# Define a new Gate type using which gates are represnted as matrices
type Gate
  qubit::Int64
  name::String
  repr::Array
end

# Define a new Gates type: the number of gates and qubits
# Represents total gates per stage
type Gates
  gates::Array{Gate}
  nqubits::Int64
end


export eigen2x2
function eigen2x2(A)
	Tp = eltype(A)
	a = A[1, 1]; b = A[1, 2]; c = A[2, 1]; d = A[2, 2];
	T = tr(A); D = det(A);

	L₁ = T / 2 + sqrt((T^2 / 4) - D)
	L₂ = T / 2 - sqrt((T^2 / 4) - D)
	values = [L₁, L₂]

	global vectors = []

	if !iszero(c)
		push!(vectors, [L₁ - d, c])
		push!(vectors, [L₂ - d, c])
	elseif !iszero(b)
		push!(vectors, [b, L₁ - a])
		push!(vectors, [b, L₂ - a])
	elseif iszero(b) && iszero(c)
		push!(vectors, [one(Tp), zero(Tp)])
		push!(vectors, [zero(Tp), one(Tp)])
	end

	return values, vectors
end


using LinearAlgebra

function polar_force(force,angle)
    return force.*[sin(deg2rad(angle)), cos(deg2rad(angle))]
end

function moments(forces, locations)
    sum = zeros(3)
    for force_loc in zip(forces, locations)
        force, location = force_loc
        force = vcat(force,[0])
        location = vcat(location,[0])
        res = cross(force, location)
        print(res)
        sum+=res
    end
    sum
end

#forces = [polar_force(718.4, 180 - 30), polar_force(879.54, 45), polar_force(100, -90)]
forces = [[0, -2000], [0, -1200], [0, 15600], [0, -12400]]
location = [[0, 0], [6, 0], [10, 0], [12, 0]]
print(moments(forces, location))


# find expression of projections on cones and their derivatives here:
#   https://stanford.edu/~boyd/papers/pdf/cone_prog_refine.pdf

"""
    projection_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.Zeros) where {T}

projection of vector `v` on zero cone i.e. K = {0}
"""
function projection_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.Zeros) where {T}
    return FillArrays.Zeros{T}(size(v))
end

"""
    projection_on_set(::AbstractDistance, ::MOI.Reals, v::Array{T}) where {T}

projection of vector `v` on real cone i.e. K = R
"""
function projection_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.Reals) where {T}
    return v
end

function projection_on_set(::DefaultDistance, v::T, set::MOI.EqualTo) where {T}
    return zero(T) + set.value
end

function projection_on_set(::DefaultDistance, v::T, set::MOI.LessThan) where {T}
    return min(v, MOI.constant(set))
end

function projection_on_set(::DefaultDistance, v::T, set::MOI.GreaterThan) where {T}
    return max(v, MOI.constant(set))
end

"""
    projection_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.Nonnegatives) where {T}

projection of vector `v` on Nonnegative cone i.e. K = R^n+
"""
function projection_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.Nonnegatives) where {T}
    return max.(v, zero(T))
end

"""
    projection_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.Nonpositives) where {T}

projection of vector `v` on Nonpositive cone i.e. K = R^n-
"""
function projection_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.Nonpositives) where {T}
    return min.(v, zero(T))
end

"""
    projection_on_set(::NormedEpigraphDistance{p}, v::AbstractVector{T}, ::MOI.SecondOrderCone) where {T}

projection of vector `v` on second order cone i.e. K = {(t, x) ∈ R+ × Rn |  ||x|| ≤ t }
"""
function projection_on_set(::NormedEpigraphDistance{p}, v::AbstractVector{T}, ::MOI.SecondOrderCone) where {p, T}
    t = v[1]
    x = v[2:length(v)]
    norm_x = LinearAlgebra.norm(x, p)
    if norm_x <= t
        return copy(v)
    elseif norm_x <= -t
        return zeros(T, size(v))
    end
    result = zeros(T, size(v))
    result[1] = one(T)
    result[2:length(v)] = x / norm_x
    result *= (norm_x + t) / 2
    return result
end

function projection_on_set(::DefaultDistance, v::AbstractVector{T}, cone::MOI.SecondOrderCone) where {T}
    return projection_on_set(NormedEpigraphDistance{2}(), v, cone)
end

"""
    projection_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.PositiveSemidefiniteConeTriangle) where {T}

projection of vector `v` on positive semidefinite cone i.e. K = S^n⨥
"""
function projection_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.PositiveSemidefiniteConeTriangle) where {T}
    dim = isqrt(2*length(v))
    X = unvec_symm(v, dim)
    λ, U = LinearAlgebra.eigen(X)
    D = LinearAlgebra.Diagonal(max.(λ, 0))
    return vec_symm(U * D * U')
end

"""
    unvec_symm(x, dim)

Returns a dim-by-dim symmetric matrix corresponding to `x`.

`x` is a vector of length dim*(dim + 1)/2, corresponding to a symmetric matrix
```
X = [ X11  X12 ... X1k
      X21  X22 ... X2k
      ...
      Xk1  Xk2 ... Xkk ],
```
where `vec(X) = (X11, X12, X22, X13, X23, X33, ..., Xkk)`.

### Note on inner products

Note that the scalar product for the symmetric matrix in its vectorized form is
the sum of the pairwise product of the diagonal entries plus twice the sum of
the pairwise product of the upper diagonal entries; see [p. 634, 1].
Therefore, this transformation breaks inner products:
```
dot(unvec_symm(x, dim), unvec_symm(y, dim)) != dot(x, y).
```

### References

[1] Boyd, S. and Vandenberghe, L.. *Convex optimization*. Cambridge university press, 2004.
"""
function unvec_symm(x, dim=isqrt(2length(x)))
    X = zeros(eltype(x), dim, dim)
    idx = 1
    for i in 1:dim
        for j in 1:i
            X[j,i] = X[i,j] = x[idx]
            idx += 1
        end
    end
    return X
end

"""
    vec_symm(X)

Returns a vectorized representation of a symmetric matrix `X`.
`vec(X) = (X11, X12, X22, X13, X23, X33, ..., Xkk)`

### Note on inner products

Note that the scalar product for the symmetric matrix in its vectorized form is
the sum of the pairwise product of the diagonal entries plus twice the sum of
the pairwise product of the upper diagonal entries; see [p. 634, 1].
Therefore, this transformation breaks inner products:
```
dot(vec_symm(X), vec_symm(Y)) != dot(X, Y).
```

### References

[1] Boyd, S. and Vandenberghe, L.. *Convex optimization*. Cambridge university press, 2004.

"""
function vec_symm(X)
    return X[LinearAlgebra.triu(trues(size(X)))]
end

"""
    projection_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.ExponentialCone) where {T}

projection of vector `v` on closure of the exponential cone
i.e. `cl(Kexp) = {(x,y,z) | y e^(x/y) <= z, y>0 } U {(x,y,z)| x <= 0, y = 0, z >= 0}`.

References:
* [Proximal Algorithms, 6.3.4](https://web.stanford.edu/~boyd/papers/pdf/prox_algs.pdf)
by Neal Parikh and Stephen Boyd.
* [Projection, presolve in MOSEK: exponential, and power cones]
(https://docs.mosek.com/slides/2018/ismp2018/ismp-friberg.pdf) by Henrik Friberg
* [Projection onto the exponential cone: a univariate root-finding problem]
(https://docs.mosek.com/whitepapers/expcone-proj.pdf) by Henrik Friberg
"""
function projection_on_set(::DefaultDistance, v::AbstractVector{T}, s::MOI.ExponentialCone; tol=1e-8) where {T}
    _check_dimension(v, s)

    if _in_exp_cone(v; dual=false)
        return SVector{3}(v)
    end
    if _in_exp_cone(-v; dual=true)
        # if in polar cone Ko = -K*
        return zeros(SVector{3,T})
    end
    if v[1] <= 0 && v[2] <= 0
        return @SVector([v[1], 0, max(v[3],0)])
    end

    return _exp_cone_proj_case_4(v; tol=tol)
end

function _in_exp_cone(v::AbstractVector{T}; dual=false, tol=1e-8) where {T}
    if dual
        return (
            (isapprox(v[1], 0, atol=tol) && v[2] >= 0 && v[3] >= 0) ||
            (v[1] < 0 && v[1]*exp(v[2]/v[1]) + ℯ * v[3] >= tol)
        )
    else
        return (
            (v[1] <= 0 && isapprox(v[2], 0, atol=tol) && v[3] >= 0) ||
            (v[2] > 0 && v[2] * exp(v[1] / v[2]) - v[3] <= tol)
        )
    end
end

function _exp_cone_proj_case_4(v::AbstractVector{T}; tol=1e-8) where {T}
    # Try Heuristic solutions [Friberg 2021, Lemma 5.1]
    # vp = proj onto primal cone, vd = proj onto polar cone
    vp = SVector{3,T}(min(v[1], 0), zero(T), max(v[3], 0))
    vd = SVector{3,T}(zero(T), min(v[2], 0), min(v[3], 0))
    if v[2] > 0
        zp = max(v[3], v[2]*exp(v[1]/v[2]))
        if zp - v[3] < norm(vp - v)
            vp = SVector{3,T}(v[1], v[2], zp)
        end
    end
    if v[1] > 0
        zd = min(v[3], -v[1]*exp(v[2]/v[1] - 1))
        if v[3] - zd < norm(vd - v)
            vd = SVector{3,T}(v[1], v[2], zd)
        end
    end

    # Check if heuristics above approximately satisfy the optimality conditions
    opt_norm = norm(vp + vd - v)
    opt_ortho = abs(dot(vp, vd))
    if norm(v - vp) < tol || norm(v - vd) < tol || (opt_norm < tol && opt_ortho < tol)
        return vp
    end

    # Failure of heuristics -> non heuristic solution
    # Ref: https://docs.mosek.com/slides/2018/ismp2018/ismp-friberg.pdf, p47-48
    # Thm: h(x) is smooth, strictly increasing, and changes sign on domain
    r, s, t = v[1], v[2], v[3]
    h(x) = (((x-1)*r + s) * exp(x) - (r - x*s)*exp(-x))/(x^2 - x + 1) - t

    # Note: won't both be Inf by case 3 of projection
    lb = r > 0 ? 1 - s/r : -Inf
    ub = s > 0 ? r/s : Inf

    # Deal with ±Inf bounds
    if isinf(lb)
        lb = min(ub-0.125, -0.125)
        for _ in 1:10
            h(lb) < 0 && break
            ub = lb
            lb *= 2
        end
    end
    if isinf(ub)
        ub = max(lb+0.125, 0.125)
        for _ in 1:10
            h(ub) > 0 && break
            lb = ub
            ub *= 2
        end
    end

    # Check bounds
    if !(h(lb) < 0 && h(ub) > 0)
        error("Failure to find bracketing interval for exp cone projection.")
    end

    x = _bisection(h, lb, ub)
    if x === nothing
        error("Failure in root-finding for exp cone projection with boundaries ($lb, $ub).")
    end

    return ((x - 1) * r + s)/(x^2 - x + 1) * SVector{3,T}(x, 1, exp(x))
end

"""
    projection_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.DualExponentialCone) where {T}

projection of vector `v` on the dual exponential cone
i.e. `Kexp^* = {(u,v,w) | u < 0, -u*exp(v/u) <= ew } U {(u,v,w)| u == 0, v >= 0, w >= 0}`.

References:
* [Proximal Algorithms, 6.3.4](https://web.stanford.edu/~boyd/papers/pdf/prox_algs.pdf)
by Neal Parikh and Stephen Boyd.
* [Projection, presolve in MOSEK: exponential, and power cones](https://docs.mosek.com/slides/2018/ismp2018/ismp-friberg.pdf)
by Henrik Friberg
"""
function projection_on_set(d::DefaultDistance, v::AbstractVector{T}, ::MOI.DualExponentialCone) where {T}
    p = projection_on_set(d, -v, MOI.ExponentialCone())
    return SVector{3,T}(v[1] + p[1], v[2] + p[2], v[3] + p[3])
end

"""
    projection_on_set(::DefaultDistance, v::AbstractVector{T}, sets::Array{<:MOI.AbstractSet})

Projection onto `sets`, a product of sets
"""
function projection_on_set(::DefaultDistance, v::AbstractVector{T}, sets::Array{<:MOI.AbstractSet}) where {T}
    length(v) == length(sets) || throw(DimensionMismatch("Mismatch between value and set"))
    return reduce(vcat, (projection_on_set(DefaultDistance(), v[i], sets[i]) for i in eachindex(sets)))
end

"""
    projection_gradient_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.Zeros) where {T}

derivative of projection of vector `v` on zero cone i.e. K = {0}^n
"""
function projection_gradient_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.Zeros) where {T}
    return FillArrays.Zeros(length(v), length(v))
end

"""
    projection_gradient_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.Reals) where {T}

derivative of projection of vector `v` on real cone i.e. K = R^n
"""
function projection_gradient_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.Reals) where {T}
    return FillArrays.Eye(length(v))
end

"""
    projection_gradient_on_set(::DefaultDistance, v::T, ::MOI.EqualTo)
"""
function projection_gradient_on_set(::DefaultDistance, ::T, ::MOI.EqualTo) where {T}
    return zero(T)
end

function projection_gradient_on_set(::DefaultDistance, v::T, s::MOI.LessThan) where {T}
    return oneunit(T) * (v <= MOI.constant(s))
end

function projection_gradient_on_set(::DefaultDistance, v::T, s::MOI.GreaterThan) where {T}
    return oneunit(T) * (v >= MOI.constant(s))
end

"""
    projection_gradient_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.Nonnegatives) where {T}

derivative of projection of vector `v` on Nonnegative cone i.e. K = R^n+
"""
function projection_gradient_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.Nonnegatives) where {T}
    y = (sign.(v) .+ one(T))/2
    return LinearAlgebra.Diagonal(y)
end

"""
    projection_gradient_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.Nonpositives) where {T}

derivative of projection of vector `v` on Nonpositives cone i.e. K = R^n-
"""
function projection_gradient_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.Nonpositives) where {T}
    y = @. (-sign(v) + one(T))/2
    return LinearAlgebra.Diagonal(y)
end

"""
    projection_gradient_on_set(::NormedEpigraphDistance{p}, v::AbstractVector{T}, ::MOI.SecondOrderCone) where {T}

derivative of projection of vector `v` on second order cone i.e. K = {(t, x) ∈ R+ × Rn |  ||x|| ≤ t }
"""
function projection_gradient_on_set(::NormedEpigraphDistance{p}, v::AbstractVector{T}, ::MOI.SecondOrderCone) where {p,T}
    n = length(v)
    t = v[1]
    x = v[2:n]
    norm_x = LinearAlgebra.norm(x, p)
    if norm_x <= t
        return Matrix{T}(LinearAlgebra.I,n,n)
    elseif norm_x <= -t
        return zeros(T, n, n)
    end
    result = [
        norm_x     x';
        x          (norm_x + t)*Matrix{T}(LinearAlgebra.I,n-1,n-1) - (t/(norm_x^2))*(x*x')
    ]
    result ./= (2 * norm_x)
    return result
end

function projection_gradient_on_set(::DefaultDistance, v::AbstractVector{T}, cone::MOI.SecondOrderCone) where {T}
    return projection_gradient_on_set(NormedEpigraphDistance{2}(), v, cone)
end

"""
    projection_gradient_on_set(::DefaultDistance, v::AbstractVector{T}, cone::MOI.PositiveSemidefiniteConeTriangle) where {T}

derivative of projection of vector `v` on positive semidefinite cone i.e. K = S^n⨥
"""
function projection_gradient_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.PositiveSemidefiniteConeTriangle) where {T}
    n = length(v)
    dim = isqrt(2n)
    X = unvec_symm(v, dim)
    λ, U = LinearAlgebra.eigen(X)
    Tp = promote_type(T, Float64)

    # if all the eigenvalues are >= 0
    if all(λi ≥ zero(λi) for λi in λ)
        return Matrix{Tp}(LinearAlgebra.I, n, n)
    end

    # k is the number of negative eigenvalues in X minus ONE
    k = count(λi < 1e-4 for λi in λ)

    y = zeros(Tp, n)
    D = zeros(Tp, n, n)

    for idx in 1:n
        # set eigenvector
        y[idx] = 1

        # defining matrix B
        X̃ = unvec_symm(y, dim)
        B = U' * X̃ * U

        for i in 1:size(B)[1] # do the hadamard product
            for j in 1:size(B)[2]
                if (i <= k && j <= k)
                    @inbounds B[i, j] = 0
                elseif (i > k && j <= k)
                    λpi = max(λ[i], zero(Tp))
                    λmj = -min(λ[j], zero(Tp))
                    @inbounds B[i, j] *= λpi / (λmj + λpi)
                elseif (i <= k && j > k)
                    λmi = -min(λ[i], zero(Tp))
                    λpj = max(λ[j], zero(Tp))
                    @inbounds B[i, j] *= λpj / (λmi + λpj)
                end
            end
        end
        @inbounds D[idx, :] = vec_symm(U * B * U')
        # reset eigenvector
        @inbounds y[idx] = 0
    end
    return D
end

"""
    projection_gradient_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.ExponentialCone) where {T}

derivative of projection of vector `v` on closure of the exponential cone,
i.e. `cl(Kexp) = {(x,y,z) | y e^(x/y) <= z, y>0 } U {(x,y,z)| x <= 0, y = 0, z >= 0}`.

References:
* [Solution Refinement at Regular Points of Conic Problems](https://stanford.edu/~boyd/papers/cone_prog_refine.html)
by Enzo Busseti, Walaa M. Moursi, and Stephen Boyd
"""
function projection_gradient_on_set(::DefaultDistance, v::AbstractVector{T}, s::MOI.ExponentialCone) where {T}
    _check_dimension(v, s)

    if _in_exp_cone(v; dual=false)
        return SMatrix{3,3,T}(I)
    end
    if _in_exp_cone(-v; dual=true)
        # if in polar cone Ko = -K*
        return zeros(SMatrix{3,3,T})
    end
    if v[1] <= 0 && v[2] <= 0
        return @SMatrix(T[
            1 0 0
            0 0 0
            0 0 (v[3] >= 0)
        ])
    end

    z1, z2, z3 = _exp_cone_proj_case_4(v)
    nu = z3 - v[3]
    rs = z1/z2
    exp_rs = exp(rs)

    mat = inv(@SMatrix([
        1+nu*exp_rs/z2     -nu*exp_rs*rs/z2       0     exp_rs;
        -nu*exp_rs*rs/z2   1+nu*exp_rs*rs^2/z2    0     (1-rs)*exp_rs;
        0                  0                      1     -1
        exp_rs             (1-rs)*exp_rs          -1    0
    ]))
    return SMatrix{3,3}(@view(mat[1:3,1:3]))
end

"""
    projection_gradient_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.DualExponentialCone) where {T}

derivative of projection of vector `v` on the dual exponential cone,
i.e. `Kexp^* = {(u,v,w) | u < 0, -u*exp(v/u) <= ew } U {(u,v,w)| u == 0, v >= 0, w >= 0}`.

References:
* [Solution Refinement at Regular Points of Conic Problems]
(https://stanford.edu/~boyd/papers/cone_prog_refine.html)
by Enzo Busseti, Walaa M. Moursi, and Stephen Boyd
"""
function projection_gradient_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.DualExponentialCone) where {T}
    # from Moreau decomposition: x = P_K(x) + P_-K*(x)
    return I - projection_gradient_on_set(DefaultDistance(), -v, MOI.ExponentialCone())
end

"""
    projection_gradient_on_set(::DefaultDistance, v::AbstractVector{T}, sets::AbstractVector{<:MOI.AbstractSet})

Derivative of the projection of vector `v` on product of `sets`
projection_gradient_on_set[i,j] = ∂projection_on_set[i] / ∂v[j] where `projection_on_set` denotes projection of `v` on `cone`

Find expression of projections on cones and their derivatives here: https://stanford.edu/~boyd/papers/pdf/cone_prog_refine.pdf
"""
function projection_gradient_on_set(::DefaultDistance, v::AbstractVector{T}, sets::AbstractVector{<:MOI.AbstractSet}) where {T}
    length(v) == length(sets) || throw(DimensionMismatch("Mismatch between value and set"))
    return BlockDiagonal([projection_gradient_on_set(DefaultDistance(), v[i], sets[i]) for i in eachindex(sets)])
end

"""
    projection_on_set(::DefaultDistance, V::AbstractVector{T}, s::NormBallNuclear{T}) where {T}

projection of matrix `V` onto nuclear norm ball
"""
function projection_on_set(d::DefaultDistance, V::AbstractMatrix{T}, s::NormNuclearBall{T}) where {T}
    U, sing_val, Vt = LinearAlgebra.svd(V)
    if (sum(sing_val) <= s.radius)
        return V
    end
    sing_val_proj = projection_on_set(d, sing_val, ProbabilitySimplex(length(sing_val), s.radius))
    return U * Diagonal(sing_val_proj) * Vt'
end

# initial implementation in FrankWolfe.jl
function projection_on_set(::DefaultDistance, v::AbstractVector{T}, s::ProbabilitySimplex{T}) where {T}
    _check_dimension(v, s)
    # TODO: allocating a ton, should implement the recent non-sorting alg
    n = length(v)
    if sum(v) ≈ s.radius && all(>=(0), v)
        return v
    end
    rev = v .- maximum(v)
    u = sort(rev, rev=true)
    cssv = cumsum(u)
    rho = sum(eachindex(u)) do idx
        u[idx] * idx > (cssv[idx] - s.radius)
    end - 1
    theta = (cssv[rho+1] - s.radius) / (rho + 1)
    w = clamp.(rev .- theta, 0.0, Inf)
    return w
end

function projection_on_set(::DefaultDistance, v::AbstractVector{T}, s::StandardSimplex{T}) where {T}
    _check_dimension(v, s)
    n = length(v)
    if sum(v) ≤ s.radius && all(>=(0), v)
        return v
    end
    x = copy(v)
    sum_pos = zero(T)
    for idx in eachindex(x)
        if x[idx] < 0
            x[idx] = 0
        else
            sum_pos += x[idx]
        end
    end
    # at least one positive element
    if sum_pos > 0
        @. x = x / sum_pos * s.radius
    end
    return x
end

function projection_on_set(::DefaultDistance, v::AbstractVector{T}, s::NormInfinityBall{T}) where {T}
    if norm(v, Inf) <= s.radius
        return v
    end
    return clamp.(v, -s.radius, s.radius)
end

function projection_on_set(::DefaultDistance, v::AbstractVector{T}, s::NormTwoBall{T}) where {T}
    nv = norm(v)
    if nv <= s.radius
        return v
    end
    return v .* s.radius ./ nv
end

# inspired by https://github.com/MPF-Optimization-Laboratory/ProjSplx.jl
function projection_on_set(::DefaultDistance, v::AbstractVector{T}, s::NormOneBall{T}) where {T}
    n = length(v)
    if norm(v, 1) ≤ τ
        return v
    end
    u = abs.(v)
    # simplex projection
    bget = false
    s_indices = sortperm(u, rev=true)
    tsum = zero(τ)

    @inbounds for i in 1:n-1
        tsum += u[s_indices[i]]
        tmax = (tsum - τ) / i
        if tmax ≥ u[s_indices[i+1]]
            bget = true
            break
        end
    end
    if !bget
        tmax = (tsum + u[s_indices[n]] - τ) / n
    end

    @inbounds for i in 1:n
        u[i] = max(u[i] - tmax, 0)
        u[i] *= sign(v[i])
    end
    return u
end


"""
    src

A general module name. In this demo, it contains:
2. DiscreteMarkovChain, a type to define a discrete Markov chain
3. stationary_distribution, a function to solve stationary distribution of a `DiscreteMarkovChain`
4. approx_ar1, approximates an AR(1) process with a discrete Markov Chain
5. golden_section, Golden Section search
6. golden_section_improved, improved 3-point Golden Section search
7. DiscreteMarkovChain, discrete state Markov Chain
"""
module src
    import Distributions
    # ----------------------
    export approx_ar1
    export DiscreteMarkovChain, stationary_distribution
    export DiscreteTimeAR1
    export golden_section, golden_section_improved

# ------------------ TYPE: DISCRETE MARKOV CHAIN
"""
    DiscreteMarkovChain( P::Matrix{T} where T <: Real ; w::Vector = Array(1:size(T,1)) )

immutable type for discrete markov chains,
including:
1. K::Int, the number of states
2. w::Vector, state values, can be anything
2. P::Matrix{T} where T <: Real, transition matrix, square matrix, P[i,j], prob from state i to j
"""
mutable struct DiscreteMarkovChain <: Any
    K::Int  # number of states
    w::Vector # state values, can be anything
    P::Matrix{T} where T <: Real  # transition matrix
    function DiscreteMarkovChain( P::Matrix{T} where T <: Real ; w::Vector = Array(1:size(P,1)) )
        (size(P)[1] == size(P)[2])  ?  nothing  :  throw(ErrorException("non-square transition matrix received"))
        all(isapprox.( sum(P,dims = 2), 1.0))  ?  nothing  : throw(ErrorException("row-sum == 1 transition matrix received!"))
        # adjust P, making row-sum EXACT 1
        local newP = copy(P)
        for x in 1:size(P,1)
            newP[x,:] ./= sum(newP[x,:])
        end
        return new( size(P,1), w, newP )
    end
end
# -----
"""
    stationary_distribution( mc::DiscreteMarkovChain )

get the stationary distribution of a given discrete Markov Chain.
returns a `Vector{Float64}`, elements are in the order of `mc.w`.
"""
function stationary_distribution( mc::DiscreteMarkovChain )
    local tmpP = copy(mc.P)
    while true
        (sum(abs.(tmpP .- tmpP^2)) < 1E-8) ? break : nothing
        tmpP ^= 2
    end # while
    return tmpP[1,:]::Vector{Float64}
end # stationary_distribution

abstract type AbstractARProcess <: Any end
# ------------------------------- AR(1) process
"""
    DiscreteTimeAR1( ρ::Real, σ::Real ; Z̄::Real = 0.0, Z0::Union{Real,Distributions.UnivariateDistribution} = 0.0 )

An AR(1) process defined as: `` Z_{t+1} = (1-\\rho)\\cdot \\bar{Z} + \\rho Z_t + \\epsilon_t ``.
where `ρ` is the auto-reg coefficient, `Z̄` is the mean value, `σ` is the std of the error term ``\\epsilon_t``
which follows i.i.d. `N(0,σ)`, and `Z0` is the initial value/distribution of this process.
It can be a constant, or an instance of one of the subtypes of `Distributions.UnivariateDistribution`, e.g. `Normal()`.
"""
struct DiscreteTimeAR1 <: AbstractARProcess
    ρ::Real
    Z̄::Real
    Z0::Union{Real,Distributions.UnivariateDistribution}
    σ::Real
    # ---------------
    function DiscreteTimeAR1( ρ::Real, σ::Real ; Z̄::Real = 0.0, Z0::Union{Real,Distributions.UnivariateDistribution} = 0.0 )
        @assert(isfinite(ρ), "rho must be finite"); @assert(isfinite(Z̄), "the mean value must be finite")
        @assert(0<σ<Inf, "sigma must be non-negative and finite")
        isa(Z0, Real) ?  @assert(isfinite(Z0), "Z0 must be finite") : nothing
        # -----
        return new(ρ,Z̄,Z0,σ)
    end # AR1()
end # AR1

# ------------------ APPROXIMATE AR(1) WITH MARKOV CHAIN
"""
    approx_ar1(ρ::Real, σ::Real ; m::Int = 9, λ::Int = 3, Z̄::Real = 0.0 )

approximates a zero-mean AR(1) process with a discrete finite-state Markov Chain
whose states are equally divided.
The zero-mean AR(1) process is defined as:
`` Z_{t+1} = (1-\\rho)\\cdot \\bar{Z} + \\rho Z_t + \\epsilon_t ``,
where `ρ` is the auto-regression coefficient, `σ` is the standard error of ``\\epsilon_t``,
`m` is the number of the new Markov chain's states, `Z̄` is the mean of `Zt`.
`λ>0` is the absolute maximum deviation of the Markov chain's states,
it is defined as `max(z) = λσ`, where `z` is the state value of the new Markov Chain.

returns an instance of `DiscreteMarkovChain`.

This algorithm is from `(Tauchen, 1986)` and developed by `(Heer & Maussner, 2015)`.
In most cases, `m=9` and `λ=3` are enough to approximate an AR(1).
"""
function approx_ar1(ρ::Real, σ::Real ; m::Int = 9, λ::Real = 3, Z̄::Real = 0.0 )
    # assertions
    abs(ρ) < 1  ?  nothing : @warn("the AR(1) process is not stationary becuase abs(ρ) >=1")
    @assert(0 < σ < Inf, "the standard error must be greater than 0 and finite")
    @assert(1 < m, "there must be at least two states to make a Markov chain")
    @assert(0 < λ < Inf, "the absolute deviation must be greater than 0 and finite")
    @assert(isfinite(Z̄), "the mean value of Zt must be finite")
    # 1. states
    local w::Vector = zeros(m)
    w[1] = -λ * σ / √(1-ρ^2)  # minimum state value (lower bound, finite)
    local steplen::Real = -2 * w[1] / (m-1)  # step length between two states
    for x in 2:m; w[x] = w[x-1] + steplen; end;  # fill
    # 2. transition matrix
    local P::Matrix = zeros(m,m)
    local tmpNorm = Distributions.Normal() # std norm distrib
    for x in 1:m
        P[x,1] = Distributions.cdf(tmpNorm, (w[1]-ρ*w[x])/σ + 0.5 * steplen / σ )
        if m>2
            for y in 2:(m-1)
                P[x,y] = Distributions.cdf(tmpNorm, (w[y]-ρ*w[x])/σ + 0.5 * steplen / σ ) -
                    Distributions.cdf(tmpNorm, (w[y]-ρ*w[x])/σ - 0.5 * steplen / σ )
            end # for
        end # if
        P[x,m] = 1 - sum(P[x,1:(m-1)])
    end # x
    return DiscreteMarkovChain(P,w = w .+ Z̄)::DiscreteMarkovChain # add mean value to state
end # approx_ar1
# -------------------
approx_ar1(ar1::DiscreteTimeAR1 ; m::Int = 9, λ::Real = 3) = begin
    return approx_ar1(ar1.ρ, ar1.σ, m = m, λ = λ, Z̄ = ar1.Z̄)::DiscreteMarkovChain
end # approx_ar1

# ----------------------- Linear interpolation & expolation
"""
    linear_function_interpolation( x0::Real, xsamples::Vector{T} where T <: Real, fsamples::Vector{T} where T <: Real )

linear interpolation on one-parameter function `f(x)` which is defined by a vector of sample `x` points
and a vector of sample `f(x)` corresponding to sample `x`;
`xsamples` should be increasingly sorted.
returns the function value evaluated at the given point `x0`.
if `x0<findmin(x)` or `x0>findmax(x)`, this function uses the linear expolation
of minimum/maximum sample `x` and the second minimum/maximum sample `x`.
"""
function linear_function_interpolation( x0::Real, xsamples::Vector{T} where T <: Real, fsamples::Vector{T} where T <: Real )
    issorted(xsamples) ? nothing : throw(ErrorException("xsamples should be pre-sorted"))
    local N = length(fsamples); local xmax = xsamples[end]; local xmin = xsamples[1]
    local x0_location = (x0 - xmin) / (xmax - xmin) * (N - 1) + 1  # locating x0
    local nearest_xsample_loc = floor(Int,x0_location)  # the nearest sample x point (on the left side) to interpolate x0
    local dist_from_nearest_xsample = x0_location - nearest_xsample_loc # the fractional distance between x0 and the nearest x point
    # NOTE: relationship: x[near] < x0 < x[near+1]

    # case: if touching the upper bound
    if (x0 > xmax) | (nearest_xsample_loc == N)  # to prevent BoundsError at index N
        res = fsamples[end]
        return res::Float64
    # case: if touching the lower bound
    elseif (x0 <= xmin) | (nearest_xsample_loc == 1)
        res = fsamples[1] - (1 - x0_location) * ( fsamples[2] - fsamples[1] )
        return res::Float64
    # case: normal case
    else
        res = (1 - dist_from_nearest_xsample) * fsamples[nearest_xsample_loc] +
            dist_from_nearest_xsample * fsamples[nearest_xsample_loc + 1]
        return res::Float64
    end # if

end # linear_function_interpolation

# ------------------------------ conventional Golden Section for 1-parameter function
"""
    golden_section(f::Function, lb::Real, ub::Real ; atol::Real = 1E-8, maxiter::Int = 5000)

conventional Golden Section search for MINIMIZATION problem.
`f` must only receive one `Real` type location parameter;
`lb` is the finite lower bound of searching; `ub` is the finite upper bound of searching;
`atol` is the absolute error to converge,
and `maxiter` sets the maximum loops to prevent dead loop.
returns a `Tuple` (in order) consisting of optimal x, optimal `f` value, a flag indicating if converged,
and in which round the algorithm ends.

testing code:
```
g(x::Real) = (x-1)^2
BenchmarkTools.@benchmark tmp = src.golden_section(g,-1.0,3.0,atol=1E-12)
```

performance tips:
```
BenchmarkTools.Trial:
  memory estimate:  96 bytes
  allocs estimate:  1
  --------------
  minimum time:     144.933 ns (0.00% GC)
  median time:      150.576 ns (0.00% GC)
  mean time:        181.232 ns (7.65% GC)
  maximum time:     71.892 μs (99.47% GC)
  --------------
  samples:          10000
  evals/sample:     852
```
"""
function golden_section(f::Function, lb::Real, ub::Real ; atol::Real = 1E-8, maxiter::Int = 5000)
    @assert(-Inf < lb < ub < Inf, "lb must be less than ub, both must be finite")
    local goldnum = 0.618033988749895  # the golden number
    local bounds = [lb,ub]  # changable bounds to update
    local left_x = bounds[1] + (1-goldnum) * (bounds[2] - bounds[1])  # left trial
    local right_x = bounds[1] + goldnum * (bounds[2]-bounds[1])  # right trial
    # initial evaluation (NOTE: using type assertion to test if `f` only returns one `Real` value)
    local left_fval::Real = f(left_x)
    local right_fval::Real = f(right_x)
    # test if both trials are defined & finite (because NaN & Inf cannot be operated rationally)
    if isnan(left_fval) | isinf(left_fval) | isnan(left_fval) | isinf(left_fval)
        throw(DomainError("NaN or Inf function value(s) found in the initial evaluation of Golden section search"))
    end # if
    # begin search
    for j in 1:maxiter
        # check convergency
        if abs(bounds[2]-bounds[1]) < atol
            local finalx::Real = (bounds[1] + bounds[2]) / 2
            return ( finalx, f(finalx), true, j )::Tuple
        end # if
        # if not converge, go on to update
        if left_fval > right_fval
            bounds[1] = left_x; bounds[2] = bounds[2]; left_x = right_x
            right_x = bounds[1] + goldnum * (bounds[2] - bounds[1])
            left_fval = f(left_x)
            right_fval = f(right_x)
        else # i.e. if left_fval <= right_fvals
            bounds[1] = bounds[1]; bounds[2] = right_x; right_x = left_x
            left_x = bounds[1] + (1-goldnum) * (bounds[2] - bounds[1])
            left_fval = f(left_x)
            right_fval = f(right_x)
        end # if
    end # for j
    # if loop normally ends, it means that the algorithm did not converge
    local finalx::Real = (left_x+right_x)/2
    return ( finalx, f(finalx), false, maxiter )::Tuple
end # golden_section

# ---------------------------- improved Golden Section search
"""
    golden_section_improved(f::Function, LLB::Float64, LB::Float64, RB::Float64 ;  TOL::Float64=1E-08, ITER=300)

improved Golden Section search, using a third point x_mid between (x_low,x_high) to improve performance.
designed for one-dim function MAXIMIZATION problem

Par:
    1. f [annonymous func]: target to search (a maximization problem), it MUST HAVE ONLY ONE PARAMETER f(x)
    1. LLB [num]: the (initial) very left bound of searching (initial)
    1. LB [num]: the (initial) left bound of searching
    1. RB [num]: the (initial) right bound of searching
    1. TOL [num]: tolerance, 1E-4 or 1E-5 is enough
    1. ITER [int]: maximum iteration times in searching

Ret:
    1. Xmin [num]: the solution

Depend:
    1. func
"""
function golden_section_improved(f::Function, LLB::Real, LB::Real, RB::Real ;  TOL::Float64=1E-08)
    # -------- DATA PROCESS ----------
    # NOTE: in Julia, the golden numebr is an integrated const (before v0.6), but we redefine a golden number for case in v1.0 and later
    GoldenNumber = 1.618033988749895
    r1 = GoldenNumber - 1; r2 = 1 - r1
    x0 = LLB; x3 = RB # initia-lize bounds
    if abs(RB-LB)<=abs(LB-LLB)
        x1 = LB; x2 = LB+r2*(RB-LB)
    else
        x2 = LB; x1 = LB-r2*(LB-LLB)
    end
    # initialization of function value (and turns a maximization to a minimization)
    f1 = - f(x1); f2 = - f(x2)
    # Searching
    while true
        if f2<f1
            x0=x1 # Update the very lower bound
            x1=x2; x2=r1*x1+r2*x3 # new left golden position
            f1=f2; f2= - f(x2)
        else
            x3=x2; x2=x1; x1=r1*x2+r2*x0; f2=f1; f1= - f(x1)
        end # if
        if abs(x3-x0) > TOL * (abs(x1)+abs(x2))
            break
        end # if
    end # while
    # Post-convergence
    Xmin = f1<=f2 ? x1 : x2

    return Xmin::Float64
end

end # src
#


using Test
using DiffEqPhysics, ForwardDiff, OrdinaryDiffEq
using StaticArrays, LinearAlgebra, Random

test_solve(prob...) = mapreduce(p->solve(p, VelocityVerlet(), dt=1//2).u, ==, prob)

H(dθ, θ, p) = dθ / 2 - 9.8 * cos(θ)
p0, q0   = rand(2)
acc      = (v, x, p, t) -> ForwardDiff.derivative(x -> -H(v[1], x, p), x[1])
vel      = (v, x, p, t) -> ForwardDiff.derivative(v -> H(v, x[1], p), v[1])
prob1    = HamiltonianProblem(H, p0, q0, (0., 10.))
prob_1   = DynamicalODEProblem(acc, vel, p0, q0, (0., 10.))
@test test_solve(prob1, prob_1)

p0 = @SVector rand(2)
q0 = @SVector rand(2)
H(dθ, θ, p) = dθ[1] / 2 + dθ[2] / 2 - 9.8 * cos(θ[1]) - 9.8 * cos(θ[2])
prob2    = HamiltonianProblem(H, p0, q0, (0., 10.))
acc      = (v, x, p, t) -> ForwardDiff.gradient(x -> -H(v, x, p), x)
vel      = (v, x, p, t) -> ForwardDiff.gradient(v -> H(v, x, p), v)
prob_2   = DynamicalODEProblem(acc, vel, p0, q0, (0., 10.))
@test test_solve(prob2, prob_2)

H(dθ, θ, p) = (dθ / 2 - 9.8 * cos.(θ))[1]
p0, q0 = [rand(5) for i in 1:2]
acc = (dv, v, x, p, t) -> ForwardDiff.gradient!(dv, x -> -H(v, x, p), x)
vel = (dx, v, x, p, t) -> ForwardDiff.gradient!(dx, v -> H(v, x, p), v)
prob3  = HamiltonianProblem(H, p0, q0, (0., 10.))
prob_3 = DynamicalODEProblem(acc, vel, p0, q0, (0., 10.))
@test test_solve(prob3, prob_3)


using LinearAlgebra

eye(n) = Matrix{Float64}(I,n,n)
eye(n,m) = Matrix{Float64}(I,n,m)

function e(i, dims)
    E = Matrix{Float64}(I, dims, dims)
    return E[:,i]
end


# canonical data version: 1.3.0

using Test

include("complex-numbers.jl")

@test ComplexNumber <: Number

@test ComplexNumber(0, 1)^2 == ComplexNumber(-1, 0)

@testset "Arithmetic" begin
    @testset "Addition" begin
        @test ComplexNumber(1, 0) + ComplexNumber(2, 0) == ComplexNumber(3, 0)
        @test ComplexNumber(0, 1) + ComplexNumber(0, 2) == ComplexNumber(0, 3)
        @test ComplexNumber(1, 2) + ComplexNumber(3, 4) == ComplexNumber(4, 6)
    end

    @testset "Subtraction" begin
        @test ComplexNumber(1, 0) - ComplexNumber(2, 0) == ComplexNumber(-1, 0)
        @test ComplexNumber(0, 1) - ComplexNumber(0, 2) == ComplexNumber(0, -1)
        @test ComplexNumber(1, 2) - ComplexNumber(3, 4) == ComplexNumber(-2, -2)
    end

    @testset "Multiplication" begin
        @test ComplexNumber(1, 0) * ComplexNumber(2, 0) == ComplexNumber(2, 0)
        @test ComplexNumber(0, 1) * ComplexNumber(0, 2) == ComplexNumber(-2, 0)
        @test ComplexNumber(1, 2) * ComplexNumber(3, 4) == ComplexNumber(-5, 10)
    end

    @testset "Division" begin
        @test ComplexNumber(1, 0) / ComplexNumber(2, 0) == ComplexNumber(0.5, 0)
        @test ComplexNumber(0, 1) / ComplexNumber(0, 2) == ComplexNumber(0.5, 0)
        @test ComplexNumber(1, 2) / ComplexNumber(3, 4) ≈ ComplexNumber(0.44, 0.08)
    end
end

@testset "Absolute value" begin
    @test abs(ComplexNumber(5, 0))  == 5
    @test abs(ComplexNumber(-5, 0)) == 5
    @test abs(ComplexNumber(0, 5))  == 5
    @test abs(ComplexNumber(0, -5)) == 5
    @test abs(ComplexNumber(3, 4))  == 5
end

@testset "Complex conjugate" begin
    @test conj(ComplexNumber(5, 0))  == ComplexNumber(5, 0)
    @test conj(ComplexNumber(0, 5))  == ComplexNumber(0, -5)
    @test conj(ComplexNumber(1, 1))  == ComplexNumber(1, -1)
end

@testset "Real part" begin
    @test real(ComplexNumber(1, 0)) == 1
    @test real(ComplexNumber(0, 1)) == 0
    @test real(ComplexNumber(1, 2)) == 1
end

@testset "Imaginary part" begin
    @test imag(ComplexNumber(1, 0)) == 0
    @test imag(ComplexNumber(0, 1)) == 1
    @test imag(ComplexNumber(1, 2)) == 2
end

# Bonus A
@testset "Complex exponential" begin
    @test_skip exp(ComplexNumber(0, π)) ≈ ComplexNumber(-1, 0)
    @test_skip exp(ComplexNumber(0, 0)) == ComplexNumber(1, 0)
    @test_skip exp(ComplexNumber(1, 0)) ≈ ComplexNumber(ℯ, 0)
    @test_skip exp(ComplexNumber(log(2), π)) ≈ ComplexNumber(-2, 0)
end

# Bonus B
@testset "Syntax sugar jm" begin
    @test_skip ComplexNumber(0, 1)  == jm
    @test_skip ComplexNumber(1, 0)  == 1 + 0jm
    @test_skip ComplexNumber(1, 1)  == 1 + 1jm
    @test_skip ComplexNumber(-1, 0) == jm^2
end


#Calculate the data transformations to the output
function y_transform(ys::Array{Float64},ml::ML_model)
    ml.yscale  = mean(ys); ys[:]=ys./ ml.yscale ; ys[:] = log.(ys);     # Scaling and log transform
 #   ml.yscale = std(ys,1); ml.ysmean  =mean(ys,1); ys[:]= (ys.-ml.ysmean)./ml.yscale; # Centering and scaling
    return ys
end

#Undo the data transformations to the output
function y_undo_transform(ys::Array{Float64}, ml::ML_model)
    ys[:] = exp.(ys); ys[:] = ys.*ml.yscale;    # Scaling and log transform
  #  ys[:]= ys.*ml.yscale.+ml.ysmean;     # Centering and scaling
    return ys
end

#Calculate the data transformations to the input x (feature vector)
function X_transform(X::Array{Float64}, ml::ML_model)
    ml.Xmean = mean(X,1);
  #  ml.Xstd = std(X,1);
  #  ind = (0.==ml.Xstd); ml.Xstd[ind] =1.0;  #replace 0 in std by 1 incase there is a constant feature
  #  X= (X.-ml.Xmean)./ml.Xstd; # Centering and scaling
    X= X.-ml.Xmean;
    sX = size(X);
    X = [X ones(sX[1],1)]; # Adding the bias term
    return X
end

#Apply the saved data transformations to a single data vector x (feature vector)
function x_apply_transform(x::Array{Float64}, ml::ML_model)
  #  x= (x.-ml.Xmean)./ml.Xstd; # Centering and scaling
    x =x.-ml.Xmean;
    x = [x 1];# Adding the bias term
   return x
end


# Solves the scalar, purely nonlinear, eigenvalue problem
# A(λ) = 0.2*sqrt(λ) - 0.6*sin(2*λ) with both a polynomial
# and fully rational approach

using NonlinearEigenproblemsTest
using NonlinearEigenproblems
using Test

function nleigs_scalar()
    as_matrix(x::Number) = (M = Matrix{eltype(x)}(undef,1,1); M[1] = x; M)
    n = 1
    C = [as_matrix(0.2), as_matrix(-0.6)]
    f = [λ -> sqrt(λ), λ -> sin(2*λ)]
    nep = SPMF_NEP([C[1], C[2]], [f[1], f[2]])

    Σ = complex([0.01, 4])

    @bench @testset "Polynomial" begin
        lambda, X = nleigs(nep, Σ, logger=displaylevel, maxit=100, v=ones(n).+0im, leja=2, isfunm=false)

        # single eigenvalue converges
        verify_lambdas(1, nep, lambda, X)
    end

    @bench @testset "Fully rational" begin
        # set of poles candidates
        Ξ = -10 .^ range(-6, stop = 5, length = 10000)

        lambda, X = nleigs(nep, Σ, Ξ=Ξ, logger=displaylevel, maxit=100, v=ones(n).+0im, leja=2, isfunm=false)

        # three eigenvalues converge
        verify_lambdas(3, nep, lambda, X)
    end
end

@testset "NLEIGS: Scalar" begin
    nleigs_scalar()
end


using SMM
using Optim
using HTTP
using GLM
using LinearAlgebra
using Random
using Statistics
using DataFrames
using DataFramesMeta
using CSV

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# question 1
#:::::::::::::::::::::::::::::::::::::::::::::::::::
url = "https://raw.githubusercontent.com/OU-PhD-Econometrics/fall-2020/master/ProblemSets/PS1-julia-intro/nlsw88.csv"
df = CSV.read(HTTP.get(url).body,DataFrame)
X = [ones(size(df,1),1) df.age df.race.==1 df.collgrad.==1]
y = df.married.==1

#GMM
function ols_gmm(b, X, y)
    g = y .- X*b
    J = g'*I*g
    return J
end
res =  optimize(b -> ols_gmm(b, X, y), rand(size(X,2)), LBFGS(), Optim.Options(g_tol=1e-6, iterations=100_000, show_trace=true))
println(res.minimizer)

#closed form solution
res1 = inv(X'*X)*X'*y
println(res1)

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# question 2
#:::::::::::::::::::::::::::::::::::::::::::::::::::
df = dropmissing(df, :occupation)
df[df.occupation.==8 ,:occupation] .= 7
df[df.occupation.==9 ,:occupation] .= 7
df[df.occupation.==10,:occupation] .= 7
df[df.occupation.==11,:occupation] .= 7
df[df.occupation.==12,:occupation] .= 7
df[df.occupation.==13,:occupation] .= 7
X = [ones(size(df,1),1) df.age df.race.==1 df.collgrad.==1]
y = df.occupation

#MlE
function mlogit(alpha, X, y)

        K = size(X,2)
        J = length(unique(y))
        N = length(y)
        bigY = zeros(N,J)
        for j=1:J
            bigY[:,j] = y.==j
        end
        bigAlpha = [reshape(alpha,K,J-1) zeros(K)]

        num = zeros(N,J)
        dem = zeros(N)
        for j=1:J
            num[:,j] = exp.(X*bigAlpha[:,j])
            dem .+= num[:,j]
        end

        P = num./repeat(dem,1,J)

        loglike = -sum( bigY.*log.(P) )

        return loglike
    end

alpha_zero = zeros(6*size(X,2))
alpha_rand = rand(6*size(X,2))
alpha_true = [.1910213,-.0335262,.5963968,.4165052,-.1698368,-.0359784,1.30684,-.430997,.6894727,-.0104578,.5231634,-1.492475,-2.26748,-.0053001,1.391402,-.9849661,-1.398468,-.0142969,-.0176531,-1.495123,.2454891,-.0067267,-.5382892,-3.78975]
alpha_start = alpha_true.*rand(size(alpha_true))
#println(size(alpha_true))
alpha_hat_optim = optimize(a -> mlogit(a, X, y), alpha_start, LBFGS(), Optim.Options(g_tol = 1e-5, iterations=100_000, show_trace=true, show_every=50))
alpha_hat_mle = alpha_hat_optim.minimizer
println(alpha_hat_mle)

#GMM
function mlogit_gmm(alpha, X, y)
    K = size(X,2)
    J = length(unique(y))
    N = length(y)
    bigY = zeros(N,J)
    for j=1:J
        bigY[:,j] = y.==j
    end
    bigAlpha = [reshape(alpha,K,J-1) zeros(K)]

    num = zeros(N,J)
    dem = zeros(N)
    for j=1:J
        num[:,j] = exp.(X*bigAlpha[:,j])
        dem .+= num[:,j]
    end

    P = num./repeat(dem,1,J)

    g = bigY .-P
    J = sum(g'*I*g)
    return J
end

alpha_rand = rand(6*size(X,2))
alpha_hat_optim1 = optimize(a -> mlogit(a, X, y), alpha_start, LBFGS(), Optim.Options(g_tol = 1e-5, iterations=100_000, show_trace=true, show_every=50))
alpha_hat_gmm = alpha_hat_optim1.minimizer
println(alpha_hat_gmm)

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# question 3
#:::::::::::::::::::::::::::::::::::::::::::::::::::
function mlogit_sim(alpha,X,K,J,N)
    bigAlpha = [reshape(alpha,K,J-1) zeros(K)]
    num = zeros(N,J)
    dem = zeros(N)
    for j=1:J
        num[:,j] = exp.(X*bigAlpha[:,j])
        dem .+= num[:,j]
    end

    P = num./repeat(dem,1,J)
    epsilon = rand(N)

    Y = zeros(N)
    for i =1:N
        temp = zeros(J)
        for j = 1:J
            temp[j] = convert(Int,(sum(P[i,j:end])>epsilon[i]))
        end
        Y[i] = sum(temp)
    end
    return Y
end

K = 3
J = 7
N = 10000
X = randn((N,K))
alpha_true = rand((J-1)*K)
println("beta true:",alpha_true)
y_sim = mlogit_sim(alpha_true,X,K,J,N)
alpha_start = alpha_true.*rand(size(alpha_true))
alpha_hat_optim2 = optimize(a -> mlogit(a, X, y_sim), alpha_start, LBFGS(), Optim.Options(g_tol = 1e-5, iterations=100_000, show_trace=true, show_every=50))
alpha_hat_sim = alpha_hat_optim2.minimizer
println("beta estimated:",alpha_hat_sim)

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# question 4
#:::::::::::::::::::::::::::::::::::::::::::::::::::
MA = SMM.parallelNormal()
dc = SMM.history(MA.chains[1])
dc = dc[dc[:accepted].==true,:]
println(describe(dc))

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# question 5
#:::::::::::::::::::::::::::::::::::::::::::::::::::
function mlogit_smm(alpha, X, y, D)
    K = size(X,2)
    J = length(unique(y))
    N = length(y)

   # N+1 moments in both model and data
    gmodel = zeros(N+1,D)
    # data moments are just the y vector itself
    # and the variance of the y vector
    gdata  = vcat(y,var(y))

    Random.seed!(1234)

    # simulated model moments

    for d = 1:D
        y_sim = mlogit_sim(alpha,X,K,J,N)
        gmodel[1:end-1,d] = y_sim
        gmodel[  end  ,d] = var(y_sim)
    end
    # criterion function
    err = vec(gdata .- mean(gmodel; dims=2))
    # weighting matrix is the identity matrix
    # minimize weighted difference between data and moments
    J = err'*I*err

    return J
end

alpha_true = [.1910213,-.0335262,.5963968,.4165052,-.1698368,-.0359784,1.30684,-.430997,.6894727,-.0104578,.5231634,-1.492475,-2.26748,-.0053001,1.391402,-.9849661,-1.398468,-.0142969,-.0176531,-1.495123,.2454891,-.0067267,-.5382892,-3.78975]
alpha_start = alpha_true.*rand(size(alpha_true))
D = 2000
#println(size(alpha_true))
alpha_hat_optim = optimize(a -> mlogit_smm(a, X, y,D), alpha_start, LBFGS(), Optim.Options(g_tol = 1e-5, iterations=100_000, show_trace=true, show_every=50))
alpha_hat_smm = alpha_hat_optim.minimizer
println(alpha_hat_smm)


const PointTypes = Union{VoronoiDelaunay.AbstractPoint2D, GeometryBasics.Point2}

function mean(pts::Vector{T}) where T <: PointTypes
	ax = 0.0
	ay = 0.0

	@simd for p in pts
		ax += getx(p)
		ay += gety(p)
	end

	Np = length(pts)
	T(ax/Np, ay/Np)
end

for name in [:sort, :sort!, :issorted, :sortperm]
	@eval begin
		function Base.$name(pts::Vector{T}) where T <: PointTypes
			center = mean(pts)
			$name(pts, by = p -> p - center, rev = true)
		end
	end
end

# http://stackoverflow.com/questions/6989100/sort-points-in-clockwise-order
function Base.isless(p::T, q::T) where T <: PointTypes
	if getx(p) >= 0.0 && getx(q) < 0.0
		return true
	elseif getx(p) < 0.0 && getx(q) >= 0.0
		return false
	elseif getx(p) == getx(q) == 0.0
		if gety(p) >= 0.0 || gety(q) >= 0.0
			return gety(p) > gety(q)
		else
			return gety(p) < gety(q)
		end
	end

	det = getx(p)*gety(q) - getx(q)*gety(p)
	if det < 0.0
		return true
	elseif det > 0.0
		return false
	end

	# p and q are on the same line from the center; check which one is
	# closer to the origin
	origin = T(0.0, 0.0)
	abs2(p, origin) > abs2(q, origin)
end


"""
    hermite(x, p::Integer)

Evaluate the Hermite polynomial of degree `p` at `x` using the three term recursion.
"""
function hermite(x, p::Integer)
    T = typeof( 2x )
    p₀ = one(T)
    p₁ = 2x

    if p <= 0
        return p₀
    elseif p == 1
        return p₁
    end

    for n in 2:p
        p₀, p₁ = p₁, 2x*p₁ - 2*(n-1)*p₀
    end

    p₁
end


# Hello World
# variable substitution with the $ operator
word = "world"
output = "Hello $word !"
println(output)

# tab autocomplete of LateX characters
ℵ = "world"
output = "Hello $ℵ, again!"
println(output)

# this extends to function evaluations as well
α = 2
println("The sine of $α is $(sin(α))")

# sidenote:
# so if printing to a file you can directly direference the evalutioan
# without passing it to an array which can then be printed.

typeof(α)
typeof(sin)
methods(sin)

# built in type promotion
check = int8(10)
typemax(Int8)
typeof(check)
check = check *1232
typeof(check)

# defining numerical functions
# looks and feels very much like real math

f(x) = 2*x + x^2 + 5
f(3.4)

g(x, y) = 34 + 2*x^2 + y^3
g(2, 3)

# passing function results to arrays
arr = Float64[]
for x in 1.3:0.01:3.9 push!(arr, f(x)) end
println(arr)
for x in 1.3:0.01:9.9 push!(arr, f(x)) end
println(arr)

arr2 = [f(i) for i in [1:2:10, 15]]
println(arr2)

using ASCIIPlots
lineplot([1.3:0.01:3.9], [1.3:0.01:3.9])
lineplot([1:2:10], [f(i) for i in [1:2:10]])

# working with imaginary numbers
a = 23
b = (1 + 34im)
c = a+b

typeof(a)
typeof(b)
typeof(c)

# now c can be dereferenced into real and imaginary parts
d = c.re
e = c.im

typeof(d)
typeof(e)

d = c*conj(c)
typeof(d)


"""
```
find_threshold(Entropy(),counts,edges)
```
An algorithm for finding the threshold value for a gray-level histogram using
the entropy of the histogram.

# Output

Returns the point in the `AbstractRange` which corresponds to the threshold bin in the histogram.

# Details

This algorithm uses the entropy of a gray level histogram to produce a threshold
value.

Let ``f_1, f_2, \\ldots, f_I`` be the frequencies in the various bins of the
histogram and ``I`` the number of bins. With ``N = \\sum_{i=1}^{I}f_i``, let
``p_i = \\frac{f_i}{N}`` (``i = 1, \\ldots, I``) denote the probability
distribution of gray levels. From this distribution one derives two additional
distributions. The first defined for discrete values ``1`` to ``s`` and the
other, from ``s+1`` to ``I``. These distributions are

```math
A: \\frac{p_1}{P_s}, \\frac{p_2}{P_s}, \\ldots, \\frac{p_s}{P_s}
\\quad \\text{and} \\quad
B: \\frac{p_{s+1}}{1-P_s}, \\ldots, \\frac{p_n}{1-P_s}
\\quad \\text{where} \\quad
P_s = \\sum_{i=1}^{s}p_i.
```
The entropies associated with each distribution are as follows:

```math
H(A) = \\ln(P_s) + \\frac{H_s}{P_s}
```
```math
H(B) = \\ln(1-P_s) + \\frac{H_n-H_s}{1-P_s}
```
```math
\\quad \\text{where} \\quad
H_s = -\\sum_{i=1}^{s}p_i\\ln{p_i}
\\quad \\text{and} \\quad
H_n = -\\sum_{i=1}^{I}p_i\\ln{p_i}.
```

Combining these two entropy functions we have

```math
\\psi(s) = \\ln(P_s(1-P_s)) + \\frac{H_s}{P_s} + \\frac{H_n-H_s}{1-P_s}.
```
Finding the discrete value ``s`` which maximises the function ``\\psi(s)`` produces
the sought-after threshold value (i.e. the bin which determines the threshold).

See Section 4 of [1] for more details on the derivation of the entropy.

# Options

## Choices for `counts`
You can specify an `AbstractArray` which should be a 1D array of frequencies for
a histogram. You should submit the corresponding `edges` range for the bins of
the histogram. The function will throw an error if it detects that the `edges`
and `counts` have different lengths.

## Choices for `edges`
You can specify an `AbstractRange` which should be the corresponding range for
the bins of the histogram array passed into `counts`.

# Example

```julia

using TestImages, Images

img = testimage("cameraman")
# building a histogram with 256 bins
edges, counts = build_histogram(img, 256)
#=
  The `counts` array stores at index 0 the frequencies that were below the
  first bin edge. Since we are seeking a threshold over the interval
  partitioned by `edges` we need to discard the first bin in `counts`
  so that the dimensions of `edges` and `counts` match.
=#
find_threshold(EntropyThresholding(), counts[1:end], edges)
```

# References
[1] J. N. Kapur, P. K. Sahoo, and A. K. C. Wong, “A new method for gray-level picture thresholding using the entropy of the histogram,” *Computer Vision, Graphics, and Image Processing*, vol. 29, no. 1, p. 140, Jan. 1985.[doi:10.1016/s0734-189x(85)90156-2](https://doi.org/10.1016/s0734-189x%2885%2990156-2)
"""
function find_threshold(algorithm::Entropy, counts::AbstractArray, edges::AbstractRange)
    if length(edges) != length(counts)
        error("the lengths of edges and counts must match")
    end
    nbins = length(counts)
    pdf = counts ./ sum(counts)
    Ψₘₐₓ = 0
    Hₛ = H(pdf)
    Hₙ = -1*Hₛ[nbins]
    Pₛ = cumsum(pdf)
    threshold_bin = 1
    for i in 1:nbins
        Pᵢ = Pₛ[i]
        Hᵢ = -1*Hₛ[i]
        Ψᵢ = log(Pᵢ*(1-Pᵢ)) + (Hᵢ/Pᵢ) + ((Hₙ-Hᵢ)/(1-Pᵢ))
        if Ψᵢ > Ψₘₐₓ
            Ψₘₐₓ = Ψᵢ
            threshold_bin = i
        end
    end
    edges[threshold_bin]
end

function H(pdf::AbstractArray)
    h_dist = zeros(length(pdf))
    for i in 1:length(pdf)
        h_dist[i] = (pdf[i] != 0) ? pdf[i]*log(pdf[i]) : 0.0
    end
    cumsum(h_dist)
end


using LassoPath
using Test
import UnicodePlots as UnicodePlots
using Plots
using LinearAlgebra

function Test1()
    """
        Test 1, some basic test on a random matrix, copied columns of matrix.
    """
    deg = 3
    N = 2000
    B = rand(N, deg)
    A = fill(0.0, (N, 2deg))
    A[:, 1:deg] = B
    A[:, deg + 1:end] = 2B.^2 + B + 0.1.*randn(size(B))

    b = A*ones(2deg, 1)
    instance = LassoSCOP(A, b)
    results, λs = GetLassoPath(instance)
    display(results)
    display(λs)

    # Let's plot this, and a global variable will be defined.
    global Test1Results = results

    ylim = (minimum(Test1Results), maximum(Test1Results))
    ThePlot = UnicodePlots.lineplot(log2.(λs), Test1Results[1, :], ylim=ylim, title="Lasso Path")

    for II ∈ 2:size(Test1Results, 1)
        UnicodePlots.lineplot!(ThePlot, log2.(λs), Test1Results[II, :])
    end
    display(ThePlot)
    ThePlot = UnicodePlots.heatmap(results[end:-1:begin, :])
    display(ThePlot)

    # Visualize the path.
    VisualizeLassoPath(instance)

    # Finding the important weights for it.
    PrintTitle("These are the important weights: ")
    WeightsIndices = CaptureImportantWeights(instance)
    println(WeightsIndices)
    return true

end

function Test2()
    """
        Testing the iterator for the lasso path object.
    """
    PrintTitle("Testing the Iterator for the Lasso Path type. ")
    deg = 3
    N = 2000
    B = rand(N, deg)
    A = fill(0.0, (N, 2deg))
    A[:, 1:deg] = B
    A[:, deg + 1:end] = 2B.^2 + B + 0.1.*randn(size(B))

    # Use this to test the lasso iterate.
    let A = A, b = A*ones(2deg, 1)
        l = LassoSCOP(A, b )
        for (λ, x) ∈ l
            println(λ)
            println(x)
        end
    end

    return true
end

function Test3()
    """
        Test proximal Gradient, internal methods
    """
    PrintTitle("Testing the proximal operator with simple regression. ")
    N, M = 10, 100
    A = randn(N, M)
    b = randn(N, 1)
    instance = LassoPath.BuildPG2NormLasso(A, b, 1/M)
    instance.tol= 1e-4

    sol = LassoPath.OptimizeProximalGradient(instance, A\b)
    error = norm(A*sol - b)
    println("The error of the solution is: $(error)")
    println("Non-zero entries: $(sum(sol .!= 0))")

    return error != Inf

end

function Test4()
    """
        Testing the proximal lasso path method.

    """
    """
    Test 1, some basic test on a random matrix, copied columns of matrix.
"""
    deg = 3
    N = 2000
    B = rand(N, deg)
    A = fill(0.0, (N, 2deg))
    A[:, 1:deg] = B
    A[:, deg + 1:end] = 2B.^2 + B + 0.1.*randn(size(B))

    b = A*ones(2deg, 1)
    instance = Lasso2NormProx(A, b)

    results, λs = GetLassoPath(instance)
    display(results)
    display(λs)

    # Let's plot this, and a global variable will be defined.
    global Test1Results = results

    ylim = (minimum(Test1Results), maximum(Test1Results))
    ThePlot = UnicodePlots.lineplot(log2.(λs), Test1Results[1, :], ylim=ylim, title="Lasso Path")

    for II ∈ 2:size(Test1Results, 1)
        UnicodePlots.lineplot!(ThePlot, log2.(λs), Test1Results[II, :])
    end
    display(ThePlot)
    ThePlot = UnicodePlots.heatmap(results[end:-1:begin, :])
    display(ThePlot)

    # Visualize the path.
    VisualizeLassoPath(instance)

    # Finding the important weights for it.
    PrintTitle("These are the important weights: ")
    WeightsIndices = CaptureImportantWeights(instance)
    println(WeightsIndices)

    return true
end

function Test5()
    """
        Test High dimensional Array for the proximal gradient lasso method.
    """

end

# A set of basic tests for Lasso.
@testset "LassoPath.jl" begin
    PrintTitle("Test Set 1")
    @test Test1()
    @test Test2()
    @test Test3()
    @test Test4()
end


"""
    local_clustering_coefficient(g, v)
    local_clustering_coefficient(g, vs)

Return the [local clustering coefficient](https://en.wikipedia.org/wiki/Clustering_coefficient)
for node `v` in graph `g`. If a list of vertices `vs` is specified, return a vector
of coefficients for each node in the list.
"""
function local_clustering_coefficient(g::AbstractGraph, v::Integer)
    ntriang, alltriang = local_clustering(g, v)

    return alltriang == 0 ? 0. : ntriang * 1.0 / alltriang
end
local_clustering_coefficient(g::AbstractGraph, vs = vertices(g)) =
    [local_clustering_coefficient(g, v) for v in vs]

@doc_str """
    local_clustering(g, v)
    local_clustering(g, vs)

Return a tuple `(a, b)`, where `a` is the number of triangles in the neighborhood
of `v` and `b` is the maximum number of possible triangles. If a list of vertices
`vs` is specified, return two vectors representing the number of triangles and
the maximum number of possible triangles, respectively, for each node in the list.

This function is related to the local clustering coefficient `r` by ``r=\frac{a}{b}``.
"""
function local_clustering(g::AbstractGraph, v::Integer)
    k = degree(g, v)
    k <= 1 && return (0, 0)
    neighs = neighbors(g, v)
    c = 0
    for i in neighs, j in neighs
        i == j && continue
        if has_edge(g, i, j)
            c += 1
        end
    end
    return is_directed(g) ? (c , k*(k-1)) : (div(c,2) , div(k*(k-1),2))
end
function local_clustering(g::AbstractGraph, vs = vertices(g))
    ntriang = zeros(Int, length(vs))
    nalltriang = zeros(Int, length(vs))
    i = 0
    for (i, v) in enumerate(vs)
        ntriang[i], nalltriang[i] = local_clustering(g, v)
    end
    return ntriang, nalltriang
end

"""
    triangles(g[, v])
    triangles(g, vs)

Return the number of triangles in the neighborhood of node `v` in graph `g`.
If a list of vertices `vs` is specified, return a vector of number of triangles
for each node in the list. If no vertices are specified, return the number
of triangles for each node in the graph.
"""
triangles(g::AbstractGraph, v::Integer) = local_clustering(g, v)[1]
triangles(g::AbstractGraph, vs = vertices(g)) = local_clustering(g, vs)[1]

"""
    global_clustering_coefficient(g)

Return the [global clustering coefficient](https://en.wikipedia.org/wiki/Clustering_coefficient)
of graph `g`.
"""
function global_clustering_coefficient(g::AbstractGraph)
    c = 0
    ntriangles = 0
    for v in vertices(g)
        neighs = neighbors(g, v)
        for i in neighs, j in neighs
            i == j && continue
            if has_edge(g, i, j)
                c += 1
            end
        end
        k = degree(g, v)
        ntriangles += k*(k-1)
    end
    ntriangles == 0 && return 1.
    return c / ntriangles
end


function B(; T=T, λ=λ)

    k_B = 1.380649e-23
    h   = 6.62607015e-34
    c   = 2.99792458e8

    aux1 = (2.0 * h * c^2) / λ^5

    aux2 = (h * c) / (λ * k_B * T)

    aux1 / (exp(aux2) - 1.0)

end

function scaledB(; s=s, T=T, λ=λ)

    s * B(T=T, λ=λ)

end

function fitgalaxy()

    wave = [433,550,790] * 10^-9

    Y = [4.18, 8.60, 20.83]

    # wave = [0.5;4;6] * 1000.0 * 1e-9
    # Y = 1e-8*[B(T=5555.5,λ = 500*1e-9); B(T=5555.5,λ = 4000*1e-9);B(T=5555.5,λ = 6000*1e-9)]

    function loss(T)

        local V = map( w -> B(T=T,λ=w), wave)

        local s = dot(V,Y) / dot(V,V)

        local l = 0.0

        for i in 1:length(Y)
            l += (Y[i] - s*V[i])^2
        end

        return l

    end

    result = optimize(loss, 10, 10_000)

    @show T = result.minimizer

    V = map( w -> B(T=T, λ=w), wave)

    s = dot(V,Y) / dot(V,V)

    λ = collect(LinRange(100, 10_000, 100)) * 1e-9;

    plot(λ, map(λ -> s*B(T=T,λ=λ), λ), "g.")

    plot(wave, Y, "ko")

end

function galaxyposterior()

    Tmin = 1000.0
    Tmax = 10_000.0
    σ(x) = 1.0 / (exp(-x) + 1)

    wave = [433,550,790] * 10^-9

    Y = [4.18, 8.60, 20.83]

    # wave = [0.5;4;6] * 1000.0 * 1e-9
    # Y = 1e-8*[B(T=5555.5,λ = 500*1e-9); B(T=5555.5,λ = 4000*1e-9);B(T=5555.5,λ = 6000*1e-9)]

    function loss(T)

        local V = map( w -> B(T=T,λ=w), wave)

        local s = dot(V,Y) / dot(V,V)

        local l = 0.0

        for i in 1:length(Y)
            l += (Y[i] - s*V[i])^2
        end

        return l

    end

    Trange = 1000:20:8_000

    negloss = map(x -> -loss(x), Trange)

    plot(Trange, negloss, ".")

    Trange[argmax(negloss)], maximum(negloss)

end

# λ = collect(LinRange(10,10_000,1000)) * 1e-9;
# plot(λ, map(λ->B(T=4000.0,λ=λ), λ), "g.")
# plot(λ, map(λ->B(T=5000.0,λ=λ), λ), "b.")


module CoordinateSystem

using LinearAlgebra

export CSys

struct CSys
    O::Vector{Float64}
    x::Vector{Float64}
    y::Vector{Float64}
    z::Vector{Float64}
    T::Matrix{Float64} #transorm_matrix
end

function CSys(o::Vector{Float64},p₁::Vector{Float64},p₂::Vector{Float64})
    v₁=p₁-o
    v₂=p₂-o
    if abs(v₁⋅v₂/norm(v₁)/norm(v₂))==1
        error("Two vectors should not be parallel!")
    end
    x=v₁/norm(v₁)
    z=v₁×v₂
    z=z/norm(z)
    y=z×x
    T=[reshape(x,1,3);
        reshape(y,1,3);
        reshape(z,1,3)]
    CSys(o,x,y,z,T)
end

end


# Test factor analysis implementation using a test similar to the one in
# scikit-learn, so we can directly compare implementations if needed.

using MultivariateStats

n_samples = 20
n_features = 5
n_components = 3

X = readdlm("X_fa.csv", ',')'
fa, L = fit(FactorAnalysis, X, n_components)

X_t = transform(fa, X)

X_rec = fa.Vbar .+ fa.F*X_t .+ fa.Ψ.*randn(n_features,n_samples)


#=  https://www.damtp.cam.ac.uk/research/dd/files/teaching/MB17Gog.pdf#subsubsection.1.3.4  =#
using Plots

δt = 0.001
T = 1
times = 0:δt:T
ntimes = length(times)

N = 10000

S = zeros(size(times))
S[1] = N - 1
I = zeros(size(times))
I[1] = 1
R = zeros(size(times))

# Rate of infection
β = 0.005
# Rate of recovery
ν = 4

# Randomness for σ > 0
σ = 1

for i ∈ 2:ntimes
    δS = - β * I[i-1] * S[i-1] * δt
	δR = ν * I[i-1] * δt
	δI = - δS - δR
    S[i] = max(S[i-1] + δS * exp(σ * randn()), 0)
    I[i] = max(I[i-1] + δI * exp(σ * randn()), 0)
	R[i] = max(R[i-1] + δR * exp(σ * randn()), 0)
	scale = (S[i] + I[i] + R[i])/N
	S[i] /= scale; I[i] /= scale; R[i] /= scale
end

lineplot = plot(times, [S I R], label = ["Susceptible" "Infectious" "Recovered"], color = [:gold :red :blue]);
areaplot_ = areaplot(times, [S I R], label = ["Susceptible" "Infectious" "Recovered"], color = [:gold :red :blue]);

display(lineplot)
display(areaplot_)


function hard_coded_radau_3()
    s6 = sqrt(BigFloat(6))

    # https://en.wikipedia.org/wiki/List_of_Runge%E2%80%93Kutta_methods#Radau_IIA_methods
    a11 = 11 // 45 - s6 * (7 // 360)
    a12 = 37 // 225 - s6 * (169 // 1800)
    a13 = -2 // 225 + s6 * (1 // 75)
    a21 = 37 // 225 + s6 * (169 // 1800)
    a22 = 11 // 45 + s6 * (7 // 360)
    a23 = -2 // 225 - s6 * (1 // 75)
    a31 = 4 // 9 - s6 * (1 // 36)
    a32 = 4 // 9 + s6 * ( 1 // 36)
    a33 = 1 // 9
    A = [a11 a12 a13; a21 a22 a23; a31 a32 a33]

    b = A[end, :]

    c1 = 2 // 5 - s6 * (1 // 10)
    c2 = 2 // 5 + s6 * (1 // 10)
    c3 = 1
    c = [c1, c2, c3]

    # https://github.com/scipy/scipy/blob/14142ff70d84a6ce74044a27919c850e893648f7/scipy/integrate/_ivp/radau.py#L37
    bi = [ 13//3 + s6 *(7//3) -23//3 - s6 * (22 // 3)  10//3 + 5 * s6;
    13//3 - s6 * (7//3)  -23//3 + s6 * (22 // 3)  10//3 - 5 * s6;
    1//3 -8//3 10//3]

    # Practical Implementation of High-Order Multiple Precision Fully Implicit
    # Runge-Kutta Methods with Step Size Control Using Embedded Formula
    # (Tomonori Kouya)
    b̂ = b + [-(2 + 3 * s6)/6, -(2-3*s6)/6, -1//3] * find_real_eigenvalue(A)

    return A, b, c, bi, b̂
end

function hard_coded_radau_2()
    a11 = 5 // 12
    a12 = -1 // 12
    a21 = 3 // 4
    a22 = 1 // 4
    A = [a11 a12; a21 a22]

    B = A[end, :]

    C = [1 // 3,   1]

    return A, B, C
end

big_eps = eps(BigFloat(1))

n = 2
A, b, c, d, bi = radau_butcher_table_plus(n)
A2, B2, C2 = hard_coded_radau_2()

@testset "Radau Butcher Table 2" begin
    tol = 100 * big_eps * 2^2
    @test sum(abs.(A2 - A)) < tol
    @test sum(abs.(C2 - c)) < tol
    @test sum(abs.(B2 - b)) < tol
end

n = 3
A, b, c, d, bi = radau_butcher_table_plus(n)
A3, B3, C3, D3, bi_3 = hard_coded_radau_3()

@testset "Radau Butcher Table 3" begin
    tol = 100 * big_eps * 3^2
    @test sum(abs.(A3 - A)) < tol
    @test sum(abs.(B3 - b)) < tol
    @test sum(abs.(C3 - c)) < tol
    @test sum(abs.(D3 - d)) < tol
    @test sum(abs.(bi_3 - bi[2:4])) < tol
    @test abs(sum(bi) - 1) < tol
    @test bi[1] == find_real_eigenvalue(A)
end


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

using AlgebraicPetri

using OrdinaryDiffEq
using Plots

using Catlab
using Catlab.Graphics
using Catlab.WiringDiagrams
using Catlab.CategoricalAlgebra
using Catlab.Programs.RelationalPrograms

display_uwd(ex) = to_graphviz(ex, box_labels=:name, junction_labels=:variable, edge_attrs=Dict(:len=>".75"));

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

birth_petri = Open(PetriNet(1, 1=>(1,1)));
Graph(birth_petri)
#-
predation_petri = Open(PetriNet(2, (1,2)=>(2,2)));
Graph(predation_petri)
#-
death_petri = Open(PetriNet(1, 1=>()));
Graph(death_petri)

# #### Step 2: Generate models using a relational syntax

lotka_volterra = @relation (wolves, rabbits) begin
  birth(rabbits)
  predation(rabbits, wolves)
  death(wolves)
end
display_uwd(lotka_volterra)
#-
lv_dict = Dict(:birth=>birth_petri, :predation=>predation_petri, :death=>death_petri);
lotka_petri = apex(oapply(lotka_volterra, lv_dict))
Graph(lotka_petri)

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

u0 = [100, 10];
p = [.3, .015, .7];
prob = ODEProblem(vectorfield(lotka_petri),u0,(0.0,100.0),p);
sol = solve(prob,Tsit5(),abstol=1e-8);
plot(sol)

# #### Step 3: Extend your model to handle more complex phenomena
# such as a small food chain between little fish, big fish, and sharks

dual_lv = @relation (fish, Fish, Shark) begin
  birth(fish)
  predation(fish, Fish)
  death(Fish)
  predation(Fish, Shark)
  death(Shark)
end
display_uwd(dual_lv)
#-
dual_lv_petri = apex(oapply(dual_lv, lv_dict))
Graph(dual_lv_petri)

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

u0 = [100, 10, 2];
p = [.3, .015, .7, .017, .35];
prob = ODEProblem(vectorfield(dual_lv_petri),u0,(0.0,100.0),p);
sol = solve(prob,Tsit5(),abstol=1e-6);
plot(sol)


using Distributions

#export AbstractConvolutionKernel, GaussianConvolutionKernel, GaussianMixtureConvolutionKernel

"""Convolution kernel using a truncated Gaussian with zero mean and unit width.
Optionally truncate at limit_kernel_width (defaults to 10)."""
struct GaussianConvolutionKernel{T} <: AbstractConvolutionKernel  where {T<:Number}
  dist::Normal{T}
  norm::T
  limit_kernel_width::Float64

  function GaussianConvolutionKernel(width::T; limit_kernel_width::Number=limit_kernel_width_default) where {T<:Number}
      @assert width > 0
      @assert limit_kernel_width > width
      dist = Normal(zero(T), width)
      norm = one(T)/(cdf(dist,limit_kernel_width)-cdf(dist,-limit_kernel_width))
      new{T}(dist, norm, limit_kernel_width)
  end
end

"Evaluate kernel at wavelength(s) x"
function (kernel::GaussianConvolutionKernel)(x::Number)
    x < kernel.limit_kernel_width^2 ? kernel.norm*pdf(kernel.dist,x) : zero(x)
end

function (kernel::GaussianConvolutionKernel)(x::A1) where { T1<:Number, A1<:AbstractArray{T1} }
    kernel.(x)
end

"""Convolution kernel using a mixture of Gaussians with zero mean and various width and weights.
Optionally truncate at limit_kernel_width (defaults to 10)."""
struct GaussianMixtureConvolutionKernel{T1,T2} <: AbstractConvolutionKernel  where {T1<:Number,T2<:Number}
    dists::Array{Normal{T1},1}
    weights::Array{T2,1}
    norm::T2
    limit_kernel_width::Float64

    function GaussianMixtureConvolutionKernel(kernel_widths::AbstractArray{T1,1}, kernel_weights::AbstractArray{T2,1};
            limit_kernel_width::Number=limit_kernel_width_default) where {T1<:Number,T2<:Number}
        @assert size(kernel_widths) == size(kernel_weights)
        @assert all(kernel_widths .> 0)
        @assert all(limit_kernel_width .> kernel_widths)
        @assert sum(kernel_weights) ≈ one(eltype(kernel_weights))
        dists = [Normal(zero(kernel_widths[i]),kernel_widths[i]) for i in 1:length(kernel_widths) ]
        cdf_at_limit(d) = cdf(d,limit_kernel_width)
        norm = sum(kernel_weights./(cdf_at_limit.(dists))) # assumes components symmetric around zero
        new{T1,T2}(dists, kernel_weights, norm, limit_kernel_width)
    end
end

"Evaluate kernel at wavelength(s) x"
function (kernel::GaussianMixtureConvolutionKernel)(x::Number)
    abs(x) < kernel.limit_kernel_width ?  kernel.norm*sum(kernel.weights.*pdf.(kernel.dists,x)) : zero(x)
end

function (kernel::GaussianMixtureConvolutionKernel)(x::AbstractArray{T}) where T<:Number
    kernel.(x)
end


include("../optimizers/trustregion.jl")
include("../utils/functions.jl")
using Gadfly

xvals = trust_region([3.;4.], 6, 1, 0.1, fenton, fenton_g, fenton_h,
    2000, "dogleg");
cvals = trust_region([3.;4.], 6, 1, 0.1, fenton, fenton_g, fenton_h,
    2000, "cg_steihaug");

println(xvals[end])

nsamps = length(xvals)
nsamps2 = length(cvals)

fx = [fenton(xvals[i]) for i in 1:nsamps]
fx2 = [fenton(cvals[i]) for i in 1:nsamps2]

Gadfly.plot(layer(x=1:nsamps, y=fx, Geom.line, Theme(default_color=color("blue"))),
layer(x=1:nsamps2, y=fx2, Geom.line, Theme(default_color=color("red"))),
#layer(x=1:nsamps3, y=fx3, Geom.line, Theme(default_color=color("orange"))),
Guide.xlabel("iteration"), Guide.ylabel("f(x)"), Guide.title("Value of function"),
Guide.manual_color_key("Legend", ["Newton dogleg", "Steihaug-CG", "quasi-Newton SR1"], ["blue", "red", "orange"]),
Scale.x_log10, Scale.y_log10)

nsamps = length(xvals)

grads = [norm(fenton_g(xvals[i]), 2) for i in 1:nsamps]
grads2 = [norm(fenton_g(cvals[i]), 2) for i in 1:nsamps2]
#grads3 = [norm(fenton_g(cvals[i]), 2) for i in 1:nsamps3]

Gadfly.plot(
layer(x=1:nsamps-2, y=grads[2:nsamps-1,:]./grads[1:nsamps-2,:], Geom.line, Theme(default_color=color("blue"))),
layer(x=1:nsamps2-2, y=grads2[2:nsamps2-1,:]./grads2[1:nsamps2-2,:], Geom.line, Theme(default_color=color("red"))),
#layer(x=1:nsamps3-1, y=grads3[2:nsamps3,:]./grads3[1:nsamps3-1,:], Geom.line, Theme(default_color=color("orange"))),
Guide.xlabel("iteration"), Guide.ylabel("gradient norm ratios"), Guide.title("gradient norm ratios"),
Guide.manual_color_key("Legend", ["Newton dogleg", "Steihaug-CG", "quasi-Newton SR1"], ["blue", "red", "orange"]),
Scale.x_log10, Scale.y_log10)

Gadfly.plot(layer(x=1:nsamps, y=grads, Geom.line, Theme(default_color=color("blue"))),
layer(x=1:nsamps2, y=grads2, Geom.line, Theme(default_color=color("red"))),
#layer(x=1:nsamps3, y=grads3, Geom.line, Theme(default_color=color("orange"))),
    Guide.xlabel("iteration"), Guide.ylabel("gradient norm"), Guide.title("gradient norms"),
Guide.manual_color_key("Legend", ["Newton dogleg", "Steihaug-CG", "quasi-Newton SR1"], ["blue", "red", "orange"]),
    Scale.x_log10, Scale.y_log10)

xvals = trust_region(randn(100), 6, 3, 0.1, rosenbrock,
    rosenbrock_g, rosenbrock_h, 1000, "dogleg");
cvals = trust_region(randn(100), 6, 3, 0.1, rosenbrock,
    rosenbrock_g, rosenbrock_h, 1000, "cg_steihaug");

nsamps = length(xvals)
nsamps2 = length(cvals)

func = rosenbrock
func_g = rosenbrock_g

fx = [func(xvals[i]) for i in 1:nsamps]
fx2 = [func(cvals[i]) for i in 1:nsamps2]

Gadfly.plot(layer(x=1:nsamps, y=fx, Geom.line, Theme(default_color=color("blue"))),
layer(x=1:nsamps2, y=fx2, Geom.line, Theme(default_color=color("red"))),
#layer(x=1:nsamps3, y=fx3, Geom.line, Theme(default_color=color("orange"))),
Guide.xlabel("iteration"), Guide.ylabel("f(x)"), Guide.title("Value of function"),
Guide.manual_color_key("Legend", ["Newton dogleg", "Steihaug-CG", "quasi-Newton SR1"],
    ["blue", "red", "orange"]),
Scale.x_log10, Scale.y_log10)

nsamps = length(xvals)
nsamps2 = length(cvals)

grads = [norm(func_g(xvals[i]), 2) for i in 1:nsamps]
grads2 = [norm(func_g(cvals[i]), 2) for i in 1:nsamps2]
#grads3 = [norm(func_g(cvals[i]), 2) for i in 1:nsamps3]

Gadfly.plot(
layer(x=1:nsamps-2, y=grads[2:nsamps-1,:]./grads[1:nsamps-2,:], Geom.line, Theme(default_color=color("blue"))),
layer(x=1:nsamps2-2, y=grads2[2:nsamps2-1,:]./grads2[1:nsamps2-2,:], Geom.line, Theme(default_color=color("red"))),
#layer(x=1:nsamps3-1, y=grads3[2:nsamps3,:]./grads3[1:nsamps3-1,:], Geom.line, Theme(default_color=color("orange"))),
Guide.xlabel("iteration"), Guide.ylabel("gradient norm ratios"), Guide.title("gradient norm ratios"),
Guide.manual_color_key("Legend", ["Newton dogleg", "Steihaug-CG", "quasi-Newton SR1"], ["blue", "red", "orange"]),
Scale.x_log10, Scale.y_log10)

Gadfly.plot(layer(x=1:nsamps, y=grads, Geom.line, Theme(default_color=color("blue"))),
layer(x=1:nsamps2, y=grads2, Geom.line, Theme(default_color=color("red"))),
#layer(x=1:nsamps3, y=grads3, Geom.line, Theme(default_color=color("orange"))),
    Guide.xlabel("iteration"), Guide.ylabel("gradient norm"), Guide.title("gradient norms"),
Guide.manual_color_key("Legend", ["Newton dogleg", "Steihaug-CG", "quasi-Newton SR1"], ["blue", "red", "orange"]),
    Scale.x_log10, Scale.y_log10)

xvals = trust_region(ones(50)*10, 6, 3, 0.1, cute, cute_g, cute_h, 2000, "dogleg");
cvals = trust_region(ones(50)*10, 6, 3, 0.1, cute, cute_g, cute_h, 2000, "cg_steihaug");

nsamps = length(xvals)
nsamps2 = length(cvals)

func = cute
func_g = cute_g

fx = [func(xvals[i]) for i in 1:nsamps]
fx2 = [func(cvals[i]) for i in 1:nsamps2]

Gadfly.plot(layer(x=1:nsamps, y=fx, Geom.line, Theme(default_color=color("blue"))),
layer(x=1:nsamps2, y=fx2, Geom.line, Theme(default_color=color("red"))),
#layer(x=1:nsamps3, y=fx3, Geom.line, Theme(default_color=color("orange"))),
Guide.xlabel("iteration"), Guide.ylabel("f(x)"), Guide.title("Value of function"),
Guide.manual_color_key("Legend", ["Newton dogleg", "Steihaug-CG", "quasi-Newton SR1"], ["blue", "red", "orange"]),
Scale.x_log10, Scale.y_log10)

nsamps = length(xvals)

grads = [norm(func_g(xvals[i]), 2) for i in 1:nsamps]
grads2 = [norm(func_g(cvals[i]), 2) for i in 1:nsamps2]
#grads3 = [norm(func_g(cvals[i]), 2) for i in 1:nsamps3]

Gadfly.plot(
layer(x=1:nsamps-1, y=grads[2:nsamps,:]./grads[1:nsamps-1,:], Geom.line, Theme(default_color=color("blue"))),
layer(x=1:nsamps2-1, y=grads2[2:nsamps2,:]./grads2[1:nsamps2-1,:], Geom.line, Theme(default_color=color("red"))),
#layer(x=1:nsamps3-1, y=grads3[2:nsamps3,:]./grads3[1:nsamps3-1,:], Geom.line, Theme(default_color=color("orange"))),
Guide.xlabel("iteration"), Guide.ylabel("gradient norm ratios"), Guide.title("gradient norm ratios"),
Guide.manual_color_key("Legend", ["Newton dogleg", "Steihaug-CG", "quasi-Newton SR1"], ["blue", "red", "orange"]),
Scale.x_log10, Scale.y_log10)

Gadfly.plot(layer(x=1:nsamps, y=grads, Geom.line, Theme(default_color=color("blue"))),
layer(x=1:nsamps2, y=grads2, Geom.line, Theme(default_color=color("red"))),
#layer(x=1:nsamps3, y=grads3, Geom.line, Theme(default_color=color("orange"))),
    Guide.xlabel("iteration"), Guide.ylabel("gradient norm"), Guide.title("gradient norms"),
Guide.manual_color_key("Legend", ["Newton dogleg", "Steihaug-CG", "quasi-Newton SR1"], ["blue", "red", "orange"]),
    Scale.x_log10, Scale.y_log10)

xvals = trust_region(ones(50)*10, 6, 3, 0.1, cute2, cute2_g, cute2_h, 2000, "dogleg");
cvals = trust_region(ones(50)*10, 6, 3, 0.1, cute2, cute2_g, cute2_h, 2000, "cg_steihaug");

nsamps = length(xvals)
nsamps2 = length(cvals)

func = cute2
func_g = cute2_g

fx = [func(xvals[i]) for i in 1:nsamps]
fx2 = [func(cvals[i]) for i in 1:nsamps2]

Gadfly.plot(layer(x=1:nsamps, y=fx, Geom.line, Theme(default_color=color("blue"))),
layer(x=1:nsamps2, y=fx2, Geom.line, Theme(default_color=color("red"))),
#layer(x=1:nsamps3, y=fx3, Geom.line, Theme(default_color=color("orange"))),
Guide.xlabel("iteration"), Guide.ylabel("f(x)"), Guide.title("Value of function"),
Guide.manual_color_key("Legend", ["Newton dogleg", "Steihaug-CG", "quasi-Newton SR1"], ["blue", "red", "orange"]),
Scale.x_log10, Scale.y_log10)

nsamps = length(xvals)

grads = [norm(func_g(xvals[i]), 2) for i in 1:nsamps]
grads2 = [norm(func_g(cvals[i]), 2) for i in 1:nsamps2]
#grads3 = [norm(func_g(cvals[i]), 2) for i in 1:nsamps3]

Gadfly.plot(
layer(x=1:nsamps-1, y=grads[2:nsamps,:]./grads[1:nsamps-1,:], Geom.line, Theme(default_color=color("blue"))),
layer(x=1:nsamps2-1, y=grads2[2:nsamps2,:]./grads2[1:nsamps2-1,:], Geom.line, Theme(default_color=color("red"))),
#layer(x=1:nsamps3-1, y=grads3[2:nsamps3,:]./grads3[1:nsamps3-1,:], Geom.line, Theme(default_color=color("orange"))),
Guide.xlabel("iteration"), Guide.ylabel("gradient norm ratios"), Guide.title("gradient norm ratios"),
Guide.manual_color_key("Legend", ["Newton dogleg", "Steihaug-CG", "quasi-Newton SR1"], ["blue", "red", "orange"]),
Scale.x_log10, Scale.y_log10)

Gadfly.plot(layer(x=1:nsamps, y=grads, Geom.line, Theme(default_color=color("blue"))),
layer(x=1:nsamps2, y=grads2, Geom.line, Theme(default_color=color("red"))),
#layer(x=1:nsamps3, y=grads3, Geom.line, Theme(default_color=color("orange"))),
    Guide.xlabel("iteration"), Guide.ylabel("gradient norm"), Guide.title("gradient norms"),
Guide.manual_color_key("Legend", ["Newton dogleg", "Steihaug-CG", "quasi-Newton SR1"], ["blue", "red", "orange"]),
    Scale.x_log10, Scale.y_log10)


# period doubling route
using DrWatson
@quickactivate "NonlinearDynamicsTextbook"
include(srcdir("style.jl"))
using DynamicalSystems, PyPlot

fig = figure(figsize = (figx, 2figy))
axlo = fig.add_subplot(3, 1, 1)
axro = fig.add_subplot(3, 1, 3)

zoomaxs = []
for i in 1:3
	push!(zoomaxs, fig.add_subplot(3, 3, 3+i))
end

axlo.clear()
ds = Systems.logistic(0.5)
i = 1
pvalues = range(2.5, 4; length = 2001)
n = 2000
Ttr = 2000
p_index = 1
output = orbitdiagram(ds, i, p_index, pvalues; n = n, Ttr = Ttr)

L = length(pvalues)
x = Vector{Float64}(undef, n*L)
y = copy(x)
for j in 1:L
    x[(1 + (j-1)*n):j*n] .= pvalues[j]
    y[(1 + (j-1)*n):j*n] .= output[j]
end
axlo.plot(x, y, ls = "None", ms = 0.5, color = "black", marker = "o", alpha = 0.01)
axlo.set_xlim(pvalues[1], pvalues[end]); axlo.set_ylim(0,1)
axlo.set_ylabel("\$x\$", labelpad = 10)
axlo.set_xticks(2.5:0.5:4)
axlo.set_yticks(0:0.2:1)
axlo.set_xlabel("\$r\$"; labelpad = -20)
axlo.text(2.6, 0.1, "logistic map", size = 36)
r∞ = 3.56995
axlo.axvline(r∞; color = "C4", alpha = 0.7, ls = "dashed")

# Zoomin boxes
zbox1 = ((3.629, 0.295), (3.640, 0.32))
zbox2 = ((3.6338, 0.304), (3.6342, 0.3065))
zbox3 = ((3.63407, 0.30495), (3.63408, 0.3052))
allboxes = (zbox1, zbox2, zbox3)
for (k, zbox) in enumerate(allboxes)
	pval = range(zbox[1][1], zbox[2][1]; length = 1001)
	nsm = n÷3
	output = orbitdiagram(ds, i, p_index, pval;
		n=nsm, Ttr, ulims = (zbox[1][2], zbox[2][2])
	)
	L = length(pval)
	x = Vector{Float64}(undef, nsm*L)
	y = copy(x)
	for j in 1:L
	    x[(1 + (j-1)*nsm):j*nsm] .= pval[j]
	    y[(1 + (j-1)*nsm):j*nsm] .= output[j]
	end
	zoomaxs[k].clear()
	zoomaxs[k].plot(x, y, ls = "None", ms = 0.5, color = "black", marker = "o", alpha = 0.01)
	zoomaxs[k].set_xlim(zbox[1][1], zbox[2][1])
	zoomaxs[k].set_ylim(zbox[1][2], zbox[2][2])
	zoomaxs[k].axis("off")
end

zbox0 = ((pvalues[1], 0.0), (pvalues[end], 1.0))
allzoomaxs = [axlo, zoomaxs...]
for i in 1:3
	origin = allzoomaxs[i]
	zoomin = allzoomaxs[i+1]
    zbox = allboxes[i]
	c = ("C1", "C3", "C5")[i]
    axis_zoomin!(zoomin, origin, zbox, zbox, c;	connect_lines = i > 1, lw = 4.0)
end

# arrow indicator
axlo.arrow(
	(zbox1[1][1] + zbox1[2][1])/2, 0.05, 0, zbox1[2][2] - 0.1;
	color = "C1", width = 0.005, length_includes_head = true, head_length = 0.1
)

# Roessler
ro = Systems.roessler()
pvalues = range(1.5, stop = 5.18, length = 2001)
i = 1
plane = (2, 0.0)
tf = 4000.0
p_index = 3
output = produce_orbitdiagram(ro, plane, i, p_index, pvalues;
                              tfinal = tf, Ttr = 2000.0)

for (j, p) in enumerate(pvalues)
    axro.plot(fill(p, length(output[j])), output[j], lw = 0,
    marker = "o", ms = 0.2, color = "black", alpha = 0.1)
end
axro.set_xlim(pvalues[1], pvalues[end]);
axro.set_ylim(minimum(minimum(o) for o in output), maximum(maximum(o) for o in output))
axro.set_ylabel("\$x\$ (at section)", labelpad = 25)
axro.set_xlabel("\$c\$"; labelpad = -20)
axro.set_yticks(3:3:9)
axro.text(1.7, 7, "Rössler system", size = 36)
fig.tight_layout(pad=0.3)
fig.subplots_adjust(hspace = 0.3, wspace = 0.1)
add_identifiers!(fig, (axlo, axro))
wsave(plotsdir("4", "orbit_diagrams"), fig)


# Sampling from manifolds. Currently used for testing only.
abstract type Manifold{Int, T<:AbstractFloat} end

# ............................................................................ #
struct Torus{T} <: Manifold{3, T}
    R::T
    r::T
    center::Tuple{T, T, T}
end

function Torus(R, r)
    args = convert.(AbstractFloat, (R, r))
    z = zero(eltype(args))
    Torus(args..., (z, z, z))
end

(t::Torus)(θ, φ) =
    ((t.R + t.r*cos(θ)) * cos(φ), (t.R + t.r*cos(θ)) * sin(φ), t.r * sin(θ)) .+
    t.center

# ............................................................................ #
struct AsymTorus{T} <: Manifold{3, T}
    R::T
    r_large::T
    r_small::T
    center::Tuple{T, T, T}
end

function AsymTorus(R, r1, r2)
    args = convert.(AbstractFloat, (R, r1, r2))
    z = zero(eltype(args))
    AsymTorus(args..., (z, z, z))
end

function (at::AsymTorus)(θ, φ)
    R = at.R
    r_large = at.r_large
    r_small = at.r_small

    r = (r_large - r_small) * (cos(φ) + 1)/2 + r_small
    ((R + r*cos(θ)) * cos(φ), (R + r*cos(θ)) * sin(φ), r*sin(θ)) .+ at.center
end

# ............................................................................ #
struct Sphere{T} <: Manifold{3, T}
    r::T
    center::Tuple{T, T, T}
end

function Sphere(r)
    arg = convert(AbstractFloat, r)
    z = zero(typeof(arg))
    Sphere(arg, (z, z, z))
end

(s::Sphere)(θ, φ) =
    (s.r*cos(θ) * sin(φ), s.r*sin(θ) * sin(φ), s.r*cos(φ)) .+ s.center

# ............................................................................ #
struct Circle{T} <: Manifold{2, T}
    r::T
    center::Tuple{T, T}
end

function Circle(r)
    arg = convert(AbstractFloat, r)
    z = zero(typeof(arg))
    Circle(arg, (z, z))
end

(c::Circle)(θ) = (c.r * cos(θ), c.r * sin(θ)) .+ c.center

# ............................................................................ #
# Uniformly sampling n points from a manifold returning D by n matrix.
# TODO: Circular noise? Overenginnered as it is?
function Base.rand(m::Manifold{D, T}, n=1; noise=0) where {D, T}
    res = zeros(T, D, n)
    for i in 1:n
        pars = 2π .* rand(D-1)
        res[:, i] .= m(pars...) .+ (rand(D) - 0.5) * noise
    end
    res
end


#Ipopt.jl example - This is taken from the documentation available on:
# http://ipoptjl.readthedocs.org/en/latest/ipopt.html#

# Install Package if you don't have it
Pkg.add("Ipopt")

tic()
# Load the package
using Ipopt

# HS071
# min x1 * x4 * (x1 + x2 + x3) + x3
# st  x1 * x2 * x3 * x4 >= 25
#     x1^2 + x2^2 + x3^2 + x4^2 = 40
#     1 <= x1, x2, x3, x4 <= 5
# Start at (1,5,5,1)
# End at (1.000..., 4.743..., 3.821..., 1.379...)

## Set the number of arguments to solve for and bounds for these
## Aside for new Julia users: the '.0' after each number is important - it tells Julia
## that the vectors x_L, etc are of Float type.
## Note that Julia is a functional language, with methods (e.g. a function) working for the types
## for which it is defined (e.g. integer, Float). So, a function defined for an integer will not
## work if you enter a Float variable, and vice-versa.
n = 4
x_L = [1.0, 1.0, 1.0, 1.0]
x_U = [5.0, 5.0, 5.0, 5.0]

# Set the number of constraints and bounds on these - Note that equality
# constraints can be accommodated for by setting the upper bound = lower bound
m = 2
g_L = [25.0, 40.0]
g_U = [2.0e19, 40.0]

# Set the objective function - Returns the value of the objective function at the current solution x
# Note that the '::Vector{Float64}' below defines the type of x. This is a Julia-specific thing, and
# need not be defined. The function would be created even without including the type for x.
function eval_f(x::Vector{Float64})
  return x[1] * x[4] * (x[1] + x[2] + x[3]) + x[3]
end

# Set the constraint set - Sets the value of the constraint functions g at the current solution x
# Note that there is no need to create a separate vector g in which to save the constraints.
# The values of g are set 'in-place'
function eval_g(x::Vector{Float64}, g::Vector{Float64})
  g[1] = x[1]   * x[2]   * x[3]   * x[4]
  g[2] = x[1]^2 + x[2]^2 + x[3]^2 + x[4]^2
end

# Gradient of the objective function - Sets the value of the gradient of the objective function at the current solution x
# Note that there is no need to create a vector grad_f to fill in. The values of grad_f are set 'in-place'.
function eval_grad_f(x::Vector{Float64}, grad_f::Vector{Float64})
  grad_f[1] = x[1] * x[4] + x[4] * (x[1] + x[2] + x[3])
  grad_f[2] = x[1] * x[4]
  grad_f[3] = x[1] * x[4] + 1
  grad_f[4] = x[1] * (x[1] + x[2] + x[3])
end

# Gradient of the constraint set (Jacobian matrix). There are two ways of setting this:
# (1) Under the ':Structure' mode, indicate the rows and columns with non-zero values (i.e. tell Ipopt
# about the sparsity structure of the matrix)
# (2) Undet the ':Values' mode, tell it the actual values.
# Each non-zero value is indexed, starting from 1.
###########################################
# General specification is as follows:
###########################################
# function eval_jac_g(
#  x::Vector{Float64},         # Current solution
#  mode,                       # Either :Structure or :Values
#  rows::Vector{Int32},        # Sparsity structure - row indices
#  cols::Vector{Int32},        # Sparsity structure - column indices
#  values::Vector{Float64})    # The values of the Hessian

#  if mode == :Structure
    # rows[...] = ...
    # ...
    # cols[...] = ...
#  else
    # values[...] = ...
#  end
# end
#############################################
function eval_jac_g(x, mode, rows, cols, values)
  if mode == :Structure
    # Constraint (row) 1
    rows[1] = 1; cols[1] = 1
    rows[2] = 1; cols[2] = 2
    rows[3] = 1; cols[3] = 3
    rows[4] = 1; cols[4] = 4
    # Constraint (row) 2
    rows[5] = 2; cols[5] = 1
    rows[6] = 2; cols[6] = 2
    rows[7] = 2; cols[7] = 3
    rows[8] = 2; cols[8] = 4
  else
    # Constraint (row) 1
    values[1] = x[2]*x[3]*x[4]  # 1,1
    values[2] = x[1]*x[3]*x[4]  # 1,2
    values[3] = x[1]*x[2]*x[4]  # 1,3
    values[4] = x[1]*x[2]*x[3]  # 1,4
    # Constraint (row) 2
    values[5] = 2*x[1]  # 2,1
    values[6] = 2*x[2]  # 2,2
    values[7] = 2*x[3]  # 2,3
    values[8] = 2*x[4]  # 2,4
  end
end

# Hessian of the Lagrangian - this is an optional object. In terms of the Julia syntax, it is
# similar to the Jacobian.
# The general definition and types for the objects in the function below are:
# function eval_h(
# x::Vector{Float64},         # Current solution
# mode,                       # Either :Structure or :Values
# rows::Vector{Int32},        # Sparsity structure - row indices
# cols::Vector{Int32},        # Sparsity structure - column indices
# obj_factor::Float64,        # Lagrangian multiplier for objective
# lambda::Vector{Float64},    # Multipliers for each constraint
# values::Vector{Float64})    # The values of the Hessian
# if....else...end            # the function text goes here
# end                         # close the function

function eval_h(x, mode, rows, cols, obj_factor, lambda, values)
  if mode == :Structure
    # Symmetric matrix, fill the lower left triangle only
    idx = 1
    for row = 1:4
      for col = 1:row
        rows[idx] = row
        cols[idx] = col
        idx += 1
      end
    end
  else
    # Again, only lower left triangle
    # Objective
    values[1] = obj_factor * (2*x[4])  # 1,1
    values[2] = obj_factor * (  x[4])  # 2,1
    values[3] = 0                      # 2,2
    values[4] = obj_factor * (  x[4])  # 3,1
    values[5] = 0                      # 3,2
    values[6] = 0                      # 3,3
    values[7] = obj_factor * (2*x[1] + x[2] + x[3])  # 4,1
    values[8] = obj_factor * (  x[1])  # 4,2
    values[9] = obj_factor * (  x[1])  # 4,3
    values[10] = 0                     # 4,4

    # First constraint
    values[2] += lambda[1] * (x[3] * x[4])  # 2,1
    values[4] += lambda[1] * (x[2] * x[4])  # 3,1
    values[5] += lambda[1] * (x[1] * x[4])  # 3,2
    values[7] += lambda[1] * (x[2] * x[3])  # 4,1
    values[8] += lambda[1] * (x[1] * x[3])  # 4,2
    values[9] += lambda[1] * (x[1] * x[2])  # 4,3

    # Second constraint
    values[1]  += lambda[2] * 2  # 1,1
    values[3]  += lambda[2] * 2  # 2,2
    values[6]  += lambda[2] * 2  # 3,3
    values[10] += lambda[2] * 2  # 4,4
  end
end

# Creates and returns an IpoptProblem with the given options. Raises error if something goes wrong during construction.
# If you do not provide a callback for the Hessian, you must set the Hessian approximation option:
# addOption(prob, "hessian_approximation", "limited-memory")
prob = createProblem(n, x_L, x_U, m, g_L, g_U, 8, 10,
                     eval_f, eval_g, eval_grad_f, eval_jac_g, eval_h)

# Set starting solution
prob.x = [1.0, 5.0, 5.0, 1.0]

# Solve
status = solveProblem(prob)

# Display the solution
println(Ipopt.ApplicationReturnStatus[status])
println(prob.x)
println(prob.obj_val)

toc()


const CLP_SOLVER = ClpSolver()
const MAX_SENSE = -1.0
const MIN_SENSE = 1.0

## -------------------------------------------------------------------
function fba(S, b, lb, ub, c; solver = CLP_SOLVER)
    LP = linprog(
        c, # Opt sense vector
        S, # Stoichiometric matrix
        b, # row lb
        b, # row ub
        lb, # column lb
        ub, # column ub
        solver
    )
    return isempty(LP.sol) ? zeros(length(lb)) : LP.sol
end

function fba(S, b, lb, ub, c, idx, sense; kwargs...)
    bk_c = c[idx]
    c[idx] = sense
    try;
        sol = fba(S, b, lb, ub, c; kwargs...)
        return sol
    finally; c[idx] = bk_c end
end

## -------------------------------------------------------------------
function fva(S, b, lb, ub, c, idx::Integer;  kwargs...)
    bk_c = c[idx]
    c[idx] = MIN_SENSE
    try
        L = fba(S, b, lb, ub, c; kwargs...)[idx]
        c[idx] = MAX_SENSE
        U = fba(S, b, lb, ub, c; kwargs...)[idx]
        return (first(L), first(U))
    finally; c[idx] = bk_c end
end

function fva(S, b, lb, ub, c, idxs; kwargs...)
    L = zeros(length(idxs))
    U = zeros(length(idxs))
    for (i, idx) in idxs |> enumerate
        L[i], U[i] = fva(S, b, lb, ub, c, idx; kwargs...)
    end
    return (L, U)
end


"""
        mairca(decisionMat, weights, fns)

Apply MAIRCA (Multi Attributive Ideal-Real Comparative Analysis) for a given matrix and weights.

# Arguments:
 - `decisionMat::DataFrame`: n × m matrix of objective values for n alternatives and m criteria
 - `weights::Array{Float64, 1}`: m-vector of weights that sum up to 1.0. If the sum of weights is not 1.0, it is automatically normalized.
 - `fns::Array{Function, 1}`: m-vector of functions to be applied on the columns.

# Description
mairca() applies the MAIRCA method to rank n alternatives subject to m criteria which are supposed to be
either maximized or minimized.

# Output
- `::MAIRCAResult`: MAIRCAResult object that holds multiple outputs including scores, rankings, and best index.

# Examples
```julia-repl
julia> decmat = [6.952 8.000 6.649 7.268 8.000 7.652 6.316;
       7.319 7.319 6.604 7.319 8.000 7.652 5.313;
       7.000 7.319 7.652 6.952 7.652 6.952 4.642;
       7.319 6.952 6.649 7.319 7.652 6.649 5.000]
4×7 Array{Float64,2}:
 6.952  8.0    6.649  7.268  8.0    7.652  6.316
 7.319  7.319  6.604  7.319  8.0    7.652  5.313
 7.0    7.319  7.652  6.952  7.652  6.952  4.642
 7.319  6.952  6.649  7.319  7.652  6.649  5.0

 julia> df = makeDecisionMatrix(decmat)
 4×7 DataFrame
  Row │ Crt1     Crt2     Crt3     Crt4     Crt5     Crt6     Crt7
      │ Float64  Float64  Float64  Float64  Float64  Float64  Float64
 ─────┼───────────────────────────────────────────────────────────────
    1 │   6.952    8.0      6.649    7.268    8.0      7.652    6.316
    2 │   7.319    7.319    6.604    7.319    8.0      7.652    5.313
    3 │   7.0      7.319    7.652    6.952    7.652    6.952    4.642
    4 │   7.319    6.952    6.649    7.319    7.652    6.649    5.0

julia> weights = [0.172, 0.165, 0.159, 0.129, 0.112, 0.122, 0.140];

julia> fns = [maximum, maximum, maximum, maximum, maximum, maximum, minimum];

julia> result = mairca(df, weights, fns);

julia> result.scores
4-element Array{Float64,1}:
 0.12064543054088471
 0.08066456363291889
 0.14586265389012484
 0.14542366685864686

julia> result.bestIndex
2
```

# References

Pamučar, D., Lukovac, V., Božanić, D., & Komazec, N. (2018). Multi-criteria FUCOM-MAIRCA model for the evaluation of level crossings: case study in the Republic of Serbia. Operational Research in Engineering Sciences: Theory and Applications, 1(1), 108-129.

Ulutaş A.(2019),Swara Ve Mairca Yöntemleri İle Catering Firması Seçimi,BMIJ, (2019), 7(4): 1467-1479 http://dx.doi.org/10.15295/bmij.v7i4.1166
"""
function mairca(decisionMat::DataFrame, weights::Array{Float64,1}, fns::Array{Function,1})::MAIRCAResult

    row, col = size(decisionMat)

    w = unitize(weights)

    T = zeros(Float64, row, col)

    for i in 1:col
        T[:, i] .= w[i] * (1 / row)
    end

    colMax = colmaxs(decisionMat)
    colMin = colmins(decisionMat)

    A = similar(decisionMat)

    for i in 1:row
        for j in 1:col
            if fns[j] == maximum
                @inbounds A[i, j] = T[i,j] * ((decisionMat[i, j] - colMin[j]) / (colMax[j] - colMin[j]))
            elseif fns[j] == minimum
                @inbounds A[i, j] = T[i,j] * ((decisionMat[i, j] - colMax[j]) / (colMin[j] - colMax[j]))
            end
        end
    end

    S = T .- A

    scores = zeros(Float64, row)
    for i in 1:row
        scores[i] = sum(S[i, :])
    end

    rankings = sortperm(scores)

    bestIndex = rankings |> first

    result = MAIRCAResult(
        decisionMat,
        w,
        scores,
        rankings,
        bestIndex
    )

    return result
end


# Arvind Prasadan
# 2019 January
# DMF Package
# Generates a realization of an ARMA process

"""
	gen_arma_sequence(n = 100, ar_comp = [], ma_comp = [], arma_std = 1.0)

Generates a realization of an ARMA process

# Arguments
- `n`: Length of process
- `ar_comp`: List of AR Coefficients; one of this and ma_comp must be non-empty
- `ma_comp`: List of MA Coefficients; one of this and ar_comp must be non-empty
- `arma_std`: Noise Standard Deviation

# Outputs
- `x`: vector of ARMA process realization
"""
function gen_arma_sequence(n = 100, ar_comp = [], ma_comp = [], arma_std = 1.0)

    n = round(abs(n[1])) # Samples
    n = (1 <= n) ? n : 1

    arma_std = abs(arma_std[1]) # Standard Deviation

    if 0 == length(ar_comp) + length(ma_comp)
	x = arma_std * randn(n)
    else
    	# Define ARMA object, simulate from this set of parameters
    	arma = ARMA(ar_comp, ma_comp, arma_std)
    	x = simulation(arma, ts_length = n)
    end

    return x
end