export graph_denman_beavers

"""
        (graph,cref)=graph_denman_beavers(k;T=ComplexF64,cref_mode=0)

Creates the graph corresponding to the Denman–Beavers iteration for the matrix
square root, using k iterations. The kwarg `cref_mode` specifies if references
to all X and Y (`cref_mode=0`) should be stored or just Y (`cref_mode=-1`) or
just X (`cref_mode=1`).

**Reference**

1. E. Denman and A. Beavers. "The matrix sign function and computations in
   systems". Applied Mathematics and Computation, 2(1), 1976. DOI:
   [10.1016/0096-3003(76)90020-5](https://doi.org/10.1016/0096-3003(76)90020-5)
"""
function graph_denman_beavers(k; T = Float64, cref_mode = 0)
    graph = Compgraph(T)

    cref = []
    for j = 0:k
        Xj = Symbol("X$j")
        Yj = Symbol("Y$j")

        if (j == 0)
            # X0=A
            # Y0=I
            Xj = :A
            Yj = :I
        end

        # Yjinv=Yj^{-1}
        Yinvj = Symbol("Yinv$j")
        if (j > 0)
            add_ldiv!(graph, Yinvj, Yj, :I)
        else
            Yinvj = :I # First step Yinv0=inv(I)=I;
        end

        # Xjinv=Xj^{-1}
        add_ldiv!(graph, Symbol("Xinv$j"), Xj, :I)

        # X_{j+1}=0.5*Xj+0.5*inv(Yj)
        add_lincomb!(graph, Symbol("X" * string(j + 1)), 0.5, Xj, 0.5, Yinvj)
        if (cref_mode >= 0)
            push!(cref, (Symbol("X" * string(j + 1)), 1))
            push!(cref, (Symbol("X" * string(j + 1)), 2))
        end

        # Y_{j+1}=0.5*Yj+0.5*inv(Xj)
        add_lincomb!(
            graph,
            Symbol("Y" * string(j + 1)),
            0.5,
            Yj,
            0.5,
            Symbol("Xinv$j"),
        )
        if (cref_mode <= 0)
            push!(cref, (Symbol("Y" * string(j + 1)), 1))
            push!(cref, (Symbol("Y" * string(j + 1)), 2))
        end

        add_output!(graph, Symbol("X" * string(j + 1)))
    end
    # Not needed to get the output
    del_node!(graph, Symbol("Y" * string(k + 1)))
    setdiff!(
        cref,
        [(Symbol("Y" * string(k + 1)), 1), (Symbol("Y" * string(k + 1)), 2)],
    )

    del_node!(graph, Symbol("Xinv" * string(k)))
    #
    return (graph, cref)
end


"""
    x = xvec(λ,N)

Create `x` values with periodicity for box specified by length `λ`, using `N` points.
"""
xvec(L,N) = LinRange(-L/2,L/2,N+1)[2:end] |> collect

"""
    k = kvec(L,N)

Create `k` values with correct periodicity for box specified by length `λ` for number of points `N`.
"""
kvec(L,N) = fftfreq(N)*N*2*π/L

"""
    X = xvecs(L,N)

Create a tuple containing the spatial coordinate array for each spatial dimension.
"""
function xvecs(L,N)
    X = []
    for (λ,ν) in zip(L,N)
        x = xvec(λ,ν)
        push!(X,x)
    end
    return X |> Tuple
end

"""
    K = kvecs(L,N)

Create a tuple containing the spatial coordinate array for each spatial dimension.
"""
function kvecs(L,N)
    K = []
    for (λ,ν) in zip(L,N)
        k = kvec(λ,ν)
        push!(K,k)
    end
    return K |> Tuple
end

"""
    k² = k2(K)

Create the kinetic energy array `k²` on the `k`-grid defined by the tuple `K`.
"""
function k2(K)
    kind = Iterators.product(K...)
    return map(k-> sum(abs2.(k)),kind)
end

"""
    X,K,dX,dK = xk_arrays(L,N)

Create all `x` and `k` arrays for box specified by tuples `L=(Lx,...)` and `N=(Nx,...)`.
For convenience, differentials `dX`, `dK` are also reaturned. `L` and `N` must be tuples of equal length.
"""
function xk_arrays(L,N)
    @assert length(L) == length(N)
    X = xvecs(L,N)
    K = kvecs(L,N)
    dX = Float64[]; dK = Float64[]
    for j ∈ eachindex(X)
        x = X[j]; k = K[j]
        dx = x[2]-x[1]; dk = k[2]-k[1]
        push!(dX,dx)
        push!(dK,dk)
    end
    dX = dX |> Tuple
    dK = dK |> Tuple
    return X,K,dX,dK
end

"""
    Dx,Dk = dfft(x,k)

Measures that make `fft`, `ifft` 2-norm preserving.
Correct measures for mapping between `x`- and `k`-space.
"""
function dfft(x,k)
    dx = x[2]-x[1]; dk = k[2]-k[1]
    Dx = dx/sqrt(2*pi)
    Dk = length(k)*dk/sqrt(2*pi)
    return Dx, Dk
end

"""
    DX,DK = fft_differentials(X,K)

Evalutes tuple of measures that make `fft`, `ifft` 2-norm preserving for each
`x` or `k` dimension.
"""
function fft_differentials(X,K)
    M = length(X)
    DX = zeros(M); DK = zeros(M)
    for i ∈ eachindex(X)
        DX[i],DK[i] = dfft(X[i],K[i])
    end
    return DX,DK
end


using Combinatorics
using LinearAlgebra
using ArgCheck

function is_a_square_three_tensor(W::Array)

	length(size(W)) == 3 && all(size(W) .== size(W)[1])

end

function naive_tensor(W::Array)

    if is_a_square_three_tensor(W)
        n = size(W)[1]

        sum_tensor_diag(sigma,rho, W) = prod(W[i,sigma[i],rho[i]] for i = 1:n)
        return sum(sum_tensor_diag(sigma,rho, W) for sigma in permutations(collect(1:n)) for rho in permutations(collect(1:n)))
    else
        error("tensor permanent implemented only for square 3-indices tensors")
    end

end

function ryser_tensor(W::Array)

	"""implements the tensor permanent according to Ryser's factorization
	without using Gray factoring

	corresponds to Eq. 21 of https://arxiv.org/abs/1410.7687"""

	@argcheck is_a_square_three_tensor(W) "tensor permanent implemented only for square 3-indices tensors"

	n = size(U)[1]
    nstring = collect(1:n)
    sub_nstring = collect(powerset(nstring))
    sub_nstring = sub_nstring[2:length(sub_nstring)]
    res = 0

    function delta_set(S1, S2)
        if S1 == S2
            return 1
        else
            return 0
        end
    end

    for r = 1:length(sub_nstring)
        for s = r:length(sub_nstring)
            R = sub_nstring[r]
            S = sub_nstring[s]

            t = prod(
                sum(
					W[rr,ss,j] for rr in R for ss in S
                ) for j = 1:n
            )
            res +=
                (2 - delta_set(S, R)) * (-1)^(length(S) + length(R)) * real(t)
        end
    end
    return res

end

function multi_dim_ryser(U, gram_matrix)

    # https://arxiv.org/pdf/1410.7687.pdf

    n = size(U)[1]
    nstring = collect(1:n)
    sub_nstring = collect(powerset(nstring))
    sub_nstring = sub_nstring[2:length(sub_nstring)]
    res = 0

    function delta_set(S1, S2)
        if S1 == S2
            return 1
        else
            return 0
        end
    end

    for r = 1:length(sub_nstring)
        for s = r:length(sub_nstring)
            R = sub_nstring[r]
            S = sub_nstring[s]

            t = prod(
                sum(
                    U[ss, j] * conj(U[rr, j]) * gram_matrix[rr, ss] for rr in R for ss in S
                ) for j = 1:n
            )
            res +=
                (2 - delta_set(S, R)) * (-1)^(length(S) + length(R)) * real(t)
        end
    end
    return res
end


#######################################################################
#=
Section of code dealing with unconstrained minimization problems.
Some functions are reused in the equality constrained section
=#
#######################################################################

#################
#Gradient Descent
#################

#Use Calculus.jl library to compute the gradient numerically
function Compute_Gradient(_f, _x)
    return Calculus.gradient(g -> _f(g),_x)
end

#Use Calculus.jl library to compute the hessian numerically
function Compute_Hessian(_f, _x)
    return Calculus.hessian(h -> _f(h),_x)
end

#Choose t using backgracking line search (p. 464, alg. 9.2)
function Back_Line_Search(_x, _f, _fx, _∇f, _Δx, _α, _β, _n)

    _t = 1

    while _f((_x + _t*_Δx)) > _fx + _α*_t*transpose(_∇f)*_Δx
        _t *= _β
    end

    return _t
end

#Main solution method for Unconstrained Gradient Descent Method (p. 466, alg 9.3)
function Unconstrained_Grad_Descent(_f, _x, _n, _α, _β, _ϵ, maxIt, verbose)
    it = 0;
    t = 1
    ∇f = Compute_Gradient(_f, _x)
    normGrad = norm(∇f)

    val = _f(_x)

    #Stopping conditions: norm(∇f) <= η or iterations exceeds maximum iteration tolerance
    while normGrad > _ϵ && it <= maxIt
        ∇f = Compute_Gradient(_f, _x) #Evaluate gradient at current point
        normGrad = norm(∇f) #Compute euclidean norm using LinearAlgebra library

        #Perform backgracking line search to pick t
        bls = Back_Line_Search(_x, _f, val, ∇f, -∇f, _α, _β, _n)
        t = bls[1]

        #update step
        _x -= t*∇f

        val = _f(_x)

        if(verbose)
            print("\n\nIteration: ", it)
            print("\nGradient Norm = ", normGrad)
        #    print("\nDescent Direction = ", -∇fVal)
            print("\nt = ", t)
        #    print("\nx = ",_x)
            print("\nf(x) = ", val)
        end

        it += 1
    end

    if it > maxIt
        print("\n\nMax Iterations Reached!")
    end

    return _x, val
end

#################
#Newton Method
#################

#Computes Newton step (p. 487, alg. 9.5)
function Newton_Step(_∇f, _∇2f, _x, _n)
    return _∇2f\-_∇f
end

#Computes Newton decrement (p. 487, alg. 9.5)
function Newton_Decrement(_∇2f, _Δxnt)
    return transpose(_Δxnt)*_∇2f*_Δxnt
end

#Main solver method for Unconstrained Newton Method (p. 487, alg. 9.5)
function Unconstrained_Newton(_f, _x, _n, _α, _β, _ϵ, maxIt, verbose)

    λ2 = 1 #prime λ
    it = 0
    val = _f(_x)

    while λ2 > _ϵ && it <= maxIt

        #Objective, Gradient, and Hessian
        val = _f(_x)
        ∇f = Compute_Gradient(_f, _x)
        ∇2f = Compute_Hessian(_f, _x)

        #Newton step
        Δxnt = Newton_Step(∇f, ∇2f, _x, _n)

        #Newton decrement
        λ2 = Newton_Decrement(∇2f, Δxnt)

        #Choose t with backtracking line search
        t = Back_Line_Search(_x, _f, val, ∇f, Δxnt, _α, _β, _n)

        #Print data
        if verbose
            print("\n\nIteration: ", it)
        #    print("\nNewton step = ", Δxnt)
            print("\nλ^2 = ", λ2)
            print("\nt = ", t)
        #    print("\nx = ", _x)
            print("\nf(x) = ", val)
        end

        #Update x
        _x += t*Δxnt

        it+=1
    end

    if it > maxIt
        print("\n\nMax Iterations Reached!")
    end

    return _x, val
end

#=
This function combines the consistent convergent rate of gradient descent
with the fast local convergence of Newton's Method by using gradient
descent at first, and then switching to Newton's method when the solution
gets within η > ϵ. Then, Newton's method gets the solution to within ϵ
=#
function Solve_UMP(_f, _x, _n, _α, _β, _η, _ϵ, maxIt, verbose)
    sol = Unconstrained_Grad_Descent(_f, _x, _n, _α, _β, _η, maxIt, verbose)

    print("\n\nCurrent solution is ", sol[2] ,"\nSwitching to Newton's Method\n")
    #use output point of gradient descent solution as starting point for Newton's Method
    return Unconstrained_Newton(_f, sol[1], _n, _α, _β, _ϵ, maxIt, verbose)
end

#######################################################################
#=
Section of code dealing with equality constrained minimization problems.
=#
#######################################################################

#################
#Feasible Start NM
#################

#Computes Newton Step (p. 526 (10.11))
function Compute_Newton_Step(_f, _∇f, _∇2f, _A, _At, _x, _n, _m)
    return ([_∇2f _At; _A zeros(_m,_m)]  \ [-_∇f; zeros(_m)])[1:_n]
end

#Compute Newton Decrement (p. 527 (10.12))
function Compute_Newton_Decrement(_Δxn, _∇2f)
    return abs(transpose(_Δxn) * _∇2f * _Δxn)
end

#Main solver function for Feasible Start Newton Method (p. 528, alg. 10.1)
function Newton_FS(_f, _x, _A, _n, _m, _α, _β, _ϵ, maxIt, verbose)

    obj = Float64
    λ2 = 10 #prime λ with a big number to ensure first iteration is run
    it = 0
    _At = transpose(_A)

    while λ2 / 2 > _ϵ && it <= maxIt

        obj = _f(_x)
        ∇f = Compute_Gradient(_f, _x)
        ∇2f = Compute_Hessian(_f, _x)
        Δxn = Compute_Newton_Step(_f, ∇f, ∇2f, _A, _At, _x, _n, _m)

        λ2 = Compute_Newton_Decrement(Δxn, ∇2f) #returns λ^2

        t = Back_Line_Search(_x, _f, obj, ∇f, Δxn, _α, _β, _n)

        #Update x
        _x += t*Δxn

        if verbose
            print("\n\nIteration ", it)
            print("\nλ^2 =         ",λ2)
            print("\nObjective =   ", obj)
        end

        it+=1
    end

    return _x, obj
end

#################
#Infeasible Start NM
#################

# This function computes the dual. (p. 531)
function Compute_Dual(_∇f, _At)
    return _At \ -_∇f
end

#This function computes the residual see page 533 (10.21)
function Compute_Residual(_∇f, _At, _v, _vr)
    return [_∇f + _At*_v; _vr]
end

#This function computes the residual just like Compute_Residual but in a way that we can use it for backgracking line search.
function R(_x, _v, _At, _A, _b, _f)
    return [Compute_Gradient(_f, _x) + _At*_v; _A*_x - _b]
end

#This function computes the Newton and Dual Step. See page 533 (10.22)
function Compute_Steps(_∇f, _∇2f, _x, _A, _At, _vr, _p, _v, _res)
    return [_∇2f _At; _A zeros(_p,_p)] \ -_res
