mutable struct KNNUCB <: BanditPolicy
    A::Int64
    t::Int64
    traj::Trajectory
    phi::Function
    theta::Float64
end

function KNNUCB(A::Int64, traj::Trajectory, phi::Function, theta::Float64)
    return KNNUCB(A, 0, traj, phi, theta)
end

function sort_nearest(pol::KNNUCB, x::Vector{Float64}, a::Int64)
    dist = Euclidean()
    x1 = reshape(x, (1, :))
    X = get_X(pol.traj, a, false)
    r = get_r(pol.traj, a, false)
    pwdist = pairwise(dist, X, x1; dims = 1)
    pwdist = reshape(pwdist, (:))
    inds = sortperm(pwdist)
    return pwdist[inds], r[inds]
end

function compute_upper_bound(pol::KNNUCB, x::Vector{Float64}, a::Int64)
    t = pol.t
    distances, r_nearest = sort_nearest(pol, x, a)
    uvals = Vector{Float64}(undef, length(r_nearest))
    for k = 1:length(r_nearest)
        radius = distances[k]
        uvals[k] = sqrt(pol.theta * log(t) / k) + pol.phi(t) * radius
    end
    min_uncertainty, min_k = findmin(uvals)
    return mean(r_nearest[1:min_k]) + min_uncertainty
end

function choose(x::Vector{Float64}, pol::KNNUCB)
    if pol.t + 1 <= pol.A
        return pol.t + 1
    end
    u = [compute_upper_bound(pol, x, a) for a in 1:pol.A]
    return findmax(u)[2]
end

function update!(pol::KNNUCB, x::Vector{Float64}, a::Int64, r::Float64)
    pol.t += 1
    if pol.t != pol.traj.t
        throw(ErrorException("Policy out of sync with trajectory"))
    end
end


Base.@kwdef mutable struct RKstruct
  nStage = nothing
  ARKE = nothing
  bRKE = nothing
  ARKI = nothing
  bRKI = nothing
  gRKI = nothing
end
function RungeKuttaMethod(Method)
  str = Method
  RK = RKstruct(;)
if str == "ARK32"
    RK.nStage=3;
    RK.ARKE=zeros(RK.nStage,RK.nStage);
    RK.bRKE=zeros(1,RK.nStage);
    RK.ARKE[2,1]=2-sqrt(2);
    ARKE32=1/6*(3 + sqrt(2));
    RK.ARKE[3,1]=1-ARKE32;
    RK.ARKE[3,2]=ARKE32;
    RK.bRKE[1]=1/(2*sqrt(2));
    RK.bRKE[2]=1/(2*sqrt(2));
    RK.bRKE[3]=1-1/(sqrt(2));
    RK.ARKI=zeros(RK.nStage,RK.nStage);
    RK.bRKI=zeros(1,RK.nStage);
    RK.ARKI[2,1]=1-1/(sqrt(2));
    RK.ARKI[2,2]=1-1/(sqrt(2));
    RK.ARKI[3,1]=1/(2*sqrt(2));
    RK.ARKI[3,2]=1/(2*sqrt(2));
    RK.ARKI[3,3]=1-1/(sqrt(2));
    RK.bRKI[1]=1/(2*sqrt(2));
    RK.bRKI[2]=1/(2*sqrt(2));
    RK.bRKI[3]=1-1/(sqrt(2));
    RK.gRKI=1-1/(sqrt(2));
elseif str == "Trap2(2,3,2)"
    #Trap2(2,3,2)
    RK.nStage=4;
    RK.ARKE=zeros(RK.nStage,RK.nStage);
    RK.bRKE=zeros(1,RK.nStage);
    RK.ARKE[2,1]=1;
    RK.ARKE[3,1]=1/2;
    RK.ARKE[3,2]=1/2;
    RK.ARKE[4,1]=1/2;
    RK.ARKE[4,3]=1/2;
    RK.bRKE[1]=1/2;
    RK.bRKE[2]=0;
    RK.bRKE[3]=1/2;
    RK.bRKE[4]=0;
elseif str == "RK1"
    RK.nStage=1;
    RK.ARKE=zeros(RK.nStage,RK.nStage);
    RK.bRKE=zeros(1,RK.nStage);
    RK.bRKE[1]=1;
elseif str == "RK1I"
    RK.nStage=1;
    RK.ARKE=zeros(RK.nStage,RK.nStage);
    RK.ARKE[1,1]=1;
    RK.bRKE=zeros(1,RK.nStage);
    RK.bRKE[1]=1;
elseif str == "RK4"
    RK.nStage=4;
    RK.ARKE=zeros(RK.nStage,RK.nStage);
    RK.ARKE[2,1]=1/2;
    RK.ARKE[3,2]=1/2;
    RK.ARKE[4,3]=1;
    RK.bRKE=zeros(1,RK.nStage);
    RK.bRKE[1]=1/6;
    RK.bRKE[2]=1/3;
    RK.bRKE[3]=1/3;
    RK.bRKE[4]=1/6;
elseif str == "RK3"
    RK.nStage=3;
    RK.ARKE=zeros(RK.nStage,RK.nStage);
    RK.ARKE[2,1]=1/3;
    RK.ARKE[3,2]=1/2;
    RK.bRKE=zeros(1,RK.nStage);
    RK.bRKE[1]=0;
    RK.bRKE[2]=0;
    RK.bRKE[3]=1;
elseif str == "DBM453"
    RK.nStage=5;
    RK.ARKE=zeros(RK.nStage,RK.nStage);
    RK.bRKE=zeros(1,RK.nStage);
    RK.ARKE[2,1]=0.10306208811591838;
    RK.ARKE[3,1]=-0.94124866143519894;
    RK.ARKE[3,2]=1.6626399742527356;
    RK.ARKE[4,1]=-1.3670975201437765;
    RK.ARKE[4,2]=1.3815852911016873;
    RK.ARKE[4,3]=1.2673234025619065;
    RK.ARKE[5,1]=-0.81287582068772448;
    RK.ARKE[5,2]=0.81223739060505738;
    RK.ARKE[5,3]=0.90644429603699305;
    RK.ARKE[5,4]=0.094194134045674111;

    RK.bRKE[1]=0.87795339639076672;
    RK.bRKE[2]=-0.72692641526151549;
    RK.bRKE[3]=0.7520413715737272;
    RK.bRKE[4]=-0.22898029400415090;
    RK.bRKE[5]=0.32591194130117247;

    RK.gRKI=0.32591194130117247;
    RK.ARKI=zeros(RK.nStage,RK.nStage);
    RK.bRKI=zeros(1,RK.nStage);
    RK.ARKI[2,1]=-0.22284985318525410;
    RK.ARKI[2,2]=RK.gRKI;
    RK.ARKI[3,1]=-0.46801347074080545;
    RK.ARKI[3,2]=0.86349284225716961;
    RK.ARKI[3,3]=RK.gRKI;
    RK.ARKI[4,1]=-0.46509906651927421;
    RK.ARKI[4,2]=0.81063103116959553;
    RK.ARKI[4,3]=0.61036726756832357;
    RK.ARKI[4,4]=RK.gRKI;
    RK.ARKI[5,1]=0.87795339639076675;
    RK.ARKI[5,2]=-0.72692641526151547;
    RK.ARKI[5,3]=0.75204137157372720;
    RK.ARKI[5,4]=-0.22898029400415088;
    RK.ARKI[5,5]=RK.gRKI;
    RK.bRKI=RK.bRKE;
end
return RK
end


using QuantumCumulants
using OrdinaryDiffEq
using SteadyStateDiffEq
using ModelingToolkit
using Test

@testset "diffeq" begin

hf = FockSpace(:cavity)
ha = NLevelSpace(:atom,(:g,:e))
h = tensor(hf, ha)

@qnumbers a::Destroy(h) σ::Transition(h,:g,:e)

# Single-atom laser
@cnumbers Δ g κ γ ν

