"""`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


function nlp_cvx_102_013(
    optimizer,
    objective_tol,
    primal_tol,
    dual_tol,
    termination_target = TERMINATION_TARGET_LOCAL,
    primal_target = PRIMAL_TARGET_LOCAL,
)
    model = Model(optimizer)

    @variable(model, x)
    @variable(model, y)

    @objective(model, Min, x^2 + y^2)
    @NLconstraint(model, x^2 + y^2 <= 1.0)
    @constraint(model, x + y >= 1.2)

    optimize!(model)

    check_status(model, FEASIBLE_PROBLEM, termination_target, primal_target)
    check_objective(model, 0.72, tol = objective_tol)
    return check_solution([x, y], [0.6, 0.6], tol = primal_tol)
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


using ForwardDiff

#thresh = 1e-12

@testset "Derivative checks" begin

    ###################################
    # Test Jacobians
    ###################################

    @testset "Jacobian checks" begin

        # Quaternion to rotation matrix
        @testset "Jacobian (QuatRotation -> RotMatrix)" begin
            for i = 1:10    # do some repeats
                q = rand(QuatRotation)  # a random quaternion

                # test jacobian to a Rotation matrix
                R_jac = Rotations.jacobian(RotMatrix, q)
                FD_jac = ForwardDiff.jacobian(x -> SVector{9}(QuatRotation(x[1], x[2], x[3], x[4])),
                                              Rotations.params(q))

                # compare
                @test FD_jac ≈ R_jac
            end
        end

        # MRP to UnitQuternion
        @testset "Jacobian (MRP -> QuatRotation)" begin
            for i = 1:10    # do some repeats
                p = rand(MRP)  # a random MRP

                # test jacobian to a Rotation matrix
                R_jac = Rotations.jacobian(QuatRotation, p)
                FD_jac = ForwardDiff.jacobian(x -> (q = QuatRotation(MRP(x[1],x[2],x[3]));
                                                    Rotations.params(q)),
                                              SVector(p.x, p.y, p.z))

                # compare
                @test FD_jac ≈ R_jac
            end
        end

        @testset "Jacobian (MRP -> QuatRotation) [Corner Cases]" begin
            for p = [MRP(1.0, 0.0, 0.0), MRP(0.0, 1.0, 0.0), MRP(0.0, 0.0, 1.0)]
                # test jacobian to a Rotation matrix
                R_jac = Rotations.jacobian(QuatRotation, p)
                FD_jac = ForwardDiff.jacobian(x -> (q = QuatRotation(MRP(x[1],x[2],x[3]));
                                                    Rotations.params(q)),
                                              SVector(p.x, p.y, p.z))

                # compare
                @test FD_jac ≈ R_jac
            end
        end

        # MRP to QuatRotation
        @testset "Jacobian (QuatRotation -> MRP)" begin
            for i = 1:10    # do some repeats
                q = rand(QuatRotation)  # a random QuatRotation

                # test jacobian to a Rotation matrix
                R_jac = Rotations.jacobian(MRP, q)
                FD_jac = ForwardDiff.jacobian(x -> (p = MRP(QuatRotation(x[1], x[2], x[3], x[4]));
                                                    SVector(p.x, p.y, p.z)),
                                              Rotations.params(q))

                # compare
                @test FD_jac ≈ R_jac
            end
        end

        # MRP to rotation matrix
        @testset "Jacobian (MRP -> RotMatrix)" begin
            for i = 1:10    # do some repeats
                p = rand(MRP)  # a random MRP

                # test jacobian to a Rotation matrix
                R_jac = Rotations.jacobian(RotMatrix, p)
                FD_jac = ForwardDiff.jacobian(x -> SVector{9}(MRP(x[1], x[2], x[3])),
                                              SVector(p.x, p.y, p.z))

                # compare
                @test FD_jac ≈ R_jac
            end
        end

        ##############################
        # Jacobians for rotating stuff
        ##############################
#=
        # Quaternion multiplication
        @testset "Jacobian (Quaternion muliplication w.r.t. the right quaternion)" begin

            for i = 1:10    # do some repeats

                ql = quatrand()    # a random quaternion (should work for non-unit quaternions)
                qr = quatrand()    # a random quaternion (should work for non-unit quaternions)

                # transformation function
                J(X) = vec(ql * Quaternion(X...))

                # test jacobian to a Rotation matrix
                R_jac = Rotations.jacobian(ql, qr)
                FD_jac = ForwardDiff.jacobian(J, vec(qr))

                # compare
                @test all(abs(Matrix(R_jac) - FD_jac) .< thresh)

            end
        end

        @testset "Jacobian (Quaternion muliplication w.r.t. the left quaternion)" begin

            for i = 1:10    # do some repeats

                ql = quatrand()    # a random quaternion (should work for non-unit quaternions)
                qr = quatrand()    # a random quaternion (should work for non-unit quaternions)

                # transformation function
                J(X) = vec(Quaternion(X...) * qr)

                # test jacobian to a Rotation matrix
                R_jac = Rotations.jacobian(ql, qr, Val{false})
                FD_jac = ForwardDiff.jacobian(J, vec(ql))

                # compare
                @test all(abs(Matrix(R_jac) - FD_jac) .< thresh)

            end
        end =#

        # rotate a point by a RotMatrix
        @testset "Jacobian (RotMatrix rotation)" begin
            for i = 1:10    # do some repeats
                r = rand(RotMatrix{3,Float64})    # a random quaternion
                v = randn(SVector{3,Float64})

                # test jacobian to a Rotation matrix
                R_jac = Rotations.jacobian(r, v)
                FD_jac = ForwardDiff.jacobian(x -> RotMatrix(x)*v,
                                              SVector(r))

                # compare
                @test FD_jac ≈ R_jac
            end
        end

        # rotate a point by a quaternion
        @testset "Jacobian (QuatRotation rotation)" begin
            for i = 1:10    # do some repeats
                q = rand(QuatRotation)    # a random quaternion
                v = randn(SVector{3,Float64})

                # test jacobian to a Rotation matrix
                R_jac = Rotations.jacobian(q, v)
                FD_jac = ForwardDiff.jacobian(x -> QuatRotation(x[1], x[2], x[3], x[4])*v,
                                              Rotations.params(q))

                # compare
                @test FD_jac ≈ R_jac
            end
        end

        # rotate a point by a MRP
        @testset "Jacobian (MRP rotation)" begin
            for i = 1:10    # do some repeats
                p = rand(MRP)    # a random quaternion
                v = randn(SVector{3,Float64})

                # test jacobian to a Rotation matrix
                R_jac = Rotations.jacobian(p, v)
                FD_jac = ForwardDiff.jacobian(x -> MRP(x[1], x[2], x[3])*v,
                                              SVector(p.x, p.y, p.z))

                # compare
                @test FD_jac ≈ R_jac
            end
        end
#=
        # rotate a point by an MRP
        @testset "Jacobian (MRP rotation)" begin

            for i = 1:10    # do some repeats

                spq = SpQuat(nquatrand())    # a random quaternion
                X = randn(Vec{3,Float64})

                # transformation function
                J(spq) = Vector(rotate(SpQuat(spq...), X))

                # test jacobian to a Rotation matrix
                R_jac = Rotations.jacobian(spq, X)
                FD_jac = ForwardDiff.jacobian(J, vec(spq))

                # compare
                @test all(abs(Matrix(R_jac) - FD_jac) .< thresh)

            end
        end
=#

    end

    ###################################
    # Test Hessians
    ###################################
    #=
    @testset "Hessian checks" begin

        # SpQuat to Quaternion
        @testset "Hessian (SpQuat -> Quaternion)" begin

            for i = 1:10    # do some repeats

                q = nquatrand()                           # a random quaternion
                spq = Rotations.quat_to_spquat_naive(q)   # I want to test the full domain of SpQuats, not just the one with ||.|| < 1

                # test jacobian to a Rotation matrix
                R_hess = Rotations.hessian(Quaternion, spq)
                for d = 1:4
                    H(X) = vec(Quaternion(SpQuat(X...)))[d]   # transformation function
                    FD_hess = ForwardDiff.hessian(H, vec(spq))

                    # compare
                    @test all(abs(Matrix(R_hess[d]) - FD_hess) .< thresh)
                end
            end
        end

        # Quaternion to SpQuat
        @testset "Hessian (Quaternion -> SpQuat)" begin

            for i = 1:10    # do some repeats

                q = nquatrand()                           # a random quaternion

                # test jacobian to a Rotation matrix
                R_hess = Rotations.hessian(SpQuat, q)
                for d = 1:3
                    H(X) = vec(SpQuat(Quaternion(X...)))[d]   # transformation function
                    FD_hess = ForwardDiff.hessian(H, vec(q))

                    # compare
                    @test all(abs(Matrix(R_hess[d]) - FD_hess) .< thresh)
                end
            end
        end

        @testset "Hessian (Quaternion rotation)" begin

            for i = 1:10   # do some repeats

                q = nquatrand()    # a random quaternion
                X = randn(Vec{3,Float64})

                # transformation function
                R_hess = Rotations.hessian(q, X)

                for d = 1:3
                    H(q) = Vector(rotate(Quaternion(q...), X))[d]

                    # test
                    FD_hess = ForwardDiff.hessian(H, vec(q))

                    # compare
                    @test all(abs(Matrix(R_hess[d]) - FD_hess) .< thresh)
                end
            end
        end

        @testset "Hessian (SpQuat rotation)" begin

            for i = 1:10    # do some repeats

                q = nquatrand()    # a random quaternion
                spq = SpQuat(q)
                X = randn(Vec{3,Float64})

                # transformation function
                R_hess = Rotations.hessian(spq, X)

                for d = 1:3
                    H(spq) = Vector(rotate(SpQuat(spq...), X))[d]

                    # test
                    FD_hess = ForwardDiff.hessian(H, vec(spq))

                    # compare
                    @test all(abs(Matrix(R_hess[d]) - FD_hess) .< thresh)
                end
            end
        end
    end =#
end


# tests for library routines

import SPTK: agexp, gexp, glog, mseq, theq!, theq, toeplitz!, toeplitz

function test_agexp()
    println("-- test_agexp")
    @test_approx_eq agexp(1, 1, 1) 5.0
    @test_approx_eq agexp(1, 2, 3) 18.0
    @test_approx_eq agexp(2, 3, 5) 12.206555615733702
end

function test_gexp()
    println("-- test_gexp")
    @test_approx_eq gexp(1, 1) 2.0
    @test_approx_eq gexp(2, 4) 3.0
end

function test_glog()
    println("-- test_glog")
    @test_approx_eq glog(1, 2) 1.0
    @test_approx_eq glog(2, 3) 4.0
end

function test_mseq()
    println("-- test_mseq")
    for i = 1:1000 @test !isnan(mseq()) end
end

function fill_toeplitz!{T}(A::AbstractMatrix{T}, t::AbstractVector{T})
    n = length(t)
    for j=1:n, i=1:n
        A[i,j] = i-j+1 >= 1 ? t[i-j+1] : t[j-i+1]
    end
    A
end

function fill_hankel!{T}(A::AbstractMatrix{T}, h::AbstractVector{T})
    n = length(h)>>1 + 1
    for j=1:n, i=1:n
        A[i,j] = h[i+j-1]
    end
    A
end

# (T + H)a = b
function test_theq()
    println("-- test_theq")
    srand(98765)

    n = 2
    # toeplitz elements
    t = [1.0, -1.0]
    T = zeros(2,2)
    fill_toeplitz!(T, t)

    # hankel elements
    h = zeros(2n-1)
    h[1] = 1.0
    h[3] = 1.0
    H = zeros(2,2)
    fill_hankel!(H, h)

    b = ones(n)
    a = zeros(n)

    # solve (T + H)a = b
    theq!(a, t, h, b)
    @test_approx_eq a ones(n)
    @test_approx_eq a theq(t, h, b)
    @test_approx_eq a (T + H) \ b

    # fail to solve toeplitz plus hankel matrix system
    for m in [3, 5, 10]
        a = ones(m)
        t = ones(a)
        h = ones(2length(a)-1)
        b = ones(a)
        @test_throws ErrorException theq!(a, t, h, b)
    end

    a = ones(5)
    t = ones(a)
    h = ones(2length(a)-1)
    b = ones(a)

    @test_throws ArgumentError theq!(a, t, h, b, min_det=-1.0)
    @test_throws DimensionMismatch theq!(a, t, h, ones(length(b)-1))
    @test_throws DimensionMismatch theq!(a, t, ones(length(h)-1), b)
    @test_throws DimensionMismatch theq!(a, t, ones(length(h)-1), b)
    @test_throws DimensionMismatch theq!(a, ones(length(t)-1), h, b)
    @test_throws DimensionMismatch theq!(ones(length(a)-1), t, h, b)
end

# Ta = b
function test_toeplitz()
    println("-- test_toeplitz")
    srand(98765)

    n = 2
    # toeplitz elements
    t = [2.0, -1.0]
    T = zeros(2,2)
    fill_toeplitz!(T, t)
    b = [1.0, 1.0]
    b = ones(n)
    a = zeros(n)

    # solve Ta = b
    toeplitz!(a, t, b)
    @test_approx_eq a ones(n)
    @test_approx_eq a toeplitz(t, b)
    @test_approx_eq a T \ b

    # fail to solve toeplitz set of linear equations
    for m in [3, 5, 10]
        a = ones(m)
        t = ones(a)
        b = ones(a)
        @test_throws ErrorException toeplitz!(a, t, b)
    end

    a = ones(5)
    t = ones(a)
    b = ones(a)
    @test_throws ArgumentError toeplitz!(a, t, b, eps=-1.0)
    @test_throws DimensionMismatch toeplitz!(a, t, ones(length(b)-1))
    @test_throws DimensionMismatch toeplitz!(a, ones(length(t)-1), b)
    @test_throws DimensionMismatch toeplitz!(ones(length(a)-1), t, b)
end

println("test library routines")
test_agexp()
test_gexp()
test_glog()
test_mseq()
test_theq()
test_toeplitz()


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)


using Test
using Random
using UnitTestDesign

### total_combinations
tc_trials = [
    [[2, 2, 2], 2, 12],
    [[2, 3, 2], 2, 16],
    [[2, 3, 2, 2], 2, 6+4+4+6+6+4],
    [[2, 3, 2, 2], 3, 12+12+8+12]
]
for trial in tc_trials
    res = UnitTestDesign.total_combinations(trial[1], trial[2])
    @test res == trial[3]
end

### combination_number

for cn_a in 1:6
    for cn_b in 1:cn_a
        @test UnitTestDesign.combination_number(cn_a, cn_b) == factorial(cn_a) ÷ (factorial(cn_a - cn_b) * factorial(cn_b))
    end
end

### next_multiplicative
arity0 = [3, 7, 9, 4]
vv0 = copy(arity0)
UnitTestDesign.next_multiplicative!(vv0, arity0)
@test vv0 == [1, 1, 1, 1]
vv = [1,1,1]
nmarity = [2,3,2]
for i in 1:prod(nmarity)
    UnitTestDesign.next_multiplicative!(vv, nmarity)
end
@test vv == [1, 1, 1]

### all_combinations
rng = MersenneTwister(9237425)
for ac_trial_idx in 1:5
    n_way = [2, 3, 4][rand(rng, 1:3)]
    param_cnt = rand(rng, (n_way + 1):(n_way + 3))
    arity = rand(rng, 2:4, param_cnt)
    coverage = UnitTestDesign.all_combinations(arity, n_way)
    # It has the right column dimnsion.
    @test size(coverage, 1) == length(arity)
    # Every combination is nonzero.
    @test sum(sum(coverage, dims = 1) == 0) == 0
    # The total number of combinations agrees with expectations.
    @test UnitTestDesign.total_combinations(arity, n_way) == size(coverage, 2)
    # Generate some random combinations and check that they are in there.
    for comb_idx in 1:100
        comb = [rand(rng, 1:arity[cj]) for cj in 1:param_cnt]
        comb[randperm(rng, param_cnt)[1:(param_cnt - n_way)]] .= 0
        @test sum(comb .!= 0) == n_way
        found = false
        for sidx in 1:size(coverage, 2)
            if coverage[:, sidx] == comb
                found = true
            end
        end
        @test found
    end
end

### one_parameter_combinations(arity, n_way)

minimal = UnitTestDesign.one_parameter_combinations([2, 3], 1)
@test minimal == [0 1; 0 2; 0 3]'

paired = UnitTestDesign.one_parameter_combinations([2, 3], 2)
@test paired == [1 1; 2 1; 1 2; 2 2; 1 3; 2 3]'

singled = UnitTestDesign.one_parameter_combinations([2, 2, 3], 2)
@test singled == [
    1 0 1; 2 0 1; 0 1 1; 0 2 1;
    1 0 2; 2 0 2; 0 1 2; 0 2 2;
    1 0 3; 2 0 3; 0 1 3; 0 2 3
    ]'

again = UnitTestDesign.one_parameter_combinations([2, 2, 3], 3)
@test again == [
    1 1 1; 1 2 1; 2 1 1; 2 2 1;
    1 1 2; 1 2 2; 2 1 2; 2 2 2;
    1 1 3; 1 2 3; 2 1 3; 2 2 3
    ]'


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 DrWatson
@quickactivate "StatReth"

# %%
using CSV
using DataFrames
using Turing
using Plots

include(srcdir("quap.jl"))
include(srcdir("tools.jl"))

# %%
# Load the data, don't forget to drop everything missing.
d = DataFrame(CSV.File(datadir("exp_raw/cherry_blossoms.csv"), missingstring = "NA"))
dropmissing!(d, [:temp, :doy])

# scatter(d.temp, d.doy, alpha = 0.3)

# %%
# First, a pure linear model.
flin(temp, α, β) = α .+ β .* temp

@model function cherry_linear(temp, doy)
    # β: With a higher temperature I expect the doy to be earlier. Maybe 1 month per 10°C?
    α ~ Normal(100, 10)
    β ~ Normal(-3, 2)
    σ ~ Exponential(1)
    μ = flin(temp, α, β)
    doy ~ MvNormal(μ, σ)
end

m = cherry_linear(d.temp, d.doy)

# %%
# Sample from the prior. Looks pretty reasonable.
prior = sample(m, Prior(), 100) |> DataFrame

scatter(d.year,  d.doy, alpha = 0.3)
for r in eachrow(prior)
    p = flin(d.temp, r.α, r.β)
    plot!(d.year, p, alpha = 0.2)
end
plot!(legend = false, ylims = (0, 200))

# %%
# Fit and show the fit.
q = quap(m)
post = DataFrame(rand(q.distr, 1000)', ["α", "β", "σ"])

mu = flin(d.temp, post.α', post.β')
mu_m, mu_lower, mu_upper = meanlowerupper(mu)
sim = rand.(Normal.(mu, post.σ'))
sim_m, sim_lower, sim_upper = meanlowerupper(sim)

scatter(d.year, d.doy, alpha = 0.3, legend = false)
plot!(d.year, mu_m, ribbon = (mu_m .- mu_lower, mu_upper .- mu_m))
plot!(d.year, sim_lower, fillrange = sim_upper, alpha = 0.3, linealpha = 0.0, c = 2)

# %%
# Second, a linear model enhanced with extra wiggle.
using BSplines: BSplineBasis, basismatrix

num_knots = 15
knot_list = quantile(d.year, range(0, 1, length = num_knots))

basis = BSplineBasis(4, knot_list)
B = basismatrix(basis, d.year)

# %%
# Fit the model.
@model function spline(temp, doy, B = B)
    α ~ Normal(100, 10)
    β ~ Normal(-3, 2)
    w ~ filldist(Normal(0, 10), size(B, 2))
    σ ~ Exponential(1)
    μ = flin(temp, α, β) .+ B * w
    doy ~ MvNormal(μ, σ)
end

sample(spline(d.temp, d.doy), NUTS(), 1000)

qspl = quap(spline(d.temp, d.doy))

w_str = ["w[$i]" for i in 1:length(basis)]
post = DataFrame(rand(qspl.distr, 1000)', ["α"; "β"; w_str; "σ"])

smu = fspl(d.temp, post.α', post.β', Array(post[!, w_str])', B)
smu_m, smu_lower, smu_upper = meanlowerupper(smu)

# %%
# Display both solutions together.
scatter(d.year, d.doy, alpha = 0.3, label = "data")
plot!(d.year, smu_m, ribbon = (smu_m .- smu_lower, smu_upper .- smu_m), label = "linear")
plot!(d.year, mu_m, ribbon = (mu_m .- mu_lower, mu_upper .- mu_m), label = "linear+spline")


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


using LinearAlgebra, IntervalSets, UnPack
import ClimaCore: Domains, Topologies, Meshes, Spaces, Geometry

using Test

@testset "Surface area of a sphere" begin
    FT = Float64
    radius = FT(3)
    ne = 4
    Nq = 4
    domain = Domains.SphereDomain(radius)
    mesh = Meshes.EquiangularCubedSphere(domain, ne)
    grid_topology = Topologies.Topology2D(mesh)
    quad = Spaces.Quadratures.GLL{Nq}()
    space = Spaces.SpectralElementSpace2D(grid_topology, quad)

    @test sum(ones(space)) ≈ 4pi * radius^2 rtol = 1e-3
end

@testset "Volume of a spherical shell" begin
    FT = Float64
    radius = FT(128)
    zlim = (0, 1)
    helem = 4
    zelem = 10
    Nq = 4

    vertdomain = Domains.IntervalDomain(
        Geometry.ZPoint{FT}(zlim[1]),
        Geometry.ZPoint{FT}(zlim[2]);
        boundary_tags = (:bottom, :top),
    )
    vertmesh = Meshes.IntervalMesh(vertdomain, nelems = zelem)
    vert_center_space = Spaces.CenterFiniteDifferenceSpace(vertmesh)

    horzdomain = Domains.SphereDomain(radius)
    horzmesh = Meshes.EquiangularCubedSphere(horzdomain, helem)
    horztopology = Topologies.Topology2D(horzmesh)
    quad = Spaces.Quadratures.GLL{Nq}()
    horzspace = Spaces.SpectralElementSpace2D(horztopology, quad)

    hv_center_space =
        Spaces.ExtrudedFiniteDifferenceSpace(horzspace, vert_center_space)

    # "shallow atmosphere" spherical shell: volume = surface area * height
    @test sum(ones(hv_center_space)) ≈ 4pi * radius^2 * (zlim[2] - zlim[1]) rtol =
        1e-3
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


using JuMP, GLPK, Test

function optimal_price_probabilties(
    price, demand, capacity; verbose = true)
    revenue = price .* demand
    model = Model(with_optimizer(GLPK.Optimizer))
    @variable(model, x[1:length(price)], lower_bound=0)
    # Objective: maximize profit
    @objective(model, Max, revenue' * x)
    # Constraint: can carry all
    @constraint(model, sum(x) <= 1)
    #=
    for i in 1:length(prices)
        @constraint(model, x[i] <= demand[i])
    end
    =#
    # Solve problem using MIP solver
    JuMP.optimize!(model)
    if verbose
        println("Objective is: ", JuMP.objective_value(model))
        println("Solution is:")
        for i in 1:length(price)
            print("x[$i] = ", JuMP.value(x[i]))
            println(", revenue = ", price[i] * JuMP.value(x[i]))
        end
    end
    @test JuMP.termination_status(model) == MOI.OPTIMAL
    @test JuMP.primal_status(model) == MOI.FEASIBLE_POINT
    JuMP.value.(x)
end

price = [1.99, 2.49, 2.99, 3.49, 3.99, 4.49]
demand = 50 .- 7 * price
capacity = 60

prob = optimal_price_probabilties(price, demand, capacity)
prob |> display
println()

demand2 = [33.74, 31.29, 25.81, 34.68, 31.84, 15.64]
prob2 = optimal_price_probabilties(price, demand2, capacity)
prob2 |> display


# 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 [-]"
)


using LinearAlgebra, DiffEqOperators, Random, Test

# Generate random parameters
al = rand(5)
bl = rand(5)
cl = rand(5)
dx = rand(5)
ar = rand(5)
br = rand(5)
cr = rand(5)

# Construct 5 arbitrary RobinBC operators
for i in 1:5
    Q = RobinBC((al[i], bl[i], cl[i]), (ar[i], br[i], cr[i]), dx[i])

    Q_L, Q_b = Array(Q,5i)

    #Check that Q_L is is correctly computed
    @test Q_L[2:5i+1,1:5i] ≈ Array(I, 5i, 5i)
    @test Q_L[1,:] ≈ [1 / (1-al[i]*dx[i]/bl[i]); zeros(5i-1)]
    @test Q_L[5i+2,:] ≈ [zeros(5i-1); 1 / (1+ar[i]*dx[i]/br[i])]

    #Check that Q_b is computed correctly
    @test Q_b ≈ [cl[i]/(al[i]-bl[i]/dx[i]); zeros(5i); cr[i]/(ar[i]+br[i]/dx[i])]

    # Construct the extended operator and check that it correctly extends u to a (5i+2)
    # vector, along with encoding boundary condition information.
    u = rand(5i)

    Qextended = Q*u
    CorrectQextended = [(cl[i]-(bl[i]/dx[i])*u[1])/(al[i]-bl[i]/dx[i]); u; (cr[i]+ (br[i]/dx[i])*u[5i])/(ar[i]+br[i]/dx[i])]
    @test length(Qextended) ≈ 5i+2

    # Check concretization
    @test Array(Qextended) ≈ CorrectQextended

    # Check that Q_L and Q_b correctly compute BoundaryPaddedVector
    @test Q_L*u + Q_b ≈ CorrectQextended

    @test [Qextended[1]; Qextended.u; Qextended[5i+2]] ≈ CorrectQextended

end

# Construct 5 arbitrary RobinBC operators w/non-uniform grid
al = rand(5)
bl = rand(5)
cl = rand(5)
dx = rand(5)
ar = rand(5)
br = rand(5)
cr = rand(5)
for j in 1:2
    for i in 1:5
        if j == 1
            Q = RobinBC((al[i], bl[i], cl[i]), (ar[i], br[i], cr[i]),
                        dx[i] .* ones(5 * i))
        else
            Q = RobinBC([al[i], bl[i], cl[i]], [ar[i], br[i], cr[i]],
                        dx[i] .* ones(5 * i))
        end
        Q_L, Q_b = Array(Q,5i)

        #Check that Q_L is is correctly computed
        @test Q_L[2:5i+1,1:5i] ≈ Array(I, 5i, 5i)
        @test Q_L[1,:] ≈ [1 / (1-al[i]*dx[i]/bl[i]); zeros(5i-1)]
        @test Q_L[5i+2,:] ≈ [zeros(5i-1); 1 / (1+ar[i]*dx[i]/br[i])]

        #Check that Q_b is computed correctly
        @test Q_b ≈ [cl[i]/(al[i]-bl[i]/dx[i]); zeros(5i); cr[i]/(ar[i]+br[i]/dx[i])]

        # Construct the extended operator and check that it correctly extends u to a (5i+2)
        # vector, along with encoding boundary condition information.
        u = rand(5i)

        Qextended = Q*u
        CorrectQextended = [(cl[i]-(bl[i]/dx[i])*u[1])/(al[i]-bl[i]/dx[i]); u; (cr[i]+ (br[i]/dx[i])*u[5i])/(ar[i]+br[i]/dx[i])]
        @test length(Qextended) ≈ 5i+2

        # Check concretization
        @test Array(Qextended) ≈ CorrectQextended

        # Check that Q_L and Q_b correctly compute BoundaryPaddedVector
        @test Q_L*u + Q_b ≈ CorrectQextended

        @test [Qextended[1]; Qextended.u; Qextended[5i+2]] ≈ CorrectQextended

    end
end

#3rd order RobinBC, calculated with left stencil [-11/6 3 -3/2 1/3], right stencil [-1/3 3/2 -3 11/6] and [α,β,γ] = [1 6 10]
u0 = -4/10
uend = 125/12
u = Vector(1.0:10.0)
Q = RobinBC((1.0, 6.0, 10.0), (1.0, 6.0, 10.0), 1.0, 3)
urobinextended = Q*u
@test urobinextended.l ≈ u0
@test urobinextended.r ≈ uend
# General BC should be equivalent
G = GeneralBC([-10.0, 1.0, 6.0], [-10.0, 1.0, 6.0], 1.0, 3)
ugeneralextended = G*u
@test ugeneralextended.l ≈ u0
@test ugeneralextended.r ≈ uend

#TODO: Implement tests for BC's that are contingent on the sign of the coefficient on the operator near the boundary


# 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 OrdinaryDiffEq, Test, DiffEqDevTools
using LinearAlgebra, Random

# Linear exponential solvers
A = DiffEqArrayOperator([2.0 -1.0; -1.0 2.0])
u0 = ones(2)
prob = ODEProblem(A,u0,(0.0,1.0))
sol1 = solve(prob, LinearExponential(krylov=:off))(1.0)
sol2 = solve(prob, LinearExponential(krylov=:simple))(1.0)
sol3 = solve(prob, LinearExponential(krylov=:adaptive))(1.0)
sol_analytic = exp(1.0 * Matrix(A)) * u0

@test isapprox(sol1, sol_analytic, rtol=1e-10)
@test isapprox(sol2, sol_analytic, rtol=1e-10)
@test isapprox(sol3, sol_analytic, rtol=1e-10)

# u' = A(t)u solvers

function update_func(A,u,p,t)
    A[1,1] = cos(t)
    A[2,1] = sin(t)
    A[1,2] = cos(t)*sin(t)
    A[2,2] = sin(t)^2
end
A = DiffEqArrayOperator(ones(2,2),update_func=update_func)
prob = ODEProblem(A, ones(2), (0.0, 5.))
dts = 1 ./2 .^(10:-1:1)
sol  = solve(prob,OrdinaryDiffEq.MagnusMidpoint(),dt=1/4)

dts = 1 ./2 .^(10:-1:1)
test_setup = Dict(:alg=>Vern9(),:reltol=>1e-14,:abstol=>1e-14)
sim = analyticless_test_convergence(dts,prob,MagnusMidpoint(),test_setup)
@test sim.𝒪est[:l2] ≈ 2 atol=0.2
sim = analyticless_test_convergence(dts,prob,MagnusMidpoint(krylov=true),test_setup)
@test sim.𝒪est[:l2] ≈ 2 atol=0.2

A = DiffEqArrayOperator(ones(2,2),update_func=update_func)
prob = ODEProblem(A, ones(2), (0.0, 5.))
dts = 1 ./2 .^(10:-1:1)
sol  = solve(prob,OrdinaryDiffEq.MagnusLeapfrog(),dt=1/4)

dts = 1 ./2 .^(10:-1:1)
test_setup = Dict(:alg=>Vern9(),:reltol=>1e-14,:abstol=>1e-14)
sim = analyticless_test_convergence(dts,prob,MagnusLeapfrog(),test_setup)
@test sim.𝒪est[:l2] ≈ 2 atol=0.2
sim = analyticless_test_convergence(dts,prob,MagnusLeapfrog(krylov=true),test_setup)
@test sim.𝒪est[:l2] ≈ 2 atol=0.2

function update_func(A,u,p,t)
    A[1,1] = cos(t)*u[1]
    A[2,1] = sin(t)*u[2]
    A[1,2] = -cos(t)*sin(t)
    A[2,2] = -(sin(t)^2) *u[1] * u[2]
end
A = DiffEqArrayOperator(ones(2,2),update_func=update_func)
prob = ODEProblem(A, ones(2), (0.5, 5.))
dts = 1 ./2 .^(10:-1:1)
sol  = solve(prob,OrdinaryDiffEq.LieEuler(),dt=1/4)

dts = 1 ./2 .^(10:-1:1)
test_setup = Dict(:alg=>Vern9(),:reltol=>1e-14,:abstol=>1e-14)
sim = analyticless_test_convergence(dts,prob,LieEuler(),test_setup)
@test sim.𝒪est[:l2] ≈ 1 atol=0.2
sim = analyticless_test_convergence(dts,prob,LieEuler(krylov=true),test_setup)
@test sim.𝒪est[:l2] ≈ 1 atol=0.2


using LinearAlgebra
using LinearManifoldCluster
using Distributions
using Distances
using LMCLUS
using Test
using Random

@testset "Random direction" begin
    n = 1000
    N = 10
    pinf, psup = extrema(LinearManifoldCluster.on_ball(N, n))
    @test pinf > -1.0
    @test psup <  1.0
    pinf, psup = extrema(LinearManifoldCluster.in_ball(N, n))
    @test pinf > -1.0
    @test psup <  1.0
    for d in colwise(Euclidean(), LinearManifoldCluster.on_ball(N,n), zeros(N))
        @test d ≈ 1.
    end
end

@testset "Translation vectors" begin
    n = 1000
    N = 10
    μ_bounds = 0.25 => 0.75
    μ = LinearManifoldCluster.translations(N, n, μ_bounds)
    # distance from center of unit hypercube to
    # any translation vector should be within bounds
    for d in colwise(Euclidean(), μ, ones(N)/2.)
        @test μ_bounds[1]/2. <= d <= μ_bounds[2]/2.
    end
end

@testset "Basis formation" begin
    N = 10
    M = 3
    B, B′ = LinearManifoldCluster.lm_basis(N, M)
    @test size(B) == (N,M)
    @test size(B′) == (N,N-M)
    P = B*B'
    @test P*P ≈ P
    @test P' ≈ P
    P′ = B′*(B′)'
    @test P′*P′ ≈ P′
    @test (P′)' ≈ P′
end

@testset "Manifold cluster generator" begin
    n = 1000
    N = 3
    M = 2
    μ_bounds = 0.25 => 0.75
    t = LinearManifoldCluster.translations(N, 1, μ_bounds)[:]
    D = Matrix{Float64}(I, N, N)
    B = D[:,1:2]
    B′ = D[:, 3:3]
    Φ = 0.5
    E = 0.05

    Random.seed!(230898573857)
    mfld_s = LinearManifoldCluster.generate_cluster(n, t, B, B′, Φ, E)

    Random.seed!(230898573857)
    Φ = fill(0.5, M)
    E = fill(0.05, N-M)
    mfld = LinearManifoldCluster.generate_cluster(n, t, B, B′, Φ, E)
    # one scalar bounds and array bounds return values
    @test mfld_s == mfld

    # manifold size & space dimension
    @test size(mfld) == (N,n)
    # if Φ is scalar, then manifold forms a hypercube
    # with side equal to Φ (± E)?
    mbound_max = norm(fill(Φ[1]+E[1],M),2)
    for d in colwise(Euclidean(), mfld.-t, zeros(N))
        @test 0.0 <= d <= mbound_max
    end
    # if E is scalar, then a distance from any point to
    # the manifold should be within ||[ɛ, ...,ɛ]||
    mdist_max = norm(E,2)
    for d in LMCLUS.distance_to_manifold(mfld, t, B)
        @test 0.0 <= d <= mdist_max
    end

    # manifold dimension generator
    @test LinearManifoldCluster.mdims(10,13,7) == [1, 2, 3, 4, 5, 6, 7, 7, 7, 7, 7, 7, 7]
end

@testset "Multi-manifold cluster generator" begin
    n = 1000
    N = 3
    M = [1,2,2]
    τ = 0.25 => 0.75
    Φ = Vector{Float64}[   # bounds on manifold points (per manifold wrt dims)
        [.6],
        [.2,  .73],
        [.4, .65]
    ]
    E = Vector{Float64}[   # bound of a point extent (per manifold wrt dims)
        [.05,  .03],
        [.01],
        [.04]
    ]
    lmcs, Bs, μs = generate(n,N,M,τ,Φ,E)
    @test length(lmcs) == length(M) # number of manifold
    @test size(lmcs[1]) == (N,n)      # number of manifold points
    for i in length(Φ)
        @test size(Bs[i]) == (N,length(Φ[i]))  # bases dimensions
    end
    @test size(μs) == (N,length(M))

    # multi-manifold generator (scalar)
    n = 1000
    N = 3
    m = 12
    τ = 0.25 => 0.75
    κ = 1.5
    lmcs, Bs, μs = generate(n,N,m,τ,κ)
    @test length(lmcs) == m
end


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


using Statistics

#mean
x = [20:100];
mean(x)

#median
median(skipmissing([10, 32, missing, 41, 75, 96]))

quantile(x, [0.1, 0.5, 0.9])

#standard deviation
function sd(x):
  expectation = mean(x)
  x = map(i -> (i - expectation)^2, x)
  std_dev = sqrt(sum(x)/(length(x) - 1))
  return std_dev


function compute_seat_id(spec)
  front::UInt8 = 0
  back::UInt8 = 127
  left::UInt8 = 0
  right::UInt8 = 7
  for c in spec
    if c == 'F'
      back = div(back - front, 2) + front
    elseif c == 'B'
      front = div(back - front, 2) + front
    elseif c == 'L'
      right = div(right - left, 2) + left
    elseif c == 'R'
      left = div(right - left, 2) + left
    end
  end

  back * 8 + right
end

function find_seat(data)
  seats = sort(map(compute_seat_id, data))
  max_id = maximum(seats)
  seat_id = -1
  for (idx, seat) in enumerate(seats)
    if seats[idx + 1] == seat + 2
      seat_id = seat + 1
      break
    end
  end

  max_id, seat_id
end

function solve(data)
  println(@time find_seat(data))
end

data = readlines("input.txt")
solve(data)


"""
    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