using DiseaseOutbreak
using Plots

#####################################
# more infectious
#####################################
n=10000
status,d,pos = makeCitizens(n);
#TODO: confirm with Saunak: infectCitizens needs second argument: nInfect
infectCitizens!(status, 10)
followCitizens!(status,d,365,0.05,14,0.0025)

uninfected(status)

newInfections(status)

cumulativeInfections(status)

gif(infectionSpread(status,pos),fps=15)

#####################################
# less infectious
#####################################
n=10000
status1,d1,pos1 = makeCitizens(n);
infectCitizens!(status1, 10)
followCitizens!(status1,d1,365,0.05,14,0.002)
uninfected(status1)

newInfections(status1)

cumulativeInfections(status1)

gif(infectionSpread(status1,pos1),fps=15)


function master_formulation(rd::RoutingData, type::FlowFormulation)
    m = _create_model(rd)
    rm = basic_routing_model_unitary(m, rd, type) # Includes flow-conservation constraints.

    @variable(m, mu >= 0)
    @objective(m, Min, mu)

    @variable(m, dual_alpha[e in edges(rd), d in demands(rd), v in vertices(rd)])
    @variable(m, dual_beta[e in edges(rd), e2 in edges(rd)] >= 0)

    # Robust uncertainty sets (one per edge).
    if rd.model_robust_reformulation_traffic_matrices
        dual_constraints = Dict{Edge{Int}, Dict{Edge{Int}, ConstraintRef}}(
            e => Dict{Edge{Int}, ConstraintRef}() for e in edges(rd)
        ) # Edge -> demand -> constraint reference.
    end
    for e in edges(rd)
        for d in demands(rd)
            # Dual constraint from the flow variables.
            for e2 in edges(rd)
                @constraint(m, dual_beta[e, e2] + dual_alpha[e, d, dst(e2)] - dual_alpha[e, d, src(e2)] >= 0)
            end

            # Dual constraint from the demand variables.
            c = @constraint(m, dual_alpha[e, d, src(d)] - dual_alpha[e, d, dst(d)] >= rm.routing[d, e])
            if rd.model_robust_reformulation_traffic_matrices
                dual_constraints[e][d] = c
            end
        end

        # Relate the main decision variables to the uncertainty sets.
        lhs = sum(dual_beta[e, e2] * capacity(rd, e2) for e2 in edges(rd))
        @constraint(m, lhs / capacity(rd, e) <= mu)
    end

    return RoutingModel(rd, m, UnitaryFlows, rm.routing, mu=mu, dual_alpha=dual_alpha, dual_beta=dual_beta)
end


export getData