H = Δ*a'*a + g*(a'*σ + σ'*a)
J = [a,σ,σ']
he_avg = meanfield([a'*a,σ'*σ,a*σ'],H,J;rates=[κ,γ,ν])

he_exp = cumulant_expansion(he_avg,2)
@test isequal(he_exp.equations, meanfield([a'*a,σ'*σ,a*σ'],H,J;rates=[κ,γ,ν],order=2).equations)

ps = [Δ,g,κ,γ,ν]
missed = find_missing(he_exp)
@test !any(p ∈ Set(missed) for p=ps)

# Exploit phase invariance
subs = Dict([missed; QuantumCumulants._conj.(missed)] .=> 0)
he_nophase = substitute(he_exp, subs)
@test isempty(find_missing(he_nophase))

@named sys = ODESystem(he_nophase)

# Numerical solution
p0 = ps .=> [0.0,0.5,1.0,0.1,0.9]
u0 = zeros(ComplexF64,3)
tmax = 10.0

prob = ODEProblem(sys,u0,(0.0,tmax),p0,jac=true)
sol = solve(prob,RK4())
n = sol[average(a'*a)]
pe = getindex.(sol.u,2)

@test all(iszero.(imag.(n)))
@test all(iszero.(imag.(pe)))
@test all(real.(n) .>= 0.0)
@test all(1.0 .>= real.(pe) .>= 0.0)

prob_ss = SteadyStateProblem(prob)
sol_ss = solve(prob_ss, DynamicSS(Tsit5(), abstol=1e-5, reltol=1e-5), abstol=1e-6, reltol=1e-6)

@test isapprox(sol_ss[a'*a], sol[a'a][end], atol=1e-4)

end # testset


module FractionalTimeDG

using ArgCheck
using OffsetArrays
import GaussQuadrature
import SpecialFunctions
import LinearAlgebra
import SparseArrays

export Store, coef_G, coef_K, coef_H0, coef_H1
export coef_H_uniform!, coef_Hn!
export FODEdG!, evaluate_pcwise_poly!, evaluate_pcwise_poly, reconstruction
export legendre_polys!
export SpatialGrid, spatial_L2_projection, FPDEDG, evaluate_pcwise_error,
       evaluate_fem1d_soln!

Γ(x) = SpecialFunctions.gamma(x)

struct Store{T<: AbstractFloat}
    # shared data for low-level routines
    # initialise with store = setup(α, rmax, Mmax)
    # or setup(α, rmax, Mmax, ppImax)
    α::T          # fraction diffusion exponent
    rmax::Integer # (max degree of piecewise polynomials) + 1
    Mmax::Integer # maximum number of Gauss points
    unitlegendre::Vector{Matrix{T}} # w(x) = 1 on [0,1]
    legendre::Vector{Matrix{T}}     # w(x) = 1 on [-1,1]
    jacobi1::Vector{Matrix{T}}      # w(x) = (1-x)^(α-1)
    jacobi2::Vector{Matrix{T}}      # w(x) = (1+x)^(α-1)
    jacobi3::Vector{Matrix{T}}      # w(x) = (1-x)^α
    jacobi4::Vector{Matrix{T}}      # w(x) = (1+x)^α
    jacobi5::Vector{Matrix{T}}      # w(x) = (1-x)(1+x)^(α-1)
    ppImax::Integer # max points per interval (second dimension of Ψ)
    Ψ::Matrix{T}
    dΨ::Matrix{T}
    A::Vector{T}
    B::Vector{T}
    C::Matrix{T}
end

include("parts/store.jl")
include("parts/legendre_utils.jl")
include("parts/coef_uniform.jl")
include("parts/coef.jl")
include("parts/scalar_fode.jl")
include("parts/FEM1D.jl")

end # module FractionalTimeDG


module BDMAnalytic

    export BDM, prob, SSprob

    using Plots, Parameters, LinearAlgebra, Distributions, StatsBase

    """
    Define the Analytic Ising Model struct
    """
    @with_kw struct BDM
        N::Int64 = 100 # number of agents
        γ::Float64 = 1.0 # timescale
        F::Float64
        J::Float64
        α::Float64 = 0.0 # could generalise to be 2*d for NN.
        β::Float64 = 1.0
        pars::Vector{Float64} = [γ, F, J, α, β]
        A::Matrix{Float64} = make_TRM(pars, N)
        λ::Array{Complex{BigFloat}} = GetEigVals(A)
        q_arr::Vector{Vector{Complex{BigFloat}}} = [GetOrthoQ(pars,N,λ[j]) for j in 1:N+1]
        p_arr::Vector{Vector{Complex{BigFloat}}} = [GetOrthoP(pars,N,λ[j]) for j in 1:N+1]
        den_prod::Vector{Complex{BigFloat}} = [prod([λ[i]-λ[j] for j in filter!(e->e≠i,[j for j in 1:N+1])]) for i in 1:N+1]
        As::Vector{BigFloat} = [a(pars,N,j) for j in 1:N]
        Bs::Vector{BigFloat} = [b(pars,N,j) for j in 0:N-1]
    end

    """
    Calculation of the gain.
    """
    function Gain(Sᵢ::Int64, n::Int64, pars::Vector{Float64}, N::Int64)
        γ, F, J, α, β = pars;
        m = (2*n-N)/N; # order parameter.
        G = -2*Sᵢ*(F + J*m*(1+α)) + ((2*J*(1+α))/N);
        return G
    end

    """
    Define rate function for N₊ prod
    """
    function r1(pars::Vector{Float64}, N::Int64, n::Int64)
        γ = pars[1];
        β = pars[5];
        G = Gain(-1, n, pars, N)
        return convert(BigFloat,γ*(N-n)*(1+exp(-β*G))^-1)::BigFloat
    end
    # rescaled definition in solution
    a(pars,N,n) = r1(pars,N,n-1);

    """
    Define rate function for N₋ prod
    """
    function r2(pars::Vector{Float64}, N::Int64, n::Int64)
        γ = pars[1];
        β = pars[5];
        G = Gain(1, n, pars, N)
        return convert(BigFloat,γ*n*(1+exp(-β*G))^-1)::BigFloat
    end
    # rescaled definition in solution
    b(pars,N,n) = r2(pars,N,n+1);

    """
    Make the transition rate matrix
    """
    function make_TRM(pars::Vector{Float64}, N::Int64)
        A = zeros(N+1, N+1);
        for i in 1:size(A)[1]
            for j in 1:size(A)[2]
                if i == 1 && j == 1
                    A[1,1] = - a(pars,N,1)
                elseif i == j && i>1
                    A[i,i] = -(a(pars,N,i)+b(pars,N,i-2))
                elseif i == j+1
                    A[i,j] = a(pars,N,j)
                elseif i == j-1
                    A[i,j] = b(pars,N,i-1)
                else
                    continue
                end
            end
        end
        return A::Matrix{Float64}
    end

    """
    Get eigenvalues
    """
    function GetEigVals(A::Matrix{Float64})
        λ = convert(Array{Complex{BigFloat}}, reverse(eigvals(A)));
        if λ[1] == λ[2] # if get repeated zero eigenvalues from solver manually separate.
            λ[1] = 0.0 + (1E-30)im; λ[2] = 0.0 - (1E-30)im;
        end
        return λ::Array{Complex{BigFloat}}
    end

    """
    Get the p orthogonal polynomials
    """
    function GetOrthoP(pars::Vector{Float64}, N::Int64, λᵢ::Complex{BigFloat})
        p = Array{Complex{BigFloat},1}(undef, N+1);
        p[1] = 1.0; p[2] = λᵢ+a(pars,N,1);
        for i in 3:N+1
            p[i] = (λᵢ+a(pars,N,i-1)+b(pars,N,i-3))*p[i-1] - b(pars,N,i-3)*a(pars,N,i-2)*p[i-2]
        end
        return p::Vector{Complex{BigFloat}}
    end

    """
    Get the q orthogonal polynomials
    """
    function GetOrthoQ(pars::Vector{Float64}, N::Int64, λᵢ::Complex{BigFloat})
        q = Array{Complex{BigFloat},1}(undef, N+3);
        q[N+3] = 1.0; q[N+2] = λᵢ + b(pars,N,N-1);
        for i in reverse([j for j in 3:N+1])
            q[i] = (λᵢ + a(pars,N,i-1)+b(pars,N,i-3))*q[i+1] - b(pars,N,i-2)*a(pars,N,i-1)*q[i+2]
        end
        return q::Vector{Complex{BigFloat}}
    end

    """
    Define the sum of the elements from the solution
    """
    function sum_elems(λᵢ::Complex{BigFloat}, t::Float64, m::Int64, m₀::Int64, pars::Array{Float64,1}, N::Int64, p_arrᵢ::Vector{Complex{BigFloat}}, q_arrᵢ::Vector{Complex{BigFloat}}, den_prodᵢ::Complex{BigFloat})
        return exp(λᵢ*t)*p_arrᵢ[m+1]*q_arrᵢ[m₀+3]/den_prodᵢ::Complex{BigFloat}
    end

    """
    Define P(m,t|m₀)
    """
    function pm(t::Float64, BD::BDM, m::Int64, m₀::Int64)
        @unpack λ, pars, N, As, Bs, q_arr, p_arr, den_prod = BD
        if m<m₀
            return prod(Bs[m+1:m₀])*sum([sum_elems(λ[i], t, m, m₀, pars, N, p_arr[i], q_arr[i], den_prod[i]) for i in 1:N+1])::Complex{BigFloat}
        elseif m==m₀
            return sum([sum_elems(λ[i], t, m, m, pars, N, p_arr[i], q_arr[i], den_prod[i]) for i in 1:N+1])::Complex{BigFloat}
        else
            return prod(As[m₀+1:m])*sum([sum_elems(λ[i], t, m₀, m, pars, N, p_arr[i], q_arr[i], den_prod[i]) for i in 1:N+1])::Complex{BigFloat}
        end
    end

    """
    Define the probability distribution return function from a initial distribution.
    """
    function prob(BD::BDM, t::Float64, q_init_D::Distribution{Univariate, Discrete})
        @unpack N = BD
        q_init = pdf(q_init_D)
        pmt = Array{Complex{BigFloat}}(undef,N+1)
        for i in 1:N+1 # loop over the m0's
            pmt[i] = sum([q_init[n+1]*pm(t, BD, i-1, n) for n in 0:N])
        end
        return (LinRange(-1.0,1.0,N+1),real(pmt))::Tuple{LinRange{Float64}, Vector{BigFloat}}
    end

    """
    Define the probability distribution return function from a precise value of m₀.
    """
    function prob(BD::BDM, t::Float64, m₀::Int64)
        @unpack N = BD
        return (LinRange(-1.0,1.0,N+1),real([pm(t, BD, m, m₀) for m in 0:N]))::Tuple{LinRange{Float64}, Vector{BigFloat}}
    end

    """
    Define the probability distribution return function for N/2 value of m₀.
    """
    function prob(BD::BDM, t::Float64)
        @unpack N = BD
        m₀ = floor(Int64,N/2) # note that we have the n=0 state too.
        return (LinRange(-1.0,1.0,N+1),real([pm(t, BD, m, m₀) for m in 0:N]))::Tuple{LinRange{Float64}, Vector{BigFloat}}
    end

    """
    Return the steady state distribution
    """
    function SSprob(BD::BDM)
        @unpack As, Bs, N = BD
        ps = Vector{BigFloat}(undef, N+1)
        for m in 2:N # use the product rule
            ps[m] = prod(As[1:m-1])*prod(Bs[m:N])
        end
        ps[1] = prod(Bs[1:N]) # do product for the B's
        ps[N+1] = prod(As[1:N]) # do product for the A's
        return (LinRange(-1.0,1.0,N+1),ps/sum(ps))::Tuple{LinRange{Float64}, Vector{BigFloat}}
    end

end # module end


## Makes a random walk Metropolis kernel
function makeRWMKernel(logTargetDensity::F,
  Σ::SMatrix{d, d, Float64}, rng::RNG=Random.GLOBAL_RNG) where
  {F<:Function, d, RNG<:AbstractRNG}

  A::SMatrix{d, d, Float64} = cholesky(Σ).L
  scratchv::MVector{d, Float64} = MVector{d, Float64}(undef)
  prevx::MVector{d, Float64} = MVector{d, Float64}(undef)
  ldprevx = Ref(-Inf)
  accepts = Ref(0)
  calls = Ref(0)

  function P(x::SVector{d, Float64})
    calls.x += 1
    randn!(rng, scratchv)
    scratchv .= A * scratchv
    z::SVector{d, Float64} = x + scratchv
    if x == prevx
      lpi_x = ldprevx.x
    else
      lpi_x = logTargetDensity(x)
      prevx .= x
    end
    lpi_z = logTargetDensity(z)
    if -randexp(rng) < lpi_z - lpi_x
      prevx .= z
      ldprevx.x = lpi_z
      accepts.x += 1
      rval = z
    else
      rval = x
    end
    return rval
  end

  function P(s::Symbol)
    s == :acceptanceRate && return accepts.x / calls.x
  end

  return P
end


## qValues ##

function qValues{T<:AbstractFloat}(pValues::Vector{T}, pi0::T, pfdr::Bool = false)
    validPValues(pValues)
    validPValues([pi0])
    n = length(pValues)
    u = sortperm(pValues)
    v = competerank(pValues) ## ties with 'min'
    if pfdr
        qvals = (pi0 .* n .* pValues) ./ (v .* (1 - (1 - pValues) .^ n))
    else
        qvals = (pi0 .* n .* pValues) ./ v
    end
    qvals[u[n]] = min(qvals[u[n]], 1)
    for i in (n - 1):-1:1
        qvals[u[i]] = min(qvals[u[i]], qvals[u[i + 1]])
    end
    return qvals
end


# MIT license
# Copyright (c) Microsoft Corporation. All rights reserved.
# See LICENSE in the project root for full license information.

"""SphericalPolygon uses StaticArrays to represent vertices. Expect performance degradation for polygons with large numbers of vertices. Performance appears to be good up to perhaps 100 vertices, perhaps as much as 1000 vertices. By 10,000 vertices performance is terrible."""
struct SphericalPolygon{N,T<:Real}
    ptvectors::SMatrix{3,N,T}
    spherecenter::SVector{3,T}
    radius::T

    function SphericalPolygon(points::SMatrix{3,N,T},spherecenter::SVector{3,T},radius::T) where{T<:Real,N}
        normpoints = MMatrix{3,N,T}(undef)

        for i in 1:N
            normpoints[:,i] = normalize(points[:,i] - spherecenter)
        end
        return new{N,T}(SMatrix{3,N,T}(normpoints),SVector{3,T}(spherecenter),radius)
    end
end
export SphericalPolygon

"""returns the spherical angle formed by the cone with centervector at its center with neighbor1,neighbor2 the edges"""
function sphericalangle(neighbor1::SVector{3,T}, centervector::SVector{3,T}, neighbor2::SVector{3,T}) where{T<:Real}
    vec1 = normalize(neighbor1 - (dot(neighbor1,centervector)*centervector)) #subtract off the component of the neighbor vectors from the center vector. This leaves only the component orthogonal to center vector
    vec2 = normalize(neighbor2 - (dot(neighbor2,centervector)*centervector))

    return acos(dot(vec1,vec2))
end

struct SphericalTriangle{T<:Real}
    ptvectors::SMatrix{3,3,T}
    spherecenter::SVector{3,T}
    radius::T

    """Constructor which uses SMatrix to store vectors. For compatibility with LazySets this is more efficient"""
    function SphericalTriangle(vectors::SMatrix{3,3,T},spherecenter::SVector{3,T},radius::T) where{T}
        temp = MMatrix{3,3,T}(undef)
        for col in 1:3
            temp[:,col] = normalize(vectors[:,col]-spherecenter)
        end

        new{T}(temp,spherecenter,radius)
    end
end
export SphericalTriangle

vector(a::SphericalTriangle,i::Int) = a.ptvectors[:,i]

function area(tri::SphericalTriangle{T}) where{T<:Real}
    vec1,vec2,vec3 = vector(tri,1),vector(tri,2),vector(tri,3)
    sum = T(0)
    sum += sphericalangle(vec1,vec3,vec2)
    sum += sphericalangle(vec2,vec1,vec3)
    sum += sphericalangle(vec3,vec2,vec1)
    return (sum-π) * tri.radius^2
end

"""Conceptually breaks the convex spherical polygon into spherical triangles and computes the sum of the angles of all the triangles. The sum of all the angles around the centroid is 2π. Have to subtract π for each of the N triangles. Rather than compute the angles of triangles formed by taking edges from the centroid to each vertex, can instead just compute the internal angle of neighboring edges. Total polygon area is 2π -Nπ + ∑(interior angles)."""
function area(poly::SphericalPolygon{N,T}) where{T<:Real,N}
    return area(poly.ptvectors,poly.radius)
end

function area(ptvecs::SMatrix{3,N,T},radius::T) where{N,T<:Real}
    accum = T(0)

    for i in 2:N-1
        v1 = SVector{3,T}(ptvecs[1,i-1],ptvecs[2,i-1],ptvecs[3,i-1])
        v2 = SVector{3,T}(ptvecs[1,i],ptvecs[2,i],ptvecs[3,i])
        v3 = SVector{3,T}(ptvecs[1,i+1],ptvecs[2,i+1],ptvecs[3,i+1])
        accum += sphericalangle(v1,v2,v3)
    end

    # finish up first and last interior angles which have different indexing because of wraparound
    v1 = SVector{3,T}(ptvecs[1,N-1],ptvecs[2,N-1],ptvecs[3,N-1])
    v2 = SVector{3,T}(ptvecs[1,N],ptvecs[2,N],ptvecs[3,N])
    v3 = SVector{3,T}(ptvecs[1,1],ptvecs[2,1],ptvecs[3,1])

    accum += sphericalangle(v1,v2,v3)

    v1 = SVector{3,T}(ptvecs[1,2],ptvecs[2,2],ptvecs[3,2])
    v2 = SVector{3,T}(ptvecs[1,1],ptvecs[2,1],ptvecs[3,1])
    v3 = SVector{3,T}(ptvecs[1,N],ptvecs[2,N],ptvecs[3,N])

    accum += sphericalangle(v1,v2,v3)

    accum = (2π -N*π + accum)*radius^2
    return accum
end

function circlepoly(nsides; offset = [0.0,0.0,1.0])
    step = 2π/nsides
    result = MMatrix{3,nsides,Float64}(undef)
    for i in 0:nsides-1
        y,x = sincos(step*i)
        result[:,i+1] = [x,y,0.0] .+ offset
    end
    return SMatrix{3,nsides,Float64}(result)
end
export circlepoly

"""creates a circular polygon that subtends a half angle of θ"""
function sphericalcircle(θ, nsides = 10)
    temp = MMatrix{3,nsides,Float64}(undef)
    for i in 0:1:(nsides-1)
        ϕ = i*2π/nsides
        temp[1,i+1] = sin(θ)*cos(ϕ)
        temp[2,i+1] = cos(θ)
        temp[3,i+1] = sin(θ)*sin(ϕ)
    end
    return SphericalPolygon(SMatrix{3,nsides,Float64}(temp),SVector(0.0,0.0,0.0),1.0)
end
export sphericalcircle

oneeigthsphere() = SphericalTriangle(SMatrix{3,3,Float64}(
    0.0,1.0,0.0,
    1.0,0.0,0.0,
    0.0,0.0,1.0),
    SVector(0.0,0.0,0.0),
    1.0)
export oneeigthsphere

onesixteenthphere() = SphericalTriangle(SMatrix{3,3,Float64}(
    0.0,1.0,0.0,
    1.0,1.0,0.0,
    0.0,0.0,1.0),
    SVector(0.0,0.0,0.0),
    1.0)
export onesixteenthphere

threesidedpoly() = SphericalPolygon(SMatrix{3,3,Float64}(
    0.0,1.0,0.0,
    1.0,0.0,0.0,
    0.0,0.0,1.0),
    SVector(0.0,0.0,0.0),
    1.0)
export threesidedpoly

foursidedpoly() = SphericalPolygon(SMatrix{3,4,Float64}(
    0.0,1.0,0.0,
    1.0,1.0,-.9,
    1.0,0.0,0.0,
    0.0,0.0,1.0),
    SVector(0.0,0.0,0.0),
    1.0)
export foursidedpoly


export solve

function print_header(io::IO)
    println("it  primal obj    dual obj           mu        delta        alpha   it. time")
end

function print_it(io::IO, it, primobj, dualobj, mu, delta, alpha, time)
    @printf("%2i  % .3e  % .3e   % .3e    % .6f   % .3e   %f\n", it, primobj, dualobj, mu, delta, alpha, time)
end

"""
    solve(pb::SDCOContext{T}, z0::PointPrimalDual)

    Strictly feasible IP method for positive orthant and SDP cone, with NT scaling.
    z0 is the strictly feasible starting point.
"""
function solve(pb::SDCOContext{T}, z0::PointPrimalDual) where T<:Number
    z = deepcopy(z0)

    outlev = pb.options[:opt_outlev]

    ## Centering section
    θ = 1/3
    τ = 0.99
    ω = 1e-4

    cur_mu = mu(z)
    cur_delta = delta(pb, z)
    it = 0
    α = -1
    d = NesterovToddstep(pb, z, 0.)

    (outlev > 0) && print_header(stdout)
    (outlev > 0) && print_it(stdout, it, -1, -1, cur_mu, cur_delta, -1, α)

    if (cur_delta > θ)
        (outlev > 0) && println("Getting to central path...")
    end

    while (cur_delta > θ) && (it < 20)
        t1 = time()
        α = NTcentering_it!(pb, τ, ω, z, cur_mu, cur_delta)

        cur_mu = mu(z)
        cur_delta = delta(pb, z)

        it += 1
        (outlev > 0) && print_it(stdout, it, -1, -1, cur_mu, cur_delta, α, time() - t1)
    end

    ## Path following section
    (outlev > 0) && println("Following central path...")
    ε = pb.options[:opt_ε]
    maxit = pb.options[:opt_maxit]
    alphacorr = 0.999

    nc = get_nc(z.x)
    # K = ceil((1 + sqrt(1+ 13*nc/2)) * log10(mu(z) / epsilon))
    # @show K

    cur_μ = mu(z)

    # K = log(ε / cur_μ) / log(1 + 4 / (13*nc) * (1 - sqrt(1+13*nc/2)))
    # @show K

    while (it < maxit) && (cur_μ > ε)
        t1 = time()

        res = α, cur_μ = NTpredcorr_it!(pb, alphacorr, cur_μ, cur_delta, z, d)

        it+=1

        cur_μ = mu(z)
        cur_delta = delta(pb, z)
        (outlev > 0) && print_it(stdout, it, get_primobj(pb, z), get_dualobj(pb, z), cur_μ, cur_delta, α, time() - t1)
    end

    (outlev > 0) && println("Final duality gap: ", get_primobj(pb, z) - get_dualobj(pb, z))
    (outlev > 0) && println("Final mu         : ", mu(z))

    return z
end

function NTcentering_it!(pb, τ, ω, z, cur_mu, cur_delta)
    mu_z = cur_mu
    dc = NesterovToddstep(pb, z, mu_z)

    if cur_delta < sqrt( (2*τ^2) / (1+2*τ^2) )
        α = 1
    else
        i = 0
        α = 2^(-i)

        while phi_mu(z + α*dc, mu_z) > (phi_mu(z, mu_z) - 4*ω*mu_z*α*delta(pb, z)^2)
            i += 1
            α = (1/2)^(i)

            (i > 100) && error("i = $i...")
        end
    end

    add!(z, α * dc)

    return α
end

# function symmetrize!(z)
#     for i in 1:length(z.x.mats)
#         z.x.mats[i] = (z.x.mats[i] + transpose(z.x.mats[i]))/2
#     end
#
#     for i in 1:length(z.s.mats)
#         z.s.mats[i] = (z.s.mats[i] + transpose(z.s.mats[i]))/2
#     end
# end

function NTpredcorr_it!(pb, alphacorr, cur_mu, cur_delta, z, d)
    # centering step
    mu_z = mu(z)
    NesterovToddstep!(pb, z, mu_z, d)

    # zc = z + dc
    add!(z, d)
    mu_zc = mu(z)

    # affine step
    NesterovToddstep!(pb, z, 0., d)

    α = NT_getalpha(pb, z, d)

    (pb.options[:opt_outlev] > 1) && printstyled("mu(z) - mu(zc) : ", abs(mu_z - mu_zc), "\n", color=:red)

    alphacorr = min(α, 0.99999)

    add!(z, alphacorr*d)
    mu_za = mu(z)

    (pb.options[:opt_outlev] > 1) && printstyled("mu(z+) / mu(zc) - (1 - alpha): ", abs(mu_za / mu_zc - (1-α)), "\n", color=:red)

    return α, mu_za
end


using JULES, Oceananigans
using Plots
using Printf

Nx = Ny = 1
Nz = 32
L = 10e3

grid = RegularCartesianGrid(size=(Nx, Ny, Nz), halo=(2, 2, 2),
                            x=(0, L), y=(0, L), z=(0, L))

pₛ = 100000
Tₐ = 293.15
g  = 9.80665

buoyancy = IdealGas()
Rᵈ, cₚ = buoyancy.Rᵈ, buoyancy.cₚ

####
#### Isothermal atmosphere
####

H = Rᵈ * Tₐ / g    # Scale height [m]
ρₛ = pₛ / (Rᵈ*Tₐ)  # Surface density [kg/m³]

p₀(x, y, z) = pₛ * exp(-z/H)
ρ₀(x, y, z) = ρₛ * exp(-z/H)

θ₀(x, y, z) = Tₐ * exp(z/H * Rᵈ/cₚ)
Θ₀(x, y, z) = ρ₀(x, y, z) * θ₀(x, y, z)

####
#### Sponge layer forcing to damp vertically propagating acoustic waves at the top.
####

const τ⁻¹ = 1    # Damping/relaxation time scale [s⁻¹]. This is very strong damping.
const Δμ = 0.1L  # Sponge layer width [m] set to 10% of the domain height.
@inline μ(z, Lz) = τ⁻¹ * exp(-(Lz-z) / Δμ)

@inline Fw(i, j, k, grid, t, Ũ, C̃, p) = @inbounds -μ(grid.zF[k], grid.Lz) * Ũ.ρw[i, j, k]
forcing = ModelForcing(w=Fw)

####
#### Create model
####

model = CompressibleModel(
                      grid = grid,
                  buoyancy = buoyancy,
        reference_pressure = pₛ,
    thermodynamic_variable = ModifiedPotentialTemperature(),
                   tracers = (:Θᵐ,),
                   forcing = forcing
)

####
#### Set initial conditions
####

set!(model.density, ρ₀)
set!(model.tracers.Θᵐ, Θ₀)

# Initial profiles
Θ₀_prof = model.tracers.Θᵐ[1, 1, 1:Nz]
ρ₀_prof = model.density[1, 1, 1:Nz]

# Arrays we will use to store a time series of ρw(z = L/2).
times = [model.clock.time]
ρw_ts = [model.momenta.ρw[1, 1, Int(Nz/2)]]

####
#### Time step and keep plotting vertical profiles of ρθ′, ρw, and ρ′.
####

# @animate for i=1:100
while model.clock.time < 500
    time_step!(model; Δt=0.5, Nt=10)

    @show model.clock.time
    Θ_prof = model.tracers.Θᵐ[1, 1, 1:Nz] .- Θ₀_prof
    W_prof = model.momenta.ρw[1, 1, 1:Nz+1]
    ρ_prof = model.density[1, 1, 1:Nz] .- ρ₀_prof

    Θ_plot = plot(Θ_prof, grid.zC, xlim=(-2.5, 2.5), xlabel="Theta_prime", ylabel="z (m)", label="")
    W_plot = plot(W_prof, grid.zF, xlim=(-1.2, 1.2), xlabel="rho*w", label="")
    ρ_plot = plot(ρ_prof, grid.zC, xlim=(-0.01, 0.01), xlabel="rho", label="")

    t_str = @sprintf("t = %d s", model.clock.time)
    display(plot(Θ_plot, W_plot, ρ_plot, title=["" "$t_str" ""], layout=(1, 3), show=true))

    push!(times, model.clock.time)
    push!(ρw_ts, model.momenta.ρw[1, 1, Int(Nz/2)])
end

####
#### Plot the initial profile and the hydrostatically balanced profile.
####

ρ∞_prof = model.density[1, 1, 1:Nz]

θ₀_prof = Θ₀_prof ./ ρ₀_prof
θ∞_prof = model.tracers.Θᵐ[1, 1, 1:Nz] ./ ρ∞_prof

θ_plot = plot(θ₀_prof, grid.zC, xlabel="theta (K)", ylabel="z (m)", label="initial", legend=:topleft)
plot!(θ_plot, θ∞_prof, grid.zC, label="balanced")

ρ_plot = plot(ρ₀_prof, grid.zC, xlabel="rho (kg/m³)", ylabel="z (m)", label="initial")
plot!(ρ_plot, ρ∞_prof, grid.zC, label="balanced")

display(plot(θ_plot, ρ_plot, show=true))

####
#### Plot timeseries of maximum ρw
####

plot(times, ρw_ts, linewidth=2, xlabel="time (s)", ylabel="rho*w(z=$(Int(L/2)) m)", label="", show=true)


using SparseTransforms
using Test

function test_query()
    n = 4
    num_delays = 6
    locs = [4, 6, 10, 15]
    signal = TestSignal(n, locs)

    # this also implicitly does type checks

    println("Testing all the query methods")
    for method_name in all_methods["query"]
        local b = get_b(signal; method=method_name)
        local Ms = get_Ms(n, b; method=method_name)
        @test b < n
        for M in Ms
            @test size(M) == (n, b)
        end
    end

    println("Testing the delays methods")
    for method_name in all_methods["delays"]
        local D = get_D(n; method=method_name, num_delays=num_delays)
        # @ test size(D, 1) > n TODO
        @test size(D, 2) == n
    end
end


export RandomizedNewton

"""
Randomized Newton
-> for symmetric positive definite matrix

Section 3.6 in Gower, R. M., & Richtárik, P. (2015).
Randomized iterative methods for linear systems.
SIAM Journal on Matrix Analysis and Applications, 36(4), 1660-1690.
"""
function RandomizedNewton(
  stp::AbstractStopping;
  is_zero_start::Bool = false,
  r::Int = 15,
  rand_r::Bool = false,
  kwargs...,
)
  A, b = get_matrix(stp.pb), get_vector(stp.pb)
  m, n = size(A)
  state = stp.current_state
  if is_zero_start
    state.res .= b
  else
    state.res .= b .- A * state.x
  end
  OK = start!(stp, no_opt_check = true)

  while !OK
    r = rand_r ? Int(floor(rand() * n) + 1) : r
    sub = zeros(Int64, r)
    sample!(1:n, sub) #x is a subset of [1,...,n] of size r

    Ai = A[sub, :]
    resk = state.res[sub] #resk = Ai*xk - b[sub]

    if 0 ∉ resk
      state.x .-= Ai \ resk
    end

    state.res .= b .- A * state.x
    OK = cheap_stop!(stp)
  end

  return stp
end


using NWNTools
using Test
using DelimitedFiles, StaticArrays

@testset "NWNTools.jl" begin
    @testset "Benchmark1" begin
        data = readdlm("benchmark1.txt",',', BigFloat)
        dims = SA[BigFloat("3.0"),BigFloat("5.0")]
        elecs = [Line(SVector{2}(data[1,:]),SVector{2}(data[2,:]))]
        grnds = [Line(SVector{2}(data[3,:]),SVector{2}(data[4,:]))]
        Rⱼ = BigFloat("20.0")
        V₊ = BigFloat("5.0")
        volts = [V₊]
        ρₑ = BigFloat("22.63676")/BigFloat("1000")
        D = BigFloat("0.06")

        lines = Line{BigFloat,2}[]
        for i ∈ 5:2:size(data)[1]
            push!(lines, Line(SVector{2}(data[i,:]),SVector{2}(data[i+1,:])))
        end
        props = repeat([WireProp(ρₑ,D)],length(lines))

        nwn_JDA = NWN_circuit{:JDA}(lines, props, elecs, volts, grnds, dims, Rⱼ,Rⱼ)
        R_JDA = sheet_resistance(nwn_JDA)
        nwn_MNR = NWN_circuit{:MNR}(lines, props, elecs, volts, grnds, dims, Rⱼ,Rⱼ)
        R_MNR = sheet_resistance(nwn_MNR)
        @test R_JDA≈(8*Rⱼ/3)
        numerator = 72*(3+5*sqrt(big(2)))*ρₑ^2+2*D^2*big(π)*Rⱼ*((81+69*sqrt(big(2)))*ρₑ+20*D^2*big(π)*Rⱼ)
        denominator = 15*D^2*big(π)*(2*(1+sqrt(big(2)))*ρₑ+D^2*big(π)*Rⱼ)
        R_MNR_analytic = numerator/denominator
        @test R_MNR≈R_MNR_analytic
end


module BasisFunctions

using Common, StaticArrays, OffsetArrays

export Gaussian, cartesianToSphericalMxs

struct Gaussian
    ζ::Float
    R::Coordinate
end

function positiveDiagonal(l::Int8, m::Int8)
    δl0 = l == 0 ? 1 : 0
    return sqrt(2 ^ δl0 * (2*l + 1) / (2 * l + 2))
end

function cartesianToSphericalMxs(lmax::Int)
    nang = (lmax+1)^2
    ncart = (lmax+1)*(lmax+2)*(lmax+3)÷6
    M = OffsetArray(Array([Matrix(zeros(Float64, ((l+1)*(l+2)÷2, 2*l+1))) for l=0:lmax]), -1)
    #  M = zeros(Float, (ncart, nang))
    for l=0:lmax #  Over common angular momentum quantum number
        iang = 0
        println()
        for m=-l:l  #  Spherical quantum number m
            iang += 1
            vm = m < 0 ? 0.5 : 0
            am = abs(m)
            vmax = Int(floor(am/2-vm))
            mfact = m == 0 ? 0.5 : 1
            N = 1 / (2^am * factorial(l)) * sqrt(2 * mfact * (factorial(l+am) * factorial(l-am)))
            #  icart = l ^ 2 #  Take into account the cartesians of all the previous ls
            icart = 0
            for k=0:l #  Cartesian quantum number k
                for j=0:l-k #  Cartesian quantum number j
                    i=l-k-j #  Cartesian quantum number i
                    #  if m == 0
                        #  println("x,y,z:",i,j,k)
                    #  end
                    icart += 1
                    for t=0:(l-am)÷2
                        kk = l - 2 * t - am
                        if kk ≠ k continue end
                        for u=0:t
                            for vv=0:vmax
                                v=vv+vm
                                ii = 2 * t + am - 2 * (u + v)
                                jj = 2 * (u + v)
                                if ii == i && jj == j && kk == k
                                    M[l][icart,iang] += (-1)^(t+vv) * 0.25 ^ t * N * binomial(l,t) * binomial(l-t,am+t) * binomial(t,u) * binomial(am,Int(2*v))
                                end
                            end
                        end
                    end
                end
            end
        end
    end
    return M
end

end


"""
function evolve_OU(df, ycol, norm_prefix, norm_feas,
    μσs, default_FloatType,
    test_wd_idxs, test_idxs,
    filename, test_dates) where I <: Integer

https://pdfs.semanticscholar.org/dfdd/ff87a20a918e7ac101808881b926bb298e95.pdf

The Ornstein-Uhlembeck position process
dXₜ = θ(μ - Xₜ)dt + σdWₜ
OU dynamics = "mean reverting" term + "white noise" term
"""
function evolve_OU_season(test_set, ycol::Symbol, scaled_ycol::Symbol,
    acf_dres1::AbstractArray, season_table::AbstractNDSparse,
    μσs::AbstractNDSparse, _eltype::DataType,
    input_size::Integer, output_size::Integer,
    test_dates::Array{ZonedDateTime,1},
    output_prefix::String) where I <: Integer

    acf_df = DataFrame(time = collect(0:(length(acf_dres1) - 1)), corr = acf_dres1)
    org_intT = compute_inttscale(acf_df[:, :time], acf_dres1)
    # not to log negative value when curve fitting
    posIntT = max(1, Int(round(org_intT)) - 1)
    org_fit = exp_fit(acf_df[1:posIntT, :time], acf_df[1:posIntT, :corr])
    #@info "Coef a and b in a * exp(b * x) for Annual Residual.. to 0:$(posIntT)", fit[1], fit[2]

    fitted_fit = exp_fit(acf_df[1:posIntT, :time], acf_df[1:posIntT, :corr])
    fit_x = acf_df[1:posIntT, :time]
    fit_y = fitted_fit[1] .* exp.(fit_x .* fitted_fit[2])
    intT = compute_inttscale(fit_x, fit_y)

    # ycol's mean / std
    μ_ycol = μσs[String(ycol) * "_res", "μ"].value
    σ_ycol = μσs[String(ycol) * "_res", "σ"].value

    # Time scale
    # T(hour), dt = 1
    T = intT
    Θ = 1.0 / T
    # becuase it's zscored, original μ and σ is 0 and 1.
    μ = 0.0
    σ = sqrt(1.0 * 2.0 / T)
    dt = 1

    # Vasicek Model
    # https://en.wikipedia.org/wiki/Vasicek_model

    # Check
    # https://github.com/SciML/DiffEqNoiseProcess.jl/blob/master/src/ornstein_uhlenbeck.jl
    ou_res_df = DataFrame(date = DateTime[], offset = _eltype[], y = _eltype[], yhat = _eltype[])

    for df_tset in test_set
        output_dates = DateTime(df_tset[1, :date], Local):Dates.Hour(1):(DateTime(df_tset[1, :date], Local) + Dates.Hour(output_size))

        X₀ = df_tset[1, scaled_ycol]
        OU = OrnsteinUhlenbeckProcess(Θ, μ, σ, 0, X₀)
        prob = NoiseProblem(OU, (0.0, output_size))
        Xₜ = solve(prob; dt=1.0)

        org_y = df_tset[1:output_size + 1, ycol]
        org_Xt = compose_seasonality(collect(output_dates[1:output_size + 1]),
            unzscore(Xₜ[1:output_size + 1], μ_ycol, σ_ycol),
            season_table)

        ou_res_df_tmp = DataFrame(date = output_dates, offset = 0:output_size, y = org_y, yhat = org_Xt)
        append!(ou_res_df, ou_res_df_tmp)
    end

    ou_res_df
end


"""
  module Multroot

Calculate mutiple roots of a real or complex polymomial exploiting the
multiplicity structure of the root.
Using methods develped by Zhang et. al.
"""
module Multroot

using Printf
using LinearAlgebra

export multroot, gcdroot, pejroot
export sylves, sylves1, sylmat, cauchymt, scalsq
export hqrt, hessqr, forsub, backsub
export zminsv, zminsv1
export polyzeros, polymult, polytransform

include("polyzeros.jl")

include("sylves.jl")
include("sylmat.jl")
include("polymult.jl")
include("zminsv.jl")
include("hessqr.jl")
include("hqrt.jl")
include("cauchymt.jl")
include("scalsq.jl")
include("forsub.jl")
include("backsub.jl")
include("gcdgn.jl")
include("multroot.jl")
include("pejroot.jl")
include("gcdroot.jl")

include("polyscale.jl")
include("horner.jl")

end # module


#See Philippe, *Simulation of right and left truncated gamma
# distributions by mixtures* and Dagpunar *Sampling of variates from
# a truncated gamma distribution*
# for pdf, cdf and quantile only right truncation is supported, left bound is 0.0
# for rand both left and right truncation is supported

TruncatedGamma(α::Float64, θ::Float64, lower::Float64, upper::Float64) =
    Distributions.Truncated(Distributions.Gamma(α, θ), lower, upper)
TruncatedGamma(α::Float64, θ::Float64, lower::Real, upper::Real) =
    TruncatedGamma(Float64(α), Float64(θ), lower, Float64(upper))

# ### statistics

function Distributions.minimum(
    d::Distributions.Truncated{Distributions.Gamma{T},Distributions.Continuous},
) where {T<:Real}
    min(d.lower, Distributions.minimum(d.untruncated))
end
function Distributions.maximum(
    d::Distributions.Truncated{Distributions.Gamma{T},Distributions.Continuous},
) where {T<:Real}
    max(d.lower, Distributions.maximum(d.untruncated))
end

# density of the right truncated gamma distribution
function Distributions.pdf(
    d::Distributions.Truncated{Distributions.Gamma{T},Distributions.Continuous},
    x::Real,
) where {T<:Real}
    if minimum(d.untruncated) == d.lower && maximum(d.untruncated) == d.upper
        Distributions.pdf(Distributions.Gamma(d.untruncated.α, d.untruncated.θ), x)
    elseif minimum(d.untruncated) == d.lower
        # right truncated
        return Distributions.pdf(Distributions.Gamma(d.untruncated.α, d.untruncated.θ), x) /
               Distributions.cdf(
            Distributions.Gamma(d.untruncated.α, d.untruncated.θ),
            d.upper,
        )
    else
        # left truncated
        println("left truncation still not supported.")
        return nothing
    end
end

function Distributions.cdf(
    d::Distributions.Truncated{Distributions.Gamma{T},Distributions.Continuous},
    x::Real,
) where {T<:Real}
    if minimum(d.untruncated) == d.lower && maximum(d.untruncated) == d.upper
        Distributions.cdf(Distributions.Gamma(d.untruncated.α, d.untruncated.θ), x)
    elseif minimum(d.untruncated) == d.lower
        # right truncated
        if (x > d.upper)
            x = d.upper
        end
        return Distributions.cdf(d.untruncated, x) /
               Distributions.cdf(d.untruncated, d.upper)
    else
        # left truncated
        println("left truncation still not supported.")
        return nothing
    end
end

function Distributions.quantile(
    d::Distributions.Truncated{Distributions.Gamma{T},Distributions.Continuous},
    x::Real,
) where {T<:Real}
    if minimum(d.untruncated) == d.lower && maximum(d.untruncated) == d.upper
        Distributions.quantile(Distributions.Gamma(d.untruncated.α, d.untruncated.θ), x)
    elseif minimum(d.untruncated) == d.lower
        # right truncated
        if (d.untruncated.α == 0.0)
            return 0.0
        end

        val = x * Distributions.cdf(d.untruncated, d.upper)
        return Distributions.quantile(d.untruncated, val)
    else
        # left truncated
        println("left truncation still not supported.")
        return nothing
    end
end

#Sampling

# optimal number of components for p = 0.95 fixed
function ncomp_optimal(θ::Float64)
    q = 1.644853626951
    ans = 0.25 * (q * sqrt(q * q + 4.0 * θ)^2)
    return trunc(Int64, ans)::Int64
end

#random number generation from the gamma with right truncation point t = 1,
# i.e. TG^-(α,θ,0.0,1.0)
function rng_right_tgamma_t1(
    rng::Distributions.AbstractRNG,
    d::Distributions.Truncated{Distributions.Gamma{T},Distributions.Continuous},
) where {T<:Real}
    n = ncomp_optimal(d.untruncated.θ)
    wl = zeros(Float64, n + 2)
    wlc = similar(wl)

    wl[1] = 1.0
    wlc[1] = 1.0
    for i = 2:(n+1)
        wl[i] = wl[i-1] * d.untruncated.θ / (d.untruncated.α + i)
        wlc[i] = wlc[i-1] + wl[i]
    end
    wlc = wlc ./ wlc[end-1]
    y = 1.0
    yy = 1.0
    for i = 2:(n+1)
        yy *= d.untruncated.θ / i
        y += yy
    end
    cont = true
    z = 0.0
    while (cont)
        u = Distributions.rand(rng)
        j = 1
        while (u > wlc[j])
            j += 1
            x = Distributions.rand(rng, Distributions.Beta(d.untruncated.α, j))
            u = Distributions.rand(rng)
            z = 1.0
            zz = 1.0
            for i = 2:(n+1)
                zz *= (1 - x) * d.untruncated.θ / i
                z += zz
            end
            z = _exp_c(-d.untruncated.θ * x) * y / z
            if (u <= z)
                cont = false
            end
        end
    end
    return z::Float64
end

function ltgamma(
    rng::Distributions.AbstractRNG,
    d::Distributions.Truncated{Distributions.Gamma{T},Distributions.Continuous},
) where {T<:Real}
    trunc = d.lower
    shape = d.untruncated.α
    b = 1 / d.untruncated.θ * trunc

    if (shape == 1)
        return Random.randexp(rng) / (1 / d.untruncated.θ) + trunc
    end

    d1 = b - shape
    d3 = shape - 1
    c0 = 0.5 * (d1 + sqrt(d1 * d1 + 4 * b)) / b

    x = 0.0
    accept = false
    while (!accept)
        x = b + Random.randexp(rng) / c0
        u = Distributions.rand(rng)
        l_rho = d3 * _log_c(x) - x * (1 - c0)
        l_M = d3 * _log_c(d3 / (1 - c0)) - d3
        accept = _log_c(u) <= (l_rho - l_M)
    end

    return trunc * (x / b)
end

# random number generation from the gamma with right truncation point#
# i.e. TG^-(α,θ,0.0,upper)
function Distributions.rand(
    rng::Distributions.AbstractRNG,
    d::Distributions.Truncated{Distributions.Gamma{T},Distributions.Continuous},
) where {T<:Real}
    if minimum(d.untruncated) == d.lower && maximum(d.untruncated) == d.upper
        Distributions.rand(rng, Distributions.Gamma(d.untruncated.α, d.untruncated.θ))
    elseif minimum(d.untruncated) == d.lower
        # right truncated
        d1 = TruncatedGamma(d.untruncate.α, d.untruncated.θ * d.upper, 0.0, d.upper)
        return rng_right_tgamma_t1(rng, d1) * d.upper
    else
        # left truncated
        return ltgamma(rng, d)
    end
end


import Base.var
export var

function var(x::AffineExpr)
  return sum_squares(mean(x) - x) / (x.size[1] * x.size[2])
end


module Statulator

include("calculations.jl")

function input(prompt::AbstractString="")
    print(prompt)
    return chomp(readline())
end

function terminate(message::AbstractString="")
    println(message)
    exit()
 end

 q1 = input("Is this is a problem of sample proportions or means? [p/m] ")

if q1 != "p" && q1 != "m"
    terminate("You must enter 'p' for proportion or 'm' for mean.")
end

q2 = input("Do you want to run a test or create an interval? [t/i] ")

if q2 != "t" && q2 != "i"
    terminate("You must enter 't' for test or 'i' for interval.")
end

q3 = input("How many proportions are you working with? [1/2] ")

if q3 != "1" && q3 != "2"
    terminate("Only one and two proportion tests and intervals are currently supported.")
end

if q1 == "p"
    if q2 == "t"
        if q3 == "1"
            p = 0
            p0 = 0
            n = 0
            alpha = 0
            alt = ""

            try
                p = parse(Float64, input("What is your sample proportion? "))
                p0 = parse(Float64, input("What is your assumed population proportion? "))
                n = parse(Float64, input("What is your sample size? "))
                alpha = parse(Float64, input("What is your alpha level? "))
                alt = input("Is your alternate hypothesis '!=', '>', or '<'? ")
            catch
                terminate("You must enter a numeric value.")
            end

            if alt != "!=" && alt != ">" && alt != "<"
                terminate("That wasn't an option.")
            end

            p_value = one_prop_z_test(p, p0, n, alpha, alt, p_pooled)
            println("The calculated p-value is ", p_value)
            println("If the population proportion is ", p0, " then the probability of observing a sample proportion of ", p, " is ", p_value, ".")
        elseif q3 == "2"
            p1 = 0
            n1 = 0
            p2 = 0
            n2 = 0
            alpha = 0
            alt = ""

            try
                p1 = parse(Float64, input("What is your first sample proportion? "))
                n1 = parse(Float64, input("What is your first sample size? "))
                p2 = parse(Float64, input("What is your second sample proportion? "))
                n2 = parse(Float64, input("What is your second sample size? "))
                alpha = parse(Float64, input("What is your alpha level?"))
                alt = input("Is your alternate hypothesis '!=', '>', or '<'? ")
            catch
                terminate("You must enter a numeric value.")
            end

            if alt != "!=" && alt != ">" && alt != "<"
                terminate("That wasn't an option.")
            end

            p_value = two_prop_z_test(p1, n1, p2, n2, alpha, alt)
            println("The calculated p-value is ", abs(p_value))
            println("If the difference between the population proportions is 0, then the probability of observing a sample proportion difference of ", abs(p1-p2), " is ", abs(p_value), ".")
        end
    elseif q2 == "i"
        if q3 == "1"
            p = 0
            n = 0
            lvl = 0

            try
                p = parse(Float64, input("What is your sample proportion? "))
                n = parse(Float64, input("What is your sample size? "))
                lvl = parse(Float64, input("What is your needed confidence level? "))
            catch
                terminate("You must enter a numeric value.")
            end

            interval = one_prop_z_int(p, n, lvl)

            println("The interval is (", interval.lower, ", ", interval.upper, ")")
            println("We can state with ", (lvl*100), "% confidence that the true proportion lies between ", interval.lower, " and ", interval.upper, ".")
        elseif q3 == "2"
            p1 = 0
            n1 = 0
            p2 = 0
            n2 = 0
            lvl = 0

            try
                p1 = parse(Float64, input("What is your first sample proportion? "))
                n1 = parse(Float64, input("What is your first sample size? "))
                p2 = parse(Float64, input("What is your second sample proportion? "))
                n2 = parse(Float64, input("What is your second sample size? "))
                lvl = parse(Float64, input("What is your needed confidence level? (in proportion)"))
            catch
                terminate("You must enter a numeric value.")
            end

            interval = two_prop_z_int(p1, n1, p2, n2, lvl)

            println("The interval is (", interval.lower, ", ", interval.upper, ")")
            println("We can state with ", (lvl*100), "% confidence that the difference in the true proportion lies between ", interval.lower, " and ", interval.upper, ".")
        end
    end
elseif q1 == "m"
    if q2 == "t"
        if q3 == "1"
            x̄ = 0
            μ0 = 0
            sx = 0
            n = 0
            alpha = 0
            alt = ""

            try
                x̄ = parse(Float64, input("What is your sample mean? "))
                μ0 = parse(Float64, input("What is your assumed population mean? "))
                sx = parse(Float64, input("What is your sample standard deviation? "))
                n = parse(Float64, input("What is your sample size? "))
                alpha = parse(Float64, input("What is your alpha level? "))
                alt = input("Is your alternate hypothesis '!=', '>', or '<'? ")
            catch
                terminate("You must enter a numeric value.")
            end

            if alt != "!=" && alt != ">" && alt != "<"
                terminate("That wasn't an option.")
            end

            p_value = z_test(p, p0, sx, n, alpha, alt)
            println("The calculated p-value is ", p-value)
        elseif q3 == "2"
            x̄1 = 0
            sx1 = 0
            n1 = 0
            x̄2 = 0
            sx2 = 0
            n2 = 0
            alpha = 0
            alt = ""

            try
                x̄1 = parse(Float64, input("What is your first sample mean? "))
                sx1 = parse(Float64, input("What is your first sample standard deviation? "))
                n1 = parse(Float64, input("What is your first sample size? "))
                x̄2 = parse(Float64, input("What is your second sample mean? "))
                sx2 = parse(Float64, input("What is your second sample standard deviation? "))
                n2 = parse(Float64, input("What is your second sample size? "))
                alpha = parse(Float64, input("What is your alpha level? "))
                alt = input("Is your alternate hypothesis '!=', '>', or '<'? ")
            catch
                terminate("You must enter a numeric value.")
            end

            if alt != "!=" && alt != ">" && alt != "<"
                terminate("That wasn't an option.")
            end

            p_value = two_sample_z_test(x̄1, sx1, n1, x̄2, sx2, n2, alpha, alt)
            println("The calculated p-value is ", p_value)
            println("If the difference between the population proportions is 0, then the probability of observing a sample proportion difference of ", abs(x̄1 - x̄2), " is ", abs(p_value), ".")
        end
    elseif q2 == "i"
        if q3 == "1"
            x̄ = 0
            sx = 0
            n = 0
            lvl = 0

            try
                x̄ = parse(Float64, input("What is your sample mean? "))
                sx = parse(Float64, input("What is your sample standard deviation? "))
                n = parse(Float64, input("What is your sample size? "))
                lvl = parse(Float64, input("What is your needed confidence level? "))
            catch
                terminate("You must enter a numeric value.")
            end

            interval = z_interval(x̄, sx, n, lvl)

            println("The interval is (", interval.lower, ", ", interval.upper, ")")
            println("We can state with ", (lvl*100), "% confidence that the true mean lies between ", interval.lower, " and ", interval.upper, ".")
        elseif q3 == "2"
            x̄1 = 0
            sx1 = 0
            x̄2 = 0
            sx2 = 0
            n1 = 0
            n2 = 0
            lvl = 0

            try
                x̄1 = parse(Float64, input("What is your first sample mean? "))
                sx1 = parse(Float64, input("What is your first sample standard deviation? "))
                n1 = parse(Float64, input("What is your first sample size? "))
                x̄2 = parse(Float64, input("What is your second sample mean? "))
                sx2 = parse(Float64, input("What is your second sample standard deviation? "))
                n2 = parse(Float64, input("What is your second sample size? "))
                lvl = parse(Float64, input("What is your needed confidence level? "))
            catch
                terminate("You must enter a numeric value.")
            end

            interval = two_sample_z_interval(x̄1, sx1, n1, x̄2, sx2, n2, lvl)

            println("The interval is (", interval.lower, ", ", interval.upper, ")")
            println("We can state with ", (lvl*100), "% confidence that the difference between the true means lies between ", interval.lower, " and ", interval.upper, ".")
        end
    end
end

end


using DataDecomposition;
using MLKernels;
nc = 3;
p = 3;
k = 4;
y = repeat(Int64[i for i = 1:k],inner=[nc],outer=[1]);
X = vcat([rand(nc,p) .+ 10rand(1,3) for i = 1:k]...);
freq = Float64[1/k for i=1:k];
lda!(copy(X), y, k, freq)


function build_A_sp(kx,U,E,Dy,Dy2,y,parms)
    U1 = U
    h1 = E.+parms.H

    # Define Differentiation Matrices for y: FD second order

    Ny  = length(U1)-1
    g1  = parms.g1

    #  Define Matrices to set up EVP

    dU1 = Diagonal(Dy*U1);
    dh1 = Diagonal(Dy*h1);

    F    = Diagonal(parms.f0.+0*y);
    Id    = Matrix{Float64}(I, Ny+1, Ny+1);
    h1f  = Diagonal(h1);

    kx2  = kx^2;
    ikx  = 1im*kx;

    # Form 1L-RSW Matrix
    #   [u1, v1, h1]

    A = [[Diagonal(U1),      -F[:,2:Ny].+dU1[:,2:Ny],                g1*I],
        [-F[2:Ny,:]/kx2,         Diagonal(U1[2:Ny]),  -g1/kx2*Dy[2:Ny,:]],
        [Diagonal(h1),       Dy*h1f[:,2:Ny],                    Diagonal(U1)]];
    return sparse(A)
end


#using JuMP, CPLEX
using JuMP, GLPK

function jobshop(n, m, p, r)
  #model = Model(CPLEX.Optimizer)
  model = Model(GLPK.Optimizer)
  set_optimizer_attribute(model, "msg_lev", GLPK.GLP_MSG_ALL)
  @variable(model,s[1:n,1:m]>=0)
  @variable(model,C>=0)
  @objective(model, Min, C)
  M=sum(p)
  #print(M,"\n")

  for i in 1:n
    for j in 1:m-1
      @constraint(model,s[i,j]+p[i,j]<=s[i,j+1])
    end
    @constraint(model,s[i,m]+p[i,m]<=C)
  end
  for i in 1:n
    for ii in i+1:n
      for j in 1:m
        for jj in 1:m
          if ruta[i,j]==ruta[ii,jj]
            z=@variable(model,base_name ="z_$(i)_$(ii)_$(j)",binary=true)
            @constraint(model,s[i,j]+p[i,j]-s[ii,jj]<=z*M)
            @constraint(model,s[ii,jj]+p[ii,jj]-s[i,j]<=(1-z)*M)
          end
        end
      end
    end
  end
  #println(model)
  optimize!(model)
  println("Objective: ",objective_value(model))
#  for i in 1:n
#    for j in 1:m
#      println(i,"\t",j,"\t",value(s[i,j]))
#    end
#  end
end

## leo archivo y guardo datos
f = open(ARGS[1], "r")
s = readlines(f)
print("filename:", s)
n=parse(Int,s[2])
m=parse(Int,s[4])
ruta=zeros(Int, n, m)
processing=zeros(Int,n,m)
print("número piezas ",n,"\nnúmero de máquinas ",m,"\n")
print("ruta:\n")
for i in 1:n
  divided=split(s[5+i])
  for j in 1:m
    ruta[i,j]=parse(Int,divided[j])
    print(ruta[i,j],"\t")
  end
  print("\n")
end
print("duraciones:\n")
for i in 1:n
  divided=split(s[6+n+i])
  for j in 1:m
    processing[i,j]=parse(Int,divided[j])
    print(processing[i,j],"\t")
  end
  print("\n")
end
close(f)
jobshop(n,m,processing,ruta)


struct DynamicBudgetedBipartiteMatchingSolution{T, U} <:
       BudgetedBipartiteMatchingSolution{T, U}
    instance::BudgetedBipartiteMatchingInstance{T, U}
    solution::Vector{Edge{T}}
    value::Float64
    state::Dict{Tuple{Int, Int, Int, Int}, Float64} # Data structure built by the dynamic-programming recursion.
    solutions::Dict{Tuple{Int, Int, Int, Int}, Vector{Edge{T}}} # From the indices of state to the corresponding solution.
end

function matching_dp_budgeted(
    instance::BudgetedBipartiteMatchingInstance{T, U},
) where {T, U}
    # First, determine the size of a perfect matching.
    nr = Dict(k => 1.0 for (k, v) in instance.matching.reward)
    ni = BipartiteMatchingInstance(instance.matching, nr)
    ns = matching_dp_imperfect(ni)
    m = length(ns.solution)

    # Solve this budgeted matching to ensure the constraints are met.

    # Four indices:
    # - two vertices, like matching_dp_imperfect;
    # - the number of edges still to add to have a perfect matching, like matching_dp;
    # - the minimum budget to respect.
    state = Dict{Tuple{Int, Int, Int, Int}, Float64}(
        (i, j, k, l) => -Inf for i in instance.matching.vertex_left,
        j in instance.matching.vertex_right, k in 0:m,
        l in 0:(instance.budget)
    )
    solutions = Dict{Tuple{Int, Int, Int, Int}, Vector{Edge{T}}}(
        (i, j, m, zero(U)) => Edge{T}[] for
        i in instance.matching.vertex_left,
        j in instance.matching.vertex_right
    )

    # Initialise: k=m-1, l=0, i.e. take at most one edge, don't consider the budget yet.
    v1 = first(instance.matching.vertex_right)
    for i in instance.matching.vertex_left
        if has_edge(instance.matching.graph, i, v1)
            state[i, v1, m - 1, 0] = instance.matching.reward[Edge(i, v1)]
            solutions[i, v1, m - 1, 0] = [Edge(i, v1)]
        end
    end

    v1 = first(instance.matching.vertex_left)
    for j in instance.matching.vertex_right
        if has_edge(instance.matching.graph, v1, j)
            state[v1, j, m - 1, 0] = instance.matching.reward[Edge(v1, j)]
            solutions[v1, j, m - 1, 0] = [Edge(v1, j)]
        end
    end

    for (i, prev_i) in zip(
        instance.matching.vertex_left[2:end],
        instance.matching.vertex_left[1:(end - 1)],
    )
        for (j, prev_j) in zip(
            instance.matching.vertex_right[2:end],
            instance.matching.vertex_right[1:(end - 1)],
        )
            # Take (i, j).
            if has_edge(instance.matching.graph, i, j)
                state[i, j, m - 1, 0] = instance.matching.reward[Edge(i, j)]
                solutions[i, j, m - 1, 0] = [Edge(i, j)]
            end

            # Don't take (i, j): copy a previous solution. No mixing possible, as only one edge is considered.
            if state[prev_i, j, m - 1, 0] > state[i, j, m - 1, 0]
                state[i, j, m - 1, 0] = state[prev_i, j, m - 1, 0]
                solutions[i, j, m - 1, 0] = solutions[prev_i, j, m - 1, 0]
            end
            if state[i, prev_j, m - 1, 0] > state[i, j, m - 1, 0]
                state[i, j, m - 1, 0] = state[i, prev_j, m - 1, 0]
                solutions[i, j, m - 1, 0] = solutions[i, prev_j, m - 1, 0]
            end
        end
    end

    # Initialisation for l=0, k < m-1: create perfect matchings, regardless of budget.
    for k in (m - 2):-1:0
        # Don't handle the cases with just one vertex in one part of the graph:
        # this loop requires at least two edges to be taken, but the considered
        # subset of edges can lead to at most only one being taken.

        # Consider the other vertices.
        for (i, prev_i) in zip(
            instance.matching.vertex_left[2:end],
            instance.matching.vertex_left[1:(end - 1)],
        )
            for (j, prev_j) in zip(
                instance.matching.vertex_right[2:end],
                instance.matching.vertex_right[1:(end - 1)],
            )
                # Take (i, j).
                if has_edge(instance.matching.graph, i, j) &&
                   state[prev_i, prev_j, k + 1, 0] > -Inf
                    state[i, j, k, 0] =
                        instance.matching.reward[Edge(i, j)] +
                        state[prev_i, prev_j, k + 1, 0]
                    solutions[i, j, k, 0] = solutions[prev_i, prev_j, k + 1, 0]
                    push!(solutions[i, j, k, 0], Edge(i, j))
                end

                # Don't take (i, j): mix previous solutions (as a special case, maybe take a previous solution completely).
                for k2 in 0:m
                    for k3 in 0:m
                        if state[prev_i, j, k2, 0] + state[i, prev_j, k3, 0] >
                           state[i, j, k, 0]
                            potential_sol = union(
                                Set{Edge{T}}(solutions[prev_i, j, k2, 0]),
                                Set{Edge{T}}(solutions[i, prev_j, k3, 0]),
                            )
                            val = sum(
                                instance.matching.reward[Edge(src(e), dst(e))] for e in potential_sol
                            )

                            if length(potential_sol) == m - k &&
                               val > state[i, j, k, 0]
                                state[i, j, k, 0] = val
                                solutions[i, j, k, 0] = collect(potential_sol)
                            end
                        end
                    end
                end
            end
        end
    end

    # Main iteration, considering the first vertex of each side is already done.
    for l in 1:(instance.budget)
        # Repeat the loop as often as necessary to handle the corner case with zero-weight edges.
        while true
            changes = false

            # Deal with the first two vertices (one on each side). All constraints are imposed.
            # Just one edge to be taken, as one vertex is always fixed on one side:
            # no feasible solution for k < m-1.
            vl = first(instance.matching.vertex_left)
            vr = first(instance.matching.vertex_right)
            if has_edge(instance.matching.graph, vl, vr) &&
               instance.weight[Edge(vl, vr)] >= l
                state[vl, vr, m - 1, l] = instance.matching.reward[Edge(vl, vr)]
                solutions[vl, vr, m - 1, l] = [Edge(vl, vr)]
            end

            v1 = first(instance.matching.vertex_right)
            for i in instance.matching.vertex_left
                if has_edge(instance.matching.graph, i, v1) &&
                   instance.weight[Edge(i, v1)] >= l
                    state[i, v1, m - 1, l] =
                        instance.matching.reward[Edge(i, v1)]
                    solutions[i, v1, m - 1, l] = [Edge(i, v1)]
                end
            end

            v1 = first(instance.matching.vertex_left)
            for j in instance.matching.vertex_right
                if has_edge(instance.matching.graph, v1, j) &&
                   instance.weight[Edge(v1, j)] >= l
                    state[v1, j, m - 1, l] =
                        instance.matching.reward[Edge(v1, j)]
                    solutions[v1, j, m - 1, l] = [Edge(v1, j)]
                end
            end

            # Do the rest.
            for k in (m - 1):-1:0
                for (i, prev_i) in zip(
                    instance.matching.vertex_left[2:end],
                    instance.matching.vertex_left[1:(end - 1)],
                )
                    for (j, prev_j) in zip(
                        instance.matching.vertex_right[2:end],
                        instance.matching.vertex_right[1:(end - 1)],
                    )

                        # Take (i, j) if possible.
                        if has_edge(instance.matching.graph, i, j)
                            remaining_budget =
                                max(0, l - instance.weight[Edge(i, j)])

                            if state[prev_i, prev_j, k + 1, remaining_budget] >
                               -Inf
                                state[i, j, k, l] =
                                    instance.matching.reward[Edge(i, j)] +
                                    state[
                                        prev_i,
                                        prev_j,
                                        k + 1,
                                        remaining_budget,
                                    ]
                                solutions[i, j, k, l] = solutions[
                                    prev_i,
                                    prev_j,
                                    k + 1,
                                    remaining_budget,
                                ]
                                push!(solutions[i, j, k, l], Edge(i, j))

                                changes = true
                            end
                        end

                        # Don't take (i, j): mix previous solutions (as a special case, maybe take a previous solution completely).
                        for k2 in 0:m
                            for k3 in 0:m
                                if state[prev_i, j, k2, l] +
                                   state[i, prev_j, k3, l] > state[i, j, k, l]
                                    potential_sol = union(
                                        Set{Edge{T}}(
                                            solutions[prev_i, j, k2, l],
                                        ),
                                        Set{Edge{T}}(
                                            solutions[i, prev_j, k3, l],
                                        ),
                                    )
                                    val = sum(
                                        instance.matching.reward[Edge(
                                            src(e),
                                            dst(e),
                                        )] for e in potential_sol
                                    )
                                    wei = sum(
                                        instance.weight[Edge(src(e), dst(e))] for e in potential_sol
                                    )

                                    if length(potential_sol) == m - k &&
                                       val > state[i, j, k, l] &&
                                       wei >= l
                                        state[i, j, k, l] = val
                                        solutions[i, j, k, l] =
                                            collect(potential_sol)

                                        changes = true
                                    end
                                end
                            end
                        end
                    end
                end
            end

            # If this iteration made no change, stop and go on to the next budget.
            if !changes
                break
            end
        end
    end

    val = state[
        last(instance.matching.vertex_left),
        last(instance.matching.vertex_right),
        0,
        instance.budget,
    ]
    sol = solutions[
        last(instance.matching.vertex_left),
        last(instance.matching.vertex_right),
        0,
        instance.budget,
    ]

    return DynamicBudgetedBipartiteMatchingSolution(
        instance,
        sol,
        val,
        state,
        solutions,
    )
    # return SimpleBudgetedBipartiteMatchingSolution(instance, sol)
end


cfl = 0.45
c_omp = 100.0 #skin depth in cell units
delgam = 1.0e-3
Nx = 400
Rm = 2
Nt = 10000
gamma = 0.25

dx = 1/c_omp
dt = cfl*dx

#box size and lambda length
Lx = Nx*dx
lambda = Lx/Rm
lambda_per_cell = lambda/dx

#thermal speed & debye length
vth = sqrt(delgam)
deb = vth
deb_per_cell = deb/dx

#beam velocity
vb = gamma/vth

kbar = 2pi/lambda
khat = (2pi/lambda)*vth

T = Nt*dt

println("dx                   : $dx")
println("dt                   : $dt")
println("cells per skin depth : $c_omp")
println("gamma                : $gamma")
println("thermal velocity     : $vth")
println("Debye length         : $deb")
println("debye per cell       : $deb_per_cell")
println("vb                   : $vb")

println("modes                : $Rm")
println("lambda               : $lambda")
println("lambda per cell      : $lambda_per_cell")
println("kbar                 : $kbar")
println("khat                 : $khat")

println("Lx                   : $Lx")
println("Nx                   : $Nx")
println("T                    : $T")
println("Nt                   : $Nt")

ompr = 1 + (3/2)*khat^2 + (15/8)*khat^4 + (147/16)*khat^6
ompi = -0.5*sqrt(pi/2)*(1/khat^3 - 6*khat)*exp(-0.5/khat^2 - 3/2 -3*khat^2 -12*khat^4)

println("omp Re: $ompr")
println("omp Im: $ompi")


import LinearAlgebra: Givens, givensAlgorithm
import LinearAlgebra.ldiv!

struct FastHessenberg{T<:AbstractMatrix}
    H::T # H is assumed to be Hessenberg of size (m + 1) × m
end

@inline Base.size(H::FastHessenberg, args...) = size(H.H, args...)

"""
Solve Hy = rhs for a non-square Hessenberg matrix.
Note that `H` is also modified as is it converted
to an upper triangular matrix via Given's rotations
"""
function ldiv!(H::FastHessenberg, rhs)
    # Implicitly computes H = QR via Given's rotations
    # and then computes the least-squares solution y to
    # |Hy - rhs| = |QRy - rhs| = |Ry - Q'rhs|

    width = size(H, 2)

    # Hessenberg -> UpperTriangular; also apply to r.h.s.
    @inbounds for i = 1 : width
        c, s, _ = givensAlgorithm(H.H[i, i], H.H[i + 1, i])

        # Skip the first sub-diagonal since it'll be zero by design.
        H.H[i, i] = c * H.H[i, i] + s * H.H[i + 1, i]

        # Remaining columns
        @inbounds for j = i + 1 : width
            tmp = -conj(s) * H.H[i, j] + c * H.H[i + 1, j]
            H.H[i, j] = c * H.H[i, j] + s * H.H[i + 1, j]
            H.H[i + 1, j] = tmp
        end

        # Right hand side
        tmp = -conj(s) * rhs[i] + c * rhs[i + 1]
        rhs[i] = c * rhs[i] + s * rhs[i + 1]
        rhs[i + 1] = tmp
    end

    # Solve the upper triangular problem.
    U = UpperTriangular(view(H.H, 1 : width, 1 : width))
    ldiv!(U, view(rhs, 1 : width))
    nothing
end


using JuMP, Gurobi, Cbc
import XLSX
using BenchmarkTools

# make sure u set the working directory to current files folder
data = XLSX.readxlsx("../data/part1/data_numP100_numC25_10_05_23_17.xlsx")

#General
General = data["General"];
l = General["B1"]
num_providers = General["B2"]
num_clients = General["B3"]

#Providers
Providers = data["Providers"];
provCostString  = string("A2:A", num_providers+1)
provStartString = string("B2:B", num_providers+1)
f = Providers[provCostString]
w = Providers[provStartString]

#Providers
Clients = data["Clients"];
clientStartString  = string("A2:A", num_clients+1)
clientDurationString  = string("B2:B", num_clients+1)
s_times = Clients[clientStartString]
d = Clients[clientDurationString]

#Distances
Distances = data["Distances"];
t = Distances[Distances.dimension]

# need to convert to floats to use 'sortperm' below
w = convert(Array{Float64,1}, w[:,1])
s_times = convert(Array{Float64,1}, s_times[:,1])
d = convert(Array{Float64,1}, d[:,1])
f = convert(Array{Float64,1}, f[:,1])

client_sorted = sortperm(s_times)  # client start times sorted (in index)
arr = [[f[i], w[i]] for i in 1:length(w)]
provider_sorted = sortperm(f) # providers cost and start time sorted

# heuristic function that computes feasible routes (but likely not optimal, i.e. minimum cost)
function heuristics_result(provider_sorted, client_sorted, w, s_times, t, d)
    past_client = zeros(0)
    num_served = zeros(num_providers, 2)
    num_served[:,1] = provider_sorted

    for m in provider_sorted  # for each index in sorted provider, and then for each client
        current_location = 1  # 1 is at depot(14HS) currently
        for n in client_sorted
            # only proceed if this client has not been assigned to a provider yet.
            if n ∉ past_client
                currentTime_provider = w[m]
                currentTime_client = s_times[n]
                travelTime = t[current_location, n+1] # n+1 is nth client in the t matrix

                #time to travel from provider's current location to the client plus the time when the provider is free.
                time_to_arrive = travelTime + currentTime_provider

                #Time to serve client "n".
                duration = d[n]
                #Time finished serving client "n".
                finish_time = currentTime_client + duration

                #Ensure the provider can arrive in time (at or earlier than service start times)
                # Must return back the origin (14HS) by hour l at latest.
                if currentTime_client >= time_to_arrive && finish_time <= l - t[n+1, 1]
                        println("Provider: ", m, " is assigned to client: ", n)
                        num_served[findfirst(isequal(m), num_served[:,1]), 2] += 1
                        w[m] = finish_time # update the time of provider m
                        current_location = n+1  # n+1 is nth client in the t-matrix
                        append!(past_client, n) # record service provided for client n
                end
            end
        end
        # end the algorithm if we have served all clients. so O(N*<M)?
        if length(past_client) == length(client_sorted)
            break
        end
    end
    println("\n")
    return convert(Array{Int,2}, num_served), convert(Array{Int,1}, past_client)
end

function compute_cost(num_served)
    cost = 0
    for i in 1:length(num_served[:,2])
        if num_served[:,2][i] > 0
            cost += f[num_served[:,1][i]]
        end
    end
    return cost
end

function is_feasible(past_client)
    return length(past_client) == length(client_sorted)
end

function get_heuristics_path(num_served, past_client, k)
    path = zeros(1)
    k_ind = findfirst(isequal(k), num_served[:,1])
    if k_ind == 1
        for i in 1:num_served[1,2]
            append!(path, past_client[i])
        end
    end
    if k_ind != 1
        start_ind = sum(num_served[1:k_ind-1,2]) + 1
        for i in start_ind:start_ind+num_served[k_ind,2]-1
            append!(path, past_client[i])
        end
    end
    return convert(Array{Int,1}, append!(path, zeros(1))) # convert to integer array
end

num_served, past_client = heuristics_result(provider_sorted, client_sorted, copy(w), s_times, t, d)
num_served
println(client_sorted)

#@btime heuristics_result(provider_sorted, client_sorted, copy(w), s_times, t, d)

length(past_client)
is_feasible(past_client)
compute_cost(num_served)

function get_path(num_served, past_client, k)
    path = get_heuristics_path(num_served, past_client, k)
    path += convert(Array{Int,1}, ones(length(get_heuristics_path(num_served, past_client, k))))
    return (path)
end

function get_workers(num_served)
    workers = zeros(0)
    for i in 1:size(num_served)[1]
        if num_served[i,:][2] != 0
            append!(workers, num_served[i,:][1])
        end
    end
    return workers
end

path_dict = Dict(trunc(Int, get_workers(num_served)[i]) => get_path(num_served, past_client, get_workers(num_served)[i]) for i in 1:length(get_workers(num_served)))

################### MIP ######################################
################### MIP ######################################
################### MIP ######################################
################### MIP ######################################
################### MIP ######################################
using JuMP, Gurobi, Cbc
import XLSX
using BenchmarkTools

#### read data ####
#xf = XLSX.readxlsx("../data/part1/data_numP10_numC10_09_05_13_29.xlsx")
xf = XLSX.readxlsx("../data/part1/data_numP200_numC150_11_05_34_00.xlsx")

l = xf["General!B1"]     # work day duration
N = xf["General!B3"] + 1  # number of clients + 1
M = xf["General!B2"]  # number of providers
# + 1 is for including the headquarter (HS14)

# Clients
s = xf["Clients!A2:A$N"]   # client service start times
d = vcat([0], xf["Clients!B2:B$N"])   # durations for clients, set d[1] = 0 for HS14

# Providers
w = xf["Providers!B2:B$(M+1)"]  # start time of providers
f = xf["Providers!A2:A$(M+1)"]  # hiring cost of providers

# make separate s vector for each provider, inserting into the first element the provider's work start time
s_test = ones(M, N)
for k in 1:M
    global s_test[k,:] = vcat([w[k]], s)
end
s = s_test

# Distances
t = xf["Distances"][:]  # 2-d array of travel times
##########

model = Model();
set_optimizer(model, Gurobi.Optimizer);
set_time_limit_sec(model, 300.0);

#(13)
@variable(model, x[1:N, 1:N, 1:M], Bin)

#################### setting start values from heursitic solution ###########

for i in 1:N
    for j in 1:N
        for k in 1:M
            set_start_value(x[i,j,k], 0)
        end
    end
end

for key in keys(path_dict)
    path = path_dict[key]
    for i in 1:length(path)-1
        set_start_value(x[path[i], path[i+1], key], 1)
    end
end

################################################################################

# subtour constraint variables using u
@variable(model, 1 <= u[1:N, 1:M] <= N-1)

######################## routing constraints ######################

# (2) each client served once by one provider
@constraint(model, served[j in 2:N], sum(x[i,j,k] for i in 1:N for k in 1:M) == 1)

# (3) route continuity
@constraint(model, continuity[p in 2:N, k in 1:M], sum(x[i,p,k] for i in 1:N) == sum(x[p,j,k] for j in 1:N) )

# (4) not all providers need to leave 14HS (and start working)
@constraint(model, leave_depot[k in 1:M], sum(x[1,j,k] for j in 2:N) <= 1)

# (7) at least 1 provider must leave 14HS and start working to serve clients
@constraint(model, sum(x[1,j,k] for j in 2:N for k in 1:M) >= 1)

#(9) subtour elimination for each provider
for k in 1:M
    for i in 2:N
        for j in 2:N
            @constraint(model, u[i,k] - u[j,k] + (N-1) * x[i,j,k] <= (N-1) - 1 )
        end
    end
end

######################## end routing constraints ######################

######################## time constraints ######################

# (10) initial client service times
@constraint(model, start[j in 2:N, k in 1:M], x[1,j,k]*(s[k,1] + t[1,j] + d[1]) <= s[k,j])

# (11) in between start and end of work, service times
@constraint(model, bet[i in 2:N, j in 2:N, k in 1:M], x[i,j,k]*(s[k,i] + t[i,j] + d[i]) <= s[k,j])

# (12) come back to 14HS by hour l.
@constraint(model, endofwork[i in 2:N, k in 1:M], x[i,1,k]*(s[k,i] + t[i,1] + d[i]) <= l)

######################## end time constraints ######################

# newer objective to minimise total cost based on hours of work.
### Note. the max() function inside the objective can only be solved by Gurobi solver (Cbc cannot solve)
#@objective(model, Min, sum(x[1,j,k]*f[k]*(sum(x[p,q,k]*(t[p,q] + d[p] + max(0, s[k,q] - s[k,p] - t[p,q] - d[p]) ) for p in 1:N for q in 1:N)) for j in 2:N for k in 1:M))

# original non-hourly objective for part 1
@objective(model, Min, sum(f[k]*x[1,j,k] for j in 1:N for k in 1:M))

optimize!(model)

@btime optimize!(model)

x[1,:,1]
# see which providers are chosen to provide service.
for k in 1:M
    if sum(value.(x[1,:,k])) == 1
        println("Provider $k provides service.")
        println("Provider $k costs $(f[k]) per hour.")
        # come back hr
        for i in 2:N
            if value.(x[i,1,k]) == 1
                println("Provider $k comes back to 14HS at $(t[i,1] + s[k,i] + d[i])hr. \n")
                break # added
            end
        end
    end
end

# Given N: number of clients+1, and provider number k, find the path
### So if you found that providers 2, 3, 4 and 5 are chosen to work, run get_path(N, 2), get_path(N, 3) etc.
function get_path(N, k)
    path = []
    for j in 2:N
        if value.(x[1,j,k]) == 1
            println("\nstart of path")
            println(1, " => ")
            path = vcat(path, [1])
            println(j, " => ")
            path = vcat(path, [j])
            i = j
            while value.(x[i,1,k]) == 0
                for m in 2:N
                    if value.(x[i,m,k]) == 1
                        println(m, " => ")
                        path = vcat(path, [m])
                        i = m
                        break
                    end
                end
            end
            println(1)
            path = vcat(path, [1])
            println("end of path.... or in client index below (not location)")
        end
    end

    for j in 2:N
        if value.(x[1,j,k]) == 1
            println("\nstart of path")
            println(0, " => ")
            println(j-1, " => ")
            i = j
            while value.(x[i,1,k]) == 0
                for m in 2:N
                    if value.(x[i,m,k]) == 1
                        println(m-1, " => ")
                        i = m
                        break
                    end
                end
            end
            println(0)
            println("end of path (in client index)")
        end
    end
    return path
end

k = 8
path = get_path(N, k)

function overview(path, k)
    println("$(s[k,1])hr: Becomes available, leaves to location $(path[2]).")
    for i in 2:length(path)-1
        println("$(s[k,path[i-1]]+t[path[i-1], path[i]] + d[path[i-1]])hr: Arrived at location $(path[i]).")
        if s[k,path[i-1]]+t[path[i-1], path[i]] + d[path[i-1]] != s[k,path[i]]
            println("   Arrived early... Wait until service start time.")
        end
        println("$(s[k,path[i]])hr: Serve for $(d[path[i]]) hrs.")
        println("$(s[k,path[i]]+d[path[i]])hr: Finished service. Leave to location $(path[i+1]).")
    end
    println("$(s[k,path[length(path)-1]] + d[path[length(path)-1]] + t[path[length(path)-1], 1])hr: Arrived at location 1 (14HS).\n")
end

# print the overview of the path for provider k, where 'path' is found above using get_path(N,k)
path = [1 7 1]
k = 2
overview(path, k)

# find the number of worked hours = provider come back time (to 14HS) -  provider start time (w_k)
worked_hr = s[k,path[length(path)-1]] + d[path[length(path)-1]] + t[path[length(path)-1], 1] - s[k,1]
println("Provider $k worked for $worked_hr hours.")


using PCCA
using Test

@testset "PCCA.jl" begin
    # Write your tests here.
	x = rand(10, 10)
	x = x ./ sum(x, dims=2)
	c = pcca(x, 2)

	function randomstochasticmatrix(n, reversible=true)
		P = rand(n,n)
		if reversible
			P = (P + P') / 2
		end
		P ./= sum(P, dims=2)
	end

	@testset "Reversible = $rev" for rev in [true, false]
		P = [randomstochasticmatrix(3+mod(i,12), rev) for i in 1:10]
		@testset "Method $method" for method in [PCCA.BaseSolver, PCCA.ArnoldiSolver, PCCA.KrylovSolver]
			for P in P, n in 2:size(P,1)-1
				χ = pcca(P, n, solver=method(), optimize=true)
				a = PCCA.crispassignments(χ)

				#@show χ
				@test all(sum(χ, dims=2) .≈ 1)
			end
		end
	end

end


# BUG: The current result does not align with R's psych package
using Distributions

function kappa(x, α = 0.05)
    tot = sum(x)
    x = x ./ tot        # convert to probabilities
    rs = sum(x, 2)      # row sum
    cs = sum(x, 1)      # column sum

    prob = rs * cs
    po = trace(x)
    pc = trace(prob)
    kappa = (po - pc)/(1 - pc)

    var_kappa = (1 / (tot*(1-pc)^4)) * (trace(x * (I * (1-pc) - (rs .+ (cs))*(1-po))^2 ) + (1-po)^2 * (sum(x * (cs .+ (rs))^2) - trace(x * (cs .+ (rs))^2))  -(po*pc - 2*pc +po)^2)

    lower_ci = kappa + quantile(Normal(0, 1), (α/2)) * sqrt(var_kappa)
    upper_ci = kappa - quantile(Normal(0, 1), (α/2)) * sqrt(var_kappa)

    return (kappa, var_kappa, lower_ci, upper_ci)
end

# kappa(cm1.matrix)

# a = cm1.matrix


include("utils.jl")

@testset "Cholesky Space" begin
    M = Manifolds.CholeskySpace(3)

    types = [
        Matrix{Float64},
        MMatrix{3, 3, Float64},
        Matrix{Float32},
    ]
    for T in types
        A(α) = [1. 0. 0.; 0. cos(α) sin(α); 0. -sin(α) cos(α)]
        ptsF = [#
            cholesky([1. 0. 0.; 0. 1. 0.; 0. 0. 1]).L,
            cholesky([2. 0. 0.; 0. 2. 0.; 0. 0. 1]).L,
            cholesky(A(π/6) * [1. 0. 0.; 0. 2. 0.; 0. 0. 1] * transpose(A(π/6))).L,
            ]
        pts = [convert(T, a) for a in ptsF]
        test_manifold(M, pts;
            test_injectivity_radius = false,
            test_vector_transport = true,
            test_forward_diff = false,
            test_reverse_diff = false,
            exp_log_atol_multiplier = 8.0,
        )
    end
    @testset "Test Error cases in is_manifold_point and is_tangent_vector" begin
        pt1f = zeros(2,3); # wrong size
        pt2f = [1. 0. 0.; 0. -1. 0.; 0. 0. 1.]; # nonpos diag
        pt3f = [2. 0. 1.; 0. 1. 0.; 0. 0. 4.]; # no lower and nonsym
        pt4 = [2. 0. 0.; 1. 2. 0.; 0. 0. 4.]
        @test !is_manifold_point(M, pt1f)
        @test_throws DomainError is_manifold_point(M, pt1f, true)
        @test !is_manifold_point(M, pt2f)
        @test_throws DomainError is_manifold_point(M, pt2f, true)
        @test !is_manifold_point(M,pt3f)
        @test_throws DomainError is_manifold_point(M,pt3f, true)
        @test is_manifold_point(M, pt4)
        @test !is_tangent_vector(M,pt3f, pt1f)
        @test_throws DomainError is_tangent_vector(M,pt3f, pt1f, true)
        @test !is_tangent_vector(M,pt4, pt1f)
        @test_throws DomainError is_tangent_vector(M,pt4, pt1f, true)
        @test !is_tangent_vector(M,pt4, pt3f)
        @test_throws DomainError is_tangent_vector(M,pt4, pt3f, true)
        @test is_tangent_vector(M,pt4, pt2f)
    end
end


using SparseArrays
using LinearAlgebra
using SparseArrays
using ForwardDiff: derivative

include("sim_funs.jl")
include("square_circle.jl")

# Define the mms solution
function ue(x, y, t, dom, A1 = 5 * exp(1) / (1 + exp(1)), A2 = 5)
    r = hypot(x, y)
    θ = atan(y, x)
    if dom == 1
        return A1 * sin(t) * (1 - exp(-1 * r^2)) * r * sin(θ)
    else
        return A2 * sin(t) * ((r - 1)^2 * cos(θ) + (r - 1) * sin(θ))
    end
end

∂t_ue(x, y, t, dom) = derivative(t -> ue(x, y, t, dom), t)
∂tt_ue(x, y, t, dom) = derivative(t -> ∂t_ue(x, y, t, dom), t)

∂x_ue(x, y, t, dom) = derivative(x -> ue(x, y, t, dom), x)
∂y_ue(x, y, t, dom) = derivative(y -> ue(x, y, t, dom), y)

∂xx_ue(x, y, t, dom) = derivative(x -> ∂x_ue(x, y, t, dom), x)
∂yy_ue(x, y, t, dom) = derivative(y -> ∂y_ue(x, y, t, dom), y)

function force(x, y, t, dom)
    return ∂tt_ue(x, y, t, dom) - (∂xx_ue(x, y, t, dom) + ∂yy_ue(x, y, t, dom))
end

# Run the convergence results
function main(
    sbp_order,
    refinement_levels,
    N0;
    characteristic_method = true,
    cfl = 2,
    friction = (V) -> asinh(V),
    tspan = (0.0, 1),
    do_output = true,
)
    # Define the friction law to use

    # The element to domain map is needed here in order to define the boundary
    # functions (we have a circular argument going on that we need EToDomain to
    # evaluate the boundary conditions, but we also need the boundary conditions
    # to build the square_circle operators)
    bc_map = [
        BC_DIRICHLET,
        BC_DIRICHLET,
        BC_NEUMANN,
        BC_NEUMANN,
        characteristic_method ? BC_JUMP_INTERFACE : -BC_JUMP_INTERFACE,
    ]
    if do_output
        @show bc_map
        @show [BC_DIRICHLET, BC_NEUMANN, BC_JUMP_INTERFACE]
    end

    (_, _, _, _, EToDomain) = read_inp_2d("square_circle.inp"; bc_map = bc_map)

    # Boundary condition function defined from the mms solution
    gDfun(x, y, t, e) = ue(x, y, t, EToDomain[e])
    gDdotfun(x, y, t, e) = ∂t_ue(x, y, t, EToDomain[e])
    function gNfun(nx, ny, xf, yf, t, e)
        dom = EToDomain[e]
        return nx .* (∂x_ue(xf, yf, t, dom)) + ny .* (∂y_ue(xf, yf, t, dom))
    end

    body_force(x, y, t, e) = force(x, y, t, EToDomain[e])

    # Loop over the levels of refinement
    ϵ = zeros(refinement_levels)
    for lvl in 1:length(ϵ)
        # Set up the local grid dimensions
        Nr = N0 * (2^(lvl - 1))
        Ns = N0 * (2^(lvl - 1))

        Nqr = Nr + 1
        Nqs = Ns + 1

        Np = Nqr * Nqs

        # Create the operators for the problem
        (metrics, rhsops, EToDomain, EToF, EToO, EToS, FToB, FToE, FToLF) =
            build_square_circle(
                sbp_order,
                Nqr,
                Nqs,
                gDfun,
                gDdotfun,
                gNfun,
                body_force,
                bc_map,
                (lvl == 1) && do_output,
            )
        nelems = length(rhsops)

        # initial conditions
        lenq0 = 2 * Np + 2 * Nqr + 2 * Nqs
        q = Array{Float64, 1}(undef, nelems * lenq0)
        for e in 1:nelems
            (xf1, xf2, xf3, xf4) = metrics[e].facecoord[1]
            (yf1, yf2, yf3, yf4) = metrics[e].facecoord[2]

            dom = EToDomain[e]
            u0 = ue.(metrics[e].coord[1][:], metrics[e].coord[2][:], 0, dom)
            v0 = ∂t_ue.(metrics[e].coord[1][:], metrics[e].coord[2][:], 0, dom)
            û10 = ue.(xf1[:], yf1[:], 0, dom)
            û20 = ue.(xf2[:], yf2[:], 0, dom)
            û30 = ue.(xf3[:], yf3[:], 0, dom)
            û40 = ue.(xf4[:], yf4[:], 0, dom)

            q[(lenq0 * (e - 1) + 1):(e * lenq0)] =
                [u0; v0; û10; û20; û30; û40]
        end

        # solve the ODE
        hmin = mapreduce(m -> m.hmin, min, values(metrics))
        dt = cfl * hmin
        if do_output
            @show dt
        end
        params = (
            Nqr = Nqr,
            Nqs = Nqs,
            rhsops = rhsops,
            EToF = EToF,
            EToO = EToO,
            EToS = EToS,
            FToB = FToB,
            FToE = FToE,
            FToLF = FToLF,
            friction = friction,
        )
        timestep!(q, waveprop!, params, dt, tspan)

        # Loop over the blocks and compute the error in each block
        for e in 1:nelems
            qe = @view q[(lenq0 * (e - 1) + 1):(e * lenq0)]

            qe_u = qe[1:Np]

            dom = EToDomain[e]
            qexact_u =
                ue.(
                    metrics[e].coord[1][:],
                    metrics[e].coord[2][:],
                    tspan[2],
                    dom,
                )
            qexact_v =
                ∂t_ue.(
                    metrics[e].coord[1][:],
                    metrics[e].coord[2][:],
                    tspan[2],
                    dom,
                )
            qexact_û1 =
                ue.(
                    metrics[e].facecoord[1][1],
                    metrics[e].facecoord[2][1],
                    tspan[2],
                    dom,
                )
            qexact_û2 =
                ue.(
                    metrics[e].facecoord[1][2],
                    metrics[e].facecoord[2][2],
                    tspan[2],
                    dom,
                )
            qexact_û3 =
                ue.(
                    metrics[e].facecoord[1][3],
                    metrics[e].facecoord[2][3],
                    tspan[2],
                    dom,
                )
            qexact_û4 =
                ue.(
                    metrics[e].facecoord[1][4],
                    metrics[e].facecoord[2][4],
                    tspan[2],
                    dom,
                )

            Δu = qe_u - qexact_u
            Mv = rhsops[e].JH

            ϵ[lvl] += Δu' * Mv * Δu
        end # end compute error at lvl

        ϵ[lvl] = sqrt(ϵ[lvl])
        if do_output
            @show (lvl, ϵ[lvl])
            if lvl > 1
                rate = (log.(ϵ[lvl - 1]) - log.(ϵ[lvl])) / log(2)
                @show rate
            end
            println()
        end
    end #loop over levels
    return ϵ
end


#------------------------------------------------------------------------------
function egarch11LL(par::Vector,y,x)

  (T,k) = (size(x,1),size(x,2))

  #lns2(t) = omega + alpha*abs(u(t-1)/s[t-1]) + beta1*lns2(t-1) + gamma1*u(t-1)/s[t-1]
  b = par[1:k]                                     #mean equation, y = x'*b
  (omega,alpha,beta1,gamma1) = par[k+1:k+4]
  yhat = x*b
  u    = y - yhat
  s2_0 = var(u)

  lns2    = zeros(typeof(alpha),T)
  lns2[1] = log(s2_0)
  for t = 2:T
    s_t1   = sqrt(exp(lns2[t-1]))
    lns2[t] = omega + alpha*abs(u[t-1])/s_t1 + beta1*lns2[t-1] + gamma1*u[t-1]/s_t1
  end
  s2 = exp.(lns2)

  LL_t    = -(1/2)*log(2*pi) .- (1/2)*log.(s2) - (1/2)*(u.^2)./s2
  LL_t[1] = 0

  LL = sum(LL_t)

  return LL,LL_t,s2,yhat,u

end
#------------------------------------------------------------------------------

#------------------------------------------------------------------------------

function DccLL(par,v,s2,Qbar)

  #(α,β) = par
  (α,β) = DccParTrans(par)             #restrict

  (T,n) = (size(v,1),size(v,2))

  u = v .* sqrt.(s2)

  Sigma = fill(NaN,(n,n,T))
  Q_t   = copy(Qbar)                    #starting guess
  LL_t  = zeros(T)
  for t = 2:T
    Q_t     = (1-α-β)*Qbar + α*v[t-1,:]*v[t-1,:]' + β*Q_t
    q_t     = diag(Q_t)
    R_t     = Q_t./sqrt.(q_t*q_t')
    d_t     = s2[t,:]
    Sigma_t = R_t.*sqrt.(d_t*d_t')
    LL_t[t] = -n*log(2*pi) - logdet(Sigma_t) - u[t,:]'*inv(Sigma_t)*u[t,:]
    Sigma[:,:,t] = Sigma_t
  end

  LL_t = 0.5*LL_t
  LL = sum(LL_t)

  return LL, LL_t, Sigma

end
#------------------------------------------------------------------------------

function DccParTrans(par)
    (a,b) = par
    α      = exp(a)/(1+exp(a)+exp(b))
    β      = exp(b)/(1+exp(a)+exp(b))
  return α,β
end
#------------------------------------------------------------------------------


using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

                     @variable(m, -1 <= x[i=1:4] <= 1)

                     @variable(m, -12.503038214935698 <= q <= 11.138318577023822)

                     add_NL_constraint(m, :(sigmoid(-0.2825986234777189 + -0.06477303738755547*sigmoid(0.91111659897616 + 0.32194568310871174*$(x[1]) + -0.04079743189253193*$(x[2]) + -0.7834407613653229*$(x[3]) + -0.5105294854228863*$(x[4])) + 0.23435836886254213*sigmoid(-0.27698098400696525 + 0.765733032903245*$(x[1]) + 0.6374561618281618*$(x[2]) + -0.006430667300674209*$(x[3]) + -0.6595935184751247*$(x[4])) + 0.7516401137686608*sigmoid(0.7789283911146407 + -0.8853253179157212*$(x[1]) + -0.6613182381252614*$(x[2]) + 0.5007525575333553*$(x[3]) + 0.7960733427996214*$(x[4]))) + sigmoid(0.12692347233027856 + 0.060984281016326136*sigmoid(0.91111659897616 + 0.32194568310871174*$(x[1]) + -0.04079743189253193*$(x[2]) + -0.7834407613653229*$(x[3]) + -0.5105294854228863*$(x[4])) + 0.4767325000240725*sigmoid(-0.27698098400696525 + 0.765733032903245*$(x[1]) + 0.6374561618281618*$(x[2]) + -0.006430667300674209*$(x[3]) + -0.6595935184751247*$(x[4])) + -0.9515049437955176*sigmoid(0.7789283911146407 + -0.8853253179157212*$(x[1]) + -0.6613182381252614*$(x[2]) + 0.5007525575333553*$(x[3]) + 0.7960733427996214*$(x[4]))) + sigmoid(-0.022625917828582143 + -0.9740734403326825*sigmoid(0.91111659897616 + 0.32194568310871174*$(x[1]) + -0.04079743189253193*$(x[2]) + -0.7834407613653229*$(x[3]) + -0.5105294854228863*$(x[4])) + -0.9306328050206538*sigmoid(-0.27698098400696525 + 0.765733032903245*$(x[1]) + 0.6374561618281618*$(x[2]) + -0.006430667300674209*$(x[3]) + -0.6595935184751247*$(x[4])) + 0.61848968844355*sigmoid(0.7789283911146407 + -0.8853253179157212*$(x[1]) + -0.6613182381252614*$(x[2]) + 0.5007525575333553*$(x[3]) + 0.7960733427996214*$(x[4]))) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m


blockvalid(x, v) = count(isequal(v), x) ≤ 1

function backtrack!(x)
    pos = findfirst(isequal(0), x)
    isa(pos, Nothing) && return true
    iloc = 3div(pos[1]-1, 3) .+ (1:3)
    jloc = 3div(pos[2]-1, 3) .+ (1:3)
    for k in 1:9
        x[pos] = k
        blockvalid(view(x, pos[1], :), k) || continue
        blockvalid(view(x, :, pos[2]), k) || continue
        blockvalid(view(x, iloc, jloc), k) || continue
        backtrack!(x) && return true
    end
    x[pos] = 0
    return false
end

function ssolve(lines, i)
    t = [lines[10i-j][k] - '0' for j in 8:-1:0, k in 1:9]
    backtrack!(t)
    sum([100, 10, 1] .* t[1, 1:3])
end

lines = readlines("p096_sudoku.txt")
@time sum(ssolve(lines, i) for i in 1:50)
@time sum(ssolve(lines, i) for i in 1:50)

Chapter04/13. Using array views to avoid memory allocation/sudoku.jl


#
# this is a trivial first write-up of NW, Algorithm 17.4
# with the new `update!` mechanism, I'd like to incorporate the
# Lagrange-multiplier and penalty updates into the minimisation
# procedure. But maybe this is not a good idea, to be tested.
#

function optimize( F::DifferentiableFunction,
                   C::EqualityConstraint,
                   x0::AbstractVector;
                   iterations = 5000,
                   c_tol = 1e-6,
                   g_tol = 1e-6,
                   verbose = 1,
                   Optimizer = Optim.ConjugateGradient )

   # initialise
   al = AugmentedLagrangian(F, C, x0)
   eta = al.mu^(-0.1)
   # TODO: this needs to distinguish whether F, C are DifferentiableFunction
   # or TwiceDifferentiableFunction!!!!
   ALobj = DifferentiableFunction( x_ -> evaluate(x_, al),
                                   (x_,g_) -> gradient!(g_, x_, al),
                                   (x_,g_) -> eval_and_grad!(g_, x_, al) )

   x = copy(x0)
   iteration = 0
   if verbose >= 1
      @printf("  it  |     |c|_∞        |g_AL|_∞      μ         η  \n")
      @printf("------|----------------------------------------------\n")
   end
   while iteration < iterations

      # solve the sub-problem
      options = OptimizationOptions(g_tol=g_tol, iterations=iterations - iteration, store_trace = true)
      result = Optim.optimize(ALobj, x, Optimizer(), options)
      # TODO: unclear what to do if this step fails? it could still be ok after
      # we update the lagrange multipliers? For now just print a warning and continue.
      if !Optim.converged(result)
         warn("an inner AL iteration has not converged")
         # return minimizer(result)
      end

      # check for convergence, perform update for multipliers and tolerances
      iteration += Optim.iterations(result)
      x = Optim.minimizer(result)
      nrm_g = Optim.g_norm_trace(result)[end]

      c = C.c(x)    # TODO: this is a superfluous call to C (but #282 would fix this)

      if verbose >= 1
         @printf(" %4d |  %1.4e   %1.4e   %1.2e   %1.2e\n",
                  iteration, norm(x, Inf), nrm_g, al.mu, eta)
      end

      if norm(c, Inf) < c_tol
         return x, al
      elseif norm(c, Inf) <= eta
         al.lambda = al.lambda - al.mu * c
         eta /= al.mu^0.9  # tighten constraint tolerance
      else
         al.mu *= 100.0
         eta = al.mu^(-0.1)
      end
      if al.mu > 1e10
         warn("something horrible is happening")
         return x, al
      end

   end

   # iteration > iterations
   warn("too many iterations in `ConstrainedOptim.optimize` ")
   return x, al
end


using OverlapIntegrals, Test

@testset "fact2" begin
    @test fact2(0) == 1
    @test fact2(1) == 1
    @test fact2(2) == 2
    @test fact2(3) == (3 * 1)
    @test fact2(4) == (4 * 2)
    @test fact2(5) == (5 * 3 * 1)
    @test fact2(6) == (6 * 4 * 2)
    @test fact2(7) == (7 * 5 * 3 * 1)
end

@testset "binomial_prefactor" begin
    @test isapprox(binomial_prefactor(1, 1, 1, 0.1, 0.2), 0.3; atol = 1e-7)
    @test isapprox(binomial_prefactor(1, 1, 1, 0.3, 0.4), 0.7; atol = 1e-7)
    @test isapprox(binomial_prefactor(1, 3, 1, 0.1, 0.2), 0.007; atol = 1e-7)
    @test isapprox(binomial_prefactor(2, 3, 1, 0.1, 0.2), 0.09; atol = 1e-7)
end

@testset "overlap1d" begin
    @test isapprox(overlap1d(1, 1, 0.1, 0.2, 1.0), 0.52; atol = 1e-5)
    @test isapprox(overlap1d(3, 1, 0.1, 0.2, 1.0), 0.7952; atol = 1e-5)
end

@testset "tho66" begin
    za = 1.8
    zb = 2.8
    ra = [0.0, 0.0, 0.0]
    rb = [0.5, 0.8, -0.2]

    integral = tho66(za, zb, ra, rb, [0, 0, 0], [0, 0, 0])
    @test isapprox(integral, 0.20373275913014607; atol = 1.0e-16)

    integral = tho66(za, zb, ra, rb, [1, 0, 0], [0, 0, 0])
    @test isapprox(integral, 0.062005622343957505; atol = 1.0e-16)

    integral = tho66(za, zb, ra, rb, [1, 1, 0], [1, 1, 0])
    @test isapprox(integral, -0.00043801221837779696; atol = 1.0e-16)

    integral = tho66(za, zb, ra, rb, [2, 1, 0], [1, 1, 0])
    @test isapprox(integral, -0.0002385994651113168; atol = 1.0e-16)
end