#' ---
#' title: Conic Programming
#' ---

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

#' This tutorial is aimed at providing a simplistic introduction to conic programming using JuMP.

#' ## What is a Cone?
#' A subset $C$ of a vector space $V$ is a cone if $\forall x \in C$ and positive scalars $\alpha$,
#' the product $\alpha x \in C$. A cone C is a convex cone if $\alpha x + \beta y \in C$,
#' for any positive scalars $\alpha, \beta$, and any $x, y \in C$.

#' ## Conic Programming
#' Conic programming problems are convex optimization problems in which a convex function is minimized
#' over the intersection of an affine subspace and a convex cone.
#' An example of a conic-form minimization problems, in the primal form is:

#' $$
#' \begin{align}
#' & \min_{x \in \mathbb{R}^n} & a_0^T x + b_0 \\
#' & \;\;\text{s.t.} & A_i x + b_i & \in \mathcal{C}_i & i = 1 \ldots m
#' \end{align}
#' $$

#' The corresponding dual problem is:

#' $$
#' \begin{align}
#' & \max_{y_1, \ldots, y_m} & -\sum_{i=1}^m b_i^T y_i + b_0 \\
#' & \;\;\text{s.t.} & a_0 - \sum_{i=1}^m A_i^T y_i & = 0 \\
#' & & y_i & \in \mathcal{C}_i^* & i = 1 \ldots m
#' \end{align}
#' $$

#' where each $\mathcal{C}_i$ is a closed convex cone and $\mathcal{C}_i^*$ is its dual cone.

#' ## Some of the Types of Cones Supported by JuMP

using JuMP
using ECOS
using LinearAlgebra
using Random

Random.seed!(1234);

#' By this point we have used quite a few different solvers.
#' To find out all the different solvers and their supported problem types, check out the
#' [solver table](http://www.juliaopt.org/JuMP.jl/v0.19.0/installation/#Getting-Solvers-1) in the docs.

#' ### Second-Order Cone
#' The Second-Order Cone (or Lorenz Cone) of dimension $n$ is of the form:

#' $$
#' Q^n = \{ (t,x) \in \mathbb{R}^\mbox{n} : t \ge ||x||_2 \}
#' $$

#' A Second-Order Cone rotated by $\pi/4$ in the $(x_1,x_2)$ plane is called a Rotated Second-Order Cone.
#' It is of the form:

#' $$
#' Q_r^n = \{ (t,u,x) \in \mathbb{R}^\mbox{n} : 2tu \ge ||x||_2^2, t,u \ge 0 \}
#' $$

#' These cones are represented in JuMP using the MOI sets `SecondOrderCone` and `RotatedSecondOrderCone`.

#' #### Example: Euclidean Projection on a Hyperplane
#' For a given point $u_{0}$ and a set $K$, we refer to any point $u \in K$
#' which is closest to $u_{0}$ as a projection of $u_{0}$ on $K$.
#' The projection of a point $u_{0}$ on a hyperplane $K = \{u | p' \cdot u = q\}$ is given by

#' $$
#' \begin{align*}
#' & \min & ||u - u_{0}|| \\
#' & \;\;\text{s.t.} & p' \cdot u = q  \\
#' \end{align*}
#' $$

u0 = rand(10)
p = rand(10)
q = rand();

#' We can model the above problem as the following conic program:

#' $$
#' \begin{align*}
#' & \min & t \\
#' & \;\;\text{s.t.} & p' \cdot u = q \\
#' & & (t, u - u_{0}) \in Q^{n+1}
#' \end{align*}
#' $$

#' On comparing this with the primal form of a conic problem we saw above,

#' $$
#' \begin{align*}
#' & x = (t , u) &\\
#' & a_0 = e_1 &\\
#' & b_0 = 0 &\\
#' & A_1 = (0, p) &\\
#' & b_1 = -q &\\
#' & C_1 = \mathbb{R}_- &\\
#' & A_2 = 1 &\\
#' & b_2 = -(0, u_0) &\\
#' & C_2 = Q^{n+1} &
#' \end{align*}
#' $$

#' Thus, we can obtain the dual problem as:

#' $$
#' \begin{align*}
#' & \max & y_1 + (0, u_0)^T y_2 \\
#' & \;\;\text{s.t.} & e_1 - (0,p)^T y_1 - y_2 = 0 \\
#' & & y_1 \in \mathbb{R}_- \\
#' & & y_2 \in Q^{n+1}
#' \end{align*}
#' $$