import jInv.ForwardShare.getData
function getData(m,pFor::BasicFWIparam,doClear::Bool=false)
    tstart = time_ns();
    # extract pointers
    Mesh  = pFor.Mesh
    omega = pFor.omega
    gamma = pFor.gamma
    Q     = pFor.Sources
    P     = pFor.Receivers

    nrec  = size(P,2)
    nsrc  = size(Q,2)
    nfreq = length(omega)

    # allocate space for data and fields
    D  = zeros(nrec,nsrc,nfreq)
    U  = zeros(ComplexF64,prod(Mesh.n.+1),nsrc,nfreq)
    LU = Array{Any}(undef, nfreq)
    HinvtPs = Array{Any}(undef, nfreq)
    for i=1:length(omega)
        t1 = time_ns();
        H = getHelmholtzOperator(m,gamma,omega[i],Mesh)
        t2 = time_ns();
        LU[i] = lu(H)
        H = nothing
        for k=1:nsrc
            U[:,k,i] = LU[i]\Vector(Q[:,k])
            D[:,k,i] = real(P'*U[:,k,i])
        end
    end
    pFor.Ainv   = LU
    pFor.Fields = U # will need to save to file
    tend = time_ns();
    println("Runtime of getData:");
    println((tend - tstart)/1.0e9);

    return D ,pFor
end


module IfSimpleLoadBearing

include("./../../constraints/geometric/constants.jl")
using JuMP
using .Constants

export if_simple_lb

function if_simple_lb(m)

    #Basically just check whether box j is located right above box i

    # Zj >= Zi + hobi
    @variable(m, if1_zj_geq_zi_add_hobi_ij[i=1:length(boxes), j=1:length(boxes); i!=j], Bin)
    @constraint(m, c_if1_zj_geq_zi_add_hobi_ij1[i=1:length(boxes), j=1:length(boxes); i!=j],
        H * (1 - m[:if1_zj_geq_zi_add_hobi_ij][i, j]) >=
        (m[:flbz_front_left_bot_z][i] + m[:hob_height_of_box][i]) - m[:flbz_front_left_bot_z][j]
    )

    @constraint(m, c_if1_zj_geq_zi_add_hobi_ij2[i=1:length(boxes), j=1:length(boxes); i!=j],
        H * m[:if1_zj_geq_zi_add_hobi_ij][i, j] >=
        m[:flbz_front_left_bot_z][j] - (m[:flbz_front_left_bot_z][i] + m[:hob_height_of_box][i]) + 0.001
    )

    # Xj < Xi + Xcoveri
    @variable(m, if2_xj_l_xi_xcoveri_ij[i=1:length(boxes), j=1:length(boxes); i!=j], Bin)
    @constraint(m, c_if2_xj_l_xi_xcoveri_ij1[i=1:length(boxes), j=1:length(boxes); i!=j],
        L * (1 - m[:if2_xj_l_xi_xcoveri_ij][i, j]) >=
        m[:flbx_front_left_bot_x][j] + 0.001 - (m[:flbx_front_left_bot_x][i] + m[:box_x_axis_cover][i])
    )

    @constraint(m, c_if2_xj_l_xi_xcoveri_ij2[i=1:length(boxes), j=1:length(boxes); i!=j],
        L * m[:if2_xj_l_xi_xcoveri_ij][i, j] >=
        (m[:flbx_front_left_bot_x][i] + m[:box_x_axis_cover][i]) - m[:flbx_front_left_bot_x][j] - 0.001
    )

    # Xj + Xcoverj < Xi
    @variable(m, if3_xj_xcoverj_g_xi_ij[i=1:length(boxes), j=1:length(boxes); i!=j], Bin)
    @constraint(m, c_if3_xj_xcoverj_g_xi_ij1[i=1:length(boxes), j=1:length(boxes); i!=j],
        L * (1 - m[:if3_xj_xcoverj_g_xi_ij][i, j]) >=
        m[:flbx_front_left_bot_x][i] + 0.001 - (m[:flbx_front_left_bot_x][j] + m[:box_x_axis_cover][j])
    )

    @constraint(m, c_if3_xj_xcoverj_g_xi_ij2[i=1:length(boxes), j=1:length(boxes); i!=j],
        L * m[:if3_xj_xcoverj_g_xi_ij][i, j] >=
        (m[:flbx_front_left_bot_x][j] + m[:box_x_axis_cover][j]) - m[:flbx_front_left_bot_x][i] - 0.001
    )

    # Yj < Yi + Ycoveri
    @variable(m, if4_yj_l_yi_ycoveri_ij[i=1:length(boxes), j=1:length(boxes); i!=j], Bin)
    @constraint(m, c_if4_yj_l_yi_ycoveri_ij1[i=1:length(boxes), j=1:length(boxes); i!=j],
        W * (1 - m[:if4_yj_l_yi_ycoveri_ij][i, j]) >=
        m[:flby_front_left_bot_y][j] + 0.001 - (m[:flby_front_left_bot_y][i] + m[:box_y_axis_cover][i])
    )

    @constraint(m, c_if4_yj_l_yi_ycoveri_ij2[i=1:length(boxes), j=1:length(boxes); i!=j],
        W * m[:if4_yj_l_yi_ycoveri_ij][i, j] >=
        (m[:flby_front_left_bot_y][i] + m[:box_y_axis_cover][i]) - m[:flby_front_left_bot_y][j] - 0.001
    )

    # Yj < Yi + Ycoveri
    @variable(m, if5_yj_ycoverj_g_yi_ij[i=1:length(boxes), j=1:length(boxes); i!=j], Bin)
    @constraint(m, c_if5_yj_ycoverj_g_yi_ij1[i=1:length(boxes), j=1:length(boxes); i!=j],
        W * (1 - m[:if5_yj_ycoverj_g_yi_ij][i, j]) >=
        m[:flby_front_left_bot_y][i] + 0.001 - (m[:flby_front_left_bot_y][j] + m[:box_y_axis_cover][j])
    )

    @constraint(m, c_if5_yj_ycoverj_g_yi_ij2[i=1:length(boxes), j=1:length(boxes); i!=j],
        W * m[:if5_yj_ycoverj_g_yi_ij][i, j] >=
        (m[:flby_front_left_bot_y][j] + m[:box_y_axis_cover][j]) - m[:flby_front_left_bot_y][i] - 0.001
    )

    #attempt at advanced AND
    @variable(m, if_simple_lb_var[i=1:length(boxes), j=1:length(boxes); i!=j], Bin)
    @constraint(m, c_if_simple_lb1[i=1:length(boxes), j=1:length(boxes); i!=j],
        m[:if_simple_lb_var][i, j] >=
        m[:if1_zj_geq_zi_add_hobi_ij][i, j] +
        m[:if2_xj_l_xi_xcoveri_ij][i, j] +
        m[:if3_xj_xcoverj_g_xi_ij][i, j] +
        m[:if4_yj_l_yi_ycoveri_ij][i, j] +
        m[:if5_yj_ycoverj_g_yi_ij][i, j] - 4
    )

    #@expression(m, if_simple_lb_objective[i=1:length(boxes), j=1:length(boxes); i!=j],
    #    (1 - m[:if_simple_lb_var][i, j])
    #)
    return m
end

end


export convolutionProductFun

#
# A ProductFun constructor for bivariate functions on Intervals
# defined as the distance of their arguments.
#

function convolutionProductFun{U<:UnivariateSpace,V<:UnivariateSpace}(f::Function,u::U,v::V;tol=eps())
    du,dv = domain(u),domain(v)
    ext = extrema(du,dv)
    if ext[1] == 0
        ff = Fun(z->f(0,z),Chebyshev(Interval(-ext[2]/2,ext[2]/2)))
        fd,T = ff(0),eltype(ff)
        c = chop(coefficients(ff),norm(coefficients(ff),Inf)*100eps(T))
        N = length(c)
        N1 = isa(du,PeriodicDomain) ? 2N : N
        N2 = isa(dv,PeriodicDomain) ? 2N : N
        return ProductFun((x,y)->x==y?fd:f(x,y),u⊗v,N1,N2;tol=tol)
    else
        ff = Fun(z->f(0,z),Chebyshev(Interval(ext...)))
        c = chop(coefficients(ff),norm(coefficients(ff),Inf)*100eps(eltype(ff)))
        N = length(c)
        N1 = isa(du,PeriodicDomain) ? 2N : N
        N2 = isa(dv,PeriodicDomain) ? 2N : N
        return ProductFun(f,u⊗v,N1,N2;tol=tol)
    end
end

convolutionProductFun{U<:UnivariateSpace,
                      V<:UnivariateSpace,T}(f::Function,
                                            ss::TensorSpace{Tuple{U,V},T,2};kwds...) = convolutionProductFun(f,ss[1],ss[2];kwds...)

#
# ProductFun constructors for functions on periodic intervals.
#

#
# Suppose we are interested in K(ϕ-θ). Then, K(⋅) is periodic
# whether it's viewed as bivariate or univariate.
#

function convolutionProductFun{DD,T,DU,DV}(f::Fun{Fourier{DD},T},u::Fourier{DU},v::Fourier{DV};tol=eps())
    df,du,dv = domain(f),domain(u),domain(v)
    @assert df == du == dv && isa(df,PeriodicInterval)
    c = coefficients(f)
    N = length(c)
    X = zeros(T,N,N)
    X[1,1] += c[1]
    @inbounds for i=2:2:N-1
        X[i,i] += c[i+1]
        X[i+1,i] += c[i]
        X[i,i+1] -= c[i]
        X[i+1,i+1] += c[i+1]
    end
    if mod(N,2)==0 X[N,N-1],X[N-1,N] = c[N],-c[N] end
    ProductFun(X,u⊗v)
end

function convolutionProductFun{S<:CosSpace,T,DU,DV}(f::Fun{S,T},u::Fourier{DU},v::Fourier{DV};tol=eps())
    df,du,dv = domain(f),domain(u),domain(v)
    @assert df == du == dv && isa(df,PeriodicInterval)
    c = coefficients(f)
    N = 2length(c)-1
    X = zeros(T,N,N)
    X[1,1] += c[1]
    @inbounds for i=2:2:N
        X[i,i] += c[div(i,2)+1]
        X[i+1,i+1] += c[div(i,2)+1]
    end
    ProductFun(X,u⊗v)
end

function convolutionProductFun{S<:SinSpace,T,DU,DV}(f::Fun{S,T},u::Fourier{DU},v::Fourier{DV};tol=eps())
    df,du,dv = domain(f),domain(u),domain(v)
    @assert df == du == dv && isa(df,PeriodicInterval)
    c = coefficients(f)
    N = 2length(c)+1
    X = zeros(T,N,N)
    @inbounds for i=2:2:N
        X[i+1,i] += c[div(i,2)]
        X[i,i+1] -= c[div(i,2)]
    end
    ProductFun(X,u⊗v)
end

function convolutionProductFun{DS,T,DU,DV}(f::Fun{Laurent{DS},T},u::Laurent{DU},v::Laurent{DV};tol=eps())
    df,du,dv = domain(f),domain(u),domain(v)
    @assert df == du == dv && isa(df,PeriodicInterval)
    c = coefficients(f)
    N = length(c)
    X = mod(N,2) == 0 ? zeros(T,N+1,N) : zeros(T,N,N)
    X[1,1] += c[1]
    @inbounds for i=2:2:N-1
        X[i+1,i] += c[i]
        X[i,i+1] += c[i+1]
    end
    if mod(N,2)==0 X[N+1,N] = c[N] end
    ProductFun(X,u⊗v)
end

function convolutionProductFun{DS,T,DU,DV}(f::Fun{Taylor{DS},T},u::Laurent{DU},v::Laurent{DV};tol=eps())
    df,du,dv = domain(f),domain(u),domain(v)
    @assert df == du == dv && isa(df,PeriodicInterval)
    c = coefficients(f)
    N = 2length(c)-1
    X = zeros(T,N-1,N)
    X[1,1] += c[1]
    @inbounds for i=2:2:N-1
        X[i,i+1] += c[div(i,2)+1]
    end
    ProductFun(X,u⊗v)
end

function convolutionProductFun{DS,T,DU,DV}(f::Fun{Hardy{false,DS},T},u::Laurent{DU},v::Laurent{DV};tol=eps())
    df,du,dv = domain(f),domain(u),domain(v)
    @assert df == du == dv && isa(df,PeriodicInterval)
    c = coefficients(f)
    N = 2length(c)
    X = zeros(T,N+1,N)
    @inbounds for i=2:2:N
        X[i+1,i] += c[div(i,2)]
    end
    ProductFun(X,u⊗v)
end


using OptimJulia
using Base.Test
using SymPy

x1,x2 = Sym("x1,x2")

f1 = ((((1 - 8x1) + 2*x1^2) - (10/3)*x1^3) + (1/4)*x1^4 + (4/5)*x1^5) - (1/6)*x1^6

Newton(f1, x1, 3)

vars = [x1,x2]
f2 = -log(1 - x1 - x2) - log(x1) - log(x2)
x0 = [0.85,0.05]

Newton(f2, vars, x0)

for i = 2:10
      QuasiNewton(f1, x1, x0 = [i i+4])
      println("=========================================")
      println("Test $i with QuasiNewton algorithm passed")
end


using Test, SimpleTensorNetworks.ContractionOrder
using LightGraphs
using LightGraphs: SimpleEdge
using SimpleTensorNetworks.ContractionOrder: abstract_contract_edge!, log2sumexp2

function get_shapes(edges, sizes, neighbors)
    log2_shapes = Vector{Float64}[]
    for i=1:length(neighbors)
        si = Float64[]
        for nb in neighbors[i]
            ie = findfirst(e->e == SimpleEdge(i=>nb) || e == SimpleEdge(nb=>i), edges)
            push!(si, sizes[ie])
        end
        push!(log2_shapes, si)
    end
    return log2_shapes
end

tolog2(shapes) = [log2.(x) for x in shapes]

@testset "abstract contract" begin
    es = [SimpleEdge(1, 2), SimpleEdge(2, 3), SimpleEdge(1, 3), SimpleEdge(3, 4)]
    log2_sizes = log2.([2, 3, 4, 5])
    neighbors = [[2,3], [3,1], [1,2,4], [3]]
    log2_shapes = get_shapes(es, log2_sizes, neighbors)
    strategy = "min_reduce"
    edge_info = edgeinfo(es, log2_shapes, neighbors, strategy)

    edge = SimpleEdge(2, 3)
    tc, sc = abstract_contract_edge!(edge, log2_shapes, edge_info,
        neighbors, strategy)
    @test edge_info.data[SimpleEdge(2, 4)] ≈ 8.0
    @test edge_info.data[SimpleEdge(1, 2)] ≈ 5.0
    @test log2_shapes == [[8], [8, 5], Int[], [5]] |> tolog2
    @test neighbors == [[2], [1, 4], Int[], [2]]
    @test tc ≈ log2(120)

    es = [SimpleEdge(1, 2), SimpleEdge(2, 3), SimpleEdge(1, 3), SimpleEdge(3, 4)]
    log2_sizes = log2.([2, 3, 4, 5])
    neighbors = [[2,3], [3,1], [1,2,4], [3]]
    log2_shapes = get_shapes(es, log2_sizes, neighbors)
    tn = TensorNetworkLayout(SimpleGraph(4, neighbors), log2_shapes)
    order = [SimpleEdge(1, 2),SimpleEdge(1, 3), SimpleEdge(1, 4)]
    tcs, scs = abstract_contract(tn, order)
    @test log2sumexp2(tcs) ≈ log2(89)
    @test maximum(scs) ≈ log2(12*5+12+5)
end

@testset "greedy order" begin
    g = random_regular_graph(10, 3)
    log2_shapes = [[2,2,2] for i=1:10] |> tolog2
    tn = TensorNetworkLayout(g, log2_shapes)
    tcs, scs, order = order_greedy(tn; strategy="min_dim")
    @test length(order) <= length(edges(g))
end

@testset "disconnected graph" begin
    n = 5
    g = SimpleGraph(n)
    add_edge!(g, 2, 3)
    tn = TensorNetworkLayout(g, [ones(degree(g, i)) for i=1:n])
    tc, sc, orders = order_greedy(tn)
    @test orders isa Vector
    tcs, scs = abstract_contract(tn, orders)
    @test maximum(tcs) == 1
end


"""
This module provides high-level functionality to generate, manage and refine
meshes based on constrained Delaunay triangulations.
"""

module DelaunayMeshes

import DiffEqPDEBase
import GeometricalPredicates

import Base.convert
import Base.push!
import Base.next
import Base.done
import Base.start

include("QuadEdge.jl")
include("Triangulation.jl")
include("DelaunayMesh.jl")
include("PostProcess.jl")

export getvoronoivertices, setboundingbox
export scalepoints, unscalepoints, addconstraint!
export getdelaunaycoordinates
end # module


function test_apply_rhs()
    grid = generate_grid(Quadrilateral, (20, 20))
    dim = 2
    ip = Lagrange{dim,RefCube,1}()
    qr = QuadratureRule{dim,RefCube}(2)
    cellvalues = CellScalarValues(qr, ip)

    dh = DofHandler(grid)
    push!(dh, :u, 1)
    close!(dh)

    K = create_sparsity_pattern(dh)

    ch = ConstraintHandler(dh)

    ∂Ω = union(getfaceset.((grid,), ["left", "right"])...)
    dbc = Dirichlet(:u, ∂Ω, (x, t) -> 0)
    add!(ch, dbc);

    ∂Ω = union(getfaceset.((grid,), ["top", "bottom"])...)
    dbc = Dirichlet(:u, ∂Ω, (x, t) -> 2)
    add!(ch, dbc);

    close!(ch)
    update!(ch, 0.0);

    function doassemble!(
        cellvalues::CellScalarValues{dim},
        K::SparseMatrixCSC,
        dh::DofHandler,
    ) where {dim}

        n_basefuncs = getnbasefunctions(cellvalues)
        Ke = zeros(n_basefuncs, n_basefuncs)
        fe = zeros(n_basefuncs)

        f = zeros(ndofs(dh))
        assembler = start_assemble(K, f)

        @inbounds for cell in CellIterator(dh)
            fill!(Ke, 0)
            fill!(fe, 0)

            reinit!(cellvalues, cell)

            for q_point = 1:getnquadpoints(cellvalues)
                dΩ = getdetJdV(cellvalues, q_point)

                for i = 1:n_basefuncs
                    v = shape_value(cellvalues, q_point, i)
                    ∇v = shape_gradient(cellvalues, q_point, i)
                    fe[i] += v * dΩ
                    for j = 1:n_basefuncs
                        ∇u = shape_gradient(cellvalues, q_point, j)
                        Ke[i, j] += (∇v ⋅ ∇u) * dΩ
                    end
                end
            end

            assemble!(assembler, celldofs(cell), fe, Ke)
        end
        return K, f
    end

    K, f = doassemble!(cellvalues, K, dh)
    A = create_sparsity_pattern(dh)
    A, g = doassemble!(cellvalues, A, dh)
    rhsdata = get_rhs_data(ch, A)

    apply!(K, f, ch)
    apply!(A, ch) # need to apply bcs to A once
    apply_rhs!(rhsdata, g, ch)
    u₁ = K \ f
    u₂ = A \ g
    return u₁, u₂
end

u1, u2 = test_apply_rhs()
@test u1 == u2


using Gen
using Distributions
using LinearAlgebra
include("structs.jl")
include("distributions.jl")

@gen function gen_scene_and_observations(
    max_objects::Int64,
    num_object_categories::Int64,
    num_observations::Int64,
    room::RoomParams,
    camera::CameraParams,
    detection_sd::Float64
    )

    # Generates scene
    num_objects = @trace(uniform_discrete(1, max_objects), :num_objects)
    println("num_objects: ", num_objects)
    categories = [@trace(uniform_discrete(1, num_object_categories),
        :objects=>i=>:category) for i in 1:num_objects]
    locations = [@trace(mvuniform(room.bottom_left_lim, room.top_right_lim),
        :objects=>i=>:location) for i in 1:num_objects]
    # Generates observations
    observations = []
    for i in 1:num_observations
        camera_position = @trace(mvuniform(room.bottom_left_lim, room.top_right_lim),
            :observations=>i=>:camera_location)

        object = @trace(uniform_discrete(1, num_objects), :observations=>i=>:object)

        category_dist = [k == object ? 1 : 0 for k in 1:num_objects]

        category = @trace(categorical(category_dist), :observations=>i=>:category)

        location = locations[object]

        direction = @trace(choose_direction(camera_position, location, detection_sd), :observations=>i=>:direction)

        append!(observations, Dict(:camera_location => camera_position,
            :category => category,
            :direction => direction,
            :object => object))
    end
    return num_objects, categories, locations, observations
end

args = (5, 10, 20, RoomParams(), CameraParams(), 0.5)

(trace, _) = Gen.generate(gen_scene_and_observations, args)

Gen.get_choices(trace)


module ClassificatinoInerseValueProblem

using LIBLINEAR: linear_train

"""Minimizes f(δi, πi) = δi + 1/(2λ) ‖ πi - πk ‖^2, s.t., πi'bi - zi ≥ δi, δi ≥ 0."""
function min_fpidelta(
    πk::AbstractVector{T}, bi::AbstractVector{T}, zi::T, λ::T
) where {T <: AbstractFloat}
    res = max(zi - bi'πk, 0.0) / (bi'bi) / λ
    μi = min(res, 1)
    πi = πk + λ * μi * bi
    return πi
end

"""Minimizes f(ϵ, c; ck, cα, λ, xi) = ϵ'1 + 1/(2λ) ‖c - ck‖^2 + α/2 ‖c - cα‖^2,
s.t., c'xi - zi ≤ ϵi, ϵi ≥ 0, for all i, c ∈ C.

"""
function min_fcepsilon(
    xi::Matrix{T}, zi::Vector{T}, ck::Vector{T}, cα::Vector{T}, λ::T, α::T
) where {T <: AbstractFloat}
    ckα = 1 / (1 + λ * α) * (ck + λ * α * cα)
    cost = λ / (1 + λ * α)

    # Create new array bi = xi'ckα - zi for all i
    b = zero(zi)
    for j in 1:size(xi, 1), i in 1:size(xi, 2)
        b[i] += xi[j, i] * ckα[j]
    end
    b .-= zi
    W = b[b .> 0.0]

    if length(W) == 0
        return zero(ck)
    else
        j = 0
        a = Matrix{Float64}(undef, size(xi, 1), length(W))
        for i in 1:length(b)
            if b[i] > 0.0
                j += 1
                a[:, i] = -1 / W[j] * xi[:, i]
            end
        end
        labels = ones(length(W))
        model = linear_train(labels, a; eps=1e-4, solver_type=Cint(3), C=cost, bias=-1, W=W)
        return model.w + ckα
    end
end

end # module


using Compat, ClassicalLaminateTheory

forcesAndMoments =
 @compat Dict(
    :Nx => 300,
    :Ny => 100,
    :Nxy => 200,
    :Mx => 5,
    :My => 3,
    :Mxy => 1)

materialProperties =@compat Dict(
  :E1 => [20.0*10^6],
  :E2 => [1.5*10^6],
  :G12 => [1.0*10^6],
  :ν12 => [.34],
  :alpha1 => [0.2e-6],
  :alpha2 => [20e-6],
  :beta1 => [0.2e-4],
  :beta2 => [20e-4],
  :thickness => [0.0075])

laminateProperties =@compat Dict(
  :nplies => 5,
  :materials => [1, 1, 1, 1, 1],
  :repeats => 1,
  :symmetric => false,
  :symmetricrepeats => 1,
  :orientation => [0, 45, 90, 45, -45],
  :deltaTemp => -270.0,
  :deltaMoisture => 1
)

qm = qmat(materialProperties)
materialProperties |> display

println("\nQ:")
qm |> display

qb0 = qbarmat(qm, 0.0)
qb90 = qbarmat(qm, 90.0)
qb45 = qbarmat(qm, 45.0)
qbm45 = qbarmat(qm, -45.0)
println("\nQbar 0 Degrees:")
qb0 |> display
println("\nQbar 90 Degrees:")
qb90 |> display
println("\nQbar 45 Degrees:")
qb45 |> display
println("\nQbar -45 Degrees:")
qbm45 |> display

println()
createLaminate!(laminateProperties, materialProperties, forcesAndMoments)
show(laminateProperties)

println("ABD matrix:")
round(laminateProperties[:ABD], 3) |> display

println()
println("Total forces and moments:")
round(laminateProperties[:fTotal], 7) |> display


module OrnsteinZernike

using LinearAlgebra
using FFTW
using Romberg

include("potentials.jl")
export Potential, PseudoHS, HardSphere, SquareWell, LennardJones
include("closures.jl")
export HypernettedChain,
    PercusYevick,
    ModifiedVerlet,
    MeanSpherical,
    SoftMeanSpherical,
    HMSA,
    Parameters,
    Structure,
    Result,
    Interaction
include("types.jl")
include("fft.jl")
include("utils.jl")
include("solver.jl")
export solve

end


module HermiteGF

  export Hermite, Radial
  export Chebyshev, Uniform
  export interpolate

  " Interpolation type (Hermite or Radial) "
  abstract type InterpolationType end

  " Node positions (Uniform or Chebyshev) "
  abstract type NodesType end

  include("chebyshev.jl")
  include("uniform.jl")
  include("hermite.jl")
  include("radial.jl")
  include("interpolation.jl")

end


using Base.LinAlg: Eigen

function Base.eigfact{T}(::Type{Positive{T}}, A::AbstractMatrix{T}, args...; tol=default_tol(A))
    F = eigfact(A, args...)
    for i = 1:size(A,1)
        tmp = abs(F.values[i])
        if tmp < tol
            tmp = one(tmp)
        end
        F.values[i] = tmp
    end
    F
end
Base.eigfact(::Type{Positive}, A::AbstractMatrix, args...; tol=default_tol(A)) = eigfact(Positive{floattype(eltype(A))}, A, args...; tol=tol)
Base.eigfact{T}(::Type{Positive{T}}, A::AbstractMatrix, args...; tol=default_tol(A)) = eigfact(Positive{T}, convert(Matrix{T}, A), args...; tol=tol)


"""
    Auroral()

Aurora with linear regression.
"""
struct Auroral <: AbstractAurora end

Base.@kwdef struct FittedAuroral{A, Ts, Ms, Ss} <: AbstractFittedAurora
    βs::Ts
    βs_list::Vector{Ts}
    μs_mat::Ms
    μs::Ss
    method::A
end

function StatsBase.fit(auroral::Auroral, Zs::AbstractVector{<:ReplicatedSample})
    K = nobs(Zs[1])
    n = length(Zs)
    any(nobs.(Zs) != K) || throw("All Zs should have the same number of replicates.")
    μs_mat = zeros(n,K)

    # See https://github.com/dmbates/CopenhagenEcon/blob/master/jmd/03-LinearAlgebra.jmd
    cache_vec = zeros(K+1)
    cache_mat = zeros(K+1,n)
    cache_cov_mat = zeros(K+1, K+1)

    βs_list = Vector{typeof(cache_vec)}(undef, K)

    for j in Base.OneTo(K)
        fill!(cache_mat, 0)
        fill!(cache_cov_mat, 0)
        fill!(cache_vec, 0)

        for (i,Z) in enumerate(Zs)
            design_response_row!(view(cache_mat, :, i), Z, j)
        end

        mul!(cache_cov_mat, cache_mat, cache_mat')
        chr = cholesky!(cache_cov_mat)

        RXX = UpperTriangular(view(chr.U, 1:K, 1:K))
        βs_list[j] = ldiv!(RXX, copy(chr.U[1:K, end]))

        for (i,Z) in enumerate(Zs)
            design_response_row!(cache_vec, Z, j)
            μs_mat[i,j] = dot(βs_list[j], view(cache_vec,1:K))
        end
    end
    μs = vec(mean(μs_mat; dims=2))
    βs = vec(mean(hcat(βs_list...), dims=2))
    FittedAuroral(βs = βs, βs_list = βs_list,
                μs = μs, μs_mat = μs_mat,
                method = auroral)
end

struct CoeyCunningham <: AbstractAurora end

function StatsBase.fit(auroral::CoeyCunningham, Zs::AbstractVector{<:ReplicatedSample})
    K = nobs(Zs[1])
    n = length(Zs)
    any(nobs.(Zs) != K) || throw("All Zs should have the same number of replicates.")
    μs_mat = zeros(n,K)

    cache_means = zeros(n)
    cache_responses = zeros(n)
    cache_vec = zeros(K-1)

    βs_list = Vector{typeof(cache_vec)}(undef, K)

    for j in Base.OneTo(K)
        cache_responses .= StatsBase.response.(Zs, j)
        fill!(cache_means, 0)

        for (i,Z) in enumerate(Zs)
            cache_means[i] = mean(design_row!(cache_vec, Z, j))
        end

        β_slope = cov(cache_means, cache_responses)/var(cache_means)
        β_intercept = mean(cache_responses) - β_slope*mean(cache_means)

        βs_list[j] = [β_intercept; β_slope]

        μs_mat[:,j] = β_intercept .+ β_slope .* cache_means
    end
    μs = vec(mean(μs_mat; dims=2))
    βs = vec(mean(hcat(βs_list...), dims=2))
    FittedAuroral(βs = βs, βs_list = βs_list,
                μs = μs, μs_mat = μs_mat,
                method = auroral)
end

# Plotting

@recipe function f(aur::Aurora.FittedAuroral{Auroral}; stick_color=:purple, x_nudge = 0.0)

    βs = aur.βs
    K = length(βs) - 1

    xs = (0:K) .+ x_nudge

    xlabs = ["Int."; [L"X^{(%$k)}" for k in Base.OneTo(K)]]

	background_color_legend --> :transparent
	foreground_color_legend --> :transparent
    grid --> nothing
    label --> ""
    linewidth --> 2.5
    seriestype := :sticks
    seriescolor := stick_color

    @series begin
        seriestype := :scatter
        seriescolor := stick_color
        markerstrokewidth := 0
        markersize :=5
        label := ""
        xs, βs
    end

    xticks := (xs, xlabs)
    yguide --> L"\hat{\beta}"
    xs, βs
end


using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

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

                     @variable(m, -2.9922365365158297 <= q <= 1.3959867501971397)

                     add_NL_constraint(m, :(log(1 + exp(-0.7172866311917714 + 0.8530094465195428*log(1 + exp(0.49975691401319544 + 0.817850285997455*log(1 + exp(-0.9044530867886094 + -0.3504940249833268*log(1 + exp(0.9239332442221326 + 0.41157080299867177*$(x[1]) + 0.6856618842893334*$(x[2]) + 0.23250449051531152*$(x[3]) + -0.8171094215290706*$(x[4]))) + -0.925387475816613*log(1 + exp(-0.1441575648301785 + -0.7324773919132506*$(x[1]) + 0.1666248349273194*$(x[2]) + 0.5034013066229974*$(x[3]) + -0.3351866390106042*$(x[4]))))) + -0.15011340339301915*log(1 + exp(-0.49866226232151334 + 0.31623027699891626*log(1 + exp(0.9239332442221326 + 0.41157080299867177*$(x[1]) + 0.6856618842893334*$(x[2]) + 0.23250449051531152*$(x[3]) + -0.8171094215290706*$(x[4]))) + 0.1823838704363938*log(1 + exp(-0.1441575648301785 + -0.7324773919132506*$(x[1]) + 0.1666248349273194*$(x[2]) + 0.5034013066229974*$(x[3]) + -0.3351866390106042*$(x[4]))))))) + -0.2302663292873568*log(1 + exp(0.6056562038214484 + 0.017530899400252764*log(1 + exp(-0.9044530867886094 + -0.3504940249833268*log(1 + exp(0.9239332442221326 + 0.41157080299867177*$(x[1]) + 0.6856618842893334*$(x[2]) + 0.23250449051531152*$(x[3]) + -0.8171094215290706*$(x[4]))) + -0.925387475816613*log(1 + exp(-0.1441575648301785 + -0.7324773919132506*$(x[1]) + 0.1666248349273194*$(x[2]) + 0.5034013066229974*$(x[3]) + -0.3351866390106042*$(x[4]))))) + 0.9141090581590978*log(1 + exp(-0.49866226232151334 + 0.31623027699891626*log(1 + exp(0.9239332442221326 + 0.41157080299867177*$(x[1]) + 0.6856618842893334*$(x[2]) + 0.23250449051531152*$(x[3]) + -0.8171094215290706*$(x[4]))) + 0.1823838704363938*log(1 + exp(-0.1441575648301785 + -0.7324773919132506*$(x[1]) + 0.1666248349273194*$(x[2]) + 0.5034013066229974*$(x[3]) + -0.3351866390106042*$(x[4]))))))))) + log(1 + exp(0.11802092343800519 + -0.23198020992542823*log(1 + exp(0.49975691401319544 + 0.817850285997455*log(1 + exp(-0.9044530867886094 + -0.3504940249833268*log(1 + exp(0.9239332442221326 + 0.41157080299867177*$(x[1]) + 0.6856618842893334*$(x[2]) + 0.23250449051531152*$(x[3]) + -0.8171094215290706*$(x[4]))) + -0.925387475816613*log(1 + exp(-0.1441575648301785 + -0.7324773919132506*$(x[1]) + 0.1666248349273194*$(x[2]) + 0.5034013066229974*$(x[3]) + -0.3351866390106042*$(x[4]))))) + -0.15011340339301915*log(1 + exp(-0.49866226232151334 + 0.31623027699891626*log(1 + exp(0.9239332442221326 + 0.41157080299867177*$(x[1]) + 0.6856618842893334*$(x[2]) + 0.23250449051531152*$(x[3]) + -0.8171094215290706*$(x[4]))) + 0.1823838704363938*log(1 + exp(-0.1441575648301785 + -0.7324773919132506*$(x[1]) + 0.1666248349273194*$(x[2]) + 0.5034013066229974*$(x[3]) + -0.3351866390106042*$(x[4]))))))) + 0.3694431087223533*log(1 + exp(0.6056562038214484 + 0.017530899400252764*log(1 + exp(-0.9044530867886094 + -0.3504940249833268*log(1 + exp(0.9239332442221326 + 0.41157080299867177*$(x[1]) + 0.6856618842893334*$(x[2]) + 0.23250449051531152*$(x[3]) + -0.8171094215290706*$(x[4]))) + -0.925387475816613*log(1 + exp(-0.1441575648301785 + -0.7324773919132506*$(x[1]) + 0.1666248349273194*$(x[2]) + 0.5034013066229974*$(x[3]) + -0.3351866390106042*$(x[4]))))) + 0.9141090581590978*log(1 + exp(-0.49866226232151334 + 0.31623027699891626*log(1 + exp(0.9239332442221326 + 0.41157080299867177*$(x[1]) + 0.6856618842893334*$(x[2]) + 0.23250449051531152*$(x[3]) + -0.8171094215290706*$(x[4]))) + 0.1823838704363938*log(1 + exp(-0.1441575648301785 + -0.7324773919132506*$(x[1]) + 0.1666248349273194*$(x[2]) + 0.5034013066229974*$(x[3]) + -0.3351866390106042*$(x[4]))))))))) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m


module IfOnTop

include("./../../constraints/geometric/constants.jl")
using JuMP
using .Constants

export if_on_top

function if_on_top(m)

    #if (flbz_front_left_bot_z[i] == flbz_front_left_bot_z[j] + hob_height_of_box[j] &&
    #    flbz_front_left_bot_z[i] - hob_height_of_box[j] >= 0 &&
    #    flbx_front_left_bot_x[j] > flbx_front_left_bot_x[i] - box_i_x_axis_cover &&
    #    flbx_front_left_bot_x[j] < flbx_front_left_bot_x[i] + box_i_x_axis_cover &&
    #    flby_front_left_bot_y[j] > flby_front_left_bot_y[i] - box_i_y_axis_cover &&
    #    flby_front_left_bot_y[j] < flby_front_left_bot_y[i] + box_i_y_axis_cover)

    # IF 1
    @variable(m, if11_on_top[i=1:length(boxes), j=1:length(boxes); i!=j], Bin)
    @variable(m, if12_on_top[i=1:length(boxes), j=1:length(boxes); i!=j], Bin)
    @variable(m, if1_on_top[i=1:length(boxes), j=1:length(boxes); i!=j] >= 0, Bin)

    #Hz1 >= x-b+e
    #H(1-z1) >= b-x-e
    @constraint(m, c1_if11_on_top[i=1:length(boxes), j=1:length(boxes); i!=j],
        H * m[:if11_on_top][i, j] >=
        m[:flbz_front_left_bot_z][i] - (m[:flbz_front_left_bot_z][j] + m[:hob_height_of_box][j]) + 0.001
    )
    @constraint(m, c2_if11_on_top[i=1:length(boxes), j=1:length(boxes); i!=j],
        H * (1 - m[:if11_on_top][i, j]) >=
        (m[:flbz_front_left_bot_z][j] + m[:hob_height_of_box][j]) - m[:flbz_front_left_bot_z][i] - 0.001
    )

    #Hz2 >= b-x+e
    #H(1-z2) >= x-b-e
    @constraint(m, c1_if12_on_top[i=1:length(boxes), j=1:length(boxes); i!=j],
        H * m[:if12_on_top][i, j] >=
        (m[:flbz_front_left_bot_z][j] + m[:hob_height_of_box][j]) - m[:flbz_front_left_bot_z][i] + 0.001
    )
    @constraint(m, c2_if12_on_top[i=1:length(boxes), j=1:length(boxes); i!=j],
        H * (1 - m[:if12_on_top][i, j]) >=
        m[:flbz_front_left_bot_z][i] - (m[:flbz_front_left_bot_z][j] + m[:hob_height_of_box][j]) - 0.001
    )

    #If if11, if12 = 1, then box i and j have z-bot/top value within interval -e <= z <= +e, hence, expression = 1
    @constraint(m, c_if1_on_top[i=1:length(boxes), j=1:length(boxes); i!=j],
        m[:if11_on_top][i, j] + m[:if12_on_top][i, j] >=
        2 * m[:if1_on_top][i, j]
    )

    #Xj < Xi + Xcoveri
    @variable(m, if2_on_top[i=1:length(boxes), j=1:length(boxes); i!=j], Bin)
    @constraint(m, c1_if2_on_top[i=1:length(boxes), j=1:length(boxes); i!=j],
        L * (1 - m[:if2_on_top][i, j]) >=
        m[:flbx_front_left_bot_x][j] + 0.001 - (m[:flbx_front_left_bot_x][i] + m[:box_x_axis_cover][i])
    )

    @constraint(m, c2_if2_on_top[i=1:length(boxes), j=1:length(boxes); i!=j],
        L * m[:if2_on_top][i, j] >=
        (m[:flbx_front_left_bot_x][i] + m[:box_x_axis_cover][i]) - m[:flbx_front_left_bot_x][j] - 0.001
    )

    # Xj + Xcoverj > Xi
    @variable(m, if3_on_top[i=1:length(boxes), j=1:length(boxes); i!=j], Bin)
    @constraint(m, c1_if3_on_top[i=1:length(boxes), j=1:length(boxes); i!=j],
        L * (1 - m[:if3_on_top][i, j]) >=
        m[:flbx_front_left_bot_x][i] + 0.001 - (m[:flbx_front_left_bot_x][j] + m[:box_x_axis_cover][j])
    )

    @constraint(m, c2_if3_on_top[i=1:length(boxes), j=1:length(boxes); i!=j],
        L * m[:if3_on_top][i, j] >=
        (m[:flbx_front_left_bot_x][j] + m[:box_x_axis_cover][j]) - m[:flbx_front_left_bot_x][i] - 0.001
    )

    # Yj < Yi + Ycoveri
    @variable(m, if4_on_top[i=1:length(boxes), j=1:length(boxes); i!=j], Bin)
    @constraint(m, c1_if4_on_top[i=1:length(boxes), j=1:length(boxes); i!=j],
        W * (1 - m[:if4_on_top][i, j]) >=
        m[:flby_front_left_bot_y][j] + 0.001 - (m[:flby_front_left_bot_y][i] + m[:box_y_axis_cover][i])
    )

    @constraint(m, c2_if4_on_top[i=1:length(boxes), j=1:length(boxes); i!=j],
        W * m[:if4_on_top][i, j] >=
        (m[:flby_front_left_bot_y][i] + m[:box_y_axis_cover][i]) - m[:flby_front_left_bot_y][j] - 0.001
    )

    # Yj + Ycoverj > Yi
    @variable(m, if5_on_top[i=1:length(boxes), j=1:length(boxes); i!=j], Bin)
    @constraint(m, c1_if5_on_top[i=1:length(boxes), j=1:length(boxes); i!=j],
        W * (1 - m[:if5_on_top][i, j]) >=
        m[:flby_front_left_bot_y][i] + 0.001 - (m[:flby_front_left_bot_y][j] + m[:box_y_axis_cover][j])
    )

    @constraint(m, c2_if5_on_top[i=1:length(boxes), j=1:length(boxes); i!=j],
        W * m[:if5_on_top][i, j] >=
        (m[:flby_front_left_bot_y][j] + m[:box_y_axis_cover][j]) - m[:flby_front_left_bot_y][i] - 0.001
    )

    @variable(m, if_on_top[i=1:length(boxes), j=1:length(boxes); i!=j], Bin)
    @constraint(m, c_if_on_top[i=1:length(boxes), j=1:length(boxes); i!=j],
        m[:if1_on_top][i, j] +
        m[:if2_on_top][i, j] +
        m[:if3_on_top][i, j] +
        m[:if4_on_top][i, j] +
        m[:if5_on_top][i, j] >=
        5 * m[:if_on_top][i, j]
    )

    return m
end

end


using NewsvendorModel, Distributions
using Test, Documenter

@testset "Testing NVModel initialization" begin include("nvmodel_test.jl") end

@testset "Testing metrics" begin include("newsvendorproblem_test.jl") end

doctest(NewsvendorModel)


# ---
# title: 879. Profitable Schemes
# id: problem879
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Dynamic Programming
# link: <https://leetcode.com/problems/profitable-schemes/description/>
# hidden: true
# ---
#
# There is a group of `G` members, and a list of various crimes they could
# commit.
#
# The `ith` crime generates a `profit[i]` and requires `group[i]` members to
# participate in it.
#
# If a member participates in one crime, that member can't participate in
# another crime.
#
# Let's call a _profitable  scheme_ any subset of these crimes that generates at
# least `P` profit, and the total number of members participating in that subset
# of crimes is at most `G`.
#
# How many schemes can be chosen?  Since the answer may be very large, **return
# it modulo** `10^9 + 7`.
#
#
#
# **Example 1:**
#
#
#
#     Input: G = 5, P = 3, group = [2,2], profit = [2,3]
#     Output: 2
#     Explanation:
#     To make a profit of at least 3, the group could either commit crimes 0 and 1, or just crime 1.
#     In total, there are 2 schemes.
#
#
# **Example 2:**
#
#
#
#     Input: G = 10, P = 5, group = [2,3,5], profit = [6,7,8]
#     Output: 7
#     Explanation:
#     To make a profit of at least 5, the group could commit any crimes, as long as they commit one.
#     There are 7 possible schemes: (0), (1), (2), (0,1), (0,2), (1,2), and (0,1,2).
#
#
#
#
# **Note:**
#
#   1. `1 <= G <= 100`
#   2. `0 <= P <= 100`
#   3. `1 <= group[i] <= 100`
#   4. `0 <= profit[i] <= 100`
#   5. `1 <= group.length = profit.length <= 100`
#
#
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


using PyPlot
using Jagot.plotting
plot_style("ggplot")
using PyPlotRecipes
using PyCall
PathEffects = pyimport("matplotlib.patheffects")
using Statistics
using Random
using Colors

using LinearAlgebra
using ArnoldiMethod

using CompactBases
import CompactBases: applied
using IntervalSets

function mean_color(color::String)
    c = parse(Colorant, color)
    mean([c.r,c.g,c.b])
end

lerp(a,b,t) = (1-t)*a + t*b

mean_position(x, ϕ) = ϕ'*Diagonal(x)*ϕ/(ϕ'ϕ)

function savedocfig(name,dir="figures")
    fig = gcf()
    filename = joinpath(@__DIR__, "src", dir, "$(name).svg")
    savefig(filename,
            transparent=true,
            facecolor=fig.get_facecolor())
    close(fig)
    if isfile(filename)
        println("Saved $(name) to $(filename)")
    else
        @warn "Saving $(name) to $(filename) failed"
    end
end

function logo()
    t = ArbitraryKnotSet(3, [0.0, 1, 1, 3, 4, 6], 1, 3)
    r = range(first(t), stop=last(t), length=301)
    R = BSpline(t)
    χ = R[r, :]

    cfigure("logo",figsize=(7,3)) do
        plot(r, χ, length(r) < 30 ? ".-" : "-", linewidth=4)
        margins(0.1, 0.1)
        axis("off")
    end
    savefig("docs/src/assets/logo.svg", transparent=true)
end

function simple_example()
    a,b = 0.0,1.0 # Extents
    N = 13 # Number of nodes
    k = 5 # Order of FE-DVR/B-splines
    x = range(a, stop=b, length=1000)

    Δx = (b-a)/(N+1) # Grid spacing
    # Standard, uniform finite-differences
    fd = FiniteDifferences(N, Δx)
    # Staggered, uniform finite-differences
    sfd_uni = StaggeredFiniteDifferences(N, Δx)
    # Staggered, non-uniform finite-differences
    sfd_nonuni = StaggeredFiniteDifferences(0.01, 0.5, 0.1, b)

    # Finite-element boundaries
    tf = range(a, stop=b, length=N+2)
    # We can vary the polynomial order in each element
    forder = vcat(7, fill(4,length(tf)-2))
    # By indexing the second dimension, we can implement Dirichlet0
    # boundary conditions.
    fem = FEDVR(tf, forder)[:,2:end-1]

    tb = ExpKnotSet(k, -2.0, log10(b), N+1)
    splines = BSpline(tb)[:,2:end-1]

    bases = [fd, sfd_uni, sfd_nonuni, fem, splines]
    m = length(bases)

    cfigure("simple example", figsize=(6,9)) do
        for (i,basis) in enumerate(bases)
            csubplot(m,1,i,nox=i<m) do
                plot(x, basis[x,:])
            end
        end
    end
    savedocfig("simple_example")
end

function restricted_bases()
    B = FEDVR(range(0, stop=10.0, length=5), 7)
    B̃ = B[:, 3:17]

    x = range(0, stop=10.0, length=1000)
    χ = B[x, :]
    χ̃ = B̃[x, :]

    cfigure("restricted bases") do
        csubplot(211,nox=true) do
            plot(x, χ)
            ylabel("Original basis")
        end
        csubplot(212) do
            plot(x, χ̃)
            ylabel("Restricted basis")
        end
    end
    savedocfig("restricted_bases")
end

function densities()
    f = x -> sin(2π*x)
    g = x -> x*exp(-x)
    h = x -> f(x)*g(x)

    rmin = 0.01
    rmax = 10.0
    ρ = 0.05
    N = ceil(Int, rmax/ρ)
    k = 7
    Nn = 71

    ρmin=rmin
    ρmax=0.5
    α=0.01

    coeff_centers = Vector{Vector{Float64}}()
    density_coeffs = Vector{Vector{Float64}}()
    errors = Vector{Vector{Float64}}()
    ϕ = Vector{Vector{Float64}}()

    rr = 10.0 .^ range(-2, stop=log10(rmax), length=3000)

    for R in [FiniteDifferences(N, ρ),
              StaggeredFiniteDifferences(ρmin, ρmax, α, rmax),
              FEDVR(range(0, stop=rmax, length=Nn), k),
              BSpline(LinearKnotSet(k, 0, rmax, Nn))]
        r = axes(R,1)

        cf = R \ f.(r)
        cg = R \ g.(r)
        ch = R \ h.(r)

        ρ = Density(applied(*,R,cf), applied(*,R,cg))

        push!(coeff_centers, centers(R))
        push!(density_coeffs, ρ.ρ)
        push!(errors, ρ.ρ-ch)
        push!(ϕ, R[rr,:]*ρ.ρ)
    end

    cfigure("mutual density", figsize=(7,10)) do
        csubplot(411,nox=true) do
            semilogx(rr, f.(rr), linewidth=1.0)
            semilogx(rr, g.(rr), linewidth=1.0)
            semilogx(rr, h.(rr), linewidth=2.0)
            legend([L"f(x)", L"g(x)", L"h(x)"])
        end
        csubplot(412, nox=true) do
            for ϕ in ϕ
                loglog(rr, abs.(ϕ-h.(rr)))
            end
            yl = ylim()
            ylim(max(yl[1],1e-15), yl[2])
            legend(["Finite-differences",
                    "Log–linear finite-differences",
                    "FE-DVR",
                    "B-splines"])
            ylabel("Reconstruction error")
        end
        csubplot(413, nox=true) do
            for (r,c) in zip(coeff_centers, density_coeffs)
                semilogx(r, c, ".-", linewidth=1.0)
            end
            legend(["Finite-differences",
                    "Log–linear finite-differences",
                    "FE-DVR",
                    "B-splines"])
            ylabel("Expansion coefficients")
        end
        csubplot(414) do
            for (r,e) in zip(coeff_centers, errors)
                loglog(r, abs.(e), ".-", linewidth=1.0)
            end
            legend(["Finite-differences",
                    "Log–linear finite-differences",
                    "FE-DVR",
                    "B-splines"])
            ylabel("Coefficient error")
            xlabel(L"r")
        end
    end

    savedocfig("mutual_densities")
end

function diagonal_operators()
    Z = 1.0
    rmax = 20*(2^1.05)/Z

    ρ = 0.1/Z
    N = ceil(Int, rmax/ρ)
    ufd = StaggeredFiniteDifferences(N, ρ)

    ρmin=0.1/Z
    ρmax=0.6
    α=0.002
    nufd = StaggeredFiniteDifferences(ρmin, ρmax, α, rmax)

    fedvr = FEDVR(range(0, stop=rmax, length=20),
                  vcat(10,fill(7,18)))[:,2:end-1]

    bsplines = BSpline(ExpKnotSet(5, -2.0, log10(rmax), 100))[:,2:end-1]

    R = bsplines

    r = axes(R,1)
    rr = clamp.(10.0 .^ range(min(-2,log10(r[1])), stop=log10(rmax), length=3000),
                0, rightendpoint(R))
    χ = R[rr,:]

    r = axes(R,1)
    V = r -> (0.314269680527354*r^2 + 0.20951312035157*r)*exp(-3*r/2);
    Vc = R \ V.(r);
    rc = R \ identity.(r);

    L = DiagonalOperator(applied(*, R, Vc));

    Vr = L*rc

    Vr_exact = rr .* V.(rr)

    cfigure("diagonal operators") do
        csubplot(211,nox=true) do
            semilogx(rr, χ*Vr, label=L"V(r)r")
            semilogx(rr, Vr_exact, "--", label=L"$V(r)r$ exact")
            semilogx(rr, V.(rr), label=L"V(r)")
            yl = ylim()
            semilogx(rr, rr, label=L"r")
            ylim(yl)
            legend()
        end
        csubplot(212) do
            for (R,label) in [
                (ufd, "Uniform finite-differences"),
                (nufd, "Non-uniform finite-differences"),
                (fedvr, "FE-DVR"),
                (bsplines, "B-splines")
            ]
                r = axes(R,1)
                Vc = R \ V.(r)
                rc = R \ identity.(r)

                L = DiagonalOperator(applied(*, R, Vc))

                Vr = L*rc
                χ = R[rr,:]
                loglog(rr, abs.(χ*Vr - Vr_exact), label=label)
            end
            xlabel(L"r")
            ylabel("Error")
            legend()
        end
    end

    savedocfig("diagonal_operators")
end

macro echo(expr)
    println(expr)
    :(@time $expr)
end

@info "Documentation plots"
fig_dir = joinpath(@__DIR__, "src", "figures")
mkpath(fig_dir)
@assert isdir(fig_dir)
@echo logo()
@echo simple_example()
@echo restricted_bases()
include("bspline_plots.jl")
include("fd_plots.jl")
@echo densities()
@echo diagonal_operators()


# Test conversions ta various normal forms
@test to_nnf(@julog(and(a))) == @julog a
@test to_nnf(@julog(true => not(and(not(!a), b, or(not(c), false))))) ==
    @julog(or(not(a), not(b), c))

@test to_cnf(@julog(and(a))) == @julog and(or(a))
@test to_cnf(@julog(and(and(or(a, and(b, or(c, d))), or(e, f)), and(not(x), or(y, z))))) ==
    @julog(and(or(a, b), or(a, c, d), or(e, f), or(not(x)), or(y, z)))
@test to_cnf(@julog(true => not(and(not(!a), b, or(not(c), false))))) ==
    @julog(and(or(not(a), not(b), c)))

@test to_dnf(@julog(and(a))) == @julog or(and(a))
@test to_dnf(@julog(or(or(and(a, or(b, and(c, d))), and(e, f)), or(not(x), and(y, z))))) ==
    @julog(or(and(a, b), and(a, c, d), and(e, f), and(not(x)), and(y, z)))
@test to_dnf(@julog(true => not(and(not(!a), b, or(not(c), false))))) ==
    @julog(or(and(not(a)), and(not(b)), and(c)))

# Test flattening of conjuctions and disjunctions
@test flatten_conjs(@julog(and(a, and(b, c)))) == @julog Const[a, b, c]
@test flatten_conjs(@julog([a, and(b, c)])) == @julog Const[a, b, c]
@test flatten_disjs(@julog(or(a, or(b, c)))) == @julog Const[a, b, c]
@test flatten_disjs(@julog([a, or(b, c)])) == @julog Const[a, b, c]

# Test instantiation of universal quantifiers
clauses = @julog [
    block(a) <<= true,
    block(b) <<= true,
    block(c) <<= true,
    holding(a) <<= true
]
universal_term = @julog(forall(block(X), not(holding(X))))
@test deuniversalize(universal_term, clauses) ==
    @julog and(not(holding(a)), not(holding(b)), not(holding(c)))
universal_clause = Clause(@julog(handempty), [universal_term])
@test deuniversalize(universal_clause, clauses) ==
    @julog handempty <<= and(not(holding(a)), not(holding(b)), not(holding(c)))

# Test regularization of clause bodies
clauses = @julog [
    binary(X) <<= or(woman(X), man(X)) & not(nonbinary(X)),
    bigender(X) <<= and(woman(X), man(X)),
    nonbinary(X) <<= or(genderqueer(X), agender(X), thirdgender(X))
]

regularized = @julog [
    binary(X) <<= woman(X) & not(nonbinary(X)),
    binary(X) <<= man(X) & not(nonbinary(X)),
    bigender(X) <<= woman(X) & man(X),
    nonbinary(X) <<= genderqueer(X),
    nonbinary(X) <<= agender(X),
    nonbinary(X) <<= thirdgender(X)
]

@test regularize_clauses(clauses) == regularized


# Graph
mutable struct Graph
  nN     :: Int
  edges  :: Vector{Vector{Int}}
  big    :: Vector{Bool}
  visits :: Vector{Int}
  tags   :: Dict{String,Int}
  paths  :: Vector{Vector{Int}}
  hasVisitedSmallTwice :: Bool
end

Graph() = Graph(0, [], [], [], Dict{String,Int}(), [], false)

# Adds a new node to the Graph
function addNode!(G::Graph, n::String) :: Int
  haskey(G.tags,n) && return G.tags[n]
  G.nN += 1
  push!(G.edges,[])
  islowercase(n[1]) ? push!(G.big,false) : push!(G.big,true)
  push!(G.visits,0)
  G.tags[n] = G.nN
  return G.tags[n]
end

# Adds a new edge to the Graph
function addEdge!(G::Graph, v::Tuple{String,String})
  idx1 = haskey(G.tags,v[1]) ? G.tags[v[1]] : addNode!(G,v[1])
  idx2 = haskey(G.tags,v[2]) ? G.tags[v[2]] : addNode!(G,v[2])
  push!(G.edges[idx1],idx2)
  push!(G.edges[idx2],idx1)
end

# Recursive function that computes the number of paths that go from 'istart' to 'iend'.
function paths!(G::Graph, istart::Int, iend::Int) :: Int
  (istart == iend) && return 1
  G.visits[istart] += 1
  res = 0
  for e in G.edges[istart]
    if G.visits[e] == 0 || G.big[e]
      res += paths!(G,e,iend)
    end
  end
  G.visits[istart] -= 1
  return res
end

paths!(G::Graph, istart::String, iend::String) = paths!(G,G.tags[istart],G.tags[iend])

# Same as paths, modified for problem 2.
function paths2!(G::Graph, istart::Int, iend::Int) :: Int
  (istart == iend) && return 1
  G.visits[istart] += 1
  res = 0
  for e in G.edges[istart]
    if G.visits[e] == 0 || G.big[e]
      res += paths2!(G,e,iend)
    elseif (G.tags["start"] !== e) && (G.visits[e] == 1) && !(G.hasVisitedSmallTwice)
      G.hasVisitedSmallTwice = true
      res += paths2!(G,e,iend)
      G.hasVisitedSmallTwice = false
    end
  end
  G.visits[istart] -= 1
  return res
end

paths2!(G::Graph, istart::String, iend::String) = paths2!(G,G.tags[istart],G.tags[iend])

# Same as paths, but also saves the actual paths in G.paths
function paths_visual!(G::Graph, istart::Int, iend::Int, p::Vector{Int}) :: Int
  if istart == iend
    push!(G.paths,deepcopy(p))
    return 1
  end

  G.visits[istart] += 1
  res = 0
  for e in G.edges[istart]
    if G.visits[e] == 0 || G.big[e]
      res += paths_visual!(G,e,iend,vcat(p,e))
    end
  end
  G.visits[istart] -= 1
  return res
end

paths_visual!(G::Graph, istart::String, iend::String) = paths_visual!(G,G.tags[istart],G.tags[iend],[G.tags[istart]])

# Prints the paths saved in G.paths
function printPaths(G)
  invMap = Dict(value => key for (key, value) in G.tags)
  for p in G.paths
    println(map(x -> invMap[x], p))
  end
end

parseLine(line::String) = tuple(map(s -> String(s),split(line,'-'))...)

function readInput()
  f = open("./inputs/day12.txt")
  G = Graph()

  edges = map( line -> parseLine(line), eachline(f))
  map(e -> addEdge!(G,e), edges)
  return G
end

println("Problem 1: ")
G = readInput()
res = paths!(G,"start","end")
printPaths(G)
println(res)

println("Problem 2: ")
G = readInput()
res = paths2!(G,"start","end")
println(res)


import Base: @propagate_inbounds

const IsingSpin = Int64
const IsingDistribution = IsingSpin[-1,1]
const IsingConf = Array{IsingSpin}

const IsingTc = 1/(1/2*log(1+sqrt(2)))

"""
Famous Ising model on a cubic lattice.

    IsingModel(; dims, L)

Create Ising model on `dims`-dimensional cubic lattice
with linear system size `L`.
"""
@with_kw_noshow struct IsingModel{C<:AbstractLattice} <: Model # noshow because we override it below
    L::Int
    dims::Int
    l::C = choose_lattice(IsingModel, dims, L)
    energy::Ref{Float64} = Ref(0.0)
end

init!(mc::MC, m::IsingModel) = energy(mc, m, mc.conf)

function choose_lattice(::Type{IsingModel}, dims::Int, L::Int)
    if dims == 1
        return Chain(L)
    elseif dims == 2
        return SquareLattice(L)
    else
        return CubicLattice(dims, L)
    end
end

"""
    IsingModel(params::Dict)
    IsingModel(params::NamedTuple)

Create an Ising model with (keyword) parameters as specified in the
dictionary/named tuple `params`.
"""
IsingModel(params::Dict{Symbol, T}) where T = IsingModel(; params...)
IsingModel(params::NamedTuple) = IsingModel(; params...)

# convenience
@inline Base.ndims(m::IsingModel) = m.dims

# cosmetics
import Base.summary
import Base.show
Base.summary(model::IsingModel) = "$(model.dims)D-Ising model"
# Base.show(io::IO, model::IsingModel{LT}) where LT<:AbstractCubicLattice =
    # print(io, "$(model.dims)D-Ising model, L=$(model.L) ($(model.l.sites) sites)")
function Base.show(io::IO, model::IsingModel{LT}) where LT<:AbstractLattice
    print(io,
        "Ising model on $(replace(string(LT), "MonteCarlo."=>"")), " *
        "L=$(model.L) ($(model.l.sites) sites)"
    )
end
Base.show(io::IO, m::MIME"text/plain", model::IsingModel) = print(io, model)

# implement `Model` interface
@inline lattice(m::IsingModel) = m.l

# implement `MC` interface
Base.rand(::Type{MC}, m::IsingModel) = rand(IsingDistribution, fill(m.L, ndims(m))...)

@propagate_inbounds @bm function propose_local(mc::MC, m::IsingModel, i::Int, conf::IsingConf)
    field = 0.0
    @inbounds for trg in neighbors(m.l, i)
        field += conf[trg]
    end
    delta_E = 2.0 * conf[i] * field
    return delta_E, nothing
end

@propagate_inbounds @bm function accept_local!(
        mc::MC, m::IsingModel, i::Int, conf::IsingConf, delta_E::Float64, passthrough
    )
    m.energy[] += delta_E
    conf[i] *= -1
    nothing
end

# optimized for 2D case
@propagate_inbounds @bm function propose_local(
        mc::MC, m::IsingModel{SquareLattice}, i::Int, conf::IsingConf
    )
    neighs = m.l.neighs
    @inbounds delta_E = 2.0 * conf[i] * (
        conf[neighs[1, i]] + conf[neighs[2, i]] +
        conf[neighs[3, i]] + conf[neighs[4, i]]
    )
    return delta_E, nothing
end

# optional functions/methods
"""
    global_move(mc::MC, m::IsingModel, conf::IsingConf) -> accepted::Bool

Constructs a Wolff cluster spinflip for configuration `conf`.
Returns wether a cluster spinflip has been performed (any spins have been flipped).
"""
@bm function global_move(mc::MC, m::IsingModel, conf::IsingConf)
    N = length(lattice(m))
    beta = mc.p.beta

    cluster = Array{Int, 1}()
    tocheck = Array{Int, 1}()

    s = rand(1:N)
    push!(tocheck, s)
    push!(cluster, s)

    while !isempty(tocheck)
        cur = pop!(tocheck)
        @inbounds for n in neighbors(m.l, cur)

            if conf[cur] == conf[n] && !(n in cluster) && rand() < (1 - exp(- 2.0 * beta))
                push!(tocheck, n)
                push!(cluster, n)
            end

        end
    end

    for spin in cluster
        conf[spin] *= -1
    end
    model.energy[] = energy(mc, m, conf)

    return length(cluster)>1
end

"""
    energy(mc::MC, m::IsingModel, conf::IsingConf)

Calculate energy of Ising configuration `conf` for Ising model `m`.
"""
function energy(mc::MC, m::IsingModel, conf::IsingConf)
    E = 0.0
    for (src, trg) in neighbors(m.l)
        E -= conf[src]*conf[trg]
    end
    m.energy[] = E
    return E
end

function energy(mc::MC, m::IsingModel{LT}, conf::IsingConf) where {
        LT <: Union{Chain, SquareLattice, CubicLattice}
    }
    E = 0.0
    @inbounds for (src, trg) in neighbors(m.l, Val(false))
        E -= conf[src]*conf[trg]
    end
    m.energy[] = E
    return E
end

"""
    energy(mc::MC, m::IsingModel{SquareLattice}, conf::IsingConf)

Calculate energy of Ising configuration `conf` for 2D Ising model `m`.
This method is a faster variant of the general method for the
square lattice case. (It is roughly twice as fast in this case.)
"""
function energy(mc::MC, m::IsingModel{SquareLattice}, conf::IsingConf)
    neighs = m.l.neighs
    E = 0.0
    @inbounds @simd for i in 1:length(lattice(m))
        E -= conf[i]*conf[neighs[1,i]] + conf[i]*conf[neighs[2,i]]
    end
    m.energy[] = E
    return E
end

function save_model(
        file::JLDFile,
        m::IsingModel,
        entryname::String="Model"
    )
    write(file, entryname * "/VERSION", 1)
    write(file, entryname * "/tag", "IsingModel")

    write(file, entryname * "/L", m.L)
    write(file, entryname * "/dims", m.dims)
    save_lattice(file, m.l, entryname * "/l")
    write(file, entryname * "/energy", m.energy[])
    nothing
end

#     load_model(data, ::Type{<: IsingModel})
#
# Loads an IsingModel from a given `data` dictionary produced by
# `JLD.load(filename)`.
function _load(data, ::Val{:IsingModel})
    if !(data["VERSION"] == 1)
        throw(ErrorException("Failed to load IsingModel version $(data["VERSION"])"))
    end

    l = _load(data["l"], to_tag(data["l"]))
    model = IsingModel(
        L = data["L"],
        dims = data["dims"],
        l = l
    )
    model.energy[] = data["energy"]
    model
end
to_tag(::Type{<: IsingModel}) = Val(:IsingModel)

compress(mc::MC, ::IsingModel, c) = BitArray(c .== 1)
compressed_conf_type(::Type{<: MC}, ::Type{<: IsingModel}) = BitArray
function decompress(mc::MC{M, CT}, ::IsingModel, c) where {M, CT}
    CT(2c .- 1)
end

include("measurements.jl")


# polyhelper function guesses what type of poly you want
export polyhelper

function polyhelper(y)
    LGRPoly(y)
end

function polyhelper(x,y)
    if length(x) < 6
        vandpoly(x,y)
    else
        LagrangePoly(x,y) # test speed of this against fitting function of Polynomials package
    end
end

#function polyhelper(x,y,b::Bound)
#    if length(x) < 6
#        vandpoly(x,y,b)
#    else
#        LagrangePoly(x,y,b)
#    end
#end

# add automatically gen GlobalPoly of local polys


using DelimitedFiles
bingodata = readdlm("bingodata.txt",',',Int32)
@show bingodata

bingoboards = readdlm("bingoboards.txt",' ',Int32,'\n')
@show bingoboards
@show typeof(bingoboards)

function get_bingo_boards(bingoboarddata)
    index = 1

    global boardcount = Int(length(bingoboarddata)/25)

    boards = Vector{Matrix{Int32}}(undef,boardcount)
    println( "# of boards: $boardcount")

    for b = 1:boardcount
        #println("Get board $b")
        boards[b] = bingoboarddata[index:index+4, :]

        #println("boards[$b] = $(boards[b])")
        index += 5
    end

    return boards
end

boards = get_bingo_boards(bingoboards)
@show boards

function draw_ball(ballnumber, boards)
    for bi in eachindex(boards)
        if boards[bi][1] == 10_000
            continue
        end

        #println("boards[$bi] = $(boards[bi])")
        for bb in eachindex(boards[bi])
            if boards[bi][bb] == ballnumber
                #println("Replacing $ballnumber")
                boards[bi][bb] = -1
            end
        end
    end
end

function check_for_bingo(boards)
    bingo = fill(-1,5)
    for bi in eachindex(boards)
        if boards[bi][1] == 10_000
            continue
        end

        for a in 1:5
            if boards[bi][a,:] == bingo
                #println("BINGO! $(boards[bi])")
                return bi # boards[bi]
            end
        end
        for a in 1:5
            if boards[bi][:,a] == bingo
                #println("BINGO! $(boards[bi])")
                return bi #boards[bi]
            end
        end
    end

    return nothing
end

#function find_winner(boardcount, boards)
    losers = falses(boardcount)
    winners = trues(boardcount)

    for ball in bingodata
        draw_ball(ball, boards)
        #println("ball= $ball, boards = $boards")
        while   (bx = check_for_bingo(boards)) != nothing
            println("We found it! board # $bx")
            losers[bx] = true
            #if winners == losers
                @show "found last one: $bx"
                bingo = boards[bx]
                bingo[bingo.==-1] .= 0
                global lastbingo = deepcopy(bingo)
            @show lastbingo, bx
                global lastball = ball
            #return
            #end
            boards[bx][1] = 10_000
        end
    end
#end

#@show boards[2]
@show lastbingo
println("Sum of bingo card: $(sum(lastbingo))")
println("Times ball number: $(sum(lastbingo)*lastball)")


using DifferentialEquations, LinearAlgebra, SparseArrays

const N = 32
const xyd_brusselator = range(0,stop=1,length=N)
brusselator_f(x, y, t) = (((x-0.3)^2 + (y-0.6)^2) <= 0.1^2) * (t >= 1.1) * 5.
limit(a, N) = a == N+1 ? 1 : a == 0 ? N : a
function brusselator_2d_loop(du, u, p, t)
  A, B, alpha, dx = p
  alpha = alpha/dx^2
  @inbounds for I in CartesianIndices((N, N))
    i, j = Tuple(I)
    x, y = xyd_brusselator[I[1]], xyd_brusselator[I[2]]
    ip1, im1, jp1, jm1 = limit(i+1, N), limit(i-1, N), limit(j+1, N), limit(j-1, N)
    du[i,j,1] = alpha*(u[im1,j,1] + u[ip1,j,1] + u[i,jp1,1] + u[i,jm1,1] - 4u[i,j,1]) +
                B + u[i,j,1]^2*u[i,j,2] - (A + 1)*u[i,j,1] + brusselator_f(x, y, t)
    du[i,j,2] = alpha*(u[im1,j,2] + u[ip1,j,2] + u[i,jp1,2] + u[i,jm1,2] - 4u[i,j,2]) +
                A*u[i,j,1] - u[i,j,1]^2*u[i,j,2]
    end
end
p = (3.4, 1., 10., step(xyd_brusselator))

function init_brusselator_2d(xyd)
  N = length(xyd)
  u = zeros(N, N, 2)
  for I in CartesianIndices((N, N))
    x = xyd[I[1]]
    y = xyd[I[2]]
    u[I,1] = 22*(y*(1-y))^(3/2)
    u[I,2] = 27*(x*(1-x))^(3/2)
  end
  u
end
u0 = init_brusselator_2d(xyd_brusselator)
prob_ode_brusselator_2d = ODEProblem(brusselator_2d_loop,u0,(0.,11.5),p)

using SparseArrays
#f = ODEFunction(f, jac_prototype=sparsematrix)

using Symbolics
du0 = copy(u0)
jac_sparsity = Symbolics.jacobian_sparsity((du,u)->brusselator_2d_loop(du,u,p,0.0),du0,u0)

f = ODEFunction(brusselator_2d_loop;jac_prototype=float.(jac_sparsity))
prob_ode_brusselator_2d_sparse = ODEProblem(f,u0,(0.,11.5),p)

using BenchmarkTools: @btime
@btime solve(prob_ode_brusselator_2d,TRBDF2(),save_everystep=false)
@btime solve(prob_ode_brusselator_2d_sparse,TRBDF2(),save_everystep=false)
@btime solve(prob_ode_brusselator_2d,KenCarp47(linsolve=KrylovJL_GMRES()),save_everystep=false)

using IncompleteLU
function incompletelu(W,du,u,p,t,newW,Plprev,Prprev,solverdata)
  if newW === nothing || newW
    Pl = ilu(convert(AbstractMatrix,W), τ = 50.0)
  else
    Pl = Plprev
  end
  Pl,nothing
end
Base.eltype(::IncompleteLU.ILUFactorization{Tv,Ti}) where {Tv,Ti} = Tv

@time solve(prob_ode_brusselator_2d_sparse,KenCarp47(linsolve=KrylovJL_GMRES(),precs=incompletelu,concrete_jac=true),save_everystep=false);

using AlgebraicMultigrid
function algebraicmultigrid(W,du,u,p,t,newW,Plprev,Prprev,solverdata)
  if newW === nothing || newW
    Pl = aspreconditioner(ruge_stuben(convert(AbstractMatrix,W)))
  else
    Pl = Plprev
  end
  Pl,nothing
end
Base.eltype(::AlgebraicMultigrid.Preconditioner) = Float64

@btime solve(prob_ode_brusselator_2d_sparse,KenCarp47(linsolve=KrylovJL_GMRES(),precs=algebraicmultigrid,concrete_jac=true),save_everystep=false);

function algebraicmultigrid2(W,du,u,p,t,newW,Plprev,Prprev,solverdata)
  if newW === nothing || newW
    A = convert(AbstractMatrix,W)
    Pl = AlgebraicMultigrid.aspreconditioner(AlgebraicMultigrid.ruge_stuben(A, presmoother = AlgebraicMultigrid.Jacobi(rand(size(A,1))), postsmoother = AlgebraicMultigrid.Jacobi(rand(size(A,1)))))
  else
    Pl = Plprev
  end
  Pl,nothing
end

@btime solve(prob_ode_brusselator_2d_sparse,KenCarp47(linsolve=KrylovJL_GMRES(),precs=algebraicmultigrid2,concrete_jac=true),save_everystep=false);

#=
using Sundials
@btime solve(prob_ode_brusselator_2d,CVODE_BDF(),save_everystep=false)
@btime solve(prob_ode_brusselator_2d,CVODE_BDF(linear_solver=:LapackDense),save_everystep=false)
@btime solve(prob_ode_brusselator_2d,CVODE_BDF(linear_solver=:GMRES),save_everystep=false)

using ModelingToolkit
prob_ode_brusselator_2d_mtk = ODEProblem(modelingtoolkitize(prob_ode_brusselator_2d_sparse),[],(0.0,11.5),jac=true,sparse=true);
@btime solve(prob_ode_brusselator_2d_mtk,CVODE_BDF(linear_solver=:KLU),save_everystep=false)

using LinearAlgebra
u0 = prob_ode_brusselator_2d_mtk.u0
p  = prob_ode_brusselator_2d_mtk.p
const jaccache = prob_ode_brusselator_2d_mtk.f.jac(u0,p,0.0)
const W = I - 1.0*jaccache

prectmp = ilu(W, τ = 50.0)
const preccache = Ref(prectmp)

function psetupilu(p, t, u, du, jok, jcurPtr, gamma)
  if jok
    prob_ode_brusselator_2d_mtk.f.jac(jaccache,u,p,t)
    jcurPtr[] = true

    # W = I - gamma*J
    @. W = -gamma*jaccache
    idxs = diagind(W)
    @. @view(W[idxs]) = @view(W[idxs]) + 1

    # Build preconditioner on W
    preccache[] = ilu(W, τ = 5.0)
  end
end

function precilu(z,r,p,t,y,fy,gamma,delta,lr)
  ldiv!(z,preccache[],r)
end
@btime solve(prob_ode_brusselator_2d_sparse,CVODE_BDF(linear_solver=:GMRES,prec=precilu,psetup=psetupilu,prec_side=1),save_everystep=false);

prectmp2 = aspreconditioner(ruge_stuben(W, presmoother = AlgebraicMultigrid.Jacobi(rand(size(W,1))), postsmoother = AlgebraicMultigrid.Jacobi(rand(size(W,1)))))
const preccache2 = Ref(prectmp2)
function psetupamg(p, t, u, du, jok, jcurPtr, gamma)
  if jok
    prob_ode_brusselator_2d_mtk.f.jac(jaccache,u,p,t)
    jcurPtr[] = true

    # W = I - gamma*J
    @. W = -gamma*jaccache
    idxs = diagind(W)
    @. @view(W[idxs]) = @view(W[idxs]) + 1

    # Build preconditioner on W
    preccache2[] = aspreconditioner(ruge_stuben(W, presmoother = AlgebraicMultigrid.Jacobi(rand(size(W,1))), postsmoother = AlgebraicMultigrid.Jacobi(rand(size(W,1)))))
  end
end

function precamg(z,r,p,t,y,fy,gamma,delta,lr)
  ldiv!(z,preccache2[],r)
end

@btime solve(prob_ode_brusselator_2d_sparse,CVODE_BDF(linear_solver=:GMRES,prec=precamg,psetup=psetupamg,prec_side=1),save_everystep=false);

=#


using Tullio, Test
using CUDA, KernelAbstractions
CUDA.allowscalar(false)
using Tracker, ForwardDiff
@tullio grad=Base

# matmul
mul(A, B) = @tullio C[i,k] := A[i,j] * B[j,k]

A = rand(3,40); B = rand(40,500);
@test A * B ≈ mul(A, B)
@test cu(A * B) ≈ mul(cu(A), cu(B))

# gradient
ΔA = Tracker.gradient((A,B) -> sum(mul(A, B)), A, B)[1]
@test ΔA ≈ ones(3,500) * B'
@test cu(ΔA) ≈ Tracker.gradient((A,B) -> sum(mul(A, B)), cu(A), cu(B))[1]

# shifts
@tullio D[i,j] := A[i,j+k]  k in 0:10
@test axes(D) == (1:3, 1:30)
@tullio cD[i,j] := cu(A)[i,j+k]  k in 0:10
@test cD isa CuArray
@test cD ≈ cu(D)

#=
# ranges
@tullio E[i,j] := A[i,j+k-1] + (-1:0.5:1)[k]
@test axes(E) == (1:3, 1:36)
@tullio cE[i,j] := cu(A)[i,j+k-1] + (-1:0.5:1)[k]
@test cE isa CuArray
@test cE ≈ cu(E)
=#

# product
@tullio (*) F[j] := A[i,j]
@test F ≈ vec(prod(A, dims=1))
@tullio (*) cF[j] := cu(A)[i,j]
@test cF ≈ cu(F)

# maximum
g(A) = @tullio (max) G[j] := A[i,j]
@test g(A) == vec(maximum(A, dims=1))
A0 = zero(A);
A0[findmax(A, dims=1)[2]] .= 1
@test A0 ≈ Tracker.gradient(sum∘g, A)[1]
@test g(cu(A)) isa CuArray
@test g(cu(A)) ≈ cu(g(A))
@test cu(A0) ≈ Tracker.gradient(sum∘g, cu(A))[1]

# functions
h(A) = @tullio H[j] := exp(A[i,j]) / log(A[i,j])
@test h(cu(A)) isa CuArray
@test h(cu(A)) ≈ cu(h(A))
A1 = Tracker.gradient(sum∘h, A)[1]
@test cu(A1) ≈ Tracker.gradient(sum∘h, cu(A))[1]

#= # broken by https://github.com/mcabbott/Tullio.jl/pull/31
# scalar
@tullio s := cu(A)[i,j]^2
@test s ≈ sum(abs2, A)
@tullio s += cu(B)[i,j]^2
@test s ≈ sum(abs2, A) + sum(abs2, B)
=#


struct BudgetedLongestPathInstance{T}
  graph::AbstractGraph{T}
  rewards::Dict{Edge{T}, Float64}
  weights::Dict{Edge{T}, Int}
  src::T
  dst::T

  budget::Int # weights * solution >= budget
  max_weight::Int

  function BudgetedLongestPathInstance(graph::AbstractGraph{T}, rewards::Dict{Edge{T}, Float64},
                                       weights::Dict{Edge{T}, Int}, src::T, dst::T;
                                       budget::Union{Nothing, Int}=nothing,
                                       max_weight::Union{Nothing, Int}=nothing) where T
    # Error checking.
    if src == dst
      error("Source node is the same as destination node.")
    end

    if ! isnothing(budget) && budget < 0
      error("Budget is present and negative.")
    end

    if any(collect(values(weights)) .< 0)
      error("At least a weight is negative.")
    end

    if length(rewards) != length(weights)
      error("Not the same number of values and weights; these two vectors must have the same size.")
    end

    # Complete the optional parameters.
    if isnothing(max_weight)
      max_weight = maximum(collect(values(weights)))
    end

    if isnothing(budget)
      d = length(values) # Dimension of the problem.
      budget = d * max_weight
    end

    # Return a new instance.
    new{T}(graph, rewards, weights, src, dst, budget, max_weight)
  end
end

graph(i::BudgetedLongestPathInstance{T}) where T = i.graph
rewards(i::BudgetedLongestPathInstance{T}) where T = i.rewards
weights(i::BudgetedLongestPathInstance{T}) where T = i.weights
src(i::BudgetedLongestPathInstance{T}) where T = i.src
dst(i::BudgetedLongestPathInstance{T}) where T = i.dst
budget(i::BudgetedLongestPathInstance{T}) where T = i.budget
max_weight(i::BudgetedLongestPathInstance{T}) where T = i.max_weight

dimension(i::BudgetedLongestPathInstance{T}) where T = ne(graph(i))
reward(i::BudgetedLongestPathInstance{T}, u::T, v::T) where T = rewards(i)[Edge(u, v)]
weight(i::BudgetedLongestPathInstance{T}, u::T, v::T) where T = weights(i)[Edge(u, v)]

struct BudgetedLongestPathSolution{T}
  instance::BudgetedLongestPathInstance{T}
  path::Vector{Edge{T}}
  states::Dict{Tuple{T, Int}, Float64}
  solutions::Dict{Tuple{T, Int}, Vector{Edge{T}}}
end

function paths_all_budgets(s::BudgetedLongestPathSolution{T}, max_budget::Int) where T
  return Dict{Int, Vector{Edge{T}}}(
    budget => s.solutions[s.instance.dst, budget] for budget in 0:max_budget)
end

function paths_all_budgets_as_tuples(s::BudgetedLongestPathSolution{T}, max_budget::Int) where T
  return Dict{Int, Vector{Tuple{T, T}}}(
    budget => [(src(e), dst(e)) for e in s.solutions[s.instance.dst, budget]]
    for budget in 0:max_budget)
end

function budgeted_lp_dp(i::BudgetedLongestPathInstance{T}) where T
  V = Dict{Tuple{T, Int}, Float64}()
  S = Dict{Tuple{T, Int}, Vector{Edge{T}}}()

  # Initialise. For β = 0, this is exactly Bellman-Ford algorithm with costs
  # (instead of rewards). Otherwise, use the same initialisation as
  # Bellman-Ford.
  β0 = lp_dp(ElementaryPathInstance(graph(i), rewards(i), src(i), dst(i)))
  for v in vertices(graph(i))
    S[v, 0] = β0.solutions[v]
    V[v, 0] = length(S[v, 0]) == 0 ? 0 : sum(rewards(i)[e] for e in S[v, 0])
  end

  for β in 1:budget(i)
    for v in vertices(graph(i))
      V[v, β] = -Inf
      S[v, β] = Edge{T}[]
    end

    V[src(i), β] = 0.0
  end

  # Dynamic part.
  for β in 1:budget(i)
    while true
      changed = false

      for v in vertices(graph(i))
        for w in inneighbors(graph(i), v)
          # Compute the remaining part of the budget still to use.
          remaining_budget = max(0, β - weight(i, w, v))

          # If the explored subproblem has no solution, skip it.
          if V[w, remaining_budget] == -Inf
            continue
          end

          # If using the solution to the currently explored subproblem would
          # lead to a cycle, skip it.
          if any(src(e) == v for e in S[w, remaining_budget])
            continue
          end

          # Compute the amount of budget already used by the solution to the currently explored subproblem.
          used_budget = 0
          if length(S[w, remaining_budget]) > 0
            used_budget = sum(weight(i, src(e), dst(e)) for e in S[w, remaining_budget])
          end
          if used_budget < remaining_budget
            continue
          end

          # Compute the maximum: is passing through w advantageous?
          if V[w, remaining_budget] + reward(i, w, v) > V[v, β] && used_budget >= remaining_budget
            changed = true

            V[v, β] = V[w, remaining_budget] + reward(i, w, v)
            S[v, β] = vcat(S[w, remaining_budget], Edge(w, v))
          end
        end
      end

      if ! changed
        break
      end
    end
  end

  return BudgetedLongestPathSolution(i, S[dst(i), budget(i)], V, S)
end


"""
```
relaxation!(ral, xₙ₋₁, Ψₙ₋₁; tol = 1e-10, max_iters = 1000, damping = .5, pnorm = Inf,
            schur_fnct = schur!, autodiff = :central, use_anderson = false, m = 5,
            verbose = :none, kwargs...)
```

solves for the coefficients ``(z, y, \\Psi)`` of a risk-adjusted linearization by the following relaxation algorithm:

1. Initialize guesses for ``(z, y, \\Psi)``
2. Do until convergence

    a) Solve for ``(z, y)`` using the expectational and state transition equations and fixing ``\\Psi``.

    b) Use a QZ decomposition to solve for ``\\Psi`` while fixing ``(z, y)``.

### Types:
- `S1 <: Number`
- `S2 <: Real`
- `S3 <: Real`

### Inputs
- `m::RiskAdjustedLinearization`: object holding functions needed to calculate
    the risk-adjusted linearization
- `xₙ₋₁::AbstractVector{S1}`: initial guess for ``(z, y)``
- `Ψₙ₋₁::AbstractVector{S1}`: initial guess for ``\\Psi``

### Keywords
- `tol::S2`: convergence tolerance of residual norm for relaxation algorithm
- `max_iters::Int`: maximumm number of iterations
- `damping::S2`: guesses are updated as the weighted average
    `xₙ = damping * proposal + (1 - damping) * xₙ₋₁`.
- `pnorm::S3`: norm for residual tolerance
- `schur_fnct::Function`: function for calculating the Schur factorization during QZ decomposition
- `autodiff::Symbol`: specifies whether to use autoamtic differentiation in `nlsolve`
    (and is the same keyword as the `autodiff` keyword for `nlsolve`)
- `use_anderson::Bool`: set to true to apply Anderson acceleration to the
    fixed point iteration of the relaxation algorithm
- `m::Int`: `m` coefficient if using Anderson acceleration
- `sparse_jacobian::Bool = false`: if true, exploit sparsity in the Jacobian in calls to `nlsolve` using SparseDiffTools.jl.
    If `jac_cache` and `sparsity` are `nothing`, then `relaxation!` will attempt to determine the sparsity pattern.
- `sparsity::Union{AbstractArray, Nothing} = nothing`: sparsity pattern for the Jacobian in calls to `nlsolve`
- `colorvec = nothing`: matrix coloring vector for sparse Jacobian in calls to `nlsolve`
- `jac_cache = nothing`: pre-allocated Jacobian cache for calls to `nlsolve` during the numerical algorithms
- `sparsity_detection::Bool = false`: If true, use SparsityDetection.jl to detect sparsity pattern (only relevant if
    both `jac_cache` and `sparsity` are `nothing`). If false,  then the sparsity pattern is
    determined by using finite differences to calculate a Jacobian and assuming any zeros will always be zero.
    Currently, SparsityDetection.jl fails to work.
- `verbose::Symbol`: verbosity of information printed out during solution.
    a) `:low` -> statement when homotopy continuation succeeds
    b) `:high` -> statement when homotopy continuation succeeds and for each successful iteration
"""
function relaxation!(ral::RiskAdjustedLinearization, xₙ₋₁::AbstractVector{S1}, Ψₙ₋₁::AbstractMatrix{S1};
                     tol::S2 = 1e-10, max_iters::Int = 1000, damping::S2 = .5, pnorm::S3 = Inf,
                     schur_fnct::Function = schur!, autodiff::Symbol = :central,
                     use_anderson::Bool = false, m::Int = 5,
                     sparse_jacobian::Bool = false,  sparsity::Union{AbstractArray, Nothing} = nothing,
                     colorvec = nothing, jac_cache = nothing,
                     sparsity_detection::Bool = true, verbose::Symbol = :none,
                     kwargs...) where {S1 <: Number, S2 <: Real, S3 <: Real}
    # Set up
    err = 1.
    nl  = nonlinear_system(ral)
    li  = linearized_system(ral)
    Nzy = ral.Nz + ral.Ny
    AA  = Matrix{Complex{S1}}(undef, Nzy, Nzy) # pre-allocate these matrices to calculate QZ decomp for Ψ
    BB  = similar(AA)

    # Initialize system of equations
    _my_eqn = (F, x, Ψ, 𝒱) -> _relaxation_equations(F, x, ral, Ψ, 𝒱)

    if use_anderson
        # Some aliases/views will be useful
        zₙ    = ral.z
        yₙ    = ral.y
        Ψₙ    = ral.Ψ
        𝒱ₙ₋₁  = nl[:𝒱_sss]
        J𝒱ₙ₋₁ = li[:JV]

        _anderson_f = function _my_anderson(F::AbstractArray{T}, xₙ₋₁::AbstractVector{T}) where {T <: Number}
            zₙ₋₁  = @view xₙ₋₁[1:ral.Nz]
            yₙ₋₁  = @view xₙ₋₁[(ral.Nz + 1):Nzy]
            Ψₙ₋₁  = @view xₙ₋₁[(Nzy + 1):end]
            Ψₙ₋₁  = reshape(Ψₙ₋₁, ral.Ny, ral.Nz)

            # Calculate entropy terms 𝒱ₙ₋₁, J𝒱ₙ₋₁
            update!(nl, zₙ₋₁, yₙ₋₁, Ψₙ₋₁; select = Symbol[:𝒱]) # updates nl.𝒱_sss
            update!(li, zₙ₋₁, yₙ₋₁, Ψₙ₋₁; select = Symbol[:JV]) # updates li.JV

            # Solve state transition and expectational equations for (zₙ, yₙ), taking 𝒱ₙ₋₁ and Ψₙ₋₁ as given
            solve_steadystate!(ral, vcat(zₙ₋₁, yₙ₋₁), _my_eqn, Ψₙ₋₁, 𝒱ₙ₋₁; autodiff = autodiff, # updates ral.z and ral.y
                               sparse_jacobian = sparse_jacobian,
                               sparsity = sparsity, colorvec = colorvec,
                               jac_cache = jac_cache,
                               sparsity_detection = sparsity_detection,
                               verbose = verbose, kwargs...)

            # Update Γ₁, Γ₂, Γ₃, Γ₄, given (zₙ, yₙ)
            update!(li, zₙ, yₙ, Ψₙ₋₁; select = Symbol[:Γ₁, :Γ₂, :Γ₃, :Γ₄]) # updates li.Γᵢ

            # QZ decomposition to get Ψₙ, taking Γ₁, Γ₂, Γ₃, Γ₄, and J𝒱ₙ₋₁ as given
            Ψₙ .= compute_Ψ!(AA, BB, li; schur_fnct = schur_fnct)

            # Update zₙ, yₙ, and Ψₙ; then calculate error for convergence check
            zₙ .= (1 - damping) .* zₙ + damping .* zₙ₋₁
            yₙ .= (1 - damping) .* yₙ + damping .* yₙ₋₁
            Ψₙ .= (1 - damping) .* Ψₙ + damping .* Ψₙ₋₁
            err = norm(vcat(zₙ - zₙ₋₁, yₙ - yₙ₋₁, vec(Ψₙ - Ψₙ₋₁)), pnorm)

            # Calculate residual
            F[1:ral.Nz] = zₙ - zₙ₋₁
            F[(ral.Nz + 1):Nzy] = yₙ - yₙ₋₁
            F[(Nzy + 1):end] = vec(Ψₙ - Ψₙ₋₁)

            return F
        end

        out   = nlsolve(_anderson_f, vcat(xₙ₋₁, vec(Ψₙ₋₁)); m = m, ftol = tol, iterations = max_iters)
        count = out.iterations
        if out.f_converged
            update!(ral, out.zero[1:ral.Nz], out.zero[(ral.Nz + 1):Nzy],
                    reshape(out.zero[(Nzy + 1):end], ral.Ny, ral.Nz); update_cache = false)
        end
    else
        count = 1

        # Some aliases/views will be useful
        zₙ₋₁  = @view xₙ₋₁[1:ral.Nz]
        yₙ₋₁  = @view xₙ₋₁[(ral.Nz + 1):end]
        zₙ    = ral.z
        yₙ    = ral.y
        Ψₙ    = ral.Ψ
        𝒱ₙ₋₁  = nl[:𝒱_sss]
        J𝒱ₙ₋₁ = li[:JV]

        while (err > tol) && (count < max_iters)

            # Calculate entropy terms 𝒱ₙ₋₁, J𝒱ₙ₋₁
            update!(nl, zₙ₋₁, yₙ₋₁, Ψₙ₋₁; select = Symbol[:𝒱]) # updates nl.𝒱_sss
            update!(li, zₙ₋₁, yₙ₋₁, Ψₙ₋₁; select = Symbol[:JV]) # updates li.JV

            # Solve state transition and expectational equations for (zₙ, yₙ), taking 𝒱ₙ₋₁ and Ψₙ₋₁ as given
            solve_steadystate!(ral, xₙ₋₁, _my_eqn, Ψₙ₋₁, 𝒱ₙ₋₁; autodiff = autodiff, # updates ral.z and ral.y
                               sparse_jacobian = sparse_jacobian,
                               sparsity = sparsity, colorvec = colorvec,
                               jac_cache = jac_cache,
                               sparsity_detection = sparsity_detection,
                               verbose = verbose, kwargs...)

            # Update Γ₁, Γ₂, Γ₃, Γ₄, given (zₙ, yₙ)
            update!(li, zₙ, yₙ, Ψₙ₋₁; select = Symbol[:Γ₁, :Γ₂, :Γ₃, :Γ₄]) # updates li.Γᵢ

            # QZ decomposition to get Ψₙ, taking Γ₁, Γ₂, Γ₃, Γ₄, and J𝒱ₙ₋₁ as given
            Ψₙ .= compute_Ψ!(AA, BB, li; schur_fnct = schur_fnct)

            # Update zₙ, yₙ, and Ψₙ; then calculate error for convergence check
            zₙ .= (1 - damping) .* zₙ + damping .* zₙ₋₁
            yₙ .= (1 - damping) .* yₙ + damping .* yₙ₋₁
            Ψₙ .= (1 - damping) .* Ψₙ + damping .* Ψₙ₋₁
            err = norm(vcat(zₙ - zₙ₋₁, yₙ - yₙ₋₁, vec(Ψₙ - Ψₙ₋₁)), pnorm)

            # Update zₙ₋₁, yₙ₋₁, and Ψₙ₋₁ (without reallocating them)
            zₙ₋₁ .= zₙ
            yₙ₋₁ .= yₙ
            Ψₙ₋₁ .= Ψₙ

            if verbose == :high
                println("Iteration $(count): error under norm=$(pnorm) is $(err)")
            end

            count += 1
        end
    end

    if count == max_iters
        throw(RALRelaxationError("Relaxation method to find the risk-adjusted linearization did not converge."))
    else
        update!(ral)

        if verbose == :low
            errvec = steady_state_errors(ral)
            println("Convergence achieved after $(count) iterations! Error under norm = $(pnorm) is " *
                    "$(norm(errvec, pnorm)).")
        elseif verbose == :high
            errvec = steady_state_errors(ral)
            println("")
            println("Convergence achieved after $(count) iterations! Error under norm = $(pnorm) is " *
                    "$(norm(errvec, pnorm)).")
        end

        return ral
    end
end

function solve_steadystate!(m::RiskAdjustedLinearization, x0::AbstractVector{S1},
                            f::Function, Ψ::AbstractMatrix{<: Number}, 𝒱::AbstractVector{<: Number};
                            sparse_jacobian::Bool = false, sparsity::Union{AbstractArray, Nothing} = nothing,
                            colorvec = nothing, jac_cache = nothing,
                            sparsity_detection::Bool = true, autodiff::Symbol = :central,
                            verbose::Symbol = :none, kwargs...) where {S1 <: Real, S2 <: Real}

    # Exploit sparsity?
    if sparse_jacobian
        nlsolve_jacobian!, jac =
            construct_sparse_jacobian_function(m, (F, x) -> f(F, x, Ψ, 𝒱), :relaxation, autodiff;
                                               sparsity = sparsity, colorvec = colorvec,
                                               jac_cache = jac_cache, sparsity_detection = sparsity_detection)
        out = nlsolve(OnceDifferentiable((F, x) -> f(F, x, Ψ, 𝒱), nlsolve_jacobian!, x0, copy(x0), jac), x0; kwargs...)
    else
        out = nlsolve(OnceDifferentiable((F, x) -> f(F, x, Ψ, 𝒱), x0, copy(x0), autodiff,
                                         ForwardDiff.Chunk(ForwardDiff.pickchunksize(min(m.Nz, m.Ny)))), x0; kwargs...)
    end

    if out.f_converged
        m.z .= out.zero[1:m.Nz]
        m.y .= out.zero[(m.Nz + 1):end]
    else
        if verbose == :high
            println(out)
        end
        throw(RALRelaxationError())
    end
end

function _relaxation_equations(F::AbstractArray, x::AbstractArray, m::RiskAdjustedLinearization,
                               Ψ::AbstractMatrix{<: Number}, 𝒱::AbstractVector{<: Number})
    # Unpack
    z = @view x[1:m.Nz]
    y = @view x[(m.Nz + 1):end]

    # Update μ(z, y) and ξ(z, y)
    update!(m.nonlinear, z, y, Ψ; select = Symbol[:μ, :ξ])

    # Calculate residuals
    μ_sss             = get_tmp(m.nonlinear.μ.cache, z, y, (1, 1)) # select the first DiffCache b/c that one
    ξ_sss             = get_tmp(m.nonlinear.ξ.cache, z, y, (1, 1)) # corresponds to autodiffing both z and y
    F[1:m.Nz]         = μ_sss - z
    F[(m.Nz + 1):end] = ξ_sss + m.linearization[:Γ₅] * z + m.linearization[:Γ₆] * y + 𝒱
end

mutable struct RALRelaxationError <: Exception
    msg::String
end
RALRelaxationError() =
    RALRelaxationError("A solution for (z, y), given Ψ and 𝒱, to the state transition and expectational equations could not be found.")
Base.showerror(io::IO, ex::RALRelaxationError) = print(io, ex.msg)


function test_okl_reduction(backend::TestBackend)
    linfo=get_info_string(backend);
    entries=10000;
    nred   =256;

    reducedentries = ceil(Int64,(entries+nred-1)/nred);
    a=Array(Float32,(entries,)); a[:]=1.0;
    ared=Array(Float32,(reducedentries,)); ared[:]=0.0;

    kernelinfo = OCCA.KernelInfo();
    OCCA.adddefine!(kernelinfo,"p_Nred",@sprintf("%d",nred));

    device  = OCCA.Device(linfo);
    o_a     = OCCA.malloc(device, a);
    o_ared  = OCCA.malloc(device, ared);

    reduction = OCCA.buildkernel(device,"reduction.okl","reduction",kernelinfo);

    OCCA.runkernel!(reduction,entries,o_a,o_ared);

    OCCA.memcpy!(ared,o_ared);

    val=sum(ared);

    relerr= abs(sum(a)-val)/sum(a);

    return relerr<1e-5;

end