end

#Backgracking line search for infeasible start method. (p. 534, Alg. 10.2)
function BLS_IF(_x, _v, _A, _At, _b, _vr, _f, _∇f, _Δxnt, _Δvnt, _α, _β, _nrm)
    _t = 1

    while norm(R(_x + _t*_Δxnt, _v + _t*_Δvnt, _At, _A, _b, _f)) > (1 - _α*_t)* _nrm
        _t *= _β
    end

    return _t
end

#Main solver function
function Newton_IFS(_f, _x, _A, _b, _n, _m, _α, _β, _ϵ, maxIt, verbose)

    #Initialize necessary values for first iteration
    At = transpose(_A)
    vr = _A*_x - _b
    ∇f = Compute_Gradient(_f, _x)
    _v = Compute_Dual(∇f, At)
    it = 0
    nrm = 100

    while  (norm(vr) > _ϵ || nrm > _ϵ) && it < maxIt

        it += 1

        ∇f = Compute_Gradient(_f, _x)
        ∇2f = Compute_Hessian(_f, _x)
        _v = Compute_Dual(∇f, At)

        vr = _A*_x - _b
        res = Compute_Residual(∇f, At, _v, vr)

        Δ = Compute_Steps(∇f, ∇2f, _x, _A, At, vr, _m, _v, res)
        Δxnt = Δ[1:_n]
        Δvnt = Δ[_n+1:_n+_m]

        nrm = norm(res)
        t = BLS_IF(_x, _v, _A, At, _b, vr, _f, ∇f, Δxnt, Δvnt, _α, _β, nrm)

        _x += t*Δxnt
        _v += t*Δvnt

        if verbose
            print("\n\nIteration: ", it)
            print("\nt = ", t)
            print("\nnorm(r) = ", nrm)
    #        print("\nx = ", x)
            print("\nObjective = ", _f(_x))
    #        print("\nΔx = ", Δxnt)
        end

    end

    return _x, _f(_x)
end


## Coordinate relaxation  in general

""" Set x[j] s.t. we solve for row i in the
linear system Ax = b. This updates x
in place"""
function rowsolve!(A,b,x,i,j)
  rhs = b[i]
  for k=1:size(A,2)
    if k!=j
      rhs -= A[i,k]*x[k]
    end
  end
  x[j] = rhs/A[i,j]
end
""" Solve Ax=b with coordinate relaxation sweeps """
function solve_linsys_coord_relax(A,b,nsweeps)
  x = copy(b)
  for iter=1:nsweeps
    for j=1:size(A,1)
      rowsolve!(A,b,x,j,j)
    end
  end
  return x
end
A = [2 1; 1 2]
b = [1,1.5]
x = solve_linsys_coord_relax(A,b,5)
##
xtrue = A\b
norm(x-xtrue)
##
for nsteps = 1:10
  x = solve_linsys_coord_relax(A,b,nsteps)
  @show x
end


#
# Quadratic fit
#

struct Quadratic
  a :: Float64
  b :: Float64
  c :: Float64
  R :: Float64
  x :: Vector{Float64}
  y :: Vector{Float64}
  ypred :: Vector{Float64}
  residues :: Vector{Float64}
end

function fitquadratic(X :: AbstractVector, Y :: AbstractVector, options :: Options)
  X, Y = checkdata(X,Y)
  @. model(x,p) = p[1]*x^2 + p[2]*x + p[3]
  fit = find_best_fit(model, X, Y, 3, options)
  R = pearson(X,Y,model,fit)
  x, y, ypred = finexy(X,options.fine,model,fit)
  return Quadratic(fit.param...,R,x,y,ypred,fit.resid)
end
fitquadratic(X :: AbstractVector, Y :: AbstractVector) = fitquadratic(X,Y,Options())
fitquad(X :: AbstractVector, Y :: AbstractVector) = fitquadratic(X,Y,Options())
fitquad(X :: AbstractVector, Y :: AbstractVector, options :: Options) = fitquadratic(X,Y,options)

function Base.show( io :: IO, fit :: Quadratic )
  println("")
  println(" ------------------- Quadratic Fit ------------- ")
  println("")
  println(" Equation: y = ax^2 + bx + c ")
  println("")
  println(" With: a = ", fit.a)
  println("       b = ", fit.b)
  println("       c = ", fit.c)
  println("")
  println(" Pearson correlation coefficient, R = ", fit.R)
  println("")
  println(" Predicted Y: ypred = [",fit.ypred[1],", ",fit.ypred[2],"...")
  println(" residues = [", fit.residues[1],", ",fit.residues[2],"...")
  println("")
  println(" ----------------------------------------------- ")
end

export fitquad, fitquadratic


#
# Generate initial-velocities
#

import Random

"""

```
initial_velocities(N;T=298.15,R=8.3145,m=1.0,seed=123)
```

Generates a vector of initial velocities for `N` particles of
mass `m`, in two dimensions according to the Maxwell-Boltzmann distribution. By
default the result will be output in ms⁻¹, according to the dimensions
of the gas constant `R`.

## Example

```
v = initial_velocities(10_000)
```

Use a random seed:

```
v = initial_velocities(10_000,100,0.5,seed=-1)
```

"""
function initial_velocities(N;T=298.15,R=8.3145,m=1.0,seed=123)
  # Use a different seed if desired
  seed > 0 ?  Random.seed!(seed) : Random.seed!()
  # Create velocity vector
  v = [ initial_velocity(T,R,m) for i in 1:N ]
  return v
end

# initial_velocity for one particle
function initial_velocity(T,R,m)
  # Inverse of cumulative distribution of p(v)dv
  vmod = sqrt( -(2*R*T/m) * log(1-rand()) )
  θ = 2π*rand()
  return [ vmod*cos(θ), vmod*sin(θ) ]
end

# From the data structure
initial_velocities(data::Data;T=298.15,R=8.3145,m=1.0,seed=123) =
   initial_velocities(data.N,T=T,R=R,m=m,seed=seed)


# 1D Linear interpolation operator

# helper module
module joLinInterp1D_etc
    #1D nearest neighbourhood operator, returns a vector I of size length(xout)
    #such that for every i in 1:length(xout), I[i] is the index such that
    # |xout[i] - xin[I[i]]| <= |xout[i]-xin[j]| for all indices j
    #
    #Parameters:
    #  xin - input grid
    #  xout - output grid
    #
    #Note: the interval [minimum(xout),maximum(xout)] must be contained
    #in the interval [minimum(xin),maximum(xin)]
    function nearestnbr1d(xin,xout)
        ((minimum(xin) <= minimum(xout)) && (maximum(xin) >= maximum(xout))) || throw(ArgumentError("interval defined by xout must be contained within xin"))
        Jin = sortperm(xin)
        Jout = sortperm(xout)
        xin = xin[Jin]
        xout = xout[Jout]
        xend = xin[2]
        istart = 1
        I = zeros(Int64,length(xout))
        for i=1:length(xout)
            while xout[i] > xend
                istart+=1
                if istart==length(xin)
                    xend = inf
                else
                    xend = xin[istart+1]
                end
            end
            if abs(xout[i]-xin[istart]) <= abs(xout[i]-xend)
                I[i] = istart
            else
                I[i] = istart+1
            end
        end
        Jin_inv = Jin
        Jout_inv = Jout
        I = Jin_inv[I]
        I[Jout_inv] = I
        return I
    end
end
using .joLinInterp1D_etc

export joLinInterp1D
"""
    julia> joLinInterp1D(xin,xout;[DDT=...,][RDT=...,][name=...])

1D Linear interpolation operator

# Signature

    function joLinInterp1D(xin::AbstractArray{T,1},xout::AbstractArray{T,1};
        DDT::DataType=T,RDT::DataType=DDT,name::String="joLinInterp1D") where {T<:AbstractFloat}

# Arguments

- `xin`  - input grid
- `xout` - output grid
- keywords
    - `DDT`: domain data type
    - `RDT`: range data type
    - `name`: custom name

# Notes

- The interval [minimum(xout),maximum(xout)] must be contained in the interval [minimum(xin),maximum(xin)]

# Examples

% description

    joLinInterp1D(xin,xout)

examples with DDT/RDT

    joLinInterp1D(xin,xout; DDT=Float32)
    joLinInterp1D(xin,xout; DDT=Float32,RDT=Float64)

"""
function joLinInterp1D(xin::AbstractArray{T,1},xout::AbstractArray{T,1};
    DDT::DataType=T,RDT::DataType=DDT,name::String="joLinInterp1D") where {T<:AbstractFloat}

    xin = sort(xin)
    xout = sort(xout)
    nin = length(xin)
    nout = length(xout)
    I = joLinInterp1D_etc.nearestnbr1d(xin,xout)
    I[xin[I].>xout] = I[xin[I].>xout].-1
    xin = [xin;1e12]
    b = (xout-xin[I])./(xin[I.+1]-xin[I])
    a = (xin[I.+1]-xout)./(xin[I.+1]-xin[I])
    A = sparse(1:nout,I,a,nout,nin) + sparse(1:nout,min.(I.+1,nin),b,nout,nin)
    return joMatrix(A,DDT=DDT,RDT=RDT,name=name)
end


# # Robust approximate fitting
# Section 6.4.2
# Boyd & Vandenberghe "Convex Optimization"
# Original by Lieven Vandenberghe
# Adapted for Convex by Joelle Skaf - 10/03/05
#
# Adapted for Convex.jl by Karanveer Mohan and David Zeng - 26/05/14
# Original cvx code and plots here:
# <http://web.cvxr.com/cvx/examples/cvxbook/Ch06_approx_fitting/html/fig6_15.html>
#
# Consider the least-squares problem:
#       minimize $\|(A + tB)x - b\|_2$
# where $t$ is an uncertain parameter in [-1,1]
# Three approximate solutions are found:
#
#   1. nominal optimal (i.e. letting t=0)
#   2. stochastic robust approximation:
#           minimize $\mathbb{E}\|(A+tB)x - b\|_2$
#      assuming $u$ is uniformly distributed on [-1,1].
#      (reduces to minimizing $\mathbb{E} \|(A+tB)x-b\|^2 = \|A*x-b\|^2  + x^TPx$
#        where $P = \mathbb{E}(t^2) B^TB = (1/3) B^TB$ )
#   3. worst-case robust approximation:
#           minimize $\mathrm{sup}_{-1\leq u\leq 1} \|(A+tB)x - b\|_2$
#      (reduces to minimizing $\max\{\|(A-B)x - b\|_2, \|(A+B)x - b\|_2\}$ ).
#
using Convex, LinearAlgebra, SCS
if VERSION < v"1.2.0-DEV.0"
     LinearAlgebra.diagm(v::AbstractVector) = diagm(0 => v)
end
# Input Data
m = 20;
n = 10;
A = randn(m,n);
(U,S,V) = svd(A);
S = diagm(exp10.(range(-1, stop=1, length=n)));
A = U[:, 1:n] * S * V';

B = randn(m, n);
B = B / norm(B);

b = randn(m, 1);
x = Variable(n)

# Case 1: Nominal optimal solution
p = minimize(norm(A * x - b, 2))
solve!(p, () -> SCS.Optimizer(verbose=0))
x_nom = evaluate(x)

# Case 2: Stochastic robust approximation
P = 1 / 3 * B' * B;
p = minimize(square(pos(norm(A * x - b))) + quadform(x, Symmetric(P)))
solve!(p, () -> SCS.Optimizer(verbose=0))
x_stoch = evaluate(x)

# Case 3: Worst-case robust approximation
p = minimize(max(norm((A - B) * x - b), norm((A + B) * x - b)))
solve!(p, () -> SCS.Optimizer(verbose=0))
x_wc = evaluate(x)

# Plot residuals:
parvals = range(-2, stop=2, length=100);

errvals(x) = [ norm((A + parvals[k] * B) * x - b) for k = eachindex(parvals)]
errvals_ls = errvals(x_nom)
errvals_stoch = errvals(x_stoch)
errvals_wc = errvals(x_wc)

using Plots
plot(parvals, errvals_ls, label="Nominal problem")
plot!(parvals, errvals_stoch, label="Stochastic Robust Approximation")
plot!(parvals, errvals_wc, label="Worst-Case Robust Approximation")
plot!(title="Residual r(u) vs a parameter u for three approximate solutions", xlabel="u", ylabel="r(u) = ||A(u)x-b||_2")


# This implements the L0 EM algorithm for regularized regression as described in:
#
#  Zhenqiu Liu and Gang Li. "Efficient Regularized Regression for Variable
#  Selection with L0 Penalty" Submitted to arXiv on July 29th 2014.
#
# theta = l0EM(X, y, lambda; kws...)
#
# Find coefficients 'theta' for regularized regression with L0 penalty, i.e.
# theta minimizes 0.5*norm(y - X*theta) + lambda/2*sum(theta .> 0.0).
#
# Outputs:
#  'theta' is an M*1 vector with the (sparse) coefficients.
#
# Inputs:
#  'X' is an N by M matrix.
#
#  'y' is an N by 1 array.
#
#  'lambda' is the regularization parameter. By default it is set via AIC, i.e. to 2.
#    Other, typically valid values are log(N) (BIC) and 2*log(M) (RIC).
#
# Keyword args:
#  'epsilon' is the minimum coefficient value that is allowed, below which the
#    coefficient is set to zero. Defaults to 1e-4.
#
#  'delta_treshold' is the treshold value below which the iteration is deemed to have
#    converged. Defaults to 1e-5.
#
#  'nonnegative' is true iff only non-negative coefficients are allowed. Defaults to false.
#
#  'maxIterations' is max number of iterations (in case a lambda given so it doesn't converge).
#
function l0EM(X, y, lambda = 2;
  epsilon = 1e-4, deltaTreshold = 1e-5,
  nonnegative = false, maxIterations = 10000)

  N, M = size(X)
  #lambda_eye = UniformScaling(lambda) # Same as but often faster than: lambda * eye(N)
  lambda_eye = lambda * eye(N)
  xt = X'
  theta = xt * ((X * xt .+ lambda_eye) \ y) # Same as but often faster / more stable than: theta = xt * inv(X * xt .+ lambda_eye) * y

  if nonnegative
    set_negative_to_zero!(theta)
  end

  iterations = 0

  while iterations < maxIterations
    iterations += 1

    # E-step:
    eta = theta

    # M-step:
    eta_squared = eta .^ 2
    xt_eta = broadcast(*, eta_squared, xt)
    theta = xt_eta * ((X * xt_eta .+ lambda_eye) \ y) # Same as but often faster / more stable than: theta = xt_eta * inv(X * xt_eta .+ lambda_eye) * y

    # Ensure non-negativity if required
    if nonnegative
      set_negative_to_zero!(theta)
    end

    # We have converged if change is too small => break.
    if norm(theta - eta) < deltaTreshold
      break
    end
  end

  theta[abs(theta) .< epsilon] = 0.0

  theta