model = Model(optimizer_with_attributes(ECOS.Optimizer, "printlevel" => 0))
@variable(model, u[1:10])
@variable(model, t)
@objective(model, Min, t)
@constraint(model, [t, (u - u0)...] in SecondOrderCone())
@constraint(model, u' * p == q)
optimize!(model)

#+

@show objective_value(model);
@show value.(u);

#+

e1 = [1, zeros(10)...]
dual_model = Model(optimizer_with_attributes(ECOS.Optimizer, "printlevel" => 0))
@variable(dual_model, y1 <= 0)
@variable(dual_model, y2[1:11])
@objective(dual_model, Max, q * y1 + dot(vcat(0, u0), y2))
@constraint(dual_model, e1 - [0, p...] .* y1 - y2 .== 0)
@constraint(dual_model, y2 in SecondOrderCone())
optimize!(dual_model)

#+

@show objective_value(dual_model);

#' We can also have an equivalent formulation using a Rotated Second-Order Cone:

#' $$
#' \begin{align*}
#' & \min & t \\
#' & \;\;\text{s.t.} & p' \cdot u = q \\
#' & & (t, 1/2, u - u_{0})\in Q_r^{n+2}
#' \end{align*}
#' $$

model = Model(optimizer_with_attributes(ECOS.Optimizer, "printlevel" => 0))
@variable(model, u[1:10])
@variable(model, t)
@objective(model, Min, t)
@constraint(model, [t, 0.5, (u - u0)...] in RotatedSecondOrderCone())
@constraint(model, u' * p == q)
optimize!(model)

#+

@show value.(u);

#' The difference here is that the objective in the case of the Second-Order Cone is $||u - u_{0}||_2$,
#' while in the case of a Rotated Second-Order Cone is $||u - u_{0}||_2^2$.
#' However, the value of x is the same for both.

#' ### Exponential Cone

#' An Exponential Cone is a set of the form:

#' $$
#' K_{exp} = \{ (x,y,z) \in \mathbb{R}^3 : y \exp (x/y) \le z, y > 0 \}
#' $$

#' It is represented in JuMP using the MOI set `ExponentialCone`.

#' #### Example: Entropy Maximization
#' As the name suggests, the entropy maximization problem consists of maximizing the entropy function,
#' $H(x) = -x\log{x}$ subject to linear inequality constraints.

#' $$
#' \begin{align*}
#' & \max & - \sum_{i=1}^n x_i \log x_i \\
#' & \;\;\text{s.t.} & \mathbf{1}' x = 1 \\
#' & & Ax \leq b
#' \end{align*}
#' $$

#' We can model this problem using an exponential cone by using the following transformation:

#' $$
#' t\leq -x\log{x} \iff t\leq x\log(1/x)  \iff (1, x, t) \in K_{exp}
#' $$

#' Thus, our problem becomes,

#' $$
#' \begin{align*}
#' & \max & 1^Tt \\
#' & \;\;\text{s.t.} & Ax \leq b \\
#' & & 1^T x = 1 \\
#' & & (1, x_i, t_i) \in K_{exp} && \forall i = 1 \ldots n \\
#' \end{align*}
#' $$

n = 15;
m = 10;
A = randn(m, n);
b = rand(m, 1);

model = Model(optimizer_with_attributes(ECOS.Optimizer, "printlevel" => 0))
@variable(model, t[1:n])
@variable(model, x[1:n])
@objective(model, Max, sum(t))
@constraint(model, sum(x) == 1)
@constraint(model, A * x .<= b )
# Cannot use the exponential cone directly in JuMP, hence we use MOI to specify the set.
@constraint(model, con[i = 1:n], [1, x[i], t[i]] in MOI.ExponentialCone())

optimize!(model);

#+

@show objective_value(model);

#' ### Positive Semidefinite Cone
#' The set of Positive Semidefinite Matrices of dimension $n$ form a cone in $\mathbb{R}^n$.
#' We write this set mathematically as

#' $$
#' \mathcal{S}_{+}^n = \{ X \in \mathcal{S}^n \mid z^T X z \geq 0, \: \forall z\in \mathbb{R}^n \}.
#' $$

#' A PSD cone is represented in JuMP using the MOI sets
#' `PositiveSemidefiniteConeTriangle` (for upper triangle of a PSD matrix) and
#' `PositiveSemidefiniteConeSquare` (for a complete PSD matrix).
#' However, it is prefferable to use the `PSDCone` shortcut as illustrated below.

#' #### Example: Largest Eigenvalue of a Symmetric Matrix
#' Suppose $A$ has eigenvalues $\lambda_{1} \geq \lambda_{2} \ldots \geq \lambda_{n}$.
#' Then the matrix $t I-A$ has eigenvalues $t-\lambda_{1}, t-\lambda_{2}, \ldots, t-\lambda_{n}$.
#' Note that $t I-A$ is PSD exactly when all these eigenvalues are non-negative,
#' and this happens for values $t \geq \lambda_{1} .$
#' Thus, we can model the problem of finding the largest eigenvalue of a symmetric matrix as:

#' $$
#' \begin{align*}
#' \lambda_{1} = \max t \\
#' \text { s.t. } t I-A \succeq 0
#' \end{align*}
#' $$

#+ tangle = false

using LinearAlgebra
using SCS

A = [3 2 4;
     2 0 2;
     4 2 3]

model = Model(optimizer_with_attributes(SCS.Optimizer, "verbose" => 0))
@variable(model, t)
@objective(model, Min, t)
@constraint(model, t .* Matrix{Float64}(I, 3, 3) - A in PSDCone())

optimize!(model)

#+ tangle = false

@show objective_value(model);

#' ## Other Cones and Functions
#' For other cones supported by JuMP, check out the
#' [MathOptInterface Manual](http://www.juliaopt.org/MathOptInterface.jl/dev/apimanual/#Standard-form-problem-1).
#' A good resource for learning more about functions which can be modelled using cones is the
#' MOSEK Modeling Cookbook[[1]](#c1).

#' ### References
#'
#' 1. MOSEK Modeling Cookbook — MOSEK Modeling Cookbook 3.1. Available at: https://docs.mosek.com/modeling-cookbook/index.html.


function laplace_matrix(::Type{T}, n, dims) where T
    D = second_order_central_diff(T, n);
    A = copy(D);

    for idx = 2 : dims
        A = kron(A, speye(n)) + kron(speye(size(A, 1)), D);
    end

    A
end

second_order_central_diff(::Type{T}, dim) where {T} = convert(SparseMatrixCSC{T, Int}, SymTridiagonal(fill(2 * one(T), dim), fill(-one(T), dim - 1)))


"""
Verify, that |A * Q - Q * T| ≈ 0
A, Q and T must be Matrices with appropriate sizes.
"""
function checkfactors(A::AbstractMatrix{T1}, Q::AbstractMatrix{T2}, T::AbstractMatrix{T3}; atol::Real = zero(real(eltype(A))), rtol::Real = eps(real(eltype(A)))) where {T1<:Number, T2<:Number, T3<:Number}

  res = norm(A * Q - Q * T, 1)
  tol = rtol * norm(A, 1) + atol
  # println("checkfactors() = $(Float64(res)) < ! $(Float64(tol)) $(Float64(res/tol))")
  res <= rtol * norm(A, 1) + atol
end


"""
Computes convex hull of 2D points. Assumes non-degenerate inputs.
"""
module ConvexHull

const Point{T} = Tuple{T, T}

raw"""
Classifies whether a point $r$ lies to the left or right of the oriented line $pq$.

Returns:

* -1 if $r$ lies to the left
* 0 if $r$ lies on the line
* 1 if $r$ lies to the right

**Warning**: This computation is not robust. See [Robust Predicates for Computational
Geometry][robust] for details. (If `BigFloat` is used, there is no rounding problems.)

[robust]: http://www.cs.cmu.edu/%7Equake/robust.html
"""
function orient(p::Point{T},
                q::Point{T},
                r::Point{T})::T where T<:AbstractFloat
    pr_x = r[1] - p[1]
    pr_y = r[2] - p[2]
    rq_x = q[1] - r[1]
    rq_y = q[2] - r[2]
    sign(pr_x * rq_y - pr_y * rq_x)
end

# Graham scan
#
# Computational Geometry: Algorithms and Applications, Chapter 1.
function convex_hull_graham_scan!(points::Vector{Point{T}})::Vector{Point{T}} where T<:AbstractFloat
    @assert length(points) >= 3 "Too few points"
    sort!(points)  # in lexical order

    function build_convex_hull(points::AbstractVector{Point{T}})
        n = length(points)
        hull = points[1:2]
        sizehint!(hull, n)

        for i in 3:n
            while true
                if orient(hull[end - 1], hull[end], points[i]) > 0
                    break
                else
                    pop!(hull)
                    if length(hull) < 2
                        break
                    end
                end
            end
            push!(hull, points[i])
        end

        hull
    end

    upper_hull = build_convex_hull(points)
    pop!(upper_hull)
    lower_hull = build_convex_hull(@view points[end:-1:1])
    pop!(lower_hull)

    [upper_hull; lower_hull]
end

# Unsorted version, complexity: O(n^2)
function convex_hull_incremental(points::Vector{Point{T}})::Vector{Point{T}} where T<:AbstractFloat
    n = length(points)
    @assert n >= 3 "Too few points"
    hull = points[1:3]
    if orient(hull[1], hull[2], hull[3]) < 0
        hull[[1, 2]] .= hull[[2, 1]]
    end
    sizehint!(hull, n)

    for i in 4:n
        p = points[i]
        # Vertices inside `[start, stop)` will be removed.
        start::Union{Nothing, Int} = nothing
        stop::Union{Nothing, Int} = nothing
        h_len = length(hull)
        prev_ori = orient(hull[end - 1], hull[end], p) >= 0

        # trying to find two tangent lines
        for (from, to) in ((mod1(i - 1, h_len), i) for i in 1:h_len)
            curr_ori = orient(hull[from], hull[to], p) >= 0
            if !curr_ori && prev_ori
                start = to  # inclusive
                if stop != nothing
                    break
                end
            elseif curr_ori && !prev_ori
                stop = from  # exclusive
                if start != nothing
                    break
                end
            end
            prev_ori = curr_ori
        end

        # Always being on the right side, `p` must be inside the hull.
        if start == nothing
            continue
        end
        # If `start != nothing`, then `stop != nothing`.

        # Note: `start == stop` means no vertices are removed.
        if start <= stop
            splice!(hull, start:(stop - 1), [p])
        else
            # The same as
            #
            #     hull = [hull[stop:(start - 1)]; p]
            #
            # but without allocating extra memory.
            hull[start] = p
            resize!(hull, start)
            deleteat!(hull, 1:(stop - 1))
        end
    end
    hull
end

#===== Unit Tests =====#
using Test
using Random: seed!

function convex_hull_length(hull::Vector{Point{T}})::T where T<:AbstractFloat
    hull2 = circshift(hull, -1)
    map((x, y) -> hypot((x .- y)...), hull2, hull) |> sum
end

@testset "Convex Hull" begin
    seed!()
    n = 100
    a = rand(n, 2)
    points = [tuple(a[i, :]...) for i in 1:n]

    cvh_increm = convex_hull_incremental(points)
    cvh_graham = convex_hull_graham_scan!(points)
    @test convex_hull_length(cvh_increm) ≈ convex_hull_length(cvh_graham)
end
end # module


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

using Markdown
using InteractiveUtils

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

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

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

# ╔═╡ 70179b62-1ea1-11eb-13ce-7fb4d54bad4b
begin
	alphas = [0.5, 1.5, 1]
	lam = 2

	lambda(dist::Weibull) = shape(dist)*scale(dist)^(-shape(dist))
	theta(lam,alpha) = (alpha/lam)^(1/alpha)

	dists = [Weibull.(a,theta(lam,a)) for a in alphas]

	hA(dist,x) = pdf(dist,x)/ccdf(dist,x)
	hB(dist,x) = lambda(dist)*x^(shape(dist)-1)
end

# ╔═╡ 884176a0-21c0-11eb-06f8-55964a60d533
begin
	xGrid = 0.01:0.01:10
	hazardsA = [hA.(d,xGrid) for d in dists]
	hazardsB = [hB.(d,xGrid) for d in dists]

	Text("Maximum difference between implementations: $(maximum(maximum.(hazardsA-hazardsB)))")
end

# ╔═╡ 8841a7b0-21c0-11eb-313a-bfa4eb49f351
begin
	Cl = [:blue :red :green]
	Lb = [L"\lambda=" * string(lambda(d)) * ",   " * L"\alpha =" * string(shape(d))
			for d in dists]

	plot(xGrid, hazardsA, c=Cl, label=reshape(Lb, 1,:), xlabel="x",
		ylabel="Instantaneous failure rate", xlims=(0,10), ylims=(0,10))
end

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

# ╔═╡ Cell order:
# ╟─ed174bc4-1ea0-11eb-1e2f-a32874cec549
# ╠═6fd616f6-1ea1-11eb-3814-8bfb4a096c49
# ╠═6ffe2628-1ea1-11eb-24ea-57f985146a72
# ╠═70179b62-1ea1-11eb-13ce-7fb4d54bad4b
# ╠═884176a0-21c0-11eb-06f8-55964a60d533
# ╠═8841a7b0-21c0-11eb-313a-bfa4eb49f351
# ╟─70182cd8-1ea1-11eb-094d-8d0d49cf15f3


abstract type Angle <: Number end

struct Radian <: Angle
    radians::Float64
end

struct DMS <: Angle
    seconds::Int
end

begin
    Degree(degrees::Integer) = Minute(degrees * 60)
    Degree(deg::Integer, min::Integer) = Degree(deg) + Minute(min)
    Degree(deg::Integer, min::Integer, secs::Integer) = Degree(deg, min) + Second(secs)

    function Minute(minutes::Integer)
        DMS(minutes * 60)
    end

    function Second(seconds::Integer)
        DMS(seconds)
    end

    import Base: -, +, show, convert, *, /, promote_rule
    +(Θ::DMS, α::DMS) = DMS(Θ.seconds + α.seconds)
    -(Θ::DMS, α::DMS) = DMS(Θ.seconds - α.seconds)

    +(Θ::Radian, α::Radian) = Radian(Θ.radians + α.radians)
    -(Θ::Radian, α::Radian) = Radian(Θ.radians - α.radians)

    function degrees(dms::DMS)
        minutes = dms.seconds ÷ 60
        minutes ÷ 60
    end

    function minutes(dms::DMS)
        minutes = dms.seconds ÷ 60
        minutes % 60
    end

    seconds(dms::DMS) = dms.seconds % 60

    function show(io::IO, dms::DMS)
    print(io, degrees(dms), "° ", minutes(dms), "' ", seconds(dms), "''")
    end

    function show(io::IO, rad::Radian)
    print(io, rad.radians, "rad")
    end

    # convert constructors
    Radian(dms::DMS) = Radian(deg2rad(dms.seconds/3600))
    DMS(rad::Radian) = DMS(floor(Int, rad2deg(rad.radians) * 3600))

    convert(::Type{Radian},  dms::DMS)    = Radian(dms)
    convert(::Type{DMS},     rad::Radian) = DMS(rad)

    sin(rad::Radian) = Base.sin(rad.radians)
    cos(rad::Radian) = Base.cos(rad.radians)

    sin(dms::DMS) = sin(Radian(dms))
    cos(dms::DMS) = cos(Radian(dms))

    *(coeff::Number, dms::DMS) = DMS(coeff * dms.seconds)
    *(dms::DMS, coeff::Number) = coeff * dms
    /(dms::DMS, denom::Number) = DMS(dms.seconds/denom)

    *(coeff::Number, rad::Radian) = Radian(coeff * rad.radians)
    *(rad::Radian, coeff::Number) = coeff * rad
    /(rad::Radian, denom::Number) = Radian(rad.radians/denom)

    const ° = Degree(1)
    const rad = Radian(1)

    Base.promote_rule(::Type{Radian}, ::Type{DMS}) = Radian
end
# Examples
# α = Degree(90, 30, 45)
# β = α + Degree(89, 29, 15)
# sin(Degree(90))
# sin(Radian(π/2))
# sin(90°)
# sin(1.5rad)

# Problem
# While this works:
# +(promote(90°,3.14rad/2)...)
# 3.140796326794897rad
# Doing this
# 90°+ 3.14rad/2
# gives this error
# ERROR: MethodError: no method matching +(::DMS, ::Radian)
# Closest candidates are:
#   +(::Any, ::Any, ::Any, ::Any...) at operators.jl:560
#   +(::DMS, ::DMS) at REPL[10]:1
#   +(::Radian, ::Radian) at REPL[12]:1
# Stacktrace:
#  [1] top-level scope
#    @ REPL[38]:1
# Solution
# For the solution to the problem see the discussion here
# https://discourse.julialang.org/t/type-promotion-not-working-as-expected/65944
# one solution to get the promotion to work is to subtype Angle from Number as above
# Adding +(Angle, Angle) is probably more in the spirit of what Engheim is trying to do


"""
`solve_system(f,u0,n)` #code

Solves the dynamical system

``u_{n+1} = f(u_n)`` #latex

for N steps. Returns the solution at step `n` with parameters `p`.

"""
function solve_system(f,u0,p,n)
  u = u0
  for i in 1:n-1
    u = f(u,p)
  end
  u
end

f(u,p) = u^2 - p*u
typeof(f)

function lorenz(u,p)
  α,σ,ρ,β = p
  du1 = u[1] + α*(σ*(u[2]-u[1]))
  du2 = u[2] + α*(u[1]*(ρ-u[3]) - u[2])
  du3 = u[3] + α*(u[1]*u[2] - β*u[3])
  [du1,du2,du3]
end
p = (0.02,10.0,28.0,8/3)

function solve_system_save(f,u0,p,n)
  u = Vector{typeof(u0)}(undef,n)
  u[1] = u0
  for i in 1:n-1
    u[i+1] = f(u[i],p)
  end
  u
end

a = solve_system_save(lorenz,[1.0,0.0,0.0],p,1000)

a[1]

reduce(hcat,a)

using BenchmarkTools
function solve_system_save_matrix(f,u0,p,n)
  u = Matrix{eltype(u0)}(undef,length(u0),n)
  u[:,1] = u0
  for i in 1:n-1
    u[:,i+1] = f(u[:,i],p)
  end
  u
end
@btime solve_system_save_matrix(lorenz,[1.0,0.0,0.0],p,1000)
@btime solve_system_save(lorenz,[1.0,0.0,0.0],p,1000)

function solve_system_save_matrix_view(f,u0,p,n)
  u = Matrix{eltype(u0)}(undef,length(u0),n)
  u[:,1] = u0
  for i in 1:n-1
    u[:,i+1] = f(@view(u[:,i]),p)
  end
  u
end
@btime  solve_system_save_matrix_view(lorenz,[1.0,0.0,0.0],p,1000)

function lorenz(du,u,p)
  α,σ,ρ,β = p
  du[1] = u[1] + α*(σ*(u[2]-u[1]))
  du[2] = u[2] + α*(u[1]*(ρ-u[3]) - u[2])
  du[3] = u[3] + α*(u[1]*u[2] - β*u[3])
end
p = (0.02,10.0,28.0,8/3)
function solve_system_save(f,u0,p,n)
  u = Vector{typeof(u0)}(undef,n)
  du = similar(u0) # create cache array, modify it, add to output
  u[1] = u0
  for i in 1:n-1
    f(du,u[i],p)
    u[i+1] = du
  end
  u
end
@btime solve_system_save(lorenz,[1.0,0.0,0.0],p,1000)

function solve_system_save(f,u0,p,n)
  u = Vector{typeof(u0)}(undef,n)
  du = similar(u0) # create cache array, modify it, add to output
  u[1] = u0
  for i in 1:n-1
    f(du,u[i],p)
    u[i+1] = du[:]
  end
  u
end
@btime solve_system_save(lorenz,[1.0,0.0,0.0],p,1000)

b = solve_system_save(lorenz,[1.0,0.0,0.0],p,1000)


using PenaltyFunctions, LaTeXStrings, Plots
pyplot()

_mkpath(fname) = joinpath(dirname(@__FILE__()), fname)
srand(123)

# ---------------------------------------------------------------
# README Plots

fnt = Plots.Font("ubuntu",8,:hcenter,:vcenter,0.0,
                 RGB{FixedPointNumbers.Normed{UInt8, 8}}(0.0,0.0,0.0))
default(bg_outside = RGBA(1,1,1,0),
        grid = false,
        legendfont = fnt,
        palette = :darktest,
        aspect_ratio = 1,
        size = (400, 400))

p = plot(L1Penalty(); ylim = (0, 8))
for pen in [L2Penalty(),
            ElasticNetPenalty(.5),
            MCPPenalty(1.), SCADPenalty(3.7, 0.4),
            LogPenalty(1.)]
    plot!(p, pen)
end
plot!(p, xlab = L"\theta_1", ylab = L"g(\theta_1)", xlim = (-4, 4))

p[1][1][:linestyle] = :solid
p[1][2][:linestyle] = :solid
p[1][3][:linestyle] = :dash
p[1][4][:linestyle] = :dashdot
p[1][5][:linestyle] = :dash
p[1][6][:linestyle] = :dot

Plots.svg(_mkpath("univariate"))

function cont_xy(penalty, stepx = 0.005, stepy = 0.0001)
    x = -1:stepx:1
    yp = 0:stepy:1
    ref = round(value(penalty, [1., 0]), 4)
    xs = collect(i for j in yp, i in x if round(value(penalty, [i,j]), 4) ≈ ref)
    ys = collect(j for j in yp, i in x if round(value(penalty, [i,j]), 4) ≈ ref)
    vcat(xs, reverse(xs)), vcat(ys, -ys)
end

default(size = (400, 400))
p = plot(cont_xy(L1Penalty())...)
plot!(cont_xy(L2Penalty())...)
plot!(cont_xy(ElasticNetPenalty(0.5))...)
plot!(cont_xy(MCPPenalty(1.))...)
plot!(cont_xy(SCADPenalty(3.7, 0.4))...)
plot!(cont_xy(LogPenalty(1.))...)
plot!(p, xlab = L"\theta_1", ylab = L"\theta_2",
         legend = false,
         xlim = (-1.2, 1.2), ylim = (-1.2, 1.2))

p[1][1][:linestyle] = :solid
p[1][2][:linestyle] = :solid
p[1][3][:linestyle] = :dash
p[1][4][:linestyle] = :dashdot
p[1][5][:linestyle] = :dash
p[1][6][:linestyle] = :dot

Plots.svg(_mkpath("bivariate"))


# Examples presented in class - Lecture 1
# LP Resource allocation

using JuMP, Cbc
# JuMP is for implementing math. programming models;
# Cbc is for solving them.

# Example 1 - resource allocation

# Problem data
i = 1:2 # i=1: Seattle; i=2: San Diego
j = 1:3 # j=1: New York; j=2: Chicago; j=3: Miami

C = [350 600] # Capacities of the factories
D = [325 300 275] # Demand of clients
T = [2.5 1.7 1.8
     3.5 1.9 1.4] # Transportation costs

# Model implementation
# Creates a model and informs the solver to be used
m = Model(solver = CbcSolver())

# Decision variable for the total transported
@variable(m, x[i,j] >= 0)

# Capacity constraint
@constraint(m, cap[i = 1:2], sum(x[i,j] for j = 1:3) <= C[i])
# Demand constraint
@constraint(m, dem[j = 1:3], sum(x[i,j] for i = 1:2) >= D[j])

# Total distribution cost that we want to minimise
@objective(m, Min, sum(T[i,j]*x[i,j] for i=1:2, j=1:3))

println(m) # Prints the mathematical model for debugging
solve(m) # Solve the model

# Prints the optimal solution
println("\nDistribution plan: \n", getvalue(x))


struct Connection
    weight::AbstractArray
    adjacency::AbstractArray
end

function randomConnection(n_source::Int, n_target::Int, p::AbstractFloat;
    min_weight=0, max_weight=1, amplitude=nothing, seed=nothing)
    if !isnothing(seed)
        Random.seed!(seed)
    end
    adjacency = rand(n_source, n_target) .< p
    if isnothing(amplitude)
        weights = rand(n_source, n_target) .* (max_weight - min_weight) .+ min_weight
    else
        weights = ones(n_source, n_target) .* amplitude
    end
    weights .*= adjacency
    return Connection(weights, adjacency)
end

function randomConnection(n::Int, p::AbstractFloat; min_weight=0, max_weight=1,
    amplitude=nothing, seed=nothing)
    randomConnection(n, n, p; min_weight, max_weight, amplitude, seed)
end

function EI(Nₑ::Integer, Nᵢ::Integer;
    E_E::Tuple=(0, 1), E_I::Tuple=(0, 1), I_E::Tuple=(0, 1), I_I::Tuple=(0, 1),
    adjacency::Union{Nothing,AbstractMatrix}=nothing, shuffle::Bool=false, seed=nothing
)
    excitatory = [ones(Bool, Nₑ); zeros(Bool, Nᵢ)]
    if shuffle
        if !isnothing(seed)
            Random.seed!(seed)
        end
        shuffle!(excitatory)
    end
    inhibitory = .!excitatory

    min_weight = (excitatory * transpose(excitatory)) .* E_E[1]
    min_weight += (excitatory * transpose(inhibitory)) .* E_I[1]
    min_weight -= (inhibitory * transpose(excitatory)) .* abs(I_E[2])
    min_weight -= (inhibitory * transpose(inhibitory)) .* abs(I_I[2])

    max_weight = (excitatory * transpose(excitatory)) .* E_E[2]
    max_weight += (excitatory * transpose(inhibitory)) .* E_I[2]
    max_weight -= (inhibitory * transpose(excitatory)) .* abs(I_E[1])
    max_weight -= (inhibitory * transpose(inhibitory)) .* abs(I_I[1])

    if !isnothing(adjacency)
        min_weight .*= adjacency
        max_weight .*= adjacency
    end

    return excitatory, inhibitory, min_weight, max_weight
end

function EI(n::Integer, EI_rate::AbstractFloat;
    E_E::Tuple=(0, 1), E_I::Tuple=(0, 1), I_E::Tuple=(0, 1), I_I::Tuple=(0, 1),
    adjacency::Union{Nothing,AbstractMatrix}=nothing, shuffle::Bool=false, seed=nothing
)
    if EI_rate > 1
        Nₑ = Int(ceil(EI_rate * n / (EI_rate + 1)))
    else
        Nₑ = Int(ceil(EI_rate * n))
    end
    Nᵢ = n - Nₑ
    return EI(Nₑ, Nᵢ; E_E, E_I, I_E, I_I, adjacency, shuffle, seed)
end


@doc raw"""
Compute δL = L(0)-L(h) where L is the quadratic model

```math
L(h)=f(\theta) + \langle \nabla f, h \rangle + \frac{1}{2}\langle \nabla^2 f h, h \rangle
```

with ``f(\theta)=\frac{1}{2}\| r(θ) \|_2^2``, ``\nabla f = J^t r`` and ``\nabla^2 f = J^t J``

A direct computation gives:

```math
δL = L(0)-L(h) = -\left(  \langle J^tr, h \rangle + \frac{1}{2}\langle \nabla^2 f h, h \rangle \right)
```

However one can avoid the computation of ``\nabla^2 f h`` if one uses the
fact that ``h`` is solution of:

```math
(\nabla^2 f + \mu I)h + \nabla f = 0
```

With this hypothesis, one gets:

```math
δL = L(0)-L(h) = \frac{1}{2} \langle h, \mu h - \nabla f \rangle
```
"""
function compute_δL_unconstrained(∇f::AbstractVector,
                                  μ::Real,
                                  h::AbstractVector)

    # The syntax to desctructure is unexpected...
    # https://discourse.julialang.org/t/argument-destructuring-and-anonymous-functions/24893/2
    #
    mapreduce(((h_i,∇f_i),)->h_i*(μ*h_i-∇f_i),+,zip(h,∇f))/2
end

@doc raw"""

Same idea than [`compute_δL_unconstrained`](@ref), however when bound
constraints are present ``h`` is such that:

```math
(\nabla^2 f + \mu I)h + \nabla f + \tau = 0
```
it follows that:

```math
δL = L(0)-L(h) = \frac{1}{2} \langle h, \mu h + \tau - \nabla f \rangle
```
"""
function compute_δL_constrained(∇f::AbstractVector,
                                μ::Real,
                                τ::AbstractVector,
                                h::AbstractVector)
    mapreduce(((h_i,τ_i,∇f_i),)->h_i*(μ*h_i+τ_i-∇f_i),+,zip(h,τ,∇f))/2
end

@doc raw"""

Compute true variation of the real model: ``δf = \frac{1}{2}(r^t(θ)r(θ)-r^t(θ+h)r(θ+h))``

Contrary to ``δL`` things are simpler. However a trick is to use an equivalent formulation:

```math
δf = \frac{1}{2}(r^t(θ)r(θ)-r^t(θ+h)r(θ+h)) = \frac{1}{2}(r(θ)-r(θ+h))^t(r(θ)+r(θ+h))
```

that has a better numerical behavior.
"""
function compute_δf(r::AbstractVector,
                    r_new::AbstractVector)
    mapreduce(((r_i,r_new_i),)->(r_i-r_new_i)*(r_i+r_new_i),+,zip(r,r_new))/2
end


workspace()

using IntervalArithmetic
using EAGOParametricInterval

# Miranda Passed
X = [Interval(1,3),Interval(1,3)]
P = [Interval(1,3),Interval(1,3)]
h(z,p) = [p[1]-z[1]; p[1]+z[1]]
h1(z,p) = [p[1]-z[1]; p[1]-z[1]]
test_mir1 = Miranda(h,X,P)
test_mir2 = Miranda(h1,X,P)
out_mir = h(X,P)

# Miranda Exclusion Test Passed
Eflag_in = false
ex_mir1 = MirandaExc(h,X,P,Eflag_in)
ex_mir2 = MirandaExc(h1,X,P,Eflag_in)
ex_mir_flag_out = Eflag_in

#
Y1 = [Interval(1,3),Interval(1,3)]
Y2 = [Interval(1.5,2.5),Interval(2.5,2.5)]
Y3 = [Interval(-1,2),Interval(-2,2)]
Y4 = [Interval(2,4),Interval(2,4)]
Y5 = [Interval(1.5,2.5),Interval(2.3,2.5)]
Y6 = [Interval(-1,4),Interval(-1,4)]
Y7 = [Interval(1.5,2.5),emptyinterval()]
Y8 = [emptyinterval(),emptyinterval()]
Y9 = [Interval(4,5),Interval(4,5)]
Y10 = [Interval(-2,-1),Interval(-2,-1)]

sXY_flag1 = Strict_XinY(Y1,X)
sXY_flag2 = Strict_XinY(Y2,X)
sXY_flag3 = Strict_XinY(Y3,X)
sXY_flag4 = Strict_XinY(Y4,X)
sXY_flag5 = Strict_XinY(Y5,X)
sXY_flag6 = Strict_XinY(Y6,X)
sXY_flag7 = Strict_XinY(Y7,X)
sXY_flag8 = Strict_XinY(Y8,X)
sXY_flag9 = Strict_XinY(Y9,X)
sXY_flag10 = Strict_XinY(Y10,X)

#=
partialIncTop(h,X,P,PIflag,incHigh)

partialIncBot(h,X,P,PIflag,incLow)

Strict_XinY(X,Y)

isEqual(X1,X2,atol)

extDivide(A,B,C)

extProcess(N,X,Mii,S1,S2,B,rtol)
=#


using Polylogarithms
using SpecialFunctions
using Test
using DataFrames, CSV
import Base.MathConstants: π, pi, ℯ, e, γ, eulergamma, catalan, φ, golden
include("test_defs.jl")

@testset "Bernoulli numbers" begin
    @testset "    throws errors" begin
        @test_throws DomainError bernoulli(-1)
        @test_throws DomainError bernoulli(36)
        @test_throws MethodError bernoulli(1.0)
    end

    @testset "    types" begin
        @test typeof(bernoulli(2)) == Rational{Int64}
    end

    @testset "    explicit cases" begin
        for i=2:34
            # zeta isn't quite as accurate as we would like
            @test Float64(bernoulli(i)) ≈ -i*zeta( 1 - i )
        end
    end

    @testset "    dataset 1" begin
        B = Symbol("B_n")
        data1 = CSV.read("../data/bernoulli_test_data_1.csv", DataFrame)
        m = size(data1,1)
        for i=1:m
            @test bernoulli( data1[i,:n] ) ≅ parse(Rational{Int64},data1[i,B])
        end
    end
end

@testset "Bernoulli polynomials" begin
    @testset "    throws errors" begin
        @test_throws DomainError bernoulli(-1, 0.0)
        @test_throws MethodError bernoulli(1.0, 0.0)
        @test_throws MethodError bernoulli(1.0, 1)
    end

    @testset "    types" begin
        @test typeof(bernoulli(2, 1.0)) == Float64
    end

    @testset "    consistency with Bernoulli numbers" begin
        for i=0:14
            @test Float64(bernoulli(i)) ≅ bernoulli(i, 0.0)
        end
        for i=15:34
            # somewhat reliant here on accuracy of SpecialFunctions, so can't make these better without rewriting that
            @test Float64(bernoulli(i)) ≈ bernoulli(i, 0.0)
        end
     end

    @testset "    explicit cases" begin
        X = collect(-2.0 : 0.1 : 3.0)
        for i=1:length(X)
            x = X[i]
            # println("   x = $x")
            @test bernoulli(0, x) ≅ 1.0
            @test bernoulli(1, x) ≅ x   - 0.5
            @test bernoulli(2, x) ≅ x^2 -     x   + 1.0/6.0
            @test bernoulli(3, x) ≅ x^3 - 1.5*x^2 + 0.5*x
            @test bernoulli(4, x) ≅ x^4 - 2.0*x^3 +     x^2 - 1/30.0
            @test bernoulli(5, x) ≅ x^5 - 2.5*x^4 + (5.0/3.0)*x^3 - x/6.0
            # @test bernoulli(6, x) ≈ x^6 - 3.0*x^5 + (5.0/2.0)*x^4 - 0.5*x^2 + 1/42.0
            # Hurwitz-zeta doesn't work well enough here
        end
    end

    @testset "    symmetry identities" begin
        X = collect(0.0: 0.1 : 1.0)
        for i=1:length(X)
            x = X[i]
            # for n=1:9
            for n=1:6
                # println("   n = $n, x = $x")
                # start to get errors of order 1.0e-14 for n=5, 1.0e-13 around n=10, ...
                @test bernoulli(n, 1-x)  ≈  (-1.0)^n * bernoulli(n, x)
                @test bernoulli(n, x+1)  ≈  bernoulli(n, x) + n*x^(n-1)
                # SpecialFunctions.zeta(-6, -1.0) = NaN (probably should be 2)
                @test (-1)^n * bernoulli(n, -x)  ≈   bernoulli(n, x) + n*x^(n-1)

                # Raabe (1851)
                m = 6.0
                k = 0:m-1
                @test sum( bernoulli.(n, x .+ k./m ) )/m  ≈  m^(-n)*bernoulli(n, m*x )
            end
        end
    end

    @testset " dataset 2" begin
        B = Symbol("B_n(x)")
        data2 = CSV.read("../data/bernoulli_test_data_2.csv", DataFrame; delim=",", type=String)
        data2[!,:n] = parse.(Float64, data2[!,:n] )
        data2[!,:x] = parse.(Float64, data2[!,:x] )
        data2[!,B] = parse.(Float64, data2[!,B] )
        m = size(data2,1)
        for i=1:m
            @test bernoulli( Int(data2[i,:n]),  data2[i,:x]) ≈ data2[i,B]
        end
    end

    @testset " dataset 3" begin
        B = Symbol("B_n(x)")
        data3 = CSV.read("../data/bernoulli_test_data_3.csv", DataFrame; delim=",")
        m = size(data3,1)
        for i=1:m
            @test bernoulli( Int(data3[i,:n]),  data3[i,:x]) ≈ data3[i,B]
        end
   end
end


using IntervalWavelets
using Wavelets
using Base.Test

println("Testing dyadic rationals...")

#=
Translation between indices and values for dyadic rationals
=#

p = 2
S = DaubSupport(-p+1, p)

# Level 0 = integers
supportx = dyadic_rationals( S, 0 )
for x in round(Int, supportx)
	@test supportx[ x2index(x,S) ] == x
end

for idx in 1:length(supportx)
	@test supportx[idx] == index2x(idx,S)
end

# Higher level
R = rand( 1:10 )

supportx = dyadic_rationals( S, R )
for x in supportx
	@test supportx[ x2index(x,S,R) ] == x
end

for idx in 1:length(supportx)
	@test supportx[idx] == index2x(idx,S,R)
end


# indicator of the L0 norm ball with given (integer) radius

export IndBallL0

"""
**Indicator of a ``L_0`` pseudo-norm ball**

    IndBallL0(r=1)

Returns the indicator function of the set
```math
S = \\{ x : \\mathrm{nnz}(x) \\leq r \\}.
```
Parameter `r` must be a positive integer.
"""
struct IndBallL0{I <: Integer} <: ProximableFunction
  r::I
  function IndBallL0{I}(r::I) where {I <: Integer}
    if r <= 0
      error("parameter r must be a positive integer")
    else
      new(r)
    end
  end
end

is_set(f::IndBallL0) = true

IndBallL0(r::I) where {I <: Integer} = IndBallL0{I}(r)

function (f::IndBallL0)(x::AbstractArray{T}) where T <: RealOrComplex
  if count(!isequal(0), x) > f.r
    return +Inf
  end
  return zero(T)
end

function prox!(y::AbstractArray{T}, f::IndBallL0, x::AbstractArray{T}, gamma::Real=1.0) where T <: RealOrComplex
  p = []
  if ndims(x) == 1
    p = partialsortperm(x, 1:f.r, by=abs, rev=true)
  else
    p = partialsortperm(x[:], 1:f.r, by=abs, rev=true)
  end
  sort!(p)
  idx = 1
  for i = 1:length(p)
    y[idx:p[i]-1] .= zero(T)
    y[p[i]] = x[p[i]]
    idx = p[i]+1
  end
  y[idx:end] .= zero(T)
  return zero(T)
end

fun_name(f::IndBallL0) = "indicator of an L0 pseudo-norm ball"
fun_dom(f::IndBallL0) = "AbstractArray{Real}, AbstractArray{Complex}"
fun_expr(f::IndBallL0) = "x ↦ 0 if countnz(x) ⩽ r, +∞ otherwise"
fun_params(f::IndBallL0) = "r = $(f.r)"

function prox_naive(f::IndBallL0, x::AbstractArray{T}, gamma::Real=1.0) where T <: RealOrComplex
  p = sortperm(abs.(x)[:], rev=true)
  y = similar(x)
  y[p[1:f.r]] .= x[p[1:f.r]]
  y[p[f.r+1:end]] .= 0.0
  return y, 0.0
end


using RvSpectraKitLearn

function make_B(p::Integer, q::Integer; frac_active::Float64 = 0.1, bias::Float64 = 0.0, sigma::Float64 = 1.0)
  B = zeros(p,q)
  for i in 1:p, j in 1:q
      if rand() < frac_active
        B[i,j] = bias+sigma*randn()
      end
  end
  return B
end

# Test sparse regression via  alg
n = 25
 p = 100
 q = 4
 frac_active = 0.1
 sigma_x = 0.1
 sigma_y = 0.1
 X = randn(n,p)
 B = make_B(p,q,frac_active=frac_active,bias=1.0,sigma=sigma_x)
 Y = X*B + sigma_y*randn(n,q)
 X_crossvalid = randn(n,p)
 Y_crossvalid = X_crossvalid*B + sigma_y*randn(n,q)
 println("B norms: ", length(find(B)), "  ", sumabs(B), " ", sumabs2(B))

(gamma_list, Z_list)  = fit_arp_slr_2d(X,Y,stepsize=1.04,max_itterations=1000, min_active=floor(Int64,p*q*frac_active*0.1))
(L0,L1,L2,rmsdelta) = compute_alr_path_stats(Z_list,target=B)
(chisq,chisq_refit,chisq_cv) = compute_alr_path_chisq!(Z_list,X,Y,X_crossvalid,Y_crossvalid)
print_alr_path(gamma_list, Z_list,X,Y)

#=
#Pkg.add("PyPlot")  # Run once if you don't already have PyPlot installed
using PyPlot

plt[:loglog](gamma_list,L0,"y-")
 plt[:loglog](gamma_list,L1/length(B),"r-")
 plt[:loglog](gamma_list,sqrt(L2),"g-")
 plt[:loglog](gamma_list,chisq_refit,"b-")
 plt[:loglog](gamma_list,chisq_refit,"b.")
 plt[:loglog](gamma_list,chisq_cv,"m-")
 plt[:loglog](gamma_list,chisq_cv,"m.")

plot(L0,chisq,"r-")
 plot(L0,chisq,"r.")
 plot(L0,chisq_refit,"b-")
 plot(L0,chisq_refit,"b.")
 plot(L0,chisq_cv,"m-")
 plot(L0,chisq_cv,"m.")
 plt[:loglog](L0,chisq_cv,"m.")

plot(L1,chisq,"r-")
 plot(L1,chisq,"r.")
 plot(L1,chisq_refit,"b-")
 plot(L1,chisq_refit,"b.")
 plot(L1,chisq_cv,"m-")
 plot(L1,chisq_cv,"m.")
 plt[:loglog](L1,chisq_cv,"m.")

plot(L0,L1./L0,"g-")
 plot(L0,L1./L0,"g.")

plot(L0,chisq,"r-")
 plot(L0,chisq,"r.")
 plot(L0,chisq_refit,"b-")
 plot(L0,chisq_refit,"b.")
 plot(L0,chisq_cv,"m-")
 plot(L0,chisq_cv,"m.")
 plt[:loglog](L0,chisq_cv,"m.")

=#


"The set of all natural numbers."
struct NaturalNumbers <: Domain{Int}
end

in(x, d::NaturalNumbers) = isinteger(x) && (x >= 0)
in(x::AbstractArray, d::NaturalNumbers) = false
==(::NaturalNumbers, ::NaturalNumbers) = true

approx_in(x::Real, d::NaturalNumbers, tol) =
    (abs(x-round(x)) < tol) && (round(Int, x) ∈ d)

"The set of all integers."
struct Integers <: Domain{Int}
end

in(x, d::Integers) = isinteger(x)
in(x::AbstractArray, d::Integers) = false
==(::Integers, ::Integers) = true

approx_in(x::Real, d::Integers, tol) =
    (abs(x-round(x)) < tol) && (round(Int, x) ∈ d)

"The set of all real numbers."
struct RealNumbers <: Domain{Float64}
end

in(x::Number, d::RealNumbers) = isreal(x)
# isreal also allows real arrays, we want to disallow that here:
in(x::AbstractArray, d::RealNumbers) = false

approx_in(x::Complex, d::RealNumbers, tol) = (imag(x) < tol) && (real(x) ∈ d)
==(::RealNumbers, ::RealNumbers) = true

"The set of all rationals."
struct Rationals <: Domain{Rational{Int}}
end

in(x, d::Rationals) = x ∈ Integers()
in(x::Rational, d::Rationals) = true
==(::Rationals, ::Rationals) = true

"The set of all complex numbers whose real and imaginary parts are real numbers."
struct ComplexNumbers <: Domain{Complex{Float64}}
end

in(x::Complex{T}, d::ComplexNumbers) where {T} =
    isreal(real(x)) && isreal(imag(x))
in(x::Complex{T}, d::ComplexNumbers) where {T<:Real} = true
in(x, d::ComplexNumbers) = x ∈ RealNumbers()
==(::ComplexNumbers, ::ComplexNumbers) = true

"The set of natural numbers."
const ℕ = NaturalNumbers()
"The set of integers."
const ℤ = Integers()
"The set of rational numbers."
const ℚ = Rationals()
"The set of real numbers."
const ℝ = RealNumbers()
"The set of complex numbers."
const ℂ = ComplexNumbers()

"The space ℝ^1."
const ℝ1 = VcatDomain(ℝ)
"The space ℝ^2."
const ℝ2 = VcatDomain(ℝ, ℝ)
"The space ℝ^3."
const ℝ3 = VcatDomain(ℝ, ℝ, ℝ)
"The space ℝ^4."
const ℝ4 = VcatDomain(ℝ, ℝ, ℝ, ℝ)


using Base.Test
using PValueAdjust

newPvalues = padjust([0.05, 0.03, 0.01, 0.5], BenjaminiHochberg)
@test_approx_eq newPvalues [0.06666666667, 0.06, 0.04, 0.5]

@test_approx_eq padjust([1, 0.5, 0.1, 0.001, 0.0001, 0.002], BenjaminiHochberg) [1.0, 0.6, 0.15, 0.003, 0.0006, 0.004]
@test_approx_eq padjust([0.05, 0.05, 0.05], BenjaminiHochberg) [0.05, 0.05, 0.05]

@test all(padjust([1,1,1], BenjaminiHochberg) .<= 1)

@test_throws DomainError padjust([-1], BenjaminiHochberg)
@test_throws DomainError padjust([2], BenjaminiHochberg)

@test padjust([0.5], BenjaminiHochberg) == [0.5]


function tinterval(x, δ, op)
    μ = mean(x)
    σ = std(x, mean=μ)
    t = 0.0
    n = length(x)
    Zygote.ignore() do
        t = quantile(TDist(n-1), δ)
    end
    s = t * σ / √n
    return op(μ,s)
end

function tdist_ci(;δ=0.05, tail=:both)
    if tail == :right
        return x->tinterval(x,δ,(μ,s)->μ-s) # quantile of 0.05 is negative so need to subtract to get upper
    elseif tail == :left
        return x->tinterval(x,δ,(μ,s)->μ+s) # need to add to get lower bound
    else
        return x->tinterval(x,δ/2,(μ,s)->(μ+s,μ-s))
    end
end

function andersons_ucb(x, b, δ)
    n = length(x)
    ϵ = √(log(1/δ)/(2n))
    r = floor(Int, n*ϵ)
    upper = (1/n) * ((r+1)*x[r+1] + sum(view(x, r+1:n))) + ϵ * (b - x[r+1])
    return upper
end

function andersons_lcb(x, a, δ)
    n = length(x)
    ϵ = √(log(1/δ)/(2n))
    s = floor(Int, n*ϵ)
    lower = (1/n) * (sum(view(x, 1:n-s-1)) + (s+1)*x[n-s]) - ϵ * (x[n-s] - a)
    return lower
end

function andersons_ci(a, b; δ=0.05, tail=:both)
    if tail == :both
        function f(x)
            z = sort(x)
            upper = andersons_ucb(z, b, δ/2)
            lower = andersons_lcb(z, a, δ/2)
            return lower, upper
        end
        return f
    elseif tail == :left
        return x->andersons_lcb(sort(x), a, δ)
    elseif tail == :right
        return x->andersons_ucb(sort(x), b, δ)
    else
        throw(ArgumentError("tail=$(tail) is invalid"))
    end
end


# Squared Exponential Function with ARD

"""
    SEArd <: StationaryARD{WeightedSqEuclidean}

ARD Squared Exponential kernel (covariance)
```math
k(x,x') = σ²\\exp(- (x - x')ᵀL⁻²(x - x')/2)
```
with length scale ``ℓ = (ℓ₁, ℓ₂, …)`` and signal standard deviation ``σ`` where
``L = diag(ℓ₁, ℓ₂, …)``.
"""
mutable struct SEArd{T<:Real} <: StationaryARD{WeightedSqEuclidean}
    "Inverse squared length scale"
    iℓ2::Vector{T}
    "Signal variance"
    σ2::T
    "Priors for kernel parameters"
    priors::Array
end

"""
Squared Exponential Function with ARD

    SEArd(ll::Vector{Real}, lσ::Real)

# Arguments
  - `ll::Vector{Real}`: vector of length scales (given on log scale)
  - `lσ::Real`: signal standard deviation (given on log scale)
"""
SEArd(ll::Vector{T}, lσ::T) where T = SEArd{T}(exp.(-2 .* ll), exp(2 * lσ), [])

function set_params!(se::SEArd, hyp::AbstractVector)
    length(hyp) == num_params(se) || throw(ArgumentError("SEArd has $(num_params(se)) parameters, received $(length(hyp))."))
    @views @. se.iℓ2 = exp(-2 * hyp[1:(end-1)])
    se.σ2 = exp(2 * hyp[end])
end

get_params(se::SEArd) = [-log.(se.iℓ2) / 2 ; log(se.σ2) / 2]
get_param_names(k::SEArd) = [get_param_names(k.iℓ2, :ll); :lσ]
num_params(se::SEArd) = length(se.iℓ2) + 1

cov(se::SEArd, r::Number) = se.σ2*exp(-r / 2)

@inline dk_dll(se::SEArd, r::Real, wdiffp::Real) = wdiffp*cov(se,r)
@inline function dKij_dθp(se::SEArd, X1::AbstractMatrix, X2::AbstractMatrix, i::Int, j::Int, p::Int, dim::Int)
    if p <= dim
        return dk_dll(se, distij(metric(se),X1,X2,i,j,dim), distijk(metric(se),X1,X2,i,j,p))
    elseif p==dim+1
        return dk_dlσ(se, distij(metric(se),X1,X2,i,j,dim))
    else
        return NaN
    end
end


function bitCount(n::Int)
	count = 0
	while n>0
		count += n & 1
		n >>= 1
	end
	return count
end

function findFirstBit(n::Int)
	n > 0 || throw(DomainError(" n needs to be larger than 0"))

	breakcondition = 0
	searchposition = 1
	idx = 0
	while breakcondition==0
		breakcondition = n & searchposition
		searchposition <<= 1
		idx += 1
	end

	return idx
end

function getBitPos(n::Int)

	Pos = Vector{Int}([])
	while n > 0

		idx = findFirstBit(n)
		push!(Pos,idx)

		n -= 2^(idx-1)

	end

	return Pos
end

"""
    bitSwap(n, i, j)

swaps bit i with j in n, for e.g.

n = 001001 , i = 1, j = 2 results in bitSwap(n,i,j) = 001010

#Arguments
- n: bit-string represented as an integer
- i: position of the first bit which should be swaped with the second
- j: position of the second bit which should be swaped with the first

"""
function bitSwap(n::Int, i::Int, j::Int)
	i -= 1
	j -= 1

	bit1 = (n >> i) & 1
	bit2 = (n >> j) & 1

	x = bit1 ⊻ bit2

	x = (x << i) | (x << j)

	return (n ⊻ x)
end

function FermionicBitSwap(n::Int, j::Int, k::Int)
	n_new = bitSwap(n,j,k)
	untouchedBit =  n & n_new
	inbetween = j<k ? (j+1:k-1) : (k+1:j-1)
	c = Int(sum(2.0.^[x-1 for x in inbetween]))
	overjumped_fermions = bitCount(untouchedBit & c)
	return n_new, (-1)^overjumped_fermions
end


# Function to normalize values
function normalize(vec)
    [(x - minimum(vec))/(maximum(vec) - minimum(vec)) for x in vec]
end


using Mdp.VI
@testset "Dict inputs" begin
    # Define simple domains (from Massau Kolobov 2016 book, p.41)
    S = 1:6 # States
    A = 1:2
    T = Dict(   # Transitions, prob = T[(state, action, next_state)]
        (1,1,3)=>1.,
        (1,2,2)=>1.,
        (2,1,3)=>1.,
        (3,1,5)=>1.,
        (3,2,2)=>1.,
        (4,1,5)=>1.,
        (5,1,6)=>1.,
        (5,2,6)=>.6,
        (5,2,4)=>.4,
        (6,1,6)=>1.,    # Manually added goal state transition
    )
    C = Dict( # Costs, cost = C[(state, action, next_state)]
        (1,1,3)=>1,
        (1,2,2)=>1,
        (2,1,3)=>1,
        (3,1,5)=>1,
        (3,2,2)=>1,
        (4,1,5)=>1,
        (5,1,6)=>5,
        (5,2,6)=>2,
        (5,2,4)=>2,
        (6,1,6)=>0,   # Manually added goal state transition cost
    )

    V_init = [3.,3.,2.,2.,1.,0.]
    @testset "Bellman backup" begin
        @test bellman_backup(1, S, A, T, C, V_init)==3.
        @test bellman_backup(2, S, A, T, C, V_init)==3.
        @test bellman_backup(3, S, A, T, C, V_init)==2.
        @test bellman_backup(4, S, A, T, C, V_init)==2.
        @test bellman_backup(5, S, A, T, C, V_init)==2.8
        @test bellman_backup(6, S, A, T, C, V_init)==0.
    end

    @test isapprox(value_iteration(S,A,T,C, ϵ=0.001, V_old=[3.,3.,2.,2.,1.,0.]), [5.99921,5.99921,4.99969, 4.99969,3.99969,0.], atol=0.001)
end

@testset "Function inputs" begin
    # Define simple domains (from Massau Kolobov 2016 book, p.41)
    S = 1:6 # States
    A = 1:2
    T = Dict(   # Transitions, prob = T[(state, action, next_state)]
        (1,1,3)=>1.,
        (1,2,2)=>1.,
        (2,1,3)=>1.,
        (3,1,5)=>1.,
        (3,2,2)=>1.,
        (4,1,5)=>1.,
        (5,1,6)=>1.,
        (5,2,6)=>.6,
        (5,2,4)=>.4,
        (6,1,6)=>1.,    # Manually added goal state transition
    )
    C = Dict( # Costs, cost = C[(state, action, next_state)]
        (1,1,3)=>1,
        (1,2,2)=>1,
        (2,1,3)=>1,
        (3,1,5)=>1,
        (3,2,2)=>1,
        (4,1,5)=>1,
        (5,1,6)=>5,
        (5,2,6)=>2,
        (5,2,4)=>2,
        (6,1,6)=>0,   # Manually added goal state transition cost
    )

    # V_init = [3.,3.,2.,2.,1.,0.]
    # @testset "Bellman backup" begin
    #     @test bellman_backup(1, S, A, T, C, V_init)==3.
    #     @test bellman_backup(2, S, A, T, C, V_init)==3.
    #     @test bellman_backup(3, S, A, T, C, V_init)==2.
    #     @test bellman_backup(4, S, A, T, C, V_init)==2.
    #     @test bellman_backup(5, S, A, T, C, V_init)==2.8
    #     @test bellman_backup(6, S, A, T, C, V_init)==0.
    # end
    #
    # @test isapprox(value_iteration(S,A,T,C, ϵ=0.001, V_old=[3.,3.,2.,2.,1.,0.]), [5.99921,5.99921,4.99969, 4.99969,3.99969,0.], atol=0.001)
end


"""
  Rudimentary adaptive function sampling
  Adapted from https://stackoverflow.com/a/14084715
"""
function sample_function(func, nodes, values=func.(nodes), mask=Colon();
    depth=0, min_nodes=16, max_level=16, tol=0.01)

  if depth > max_level
    return nodes, values
  end

  # Calculate the function in between the nodes
  x_c = 0.5 * (nodes[1:end-1][mask] + nodes[2:end][mask])
  y_c = func.(x_c)

  x₂ = vcat(nodes, x_c)
  y₂ = vcat(values, y_c)
  is = sortperm(x₂)

  x₂ = x₂[is]
  y₂ = y₂[is]

  # Determine where to refine
  if length(x₂) < min_nodes
    mask = trues(length(x₂) - 1)
  else
    # Rescale and compute the length of each line segment in the path
    dx = diff(x₂) / ptp(x₂)
    dy = diff(y₂) / abs(ptp(y₂))
    ds = sqrt.(abs2.(dx) + abs2.(dy))

    # Compute the angle between consecutive line segments θ = acos(a⋅b / (|a||b|))
    dx ./= ds
    dy ./= ds

    if eltype(dy) <: Complex
      @views dy_re = reinterpret(Float64, dy)[1:2:end]
      @views dy_im = reinterpret(Float64, dy)[2:2:end]
      dθ = acos.(clamp.(dx[2:end] .* dx[1:end-1] + dy_re[2:end] .* dy_re[1:end-1] + dy_im[2:end] .* dy_im[1:end-1], -1, 1))
    else
      dθ = acos.(clamp.(dx[2:end] .* dx[1:end-1] + dy[2:end] .* dy[1:end-1], -1, 1))
    end

    # Determine where to subdivide (total length of the path is computed ≈ to accuracy `tol`)
    mask = 0.5 * dθ .* (ds[2:end] .+ ds[1:end-1]) .> tol * sum(ds)
    mask = [mask; false]
    mask[2:end] .|= mask[1:end-1]
  end

  # Refine!
  if any(mask)
    return sample_function(func, x₂, y₂, mask; tol=tol, depth=depth+1, min_nodes=min_nodes, max_level=max_level)
  else
    return x₂, y₂
  end
end


using myJuliaUtils
using Test
using LinearAlgebra

N = rand(1:100)
A = genPDmat(N)
cholA = cholPSD(A)
@test cholA == cholesky(A).U
@test cholA'cholA ≈ A

# test1=0
# test2=0
# NoTimes=20
# for i=1:NoTimes
# N = round(Int, 1 + rand() * 100)
# A = randn(N, N)
# A = A'A
# cholA = cholPSD(A)
# test1+=cholA == cholesky(A).U
# test2+=cholA'cholA ≈ A
# end
# @test test1==NoTimes
# @test test2==NoTimes

N = rand(2:100)
K = rand(1:N - 1)
A = genPDmat(N)
S, Q = eigen(A)
# S[1:K] = ones(K) * eps(0.0)
S[1:K] = zeros(K)
# A = Q * Diagonal(S) * inv(Q)
A = Q * Diagonal(S)
A[:,1:K] = zeros(N, K)
A .= A / Q
cholA = cholPSD(A)
if K > 1
    @test size(cholA) == (N - K, N)
else
    @test all(size(cholA) .>= (N - K, N)) # this is done to avoid approximation error
end

@test cholA'cholA ≈ A

# Nseed=round(Int, 1 + rand() * 5000)
# srand(Nseed)
# a=randn(1, N) * cholx(A)
# srand(Nseed)
# a=randn(1, N) * cholx(A)