"""`care(A, B, Q, R)`

Compute 'X', the solution to the continuous-time algebraic Riccati equation,
defined as A'X + XA - (XB)R^-1(B'X) + Q = 0, where R is non-singular.

Algorithm taken from:
Laub, "A Schur Method for Solving Algebraic Riccati Equations."
http://dspace.mit.edu/bitstream/handle/1721.1/1301/R-0859-05666488.pdf
"""
function care(A, B, Q, R)
    G = try
        B*inv(R)*B'
    catch y
        if y isa SingularException
            error("R must be non-singular in care.")
        else
            throw(y)
        end
    end

    Z = [A  -G;
        -Q  -A']

    S = schur(Z)
    S = ordschur(S, real(S.values).<0)
    U = S.Z

    (m, n) = size(U)
    U11 = U[1:div(m, 2), 1:div(n,2)]
    U21 = U[div(m,2)+1:m, 1:div(n,2)]
    return U21/U11
end

"""`dare(A, B, Q, R)`

Compute `X`, the solution to the discrete-time algebraic Riccati equation,
defined as A'XA - X - (A'XB)(B'XB + R)^-1(B'XA) + Q = 0, where Q>=0 and R>0

Algorithm taken from:
Laub, "A Schur Method for Solving Algebraic Riccati Equations."
http://dspace.mit.edu/bitstream/handle/1721.1/1301/R-0859-05666488.pdf
"""
function dare(A, B, Q, R)
    if (!ishermitian(Q) || minimum(eigvals(real(Q))) < 0)
        error("Q must be positive-semidefinite.");
    end
    if (!isposdef(R))
        error("R must be positive definite.");
    end

    n = size(A, 1);

    E = [
        Matrix{Float64}(I, n, n) B/R*B';
        zeros(size(A)) A'
    ];
    F = [
        A zeros(size(A));
        -Q Matrix{Float64}(I, n, n)
    ];

    QZ = schur(F, E);
    QZ = ordschur(QZ, abs.(QZ.alpha./QZ.beta) .< 1);

    return QZ.Z[(n+1):end, 1:n]/QZ.Z[1:n, 1:n];
end

"""`dlyap(A, Q)`

Compute the solution `X` to the discrete Lyapunov equation
`AXA' - X + Q = 0`.
"""
function dlyap(A::AbstractMatrix, Q)
    lhs = kron(A, conj(A))
    lhs = I - lhs
    x = lhs\reshape(Q, prod(size(Q)), 1)
    return reshape(x, size(Q))
end

"""`gram(sys, opt)`

Compute the grammian of system `sys`. If `opt` is `:c`, computes the
controllability grammian. If `opt` is `:o`, computes the observability
grammian."""
function gram(sys::AbstractStateSpace, opt::Symbol)
    if !isstable(sys)
        error("gram only valid for stable A")
    end
    func = iscontinuous(sys) ? lyap : dlyap
    if opt == :c
        # TODO probably remove type check in julia 0.7.0
        return func(sys.A, sys.B*sys.B')#::Array{numeric_type(sys),2} # lyap is type-unstable
    elseif opt == :o
        return func(Matrix(sys.A'), sys.C'*sys.C)#::Array{numeric_type(sys),2} # lyap is type-unstable
    else
        error("opt must be either :c for controllability grammian, or :o for
                observability grammian")
    end
end

"""`obsv(A, C)` or `obsv(sys)`

Compute the observability matrix for the system described by `(A, C)` or `sys`.

Note that checking for observability by computing the rank from `obsv` is
not the most numerically accurate way, a better method is checking if
`gram(sys, :o)` is positive definite."""
function obsv(A::AbstractMatrix, C::AbstractMatrix)
    T = promote_type(eltype(A), eltype(C))
    n = size(A, 1)
    ny = size(C, 1)
    if n != size(C, 2)
        throw(ArgumentError("C must have the same number of columns as A"))
    end
    res = fill(zero(T), n*ny, n)
    res[1:ny, :] = C
    for i=1:n-1
        res[(1 + i*ny):(1 + i)*ny, :] = res[((i - 1)*ny + 1):i*ny, :] * A
    end
    return res
end
obsv(sys::StateSpace) = obsv(sys.A, sys.C)

"""`ctrb(A, B)` or `ctrb(sys)`

Compute the controllability matrix for the system described by `(A, B)` or
`sys`.

Note that checking for controllability by computing the rank from
`ctrb` is not the most numerically accurate way, a better method is
checking if `gram(sys, :c)` is positive definite."""
function ctrb(A::AbstractMatrix, B::AbstractMatrix)
    T = promote_type(eltype(A), eltype(B))
    n = size(A, 1)
    nu = size(B, 2)
    if n != size(B, 1)
        throw(ArgumentError("B must have the same number of rows as A"))
    end
    res = fill(zero(T), n, n*nu)
    res[:, 1:nu] = B
    for i=1:n-1
        res[:, (1 + i*nu):(1 + i)*nu] = A * res[:, ((i - 1)*nu + 1):i*nu]
    end
    return res
end
ctrb(sys::StateSpace) = ctrb(sys.A, sys.B)

"""`P = covar(sys, W)`

Calculate the stationary covariance `P = E[y(t)y(t)']` of an lti-model `sys`, driven by gaussian
white noise 'w' of covariance `E[w(t)w(τ)]=W*δ(t-τ)` where δ is the dirac delta.

The ouput is if Inf if the system is unstable. Passing white noise directly to
the output will result in infinite covariance in the corresponding outputs
(D*W*D' .!= 0) for contunuous systems."""
function covar(sys::AbstractStateSpace, W)
    (A, B, C, D) = (sys.A, sys.B, sys.C, sys.D)
    if !isa(W, UniformScaling) && (size(B,2) != size(W, 1) || size(W, 1) != size(W, 2))
        error("W must be a square matrix the same size as `sys.B` columns")
    end
    if !isstable(sys)
        return fill(Inf,(size(C,1),size(C,1)))
    end
    func = iscontinuous(sys) ? lyap : dlyap
    Q = try
        func(A, B*W*B')
    catch
        error("No solution to the Lyapunov equation was found in covar")
    end
    P = C*Q*C'
    if iscontinuous(sys)
        #Variance and covariance infinite for direct terms
        directNoise = D*W*D'
        for i in 1:size(C,1)
            if directNoise[i,i] != 0
                P[i,:] .= Inf
                P[:,i] .= Inf
            end
        end
    else
        P += D*W*D'
    end
    return P
end

covar(sys::TransferFunction, W) = covar(ss(sys), W)

"""
    covar(C,W)
If `C` is a matrix, return CWC'
"""
covar(C::Union{AbstractMatrix,UniformScaling}, R) = C*R*C'

# Note: the H∞ norm computation is probably not as accurate as with SLICOT,
# but this seems to be still reasonably ok as a first step
"""
`..  norm(sys, p=2; tol=1e-6)`

`norm(sys)` or `norm(sys,2)` computes the H2 norm of the LTI system `sys`.

`norm(sys, Inf)` computes the L∞ norm of the LTI system `sys`.
The H∞ norm is the same as the L∞ for stable systems, and Inf for unstable systems.
If the peak gain frequency is required as well, use the function `norminf` instead.

`tol` is an optional keyword argument, used only for the computation of L∞ norms.
It represents the desired relative accuracy for the computed L∞ norm
(this is not an absolute certificate however).

sys is first converted to a state space model if needed.

The L∞ norm computation implements the 'two-step algorithm' in:
N.A. Bruinsma and M. Steinbuch, 'A fast algorithm to compute the H∞-norm
of a transfer function matrix', Systems and Control Letters 14 (1990), pp. 287-293.
For the discrete-time version, see, e.g.,: P. Bongers, O. Bosgra, M. Steinbuch, 'L∞-norm
calculation for generalized state space systems in continuous and discrete time',
American Control Conference, 1991.
"""
function LinearAlgebra.norm(sys::AbstractStateSpace, p::Real=2; tol=1e-6)
    if p == 2
        return sqrt(tr(covar(sys, I)))
    elseif p == Inf
        if sys.Ts == 0
            return normLinf_twoSteps_ct(sys,tol)[1]
        else
            return normLinf_twoSteps_dt(sys,tol)[1]
        end
    else
        error("`p` must be either `2` or `Inf`")
    end
end

function LinearAlgebra.norm(sys::TransferFunction, p::Real=2; tol=1e-6)
    return norm(ss(sys), p, tol=tol)
end

"""
`.. (peakgain, peakgainfrequency) = norminf(sys; tol=1e-6)`

Compute the L∞ norm of the LTI system `sys`, together with the frequency
`peakgainfrequency` (in rad/TimeUnit) at which the gain achieves its peak value `peakgain`.
The H∞ norm is the same as the L∞ for stable systems, and Inf for unstable systems.

`tol` is an optional keyword argument representing the desired relative accuracy for
the computed L∞ norm (this is not an absolute certificate however).

sys is first converted to a state space model if needed.

The L∞ norm computation implements the 'two-step algorithm' in:
N.A. Bruinsma and M. Steinbuch, 'A fast algorithm to compute the H∞-norm
of a transfer function matrix', Systems and Control Letters 14 (1990), pp. 287-293.
For the discrete-time version, see, e.g.,: P. Bongers, O. Bosgra, M. Steinbuch, 'L∞-norm
calculation for generalized state space systems in continuous and discrete time',
American Control Conference, 1991.
"""
function norminf(sys::AbstractStateSpace; tol=1e-6)
    if sys.Ts == 0
        return normLinf_twoSteps_ct(sys,tol)
    else
        return normLinf_twoSteps_dt(sys,tol)
    end
end

function norminf(sys::TransferFunction, ; tol=1e-6)
    return norminf(ss(sys), tol=tol)
end

function normLinf_twoSteps_ct(sys::AbstractStateSpace, tol=1e-6, maxIters=1000, approximag=1e-10)
    # `maxIters`: the maximum  number of iterations allowed in the algorithm (default 1000)
    # approximag is a tuning parameter: what does it mean for a number to be on the imaginary axis
    # Because of this tuning for example, the relative precision that we provide on the norm computation
    # is not a true guarantee, more an order of magnitude
    # outputs: pair of Float64, namely L∞ norm approximation and frequency fpeak at which it is achieved
    T = promote_type(numeric_type(sys), Float64)
    if sys.nx == 0  # static gain
        return (svdvals(sys.D)[1], T(0))
    end
    p = pole(sys)
    # Check if there is a pole on the imaginary axis
    pidx = findfirst(map(x->isapprox(x,0.0),real(p)))
    if !(pidx isa Nothing)
        return (T(Inf), imag(p[pidx]))
        # note: in case of cancellation, for s/s for example, we return Inf, whereas Matlab returns 1
    else
        # Initialization: computation of a lower bound from 3 terms
        lb = maximum(svdvals(sys.D))
        fpeak = T(Inf)
        (lb, idx) = findmax([lb, T(maximum((svdvals(evalfr(sys,0)))))]) #TODO remove T() in julia 0.7.0
        if idx == 2
            fpeak = T(0)
        end
        if isreal(p)  # only real poles
            omegap = minimum(abs.(p))
        else  # at least one pair of complex poles
            tmp = maximum(abs.(imag.(p)./(real.(p).*abs.(p))))
            omegap = abs(p[argmax(tmp)])    # TODO This is highly suspicious
        end
        (lb, idx) = findmax([lb, T(maximum(svdvals(evalfr(sys, omegap*1im))))]) #TODO remove T() in julia 0.7.0
        if idx == 2
            fpeak = omegap
        end

        # Iterations
        iter = 1;
        while iter <= maxIters
            res = (1+2*T(tol))*lb
            R = sys.D'*sys.D - res^2*I
            S = sys.D*sys.D' - res^2*I
            M = sys.A-sys.B*(R\sys.D')*sys.C
            H = [         M              -res*sys.B*(R\sys.B') ;
                   res*sys.C'*(S\sys.C)            -M'            ]
            omegas = eigvals(H) .+ 0im # To make type stable
            omegaps = imag.(omegas[ (abs.(real.(omegas)).<=approximag) .& (imag.(omegas).>=0) ])
            sort!(omegaps)
            if isempty(omegaps)
                return (1+T(tol))*lb, fpeak
            else  # if not empty, omegaps contains at least two values
                ms = [(x+y)/2 for x=omegaps[1:end-1], y=omegaps[2:end]]
                for mval in ms
                    (lb, idx) = findmax([lb, T(maximum(svdvals(evalfr(sys,mval*1im))))]) #TODO remove T() in julia 0.7.0
                    if idx == 2
                        fpeak = mval
                    end
                end
            end
            iter += 1
        end
        error("In norminf: The computation of the H-infinity norm did not converge in $maxIters iterations")
    end
end

# discrete-time version of normHinf_twoSteps_ct above
# The value fpeak returned by the function is in the range [0,pi)/sys.Ts (in rad/s)
function normLinf_twoSteps_dt(sys::AbstractStateSpace,tol=1e-6, maxIters=1000, approxcirc=1e-8)
    T = promote_type(numeric_type(sys), Float64)
    if sys.nx == 0  # static gain
        return (svdvals(sys.D)[1], T(0))
    end
    p = pole(sys)
    # Check first if there is a pole on the unit circle
    pidx = findfirst(map(x->isapprox(x,1.0),abs.(p)))
    if !(pidx isa Nothing)
        return (T(Inf), angle(p[pidx])/abs(T(sys.Ts)))
    else
        # Initialization: computation of a lower bound from 3 terms
        lb = T(maximum(svdvals(evalfr(sys,1))))  #TODO remove T() in julia 0.7.0
        fpeak = T(0)
        (lb, idx) = findmax([lb, T(maximum(svdvals(evalfr(sys,-1))))]) #TODO remove T() in julia 0.7.0
        if idx == 2
            fpeak = T(pi)
        end

        p = p[imag(p).>0]
        if ~isempty(p)  # not just real poles
            # find frequency of pôle closest to unit circle
            omegap = angle(p[findmin(abs.(abs.(p).-1))[2]])
        else
            omegap = T(pi)/2
        end
        (lb, idx) = findmax([lb, T(maximum(svdvals(evalfr(sys, exp(omegap*1im)))))]) #TODO remove T() in julia 0.7.0
        if idx == 2
            fpeak = omegap
        end

        # Iterations
        iter = 1;
        while iter <= maxIters
            res = (1+2*T(tol))*lb
            R = res^2*I - sys.D'*sys.D
            RinvDt = R\sys.D'
            L = [ sys.A+sys.B*RinvDt*sys.C  sys.B*(R\sys.B');
                  zeros(T, sys.nx,sys.nx)      I]
            M = [ I                                 zeros(T, sys.nx,sys.nx);
                  sys.C'*(I+sys.D*RinvDt)*sys.C     L[1:sys.nx,1:sys.nx]']
            # +0im to make type stable
            zs = eigvals(L,M) .+ 0im # generalized eigenvalues
            # are there eigenvalues on the unit circle?
            omegaps = angle.(zs[ (abs.(abs.(zs).-1) .<= approxcirc) .& (imag(zs).>=0)])
            sort!(omegaps)
            if isempty(omegaps)
                return (1+T(tol))*lb, fpeak/T(sys.Ts)
            else  # if not empty, omegaps contains at least two values
                ms = [(x+y)/2 for x=omegaps[1:end-1], y=omegaps[2:end]]
                for mval in ms
                    (lb, idx) = findmax([lb, T(maximum(svdvals(evalfr(sys,exp(mval*1im)))))]) #TODO remove T() in julia 0.7.0
                    if idx == 2
                        fpeak = mval
                    end
                end
            end
            iter += 1
        end
        error("In norminf: The computation of the H-infinity norm did not converge in $maxIters iterations")
    end
end

"""`S, P, B = balance(A[, perm=true])`

Compute a similarity transform `T` resulting in `B = T\\A*T` such that the row
and column norms of `B` are approximately equivalent. If `perm=false`, the
transformation will only scale `A` using diagonal `S`, and not permute `A` (i.e., set `P=I`)."""
function balance(A, perm::Bool=true)
    n = LinearAlgebra.checksquare(A)
    B = copy(A)
    job = perm ? 'B' : 'S'
    ilo, ihi, scaling = LAPACK.gebal!(job, B)

    S = diagm(0 => scaling)
    for j = 1:(ilo-1)   S[j,j] = 1 end
    for j = (ihi+1):n   S[j,j] = 1 end

    P = Matrix{Int}(I,n,n)
    if perm
        if ilo > 1
            for j = (ilo-1):-1:1 cswap!(j, round(Int, scaling[j]), P) end
        end
        if ihi < n
            for j = (ihi+1):n    cswap!(j, round(Int, scaling[j]), P) end
        end
    end
    return S, P, B
end

function cswap!(i::Integer, j::Integer, X::StridedMatrix)
    for k = 1:size(X,1)
        X[i, k], X[j, k] = X[j, k], X[i, k]
    end
end

"""
`sysr, G = balreal(sys::StateSpace)`

Calculates a balanced realization of the system sys, such that the observability and reachability gramians of the balanced system are equal and diagonal `G`

See also `gram`, `baltrunc`

Glad, Ljung, Reglerteori: Flervariabla och Olinjära metoder
"""
function balreal(sys::ST) where ST <: AbstractStateSpace
    P = gram(sys, :c)
    Q = gram(sys, :o)

    Q1 = try
        cholesky(Hermitian(Q)).U
    catch
        throw(ArgumentError("Balanced realization failed: Observability grammian not positive definite, system needs to be observable"))
    end
    U,Σ,V = svd(Q1*P*Q1')
    Σ .= sqrt.(Σ)
    Σ1 = diagm(0 => sqrt.(Σ))
    T = Σ1\(U'Q1)

    Pz = T*P*T'
    Qz = inv(T')*Q*inv(T)
    if norm(Pz-Qz) > sqrt(eps())
        @warn("balreal: Result may be inaccurate")
        println("Controllability gramian before transform")
        display(P)
        println("Controllability gramian after transform")
        display(Pz)
        println("Observability gramian before transform")
        display(Q)
        println("Observability gramian after transform")
        display(Qz)
        println("Singular values of PQ")
        display(Σ)
    end

    sysr = ST(T*sys.A/T, T*sys.B, sys.C/T, sys.D, sys.Ts), diagm(0 => Σ)
end

"""
`sysr, G = baltrunc(sys::StateSpace, atol = √ϵ, rtol=1e-3, unitgain=true)`

Reduces the state dimension by calculating a balanced realization of the system sys, such that the observability and reachability gramians of the balanced system are equal and diagonal `G`, and truncating it such that all states corresponding to singular values less than `atol` and less that `rtol σmax` are removed. If `unitgain=true`, the matrix `D` is chosen such that unit static gain is achieved.

See also `gram`, `balreal`

Glad, Ljung, Reglerteori: Flervariabla och Olinjära metoder
"""
function baltrunc(sys::ST; atol = sqrt(eps()), rtol = 1e-3, unitgain = true) where ST <: AbstractStateSpace
    sysbal, S = balreal(sys)
    S = diag(S)
    S = S[S .>= atol]
    S = S[S .>= S[1]*rtol]
    n = length(S)
    A = sysbal.A[1:n,1:n]
    B = sysbal.B[1:n,:]
    C = sysbal.C[:,1:n]
    D = sysbal.D
    if unitgain
        D = D/(C*inv(-A)*B)
    end

    return ST(A,B,C,D,sys.Ts), diagm(0 => S)
end

"""
    syst = similarity_transform(sys, T)
Perform a similarity transform `T : Tx̃ = x` on `sys` such that
```
Ã = T⁻¹AT
B̃ = T⁻¹ B
C̃ = CT
D̃ = D
```
"""
function similarity_transform(sys::ST, T) where ST <: AbstractStateSpace
    Tf = factorize(T)
    A = Tf\sys.A*T
    B = Tf\sys.B
    C = sys.C*T
    D = sys.D
    ST(A,B,C,D,sys.Ts)
end

"""
sysi = innovation_form(sys, R1, R2)
sysi = innovation_form(sys; sysw=I, syse=I, R1=I, R2=I)

Takes a system
```
x' = Ax + Bu + w ~ R1
y  = Cx + e ~ R2
```
and returns the system
```
x' = Ax + Kv
y  = Cx + v
```
where `v` is the innovation sequence.

If `sysw` (`syse`) is given, the covariance resulting in filtering noise with `R1` (`R2`) through `sysw` (`syse`) is used as covariance.

See Stochastic Control, Chapter 4, Åström
"""
function innovation_form(sys::ST, R1, R2) where ST <: AbstractStateSpace
    K = kalman(sys, R1, R2)
    ST(sys.A, K, sys.C, Matrix{eltype(sys.A)}(I, sys.ny, sys.ny), sys.Ts)
end
# Set D = I to get transfer function H = I + C(sI-A)\ K
function innovation_form(sys::ST; sysw=I, syse=I, R1=I, R2=I) where ST <: AbstractStateSpace
	K = kalman(sys, covar(sysw,R1), covar(syse, R2))
	ST(sys.A, K, sys.C, Matrix{eltype(sys.A)}(I, sys.ny, sys.ny), sys.Ts)
end


using DifferentialEquations
using Plots

function seir_ode(dY,Y,p,t)
 dY[1] = p[4]-p[1]*Y[1]*Y[3]-p[4]*Y[1]
 dY[2] = p[1]*Y[1]*Y[3]-(p[2]+p[4])*Y[2]
 dY[3] = p[2]*Y[2] - (p[3]+p[4])*Y[3]
end

par=[520/365,1/60,1/30,774835/(65640000*365)]
init=[0.8,0.1,0.1]
tspan=(0.0,365.0)

seir_prob = ODEProblem(seir_ode,init,tspan,par)

sol=solve(seir_prob);

# Plot
R=ones(1,size(sol,2))-sum(sol,1);

plot(sol.t,[sol',R'],xlabel="Time",ylabel="Proportion")


# Represent as a permtuation matrix.
"""
    Matrix(p::Permutation)

Return the permutation matrix for `p`.
"""
function Matrix{T}(p::Permutation) where {T}
    n = length(p)
    A = Matrix{T}(I, n, n)     # A = eye(T,n)   #  int(eye(n))
    return A[:, p.data]
end

Matrix(p::Permutation) = Matrix{Int}(p)
Array(p::Permutation) = Matrix(p)
AbstractMatrix(p::Permutation) = Matrix(p)
AbstractArray(p::Permutation) = Matrix(p)

Array{T}(p::Permutation) where {T} = Matrix{T}(p)
AbstractMatrix{T}(p::Permutation) where {T} = Matrix{T}(p)
AbstractArray{T}(p::Permutation) where {T} = Matrix{T}(p)

function Permutation(M::Matrix)
    n, c = size(M)
    @assert n == c "Matrix must be square"
    v = [findfirst(Bool.(M[:, k])) for k = 1:n]
    return Permutation(v)
end


"""
    owens_t(x, a)

Owen's T function

Ported from Matlab implementation of John Burkardt, which is distributed under the GNU LGPL license.

Reference:
    MA Porter, DJ Winstanley, Remark AS R30: A Remark on Algorithm AS76: An Integral Useful in
    Calculating Noncentral T and Bivariate Normal Probabilities, Applied Statistics, Volume 28,
    Number 1, 1979, page 113.

    JC Young, Christoph Minder, Algorithm AS 76: An Algorithm Useful in Calculating Non-Central T and
    Bivariate Normal Distributions, Applied Statistics, Volume 23, Number 3, 1974, pages 455-457.
"""
function owens_t(x::Real, a::Real)
    ng = 5
    r = [0.1477621, 0.1346334, 0.1095432, 0.0747257, 0.0333357]
    tp = 0.159155
    tv1 = 1.0E-35
    tv2 = 15.0
    tv3 = 15.0
    tv4 = 1.0E-05
    u = [0.0744372, 0.2166977, 0.3397048, 0.4325317, 0.4869533]
    # Test for X near zero.
    if abs(x) < tv1
        return tp * atan(a)
    end
    # Test for large values of abs(X).
    if tv2 < abs(x)
        return 0.
    end
    # Test for `a` near zero.
    if abs(a) < tv1
        return 0.
    end
    # Test whether abs(`a`) is so large that it must be truncated.
    xs = -0.5 * x * x
    x2 = a
    fxs = a * a
    #  Computation of truncation point by Newton iteration.
    if tv3 <= log(1.0 + fxs) - xs * fxs
        x1 = 0.5 * a
        fxs = 0.25 * fxs
        while true
            rt = fxs + 1.
            x2 = x1 + (xs * fxs + tv3 - log(rt)) / (2. * x1 * ( 1. / rt - xs))
            fxs = x2 * x2
            if abs(x2 - x1) < tv4
                break
            end
            x1 = x2
        end
    end
    # Gaussian quadrature.
    rt = 0.0
    for i in 1:ng
        r1 = 1.0 + fxs * (0.5 + u[i])^2
        r2 = 1.0 + fxs * (0.5 - u[i])^2
        rt = rt + r[i] * (exp(xs * r1) / r1 + exp(xs * r2) / r2)
    end
    return rt * x2 * tp
end


# functions to create regular polygons

export equilateral, equiangular

"""
`_solver(f,goal,lo,hi)` assumes `f` is increasing on the interval `[lo,hi]`
"""
function _solver(f::Function, goal::Real, lo::Real, hi::Real)
    mid = (hi+lo)/2
    y = f(mid)

    if abs(y-goal) <= THRESHOLD*eps(1.0)
        return mid
    end
    if y > goal
        return _solver(f,goal,lo,mid)
    end
    return _solver(f,goal,mid,hi)
end

"""
`_solver(f,goal)` assumes that `f` is increasing and there is a nonnegative `x`
so that `f(x)==goal`.
"""
function _solver(f::Function, goal::Real)
    x::Float64 = 1.0
    while f(x) < goal
        x *= 2.0
    end
    _solver(f,goal,0,x)
end

"""
`equilateral(n,s)` creates a regular `n`-gon with side lengths `s`
centered at the origin. First point is on the positive x-axis.
"""
function equilateral(n::Int, s::Real)
    @assert n>2 "Number of sides must be at least 3"
    @assert s>0 "Side length must be positive"

    tlist = [ k*2*pi/n for k=0:n-1 ]   # list of angles
    t = tlist[2]

    f(x)  = dist(HPoint(x,0), HPoint(x,t))
    r = _solver(f,s)

    pts = [ HPoint(r,theta) for theta in tlist]

    return HPolygon(pts)
end

"""
`equiangular(n,theta)` creates a regular `n`-gon where the vertex
angles equal `theta`
"""
function equiangular(n::Int, theta::Real)
    @assert n>2 "Number of sides must be at least 3"
    if theta < 0
        theta = -theta
    end
    @assert theta!=0 "Angle cannot be zero"

    tmax = (n-2)*pi/n
    @assert theta<tmax "Angle must be less than $(n-2)*pi/$n = $tmax"

    tlist = [ k*2*pi/n for k=0:n-1 ]   # list of angles
    t = tlist[2]

    f(x) = -angle(HPoint(x,0), HPoint(x,t), HPoint(x,2t))
    r = _solver(f,-theta)

    pts = [ HPoint(r,theta) for theta in tlist]

    return HPolygon(pts)

end


"""
    ifft_times(fs)

    Returns an array of DFT time points associated with the set of
    frequency-domain points.

    Assumes unit of argument `x` is Hz (not rad/sec). Output unit is in s.
"""
function ifft_times(fs)
    deltaf = fs[2]-fs[1]
    N = Int(length(fs))
    if N%2 == 0
        delta_t = 1/(-2*fs[1])
    else
        delta_t = 1/(2*fs[end]+deltaf)
    end
    ts = (0:N-1)*delta_t
    return ts
end


export calc_k, calc_λ, calc_κ
export calc_NA
export fresnel_number

"""
    calc_NA(focal_length, diameter[, n])

Calculate the numerical aperture of a system with `focal_length`
and `diameter` of the lens. Per default `n=1` is the refractive index.

 # Examples
```julia-repl
julia> calc_NA(100e-3, 200e-3)
1.0

julia> calc_NA(100e-3, 200e-3, 1.33)
1.33
```
"""
function calc_NA(focal_length, diameter, n=1)
    return n * diameter / 2 / focal_length
end

"""
    calc_k(λ, n=1)

Calculate the value of the (angular) wave number with vacuum wavelength `λ` in
medium with refractive index `n`.
It holds: \$ k = \\kappa \\cdot 2 \\pi \$
"""
function calc_k(λ, n=1)
    return 2 * oftype(λ, π) * n / λ
end

"""
    calc_κ(λ, n=1)

Calculate the value of the (non angular) wave number with vacuum wavelength `λ` in
medium with refractive index `n`.
It holds: \$ \\kappa = \\frac{k}{2 \\pi} \$
"""
function calc_κ(λ, n=1)
    return calc_k(λ, n) / 2 / oftype(λ, π)
end

"""
    calc_λ(k, n=1)

Calculate the vacuum wavelength `λ` from the angular wave number `k` in
medium with refractive index `n`.
It holds: \$ \\lambda = \\frac{n \\cdot 2 \\pi}{k} \$
"""
function calc_λ(k, n=1)
    return 2 * oftype(k, π) * n / k
end

"""
    fresnel_number(a, L, λ=λ0)

Calculate the Fresnel number where `a` is the characteristic size,
`L` the distance from screen to aperture and `λ` the wavelength.
"""
function fresnel_number(a, L, λ=λ0)
    return a^2 / L / λ
end


function higibbs(Y, T, mu0 = 50.0, gamma20 = 25.0, nu0 = 1.0, sigma2 = 100.0, eta0 = 1.0, tau20 = 100.0)
    m = size(unique(Y[:,1]), 1)
    # starting value
    ybar = ones(m)
    sv = ones(m)
    n = ones(m)
    for j = 1:m
        yj = Y[ [Y[i,1] == j for i = 1:end], 2]
        ybar[j] = mean(yj)
        sv[j] = var(yj)
        n[j] = size(yj, 1)
    end
    theta = ybar
    sigma2 = mean(sv)
    mu = mean(theta)
    tau2 = var(theta)

    THETA = ones(T, m)
    # sigma2 mu tau2
    SMT = ones(T, 3)

    for t = 1:T
        # sample theta
        for j = 1:m
            vartheta = 1 / (n[j] / sigma2 + 1 / tau2)
            meantheta = vartheta * ( n[j] * ybar[j]) / sigma2 + mu / tau2)
            rnorm = Normal(meantheta, sqrt(vartheta))
            theta[j] = rand(rnorm, 1)[1]
        end
        THETA[t, :] .= theta

        # sample new sigma2
        nun = nu0 + sum(n)
        ss = nu0 * s20
        for j = 1:m
            ss += sum((Y[ [Y[i, 1] == j for i = 1:end], 2] .- theta[j]).^2)
        end
        rgamma = Gamma(nun/2, 2/ss)
        sigma2 = 1 / rand(rgamma, 1)[1]

        # sample mu
        varmu = 1 / (m / tau2 + 1 / gamma20)
        meanmu = varmu * (m * mean(theta) / tau2 + mu0 / gamma20)
        rnorm = Normal(meanmu, sqrt(varmu))
        mu = rand(rnorm, 1)[1]

        # sample 1/tau2
        shapetau = (eta0 + m) / 2
        ratetau = ( eta0 * tau20 + sum((theta .- mu).^2) ) / 2
        rgamma = Gamma(shapetau, 1/ratetau)
        tau2 = 1 / rand(rgamma, 1)[1]

        # store results
        SMT[t, :] .= [sigma2, mu, tau2]
    end
    return THETA, SMT
end

# run
Y = readdlm("math-score-Y.csv")
THETA, SMT = higibbs(Y, 5000)


function FPU(omega)

# returns functions for the fast potential, the slow potential, kineticTerm, Lagragian, Hamiltonian and Oscillatory Energy
# of the FPUT problem, when provided with the stiffness parameter omega

# FPUT problem: see I.5.1 A Fermi–Pasta–Ulam Problem in Geometric Numerical Integration by Hairer, Lubich, Wanner (2006)

    function fastPotential(qf)
    omega^2/2*sum(qf.^2)
    end

    function slowPotential(qs,qf)
        qd= qs-qf
        1/4*(qd[1]^4 + sum((qd[2:end]-qs[1:end-1]-qf[1:end-1]).^4) + qd[end]^4)
    end

    function kineticTerm(qsdot,qfdot)
        1/2*(sum(qsdot.^2)+sum(qfdot.^2))
    end

    function L(qs,qf,qsdot,qfdot)
        kineticTerm(qsdot,qfdot)-slowPotential(qs,qf)-fastPotential(qf)
    end

    function Hamiltonian(qs,qf,qsdot,qfdot)
        kineticTerm(qsdot,qfdot)+slowPotential(qs,qf)+fastPotential(qf)
    end

    function OscillatoryEnergy(qf,pf)
    # adiabatic invariant
        return sum(1/2*(pf.^2 .+omega^2*qf.^2),dims=1)
    end

    return fastPotential, slowPotential, kineticTerm, L, Hamiltonian, OscillatoryEnergy

end


using ModelingToolkit
using DiffEqJump
using Random
using DataFrames
using Tables
using StatsPlots
using BenchmarkTools

@parameters β c γ
@variables t S(t) I(t) R(t)
N = S+I+R
rate₁   = β*c*S*I/N
affect₁ = [S ~ S - 1, I ~ I + 1]
rate₂   = γ*I
affect₂ = [I ~ I - 1, R ~ R + 1]
j₁      = ConstantRateJump(rate₁,affect₁)
j₂      = ConstantRateJump(rate₂,affect₂)
@named sir_js = JumpSystem([j₁,j₂], t, [S,I,R], [β,c,γ])

tmax = 40.0
tspan = (0.0,tmax)
δt = 1.0;

u0 = [S => 990, I => 10, R => 0];

p = [β => 0.05, c => 10.0, γ => 0.25];

Random.seed!(1234);

sir_dprob = DiscreteProblem(sir_js, u0, tspan,p);

sir_jprob = JumpProblem(sir_js, sir_dprob,Direct());

sol_jump = solve(sir_jprob,SSAStepper());

out_jump = sol_jump(0:δt:tmax);

df_jump = DataFrame(Tables.table(out_jump'))
rename!(df_jump,["S","I","R"])
df_jump[!,:t] = out_jump.t;

@df df_jump plot(:t,
    [:S :I :R],
    xlabel="Time",
    ylabel="Number")

@benchmark solve(sir_jprob,SSAStepper())


using CompEcon
using Dierckx
using Statistics

# Create functions to test against
f1(x) = sin.(x)
fp1(x) = cos.(x)
f2(x) = exp.(x)
fp2(x) = exp.(x)
f3(x) = 1.5 .+ log.(x)
fp3(x) = 1.0 ./ x
fs = [f1, f2, f3]
fps = [fp1, fp2, fp3]
nf = length(fs)

x = collect(range(.5; stop=2, length=35))
finex = collect(range(1e-2; stop=2*pi - 1e-2, length=150))
ys = map(f->f(x), fs)
finey = map(f->f(finex), fs)

function compare_D_CE_levels(f::Function, x::Array{Float64, 1}, order::Int)

    # Create Dierckx Spline (and data)
    n = length(x)
    y = f(x)
    dspl = Spline1D(x, y; s=0., k=order)

    # Create CompEcon Spline (and data)
    cebasis = Basis(SplineParams(x, 0, order))
    xx = nodes(cebasis)[1]
    yy = f(xx)
    c, bs = funfitxy(cebasis, xx, yy)

    # Evaluate Splines on finer grid
    xfine = collect(range(x[1]; stop=x[end], length=2*n + 1))
    yfine = f(xfine)

    deval = Dierckx.evaluate(dspl, xfine)
    derr = abs.(yfine - deval)

    ceeval = CompEcon.funeval(c, cebasis, xfine)
    ceerr = abs.(yfine - ceeval)

    # # Evaluate Derivative
    # dderiv1 = Dierckx.derivative(dspl1, finex)
    # cederiv1 = CompEcon.evaluate(ceinterp1, finex, order=1)

    return derr, ceerr
end

function compare_D_CE_deriv(f::Function, fp::Function, x::Array{Float64, 1}, order::Int)

    # Create Dierckx Spline (and data)
    n = length(x)
    y = f(x)
    dspl = Spline1D(x, y; s=0., k=order)

    # Create CompEcon Spline (and data)
    cebasis = Basis(SplineParams(x, 0, order))
    xx = nodes(cebasis)[1]
    yy = f(xx)
    c, bs = funfitxy(cebasis, xx, yy)

    # Evaluate Splines on finer grid
    xfine = collect(range(x[1]; stop=x[end], length=2*n + 1))
    ypfine = fp(xfine)

    devalderiv = Dierckx.derivative(dspl, xfine)
    derr = abs.(ypfine - devalderiv)

    ceevalderiv = funeval(c, cebasis, xfine, ones(Int64, 1,1))
    ceerr = abs.(ypfine .- ceevalderiv)

    return derr, ceerr
end

for i=1:nf

    curr_f, curr_fp = fs[i], fps[i]

    derr1_levels, ceerr1_levels = compare_D_CE_levels(curr_f, x, 1)
    derr2_levels, ceerr2_levels = compare_D_CE_levels(curr_f, x, 2)
    derr3_levels, ceerr3_levels = compare_D_CE_levels(curr_f, x, 3)

    # derr1_deriv, ceerr1_deriv = compare_D_CE_deriv(curr_f, curr_fp, x, 1)
    derr2_deriv, ceerr2_deriv = compare_D_CE_deriv(curr_f, curr_fp, x, 2)
    derr3_deriv, ceerr3_deriv = compare_D_CE_deriv(curr_f, curr_fp, x, 3)

    println("\nLinear Approximation to Function\n")
    println("Dierckx Level Errors \n\t Max: $(maximum(derr1_levels)) \n\t Min: $(minimum(derr1_levels)) \n\t Mean: $(mean(derr1_levels))")
    # println("Dierckx Deriv Errors \n\t Max: $(maximum(derr1_deriv)) \n\t Min: $(minimum(derr1_deriv)) \n\t Mean: $(mean(derr1_deriv))")
    println("CompEcon Level Errors \n\t Max: $(maximum(ceerr1_levels)) \n\t Min: $(minimum(ceerr1_levels)) \n\t Mean: $(mean(ceerr1_levels))")
    # println("CompEcon Deriv Errors \n\t Max: $(maximum(ceerr1_deriv)) \n\t Min: $(minimum(ceerr1_deriv)) \n\t Mean: $(mean(ceerr1_deriv))")

    println("\n\nQuadratic Approximation to Function\n")
    println("Dierckx Level Errors \n\t Max: $(maximum(derr2_levels)) \n\t Min: $(minimum(derr2_levels)) \n\t Mean: $(mean(derr2_levels))")
    println("Dierckx Deriv Errors \n\t Max: $(maximum(derr2_deriv)) \n\t Min: $(minimum(derr2_deriv)) \n\t Mean: $(mean(derr2_deriv))")
    println("CompEcon Level Errors \n\t Max: $(maximum(ceerr2_levels)) \n\t Min: $(minimum(ceerr2_levels)) \n\t Mean: $(mean(ceerr2_levels))")
    println("CompEcon Deriv Errors \n\t Max: $(maximum(ceerr2_deriv)) \n\t Min: $(minimum(ceerr2_deriv)) \n\t Mean: $(mean(ceerr2_deriv))")

    println("\n\nCubic Approximation to Function\n")
    println("Dierckx Level Errors \n\t Max: $(maximum(derr3_levels)) \n\t Min: $(minimum(derr3_levels)) \n\t Mean: $(mean(derr3_levels))")
    println("Dierckx Deriv Errors \n\t Max: $(maximum(derr3_deriv)) \n\t Min: $(minimum(derr3_deriv)) \n\t Mean: $(mean(derr3_deriv))")
    println("CompEcon Level Errors \n\t Max: $(maximum(ceerr3_levels)) \n\t Min: $(minimum(ceerr3_levels)) \n\t Mean: $(mean(ceerr3_levels))")
    println("CompEcon Deriv Errors \n\t Max: $(maximum(ceerr3_deriv)) \n\t Min: $(minimum(ceerr3_deriv)) \n\t Mean: $(mean(ceerr3_deriv))")

end


function linear_interpolation(x1, xT, T)
    n = length(x1)
    X = [copy(Array(x1)) for t = 1:T]
    for t = 1:T
        for i = 1:n
            X[t][i] = (xT[i] - x1[i]) / (T - 1) * (t - 1) + x1[i]
        end
    end
    return X
end


# Each new term in the Fibonacci sequence is generated by adding the previous
# two terms. By starting with 1 and 2, the first 10 terms will be:
#
# 1, 2, 3, 5, 8, 13, 21, 34, 55, 89, ...
#
# By considering the terms in the Fibonacci sequence whose values do not exceed
# four million, find the sum of the even-valued terms.

using ProjectEulerSolutions

# Brute force solution, cycle through Fibonacci numbers up to 4 mil and sum the
# even ones.
function p002solution_running_sum(n::Integer=3)::Integer
    fibs = fibonacci_bound(BigInt(n))
    running_sum = 0
    @simd for val in fibs
        if val % 2 == 0
            running_sum += val
        end
    end
    return running_sum
end

# With a little thought, we see that it is only the 3*k+2 Fibonacci numbers that
# are even, so just sum them.
function p002solution_sum(n::Integer=3)::Integer
    fibs = fibonacci_bound(BigInt(n))
    return sum(fibs[2:3:length(fibs)])
end

p002 = Problems.Problem(Dict("Running sum" => p002solution_running_sum, "Sum" => p002solution_sum))

Problems.benchmark(p002, 4_000_000 - 1)


# Worked with Kasra and Waleed #

cd("D:\\Semester 3_Fall 2020\\ECON 6453 - Econometrics III\\ProblemSets\\PS4-mixture")
using Random, Distributions, JLD2, DataFrames, CSV, LinearAlgebra, Statistics, FreqTables, TexTables, Optim, HTTP, GLM, ForwardDiff

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# Question 1: Multinomial logit with Alternate Specific Covariates
#:::::::::::::::::::::::::::::::::::::::::::::::::::

url = "https://raw.githubusercontent.com/OU-PhD-Econometrics/fall-2020/master/ProblemSets/PS4-mixture/nlsw88t.csv"
df = CSV.read(HTTP.get(url).body)

X = [df.age df.white df.collgrad]
Z = hcat(df.elnwage1, df.elnwage2, df.elnwage3, df.elnwage4, df.elnwage5, df.elnwage6, df.elnwage7, df.elnwage8)
y = df.occ_code
# Using PS3's code for Multinomial Logit #
function mlogit(alpha, X, Z, y)
    beta = alpha[1:end-1] # Indexing the parameter so coefficient on X is the first element
    gamma = alpha[end] # Indexing the parameter so coefficient on Z is the last element
    K = size(X,2) # K is the no. of covariates in X, which in this case are 3
    J = length(unique(y)) # J is the no. of coefficients, which are 8 in this case (So J-1 = 7)
    N = length(y) # N is a vector with the same length as y (not sure why we are defining N)
    bigY = zeros(N,J) # Simply defining the matrix of dependent variables Y, which are going to be dummies in case of Multinomial Logit

    for j=1:J # Simply defining j to be something, next step, j represents column
        bigY[:,j] = y.==j # Making a dummies for each of the corresponding occupation in y
    end
    bigBeta = [reshape(beta,K,J-1) zeros(K)] # Reshaping into a matrix with K rows and J-1 columns
    # Adding a column of zeros in the end to make it conform with bigY and also to normalize the parameters of the last occupation to be zero

    T = promote_type(eltype(X),eltype(alpha)) # To get the standard errors we promote the type
    num   = zeros(T,N,J) # Define numerator T-type , N-rows, J-columns
    dem   = zeros(T,N) # Define denominator T-type , N-rows
    for j=1:J
        num[:,j] = exp.(X*bigBeta[:,j] .+ (Z[:,j] .- Z[:,J])*gamma) # Numerator Likelihood function
        dem .+= num[:,j] # Denominator Likelihood function: Simply the Sum of all the numerators
    end

    P = num./repeat(dem,1,J) # Putting together numerator and denominator

    loglike = -sum( bigY.*log.(P) ) # Since Optim minimizes the function, we take its negative to maximize

    return loglike # This is used for storing the function parameters
end

startvals = [2*rand(7*size(X,2)).-1; .1]
td = TwiceDifferentiable(alpha -> mlogit(alpha, X, Z, y), startvals; autodiff = :forward)

# Use Optim #
α_hat_multinomial = optimize(td, startvals, LBFGS(), Optim.Options(g_tol = 1e-5, iterations=100_000, show_trace=true, show_every=50))
α_hat_multinomial_ad = α_hat_multinomial.minimizer

# Find Hessian #
H  = Optim.hessian!(td, α_hat_multinomial_ad) # Finding the Hessian Matrix
α_hat_multinomial_ad_se = sqrt.(diag(inv(H))) # Standard Errors of alpha_hat

println([α_hat_multinomial_ad α_hat_multinomial_ad_se])

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# Question 2: Interpretation of gamma_hat
#:::::::::::::::::::::::::::::::::::::::::::::::::::
# gamma_hat make more sense as compared to PS3 because it's sign is positive #

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# Question 3:
#:::::::::::::::::::::::::::::::::::::::::::::::::::

### Part (A) ###
using Distributions
include("lgwt.jl") # make sure the function gets read in

# define distribution
d = Normal(0,1) # mean=0, standard deviation=1

# get quadrature nodes and weights for 7 grid points, with +_ 4SD
nodes, weights = lgwt(7,-4,4)

# now compute the integral over the density and verify it's 1
sum(weights.*pdf.(d,nodes)) # d here is normal dist. as defined earlier #

# now compute the expectation and verify it's 0
sum(weights.*nodes.*pdf.(d,nodes)) # it seems as if the nodes here is x in: intergral x phi(x) dx #

### Part (B) ###

#Practice 1#
d = Normal(0,2)
nodes, weights = lgwt(7,-5,5)
sum(weights.*(nodes.^2).*pdf.(d,nodes))

#Practice 2#
nodes, weights = lgwt(10,-5,5)
sum(weights.*(nodes.^2).*pdf.(d,nodes))

# Comment: Increasing the quadrature points gives us better estimates (of variace of f(x)) #

### Part (C) ###

#Practice 1#
Random.seed!(1234)
d = Normal(0,2)
N = 1000000
nodes, weights = lgwt(10,-5,5)
x = rand(Uniform(-10,10), N)
A = (10-(-10))*(1/N)*sum((x.^2).*pdf(d,x)) # A is approximately equal to 4 #

#Practice 2#
A1 = (10-(-10))*(1/N)*sum((x).*pdf(d,x)) # A1 is approx. equal to 0 #
#Practice 3#
A2 = (10-(-10))*(1/N)*sum(pdf(d,x)) # A2 is approx. equal to 1 #

#Practice 4#
D = 1000
x = rand(Uniform(-10,10), D)

A = (10-(-10))*(1/D)*sum((x.^2).*pdf(d,x))

A1 = (10-(-10))*(1/D)*sum((x).*pdf(d,x))

A2 = (10-(-10))*(1/D)*sum(pdf(d,x))
# Reducing D to 1000 more likely less precise estimates #

### Part (D) ###

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# Question 4: With Gauss-Lagendre Quadrature
#:::::::::::::::::::::::::::::::::::::::::::::::::::
function mixedlogit_quadrature(theta, X, Z, y, rn)

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

        nodes, weights = lgwt(rn,-4,4);
        d=Normal(0,1)

        T = promote_type(eltype(X),eltype(theta))
        num   = zeros(T,N,J)
        dem   = zeros(T,N)
        for j=1:J
            num[:,j] = sum(exp.(X*bigAlpha[:,j] .+ (Z[:,j] .- Z[:,J])*gamma).*weights.*pdf.(d,nodes))
            dem .+= num[:,j]
        end

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

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

        return loglike
    end

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# Question 5: With Monte Carlo Simulations
#:::::::::::::::::::::::::::::::::::::::::::::::::::
function mixedlogit_simulation(theta, X, Z, y,rn)

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

        d=Normal(0,1)
        rx=rand(Uniform(0,1 ),rn)

        T = promote_type(eltype(X),eltype(theta))
        num   = zeros(T,N,J)
        dem   = zeros(T,N)
        for j=1:J
            num[:,j] = sum(exp.(X*bigAlpha[:,j] .+ (Z[:,j] .- Z[:,J])*gamma).*(pdf.(d,x)))
            dem .+= num[:,j]
        end

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

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

        return loglike
    end

                                ### End ###


# # Thermal Generation Dispatch Example

#md # [![](https://img.shields.io/badge/GitHub-100000?style=for-the-badge&logo=github&logoColor=white)](@__REPO_ROOT_URL__/examples/Thermal_Generation_Dispatch_Example.jl)

# This example illustrates the sensitivity analysis of thermal generation dispatch problem.

# This problem can be described as the choice of thermal generation `g` given a demand `d`, a price for thermal generation `c` and a penalty price `c_{ϕ}` for any demand not attended ϕ.

# ```math
# \begin{split}
# \begin{array} {ll}
# \mbox{minimize} & \sum_{i=1}^{N} c_{i} g_{i} + c_{\phi} \phi \\
# \mbox{s.t.} & g_{i} \ge 0 \quad i=1..N  \\
#             & g_{i} \le G_{i} \quad i=1..N  \\
#             & \sum_{i=1}^{N} g_{i} + \phi = d\\
# \end{array}
# \end{split}
# ```
# where
# - `G_{i}` is the maximum possible generation for a thermal generator `i`

# ## Define and solve the Thermal Dispatch Problem

# First, import the libraries.

using Test
using JuMP
import DiffOpt
import LinearAlgebra: dot
import HiGHS
import MathOptInterface
import Plots
const MOI = MathOptInterface

# Define the model that will be construct given a set of parameters.

function generate_model(d::Float64; g_sup::Vector{Float64}, c_g::Vector{Float64}, c_ϕ::Float64)
    ## Creation of the Model and Parameters
    model = Model(() -> DiffOpt.diff_optimizer(HiGHS.Optimizer))
    set_silent(model)
    I = length(g_sup)

    ## Variables
    @variable(model, g[i in  1:I] >= 0.0)
    @variable(model, ϕ >= 0.0)

    ## Constraints
    @constraint(model, limit_constraints_sup[i in 1:I], g[i] <= g_sup[i])
    @constraint(model, demand_constraint, sum(g) + ϕ == d)

    ## Objectives
    @objective(model, Min, dot(c_g, g) + c_ϕ * ϕ)

    ## Solve the model
    optimize!(model)

    ## Return the solved model
    return model
end

# Define the functions that will get the primal values `g` and `\phi` and sensitivity analysis of the demand `dg/dd` and `d\phi/dd` from a optimized model.

function diff_forward(model::Model, ϵ::Float64 = 1.0)
    ## Initialization of parameters and references to simplify the notation
    vect_ref = [model[:g]; model[:ϕ]]
    I = length(model[:g])

    ## Get the primal solution of the model
    vect = MOI.get.(model, MOI.VariablePrimal(), vect_ref)

    ## Pass the perturbation to the DiffOpt Framework and set the context to Forward
    constraint_equation = convert(MOI.ScalarAffineFunction{Float64}, ϵ)
    MOI.set(model, DiffOpt.ForwardConstraintFunction(), model[:demand_constraint], constraint_equation)
    DiffOpt.forward_differentiate!(model)

    ## Get the derivative of the model
    dvect = MOI.get.(model, DiffOpt.ForwardVariablePrimal(), vect_ref)

    ## Return the values as a vector
    return [vect; dvect]
end

function diff_reverse(model::Model, ϵ::Float64 = 1.0)
    ## Initialization of parameters and references to simplify the notation
    vect_ref = [model[:g]; model[:ϕ]]
    I = length(model[:g])

    ## Get the primal solution of the model
    vect = MOI.get.(model, MOI.VariablePrimal(), vect_ref)

    ## Set variables needed for the DiffOpt Backward Framework
    dvect = Array{Float64, 1}(undef, I + 1)
    perturbation = zeros(I + 1)

    ## Loop for each primal variable
    for i in 1:I+1
        ## Set the perturbation in the Primal Variables and set the context to Backward
        perturbation[i] = ϵ
        MOI.set.(model, DiffOpt.ReverseVariablePrimal(), vect_ref, perturbation)
        DiffOpt.reverse_differentiate!(model)

        ## Get the value of the derivative of the model
        dvect[i] = JuMP.constant(MOI.get(model, DiffOpt.ReverseConstraintFunction(), model[:demand_constraint]))
        perturbation[i] = 0.0
    end

    ## Return the values as a vector
    return [vect;dvect]
end

# Initialize of Parameters

g_sup = [10.0, 20.0, 30.0]
I = length(g_sup)
d = 0.0:0.1:80
d_size = length(d)
c_g = [1.0, 3.0, 5.0]
c_ϕ = 10.0
;

# Generate models for each demand `d`
models = generate_model.(d; g_sup = g_sup, c_g = c_g, c_ϕ = c_ϕ);

# Get the results of models with the DiffOpt Forward and Backward context

result_forward = diff_forward.(models)
optimize!.(models)
result_reverse = diff_reverse.(models);

# Organization of results to plot
# Initialize data_results that will contain every result
data_results = Array{Float64,3}(undef, 2, d_size, 2*(I+1));

# Populate the data_results array
for k in 1:d_size
    data_results[1,k,:] = result_forward[k]
    data_results[2,k,:] = result_reverse[k]
end

# ## Results with Plot graphs
# ### Results for the forward context
# Result Primal Values:
Plots.plot(d,data_results[1,:,1:I+1],
    title="Generation by Demand",label=["Thermal Generation 1" "Thermal Generation 2" "Thermal Generation 3" "Generation Deficit"],
    xlabel="Demand [unit]",ylabel= "Generation [unit]"
)

# Result Sensitivity Analysis:
Plots.plot(d,data_results[1,:,I+2:2*(I+1)],
    title="Sensitivity of Generation by Demand",label=["T. Gen. 1 Sensitivity" "T. Gen. 2 Sensitivity" "T. Gen. 3 Sensitivity" "Gen. Deficit Sensitivity"],
    xlabel="Demand [unit]",ylabel= "Sensitivity [-]"
)

# ### Results for the reverse context
# Result Primal Values:
Plots.plot(d,data_results[2,:,1:I+1],
    title="Generation by Demand",label=["Thermal Generation 1" "Thermal Generation 2" "Thermal Generation 3" "Generation Deficit"],
    xlabel="Demand [unit]",ylabel= "Generation [unit]"
)

# Result Sensitivity Analysis:
Plots.plot(d,data_results[2,:,I+2:2*(I+1)],
    title="Sensitivity of Generation by Demand",label=["T. Gen. 1 Sensitivity" "T. Gen. 2 Sensitivity" "T. Gen. 3 Sensitivity" "Gen. Deficit Sensitivity"],
    xlabel="Demand [unit]",ylabel= "Sensitivity [-]"
)


# function to calcule the force between a pair of particles
# Exercise 39_40
using simulationQP934

#inputs
data = Data();
box  = initial_point(data);
particle1 = box[1];
particle2 = box[2];

function fpair(x,y,r,data)

  # r values for F calculation
  r6  = r^6
  r7  = r6 * r
  r12 = r6^2
  r13 = r12 * r

  # x and y components
  # Fz = -( dU(r) / dz) = - (dr/dz) * (dU(r)/dr), r = sqrt(dx^2,dy^2)
  # calculation for -(dr/dr)

  drdx1 = -(x[1]-y[1])/r
  drdx2 = -(x[2]-y[2])/r

  dudr1 = -12*(data.sig12/r13)
  dudr2 = -6*(data.sig6/r7)

  dfacdr =  data.eps4*(dudr1 - dudr2)

  upair = data.eps4*(data.sig12/r12 - data.sig6/r6)
  f = ( -dfacdr*drdx1, -dfacdr*drdx2 )

  return upair, f

end

r = pbcseparation(particle1,particle2,data.side)

f12 = fpair(particle1,particle2,r,data)

f21 = fpair(particle2,particle1,r,data)


"""
    newton(x0, f, j; maxiter, tol)

Solution to f(x) = 0 where f'(x) = j(x) using Newton's method.
"""
function newton(x0, f, j; maxiter = 1000, tol = 1e-14)
    x = x0
    for _ in 1:maxiter
        fx = f(x)
        if abs(fx) < tol
            return x
        else
            x = x - (fx / j(x))
        end
    end
    return x
end


# # One-Dimensional Cases
#
# The 1D case in Optimal Transport is a special case where one can
# easily obtain closed form solutions efficiently
# when the cost function is convex. In this situation,
# one does no need to use Linear Programming solvers
# to obtain the exact solution to the problem.

# ## Packages
#
# We load the following packages into our environment:

using OptimalTransport
using Distances
using Distributions
using StatsPlots

using LinearAlgebra
using Random

Random.seed!(1234);

# ## Continuous Distribution
#
# In the 1D case, when the source measure $\mu$ is continuous and the cost function
# has the form ``c(x, y) = h(|x - y|)`` where ``h`` is a convex function,
# the optimal transport plan is the Monge map
# ```math
# T = F_\nu^{-1} \circ F_\mu
# ```
# where ``F_\mu`` is the cumulative distribution function of `μ` and ``F_\nu^{-1}`` is the
# quantile function of `ν`.
# In this setting, the optimal transport cost can be computed as
# ```math
# \int_0^1 c(F_\mu^{-1}(x), F_\nu^{-1}(x)) \mathrm{d}x
# ```
# where ``F_\mu^{-1}`` and ``F_\nu^{-1}`` are the quantile functions of `μ` and `ν`,
# respectively.
#
# We start by defining the distributions.

μ = Normal(0, 1)

N = 10
ν = Poisson(N);

# Nest, we define a cost function.
c(x, y) = (abs(x - y))^2 # could have used `sqeuclidean` from `Distances.jl`

T = ot_plan(c, μ, ν);

# `T` is the Monge Map. Let's visualize it.
p1 = plot(μ; label='μ')
p1 = plot!(ν; marker=:circle, label='ν')
p2 = plot(-2:0.1:2, T(-2:0.1:2); label="Monge map", color=:green, legend=:topleft)
plot(p1, p2)

# The optimal transport cost can be computed with
ot_cost(c, μ, ν)

# If instead you want the 2-Wasserstein distance (which is the square root
# of the optimal transport with the Square Euclidean distatce, then use
wasserstein(μ, ν; p=2)

# ## Finite Discrete Distributions
#
# If the source and target measures are 1D finite discrete distributions
# (sometimes referred as empirical distributions, or as sample distributions),
# and if the cost function is convex, then the optimal
# transport plan can be written as a sorting algorithm,
# where the utmost left probability mass of the source is transported
# to the closest probability mass of the target, until everything is transported.
#
# Define your measures as DiscreteNonParametric, which is a type in Distributions.jl.
# Also, let's assume both point masses with equal weights and let's
# use the `sqeuclidean` function instead of creating our own cost function.

M = 15
μ = DiscreteNonParametric(1.5rand(M), fill(1 / M, M))

N = 10
ν = DiscreteNonParametric(1.5rand(N) .+ 2, fill(1 / N, N))

γ = ot_plan(sqeuclidean, μ, ν);

# This time γ is a sparse matrix containing the transport plan. Let's visualize the results.
# We create a function `curve` just as a helper to draw the transport plan.

function curve(x1, x2, y1, y2)
    a = min(y1, y2)
    b = (y1 - y2 + a * (x1^2 - x2^2)) / (x1 - x2)
    c = y1 + a * x1^2 - b * x1
    f(x) = -a * x^2 + b * x + c
    return f
end

p = plot(μ; marker=:circle, label='μ')
p = plot!(ν; marker=:circle, label='ν', ylims=(0, 0.2))
for i in 1:M, j in 1:N
    if γ[i, j] > 0
        transport = curve(μ.support[i], ν.support[j], 1 / M, 1 / N)
        x = range(μ.support[i], ν.support[j]; length=100)
        p = plot!(x, transport.(x); color=:green, label=nothing, alpha=0.5)
    end
end
p

# Again, the optimal transport cost can be calculated with

ot_cost(sqeuclidean, μ, ν)


# Period doubling bifurcation + cobweb for logistic map
using DrWatson
@quickactivate "NonlinearDynamicsTextbook"
include(srcdir("style.jl"))
using DynamicalSystems, PyPlot, Random

lo = Systems.logistic(0.4)
rs = (2.8, 3.3)

fig, axs = subplots(1, 2, figsize = (0.66figx, figy))

T = 1000
xs = 0:0.0001:1

function cobweb(t) # transform timeseries t into cobweb (point2D)
    cx = Float64[]; cy = Float64[]
    for i ∈ 1:length(t)-1
        push!(cx, t[i]); push!(cy, t[i])
		push!(cx, t[i]); push!(cy, t[i+1])
    end
    # add line to x axis
    pushfirst!(cx, cx[1])
    pushfirst!(cy, 0)
    return cx, cy
end

for (i, r) in enumerate(rs)
    ax = axs[i]
    ax.set_xlim(0,1)
    ax.set_ylim(0,1)
    ax.set_xticks([0, 1])
    ax.set_xlabel("\$x\$"; labelpad=-30)
	set_parameter!(lo, 1, r)
	f  = lo.f.(xs, Ref([r]), 0)
    f² = lo.f.(f, Ref([r]), 0)
	ax.plot(xs, f, color = "C0")
	ax.plot([0, 1], [0,1]; lw = 2, color = "C3", ls = "--")
	x = trajectory(lo, T)
    cx, cy = cobweb(x)
    if i>1
        ax.plot(xs, f², color = "C1")
        ax.text(0.15, 0.875, "\$f^{(2)}\$"; color = "C1", transform = ax.transAxes, size = 32)
    end
    ax.plot(cx, cy; lw = 1.0, color = "C2")
    ax.text(0.05, 0.875, "\$f\$"; color = "C0", transform = ax.transAxes, size = 32)
end
axs[2].set_yticklabels([])

fig.tight_layout(pad = 0.37)
add_identifiers!(fig)

wsave(plotsdir("4", "period_doubling_bifurcation"), fig)


using LinearAlgebra
using Plots

include("cheb.jl")
include("parameters.jl")
include("geometry.jl")
include("profile.jl")
include("build_A.jl")
include("compute_spectrum.jl")
include("mesh.jl")
include("plot_fields.jl")

params=parameters(
     H = 1,
    Lϕ = π/100,
  ϕmin = π/4,
   TwoΩ = 1,
      g = 10,
      a = 1,
    Lj = π/2000,
    Uj = 1.0,
    Nϕ = 100,
    dk = 1e-2,
    kₘ = 3
    );

params.Fr = (params.TwoΩ * params.Lj)^2/(params.g * params.H);
params.Ro =  params.Uj/(params.TwoΩ * params.Lj)

print("\n")
print("Linear-Stability-Caluculator\n")
print("============================\n\n")
print("Repo:      Linear-Stability-Calculators\n")
print("Code:      ShallowWater/Julia/Spherical/linear_stability_shallow_water.jl\n")
print("Model:     Rotating Shallow Water\n")
print("Geometry:  Spherical\n")
print("Structure: Bickley Jet\n")
print("\n")
print("Nondimensional Parameters\n")
print("=========================\n")
print("Fr = ", params.Fr, "\n")
print("Ro = ", params.Ro, "\n")

# Wavenumbers
ks = collect(params.dk:params.dk:params.kₘ) / params.Lj;
Nk = length(ks);

# Grid and basic state
Dϕ, Dϕ2, ϕ = geometry(params);
      U,  E = profile(  ϕ, params);

plot_basic_state(ϕ, U, E, "basic_state.png")

# initialize fields to store
Nmodes = 2
     σ = zeros(Nmodes, Nk);
     ω = zeros(Nmodes, Nk);
σmodes = zeros(ComplexF64, 3*params.Nϕ+1, Nmodes, Nk);

# Compute growth rates
for cnt in 1:Nk
    local k = ks[cnt]
    local A = build_A(k, U, E, Dϕ, Dϕ2, ϕ, params);
    local σ[:,cnt], ω[:,cnt], σmodes[:,:,cnt] = compute_spectrum(A, k, params, Nmodes);
end

plot_growth_rates(ks, σ, Nmodes, "growth_rates_spherical_Bickley_jet.png")

#print("σ = ", σ[1,:],"\n")

mode_number = 1;
      σ_max = maximum(σ[mode_number,:]);
    k_index = sortperm(σ[mode_number,:],rev=true)[1];
          k = ks[k_index]; # pick wavenumber

plot_1D_streamfunction(k_index, k, ϕ, σmodes, mode_number, "modes_1D_streamfunction.png")
plot_2D_streamfunction(k_index, k, ϕ, σmodes, mode_number, "modes_2D_streamfunction.png")
plot_2D_vorticity(Dϕ,  k_index, k, ϕ, σmodes, mode_number, "modes_2D_vorticity.png")


module CartesianSphericalHarmonics

# Write your package code here.
using LinearAlgebra
using Reexport

@reexport using MultivariatePolynomials
@reexport using TypedPolynomials

export YLMNorm, Schmidt, Laplace, Nonorm, Full
export rlm, ylm
#normalization factor
abstract type YLMNorm{T} end

struct Schmidt{T} <: YLMNorm{T}; end
struct Laplace{T} <: YLMNorm{T}; end
struct Nonorm{T} <: YLMNorm{T}; end
struct Full{T} <: YLMNorm{T}; end

function ylmcoeff(N::Schmidt{T},l::Integer, m::Integer) where T

	k = one(T)
	for i in (l-m+1):(l+m)
	  k *= i
	end
	fac = sqrt(1/k)
	if m!=0
		fac*=sqrt(2*one(T))
	end
	return fac
end

function ylmcoeff(N::Laplace{T}, l::Integer, m::Integer) where T

  k = one(T)
  for i in (l-m+1):(l+m)
    k *= i
  end

  return sqrt((2*l+1) / (4*T(pi)*k))
end

function ylmcoeff(N::Full{T}, l::Integer, m::Integer) where T
	fac = m==0 ? one(T) : sqrt(2*one(T))
	return ylmcoeff(Laplace{T}(),l,m)*sqrt(4T(pi))*fac
end

function ylmcoeff(N::Nonorm{T}, l::Integer, m::Integer) where T
 return one(T)
end

function cossinpoly(m::Integer, x::Variable, y::Variable)

  sum = zero(x*y)
  for j in 0:div(m,2)
	  k = 2j
    sum += ((-1)^j)*binomial(m, k)*y^Int(k)*x^Int(m-k)
  end
  return sum
end

function sinsinpoly(m::Integer, x::Variable, y::Variable)

  sum = zero(x*y)
  for j in 0:div((m-1),2)
	  k = 2j+1
    sum += ((-1)^j)*binomial(m, k)*y^Int(k)*x^Int(m-k)
  end
  return sum
end

#
function legendre_assoc(::Type{T},z,l,m) where T
	p = (z^2 - 1)^l

	for i = 1:l+abs(m)
	c = i <= l ? 1/(2one(T)*i) : one(T)
	p = c*differentiate(p, z)
	end

	p *= (-1)^m
	return p
end
# #legendre polynomial:
# P(l::BigInt,x) = differentiate((x^2 - 1)^l,x,l)
#
# #associated legendre polynomial (without sin(θ)^m)
# P(l::BigInt,m::BigInt,x) = 1//(big(2)^l*factorial(l))*differentiate(P(l,x),x,m)

"""
    ylm(l::Integer, m::Integer, x::Variable, y::Variable, z::Variable)
*Description:*  Calculation of the spherical harmonic for a given order (l,m) in Cartesian coordinates\\

*Input:*  `l`       - Degree of the spherical harmonic\\
          `m`       - Order of the spherical harmonic\\
          `x, y, z` - Cartesian coordinates\\

*Output:*  Spherical harmonic polynomial
"""
function ylm(l::Integer, m::Integer, x::Variable, y::Variable, z::Variable;
			 norm::YLMNorm{T}=Nonorm{Rational{Int}}(),real::Bool=true) where T

	if abs(m) > l
		throw(DomainError(m,"-l <= m <= l expected, but m = $m and l = $l."))
	end

	p = legendre_assoc(T,z,l,m)
	# p = (z^2 - 1)^l

	# for i = 1:l+abs(m)
	#   c = i <= l ? 1/(2one(T)*i) : one(T)
	#   p = c*differentiate(p, z)
	# end

	# p *= (-1)^m

	if real
		if m > 0
			return ylmcoeff(norm, l, m)*cossinpoly(m,x,y)*p
		elseif m < 0
			return ylmcoeff(norm, l, abs(m))*sinsinpoly(abs(m),x,y)*p
		else
			return ylmcoeff(norm, l, 0)*p
		end
	else
		return ylmcoeff(norm, l, m)*(cossinpoly(m,x,y)+im*sinsinpoly(m,x,y))*p
	end
end

# multiplying r^l*ylm(x,y,z)
function rlm(l::Integer, m::Integer, x::Variable, y::Variable, z::Variable;
			   norm::YLMNorm{T}=Nonorm{Rational{Int}}(),real::Bool=true) where T
	p = ylm(l,m,x,y,z;norm=norm,real=real)
	tout = []

	for t in terms(p)
		deg = degree(monomial(t)) # degree of monomial
		degR = l-deg # degree of r: l-deg
		push!(tout,(x^2+y^2+z^2)^div(degR,2)*t) # r² replaced by x²+y²+z²
	end

	return polynomial(tout)
end

end


"""
    poisson_rng(p, n, shape=1; seed=nothing)

Generate a `shape` element array of random variables from a Poisson(`λ`) distribution. Optionally you can set a specific seed.

# Notes

The Poisson distribution is given by:

``f(x, λ) = \\frac{λ^x}{x!} e^λ \\quad x = 0,1,2,…``

# Examples

```julia-repl
julia> poisson_rng(3)
1×1 Matrix{Int64}:
 7

julia> poisson_rng(10, 5)
5×1 Matrix{Int64}:
 13
 11
 10
  8
 15

julia> poisson_rng(10, (5,5))
5×5×1 Array{Int64, 3}:
[:, :, 1] =
 11  15   9  11   9
  8  15  13  10   9
 11  12   4  10   6
  7   9  13  11   7
 13   7  10  10  14

```

# References

R. Larson, A. Odoni. Urban operations research. Prentice-Hall, New Jersey, 1981.

G. Last, M. Penrose. Lectures on the poisson process. Cambridge University Press, 2017.
"""
function poisson_rng(λ::Real, shape::Union{Int, Tuple{Vararg{Int}}}=1; seed::Union{Int, Nothing}=nothing)
    n = ceil(Int, λ*5)  # ensure n is integer
    U = expon_rng(λ, (shape..., n), seed=seed)
    X = sum(cumsum(U, dims=ndims(U)) .< 1, dims=ndims(U))
    return X
end


immutable Beta <: ContinuousUnivariateDistribution
    α::Float64
    β::Float64

    function Beta(a::Real, b::Real)
        (a > zero(a) && b > zero(b)) || error("α and β must be positive")
        new(float64(a), float64(b))
    end

    Beta(α::Real) = Beta(α, α)
    Beta() = new(1.0, 1.0)
end

@_jl_dist_2p Beta beta

@distr_support Beta 0.0 1.0

#### Parameters

params(d::Beta) = (d.α, d.β)

#### Statistics

mean(d::Beta) = ((α, β) = params(d); α / (α + β))

function mode(d::Beta)
    (α, β) = params(d)
    (α > 1.0 && β > 1.0) || error("mode is defined only when α > 1 and β > 1.")
    return (α - 1.0) / (α + β - 2.0)
end

modes(d::Beta) = [mode(d)]

function var(d::Beta)
    (α, β) = params(d)
    s = α + β
    return (α * β) / (abs2(s) * (s + 1.0))
end

meanlogx(d::Beta) = ((α, β) = params(d); digamma(α) - digamma(α + β))

varlogx(d::Beta) = ((α, β) = params(d); trigamma(α) - trigamma(α + β))

function skewness(d::Beta)
    (α, β) = params(d)
    if α == β
        return 0.0
    else
        s = α + β
        (2.0 * (β - α) * sqrt(s + 1.0)) / ((s + 2.0) * sqrt(α * β))
    end
end

function kurtosis(d::Beta)
    α, β = params(d)
    s = α + β
    p = α * β
    6.0 * (abs2(α - β) * (s + 1.0) - p * (s + 2.0)) / (p * (s + 2.0) * (s + 3.0))
end

function entropy(d::Beta)
    α, β = params(d)
    s = α + β
    lbeta(α, β) - (α - 1.0) * digamma(α) - (β - 1.0) * digamma(β) +
        (s - 2.0) * digamma(s)
end

#### Evaluation

gradlogpdf(d::Beta, x::Float64) =
    ((α, β) = params(d); 0.0 <= x <= 1.0 ? (α - 1.0) / x - (β - 1.0) / (1 - x) : 0.0)

## Fit model

# TODO: add MLE method (should be similar to Dirichlet)

# This is a moment-matching method (not MLE)
#
function fit(::Type{Beta}, x::AbstractArray)
    x_bar = mean(x)
    v_bar = varm(x, x_bar)
    α = x_bar * (((x_bar * (1.0 - x_bar)) / v_bar) - 1.0)
    β = (1.0 - x_bar) * (((x_bar * (1.0 - x_bar)) / v_bar) - 1.0)
    Beta(α, β)
end


import Healpix
export std2fwhm, fwhm2std, gaussian_beam

@doc raw"""
    std2fwhm(stddev)

Convert a standard deviation into a Full Width Half Maximum (FWHM) value. The
measure unit for `stddev` can be arbitrary, as the transformation is linear.
See also `fwhm2std`.
"""
std2fwhm(stddev) = 2 * sqrt(2 * log(2)) * stddev

@doc raw"""
    std2fwhm(stddev)

Convert a standard deviation into a Full Width Half Maximum (FWHM) value. The
measure unit for `stddev` can be arbitrary, as the transformation is linear.
See also `std2fwhm`.
"""
fwhm2std(fwhm) = fwhm / (2 * sqrt(2 * log(2)))

@doc raw"""
    gaussian_beam!(beam_map::Healpix.Map{T,O}, angle_std_rad; normalization = 1.0) where {T,O <: Healpix.Order}

Save a map of a circual Gaussian beam in `beam_map`. The width of the beam is
provided through the parameter `angle_std_rad`, which is expressed in radians.
If you want to specify the FWHM instead, use `fwhm2std` like in the following
example:

````julia
import Healpix
nside = 1024
beam_map = Healpix.Map{Float64, RingOrder}(nside)

# Assume 2.5° of FWHM
gaussian_beam!(beam_map, 2.5 |> fwhm2std)
````

See also `gaussian_beam`.
"""
function gaussian_beam!(beam_map::Healpix.Map{T,O}, angle_std_rad; normalization = 1.0) where {T,O <: Healpix.Order}
    for pixidx in 1:beam_map.resolution.numOfPixels
        theta, _ = Healpix.pix2ang(beam_map, pixidx)
        beam_map[pixidx] = normalization * exp(-theta^2 / (2 * angle_std_rad))
    end
end

function gaussian_beam(nside::Integer, angle_std_rad; normalization = 1.0)
    beam_map = Healpix.Map{Float64,Healpix.RingOrder}(nside)
    gaussian_beam!(beam_map, angle_std_rad, normalization = normalization)

    beam_map
end

function gaussian_beam(res::Healpix.Resolution, angle_std_rad; normalization = 1.0)
    gaussian_beam(res.nside, angle_std_rad, normalization = normalization)
end

@doc raw"""
    gaussian_beam(res::Healpix.Resolution, angle_std_rad; normalization = 1.0) where {T,O <: Healpix.Order}
    gaussian_beam(nside::Integer, angle_std_rad; normalization = 1.0) where {T,O <: Healpix.Order}

Save a map of a circual Gaussian beam in `beam_map`. The width of the beam is
provided through the parameter `angle_std_rad`, which is expressed in radians.
If you want to specify the FWHM instead, use `fwhm2std` like in the following
example:

````julia
# Assume 2.5° of FWHM
beam_map = gaussian_beam(1024, 2.5 |> fwhm2std)
````

See also `gaussian_beam!`.
"""
gaussian_beam

struct BeamMoments
    S::Array{Float32,1}
    M::Array{Float32,2}
end

@doc raw"""

    beam_m(beam_map, l, m, n)

Calculate the moment `(l, m, n)` of the beam in the Healpix map `beam_map`. Beam
moments are defined in Appendix A of the paper "Planck 2013 results. LFI
calibration" (Planck collaboration, A&A, 2013),
https://dx.doi.org/10.1051/0004-6361/201321527.

This function does not take into account the factor `N` (normalization of the beam).
You can compute it using `beam_m(beam_map, 0, 0, 0)`.

"""
function beam_m(beam_map, l::Integer, m::Integer, n::Integer)
    result = 0.0
    for idx in 1:length(beam_map)
        θ, ϕ = Healpix.pix2ang(beam_map, idx)

        # We could have used `Healpix.ang2vec`, but we need the value
        # of sinθ in order to calculate dΩ = sinθ⋅dθ⋅dϕ
        sinθ, cosθ = sincos(θ)
        sinϕ, cosϕ = sincos(ϕ)

        x = sinθ * cosϕ
        y = sinθ * sinϕ
        z = cosθ

        result += x^l * y^m * z^n * sinθ * beam_map[idx]
    end
    result
end

function beam_moments(beam_map)
    N = beam_m(beam_map, 0, 0, 0)

    S = [
        beam_m(beam_map, 1, 0, 0) / N,
        beam_m(beam_map, 0, 1, 0) / N,
        beam_m(beam_map, 0, 0, 1) / N,
    ]

    M = [
        [(beam_m(beam_map, 2, 0, 0) / N)  (beam_m(beam_map, 1, 1, 0) / N)  (beam_m(beam_map, 1, 0, 1) / N)];
        [(beam_m(beam_map, 1, 1, 0) / N)  (beam_m(beam_map, 0, 2, 0) / N)  (beam_m(beam_map, 0, 1, 1) / N)];
        [(beam_m(beam_map, 1, 0, 1) / N)  (beam_m(beam_map, 0, 1, 1) / N)  (beam_m(beam_map, 0, 0, 2) / N)];
    ]

    BeamMoments(S, M)
end


# indicator of the L2 norm sphere with given radius

export IndSphereL2

"""
**Indicator of a Euclidean sphere**

    IndSphereL2(r=1.0)

Returns the indicator function of the set
```math
S = \\{ x : \\|x\\| = r \\},
```
where ``\\|\\cdot\\|`` is the ``L_2`` (Euclidean) norm. Parameter `r` must be positive.
"""
struct IndSphereL2{R <: Real} <: ProximableFunction
    r::R
    function IndSphereL2{R}(r::R) where {R <: Real}
        if r <= 0
            error("parameter r must be positive")
        else
            new(r)
        end
    end
end

is_set(f::IndSphereL2) = true

IndSphereL2(r::R=1.0) where {R <: Real} = IndSphereL2{R}(r)

function (f::IndSphereL2)(x::AbstractArray{T}) where {R <: Real, T <: RealOrComplex{R}}
    if isapprox(norm(x), f.r, atol=eps(R), rtol=sqrt(eps(R)))
        return R(0)
    end
    return R(Inf)
end

function prox!(y::AbstractArray{T}, f::IndSphereL2, x::AbstractArray{T}, gamma::R=R(1)) where {R <: Real, T <: RealOrComplex{R}}
    normx = norm(x)
    if normx > 0 # zero-zero?
        scal = f.r/normx
        for k in eachindex(x)
            y[k] = scal*x[k]
        end
    else
        normy = R(0)
        for k in eachindex(x)
            y[k] = randn()
            normy += y[k]*y[k]
        end
        normy = sqrt(normy)
        y .*= f.r/normy
    end
    return R(0)
end

fun_name(f::IndSphereL2) = "indicator of an L2 norm sphere"
fun_dom(f::IndSphereL2) = "AbstractArray{Real}, AbstractArray{Complex}"
fun_expr(f::IndSphereL2) = "x ↦ 0 if ||x|| = r, +∞ otherwise"
fun_params(f::IndSphereL2) = "r = $(f.r)"

function prox_naive(f::IndSphereL2, x::AbstractArray{T}, gamma::R=R(1)) where {R <: Real, T <: RealOrComplex{R}}
    normx = norm(x)
    if normx > 0
        y = x*f.r/normx
    else
        y = randn(size(x))
        y *= f.r/norm(y)
    end
    return y, R(0)
end


export MoreauEnvelope

"""
**Moreau envelope**

    MoreauEnvelope(f, γ=1.0)

Returns the Moreau envelope (also known as Moreau-Yosida regularization) of function `f` with parameter `γ` (positive), that is
```math
f^γ(x) = \\min_z \\left\\{ f(z) + \\tfrac{1}{2γ}\\|z-x\\|^2 \\right\\}.
```
If ``f`` is convex, then ``f^γ`` is a smooth, convex, lower approximation to ``f``, having the same minima as the original function.
"""
mutable struct MoreauEnvelope{R <: Real, T <: ProximableFunction} <: ProximableFunction
    g::T
    lambda::R
    function MoreauEnvelope{R, T}(g::T, lambda::R) where {R, T}
        if lambda <= 0 error("parameter lambda must be positive") end
        new(g, lambda)
    end
end

MoreauEnvelope(g::T, lambda::R=1.0) where {R <: Real, T <: ProximableFunction} = MoreauEnvelope{R, T}(g, lambda)

is_convex(f::MoreauEnvelope) = is_convex(f.g)
is_smooth(f::MoreauEnvelope) = is_convex(f.g)
is_quadratic(f::MoreauEnvelope) = is_generalized_quadratic(f.g)
is_strongly_convex(f::MoreauEnvelope) = is_strongly_convex(f.g)

function (f::MoreauEnvelope)(x::AbstractArray)
    buf = similar(x)
    g_prox = prox!(buf, f.g, x, f.lambda)
    return g_prox + 1/(2*f.lambda)*norm(buf .- x)^2
end

function gradient!(grad::AbstractArray, f::MoreauEnvelope, x::AbstractArray)
    g_prox = prox!(grad, f.g, x, f.lambda)
    grad .= (x .- grad)./f.lambda
    fx = g_prox + (f.lambda/2)*norm(grad)^2
    return fx
end

fun_name(f::MoreauEnvelope,i::Int64) =
"f$(i)(prox{λ$(i),f$(i)}(A$(i)x))+ 1/2 ‖x - prox{λ$(i),f$(i)}(A$(i)x)‖²"

fun_par( f::MoreauEnvelope,i::Int64)  = "λ$i = $(round(f.lambda, digits=3))"


# This script provides implementations of
# multi-step methods of different orders
# for solving a system of differential equations
# and compares them to the exact solutions and
# other methods while solving a 1-body problem

println('\n', " "^4, "> Loading the packages...")

using LinearAlgebra
using Printf

# Define the value of ϰ
const ϰ = -1

# Integrate equations of motion using the
# Euler's method, return the values of position
# and velocity on the last step
function euler(
    r₀₁::Vector{F},
    r₀₂::Vector{F},
    r₀₃::Vector{F},
    v₀₁::Vector{F},
    v₀₂::Vector{F},
    v₀₃::Vector{F},
    h::F,
    n::I,
)::Tuple{Vector{F},Vector{F},Vector{F},Vector{F},Vector{F},Vector{F}} where
{F <: AbstractFloat,I <: Unsigned}
    # Determine the length of the input vectors
    N = length(r₀₁)
    # Prepare the output vectors
    r₁ = copy(r₀₁)
    r₂ = copy(r₀₂)
    r₃ = copy(r₀₃)
    v₁ = copy(v₀₁)
    v₂ = copy(v₀₂)
    v₃ = copy(v₀₃)
    # Compute the solutions (in-place)
    for _ in 1:n
        ρ₁₂ = norm(r₁ - r₂)^3
        ρ₁₃ = norm(r₁ - r₃)^3
        ρ₂₃ = norm(r₂ - r₃)^3
        for k in 1:N
            a₁₂ = ϰ * (r₁[k] - r₂[k]) / ρ₁₂
            a₁₃ = ϰ * (r₁[k] - r₃[k]) / ρ₁₃
            a₂₃ = ϰ * (r₂[k] - r₃[k]) / ρ₂₃
            a₂₁ = -a₁₂
            a₃₁ = -a₁₃
            a₃₂ = -a₂₃
            r₁[k] += h * v₁[k]
            r₂[k] += h * v₂[k]
            r₃[k] += h * v₃[k]
            v₁[k] += h * (a₁₂ + a₁₃)
            v₂[k] += h * (a₂₁ + a₂₃)
            v₃[k] += h * (a₃₁ + a₃₂)
        end
    end
    return r₁, r₂, r₃, v₁, v₂, v₃
end

# Integrate the three-body problem using the
# two-step Adams–Bashforth's method, return the
# values of position and velocity on the last step
function ab2(
    r₀₁::Vector{F},
    r₀₂::Vector{F},
    r₀₃::Vector{F},
    v₀₁::Vector{F},
    v₀₂::Vector{F},
    v₀₃::Vector{F},
    h::F,
    n::I,
)::Tuple{Vector{F},Vector{F},Vector{F},Vector{F},Vector{F},Vector{F}} where
{F <: AbstractFloat,I <: Unsigned}
    # Determine the length of the input vectors
    N = length(r₀₁)
    # Prepare the output vectors
    r₁ = copy(r₀₁)
    r₂ = copy(r₀₂)
    r₃ = copy(r₀₃)
    v₁ = copy(v₀₁)
    v₂ = copy(v₀₂)
    v₃ = copy(v₀₃)
    # Prepare buffers for previous values
    r₁_ₖ₋₁ = copy(r₀₁)
    r₂_ₖ₋₁ = copy(r₀₂)
    r₃_ₖ₋₁ = copy(r₀₃)
    v₁_ₖ₋₁ = copy(v₀₁)
    v₂_ₖ₋₁ = copy(v₀₂)
    v₃_ₖ₋₁ = copy(v₀₃)
    # Compute the second value of the solution
    # by using the one-step Euler's method
    r₁, r₂, r₃, v₁, v₂, v₃ = euler(r₀₁, r₀₂, r₀₃, v₀₁, v₀₂, v₀₃, h, UInt(1))
    # Define a couple of independent coefficients
    k₁ = 3 / 2 * h
    k₂ = -1 / 2 * h
    # Compute the rest in two steps
    for _ in 2:n
        ρ₁₂_ₖ = norm(r₁ - r₂)^3
        ρ₁₃_ₖ = norm(r₁ - r₃)^3
        ρ₂₃_ₖ = norm(r₂ - r₃)^3
        ρ₁₂_ₖ₋₁ = norm(r₁_ₖ₋₁ - r₂_ₖ₋₁)^3
        ρ₁₃_ₖ₋₁ = norm(r₁_ₖ₋₁ - r₃_ₖ₋₁)^3
        ρ₂₃_ₖ₋₁ = norm(r₂_ₖ₋₁ - r₃_ₖ₋₁)^3
        for k in 1:N
            a₁₂_ₖ = k₁ * ϰ * (r₁[k] - r₂[k]) / ρ₁₂_ₖ
            a₁₃_ₖ = k₁ * ϰ * (r₁[k] - r₃[k]) / ρ₁₃_ₖ
            a₂₃_ₖ = k₁ * ϰ * (r₂[k] - r₃[k]) / ρ₂₃_ₖ
            a₁₂_ₖ₋₁ = k₂ * ϰ * (r₁_ₖ₋₁[k] - r₂_ₖ₋₁[k]) / ρ₁₂_ₖ₋₁
            a₁₃_ₖ₋₁ = k₂ * ϰ * (r₁_ₖ₋₁[k] - r₃_ₖ₋₁[k]) / ρ₁₃_ₖ₋₁
            a₂₃_ₖ₋₁ = k₂ * ϰ * (r₂_ₖ₋₁[k] - r₃_ₖ₋₁[k]) / ρ₂₃_ₖ₋₁

            a₂₁_ₖ = -a₁₂_ₖ
            a₃₁_ₖ = -a₁₃_ₖ
            a₃₂_ₖ = -a₂₃_ₖ
            a₂₁_ₖ₋₁ = -a₁₂_ₖ₋₁
            a₃₁_ₖ₋₁ = -a₁₃_ₖ₋₁
            a₃₂_ₖ₋₁ = -a₂₃_ₖ₋₁

            r₁_ₖ₋₁[k] = r₁[k]
            r₂_ₖ₋₁[k] = r₂[k]
            r₃_ₖ₋₁[k] = r₃[k]

            r₁[k] += k₁ * v₁[k] + k₂ * v₁_ₖ₋₁[k]
            r₂[k] += k₁ * v₂[k] + k₂ * v₂_ₖ₋₁[k]
            r₃[k] += k₁ * v₃[k] + k₂ * v₃_ₖ₋₁[k]

            v₁_ₖ₋₁[k] = v₁[k]
            v₂_ₖ₋₁[k] = v₂[k]
            v₃_ₖ₋₁[k] = v₃[k]

            v₁[k] += a₁₂_ₖ + a₁₃_ₖ + a₁₂_ₖ₋₁ + a₁₃_ₖ₋₁
            v₂[k] += a₂₁_ₖ + a₂₃_ₖ + a₂₁_ₖ₋₁ + a₂₃_ₖ₋₁
            v₃[k] += a₃₁_ₖ + a₃₂_ₖ + a₃₁_ₖ₋₁ + a₃₂_ₖ₋₁
        end
    end
    return r₁, r₂, r₃, v₁, v₂, v₃
end

# Integrate the three-body problem using the
# three-step Adams–Bashforth's method, return the
# values of position and velocity on the last step
function ab3(
    r₀₁::Vector{F},
    r₀₂::Vector{F},
    r₀₃::Vector{F},
    v₀₁::Vector{F},
    v₀₂::Vector{F},
    v₀₃::Vector{F},
    h::F,
    n::I,
)::Tuple{Vector{F},Vector{F},Vector{F},Vector{F},Vector{F},Vector{F}} where
{F <: AbstractFloat,I <: Unsigned}
    # Determine the length of the input vectors
    N = length(r₀₁)
    # Prepare the output vectors
    r₁ = copy(r₀₁)
    r₂ = copy(r₀₂)
    r₃ = copy(r₀₃)
    v₁ = copy(v₀₁)
    v₂ = copy(v₀₂)
    v₃ = copy(v₀₃)
    # Prepare buffers for previous values
    r₁_ₖ₋₁ = copy(r₀₁)
    r₂_ₖ₋₁ = copy(r₀₂)
    r₃_ₖ₋₁ = copy(r₀₃)
    v₁_ₖ₋₁ = copy(v₀₁)
    v₂_ₖ₋₁ = copy(v₀₂)
    v₃_ₖ₋₁ = copy(v₀₃)
    r₁_ₖ₋₂ = copy(r₀₁)
    r₂_ₖ₋₂ = copy(r₀₂)
    r₃_ₖ₋₂ = copy(r₀₃)
    v₁_ₖ₋₂ = copy(v₀₁)
    v₂_ₖ₋₂ = copy(v₀₂)
    v₃_ₖ₋₂ = copy(v₀₃)
    # Compute the second value of the solution
    # by using the one-step Euler's method
    r₁, r₂, r₃, v₁, v₂, v₃ = euler(r₀₁, r₀₂, r₀₃, v₀₁, v₀₂, v₀₃, h, UInt(1))
    # Compute the third value of the solution
    # by using the two-step Adams–Bashforth's method
    r₁_ₖ₋₁, r₂_ₖ₋₁, r₃_ₖ₋₁, v₁_ₖ₋₁, v₂_ₖ₋₁, v₃_ₖ₋₁ = r₁, r₂, r₃, v₁, v₂, v₃
    r₁, r₂, r₃, v₁, v₂, v₃ = ab2(r₀₁, r₀₂, r₀₃, v₀₁, v₀₂, v₀₃, h, UInt(2))
    # Define a couple of independent coefficients
    k₁ = 23 / 12 * h
    k₂ = -4 / 3 * h
    k₃ = 5 / 12 * h
    # Compute the rest in two steps
    for _ in 3:n
        ρ₁₂_ₖ = norm(r₁ - r₂)^3
        ρ₁₃_ₖ = norm(r₁ - r₃)^3
        ρ₂₃_ₖ = norm(r₂ - r₃)^3
        ρ₁₂_ₖ₋₁ = norm(r₁_ₖ₋₁ - r₂_ₖ₋₁)^3
        ρ₁₃_ₖ₋₁ = norm(r₁_ₖ₋₁ - r₃_ₖ₋₁)^3
        ρ₂₃_ₖ₋₁ = norm(r₂_ₖ₋₁ - r₃_ₖ₋₁)^3
        ρ₁₂_ₖ₋₂ = norm(r₁_ₖ₋₂ - r₂_ₖ₋₂)^3
        ρ₁₃_ₖ₋₂ = norm(r₁_ₖ₋₂ - r₃_ₖ₋₂)^3
        ρ₂₃_ₖ₋₂ = norm(r₂_ₖ₋₂ - r₃_ₖ₋₂)^3
        for k in 1:N
            a₁₂_ₖ = k₁ * ϰ * (r₁[k] - r₂[k]) / ρ₁₂_ₖ
            a₁₃_ₖ = k₁ * ϰ * (r₁[k] - r₃[k]) / ρ₁₃_ₖ
            a₂₃_ₖ = k₁ * ϰ * (r₂[k] - r₃[k]) / ρ₂₃_ₖ
            a₁₂_ₖ₋₁ = k₂ * ϰ * (r₁_ₖ₋₁[k] - r₂_ₖ₋₁[k]) / ρ₁₂_ₖ₋₁
            a₁₃_ₖ₋₁ = k₂ * ϰ * (r₁_ₖ₋₁[k] - r₃_ₖ₋₁[k]) / ρ₁₃_ₖ₋₁
            a₂₃_ₖ₋₁ = k₂ * ϰ * (r₂_ₖ₋₁[k] - r₃_ₖ₋₁[k]) / ρ₂₃_ₖ₋₁
            a₁₂_ₖ₋₂ = k₃ * ϰ * (r₁_ₖ₋₂[k] - r₂_ₖ₋₂[k]) / ρ₁₂_ₖ₋₂
            a₁₃_ₖ₋₂ = k₃ * ϰ * (r₁_ₖ₋₂[k] - r₃_ₖ₋₂[k]) / ρ₁₃_ₖ₋₂
            a₂₃_ₖ₋₂ = k₃ * ϰ * (r₂_ₖ₋₂[k] - r₃_ₖ₋₂[k]) / ρ₂₃_ₖ₋₂

            a₂₁_ₖ = -a₁₂_ₖ
            a₃₁_ₖ = -a₁₃_ₖ
            a₃₂_ₖ = -a₂₃_ₖ
            a₂₁_ₖ₋₁ = -a₁₂_ₖ₋₁
            a₃₁_ₖ₋₁ = -a₁₃_ₖ₋₁
            a₃₂_ₖ₋₁ = -a₂₃_ₖ₋₁
            a₂₁_ₖ₋₂ = -a₁₂_ₖ₋₂
            a₃₁_ₖ₋₂ = -a₁₃_ₖ₋₂
            a₃₂_ₖ₋₂ = -a₂₃_ₖ₋₂

            r₁_ₖ₋₂[k] = r₁_ₖ₋₁[k]
            r₂_ₖ₋₂[k] = r₂_ₖ₋₁[k]
            r₃_ₖ₋₂[k] = r₃_ₖ₋₁[k]
            r₁_ₖ₋₁[k] = r₁[k]
            r₂_ₖ₋₁[k] = r₂[k]
            r₃_ₖ₋₁[k] = r₃[k]

            r₁[k] += k₁ * v₁[k] + k₂ * v₁_ₖ₋₁[k] + k₃ * v₁_ₖ₋₂[k]
            r₂[k] += k₁ * v₂[k] + k₂ * v₂_ₖ₋₁[k] + k₃ * v₂_ₖ₋₂[k]
            r₃[k] += k₁ * v₃[k] + k₂ * v₃_ₖ₋₁[k] + k₃ * v₃_ₖ₋₂[k]

            v₁_ₖ₋₂[k] = v₁_ₖ₋₁[k]
            v₂_ₖ₋₂[k] = v₂_ₖ₋₁[k]
            v₃_ₖ₋₂[k] = v₃_ₖ₋₁[k]
            v₁_ₖ₋₁[k] = v₁[k]
            v₂_ₖ₋₁[k] = v₂[k]
            v₃_ₖ₋₁[k] = v₃[k]

            v₁[k] += a₁₂_ₖ + a₁₃_ₖ + a₁₂_ₖ₋₁ + a₁₃_ₖ₋₁ + a₁₂_ₖ₋₂ + a₁₃_ₖ₋₂
            v₂[k] += a₂₁_ₖ + a₂₃_ₖ + a₂₁_ₖ₋₁ + a₂₃_ₖ₋₁ + a₂₁_ₖ₋₂ + a₂₃_ₖ₋₂
            v₃[k] += a₃₁_ₖ + a₃₂_ₖ + a₃₁_ₖ₋₁ + a₃₂_ₖ₋₁ + a₃₁_ₖ₋₂ + a₃₂_ₖ₋₂
        end
    end
    return r₁, r₂, r₃, v₁, v₂, v₃
end

println(" "^4, "> Integrating the 3-body problem...")

# Define the initial values of the position and velocity
r₀₁ = [0.97, -0.2431]
r₀₂ = [-0.97, 0.2431]
r₀₃ = [0.0, 0.0]
v₀₁ = [0.4662, 0.4324]
v₀₂ = [0.4662, 0.4324]
v₀₃ = [-0.9324, -0.8647]

# Calculate total energy
E₀ = 1 / 2 * (norm(v₀₁)^2 + norm(v₀₂)^2 + norm(v₀₃)^2) - ϰ^2 *
    (1 / norm(r₀₁ - r₀₂) + 1 / norm(r₀₃ - r₀₂) + 1 / norm(r₀₃ - r₀₁))

# Print the initial values
println(
    '\n',
    " "^6, "r₀₁: ", r₀₁[1], " ", r₀₁[2], '\n',
    " "^6, "r₀₂: ", r₀₂[1], " ", r₀₂[2], '\n',
    " "^6, "r₀₃: ", r₀₃[1], " ", r₀₃[2], '\n',
    " "^6, "v₀₁: ", v₀₁[1], " ", v₀₁[2], '\n',
    " "^6, "v₀₂: ", v₀₂[1], " ", v₀₂[2], '\n',
    " "^6, "v₀₃: ", v₀₃[1], " ", v₀₃[2], '\n',
    " "^6, "E₀: ", E₀,
)

# Define the output directories
const data_dir = "$(@__DIR__)/../data/3-body"
const files = joinpath.(
    data_dir,
    [
        "euler.dat",
        "ab2.dat",
        "ab3.dat",
    ]
)

# Prepare a list of methods
const methods = (euler, ab2, ab3)

# Integrate equations of motion using all available
# methods and write the values of position and velocity
# on the last step
function integrate(h::Vector{F}, n::Vector{I}) where {F <: AbstractFloat,I <: Unsigned}
    # Open the data files
    io = open.(files, "a")

    # Prepare results vectors

    r₁ = [ Vector{Vector{F}}() for _ in 1:length(io) ]
    r₂ = [ Vector{Vector{F}}() for _ in 1:length(io) ]
    r₃ = [ Vector{Vector{F}}() for _ in 1:length(io) ]
    v₁ = [ Vector{Vector{F}}() for _ in 1:length(io) ]
    v₂ = [ Vector{Vector{F}}() for _ in 1:length(io) ]
    v₃ = [ Vector{Vector{F}}() for _ in 1:length(io) ]

    E = [ Vector{F}() for _ in 1:length(io) ]
    ΔE = [ Vector{F}() for _ in 1:length(io) ]

    # For each pair of parameters
    for i in eachindex(h)
        # Integrate and get the results of the last steps
        for j in eachindex(io)
            rᵢ₁, rᵢ₂, rᵢ₃, vᵢ₁, vᵢ₂, vᵢ₃ = methods[j](r₀₁, r₀₂, r₀₃, v₀₁, v₀₂, v₀₃, h[i], n[i])

            push!(r₁[j], rᵢ₁)
            push!(r₂[j], rᵢ₂)
            push!(r₃[j], rᵢ₃)
            push!(v₁[j], vᵢ₁)
            push!(v₂[j], vᵢ₂)
            push!(v₃[j], vᵢ₃)

            Eᵢ = 1 / 2 * (norm(vᵢ₁)^2 + norm(vᵢ₂)^2 + norm(vᵢ₃)^2) - ϰ^2 *
                (1 / norm(rᵢ₁ - rᵢ₂) + 1 / norm(rᵢ₃ - rᵢ₂) + 1 / norm(rᵢ₃ - rᵢ₁))
            ΔEᵢ = abs(E₀ - Eᵢ)

            push!(E[j], Eᵢ)
            push!(ΔE[j], ΔEᵢ)
        end
    end

    # Get the format string based on the index of the iteration
    function get_format(i)
        Printf.Format(
            "\$ 10^{-$(length("$(n[i])") - 1)} \$ " *
            "& \$ $(n[i] % 10 == 0 ? "10^$(length("$(n[i])") - 1)" : n[i]) \$ " *
            "& \$ %.14f \$ " *
            "& \$ %.14f \$ " *
            "\\\\\n"
        )
    end

    # Print the data in the files in the LaTeX format

    for j in eachindex(io)
        println(io[j], "# Positions #1")
        for i in eachindex(h)
            Printf.format(io[j], get_format(i), r₁[j][i][1], r₁[j][i][2])
        end
    end

    for j in eachindex(io)
        println(io[j], "\n# Positions #2")
        for i in eachindex(h)
            Printf.format(io[j], get_format(i), r₂[j][i][1], r₂[j][i][2])
        end
    end

    for j in eachindex(io)
        println(io[j], "\n# Positions #3")
        for i in eachindex(h)
            Printf.format(io[j], get_format(i), r₃[j][i][1], r₃[j][i][2])
        end
    end

    for j in eachindex(io)
        println(io[j], "\n# Velocities #1")
        for i in eachindex(h)
            Printf.format(io[j], get_format(i),  v₁[j][i][1], v₁[j][i][2])
        end
    end

    for j in eachindex(io)
        println(io[j], "\n# Velocities #2")
        for i in eachindex(h)
            Printf.format(io[j], get_format(i),  v₂[j][i][1], v₂[j][i][2])
        end
    end

    for j in eachindex(io)
        println(io[j], "\n# Velocities #3")
        for i in eachindex(h)
            Printf.format(io[j], get_format(i),  v₃[j][i][1], v₃[j][i][2])
        end
    end

    for j in eachindex(io)
        println(io[j], "\n# Total energy")
        for i in eachindex(h)
            Printf.format(io[j], get_format(i), E[j][i], ΔE[j][i])
        end
        println(io[j])
        close(io[j])
    end
end

# Truncate the previous results
open.(files; truncate=true)

# Integrate n = 2^m iterations
integrate(
    [1e-2, 1e-3, 1e-4, 1e-5, 1e-6, 1e-7],
    UInt.([ 10^i for i in 2:7 ]),
)

println()


# Chapter 1.1

# Visualizing a geometric series

## Julia code to generate a geometric series
using Plots:bar
p = 0.5
n = 1:10
X = p .^ n
bar(n, X, legend=false)

## compute N choose K
# Julia code to (N choose K) and K!
n = 10
k = 2
binomial(n, k)
factorial(k)

# Chapter 1.4

# Inner product of two vectors

## Julia code to perform and inner product
x = [1, 0, -1]
y = [3, 2, 0]
x'y

## Norm of a vector
# Julia code to compute the norm
using LinearAlgebra:norm
x = [1, 0, -1]
norm(x)

## Weighted norm of a vector
# Julia code to compute the weighted norm
W = [1. 2. 3.; 4. 5. 6.; 7. 8. 9.]
x = [2, -1, 1]
z = x'W*x

## System of linear equations
# Julia code to solve Xβ = y
X = [1. 3.; -2. 7.; 0. 1.]
y = [2., 1., 0.]
β = X\y


# Estimates the Klein consumption equation by GMM
using Econometrics, DelimitedFiles, Statistics
function KleinMoments(θ, y, x, z)
	e = y - x*θ
	m = e.*z
end

function main()
data = readdlm("klein.data")
# construct missing lags, and drop first row that has missing data
profits = data[:,3]
output = data[:,7]
data = [data lag(profits,1) lag(output,1)]
data = data[2:end,:]
n = size(data,1)
# define instruments
exogs = [1, 6, 8, 9, 10, 11, 12]
exogs = data[:,exogs]
exogs = [ones(n,1) exogs]
# CONSUMPTION
println("CONSUMPTION EQUATION")
# define variables in consumption equation
y = data[:,2]
profits = data[:,3]
lagprofits = data[:,11]
wp = data[:,4]
wg = data[:,8]
wages = wp + wg
# regressors in consumption equation
x = [profits lagprofits wages]
x = [ones(n,1) x]
# GMM estimation using CUE
θstart = x\y # ols start values
names = ["Constant", "Profits", "Profits-1", "Wages"]
moments = θ -> KleinMoments(θ, y, x, exogs)
# estimation results using CUE
gmmtitle = "Klein model 1 GMM example CUE"
gmmresults(moments, θstart, "",  gmmtitle, names)
nothing
end
main()


suml2loss(yhat, y) = sum((yhat .- y).^2)
sumabsloss(yhat, y) = sum(abs.(yhat .- y))
l2loss(yhat, y) = mean((yhat .- y).^2)
absloss(yhat, y) = mean(abs.(yhat .- y))
mape(yhat, y, eps = 1.0) =
    100.0 * mean(abs.(yhat .- (y .+ eps)) ./ abs.(y .+ eps))


using OrdinaryDiffEq
using DiffEqUncertainty
using DiffEqCallbacks
using Statistics
using Random
using Plots
using BenchmarkTools

function sir_ode!(du,u,p,t)
    (S,I,R) = u
    (β,c,γ) = p
    N = S+I+R
    @inbounds begin
        du[1] = -β*c*I/N*S
        du[2] = β*c*I/N*S - γ*I
        du[3] = γ*I
    end
    nothing
end;

function condition(u,t,integrator) # Event when event_f(u,t) == 0
  u[2]
end;

function affect!(integrator)
  integrator.u[2] = 0.0
end;

positive_cb = ContinuousCallback(condition,affect!);

tmax = 40.0
tspan = (0.0,tmax);

u0 = [990.0,10.0,0.0]; # S,I,R

p = [0.05,10.0,0.25]; # β,c,γ

Random.seed!(1234);

prob_ode = ODEProblem(sir_ode!, u0, tspan, p);

probint_cb_4th = AdaptiveProbIntsUncertainty(4);

num_samples = 100
ensemble_prob_ode = EnsembleProblem(prob_ode)

samples_ode = solve(ensemble_prob_ode,
                             ROS34PW3(),
                             trajectories=num_samples,
                             callback=CallbackSet(positive_cb,probint_cb_4th));

s20 = [s(20.0) for s in samples_ode]
[[mean([s[i] for s in s20]) for i in 1:3] [std([s[i] for s in s20]) for i in 1:3]]

p = plot(samples_ode[1],
     label=["S" "I" "R"],
     color=[:blue :red :green],
     xlabel="Time",
     ylabel="Number")
for i in 2:num_samples
    plot!(p,
          samples_ode[i],
          label="",
          color=[:blue :red :green])
end;

plot(p,yaxis=:log10,xlim=(15,20),ylim=(100,1000))

plot(p,yaxis=:log10,xlim=(35,40),ylim=(10,1000))

@benchmark solve(ensemble_prob_ode,
                 ROS34PW3(),
                 trajectories=100,
                 callback=CallbackSet(positive_cb,probint_cb_4th))


using Optim

struct HW
	m::Int64
    α::Float64
    β::Float64
    γ::Float64
    l0::Float64
    b0::Float64
    s0::Array{Float64, 1}

    HW(m::Int64) = new(m, 0, 0, 0, 0, 0, repeat([0], m))
	HW(m::Int64, α::Number, β::Number, γ::Number, l0::Number, b0::Number, s0::Array{Float64, 1}) =
		m == length(s0) ? new(m, Float64(α), Float64(β), Float64(γ), Float64(l0), Float64(b0),
		Float64.(s0)) : error("length(s0) must be equal to m")
end

function loss(model::HW, time_series)
	α = model.α
	β = model.β
	γ = model.γ
	l0 = model.l0
	b0 = model.b0
	m = model.m
	s = model.s0

	N = length(time_series)
	l_t = 0
	b_t = 0
	l_t_ = 0 #Variable to save l(t-1)
	b_t_ = 0 #Variable to save b(t-1)
	s_ = 0

	loss = 0

	for i in 0:(N - 1)
		if i == 0
			l_t = l0
			b_t = b0
		else
			l_t = (time_series[i] - s_) * α + (l_t_ + b_t_) * (1 - α)
			b_t = β * (l_t - l_t_) + (1 - β) * b_t_
		end
		l_t_ = l_t
		b_t_ = b_t
		s_ = s[i%m + 1]

		y_pred = l_t + b_t + s[i%m + 1]

		s[i%m + 1] = γ * (time_series[i + 1] - l_t_ - b_t_) + (1 - γ) * s[i%m + 1]

		loss += (time_series[i + 1] - y_pred)^2
	end

	return loss
end

function forecast(model::HW, time_serie, n_pred)
	α = model.α
	β = model.β
	γ = model.γ
	l0 = model.l0
	b0 = model.b0
	m = model.m
	s = model.s0

	N = length(time_serie)
	l_t = 0
	b_t = 0
	l_t_ = 0 #Variable to save l(t-1)
	b_t_ = 0 #Variable to save b(t-1)
	s_ = 0

	pred = []

	for i in 0:(N - 1)
		if i == 0
			l_t = l0
			b_t = b0
		else
			l_t = (time_serie[i] - s_) * α + (l_t_ + b_t_) * (1 - α)
			b_t = β * (l_t - l_t_) + (1 - β) * b_t_
		end
		l_t_ = l_t
		b_t_ = b_t
		s_ = s[i%m + 1]

		s[i%m + 1] = γ * (time_serie[i + 1] - l_t_ - b_t_) + (1 - γ) * s[i%m + 1]
	end

	l_t = (time_serie[end] - s_) * α + (l_t + b_t) * (1 - α)
	b_t = β * (l_t - l_t_) + (1 - β) * b_t_

	for i in N:(N+n_pred - 1)
		y_pred = l_t + b_t*(i-N+1) + s[i%m + 1]
		#The trend has to be added as many times as periods we want to forecast.
		push!(pred, y_pred)
	end

	return pred
end


using StaticArrays
using SymPy

"""
WENO reconstruction schemes following the lecture notes by Shu (1998) and
the review article by Shu (2009). WENO smoothness indicators are due to
Jiang & Shu (1996).

Shu (1998), "Essentially Non-Oscillatory and Weighted Essentially Non-Oscillatory
    Schemes for Hyperbolic Conservation Laws", In: Advanced Numerical Approximation
    of Nonlinear Hyperbolic Equations, edited by Cockburn et al., pp. 325–432.
    DOI: https://doi.org/10.1007/BFb0096355

Shu (2009) "High Order Weighted Essentially Nonoscillatory Schemes for Convection
    Dominated Problems", SIAM Review 51(1), pp. 82–126.
    DOI: https://doi.org/10.1137/070679065

Jiang & Shu (1996) "Efficient Implementation of Weighted ENO Schemes", Journal of
    Computational Physics 126, pp. 202–228.
    DOI: https://doi.org/10.1006/jcph.1996.0130
"""

#####
##### ENO reconstruction coefficients and WENO weights on a uniform grid
##### Equation (2.21) from the Shu (1998) lecture notes.
#####

"""
    UΠ(k, r, m, l)

Return the product in the numerator of equation (2.21) of Shu (1998) for the `l`th
term of the `m`th Lagrange basis polynomial of an ENO reconstruction scheme with
order `k` and left shift `r`.
"""
UΠ(k, r, l, m) = prod([r-q+1 for q in 0:k if q ∉ (m, l)])

"""
    U(k, r, m)

Return the numerator in equation (2.21) of Shu (1998) for the `m`th Lagrange basis
polynomial of an ENO reconstruction scheme with order `k` and left shift `r`.
"""
U(k, r, m) = sum([UΠ(k, r, m, l) for l in 0:k if l != m])

"""
    D(k, m)

Return the denominator in equation (2.21) of Shu (1998) for the `m`th Lagrange basis
polynomial of an ENO reconstruction scheme with order `k`.
"""
D(k, m) = prod([m - l for l in 0:k if l != m])

"""
    eno_coefficient(k, r)

Return the `j`th ENO coefficient used to reconstruct a value at the point x(i+½) with
order of accuracy `k` (stencil size) and left shift `r`.
"""
eno_coefficient(k, r, j) = sum([U(k, r, m)//D(k, m) for m in j+1:k])

"""
    eno_coefficients(k, r)

Return an array of ENO coefficients to reconstruct a value at the point x(i+½) with
order of accuracy `k` (stencil size) and left shift `r`.
"""
eno_coefficients(k, r) = [eno_coefficient(k, r, j) for j in 0:k-1]

"""
    eno_coefficients_matrix([FT=Rational], k)

Return a k×k static array containing ENO coefficients to reconstruct a value at the
point x(i+½) with order of accuracy `k` (stencil size) with element type `FT`. Note
that when combined these ENO interpolants produce a WENO scheme of order 2k-1.
"""
eno_coefficients_matrix(FT, k) =
    cat([eno_coefficients(k, r) for r in 0:k-1]..., dims=1) |> SMatrix{k,k,FT}

eno_coefficients_matrix(k) = eno_coefficients_matrix(Rational, k)

"""
    optimal_weno_weights([FT=Rational], k)

Return a static vector containing the optimal weights that can be used to weigh ENO
reconstruction schemes of order `k` to produce a WENO scheme of order 2k-1.
"""
function optimal_weno_weights(FT, k)
    C = zeros(Rational, 2k-1, k)
    b = eno_coefficients(2k-1, k-1)

    for n in 0:k-1
        C[n+1:n+k, n+1] .= eno_coefficients(k, k-1-n)
    end

    γ = C \ b
    γ = rationalize.(γ, tol=√eps(Float64)) |> reverse
    return SVector{k,FT}(γ)
end

optimal_weno_weights(k) = optimal_weno_weights(Rational, k)

#####
##### Jiang & Shu (1996) WENO smoothness indicators β
##### See equation (2.61) of Shu (1998).
#####

x(j) = j  # Assume uniform grid for now.

"""
    ℓ(ξ, k, r, j)

Return a symbolic expression for the `j`th Lagrange basis polynomial ℓ(ξ) for an
ENO reconstruction scheme of order k and left shift `r` with field values `ϕ`.
"""
ℓ(ξ, k, r, j) = prod((ξ - x(m-r)) / (x(j-r) - x(m-r)) for m in 0:k-1 if m != j)

"""
    L(ξ, k, r, ϕ)

Return a symbolic expression for the interpolating Lagrange polynomial p(ξ) for an
ENO reconstruction scheme of order k and left shift `r` with field values `ϕ`.
"""
p(ξ, k, r, ϕ) = sum(ℓ(ξ, k, r, j) * ϕ[j+1] for j in 0:k-1)

"""
    β(k, r, ϕ)

Return a symbolic expression for the smoothness indicator β for a WENO reconstruction
scheme with stencils of size `k` and left shift `r` (WENO scheme of order 2k-1). The
field values are represented by the symbols in `ϕ` which should have length k.
"""
function β(k, r, ϕ)
    @vars ξ
    return sum(integrate(diff(p(ξ, k, r, ϕ), ξ, l)^2, (ξ, Sym(-1//2), Sym(1//2))) for l in 1:k-1)
end

"""
    subscript(n)

Convert the integer `n` to a subscript in the form of a unicode string. Note
that `0x2080` is the unicode encoding for the subscript 0.
"""
subscript(n) = join(Char(0x2080 + parse(Int, d)) for d in string(n))

subscript_sign(n) = n > 0 ? "₊" : n < 0 ? "₋" : ""

subscript_index(n) = n == 0 ? "" : subscript_sign(n) * subscript(abs(n))

"""
    β_coefficients([FT=Rational], k)

Return a k×k×k static array containing the WENO smoothness indicator coefficients
described by Jiang & Shu (1998) for a WENO reconstruction with stencils of size `k`
(WENO scheme of order 2k-1) with element type `FT`.

The `B[m, n, r]` coefficient corresponds to the coefficient of the
`ϕ[r-k+m+1] * ϕ[r-k+n+1]` term where r-k+1 <= m, n <= r and `r` is the left shift of
the ENO interpolants.
"""
function β_coefficients(FT, k)
    B = zeros(Float64, k, k, k)

    for r in 0:k-1
        ϕ = [Sym("ϕᵢ" * subscript_index(n)) for n in r:-1:r-k+1]
        β_symbolic = β(k, r, ϕ) |> expand

        for m in 1:k, n in 1:k
            B[m, n, r+1] = β_symbolic.coeff(ϕ[m] * ϕ[n])
        end
    end

    B = rationalize.(B, tol=√eps(Float64))
    return SArray{Tuple{k,k,k},FT}(B)
end

β_coefficients(k) = β_coefficients(Rational, k)


# this uses catalyst reaction networks -- probably not optimised for performance
# function calculate_steady_state_of_ode(repression_threshold,
#                                        hill_coefficient,
#                                        mRNA_degradation_rate,
#                                        protein_degradation_rate,
#                                        basal_transcription_rate,
#                                        translation_rate)
#    # define Hes5 reaction network
#    hes5_network = @reaction_network begin
#       hillr(protein,basal_transcription_rate,repression_threshold,hill_coefficient), ∅ --> mRNA
#       mRNA_degradation_rate, mRNA → ∅
#       protein_degradation_rate, protein → ∅
#       translation_rate, mRNA → mRNA + protein
#    end repression_threshold hill_coefficient mRNA_degradation_rate protein_degradation_rate basal_transcription_rate translation_rate
#
#    # parameters [α,K,n,δ,γ,β,μ]
#    parameters = (repression_threshold,
#                  hill_coefficient,
#                  mRNA_degradation_rate,
#                  protein_degradation_rate,
#                  basal_transcription_rate,
#                  translation_rate)
#    # initial condition [m,p]
#    u₀ = [40.,50000.]
#    tspan = (0.0,100.0)
#    # create the ODEProblem we want to solve
#    prob = ODEProblem(hes5_network, u₀,tspan,parameters)
#    sol = solve(prob)
#
# end #function

"""
Calculate the steady state of the Hes5 ODE system, for a specific set of parameters.

# Arguments

- `repression_threshold::Float64`

- `hill_coefficient::Float64`

- `mRNA_degradation_rate::Float64`

- `protein_degradation_rate::Float64`

- `basal_transcription_rate::Float64`

- `translation_rate::Float64`

# Returns

- `steady_state_solution::Array{Float64,1}`: A 2-element array, giving the steady state for the mRNA and protein respectively.
"""
function calculate_steady_state_of_ode(repression_threshold,
                                       hill_coefficient,
                                       mRNA_degradation_rate,
                                       protein_degradation_rate,
                                       basal_transcription_rate,
                                       translation_rate)

   hill_function(protein,repression_threshold,hill_coefficient) = 1/(1 + (protein/repression_threshold)^hill_coefficient)

   function hes_ode!(du,u,p,t)
      du[1] = p[5]*hill_function(u[2],p[1],p[2]) - p[3]*u[1]
      du[2] = p[6]*u[1] - p[4]*u[2]
   end #function

   parameters = (repression_threshold,
                 hill_coefficient,
                 mRNA_degradation_rate,
                 protein_degradation_rate,
                 basal_transcription_rate,
                 translation_rate)

   # initial condition [m,p]
   initial_guess = [40.,50000.]
   # create the SteadyStateProblem we want to solve
   prob = SteadyStateProblem(hes_ode!, initial_guess, parameters)
   # solve the problem
   steady_state_solution = solve(prob,SSRootfind())
end #function


#=
Hidden layer activation functions and gradients.
Options are:
:logistic, :tanh, :rectifier
=#

################################################################################
#
# Activation Functions and Derivatives
#
################################################################################

########################################
# Logistic
########################################

# logistic function of element z
# f(z) = 1 / (1 + exp(-z))
function logistic{F<:AbstractFloat}(z::F)
    one(z) / (one(z) + exp(-z))
end

# logistic function of array Z
# replacing array A in-place
# A[i] = f(Z[i]) = 1 / (1 + exp(-Z[i])) ∀ i
function logistic!{F<:AbstractFloat}(A::Array{F}, Z::Array{F})
    size(A) == size(Z) || throw(DimensionMismatch("Inconsistent array sizes"))
    @threads for i in eachindex(Z)
        @inbounds A[i] = logistic(Z[i])
    end
end

# logistic function derivative of element z
# df/dz = f(z) (1 - f(z))
function logistic_deriv{F<:AbstractFloat}(z::F)
    s = logistic(z)
    s * (F(1) - s)
end

# logistic function derivative of array Z
# ∂f/∂Z[i] = f(Z[i]) (1 - f(Z[i]))
# replacing dest in-place with Hademaard product of dest and derivative i.e.
# dest[i] = dest[i] * ∂f/∂Z[i]
function logistic_deriv!{F<:AbstractFloat}(dest::Array{F}, Z::Array{F})
    size(dest) == size(Z) || throw(DimensionMismatch("Inconsistent array sizes"))
    @threads for i in eachindex(Z)
        @inbounds dest[i] *= logistic_deriv(Z[i])
    end
end

########################################
# Tanh
########################################

# Elementwise tanh function of array Z
# replacing array A in-place
function tanh!{F<:AbstractFloat}(A::Array{F}, Z::Array{F})
    size(A) == size(Z) || throw(DimensionMismatch("Inconsistent array sizes"))
    @threads for i in eachindex(Z)
        @inbounds A[i] = tanh(Z[i])
    end
end

# Elementwise tanh function derivative of array Z
# ∂f/∂Z[i] = 1 - f(Z[i])^2
# replacing dest in-place with Hademaard product of dest and derivative i.e.
# dest[i] = dest[i] * ∂f/∂Z[i]
function tanh_deriv!{F<:AbstractFloat}(dest::Array{F}, Z::Array{F})
    size(dest) == size(Z) || throw(DimensionMismatch("Inconsistent array sizes"))
    @threads for i in eachindex(Z)
        @inbounds dest[i] *= (F(1) - tanh(Z[i])^2)
    end
end

########################################
# Rectified linear units
########################################

# Elementwise rectifier function of array Z
# replacing array A in-place
function relu!{F<:AbstractFloat}(A::Array{F}, Z::Array{F})
    size(A) == size(Z) || throw(DimensionMismatch("Inconsistent array sizes"))
    @threads for i in eachindex(Z)
        @inbounds A[i] = max(F(0), Z[i])
    end
end

# Elementwise rectifier function derivative of array Z
# ∂f/∂Z[i] = 0 (if Z[i] <= 0) or 1 (if Z[i] > 0)
# replacing dest in-place with Hademaard product of dest and derivative i.e.
# dest[i] = dest[i] * ∂f/∂Z[i]
function relu_deriv!{F<:AbstractFloat}(dest::Array{F}, Z::Array{F})
    size(dest) == size(Z) || throw(DimensionMismatch("Inconsistent array sizes"))
    for i in eachindex(Z)
        @inbounds (Z[i] <= F(0)) && (dest[i] = F(0))
        #@inbounds dest[i] *= ifelse(Z[i] > F(0), F(1), F(0))
    end
end

########################################
# Softmax
########################################

# Softmax function of matrix Z replacing matrix A in-place
# Applies function to each column of Z i.e. each column will sum to 1
function softmax!{F<:AbstractFloat}(A::Matrix{F}, Z::Matrix{F}, u::Matrix{F})
    #size(A) == size(Z) || throw(DimensionMismatch("Inconsistent array sizes"))
    maximum!(u, Z)
    for j in indices(Z, 2)
        s = F(0)
        for i in indices(Z, 1)
            @inbounds s += (A[i, j] = exp(Z[i, j] - u[j]))
        end
        for i in indices(Z, 1)
            @inbounds A[i, j] /= s
        end
    end
end

################################################################################
#
# Activation Dictionaries
#
################################################################################

activation_dict = Dict{Symbol, Function}(:logistic => logistic!,
                                         :tanh => tanh!,
                                         :relu => relu!)

activation_deriv_dict = Dict{Symbol, Function}(:logistic => logistic_deriv!,
                                               :tanh => tanh_deriv!,
                                               :relu => relu_deriv!)