end

function set_negative_to_zero!(v)
  v[v .< 0.0] = 0.0
  v
end

# Given mild constraints on the co-linearity of features there is
# a max lambda value we can calculate.
function find_max_lambda(X, y)
  lambda_from_col_j(j) = (sum(X[:,j] .* y)^2) / (4 * sum(X[:,j].^2))
  N, M = size(X)
  maximum(map(lambda_from_col_j, 1:M))
end

num_selected(v) = sum(v .> 0.0)

# Instead of checking a fixed number of lambda values in a range from 1e-8
# to max_lambda we do a binary search to identify lambda values where there
# is a transition in the number of selected variables.
#
#  cs, ns = adaptive_lambda_regularized_regression(X, y, regressor)
#
# Outputs:
#  'cs' is a dict mapping the tried lambda values to the regressed coefficients for
#    each lambda.
#
#  'ns' is a dict mapping an integer number of selected variables to the lambda value
#    for which that number of selected variables was first encountered.
#
# Inputs:
#  'X' is an N by M matrix corresponding to N cases each with M measured features.
#
#  'y' is an N by 1 array corresponding to the dependent variable for the N cases.
#
#  'regressor' is a Julia function to call to perform regularized regression.
#    It needs to have an interface like regressor(X, y, lambda; kws...).
#    Defaults to using the L0 EM algorithm for regularized regression.
#
# Optional keyword arguments:
#  'minLambda' is the minimum lambda value to use. Defaults to 1e-7 and might not
#    converge at all if a smaller value is selected.
#
#  'numLambdas' is an indication (not exact) of the number of calls to the regressor
#    that we are allowed to do. The actual number of calls can deviate somewhat but not a lot.
#
#  'stepDivisor' determines how small lambda increments are considered when searching in
#    between the standard log-spaced values of LASSO and similar algorithms. Defaults to
#    2^12 which tends to give a sensible number of alternative models to choose from.
#    Increase it to find a large number of selected number of variables (returned in 'ns').
#
#  'logSpace' is true iff the search should be in log space i.e. if taking the logarithm
#    into account when selecting mid points in the binary search.
#
#  'includeAICetal' is true iff we ensure AIC, BIC and RIC are included among the lambdas
#    to be tested. The stepDivisor should typically be set lower (2^6) is this is set to true.
#    Defaults to false.
#
#  'kws' are any other keyword arguments that will be passed on to the regressor.
function adaptive_lambda_regularized_regression(X, y, regressor = l0EM;
  minLambda = 1e-7, numLambdas = 100,
  stepDivisor = 2^12,
  logSpace = true,
  includeAICetal = false,
  kws...)

  max_lambda = find_max_lambda(X, y)
  log_increment = (log(max_lambda + 1) - log(minLambda + 1)) / (numLambdas-1)

  num_vars_to_lambda = Dict{Int64, Float64}()
  coeffs = Dict{Float64, Matrix{Float64}}()

  num_calls_to_regularized_regressor = 0

  update_if_not_there!(l) = begin
    if haskey(coeffs, l)
      cs = coeffs[l]
    else
      coeffs[l] = cs = regressor(X, y, l; kws...)
      num_calls_to_regularized_regressor += 1
    end
    nvars = num_selected(cs)
    if !haskey(num_vars_to_lambda, nvars)
      num_vars_to_lambda[nvars] = l
    end
    return (cs, nvars)
  end

  midpoint(minL, maxL) = begin
    if logSpace
      log_minL = log(minL)
      exp( log_minL + (log(maxL) - log_minL) / 2 )
    else
      minL + (maxL - minL) / 2
    end
  end

  # Search the lambda values from minLambda to maxLambda
  binary_index_search(minLambda, maxLambda, minDelta) = begin
    #println("bis($minLambda, $maxLambda, $minDelta)")
    if num_calls_to_regularized_regressor <= numLambdas && (maxLambda - minLambda > minDelta)
      csmin, nvarmin = update_if_not_there!(minLambda)
      #println("  nvarmin = $nvarmin ($minLambda)")
      csmax, nvarmax = update_if_not_there!(maxLambda)
      #println("  nvarmax = $nvarmax ($maxLambda)")
      if (nvarmin - nvarmax > 1)
        mid = midpoint(minLambda, maxLambda)
        csmid, nvarmid = update_if_not_there!(mid)
        #println("  nvarmid = $nvarmid ($mid)")
        if nvarmin - nvarmid > 1
          binary_index_search(minLambda, mid, minDelta)
        end
        if nvarmid - nvarmax > 1
          binary_index_search(mid, maxLambda, minDelta)
        end
      end
    end
  end

  if includeAICetal

    # Always include the AIC, BIC and RIC endpoints and search them first to ensure
    # they are included.
    a, b = 2.0 < log(N) ? (2.0, log(N)) : (log(N), 2.0)
    binary_index_search(a, b, log_increment/stepDivisor)
    b, c = b < (2*log(M)) ? (b, 2*log(M)) : (2*log(M), b)
    binary_index_search(b, c, log_increment/stepDivisor)

    # And then from minLambda up to AIC and from RIC up to max lambda.
    binary_index_search(minLambda, a, log_increment/stepDivisor)
    binary_index_search(c, max_lambda, log_increment/stepDivisor)

  else
    binary_index_search(minLambda, max_lambda, log_increment/stepDivisor)
  end

  return (coeffs, num_vars_to_lambda)
end

mse(theta, X, y) = mean( (X * theta - y).^2 )

# Basic test:
function basic_l0EM_test(N, M, NumReps = 30)
  for rep in 1:NumReps
    k = rand(1:iceil(2*log(M)))
    X = randn(N, M)
    actual_theta = [linspace(1,k,k), zeros(M-k)]
    y = X * actual_theta + 0.10 * randn(N, 1)

    cs, ns = adaptive_lambda_regularized_regression(X, y, l0EM;
      stepDivisor = 2^4, includeAICetal = true);

    # Find lambda with lowest MSE on test set
    Xtest = randn(N, M)
    ytest = Xtest * actual_theta + 0.10 * randn(N, 1)
    best_mse = Inf
    local best_lambda, best_nvars
    for lambda in values(ns)
      newmse = mse(cs[lambda], Xtest, ytest)
      if newmse < best_mse
        best_nvars = num_selected(cs[lambda])
        best_lambda = lambda
        best_mse = newmse
      end
    end
    println("k = $k, Lambda = $best_lambda, num vars = $best_nvars, MSE = $best_mse")
  end
end

basic_l0EM_test(50, 100, 10)

# Using log space binary search seems to speed the search up about 10-30% depending on problem size.
# Potentially more for big data sets and with cross validation.
#@time cs2, ns2 = adaptive_lambda_regularized_regression(X, y, l0EM; stepDivisor = 2^2, includeAICetal = true);
#@time cs4, ns4 = adaptive_lambda_regularized_regression(X, y, l0EM; stepDivisor = 2^4, includeAICetal = true);
#@time cs6, ns6 = adaptive_lambda_regularized_regression(X, y, l0EM; stepDivisor = 2^6, includeAICetal = true);
#@time cs4b, ns4b = adaptive_lambda_regularized_regression(X, y, l0EM; stepDivisor = 2^4, includeAICetal = false);
#@time csf, nsf = adaptive_lambda_regularized_regression(X, y, l0EM; logSpace = false);
#sort(collect(keys(ns)))
#length(cs)

# If not using the adaptive (binary) search the standard method is to do 100 lambda values
# in a log scale:
#function log_split_lambdas(X, y, numLambdas = 10)
#  exp(linspace(0.0, log(find_max_lambda(X, y) + 1), numLambdas)) - 1.0 + 1e-6
#end
#ts100 = map((l) -> l0EM(X, y, l), log_split_lambdas(X, y, 100))
#num_vars = map((t) -> sum(t .> 0.0), ts100)


## Spectral clustering
using DelimitedFiles
using SparseArrays
using LinearAlgebra
using Random
ei,ej = copy.(eachcol(Int.(readdlm("undir-graph-14.edges"))))
xy = readdlm("undir-graph-14.xy")'
n = max(maximum(ei),maximum(ej))
A = sparse(ei,ej,1,n,n)
A = A+A'
m = length(ei) # number of edges
B = sparse(1:m,ei,1,m,n) - sparse(1:m,ej,1,m,n)
L = B'*B
d = vec(sum(A,dims=2))
D = Diagonal(d)
Di = inv(D) # Julia does this efficiently...
Pt = Di*A
P = A*Di # A is symmetric ...
## Make a vector orthogonal to degrees
Random.seed!(1)
x = randn(n)
#x .-= ((d'*x)/(d'*d)).*d
x .-= (d'*x)/(sum(d))
x ./= x'*D*x

##

##
Random.seed!(1)
function simple_spectral_eigenvector(A,nsteps=500)
  @assert issymmetric(A) # expensive, but useful...
  n = size(A,1)
  d = vec(sum(A,dims=1))
  nd2 = d'*d
  x = randn(n)
  # project x orthogonal to d
  x .-= ((d'*x)/nd2).*d
  x ./= x'*(d.*x) # normalize
  for i=1:nsteps
    x = (A*x)./d + x       # matrix-vector
    x .-= ((d'*x)/nd2).*d
    x ./= sqrt(x'*(d.*x)) # normalize
    # make sure the first component is positive
    x .*= sign(x[1])
  end

  return x
end
f = simple_spectral_eigenvector(A)
##
using MatrixNetworks
f2 = fiedler_vector(A)[1]
# make sure the first component is positive
f2 .*= sign(f2[1])
##
norm(f-f2)

##
using Plots
using GraphRecipes
pyplot(size=(300,300))
p = graphplot(ei, ej, x =xy[1,:], y=xy[2,:],
  markercolor=:white, markerstrokecolor=:black,
  markersize=13, linecolor=3, linealpha=1.0, linewidth=1.0,
  marker_z = f,
  axis_buffer=0.5, background=:transparent,dpi=300)
xlims!(0,1)
ylims!(-0.05, 1.05)
for i=1:size(A,1)
  if f[i] >= 0
    annotate!(xy[1,i],xy[2,i],text("$i",9,:black))
  else
    annotate!(xy[1,i],xy[2,i],text("$i",9,:white))
  end
end
plot!()
savefig("spectral-clustering-undir.pdf")


"""
    dot(x1,y1,x2,y2) -> Real

Calculate the dot product between two vectors.
"""
dot(x1::Real, y1::Real, x2::Real, y2::Real) = x1*x2 + y1*y2

"""
    dot(z1,z2) -> Real

Calculate the dot product between two complex number (like 2D vectors).
"""
dot(z1::Number, z2::Number) = real(z1)*real(z2) + imag(z1)*imag(z2)

"""
    perp(x,y) -> Tuple{Real,Real}

Return the pependicular vector to \$(x,y)\$.
"""
perp(x::Real, y::Real) = -y, x

"""
    perp(z) -> Complex

Return the pependicular complex number to \$z\$.
"""
perp(z::Number) = -imag(z) + im*real(z)

"""
"""
tripledet(x1::Real, y1::Real, x2::Real, y2::Real, x3::Real, y3::Real) =
  x1*(y2-y3)-y1*(x2-x3)+x2*y3-y2*y3

"""
"""
tripledet(z1::Number, z2::Number, z3::Number) =
  tripledet(real(z1), imag(z1), real(z2), imag(z2), real(z3), imag(z3))

#=
ToDo!
circumcenter
circumradius
=#


"""
The sum of the squares of the first ten natural numbers is,
1^2 +2^2 +... +10^2 = 385
The square of the sum of the first ten natural numbers is,
(1+2+...+10)^2 =55^2 =3025
Hence the difference between the sum of the squares of the first ten natural numbers and the square of the sum is 3025-385 =2640
Find the difference between the sum of the squares of the first one hundred natural numbers and the square of the sum.
"""
function sumsqdiff(n)
sum(1:n)^2 - sum((1:n).^2)
end

sumsqdiff(100)


module gss

function goldss(f,a,b;ϵ=1e-6,τ=0.618)
    fa,fb=f(a),f(b)
    x1,x2=b-(τ)*(b-a),a+(τ)*(b-a)
    err=Inf
    while err>ϵ
        if f(x2)>f(x1)
            b=x2
            x2=x1
            x1=a+(1-τ)*(b-a)
        elseif f(x2)<f(x1)
            a=x1
            x1=x2
            x2=b-(1-τ)*(b-a)
        else
            a=(x1+x2)/2
            b=a
        end
        err=2*abs((b-a)/(b+a))
    end
    xsr=(x1+x2)/2
    wart=f(xsr)
    return(round(xsr,digits=2),round(wart,digits=3))
end

export goldss

end # module


# ---
# title: 365. Water and Jug Problem
# id: problem365
# author: Indigo
# date: 2021-06-14
# difficulty: Medium
# categories: Math
# link: <https://leetcode.com/problems/water-and-jug-problem/description/>
# hidden: true
# ---
#
# You are given two jugs with capacities _x_ and _y_ litres. There is an
# infinite amount of water supply available. You need to determine whether it is
# possible to measure exactly _z_ litres using these two jugs.
#
# If _z_ liters of water is measurable, you must have _z_ liters of water
# contained within **one or both buckets** by the end.
#
# Operations allowed:
#
#   * Fill any of the jugs completely with water.
#   * Empty any of the jugs.
#   * Pour water from one jug into another till the other jug is completely full or the first jug itself is empty.
#
# **Example 1:** (From the famous [_" Die Hard"_
# example](https://www.youtube.com/watch?v=BVtQNK_ZUJg))
#
#
#
#     Input: x = 3, y = 5, z = 4
#     Output: True
#
#
# **Example 2:**
#
#
#
#     Input: x = 2, y = 6, z = 5
#     Output: False
#
#
#
#
# **Constraints:**
#
#   * `0 <= x <= 10^6`
#   * `0 <= y <= 10^6`
#   * `0 <= z <= 10^6`
#
#
## @lc code=start
using LeetCode

function can_measure_water(jug1_cap::Int, jug2_cap::Int, target::Int)
    target > jug1_cap + jug2_cap && return false
    return target % gcd(jug2_cap, jug1_cap) == 0 ? true : false
end
## @lc code=end


# Tests concerning cross-entropy

arr21 = rand(100)
arr22 = rand(100)
arr41 = rand(10000)
arr42 = rand(10000)

# The cross-entropy between two copies of an array
@test cross_entropy(arr41, arr41) ≈ log2(100) atol = 0.01
println("Cross-entropy between array and itself passed.")

# The cross-entropy between an array and itself is equal to the regular entropy of that array
@test cross_entropy(arr21, arr21) == get_entropy(arr21)
println("Equality of cross-entropy between array and itself and entropy passed.")

# The cross-entropy between two small samples from the same distribution will be somewhat larger than the regular entropy
@test cross_entropy(arr21, arr22) ≈ log2(10) atol = 0.4
println("Cross-entropy between two small sample arrays passed.")
@test cross_entropy(arr41, arr42) ≈ log2(100) atol = 0.02
println("Cross-entropy between two large sample arrays passed.")


# Algorithm 5.2
# Fourth order Runge-Kutta

function rk4(f, a, b, N, α)
	result = Tuple{Float64, Float64}[]
	k = Array{Float64}(undef, 4)
	h = (b - a)/N
	t = α
	ω = α
	push!(result, (t, ω))
	for i ∈ 1:N
		k[1] = h*f(t, ω)
		k[2] = h*f(t + h/2, ω + k[1]/2)
		k[3] = h*f(t + h/2, ω + k[2]/2)
		k[4] = h*f(t + h, ω + k[3])
		ω = ω +(k[1] + 2k[2] + 2k[3] + k[4])/6
		t = a + i*h
		push!(result, (t, ω))
	end
	return result
end

f(t, y) = y -t^2 + 1
println(rk4(f, 0, 2, 10, 0.5))


function ThomasLUSolve(A::Any,B::Vector)
      #A = the matrix of coefficients 'A' must be squared
  #B = The Vector 'B' right hand side vector Vertical
    m,n =size(A)
    v = size(B,1)
    if m != n
        error("Matrix must be squared")
    elseif m != v
        error("Vector Length must be equal to Matrix")
    end
    e=zeros(m-1)
    f=zeros(m)
    g=zeros(m-1)
    r=zeros(m)
    x=zeros(m)
    f[1]=A[1,1]
        #Decomposition
    for i=2:m
        e[i-1]=A[i,i-1]/f[i-1]
        g[i-1]=A[i-1,i]
        f[i]=A[i,i]-e[i-1]*g[i-1]

    end

    #Fordward Substitution
    r[1]=B[1]
    for i=2:m
        r[i]=B[i]-e[i-1]*r[i-1]
    end

    #Back Substitution
    x[m]=r[m]/f[m]

    for i=m-1:-1:1
       x[i]=(r[i]-g[i]*x[i+1])/f[i]
    end

    return x
end


"""
	varimax(A; gamma = 1.0, minit = 20, maxit = 1000, reltol = 1e-12)

VARIMAX perform varimax (or quartimax, equamax, parsimax) rotation to the column vectors of the input matrix.

# Input Arguments
- `A::Matrix{Float64}`: input matrix, whose column vectors are to be rotated. d, m = size(A).
- `gamma::Float64`: default is 1. gamma = 0, 1, m/2, and d(m - 1)/(d + m - 2), corresponding to quartimax, varimax, equamax, and parsimax.
- `minit::Int`: default is 20. Minimum number of iterations, in case of the stopping criteria fails initially.
- `maxit::Int`: default is 1000. Maximum number of iterations.
- `reltol::Float64`: default is 1e-12. Relative tolerance for stopping criteria.

# Output Argument
- `B::Matrix{Float64}`: output matrix, whose columns are already been rotated.

Implemented by Haotian Li, Aug. 20, 2019
"""
function varimax(A; gamma = 1.0, minit = 20, maxit = 1000, reltol = 1e-12)
	# Get the sizes of input matrix
	d, m = size(A)

	# If there is only one vector, then do nothing.
	if m == 1
		return A
	end

	if d == m && rank(A) == d
		return Matrix{Float64}(I, d, m)
	end

	# Warm up step: start with a good initial orthogonal matrix T by SVD and QR
	T = Matrix{Float64}(I, m, m)
	B = A * T
	L,_,M = svd(A' * (d*B.^3 - gamma*B * Diagonal(sum(B.^2, dims = 1)[:])))
	T = L * M'
	if norm(T-Matrix{Float64}(I, m, m)) < reltol
		T,_ = qr(randn(m,m)).Q
		B = A * T
	end

	# Iteration step: get better T to maximize the objective (as described in Factor Analysis book)
	D = 0
	for k in 1:maxit
		Dold = D
		L,s,M = svd(A' * (d*B.^3 - gamma*B * Diagonal(sum(B.^2, dims = 1)[:])))
		T = L * M'
		D = sum(s)
		B = A * T
		if (abs(D - Dold)/D < reltol) && k >= minit
			break
		end
	end

	# Adjust the sign of each rotated vector such that the maximum absolute value is positive.
	for i in 1:m
		if abs(maximum(B[:,i])) < abs(minimum(B[:,i]))
			B[:,i] .= - B[:,i]
		end
	end

	return B
end


function test_ldl_decomp()

    A = [
        4    3       1      5;
        1    2       1      0.1;
        1    2       16     3;
        -2   1       4      9
        ]

    n = 20 # 200
    A = rand(n, n)
    A = A + A'

    (L, d) = linalg.ldl_decomp(A)

    A_dec = L * diagm(d) * L'

    println("L:")
    show(stdout, MIME("text/plain"), L)
    println("\n")

    println("d:")
    show(stdout, MIME("text/plain"), d)
    println("\n")

    #=
    println("A:")
    show(stdout, MIME("text/plain"), A)
    println("\n")

    println("A_dec:")
    show(stdout, MIME("text/plain"), A_dec)
    println("\n")
    =#

    println("A_dec - A:")
    show(stdout, MIME("text/plain"), norm( A_dec - A) <= 1e-10 )
    println("\n")

end

function test_linsolve_ldl()
    A = [
        4    3       1      5;
        1    2       1      0.1;
        1    2       16     3;
        -2   1       4      9
        ]

    n = 200 # 200
    A = rand(n, n)
    A = A + A'

    # b = [3.0, 6, 2, 10]
    b = rand(n)

    x_ldl = linalg.linsolve_ldl(A, b)

    x_theo = A \ b

    #=
    println("x_ldl:")
    println(x_ldl)
    println( "residuum norm x_lu:  $(norm( A * x_ldl - b )) " )

    println("x_theo:")
    println(x_theo)
    println( "residuum norm x_theo:  $(norm( A * x_theo - b )) " )

    println("dx:")
    println(x_ldl - x_theo)
    println( "norm dx:  $(norm(x_ldl - x_theo)) " )
    =#

    println("dx:")
    println( "norm dx:  $(norm(x_ldl - x_theo)) " )

end

function benchmark_linsolve_ldl(A, b)
    b_ = @benchmark linalg.linsolve_ldl($A, $b)

    println("=== benchmark_linsolve_ldl ===")
    show(
            stdout,
            MIME("text/plain"),
            b_ )
    println("\n\n")
end


module DSToolsDemo

# # DSTOOLSDEMO    Demonstration of the Descriptor System Tools (DSTOOLS).
# #                A. Varga.

using DescriptorSystems
using LinearAlgebra
using Polynomials
using Test

function pause()
    # Comment out the following two lines for a non-interactive execution
    println("Press any key to continue")
    readline()
end

s = rtf('s'); z = rtf('z'); # define the complex variables s and z
println("s = rtf('s'); z = rtf('z'); # define the complex variables s and z ")
pause()

println("Gc = [s^2 s/(s+1); 0 1/s]   # define the 2-by-2 improper Gc(s)")
Gc = [s^2 s/(s+1); 0 1/s]   # define the 2-by-2 improper Gc(s)
Gd = [z^2 z/(z-2); 0 1/z]   # define the 2-by-2 improper Gd(z)
println(Gc)
pause()

println("Gd = [z^2 z/(z-2); 0 1/z]   # define the 2-by-2 improper Gd(z)")
println(Gd)
pause()

println("Build LTI descriptor realizations of Gc(s) and Gd(z) ")
println("sysc = dss(Gc, minimal = true);  # build continuous-time descriptor system realization")
sysc = dss(Gc, minimal = true);  # build continuous-time descriptor system realization
println(sysc)
pause()

println("dss(Gd, minimal = true)  # build discrete-time descriptor system realization")
sysd = dss(Gd, minimal = true);   # build continuous-time descriptor system realization
println(sysd)
pause()

println("Compute poles and zeros of Gc(s)")
println("gpole(sysc)   #  computes all poles (finite and infinite)")
display(gpole(sysc))
pause()

println("gzero(sysc)   #  computes all zeros (finite and infinite)")
display(gzero(sysc))
pause()

println("Descriptor systems (A-λE,B,C,D) with singular pole pencil A-λE are not supported in DSTOOLS!")
println("Non-regular A-λE can be detected in several ways.\n ")
println("(1) non-regular A-λE can be detected when building a descriptor system ")
println("A = zeros(1,1); E = zeros(1,1); B = ones(1,1); C = ones(1,1); D = zeros(1,1);")
println("dss(A,E,B,C,D,check_reg = true)")
A = zeros(1,1); E = zeros(1,1); B = ones(1,1); C = ones(1,1); D = zeros(1,1);
try
  dss(A,E,B,C,D,check_reg = true)  # check regularity when building the descriptor system
catch err
  println(err)
end
pause()

println("(2) non-regular A-λE can be explicitly detected using isregular")
println("syst = dss(A,E,B,C,D); # no regularity check performed when building a descriptor system")
println("isregular(syst) == false")
syst = dss(A,E,B,C,D);
println(@test isregular(syst) == false)
pause()

println("(3) non-regular A-λE can be detected using isproper")
println("isproper(syst) == false")
syst = dss(A,E,B,C,D);
println(@test isproper(syst) == false)
pause()

println("(4) gpole can also be used to check regularity of the pole pencil")
println("gpole(syst)    # the system has some poles equal to NaN")
display(gpole(syst))
pause()

println("Let's try some decomposition and factorization functions\n")

println("Use gsdec to compute the separation of proper and polynomial parts of Gc(s)")
println("sysp, sysi = gsdec(sysc); # Gc(s) = Gcp(s) + Gci(s)")
sysp, sysi = gsdec(sysc);
println("@test iszero(sysc-sysp-sysi)  # checking the decomposition")
println(@test iszero(sysc-sysp-sysi))
pause()
println("dss2rm(sysp)  # display proper part Gcp(s)")
println(dss2rm(sysp))
pause()
println("dss2rm(sysi)  # display polynomial part Gci(s)")
println(dss2rm(sysi))
pause()

println("Compute the stable and proper right coprime factorization of Gc(s) as")
println("Gc(s) = N(s)*inv(M(s)), with N(s) and M(s) stable and proper,")
println(" having a stability degree -1 and poles assigned to [-2,-3,-4]")
println("sysn, sysm = grcf(sysc,evals = [-2,-3,-4],sdeg = -1,mindeg = true, mininf = true);")
sysn, sysm = grcf(sysc,evals = [-2,-3,-4], sdeg = -1, mindeg = true,  mininf = true);
println("@test iszero(sysc*sysm-sysn)  # checking the factorization")
println(@test iszero(sysc*sysm-sysn))
pause()

println("@test isproper(sysm) && isproper(sysn) # checking properness of factors")
println(@test isproper(sysm) && isproper(sysn) )
pause()

println("@test (maximum(real(gpole(sysm))) < 0 && maximum(real(gpole(sysn))) < 0) # checking stability of poles")
println(@test (maximum(real(gpole(sysm))) < 0 && maximum(real(gpole(sysn))) < 0) )
pause()

println("@test isempty(gzero(gminreal([sysn;sysm]))) # checking coprimeness:  [N(s);M(s)] has no zeros")
println(@test isempty(gzero(gminreal([sysn;sysm]))))
pause()

println("Compute the right coprime factorization with inner denominator of Gd(z) as")
println("Gd(z) = N(z)*inv(M(z)), with N(z) and M(z) stable and proper, and M(z) inner")
println("sysni, sysmi = grcfid(sysd);")
sysni, sysmi = grcfid(sysd);
println("@test iszero(sysc*sysm-sysn)  # checking the factorization")
println(@test iszero(sysc*sysm-sysn))
pause()

println("@test maximum(abs.(gpole(sysmi))) < 1 && maximum(abs.(gpole(sysni))) < 1) # checking stability of poles")
println(@test maximum(abs.(gpole(sysmi))) < 1 && maximum(abs.(gpole(sysni))) < 1  )
pause()

println("@test iszero(sysmi'*sysmi-I) # checking the innerness of M(z): conj(M(z))*M(z)-I = 0")
println(@test iszero(sysmi'*sysmi-I))
pause()

println("@test isempty(gzero(gminreal([sysn;sysm]))) # checking coprimeness:  [N(z);M(z)] has no zeros")
println(@test isempty(gzero(gminreal([sysn;sysm]))))
pause()

println("Compute the inner-outer factorization of stable 3x3 G(s) with rank 2 ")
println("s = Polynomial([0, 1],'s'); ")
s = Polynomial([0, 1],'s');
println("num = [(s-1) s 1; 0 (s-2) (s-2); (s-1) (s^2+2*s-2) (2*s-1)]; # numerators")
num = [(s-1) s 1; 0 (s-2) (s-2); (s-1) (s^2+2*s-2) (2*s-1)]; # numerators
pause()
println("den = [(s+2) (s+2) (s+2); 1 (s+1)^2 (s+1)^2; (s+2) (s+1)*(s+2) (s+1)*(s+2)]; # denominators")
den = [(s+2) (s+2) (s+2); 1 (s+1)^2 (s+1)^2; (s+2) (s+1)*(s+2) (s+1)*(s+2)]; # denominators
pause()
println("sys = dss(num,den,minimal = true, atol = 1.e-7);  # build a minimal descriptor realization")
sys = dss(num,den,minimal = true, atol = 1.e-7);
# println(sys)
pause()
println("Analysis of some properties")
println("gpole(sys);  # the system is stable")
display(gpole(sys))    # the system is stable
pause()
println("gzero(sys,atol1=1.e-7)    # the system has 2 unstable zeros and an infinite zero")
display(gzero(sys,atol1=1.e-7) )   # the system has 2 unstable zeros and an infinite zero
pause()
println("nr = gnrank(sys,atol1=1.e-7)    # the normal rank of G(s) is 2")
nr = gnrank(sys,atol1=1.e-7);  # the normal rank of G(s) is 2
println(nr)
pause()

println("Compute the inner-quasi-outer factorization of G(s) as G(s) = Gi(s)*[Go(s);0],")
println("with Gi(s) inner and square, and Go(s) quasi-outer (i.e., full row rank, ")
println("and without zeros in the open right-half plane), can be computed by using")
println("sysi, syso, info = giofac(sys)")
sysi, syso, info = giofac(sys);
pause()

println("Checking the factorization: Gi(:,1:nr)(s)*Go(s) = G(s)")
println("@test iszero(sysi[:,1:nr]*syso-sys,atol1=1.e-7)")
println(@test iszero(sysi[:,1:nr]*syso-sys,atol1=1.e-7))
pause()

println("Checking the innerness of Gi(s)")
println("@test iszero(sysi'*sysi-I,atol1=1.e-7)")
println(@test iszero(sysi'*sysi-I,atol1=1.e-7))
pause()

println("If G(s) contains a so-called free inner factor, then this factor")
println("is included in Gi(s), but the realization of Go(s) is not minimal")
println("@test order(syso)-order(gir(syso)) > 0  # a free inner factor is present in G(s)")
println(@test order(syso)-order(gir(syso)) > 0 ) # a free inner factor is present in G(s)
pause()

println("Checking that Go(s) has no zeros in the open right-half plane")
println("zer = gzero(syso,atol1 = 1.e-7);  # compute zeros of Go")
zer = gzero(syso,atol1 = 1.e-7);
println("@test all(real.(zer[.!isinf.(zer)]) .< 0 )")
println(@test all(real.(zer[.!isinf.(zer)]) .< 0 ))
pause()

##
println("Solution of linear rational equations G(s)*X(s) = F(s)")
println("Consider the Wang and Davison example (IEEE Trans. Autom. Contr.,1973)")
println("to determine a right inverse of G(s) by solving G(s)*X(s) = I.")
println("s = Polynomial([0, 1],'s'); # define polynomial s")
s = Polynomial([0, 1],'s');
println("num =  [ s+1 s+3 s^2+3*s; s+2 s^2+2*s 0 ]; # numerators")
num =  [ s+1 s+3 s^2+3*s; s+2 s^2+2*s 0 ]; # numerators
# println(num)
pause()
println("den = [s^2+3*s+2 s^2+3*s+2 s^2+3*s+2; s^2+3*s+2 s^2+3*s+2 s^2+3*s+2]; # denominators")
den = [s^2+3*s+2 s^2+3*s+2 s^2+3*s+2; s^2+3*s+2 s^2+3*s+2 s^2+3*s+2]; # denominators
# println(den)
pause()
println("sysg = dss(num,den,minimal = true, atol = 1.e-7);  # build a minimal descriptor realization of G(s)")
sysg = dss(num,den,minimal = true, atol = 1.e-7);
# println(sysg)
pause()
println("sysf = dss([1 0;0 1.]);  # build a minimal descriptor realization of F(s) = I")
sysf = dss([1 0;0 1.]);
# println(sysf)
pause()

println("zer = gzero(sysg)   # G(s) has no zeros, thus stable right inverses exist")
zer = gzero(sysg)         # G(s) has no zeros, thus stable right inverses exist
println("zer = $zer")
pause()

println("Compute a stable right inverse with poles in [-1 -2 -3] using")
println("sysx, info = grsol(sysg,sysf,poles = [-1, -2, -3]);")
sysx, info = grsol(sysg, sysf, poles = [-1, -2, -3]);
# println(sysx)
pause()

println("Checking the solution: G(s)*X(s) - I = 0")
println("@test iszero(sysg*sysx-I,atol=1.e-7)")
println(@test iszero(sysg*sysx-I,atol=1.e-7))   #  G(s)*X(s) - I = 0
pause()

println("Check assigned poles: gpole(sysx) = [-1, -2, -3]")
println(@test sort(gpole(sysx)) ≈ sort([-1, -2, -3]))
pause()

println("We can also compute a right inverse of least order = 2")
println("sysxmin, = grsol(sysg,sysf,mindeg = true); order(sysxmin) == 2")
sysxmin, = grsol(sysg,sysf,mindeg = true);
println(@test iszero(sysg*sysxmin-I,atol=1.e-7) && order(sysxmin) == 2)
pause() # Press any key to continue ...

gpole(sysxmin)
println("The least order inverse is unstable: max(real(gpole(sysxmin))) > 0")
println(@test maximum(real(gpole(sysxmin))) > 0)
pause()

##
println("\nWe can try to explicitly determine a least order right inverse")
println("using stable generators (X0(s),XN(s)), where X0(s) is a particular")
println("stable right inverse satisfying G(s)*X0(s) = I and XN(s) is a stable")
println("basis of the right nullspace of G(s), satisfying G(s)*XN(s) = 0;")
println("all solutions of G(s)*X(s) = I are given by X(s) = X0(s)+XN(s)*Y(s),")
println("where Y(s) is arbitrary\n")

println("Determine the generators X0 and XN with poles assigned to [-1 -2 -3]")
println("_,_,sysgen = grsol(sysg,sysf,poles = [-1, -2, -3],solgen = true);")
_,_,sysgen = grsol(sysg,sysf,poles = [-1, -2, -3],solgen = true);
pause()

println("Compute a least order solution X2(s) = X0(s)+XN(s)*Y2(s) using")
println("order reduction based on minimal dynamic covers of Type 2")

println("sysx2,sysy2, = grmcover2(sysgen,2); order(gminreal(sysx2)) == 2")
sysx2,sysy2, = grmcover2(sysgen,2);
println(@test order(gminreal(sysx2)) == 2)
pause()

println("Checking the solution: G(s)*X2(s) = I")
println("@test iszero(sysg*sysx2-I,atol=1.e-7)" )
println(@test iszero(sysg*sysx2-I,atol=1.e-7) )  #  G(s)*X2(s) - I = 0
pause()

println("The least order inverse is unstable: max(real(gpole(sysx2))) > 0")
println(@test maximum(real(gpole(sysx2))) > 0)
pause()

println("Checking the minimal cover reduction results: X2(s) = X0(s)+XN(s)*Y2(s)")
println("@test iszero(sysx2-sysgen[:,1:2]-sysgen[:,3]*sysy2,atol=1.e-7)")
println(@test iszero(sysx2-sysgen[:,1:2]-sysgen[:,3]*sysy2,atol=1.e-7) )
pause()

println("\nSolution of an H∞ model matching problem  min||X(s)*G(s)-F(s)||")
println("\nSetting up an example taken from Francis' book (1987)")
println("s = Polynomial([0, 1],'s');")
println("W = (s+1)/(10s+1); # weighting function")
println("G = dss([ -(s-1)/(s^2+s+1)*W; (s^2-2*s)/(s^2+s+1)*W], minimal=true);")
println("F = dss([ W; 0 ]);")

s = Polynomial([0, 1],'s');
W = (s+1)/(10s+1); # weighting function
G = dss([ -(s-1)/(s^2+s+1)*W; (s^2-2*s)/(s^2+s+1)*W], minimal=true);
F = dss([ W; 0 ]);
pause()

println("\nWe employ the γ-iteration based solution proposed by Francis (1987)")

println("Xopt, info = grasol(G, F, mindeg = true, atol = 1.e-7);")
println("This will take some time at first execution!!!")
Xopt, info = grasol(G, F, mindeg = true, atol = 1.e-7);
println("Optimum of approximation error: info.mindist = $(info.mindist)")
pause()

println("This fully agrees with info.mindist: ghinfnorm(G*Xopt-F) ≈ info.mindist")
println(@test ghinfnorm(G*Xopt-F)[1] ≈ info.mindist)
pause()

println("\nThe solution in Francis' book corresponds to a suboptimal solution for γ = 0.2729")
println("Xsub, info = grasol(G, F, 0.2729; mindeg = true, atol = 1.e-7);")
Xsub, info = grasol(G, F, 0.2729; mindeg = true, atol = 1.e-7);
println("The suboptimal approximation error: info.mindist = $(info.mindist)")
pause()

println("This fully agrees with info.mindist: ghinfnorm(G*Xsub-F) ≈ info.mindist")
println(@test ghinfnorm(G*Xsub-F)[1] ≈ info.mindist)
pause()

println("\nLet's illustrate the use of the bilinear transformation")

println("s = rtf('s');")
println("G  = [s^2 s/(s+1); 0 1/s]")
println("sys = dss(G, minimal = true);")
s = rtf('s');                 # define the complex variable s
G  = [s^2 s/(s+1); 0 1/s]     # define the 2-by-2 improper G(s)
sys = dss(G, minimal = true); # build continuous-time descriptor system realization
pause()

println("Pole-zero analysis shows that G(s) has poles and zeros in the origin and at infinity")
println("pol = gpole(sys);")
println("zer = gzero(sys);")
pol = gpole(sys);
zer = gzero(sys);
println("pol = ")
display(pol)
println("zer = ")
display(zer)
pause()

println("Define a bilinear transformation with g(s) = (s+0.01)/(1+0.01*s) to make")
println("all poles and zeros stable and finite using")
println("g = (s+0.01)/(1+0.01*s);")
g = (s+0.01)/(1+0.01*s);
println("\ng = \n")
print(g)
pause()

println("Compute the transformed system and resulting stable poles and zeros")
println("syst = gbilin(sys,g, atol = 1.e-7,minimal = true)[1];")
syst = gbilin(sys,g, atol = 1.e-7,minimal = true)[1];
pol_new = gpole(syst)
zer_new = gzero(syst)
println("pol_new = ")
display(pol_new)
println("zer_new = ")
display(zer_new)
pause()

println("Compute the ν-gap distance between the two models using Vinnicombe's formula")
println("nugap = gnugap(sys,syst)[1]")
nugap = gnugap(sys,syst)[1]
println("The resulting ν-gap distance $nugap reflects well to the s-variable perturbation")

end # module


"""
    VaR(v::AbstractArray,p::Real;rev::Bool=false)

The `p`th quantile of the vector `v` is the Value at Risk. Assumes more positive values are higher risk measures, so a higher p will return a more positive number, but this can be reversed if `rev` is `true`.

Also can be called with `ValueAtRisk(...)`.
"""
function VaR(v,p;rev=false)
    if rev
        return StatsBase.quantile(v,1-p)
    else
        return StatsBase.quantile(v,p)
    end
end
"""
[VaR](@ref)
"""
ValueAtRisk = VaR

"""
    CTE(v::AbstractArray,p::Real;rev::Bool=false)

The average of the values ≥ the `p`th percentile of the vector `v` is the Conditiona Tail Expectation. Assumes more positive values are higher risk measures, so a higher p will return a more positive number, but this can be reversed if `rev` is `true`.

May also be called with `ConditionalTailExpectation(...)`.

Also known as Tail Value at Risk (TVaR), or Tail Conditional Expectation (TCE)
"""
function CTE(v,p;rev=false)
    # filter has the "or approximately equalt to quantile" because
    # of floating point path might make the quantile slightly off from the right indexing
    # e.g. if values should capture <= q, where q should be 10 but is calculated to be
    # 9.99999...
    if rev
        q = StatsBase.quantile(v,1-p)
        filter = (v .<= q) .| (v .≈ q)
    else
        q = StatsBase.quantile(v,p)
        filter = (v .>= q) .| (v .≈ q)
    end

    return sum(v[filter]) / sum(filter)

end

"""
[CTE](@ref)
"""
ConditionalTailExpectation = CTE


function norm_L1(a::Matrix{T}) where {T}
    col_abs_sums = sum(a, dims = 1)
    return maximum(col_abs_sums)
end

function norm_Linf(a::Matrix{T}) where {T}
    row_abs_sums = sum(a, dims = 2)
    return maximum(row_abs_sums)
end

function norm(a::Matrix{T}, norm_sym) where {T}
    if norm_sym === :1
        return norm_L1(a)
    elseif norm_sym === :2
        return norm_L2(a)
    elseif norm_sym === :inf
        return norm_Linf(a)
    else
        # some other p norm
        p = Float64(norm_sym)
        println("for other than 1, 2, or inf there is no implementation yet")
        return zero(T)
    end
end


"""
    epsp(t::Real, lastspike, q, taum, taus)
    epsp(t::Real, lastspike::AbstractArray{<:Real}, q::AbstractArray{<:Real}, taum::AbstractArray{<:Real}, taus::AbstractArray{<:Real})
    epsp(t::Real, lastspike::CuVecOrMat{<:Real}, q::CuVecOrMat{<:Real}, taum::CuVecOrMat{<:Real}, taus::CuVecOrMat{<:Real})

Evaluate an EPSP synapse. Modeled as `(ϵ₀ / τm - τs) * (exp(-Δ / τm) - exp(-Δ / τs)) Θ(Δ)`
  (where `Θ` is the Heaviside function and `Δ = t - lastspike`).
Specifically, this is the EPSP time course for the SRM0 model introduced by Gerstner.
Details: [Spiking Neuron Models: Single Neurons, Populations, Plasticity]
         (https://icwww.epfl.ch/~gerstner/SPNM/node27.html#SECTION02323400000000000000)

Use `CuVector` instead of `Vector` for GPU support.

# Fields
- `t`: current time
- `lastspike`: last pre-synaptic spike time
- `q`: amplitude
- `taum`: rise time constant
- `taus`: fall time constant
"""
function epsp(t::Real, lastspike, q, taum, taus)
    Δ = t - lastspike

    return (Δ >= 0 && Δ < Inf && taus != taum) * q / (1 - taus / taum) * (exp(-Δ / taum) - exp(-Δ / taus))
end
function epsp(t::Real, lastspike::AbstractArray{<:Real}, q::AbstractArray{<:Real}, taum::AbstractArray{<:Real}, taus::AbstractArray{<:Real})
    Δ = t .- lastspike
    I = @. q / (1 - taus / taum) * (exp(-Δ / taum) - exp(-Δ / taus))

    return map((ts, tm, δ, i) -> (δ >= 0) && (δ < Inf) && (ts != tm) ? i : zero(i), taus, taum, Δ, I)
end
function epsp(t::Real, lastspike::CuVecOrMat{<:Real}, q::CuVecOrMat{<:Real}, taum::CuVecOrMat{<:Real}, taus::CuVecOrMat{<:Real})
    Δ = t .- lastspike

    return @. (Δ >= 0) * (Δ < Inf) * (taus != taum) * q / (1 - taus / taum) * (exp(-Δ / taum) - exp(-Δ / taus))
end


function maprange(s, a, b)
    a₁, a₂ = minimum(a), maximum(a)
    b₁, b₂ = minimum(b), maximum(b)
    return b₁ + (s - a₁) * (b₂ - b₁) / (a₂ - a₁)
end

@show maprange(6, 1:10, -1:0)
@show maprange(0:10, 0:10, -1:0)


include("datafitting.jl")
function f(x)
    t = collect(0:3)
    y = [2.0; 0.7; 0.3; 0.1]
    r = y - x[1].*exp.(x[2].*t)
    f =  sum(r.^2)/2
    jac = [-exp.(x[2].*t) -x[1].*t.*exp.(x[2].*t)]
    g = jac'*r
    return f, jac, g, r
end

# good initial guess
x0 = [1;0]
x = gaussnewton(x0, f)
println(x)

# bad initial guess
x0 = [-3;3]
x = gaussnewton(x0, f)
println(x)

# Levenberg-Marquardt with good initial guess
x0 = [1;0]
x = levenbergmarquardt(x0, f)
println(x)

# Levenberg-Marquardt with bad initial guess
x0 = [-3;3]
x = levenbergmarquardt(x0, f)
println(x)


# generates a random positive definite matrix.
function generaterandomposdefmat(n)
    A = rand(n,n)
    out = A'*A
    return forcesymmetric(out)
end

function generaterandomsym(n)
    return extractsym(rand(n,n))
end

function generaterandomskew(n)
    return extractskew(rand(n,n))
end

function generaterandomortho(n)
    A = randn(n,n)
    B = A*A'
    s,Q = LinearAlgebra.eigen(B)
    return Q
end

function generateStiefel(n,p)
    Q = generaterandomortho(p)
    out = zeros(Float64, n,p)
    out[1:p,1:end] = Q
    return out
end

function generaterandomhankel(n)
    a = rand(2*n-1)
    out = Matrix{Float64}(undef,n,n)
    for j = 1:n
        for i = 1:n
            out[i,j] = a[i+j-1]
        end
    end
    return out
end

# function generateisotopysubgroup(n,p)
#     O = generaterandomortho(n-p)
#     out = zeros(Float64, n,n)
#     out[1:p,1:p] = eye(p)
#     out[(n-p):end, (n-p):end] = O
#     return out
# end

# generate an invertible nxn indefinite matrix, with k positive eigenvalues.
function generateindefmatrix(n,k=rand(2:n-1))
    @assert 2 <= k <= n

    U,S,V = LinearAlgebra.svd(rand(n,n))
    for i = 1:k
        S[i] = rand() # rand() is between [0,1]
    end
    for i = k+1:length(S)
        S[i] = -rand()
    end
    return U*LinearAlgebra.diagm(0=>S)*U'
end

function generateposdefmatrix(n)
    U,S,V = LinearAlgebra.svd(rand(n,n))
    S = rand(n)
    out = U*LinearAlgebra.diagm(0=>S)*U'

    return forcesymmetric(out)
end

function generatenegdefmatrix(n)
    U,S,V = LinearAlgebra.svd(rand(n,n))
    S = -rand(n)
    out = U*LinearAlgebra.diagm(0=>S)*U'

    return forcesymmetric(out)
end


module lineq
using LinearAlgebra

function linear_index(i, j, sz)
    n_row = sz[1]
    return i + (j - 1) * n_row
end

function matrix_index(k, sz)
    n_row = sz[1]
    (j, r) = divrem(k, n_row)
    j = (r>0) + j
    i = k - (j-1) * n_row
    return (i, j)
end

struct JacobiOptions
    max_iter::Int
    n_iter_norm_check::Int
    delta_x_tol::Float64
    damping::Float64
end
function JacobiOptions()
    return JacobiOptions(100, 10, 1e-6, 1.0)
end
function JacobiOptions(;
    max_iter = 100,
    n_iter_norm_check = 10,
    delta_x_tol = 1e-15,
    damping = 1.0)

    return JacobiOptions(
        max_iter,
        n_iter_norm_check,
        delta_x_tol,
        damping)
end

function delta_x_norm_check(x_new, x_old, options)
    break_flag = false
    norm_dx = norm(x_new - x_old)
    if norm_dx <= options.delta_x_tol
        break_flag = true
    end
    return break_flag
end

function jacobi_iteration(A, b, options::JacobiOptions = JacobiOptions())

    n_dim = length(b)
    x_old = zeros(eltype(b), n_dim)
    x_new = zeros(eltype(b), n_dim)

    d_vec = diag(A)

    use_damping = options.damping == 1.0

    iter_norm_check = 0
    iter = 0
    for outer iter = 1:options.max_iter
        iter_norm_check += 1

        # stepping algorithm
        x_new = ( b + d_vec .* x_old - (A * x_old) ) ./ d_vec

        if use_damping
            x_new = x_new * options.damping + x_old * (1 - options.damping)
        end

        # check for stopping conditions
        if iter_norm_check >= options.n_iter_norm_check
            iter_norm_check = 0
            break_flag = delta_x_norm_check(x_new, x_old, options)
            if break_flag
                break
            end
        end

        # update x_old
        x_old = x_new
    end

    return (x_new, iter)
end

struct GaussSeidelOptions
    max_iter::Int
    n_iter_norm_check::Int
    delta_x_tol::Float64
    damping::Float64
end
function GaussSeidelOptions()
    return GaussSeidelOptions(100, 10, 1e-6, 1.0)
end
function GaussSeidelOptions(;
    max_iter = 100,
    n_iter_norm_check = 10,
    delta_x_tol = 1e-15,
    damping = 1.0)

    return GaussSeidelOptions(
        max_iter,
        n_iter_norm_check,
        delta_x_tol,
        damping)
end

function gauss_seidel_iteration(A, b, options::GaussSeidelOptions)

    n_dim = length(b)
    x_old = zeros(eltype(b), n_dim)
    x_new = zeros(eltype(b), n_dim)

    d_vec = diag(A)

    use_damping = options.damping == 1.0

    iter_norm_check = 0
    iter = 0
    for outer iter = 1:options.max_iter
        iter_norm_check += 1

        # stepping algorithm

        if use_damping
            x_new = x_new * options.damping + x_old * (1 - options.damping)
        end

        # check for stopping conditions
        if iter_norm_check >= options.n_iter_norm_check
            iter_norm_check = 0
            break_flag = delta_x_norm_check(x_new, x_old, options)
            if break_flag
                break
            end
        end

    end

    return (x_new, iter)
end

end # dev

module dtest
using ..lineq
using LinearAlgebra

function create_matrix(n_dim = 8)

    A = zeros(Float64, n_dim, n_dim)

    main_diag = [1; 4 * ones(n_dim-2); 4]
    side_diag = -ones(n_dim-1)
    sideside_diag = -ones(n_dim-3)

    main_diag_idx = diagind(A, 0)
    side_m1_diag_idx = diagind(A, 1)
    side_p1_diag_idx = diagind(A, -1)

    side_m3_diag_idx = diagind(A, 3)
    side_p3_diag_idx = diagind(A, -3)

    A[ main_diag_idx ] = main_diag
    A[ side_m1_diag_idx ] = side_diag
    A[ side_p1_diag_idx ] = side_diag
    A[ side_m3_diag_idx ] = sideside_diag
    A[ side_p3_diag_idx ] = sideside_diag

    return A
end

function create_vector(n_dim = 8)

    b = zeros(n_dim)
    b[1]   = -1.0
    b[end] = 2.0

    return b
end

n_dim = 8 # 100
A = create_matrix(n_dim)
b = create_vector(n_dim)

x_theo = A \ b

# jac_options = lineq.JacobiOptions(delta_x_tol = 1e-10, max_iter = 10000, damping = 1.0)
# (x_jac, iter) = lineq.jacobi_iteration(A, b, jac_options)

# println("Norm diff between theo and jac: $(norm(x_theo - x_jac))")
# println("Needed iterations: $(iter)")

gs_options = lineq.GaussSeidelOptions(delta_x_tol = 1e-10, max_iter = 100, damping = 1.0)
(x_gs, iter) = lineq.gauss_seidel_iteration(A, b, gs_options)

println("Norm diff between theo and gs: $(norm(x_theo - x_gs))")
println("Needed iterations: $(iter)")

D = diagm( diag(A) )
E = -LowerTriangular(A)
F = -UpperTriangular(A)

end


#########################
# Spin Operator Methods #
#########################
immutable HalfSpin{T<:Integer}
    val::T
    HalfSpin(val) = val >= 0 ? new(val) : error("Spin must be positive")
end

HalfSpin{T}(val::T) = HalfSpin{T}(val)

spin(s::HalfSpin) = s
spin(j::Integer) = HalfSpin(2*j)
function spin(j)
    if isinteger(j)
        return spin(@compat(Int(j))) # j = 0.0, 1.0, 2.0...
    elseif isinteger(2*j)
        return HalfSpin(@compat(Int(2*j))) # j = .5, 1.5, 2.5...
    else
        error("Spin must be a multiple of 1/2.")
    end
end

spin_value(s::HalfSpin) = s.val/2

#####################
# Auxiliary Methods #
#####################
spin_coeffs(j, scalar) = [scalar*sqrt(j*(j+1)-(x+1)*x) for x in j-1:-1:-j]

function spinjyx_mat(s::HalfSpin, a, b)
    N = s.val + 1
    j = spin_value(s)
    up_coeffs = spin_coeffs(j, a)
    down_coeffs = b == 1 ? up_coeffs : b * up_coeffs
    return spdiagm((up_coeffs, down_coeffs), (1, -1), N, N)
end

spinjx_mat(j) = spinjyx_mat(spin(j), .5, 1)
spinjy_mat(j) = spinjyx_mat(spin(j), -.5im, -1)

function spinjpm_mat(s::HalfSpin, k)
    N = s.val + 1
    return spdiagm(spin_coeffs(spin_value(s), 1), k, N, N)
end

spinjp_mat(j) = spinjpm_mat(spin(j), 1)
spinjm_mat(j) = spinjpm_mat(spin(j), -1)

function spinjz_mat(s::HalfSpin)
    j = spin_value(s)
    return spdiagm(j:-1:-j)
end

spinjz_mat(j) = spinjz_mat(spin(j))

################################
# Jx, Jy, Jz, Jp, Jm operators #
################################
# These Wikipedia pages give a nice overview of spin/ladder operators:
# http://en.wikipedia.org/wiki/Spin_%28physics%29#Operator
# http://en.wikipedia.org/wiki/Ladder_operator#Angular_momentum
spinjx(j) = QuArray(spinjx_mat(j))
spinjy(j) = QuArray(spinjy_mat(j))
spinjz(j) = QuArray(spinjz_mat(j))
spinjp(j) = QuArray(spinjp_mat(j))
spinjm(j) = QuArray(spinjm_mat(j))

############################
#  Pauli spin 1/2 matrices #
############################
const sigmax = 2.0 * full(spinjx(1/2))
const sigmay = 2.0 * full(spinjy(1/2))
const sigmaz = 2.0 * full(spinjz(1/2))
const sigmap = full(sigmax) + im * full(sigmay)
const sigmam = full(sigmax) - im * full(sigmay)

# TODO : unicode sigmay, sigmaz
const σₓ = sigmax

export spinjx,
    spinjy,
    spinjz,
    spinjp,
    spinjm,
    sigmax,
    sigmay,
    sigmaz,
    sigmap,
    sigmam,
    σₓ


function nlp_cvx_108_010(
    optimizer,
    objective_tol,
    primal_tol,
    dual_tol,
    termination_target = TERMINATION_TARGET_LOCAL,
    primal_target = PRIMAL_TARGET_LOCAL,
)
    # Test Goals:
    # - convex objective
    # - intersection of constraints
    # - convex constraint, only in positive domain
    # Variants
    #   010 - no binding constraints
    #   011 - intersection point
    #   012 - intersection point
    #   013 - one binding constraint

    model = Model(optimizer)

    @variable(model, x >= 0)
    @variable(model, y >= 0)

    @objective(model, Min, (x - 1.0)^2 + (y - 0.75)^2)
    @NLconstraint(model, 2 * x^2 - 4x * y - 4 * x + 4 <= y)
    @NLconstraint(model, y^2 <= -x + 2)

    optimize!(model)

    check_status(model, FEASIBLE_PROBLEM, termination_target, primal_target)
    check_objective(model, 0, tol = objective_tol)
    return check_solution([x, y], [1, 0.75], tol = primal_tol)
end


"""
    sigma2(pk, R)

Variance of mass fluctuation, σ^2, with a top-hat filter with radius `R`

# Arguments
- `pk`(k): a function which returns a linear matter power spectrum with the argument k being the comoving wavenumber.
   - **pk times k^3 must be dimensionless**. For example, if k is in units of h/Mpc, `pk` must be in units of Mpc^3/h^3.
- `R::Real`: a top-hat smoothing scale.
   - **R times k must be dimensionless**. For example, if k is in units of h/Mpc, `R` must be in units of Mpc/h.

`sigma2` is computed as

``σ^2(R) = ∫_0^∞ k^2dk W^2(kR) pk(k) / (2π^2)``

where and W(kR) is a window function given by

``W(x) = (3/x)j_1(x) = (3/x)(sin(x)/x^2 - cos(x)/x)``

j_1(x) is the spherical Bessel function of order 1.

This function is based on [Cosmology Routine Library (CRL)](https://wwwmpa.mpa-garching.mpg.de/~komatsu/crl/).
"""
function sigma2(pk, R::Real)
   function dσ2dk(k)
      x = k * R
      W = (3 / x) * (sin(x) / x^2 - cos(x) / x)
      dσ2dk = W^2 * pk(k) * k^2 / 2 / π^2
   end
   res, err = quadde(dσ2dk, 0, 20 / R)
   σ2 = res
end

"""
    dsigma2dR(pk, R)

Derivative of variance of mass fluctuation, dσ^2/dR,  with respect to a top-hat filter with radius `R`.

# Arguments
- `pk`(k): a function which returns a linear matter power spectrum with the argument k being the comoving wavenumber.
   - **pk times k^3 must be dimensionless**. For example, if k is in units of h/Mpc, `pk` must be in units of Mpc^3/h^3.
- `R::Real`: a top-hat smoothing scale.
   - **R times k must be dimensionless**. For example, if k is in units of h/Mpc, `R` must be in units of Mpc/h.

`dsigma2/dR` is computed as

``dσ^2/dR = 2 ∫_0^∞ k^2dk W(kR) dW/dR pk(k) / (2π^2)``

where and W(kR) is a window function given by

``W(x) = (3/x)j_1(x) = (3/x)(sin(x)/x^2 - cos(x)/x)``

and dW/dR is its derivative:

``dW/dR = k dW/dx = (-3k/x)j_2(x) = (-3k/x)[(3/x^2-1)sin(x)/x - 3cos(x)/x^2]``

j_n(x) is the spherical Bessel function of order n.

This function is based on [Cosmology Routine Library (CRL)](https://wwwmpa.mpa-garching.mpg.de/~komatsu/crl/).
"""
function dsigma2dR(pk, R::Real)
   function dσ2dRdk(k)
      x = k * R
      W = (3 / x) * (sin(x) / x^2 - cos(x) / x)
      dWdx = (-3 / x) * ((3 / x^2 - 1) * sin(x) / x - 3 * cos(x) / x^2)
      dσ2dRdk = 2 * W * dWdx * pk(k) * k^3 / 2 / π^2
   end
   res, err = quadde(dσ2dRdk, 0, 20 / R)
   dσ2dR = res
end

"""
    sigma2gaus(pk, R)

Variance of mass fluctuation, σ^2(`R`), with a Gaussian filter with width `R`.

# Arguments
- `pk`(k): a function which returns a linear matter power spectrum with the argument k being the comoving wavenumber.
   - **pk times k^3 must be dimensionless**. For example, if k is in units of h/Mpc, `pk` must be in units of Mpc^3/h^3.
- `R::Real`: a top-hat smoothing scale.
   - **R times k must be dimensionless**. For example, if k is in units of h/Mpc, `R` must be in units of Mpc/h.

`sigma2gaus(R)` is computed as

``σ^2(R) = ∫_0^∞ k^2dk exp(-k^2 R^2) pk(k) / (2π^2)``

This function is based on [Cosmology Routine Library (CRL)](https://wwwmpa.mpa-garching.mpg.de/~komatsu/crl/).
"""
function sigma2gaus(pk, R::Real)
   function dσ2dk(k)
      x = k * R
      dσ2dk = exp(-x^2) * pk(k) * k^2 / 2 / π^2
   end
   res, err = quadde(dσ2dk, 0, 10 / R)
   σ2 = res
end

"""
    dsigma2gausdR(pk, R)

Derivative of variance of mass fluctuation, dσ^2/dR, with respect to a Gaussian filter with width `R`.

# Arguments
- `pk`(k): a function which returns a linear matter power spectrum with the argument k being the comoving wavenumber.
   - **pk times k^3 must be dimensionless**. For example, if k is in units of h/Mpc, `pk` must be in units of Mpc^3/h^3.
- `R::Real`: a top-hat smoothing scale.
   - **R times k must be dimensionless**. For example, if k is in units of h/Mpc, `R` must be in units of Mpc/h.

`dsigma2gaus/dR` is computed as

``dσ^2/dR = -2R ∫_0^∞ k^4dk exp(-k^2 R^2) pk(k) / (2π^2)``

This function is based on [Cosmology Routine Library (CRL)](https://wwwmpa.mpa-garching.mpg.de/~komatsu/crl/).
"""
function dsigma2gausdR(pk, R::Real)
   function dσ2dRdk(k)
      x = k * R
      dσ2dRdk = -2R * exp(-x^2) * pk(k) * k^4 / 2 / π^2
   end
   res, err = quadde(dσ2dRdk, 0, 10 / R)
   dσ2dR = res
end

"""
    d2sigma2gausdR2(pk, R)

Second derivative of variance of mass fluctuation, d^2σ^2/dR^2, with respect to a Gaussian filter with width `R`.

# Arguments
- `pk`(k): a function which returns a linear matter power spectrum with the argument k being the comoving wavenumber.
   - **pk times k^3 must be dimensionless**. For example, if k is in units of h/Mpc, `pk` must be in units of Mpc^3/h^3.
- `R::Real`: a top-hat smoothing scale.
   - **R times k must be dimensionless**. For example, if k is in units of h/Mpc, `R` must be in units of Mpc/h.

`d2sigma2gaus/dR2` is computed as

``d^2σ^2/dR^2 = -2 ∫_0^∞ (1 - 2x^2) k^4dk exp(-k^2 R^2) pk(k) / (2π^2)``

This function is based on [Cosmology Routine Library (CRL)](https://wwwmpa.mpa-garching.mpg.de/~komatsu/crl/).
"""
function d2sigma2gausdR2(pk, R::Real)
   function d2σ2dR2dk(k)
      x = k * R
      d2σ2dR2dk = (-2 + 4 * x^2) * exp(-x^2) * pk(k) * k^4 / 2 / π^2
   end
   res, err = quadde(d2σ2dR2dk, 0, 10 / R)
   d2σ2dR2 = res
end


roundc(r) = round.(real(r)) + im*round.(imag(r));
#round{Td}(r::Td) = Td<:Complex?round(real(r)) + im*round(imag(r)):round(r);

#################################################################
function scodes(M,Ctype)
#C = scodes(M,Ctype)
#  Return a constellation of M signals, where <Ctype> is 'PSK', 'QAM',
#  or 'PAM'.
#[C,d] = scodes(M,Ctype)
#  Also return the minimum distance between constellation points.
#[C,d,Cr] = scodes(M,Ctype)
#  Returns the underlying real PAM constellation for square (M=m^2)
#  QAM constellations

# By Christian Peel  (chris.peel@ieee.org)
# Last Modified: Sun 14 Apr 19, 6:53pm
#

#if nargin==0
#  error("Bad arguments. Try again.")
#end
if Ctype=="PSK"
    C=exp(im*2*pi*[1:M]'/M);
elseif Ctype=="QAM"
    if abs(sqrt(M)-round(sqrt(M)))< 1e-3
        # Make QAM constellations if M is a square
        m = round(Integer,sqrt(M))
        C = complex(zeros(m,m))
        Cr = complex([-(m-1):2:(m-1);])
        for ix = 1:m
            for jx = 1:m
                C[jx,ix] = Cr[ix] +im*Cr[jx];
            end
        end
    else
        error("Can only handle square QAM constellations.")
        #     [x y] = qaskenco(M);
        #     C = x+im*y;
    end
elseif Ctype=="PAM"
    C = [-(M-1):2:(M-1);];
    Cr = C;
else
  error("Bad args.");
end

C = C[:]; # Turn into column vector
# for QAM, gam = sqrt(2/3*(M-1))
gam = sqrt(C'*C/length(C));

C = C/gam;
#if nargout ==3
  Cr = Cr/gam;
#end

d2 = real(C[2:end] .- C[1]);
id = findall(abs.(d2).>0);
d = minimum(abs.(d2[id]));
return (C, d, Cr, gam)
end

#######################################################################
function mimo_slice(Y,H,C)
M = size(Y,1);
Td = size(Y,2);
Zd = zeros(M,Td);
Mc = length(C);
dy = zeros(M,Mc);
for t=1:Td
    for i = 1:Mc
        dd = Y[:,t]-H[:,t]*C[i];
        dy[:,i] = abs.(dd).^2
    end
    for m = 1:M
        mn,zd = findmin(dy[m,:]);
        Zd[m,t] = zd[1]-1.0;
    end
end
return Zd
end

#################################################################
function siso_demod(y,C)
Mc = length(C);
dy = zeros(Mc,1);
for ix = 1:Mc
    dy[ix] = abs(y[1]-C[ix])
end
m,zd = findmin(dy);
sd = C[zd];
return sd[1]
end


# logarithmic barrier function

export LogBarrier

"""
**Logarithmic barrier**

    LogBarrier(a=1, b=0, μ=1)

Returns the function
```math
f(x) = -μ⋅∑_i\\log(a⋅x_i+b),
```
for a nonnegative parameter `μ`.
"""
struct LogBarrier{R, S, T} <: ProximableFunction
    a::R
    b::S
    mu::T
    function LogBarrier{R, S, T}(a::R, b::S, mu::T) where {R, S, T}
        if mu <= 0
            error("parameter mu must be positive")
        else
            new(a, b, mu)
        end
    end
end

is_separable(f::LogBarrier) = true
is_convex(f::LogBarrier) = true

LogBarrier(a::R=1, b::S=0, mu::T=1) where {R, S, T} = LogBarrier{R, S, T}(a, b, mu)

function (f::LogBarrier)(x::AbstractArray{T}) where T <: Real
    sumf = T(0)
    v = T(0)
    for i in eachindex(x)
        v = f.a * x[i] + f.b
        if v <= T(0)
            return +Inf
        end
        sumf += log(v)
    end
    return -f.mu * sumf
end

function prox!(y::AbstractArray{T}, f::LogBarrier, x::AbstractArray{T}, gamma::Real=1) where T <: Real
    par = 4 * gamma * f.mu * f.a * f.a
    sumf = T(0)
    z = T(0)
    v = T(0)
    for i in eachindex(x)
        z = f.a * x[i] + f.b
        v = (z + sqrt(z * z + par)) / 2
        y[i] = (v - f.b) / f.a
        sumf += log(v)
    end
    return -f.mu * sumf
end

function prox!(y::AbstractArray{T}, f::LogBarrier, x::AbstractArray{T}, gamma::AbstractArray) where T <: Real
    par = 4 * f.mu * f.a * f.a
    sumf = T(0)
    z = T(0)
    v = T(0)
    for i in eachindex(x)
        par_i = gamma[i] * par
        z = f.a * x[i] + f.b
        v = (z + sqrt(z * z + par_i)) / 2
        y[i] = (v - f.b) / f.a
        sumf += log(v)
    end
    return -f.mu * sumf
end

function gradient!(y::AbstractArray{T}, f::LogBarrier, x::AbstractArray{T}) where T <: Real
    sum = T(0)
    for i in eachindex(x)
        logarg = f.a * x[i] + f.b
        y[i] = -f.mu * f.a / logarg
        sum += log(logarg)
    end
    sum *= -f.mu
    return sum
end

fun_name(f::LogBarrier) = "logarithmic barrier"
fun_dom(f::LogBarrier) = "AbstractArray{Real}"
fun_expr(f::LogBarrier) = "x ↦ -μ * sum( log(a*x_i+b), i=1,...,n )"
fun_params(f::LogBarrier) = "a = $(f.a), b = $(f.b), μ = $(f.mu)"

function prox_naive(f::LogBarrier, x::AbstractArray{T,1}, gamma::Union{Real, AbstractArray}=1) where T <: Real
    asqr = f.a * f.a
    z = f.a * x .+ f.b
    y = ((z .+ sqrt.(z .* z .+ 4 * gamma * f.mu * asqr)) / 2 .- f.b) / f.a
    fy = -f.mu * sum(log.(f.a .* y .+ f.b))
    return y, fy
end


using Distributions
using Plots

# Sample x_{k+1} given x_k
# Parameters:
# - x_k: The kth X sample. One of -1, 0, or 1
# Outputs:
# - x_{k+1}: The (k+1)th X sample
function sample_x(xk)
    # Matrix defining p(x_{k+1}|x_k)
    probability_matrix = [0.9 0.05 0.05; 0.05 0.9 0.05; 0.05 0.05 0.9]

    # Extract row from probability matrix
    distribution = xk + 2

    # Define cumulative sum of distribution for sampling
    cumulative_distribution = cumsum(probability_matrix[Int(distribution),:])

    # Obtain x_k
    unif_sample = rand()
    xk_values = [-1 0 1]

    for i=1:3
        if unif_sample < cumulative_distribution[i]
            return xk_values[i]
        end
    end
end

# Obtains n samples of X based on HMM
# Inputs
# - n: An integer representing the number of samples to return
# Outputs
# A vector of length n where the kth element is the kth sample of X
function get_x_samples(n)
    # Store x samples
    samples = zeros(n)

    # Randomly set x1
    samples[1] = rand(-1:1)

    # Get samples x_2 through x_n
    for i=2:n
        samples[i] = sample_x(samples[i-1])
    end

    return samples
end

# Obtains n samples of Y based on HMM
# Inputs
# - x_samples: A vector of length n representing the samples of X
# Outputs
# A vector of length n where the kth element is the kth sample of Y
function get_y_samples(x_samples)
    # Store y samples
    n = length(x_samples)
    y_samples = zeros(n)

    for i=1:n
        𝒩 = Normal(x_samples[i],1)
        y_samples[i] = rand(𝒩)
    end

    return y_samples
end

# Obtian and plot samples
n = 1000
x_samples = get_x_samples(n)
y_samples = get_y_samples(x_samples)
plot(1:n,x_samples, seriestype=:scatter, bg = RGB(247/255, 236/255, 226/255), color = RGB(0,191/255,255/255), label = "x", alpha = 0.2)
plot!(1:n,y_samples, seriestype=:scatter, color = RGB(191/255,1,0), label =
"y", alpha = 0.2)

cd("/Users/elvis/Documents/DSI/courses/content/graphical-models/images/")
savefig("bn_sampling_ex_1.svg")


module NewtonsMethod

using LinearAlgebra, ForwardDiff

function newtonroot(f, f′; x₀, tolerance=1E-7, maxiter=1000)
    x_old = x₀
    normdiff = Inf
    iter = 1
    while normdiff > tolerance && iter <= maxiter
        x_new = x_old - (f(x_old) / f′(x_old))
        normdiff = norm(x_new - x_old)
        x_old = x_new
        iter = iter + 1
    end
    if iter > maxiter && normdiff > tolerance
        return NaN
    else
        return x_old
    end
end

function newtonroot(f; x₀, tolerance=1E-7, maxiter=1000)
    x_old = x₀
    f′ = x -> ForwardDiff.derivative(f, x)
    normdiff = Inf
    iter = 1
    while normdiff > tolerance && iter <= maxiter
        x_new = x_old - (f(x_old) / f′(x_old))
        normdiff = norm(x_new - x_old)
        x_old = x_new
        iter = iter + 1
    end
    if iter > maxiter && normdiff > tolerance
        return NaN
    else
        return x_old
    end
end

export newtonroot

end


"Construction of the Kravchuk functions for at most 256 points (j ≤ 128)"

using HypergeometricFunctions
using SpecialFunctions

export kp, binom, Φ, Φ2

"`kp(n, q, N)` gives the symmetric Kravchuk polynomia for n ∈ [0, N] and q ∈ [-N,N]"
kp(n, q, N) = HypergeometricFunctions._₂F₁general2(-BigFloat(n), -q, -N, 2)

A(n, j) = ((-1.0)^n) / (2.0^j)

"`binom(n, k)` gives the binomial coefficient using the Beta function definition"
binom(n, k) = (1/n) * (1/beta(k + 1, n - k + 1))

B(n, q, j) = √(binom(2 * j, n) * binom(2 * j, j + q))

"Fundamental definition of the Kravchuk function"
ϕ(n, q, j) = Float64(A(n, j) * B(n, q, j) * kp(n, j + q, 2 * j))

"`Φ(n, q, j)` gives the Kravchuk function within the ranges n ∈ [0, 2j], q ∈ [-j, j]. It is tested to give correct values for j ≤ 128."
function Φ(n, q, j)
    if 0 ≤ n ≤ j
        return ϕ(n, q, j)
    elseif j < n ≤ 2 * j
        return (-1.0)^q * ϕ(2 * j - n, q, j)
    end
end

"`Φ2(n₁, n₂, q₁, q₂, j)` evaluates the two-dimensional Kravchuk functions over the set of lineal independent eigenmodes"
Φ2(n₁, n₂, q₁, q₂, j) = Φ(n₁, q₁, j) * Φ(n₂, q₂, j)


# p2.jl - convergence of periodic spectral method (compare p1.jl)

# For various N (even), set up grid as before:
clf(); axes([.1,.4,.8,.5]);
for N = 2:2:100
    h = 2*pi/N;  x = [ -pi + i*h for i = 1:N ];
    u = @. exp(sin(x)); uprime = @. cos(x)*u;

    # Construct spectral differentiation matrix:
    column = [ 0; @. .5*(-1)^(1:N-1)*cot((1:N-1)*h/2) ];
    D = toeplitz(column,column[[1;N:-1:2]]);

    # Plot max(abs(D*u-uprime)):
    error = norm(D*u-uprime,Inf);
    loglog(N,error,"k.",markersize=6);
end
grid(true); xlabel("N"); ylabel("error");
title("Convergence of spectral differentiation")


function hessian{T<:FloatingPoint}(f::Function,x::Array{T,1})

    m = length(x)

	e = eps(T)^(1/4)*Base.maxabs([x,one(T)])
	dh = eye(m)*e
    hess = Array(T,m,m)

    for i = 1:m
	    f1 = f(x+2*dh[:,i])
		f2 = f(x+dh[:,i])
		f3 = f(x)
		f4 = f(x-dh[:,i])
		f5 = f(x-2*dh[:,i])
		hess[i:i,i:i] = (-f1+16*f2-30*f3+16*f4-f5)/(12*e^2)
    end

    for i = 1:m-1
        for j = i+1:m
		    f1 = f(x+dh[:,i]+dh[:,j])
			f2 = f(x-dh[:,i]-dh[:,j])
			f3 = f(x-dh[:,i]+dh[:,j])
			f4 = f(x+dh[:,i]-dh[:,j])
			hess[i:i,j:j] = (f1+f2-f3-f4)/(4*e^2)
            hess[j,i] = hess[i,j]
        end
    end

    return hess

end


# # Motzkin

#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/generated/Getting started/motzkin.ipynb)
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/generated/Getting started/motzkin.ipynb)
# **Adapted from**: (3.6) and (3.19) of [BPT12]
#
# [BPT12] Blekherman, G.; Parrilo, P. A. & Thomas, R. R.
# *Semidefinite Optimization and Convex Algebraic Geometry*.
# Society for Industrial and Applied Mathematics, **2012**.

# The first explicit example of nonnegative polynomial that is not a sum of squares was found by Motzkin in 1967. By the [Arithmetic-geometric mean](https://en.wikipedia.org/wiki/Arithmetic%E2%80%93geometric_mean),
# $$ \frac{x^4y^2 + x^2y^4 + 1}{3} \ge \sqrt[3]{x^4y^2 \cdot x^2y^4 \cdot 1} = x^2y^2 $$
# hence
# $$ x^4y^2 + x^2y^4 + 1 - 3x^2y^2 \ge 0. $$
# The code belows construct the Motzkin polynomial using [DynamicPolynomials](https://github.com/JuliaAlgebra/DynamicPolynomials.jl).

using Test #src
using DynamicPolynomials
@polyvar x y
motzkin = x^4*y^2 + x^2*y^4 + 1 - 3x^2*y^2

# The Motzkin polynomial is nonnegative but is not a sum of squares as we can verify numerically as follows.
# We first need to pick an SDP solver, see [here](https://jump.dev/JuMP.jl/v0.21.6/installation/#Supported-solvers) for a list of the available choices.

using SumOfSquares
import CSDP
solver = optimizer_with_attributes(CSDP.Optimizer, MOI.Silent() => true)
model = SOSModel(solver)
@constraint(model, motzkin >= 0) # We constraint `motzkin` to be a sum of squares

optimize!(model)

# We see that the problem is detected as infeasible...

@test termination_status(model) == MOI.INFEASIBLE #src
termination_status(model)

# ... and that the dual solution is a certificate of the infeasibility of the problem.

@test dual_status(model) == MOI.INFEASIBILITY_CERTIFICATE #src
dual_status(model)

# Even if the Motzkin polynomial is not a sum of squares, it can still be certified to be nonnegative using sums of squares.
# Indeed a polynomial is certified to be nonnegative if it is equal to a fraction of sums of squares.
# The Motzkin polynomial is equal to a fraction of sums of squares whose denominator is $x^2 + y^2$.
# This can be verified numerically as follows:

model = SOSModel(solver)
@constraint(model, (x^2 + y^2) * motzkin >= 0) # We constraint the `(x^2 + y^2) * motzkin` to be a sum of squares

optimize!(model)

# Now the problem is declared feasible by the solver...

@test termination_status(model) == MOI.OPTIMAL #src
termination_status(model)

# ... and the primal solution is a feasible point, hence it is a certificate of nonnegativity of the Motzkin polynomial.

@test primal_status(model) == MOI.FEASIBLE_POINT #src
primal_status(model)

# One may consider ourself lucky to have had the intuition that $x^2 + y^2$ would work as denominator.
# In fact, the search for the denominator can be carried out in parallel to the search of the numerator.
# In the example below, we search for a denominator with monomials of degrees from 0 to 2.
# If none is found, we can increase the maximum degree 2 to 4, 6, 8, ...
# This gives a hierarchy of programs to try in order to certify the nonnegativity of a polynomial by identifying it with a fraction of sum of squares polynomials.
# In the case of the Motzkin polynomial we now that degree 2 is enough since $x^2 + y^2$ works.

model = SOSModel(solver)
X = monomials([x, y], 0:2)

# We create a quadratic polynomial that is not necessarily a sum of squares since this is implied by the next constraint: `deno >= 1`.

@variable(model, deno, Poly(X))

# We want the denominator polynomial to be strictly positive, this prevents the trivial solution deno = 0 for instance.

@constraint(model, deno >= 1)
@constraint(model, deno * motzkin >= 0)
optimize!(model)

@test termination_status(model) == MOI.OPTIMAL #src
termination_status(model)

@test primal_status(model) == MOI.FEASIBLE_POINT #src
primal_status(model)

# We can check the denominator found by the program using `JuMP.value`

value(deno)

# Because a picture is worth a thousand words let's plot the beast.
# We can easily extend `Plots` by adding a recipe to plot bivariate polynomials.

using RecipesBase
@recipe function f(x::AbstractVector, y::AbstractVector, p::Polynomial)
    x, y, (x, y) -> p(variables(p) => [x, y])
end
import Plots
Plots.plot(
    range(-2, stop=2, length=100),
    range(-2, stop=2, length=100),
    motzkin,
    st = [:surface],
    seriescolor=:heat,
    colorbar=:none,
    clims = (-10, 80)
)


# Example: Compartmental models in epidemiology
#=
- https://github.com/epirecipes/sir-julia/blob/master/markdown/function_map/function_map.md
- https://en.wikipedia.org/wiki/Compartmental_models_in_epidemiology#Deterministic_versus_stochastic_epidemic_models
=#
using ModelingToolkit, Test

@inline function rate_to_proportion(r,t)
    1-exp(-r*t)
end;

# Independent and dependent variables and parameters
@parameters t c nsteps δt β γ
D = Difference(t; dt=0.1)
@variables S(t) I(t) R(t) next_S(t) next_I(t) next_R(t)

infection = rate_to_proportion(β*c*I/(S+I+R),δt)*S
recovery = rate_to_proportion(γ,δt)*I

# Equations
eqs = [D(S) ~ S-infection,
       D(I) ~ I+infection-recovery,
       D(R) ~ R+recovery]

# System
sys = DiscreteSystem(eqs,t,[S,I,R],[c,nsteps,δt,β,γ]; controls = [β, γ])

# Problem
u0 = [S => 990.0, I => 10.0, R => 0.0]
p = [β => 0.05, c => 10.0, γ => 0.25, δt => 0.1, nsteps => 400]
tspan = (0.0,ModelingToolkit.value(substitute(nsteps,p))) # value function (from Symbolics) is used to convert a Num to Float64
prob_map = DiscreteProblem(sys,u0,tspan,p)

# Solution
using OrdinaryDiffEq
sol_map = solve(prob_map,FunctionMap());

# Direct Implementation

function sir_map!(u_diff,u,p,t)
    (S,I,R) = u
    (β,c,γ,δt) = p
    N = S+I+R
    infection = rate_to_proportion(β*c*I/N,δt)*S
    recovery = rate_to_proportion(γ,δt)*I
    @inbounds begin
        u_diff[1] = S-infection
        u_diff[2] = I+infection-recovery
        u_diff[3] = R+recovery
    end
    nothing
end;
u0 = [990.0,10.0,0.0];
p = [0.05,10.0,0.25,0.1];
prob_map = DiscreteProblem(sir_map!,u0,tspan,p);
sol_map2 = solve(prob_map,FunctionMap());

@test Array(sol_map) ≈ Array(sol_map2)


function countingsort(x::AbstractArray{<:Int})::AbstractArray{<:Int}
    l = length(x)
    n = maximum(x)+1
    count = zeros(Int, n)
    aux = copy(x)
    for i ∈ 1:l
        @inbounds count[x[i]+1]+=1
    end
    for i ∈ 2:n
        @inbounds count[i] += count[i-1]
    end
    for i ∈ l:-1:1
        @inbounds aux[count[x[i]+1]] = x[i]
        @inbounds count[x[i]+1]-=1
    end
    return aux
end


# Parts of this code were taken / derived from NetworkX. See LICENSE for
# licensing details.

"""
    pagerank(g::ADiGraph, α=0.85, n=100, ϵ = 1.0e-6)

Calculates the [PageRank](https://en.wikipedia.org/wiki/PageRank) of the graph
`g`. Can optionally specify a different damping factor (`α`), number of
iterations (`n`), and convergence threshold (`ϵ`). If convergence is not
reached within `n` iterations, an error will be returned.
"""
function pagerank(g::ADiGraph, α=0.85, n=100, ϵ = 1.0e-6)
    A = adjacency_matrix(g,:in,Float64)
    S = 1 ./ vec(sum(A, dims=1))
    S[S .== Inf] .= 0.0
    M = A' # need a separate line due to bug #17456 in julia
    M = (Diagonal(S) * M)'
    N = nv(g)
    x = fill(1/N, N)
    p = fill(1/N, N)
    dangling_weights = p
    is_dangling = findall(S .== 0)

    for _ in 1:n
        xlast = x
        x = α * (M * x + sum(x[is_dangling]) * dangling_weights) + (1 - α) * p
        err = sum(abs, x - xlast)
        err < N * ϵ && return x
    end
    error("Pagerank did not converge after $n iterations.")
end


# Continuation of ex 7.13, Least Squares WRM

using NumericalMethodsforEngineers
using Test

@sym begin
  ClearAll(xi, yi, N1, Y, Ydotdot, C11, ytilde1, ytilde2)
  xi = [0, 1//2, 1]
  yi = [0, a, 1]
  Y(x_) := LagrangePolynomial(xi, yi)
  #
  # Can be formulated as ytile(x) = F(x) + C1(a) * Ψ(x)
  #
  # F(x_) := 2*x^2 - x
  # C(a_) := -4a
  Ψ(x_) := x^2 - x
  #
  #Y(x_) := F(x) + C(a)*Ψ(x)
  #
  Ydotdot(x_) = D(Y(x), x, 2)
  R(x_) := Simplify(Expand(Ydotdot(x) - 3*x - 4*Y(x)))
  R(x_) = Simplify(Expand(R(x) ./ (a => -C1/4)))
  C11 = Solve(Integrate(R(x)*Ψ(x), [x, 0, 1]), C1)[1]
  SetJ(r, ToString(Simplify(R(x))))
  SetJ(C1, ToString(C11[1][2]))
  ytilde1 = Simplify(Y(x) ./ (a => -C1/4))
  ytilde2 = Simplify(ytilde1 ./ C11)
  SetJ(y, ToString(Simplify(Expand(ytilde2))));
end

println("\n\nExample 7.18: y''=3x + 4y, y(0)=0, y(1)=1")
println("using 1-point Galerkin Weighted Residual Method")
@sym Println("\nY(x) = $(Y(x))\n")
@sym Println("Y(x) = $(Simplify(Expand(Y(x) ./ (a => -C1/4))))\n")
@sym Println("R(x) = $(R(x))\n")
@sym Println("C1 = $(Solve(Integrate(R(x)*Ψ(x), [x, 0, 1]), C1)[1][1][2])\n")
@sym Println("ytilde_1pt_galerkin(x) = $(ytilde2)\n")
println("( Example 7.18 gives: ytilde = 1/4*x*(5x - 1) )\n")
@eval ytilde_1pt_galerkin(x) = $(Meta.parse(y))
@eval rf_1pt_galerkin(x, C1) = $(Meta.parse(r))

C1 = Meta.parse(C1)
rf_1pt_galerkin_1(x) = rf_1pt_galerkin(x, C1)
println()

@test r == "4 + 2C1 + x + 4C1*x - 8x^2 - 4C1*x^2"
@test y == "(1/4)*x*(-1 + 5x)"
#@test (quadgk(rf_1pt_galerkin_1, 0, 1))[1] < 5*eps()


module lossFunctions
  export hingeLoss, MSE, zeroOneLoss

  """
      function hingeLoss( y::Float64, t::Float64 )

      compute hinge loss for a single binary classification.
      #Arguments
        - `y::Float64`: predicted value.
        - `t::Float64`: ground truth value.
      #Return
        - `Float64`:    hingeLoss error.
  """
  function hingeLoss( y::Float64, t::Float64 )
    e = 1.0 - t * y;
    if e > 0.0
      return e;
    else
      return 0.0;
    end
  end

  """
      hingeLoss( y::Vector{Float64}, t::Vector{Float64} )

      compute hinge loss for a vector of binary classifications.
      #Arguments
        - `Vector{Float64}`: predicted values.
        - `Vector{Float64}`: ground truth values.
      #Return
        - `Vector{Float64}`: vector of hingeLoss errors.
  """
  function hingeLoss( y::Vector{Float64}, t::Vector{Float64} )
    return hingeLoss.( y, t );
  end

  """
      MSE( y::Float64, t::Float64 )

      compute mean squared error for a single binary classification.
        #Arguments
        - `Float64`: predicted value.
        - `Float64`: ground truth value.
      #Return
        - `Float64`: hingeLoss error.
  """
  function MSE( y::Float64, t::Float64 )
    return ( y - t )^2;
  end

  """
      MSE( Y::Vector{Float64}, T::Vector{Float64} )

      compute mean squared error for a vector of binary classifications.
        #Arguments
        - `Vector{Float64}`: predicted value.
        - `Vector{Float64}`: ground truth value.
      #Return
        - `Vector{Float64}`: vector of hingeLoss errors.
  """
  function MSE( Y::Vector{Float64}, T::Vector{Float64} )
    return MSE.( Y, T );
  end

    """
        zeroOneLoss( y::Float64, t::Float64 )

        compute zero-one loss for a vector for a binary classifications.
          #Arguments
          - `Float64`: predicted value.
          - `Float64`: ground truth value.
          #Return
          - `Float64`: zero-one loss error.

  """
  function zeroOneLoss( y::Float64, t::Float64 )
   if y ≈ t
     return 0.0;
   else
     return 1.0;
   end
  end

  """
        zeroOneLoss( y::Float64, t::Float64 )

        compute zero-one loss for a vector of binary classifications.
          #Arguments
          - `Vector{Float64}`: predicted values.
          - `Vector{Float64}`: ground truth values.
          #Return
          - `Vector{Float64}`: vector of zero-one loss errors.

  """
  function zeroOneLoss( Y::Vector{Float64}, T::Vector{Float64} )
    return zeroOneLoss.( Y, T );
  end
end