text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

mutable struct NetworkAccessPolicy <: SwaggerModel

    function NetworkAccessPolicy(;)
        o = new()
        o
    end
end # type NetworkAccessPolicy

const _property_map_NetworkAccessPolicy = Dict{Symbol,Symbol}()
const _property_types_NetworkAccessPolicy = Dict{Symbol,String}()
Base.propertynames(::Type{ NetworkAccessPolicy }) = collect(keys(_property_map_NetworkAccessPolicy))
Swagger.property_type(::Type{ NetworkAccessPolicy }, name::Symbol) = Union{Nothing,eval(Base.Meta.parse(_property_types_NetworkAccessPolicy[name]))}
Swagger.field_name(::Type{ NetworkAccessPolicy }, property_name::Symbol) =  _property_map_NetworkAccessPolicy[property_name]

function check_required(o::NetworkAccessPolicy)
    true
end

function validate_property(::Type{ NetworkAccessPolicy }, name::Symbol, val)
end


text:
module SpectralCRsUtility

include(joinpath("constants", "constants.jl"))
include(joinpath("integrals", "number_density.jl"))
include(joinpath("integrals", "energy.jl"))
include(joinpath("datatypes", "spectra.jl"))
include(joinpath("spectra", "norm_spectrum.jl"))
include(joinpath("io", "io.jl"))
include(joinpath("spectra", "utility.jl"))
include(joinpath("spectra", "number_density_spectrum.jl"))
include(joinpath("spectra", "energy_spectrum.jl"))
include(joinpath("debug", "get_detailled_data.jl"))
include(joinpath("debug", "read_cr_test_data.jl"))
include(joinpath("synchrotron", "synchrotron_emission.jl"))
include(joinpath("synchrotron", "min_momentum.jl"))
include(joinpath("plotting", "plot_spectra.jl"))

# datatypes and helper functions for LMB_SPECTRAL_CRs
export CRShockData,          # datatype to analyse single shocked particle
    readSingleCRShockDataFromOutputFile, # as the name says
    CRMomentumDistributionConfig, # config parameters for momentum distribution function
    CRMomentumDistribution,
    getCRMomentumDistributionFromPartID, # function to get distribution function
    calculateCREnergyInCGS,
    calculateCRNumber,
    get_detailled_shock_data,
    get_detailled_Dpp_data,
    get_detailled_radiative_data,
    get_detailled_adiabatic_data,
    synchrotron_emission,
    smallest_synch_bright_p,
    read_adiabatic_test_data,
    construct_spectrum,
    construct_bin_centers,
    plot

end # module

ext: jl


text:
s = open("input") do file
    readlines(file)
end

nums = map(x->parse(Int64,x), s);

lookingFor = Dict();
secondarySum = Set();

for num in nums
    if(haskey(lookingFor, num))
        println(num * lookingFor[num][1] * lookingFor[num][2]);
        break;
    end

    for secondSum in secondarySum
        if secondSum <= num
            continue
        end

        lookingFor[secondSum-num] = [num, 2020-secondSum];
    end

    push!(secondarySum, 2020-num);
end


text:
using Analytical
using Test

# Running Analytical approximation with default parameters
@testset "Analytical.jl" begin

	adap = Analytical.parameters(N=500,n=661,alLow=0.2)
	cnv = Analytical.binomialDict()
	Analytical.binomOp!(adap,cnv.bn)
	Analytical.setThetaF!(adap)
	Analytical.setPpos!(adap)

	@test adap.pposH == 0.0002641360146660052
	@test adap.thetaF == 1.6433217979109007e-6

end


text:
function rrule(::typeof(VAN.free_energy), β::Real, H::AbstractBlock, sampler, circuit::AbstractBlock, samples)
    free_energy(β, H, sampler, circuit, samples), function (adjy)
        adjθ = grad_θ(H, circuit, samples) * adjy
        nt = map(x->x === nothing ? NoTangent() : x .* adjy, grad_sampler(β, H, sampler, circuit, samples))
        adjsampler = Tangent{typeof(sampler)}(; nt...)
        return (NoTangent(), NoTangent(), NoTangent(), adjsampler, adjθ, NoTangent())
    end
end

@non_differentiable Optimisers.update!(x, x̄)
@non_differentiable dispatch!(x, collection)
@non_differentiable get_input_reg(nbits, samples)
@non_differentiable free_energy_local(β, H, sampler, circuit, samples)


text:
# UnitCommitment.jl: Optimization Package for Security-Constrained Unit Commitment
# Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
# Released under the modified BSD license. See COPYING.md for more details.

abstract type TransmissionFormulation end
abstract type RampingFormulation end
abstract type PiecewiseLinearCostsFormulation end
abstract type StartupCostsFormulation end
abstract type StatusVarsFormulation end
abstract type ProductionVarsFormulation end

struct Formulation
    prod_vars::ProductionVarsFormulation
    pwl_costs::PiecewiseLinearCostsFormulation
    ramping::RampingFormulation
    startup_costs::StartupCostsFormulation
    status_vars::StatusVarsFormulation
    transmission::TransmissionFormulation

    function Formulation(;
        prod_vars::ProductionVarsFormulation = Gar1962.ProdVars(),
        pwl_costs::PiecewiseLinearCostsFormulation = KnuOstWat2018.PwlCosts(),
        ramping::RampingFormulation = MorLatRam2013.Ramping(),
        startup_costs::StartupCostsFormulation = MorLatRam2013.StartupCosts(),
        status_vars::StatusVarsFormulation = Gar1962.StatusVars(),
        transmission::TransmissionFormulation = ShiftFactorsFormulation(),
    )
        return new(
            prod_vars,
            pwl_costs,
            ramping,
            startup_costs,
            status_vars,
            transmission,
        )
    end
end

"""
    struct ShiftFactorsFormulation <: TransmissionFormulation
        isf_cutoff::Float64
        lodf_cutoff::Float64
        precomputed_isf::Union{Nothing,Matrix{Float64}}
        precomputed_lodf::Union{Nothing,Matrix{Float64}}
    end

Transmission formulation based on Injection Shift Factors (ISF) and Line
Outage Distribution Factors (LODF). Constraints are enforced in a lazy way.

Arguments
---------
- `precomputed_isf::Union{Matrix{Float64},Nothing} = nothing`:
    the injection shift factors matrix. If not provided, it will be computed.
- `precomputed_lodf::Union{Matrix{Float64},Nothing} = nothing`:
    the line outage distribution factors matrix. If not provided, it will be
    computed.
- `isf_cutoff::Float64 = 0.005`:
    the cutoff that should be applied to the ISF matrix. Entries with magnitude
    smaller than this value will be set to zero.
- `lodf_cutoff::Float64 = 0.001`:
    the cutoff that should be applied to the LODF matrix. Entries with magnitude
    smaller than this value will be set to zero.
"""
struct ShiftFactorsFormulation <: TransmissionFormulation
    isf_cutoff::Float64
    lodf_cutoff::Float64
    precomputed_isf::Union{Nothing,Matrix{Float64}}
    precomputed_lodf::Union{Nothing,Matrix{Float64}}

    function ShiftFactorsFormulation(;
        isf_cutoff = 0.005,
        lodf_cutoff = 0.001,
        precomputed_isf = nothing,
        precomputed_lodf = nothing,
    )
        return new(isf_cutoff, lodf_cutoff, precomputed_isf, precomputed_lodf)
    end
end


text:
function makePlots()

println("Generating plots")

plotsDir = (pwd()[end-3:end] == "docs") ? "build/plots" : "docs/build/plots"
mkdir(plotsDir)
Plots.gr()

# LQG design
h = 0.1
A = [1 h; 0 1]
B = [0;1]
C = [1 0]
sys = ss(A,B,C,0, h)
Q = Matrix{Float64}(I,2,2)
R = Matrix{Float64}(I,1,1)
L = dlqr(A,B,Q,R) # lqr(sys,Q,R) can also be used

u(x,t)  = -L*x .+ 1.5(t>=2.5)# Form control law (u is a function of t and x), a constant input disturbance is affecting the system from t≧2.5
t=0:h:5
x0 = [1,0]
y, t, x, uout = lsim(sys,u,t,x0=x0)
plot(t,x, lab=["Position" "Velocity"], xlabel="Time [s]")
Plots.savefig(plotsDir*"/lqrplot.svg")

# PID design functions
P = tf(1,[1,1])^4
gangoffourplot(P,tf(1))
Plots.savefig(plotsDir*"/pidgofplot.svg")

ωp = 0.8
kp,ki,C = loopshapingPI(P,ωp,phasemargin=60, doplot=false)
gangoffourplot(P, [tf(1), C])
Plots.savefig(plotsDir*"/pidgofplot2.svg")
nyquistplot([P, P*C])
Plots.savefig(plotsDir*"/pidnyquistplot.svg")

ωp = 2
kp,ki,C60 = loopshapingPI(P,ωp,rl=1,phasemargin=60, doplot=true)
gangoffourplot(P, [tf(1), C60])
Plots.savefig(plotsDir*"/pidgofplot3.svg")
nyquistplot([P, P*C60])
Plots.savefig(plotsDir*"/pidnyquistplot2.svg")

# Advanced pole placement
ζ = 0.2
ω = 1
B = [1]
A   = [1, 2ζ*ω, ω^2]
P  = tf(B,A)
# Control design
ζ0 = 0.7
ω0 = 2
Am = [1, 2ζ0*ω0, ω0^2]
Ao = conv(2Am, [1/2, 1]) # Observer polynomial
AR = [1,0] # Force the controller to contain an integrator

B⁺  = [1] # The process numerator polynomial can be facored as B = B⁺B⁻ where B⁻ contains the zeros we do not want to cancel (non-minimum phase and poorly damped zeros)
B⁻  = [1]
Bm  = conv(B⁺, B⁻) # In this case, keep the entire numerator polynomial of the process

R,S,T = rstc(B⁺,B⁻,A,Bm,Am,Ao,AR) # Calculate the 2-DOF controller polynomials

Gcl = tf(conv(B,T),zpconv(A,R,B,S)) # Form the closed loop polynomial from reference to output

stepplot([P,Gcl]) # Visualize the open and closed loop responses.
Plots.savefig(plotsDir*"/ppstepplot.svg")
gangoffourplot(P, tf(-S,R)) # Plot the gang of four to check that all tranfer functions are OK
Plots.savefig(plotsDir*"/ppgofplot.svg")

P1(s) = exp(-sqrt(s))
f1 = stabregionPID(P1,exp10.(range(-5, stop=1, length=1000))); Plots.savefig(plotsDir*"/stab1.svg")
P2 = s -> 100*(s+6).^2 ./(s.*(s+1).^2 .*(s+50).^2)
f2 = stabregionPID(P2,exp10.(range(-5, stop=2, length=1000))); Plots.savefig(plotsDir*"/stab2.svg")
P3 = tf(1,[1,1])^4
f3 = stabregionPID(P3,exp10.(range(-5, stop=0, length=1000))); Plots.savefig(plotsDir*"/stab3.svg")

# PID plots

P = tf([1.],[1., 1])
ζ = 0.5 # Desired damping
ws = exp10.(range(-1, stop=2, length=8)) # A vector of closed-loop bandwidths
kp = 2*ζ*ws.-1 # Simple pole placement with PI given the closed-loop bandwidth, the poles are placed in a butterworth pattern
ki = ws.^2
pidplots(P,:nyquist,;kps=kp,kis=ki, ω= exp10.(range(-2, stop=2, length=500)))
Plots.savefig(plotsDir*"/pidplotsnyquist1.svg")
pidplots(P,:gof,;kps=kp,kis=ki, ω= exp10.(range(-2, stop=2, length=500)))
Plots.savefig(plotsDir*"/pidplotgof1.svg")

kp = range(-1, stop=1, length=8) # Now try a different strategy, where we have specified a gain crossover frequency of 0.1 rad/s
ki = sqrt.(1 .-kp.^2)./10
pidplots(P,:nyquist,;kps=kp,kis=ki)
Plots.savefig(plotsDir*"/pidplotsnyquist2.svg")
pidplots(P,:gof,;kps=kp,kis=ki)
Plots.savefig(plotsDir*"/pidplotsgof2.svg")

#AR = [1,0]
#B⁺  = [1]
#B⁻  = [1]
#Bm  = conv(B⁺, B⁻)
#R,S,T = rstc(B⁺,B⁻,A,bm,Am,Ao,AR)
#Gcl = tf(conv(B,T),zpconv(A,R,B,S))
#f1 = stepplot([P,Gcl])
#f2 = gangoffourplot(P, tf(-S,R))
#Plots.savefig(f1, "$plotsDir/rstcstepplot.svg")
#Plots.savefig(f2, "$plotsDir/rstcgofplot.svg")
end


text:
using BasisFunctions, DomainSets, StaticArrays, Test, LinearAlgebra, FrameFun

# TODO add AugmentationPlatforms tests
d1 = 0.0..0.5
plat1 = platform(extensionframe(Fourier(10),0.0..0.5)), 10
d2 = (0.0..0.5)^2
plat2 = platform(extensionframe(Fourier(10)^2,d2)), (10,10)
d3 = Disk(.45,SVector(.5,.5))
plat3 = platform(extensionframe(Fourier(10)^2,d3)), (10,10)
plat4 = platform(extensionframe((Fourier(10)→(-1.0..1.0))^2,d3)), (10,10)

wplat1 = WeightedSumPlatform(platform(Fourier(10)), x->1., sqrt), (10,10)
wplat2 = WeightedSumPlatform(platform(Fourier(10)^2), (x,y)->1., (x,y)->sqrt(x^2+y^2)), ((10,10),(10,10))

ap1 = approximationproblem(Fourier(10),d1)
ap2 = approximationproblem(Fourier(10)^2,d2)
ap3 = approximationproblem(Fourier(10)^2,d3)
ap4 = approximationproblem((Fourier(10)→(-1.0..1.0))^2,d3)
wap1 = approximationproblem(wplat1...)
wap2 = approximationproblem(wplat2...)

dict1 = dictionary(ap1)
dict2 = dictionary(ap2)
dict3 = dictionary(ap3)
dict4 = dictionary(ap4)
wdict1 = dictionary(wap1)
wdict2 = dictionary(wap2)

@testset "dictionary" begin
    for plat in (plat1,plat2,plat3,plat4)
        @test dictionary(plat...) isa ExtensionFrameSuper
        @test dualdictionary(plat...,measure(plat[1])) isa ExtensionFrameSuper
    end

    for ap in (ap1,ap2,ap3,ap4)
        @test dictionary(ap) isa ExtensionFrameSuper
    end

    for (plat,ap) in zip((wplat1,wplat2),(wap1,wap2))
        @test dictionary(plat...) isa MultiDict
        @test dictionary(ap) isa MultiDict
        @test dualdictionary(plat...,measure(plat[1])) isa MultiDict
    end
end

@testset "samplingparameter" begin
    @test samplingparameter(ap1) == 38
    @test samplingparameter(ap2) == (27,27)
    @test samplingparameter(ap3) == (18,18)
    @test samplingparameter(ap4) == (36,36)

    @test samplingparameter(wap1) == (40)
    @test samplingparameter(wap2) == (21,21)

    @test samplingparameter(plat1...) == samplingparameter(ap1)
    @test samplingparameter(plat2...) == samplingparameter(ap2)
    @test samplingparameter(plat3...) == samplingparameter(ap3)
    @test samplingparameter(plat4...) == samplingparameter(ap4)
    @test samplingparameter(wplat1...) == samplingparameter(wap1)
    @test samplingparameter(wplat2...) == samplingparameter(wap2)
end

@testset "oversampling_grid" begin
    @test oversampling_grid(ap1) isa IndexSubGrid
    @test oversampling_grid(ap2) isa ProductSubGrid
    @test oversampling_grid(ap3) isa GridArrays.MaskedGrid
    @test oversampling_grid(ap4) isa GridArrays.MaskedGrid

    @test oversampling_grid(wap1) == FourierGrid(40)
    @test oversampling_grid(wap2) == FourierGrid(21)^2

    @test oversampling_grid(plat1...) isa IndexSubGrid
    @test oversampling_grid(plat2...) isa ProductSubGrid
    @test oversampling_grid(plat3...) isa GridArrays.MaskedGrid
    @test oversampling_grid(plat4...) isa GridArrays.MaskedGrid
    @test oversampling_grid(wplat1...) == FourierGrid(40)
    @test oversampling_grid(wplat2...) == FourierGrid(21)^2
end

@testset "sampling_operator" begin
    f1 = exp
    f2 = (x,y) -> exp(x*y)
    op1 = sampling_operator(ap1)
    op2 = sampling_operator(ap2)
    op3 = sampling_operator(ap3)
    op4 = sampling_operator(ap4)

    op5 = sampling_operator(wap1)
    op6 = sampling_operator(wap2)
    # Generic operators so hard to check their structure

    @test sampling_operator(ap1)*f1≈sampling_operator(plat1...)*f1
    @test sampling_operator(ap2)*f2≈sampling_operator(plat2...)*f2
    @test sampling_operator(ap3)*f2≈sampling_operator(plat3...)*f2
    @test sampling_operator(ap4)*f2≈sampling_operator(plat4...)*f2

    @test sampling_operator(wap1)*f1≈sampling_operator(wplat1...)*f1
    @test sampling_operator(wap2)*f2≈sampling_operator(wplat2...)*f2
end

@testset "sampling_grid" begin
    for plat in (plat1,plat2,plat3,plat4,wplat1,wplat2)
        @test sampling_grid(plat...)≈oversampling_grid(plat...)
    end
    for dict in (dict1,dict2,dict3,dict4)
        @test sampling_grid(dict)≈oversampling_grid(dict)
    end
end

@testset "discretemeasure" begin
    for plat in (plat1,plat2,plat3,plat4,wplat1,wplat2)
        @test discretemeasure(plat...)≈discretemeasure(sampling_grid(plat...))
    end
    for dict in (dict1,dict2,dict3,dict4)
        @test discretemeasure(dict)≈discretemeasure(sampling_grid(dict))
    end
end

@testset "measure" begin
    @test measure(ap1)==measure(dict1)
    @test measure(ap2)==measure(dict2)
    @test measure(ap3)==measure(dict3)
    @test measure(ap4)==measure(dict4)
    @test measure(wap1)==measure(wdict1)
    @test measure(wap2)==measure(wdict2)
    @test measure(plat1[1])==measure(dict1)
    @test measure(plat2[1])==measure(dict2)
    @test measure(plat3[1])==measure(dict3)
    @test measure(plat4[1])==measure(dict4)
    @test measure(wplat1[1])==measure(wdict1)
    @test measure(wplat2[1]) ==measure(wdict2)
end

@testset "azdual" begin
    @test azdual(ap1) isa ExtensionFrame
    @test azdual(ap2) isa ExtensionFrameTensor
    @test azdual(ap3) isa ExtensionFrame
    @test azdual(ap4) isa ExtensionFrame
    @test azdual(wap1) isa MultiDict
    @test azdual(wap2) isa MultiDict

    @test azdual(plat1...) isa ExtensionFrame
    @test azdual(plat2...) isa ExtensionFrameTensor
    @test azdual(plat3...) isa ExtensionFrame
    @test azdual(plat4...) isa ExtensionFrame
    @test azdual(wplat1...) isa MultiDict
    @test azdual(wplat2...) isa MultiDict

    @test azdual(dict1) isa ExtensionFrame
    @test azdual(dict2) isa ExtensionFrameTensor
    @test azdual(dict3) isa ExtensionFrame
    @test azdual(dict4) isa ExtensionFrame
end

@testset "discretization" begin
    op = discretization(ap1)
    r = rand(src(op))
    @test op*r≈eval_expansion(dict1, r, sampling_grid(ap1))
    op = discretization(ap2)
    r = rand(src(op))
    @test op*r≈eval_expansion(dict2, r, sampling_grid(ap2))
    op = discretization(ap3)
    r = rand(src(op))
    @test op*r≈eval_expansion(dict3, r, sampling_grid(ap3))
    op = discretization(ap4)
    r = rand(src(op))
    @test op*r≈eval_expansion(dict4, r, sampling_grid(ap4))
    op = discretization(wap1)
    r = rand(src(op))
    @test op*r≈eval_expansion(wdict1, r, sampling_grid(wap1))
    op = discretization(wap2)
    r = rand(src(op))
    @test op*r≈eval_expansion(wdict2, r, sampling_grid(wap2))

    op = discretization(plat1...)
    r = rand(src(op))
    @test op*r≈eval_expansion(dict1, r, sampling_grid(ap1))
    op = discretization(plat2...)
    r = rand(src(op))
    @test op*r≈eval_expansion(dict2, r, sampling_grid(ap2))
    op = discretization(plat3...)
    r = rand(src(op))
    @test op*r≈eval_expansion(dict3, r, sampling_grid(ap3))
    op = discretization(plat4...)
    r = rand(src(op))
    @test op*r≈eval_expansion(dict4, r, sampling_grid(ap4))
    op = discretization(wplat1...)
    r = rand(src(op))
    @test op*r≈eval_expansion(wdict1, r, sampling_grid(wap1))
    op = discretization(wplat2...)
    r = rand(src(op))
    @test op*r≈eval_expansion(wdict2, r, sampling_grid(wap2))

    op = discretization(dict1)
    r = rand(src(op))
    @test op*r≈eval_expansion(dict1, r, sampling_grid(ap1))
    op = discretization(dict2)
    r = rand(src(op))
    @test op*r≈eval_expansion(dict2, r, sampling_grid(ap2))
    op = discretization(dict3)
    r = rand(src(op))
    @test op*r≈eval_expansion(dict3, r, sampling_grid(ap3))
    op = discretization(dict4)
    r = rand(src(op))
    @test op*r≈eval_expansion(dict4, r, sampling_grid(ap4))

end

@testset "dualdiscretization" begin
    op = dualdiscretization(ap1)
    r = rand(src(op))
    @test op*r≈eval_expansion(azdual(ap1), r, sampling_grid(ap1))
    op = dualdiscretization(ap2)
    r = rand(src(op))
    @test op*r≈eval_expansion(azdual(ap2), r, sampling_grid(ap2))
    op = dualdiscretization(ap3)
    r = rand(src(op))
    @test op*r≈eval_expansion(azdual(ap3), r, sampling_grid(ap3))
    op = dualdiscretization(ap4)
    r = rand(src(op))
    @test op*r≈eval_expansion(azdual(ap4), r, sampling_grid(ap4))
    op = dualdiscretization(wap1)
    r = rand(src(op))
    @test op*r≈eval_expansion(azdual(wap1), r, sampling_grid(wap1))
    op = dualdiscretization(wap2)
    r = rand(src(op))
    @test op*r≈eval_expansion(azdual(wap2), r, sampling_grid(wap2))

    op = dualdiscretization(plat1...)
    r = rand(src(op))
    @test op*r≈eval_expansion(azdual(ap1), r, sampling_grid(ap1))
    op = dualdiscretization(plat2...)
    r = rand(src(op))
    @test op*r≈eval_expansion(azdual(ap2), r, sampling_grid(ap2))
    op = dualdiscretization(plat3...)
    r = rand(src(op))
    @test op*r≈eval_expansion(azdual(ap3), r, sampling_grid(ap3))
    op = dualdiscretization(plat4...)
    r = rand(src(op))
    @test op*r≈eval_expansion(azdual(ap4), r, sampling_grid(ap4))
    op = dualdiscretization(wplat1...)
    r = rand(src(op))
    @test op*r≈eval_expansion(azdual(wap1), r, sampling_grid(wap1))
    op = dualdiscretization(wplat2...)
    r = rand(src(op))
    @test op*r≈eval_expansion(azdual(wap2), r, sampling_grid(wap2))

    op = dualdiscretization(dict1)
    r = rand(src(op))
    @test op*r≈eval_expansion(azdual(ap1), r, sampling_grid(ap1))
    op = dualdiscretization(dict2)
    r = rand(src(op))
    @test op*r≈eval_expansion(azdual(ap2), r, sampling_grid(ap2))
    op = dualdiscretization(dict3)
    r = rand(src(op))
    @test op*r≈eval_expansion(azdual(ap3), r, sampling_grid(ap3))
    op = dualdiscretization(dict4)
    r = rand(src(op))
    @test op*r≈eval_expansion(azdual(ap4), r, sampling_grid(ap4))
end

@testset "solvers" begin
    for dict in (dict1,)
        op1 = solver(dict;solverstyle=DirectStyle())
        r = rand(src(op1))
        op2 = solver(dict;solverstyle=IterativeStyle())
        @test norm(op1*r-op2*r) < 1e-4
        op3 = solver(dict;solverstyle=AZStyle())
        @test norm(op1*r-op3*r) < 1e-4
        op4 = solver(dict;solverstyle=AZSmoothStyle())
        @test norm(op1*r-op4*r) < 1e-4
        op4 = solver(dict;solverstyle=TridiagonalProlateStyle())
    end

    for dict in (dict1,dict3,dict4)
        op1 = solver(dict;solverstyle=DirectStyle())
        r = rand(src(op1))
        op2 = solver(dict;solverstyle=AZStyle())
        @test  norm(op1*r-op2*r) < 1e-1
    end

    P = FourierExtensionPlatform(0.0..0.5)
    f = exp
    A,b = full_discretization(f, P, 100)
    M = plungeoperator(P,100)*A
    S = rSVD_solver(M;threshold=1e-4)
    x1 = S*plungeoperator(P,100)*b
    x2 = AZ_Zt(P,100)*(b-A*x1)
    @test norm(A*(x1+x2)-b) <  1e-3
end

@testset "AZ_A, AZ_Z, AZ_Zt, plungeoperator, firstAZstepoperator, plungerank" begin
    for (ap, plat) in ((ap1,plat1),(ap2,plat2),(ap3,plat3),(ap4,plat4),(wap1,wplat1),(wap2,wplat2))
        @test AZ_A(ap) ≈ AZ_A(plat...)
        @test AZ_Zt(ap) ≈ AZ_Zt(plat...)
        @test plungeoperator(ap) ≈ plungeoperator(plat...)
        @test firstAZstepoperator(ap) ≈ firstAZstepoperator(plat...)
        @test abs(plungerank(ap) - plungerank(plat...)) < 3
    end

    for (ap, dict) in ((ap1,dict1),(ap2,dict2),(ap3,dict3),(ap4,dict4),)
        @test AZ_A(ap) ≈ AZ_A(dict)
        @test AZ_Zt(ap) ≈ AZ_Zt(dict)
        @test plungeoperator(ap) ≈ plungeoperator(dict)
        @test firstAZstepoperator(ap) ≈ firstAZstepoperator(dict)
        @test abs(plungerank(ap) - plungerank(dict)) < 3
    end

    for ap in (ap1,ap2,ap3,ap4,wap1,wap2)
        @test norm(plungeoperator(ap) - (I -AZ_A(ap)*AZ_Zt(ap))) < 1e-12
        @test norm(firstAZstepoperator(ap) - (AZ_A(ap) -AZ_A(ap)*AZ_Zt(ap)*AZ_A(ap) )) < 1e-12
    end

    @test plungerank(plat1[1],(100);threshold=1e-6,verbose=false) <= 30
    @test plungerank(plat2[1],(30,30);threshold=1e-6,verbose=false) <= 700
    @test plungerank(plat3[1],(30,30);threshold=1e-6,verbose=false) <= 600
    @test plungerank(plat4[1],(30,30);threshold=1e-6,verbose=false) <= 600
    @test plungerank(wplat1[1],(100,100);threshold=1e-6,verbose=false) <= 30
    @test plungerank(wplat2[1],((15,15),(15,15));threshold=1e-6,verbose=false)  <= 400
end

@testset "approximate" begin
    P = FourierExtensionPlatform(0.0..0.5)
    f = exp
    A,b = full_discretization(f, P, 100)
    M = plungeoperator(P,100)*A
    S = rSVD_solver(M;threshold=1e-4)
    x1 = S*plungeoperator(P,100)*b
    x2 = AZ_Zt(P,100)*(b-A*x1)
    @test norm(A*(x1+x2)-b) <  1e-3

    F = approximate(f, P, 100; REG=rSVD_solver)[1]
    x = .2; @test abs(F(x)-f(x)) < 1e-3
    P = FourierExtensionPlatform(0.0..0.5)
    f = exp
    A,b = full_discretization(f, P, 100)
    M = plungeoperator(P,100)*A
    S = pQR_solver(M;threshold=1e-4)
    x1 = S*plungeoperator(P,100)*b
    x2 = AZ_Zt(P,100)*(b-A*x1)
    @test norm(A*(x1+x2)-b) < 1e-3
    F = approximate(f, P, 100;REG=pQR_solver)[1]
    x = .2; @test abs(F(x)-f(x)) < 1e-3
    P = FourierExtensionPlatform(0.0..0.5)
    f = exp
    A,b = full_discretization(f, P, 100)
    M = plungeoperator(P,100)*A
    S = pSVD_solver(M;threshold=1e-4)
    x1 = S*plungeoperator(P,100)*b
    x2 = AZ_Zt(P,100)*(b-A*x1)
    @test norm(A*(x1+x2)-b) < 1e-3

    F = approximate(f, P, 100;REG=pSVD_solver)[1]
    x = .2; @test abs(F(x)-f(x)) < 1e-3
end


text:
export
    DIRCOLSolver

abstract type DirectSolver{T} <: AbstractSolver{T} end
abstract type DirectSolverOptions{T} <: AbstractSolverOptions{T} end

abstract type QuadratureRule end
abstract type HermiteSimpson <: QuadratureRule end
abstract type Midpoint <: QuadratureRule end

include("primals.jl")

"$(TYPEDEF) Solver options for the Projected Newton solver"
@with_kw mutable struct ProjectedNewtonSolverOptions{T} <: DirectSolverOptions{T}
    "Print output to console"
    verbose::Bool = true

    "Number of Newton steps"
    n_steps::Int = 1

    "Solve type, feasibile or optimal"
    solve_type::Symbol = :feasible

    "Tolerance for checking active inequality constraints. Positive values move the boundary further into the feasible region (i.e. negative)"
    active_set_tolerance = 1e-3

    "Tolerance for constraint feasibility during projection"
    feasibility_tolerance = 1e-6

end

"""
$(TYPEDEF)
Projected Newton Solver
Direct method developed by the REx Lab at Stanford University

Achieves machine-level constraint satisfaction by projecting onto the feasible subspace.
    It can also take a full Newton step by solving the KKT system.

This solver is to be used exlusively for solutions that are close to the optimal solution.
    It is intended to be used as a "solution polishing" method for augmented Lagrangian methods.
"""
mutable struct ProjectedNewtonSolver{T} <: DirectSolver{T}
    opts::ProjectedNewtonSolverOptions{T}
    stats::Dict{Symbol,Any}
    V::PrimalDual{T}
    H::SparseMatrixCSC{T,Int}      # Cost Hessian
    g::Vector{T}                   # Cost gradient
    Y::PseudoBlockArray{T,2,SparseArrays.SparseMatrixCSC{T,Int64},BlockArrays.BlockSizes{2,NTuple{2,Vector{Int64}}}}      # Constraint Jacobian
    y::PseudoBlockArray{T,1,Vector{T},BlockArrays.BlockSizes{1,Tuple{Array{Int64,1}}}}                   # Constraint Violations

    fVal::Vector{SubArray{T,1,BlockArrays.PseudoBlockArray{T,1,Vector{T},BlockArrays.BlockSizes{1,Tuple{Array{Int64,1}}}},Tuple{BlockArrays.BlockSlice{BlockArrays.Block{1,Int64}}},false}}
    # ∇F::Vector{PartedArray{T,2,SubArray{T,2,SparseMatrixCSC{T,Int},Tuple{UnitRange{Int},UnitRange{Int}},false}, P}} where P
    ∇F::PartedMatTrajectory{T}
    C::Vector{PartedArrays.PartedArray{T,1,SubArray{T,1,BlockArrays.PseudoBlockArray{T,1,Array{T,1},BlockArrays.BlockSizes{1,Tuple{Array{Int64,1}}}},Tuple{BlockArrays.BlockSlice{BlockArrays.Block{1,Int64}}},false},P} where P}
    ∇C::Vector{PartedArrays.PartedArray{T,2,SubArray{T,2,BlockArrays.PseudoBlockArray{T,2,SparseArrays.SparseMatrixCSC{T,Int64},BlockArrays.BlockSizes{2,NTuple{2,Array{Int64,1}}}},NTuple{2,BlockArrays.BlockSlice{BlockArrays.Block{1,Int64}}},false},P} where P}
    a::PartedArrays.PartedArray{Bool,1,PseudoBlockArray{Bool,1,Array{Bool,1},BlockArrays.BlockSizes{1,Tuple{Array{Int64,1}}}},NamedTuple{(:primals, :duals),NTuple{2,UnitRange{Int64}}}}
    active_set::Vector{SubArray{Bool,1,PseudoBlockArray{Bool,1,Array{Bool,1},BlockArrays.BlockSizes{1,Tuple{Array{Int64,1}}}},Tuple{BlockArrays.BlockSlice{Block{1,Int64}}},false}}
    parts::NamedTuple{(:primals,:duals),NTuple{2,UnitRange{Int}}}
end

function AbstractSolver(prob::Problem{T,D}, opts::ProjectedNewtonSolverOptions{T}) where {T<:AbstractFloat,D<:DynamicsType}
    n,m,N = size(prob)
    X_ = [zeros(T,n) for k = 1:N-1] # midpoints

    stats = Dict{Symbol,Any}(:timer=>TimerOutput())

    NN = N*n + (N-1)*m
    p = num_constraints(prob)
    pcum = insert!(cumsum(p),1,0)
    P = sum(p) + N*n

    V = PrimalDual(prob)

    part_f = create_partition2(prob.model)
    constraints = prob.constraints
    part_a = (primals=1:NN, duals=NN+1:NN+P) #, ν=NN .+ (1:N*n), λ=NN + N*n .+ (1:sum(p)))

    # Block Array partitions
    y_part = ones(Int,2,N-1)*n
    y_part[2,:] = p[1:end-1]
    y_part = vec(y_part)
    insert!(y_part,1,n)
    push!(y_part, p[N])
    c_blocks = push!(collect(3:2:length(y_part)),length(y_part))

    z_part = repeat([n+m],N-1)
    push!(z_part, n)
    d_blocks = insert!(collect(2:2:length(y_part)-1),1,1)

    # Build Blocks
    H = spzeros(NN,NN)
    g = zeros(NN)
    Y = PseudoBlockArray(spzeros(P,NN), y_part, z_part)
    y = PseudoBlockArray(zeros(sum(y_part)),y_part)
    a = PseudoBlockArray(ones(Bool,NN+P), [NN; y_part])

    # Build views
    fVal = [view(y,Block(i)) for i in d_blocks]
    ∇F = [PartedMatrix(zeros(n,n+m+1), part_f) for k = 1:N]

    C = [PartedArray(view(y, Block(c_blocks[k])),
         create_partition(constraints[k], k==N ? :terminal : :stage)) for k = 1:N]
    ∇C = [PartedArray(view(Y, Block(c_blocks[k],k)),
         create_partition2(constraints[k], n,m, k==N ? :terminal : :stage)) for k = 1:N]

    active_set = [view(a,Block(i+1)) for i in c_blocks]
    a = PartedVector(a, part_a)

    solver = ProjectedNewtonSolver{T}(opts, stats, V, H, g, Y, y, fVal, ∇F, C, ∇C, a, active_set, part_a)
    reset!(solver)
    return solver
end

ProjectedNewtonSolver(prob::Problem,
    opts::ProjectedNewtonSolverOptions=ProjectedNewtonSolverOptions{Float64}()) = AbstractSolver(prob, opts)

function reset!(solver::ProjectedNewtonSolver{T}) where T
    solver.stats[:iterations] = 0
    solver.stats[:c_max] = T[]
    solver.stats[:cost] = T[]
end

function num_active_constraints(solver::ProjectedNewtonSolver)
    sum(solver.a.duals)
end

"$(TYPEDEF) Solver options for the Direct Collocation solver. Most options are passed to the NLP through the `opts` dictionary"
@with_kw mutable struct DIRCOLSolverOptions{T} <: DirectSolverOptions{T}
    "NLP Solver to use. See MathOptInterface for available NLP solvers"
    nlp::MathOptInterface.AbstractOptimizer = Ipopt.Optimizer()

    "Options dictionary for the nlp solver"
    opts::Dict{Symbol,Any} = Dict{Symbol,Any}()

    "Print output to console"
    verbose::Bool = true

    "Feasibility tolerance"
    feasibility_tolerance::T = -1.0
end

"""
$(TYPEDEF)
Direct Collocation Solver.
Uses a commerical NLP solver to solve the Trajectory Optimization problem.
Uses the MathOptInterface to interface with the NLP.
"""
mutable struct DIRCOLSolver{T,Q} <: DirectSolver{T}
    opts::DIRCOLSolverOptions{T}
    stats::Dict{Symbol,Any}
    Z::Primals{T}
    X_::VectorTrajectory{T}
    ∇F::PartedMatTrajectory{T}
    C::PartedVecTrajectory{T}
    ∇C::PartedMatTrajectory{T}
    fVal::VectorTrajectory{T}
    p::Vector{Int}
end

DIRCOLSolver(prob::Problem, opts::DIRCOLSolverOptions=DIRCOLSolverOptions{Float64}(),
    Z::Primals{T}=Primals(prob,true)) where {T,Q} = AbstractSolver(prob, opts, Z)

type(::Primals{T}) where T = T

function AbstractSolver(prob::Problem, opts::DIRCOLSolverOptions, Z::Primals{T}=Primals(prob, true)) where T
    n,m,N = size(prob)
    X_ = [zeros(T,n) for k = 1:N-1] # midpoints

    part_f = create_partition2(prob.model)
    constraints = prob.constraints
    p = num_stage_constraints(constraints)
    c_stage = [stage(constraints[k]) for k = 1:N-1]
    c_part = [create_partition(c_stage[k]) for k = 1:N-1]
    c_part2 = [create_partition2(c_stage[k],n,m) for k = 1:N-1]

    # Create Trajectories
    ∇F         = [PartedMatrix(zeros(T,n,length(prob.model)),part_f)         for k = 1:N]
    C          = [PartedVector(T,constraints[k],:stage)     for k = 1:N-1]
    ∇C         = [PartedMatrix(T,constraints[k],n,m,:stage) for k = 1:N-1]
    C          = [C...,  PartedVector(T,constraints[N],:terminal)]
    ∇C         = [∇C..., PartedMatrix(T,constraints[N],n,m,:terminal)]

    c_term = terminal(constraints[N])
    p_N = num_constraints(c_term)
    fVal = [zeros(T,n) for k = 1:N]
    p = num_constraints(prob)

    solver = DIRCOLSolver{T,HermiteSimpson}(opts, Dict{Symbol,Any}(), Z, X_, ∇F, C, ∇C, fVal, p)
    reset!(solver)
    return solver
end

function reset!(solver::DIRCOLSolver{T,Q}) where {T, Q<:QuadratureRule}
    # state = Dict{Symbol,Any}(:iterations=>0, :c_max=>T[], :cost=>T[])
    solver.stats[:iterations] = 0
    solver.stats[:c_max] = T[]
    solver.stats[:cost] = T[]
    solver.stats[:iter_time] = T[]
end

max_issues_repo_name: MehrdadZareian/GMARA

max_forks_repo_name: MehrdadZareian/GMARA


text:
"""
    proj(u, v)

Compute the vector projection of vector `v` onto vector `u`.
"""
function proj(u::AbstractVector, v::AbstractVector)
    (u ⋅ v)/(u ⋅ u)*u
end

@generated function insertcolumn(s::StaticVector{N, T}) where {N, T}
    S = (N, N+1)
    exprs = Array{Expr}(undef, S)
    itr = [1:n for n = S]
    for k = Base.product(itr...)
        exprs[k...] = :($(k[2]) > 1 ? zero($T) : s[$(k[1])])
    end

    return quote
        Base.@_inline_meta
        @inbounds elements = tuple($(exprs...))
        @inbounds return SMatrix{N, N+1, T}(elements)
    end
end

@generated function insertcolumn(simplex::SMatrix{N, M, T}, s::StaticVector{N, T}, idx::Vararg{Int, 1}) where {N, M, T}
    S = (N, M)
    exprs = Array{Expr}(undef, S)
    itr = [1:n for n = S]
    for k = Base.product(itr...)
        exprs[k...] = :($(k[2]) > idx[1] ? zero($T) : ($(k[2]) == idx[1] ? s[$(k[1])] : simplex[$(k...)]))
    end

    return quote
        Base.@_inline_meta
        @inbounds elements = tuple($(exprs...))
        @inbounds return SMatrix{N, M, T}(elements)
    end
end

@generated function pickcolumns(simplex::SMatrix{N, M, T}, idx::Vararg{Int, J}) where {N, M, T, J}
    S = (N, M)
    exprs = Array{Expr}(undef, S)
    itr = [1:n for n = S]
    for k = Base.product(itr...)
        exprs[k...] = :($(k[2]) > $J ? zero($T) : simplex[$(k[1]), idx[$(k[2])]])
    end

    return quote
        Base.@_inline_meta
        @inbounds elements = tuple($(exprs...))
        @inbounds return SMatrix{N, M, T}(elements)
    end
end

@generated function reducecolumns(simplex::SMatrix{N, M, T}, ::Vararg{Any, J}) where {N, M, T, J}
    S = (N, J)
    exprs = Array{Expr}(undef, S)
    itr = [1:n for n = S]
    for k = Base.product(itr...)
        exprs[k...] = :(simplex[$(k...)])
    end

    return quote
        Base.@_inline_meta
        @inbounds elements = tuple($(exprs...))
        @inbounds return SMatrix{N, J, T}(elements)
    end
end

"""
    findsimplex(simplex)

Compute the new simplices from a pair of given simplices.
Return the new search direction. Return a collision flag
if the origin was enclosed by the Minkowski simplex.
"""
function findsimplex(psimplex::SMatrix, qsimplex::SMatrix, sz::Int)
    if sz == 2
        findline(psimplex, qsimplex)
    elseif sz == 3
        findtriangle(psimplex, qsimplex)
    elseif sz == 4
        findtetrahedron(psimplex, qsimplex)
    end
end

function findline(psimplex::SMatrix{N}, qsimplex::SMatrix{N}) where {N}
    simplex = psimplex - qsimplex
    AB = simplex[:, 1] - simplex[:, 2]
    AO = -simplex[:, 2]
    T = eltype(psimplex)
    ntol = eps(T)*oneunit(T)
    if AB ⋅ AO > zero(T)^2
        dir = AO - proj(AB, AO)
        collision = norm(dir) ≤ ntol
        return psimplex, qsimplex, dir, collision, 2
    else
        dir = AO
        psimplex = pickcolumns(psimplex, 2)
        qsimplex = pickcolumns(qsimplex, 2)
        collision = norm(dir) ≤ ntol
        return psimplex, qsimplex, dir, collision, 1
    end
end

function findtriangle(psimplex::SMatrix{N}, qsimplex::SMatrix{N}) where {N}
    simplex = psimplex - qsimplex
    AB = simplex[:, 2] - simplex[:, 3]
    AC = simplex[:, 1] - simplex[:, 3]
    BC = simplex[:, 1] - simplex[:, 2]
    AO = -simplex[:, 3]
    T = eltype(psimplex)
    ntol = eps(T)*oneunit(T)
    if (AC ⋅ AB * BC - AC ⋅ BC * AB) ⋅ AO > zero(T)^4
        if AC ⋅ AO > zero(T)^2
            psimplex = pickcolumns(psimplex, 1, 3)
            qsimplex = pickcolumns(qsimplex, 1, 3)
            dir = AO - proj(AC, AO)
            collision = norm(dir) ≤ ntol
            return psimplex, qsimplex, dir, collision, 2
        else
            psimplex = pickcolumns(psimplex, 2, 3)
            qsimplex = pickcolumns(qsimplex, 2, 3)
            return findline(psimplex, qsimplex)
        end
    elseif (AB ⋅ BC * AB - AB ⋅ AB * BC) ⋅ AO > zero(T)^3
        psimplex = pickcolumns(psimplex, 2, 3)
        qsimplex = pickcolumns(qsimplex, 2, 3)
        return findline(psimplex, qsimplex)
    else
        if norm(AO) ≤ ntol
            return psimplex, qsimplex, AO, true, 3
        elseif norm(AC - proj(AB, AC)) ≤ ntol
            psimplex = pickcolumns(psimplex, 2, 3)
            qsimplex = pickcolumns(qsimplex, 2, 3)
            return findline(psimplex, qsimplex)
        elseif N == 2
            return psimplex, qsimplex, AO, true, 3
        else
            ABC = AB × BC
            dir = proj(ABC, AO)
            if ABC ⋅ AO > zero(T)^3
                psimplex = pickcolumns(psimplex, 2, 1, 3)
                qsimplex = pickcolumns(qsimplex, 2, 1, 3)
                collision = norm(dir) ≤ ntol
                return psimplex, qsimplex, dir, collision, 3
            else
                collision = norm(dir) ≤ ntol
                return psimplex, qsimplex, dir, collision, 3
            end
        end
    end
end

function findtetrahedron(psimplex::SMatrix{N}, qsimplex::SMatrix{N}) where {N}
    simplex = psimplex - qsimplex
    AB = simplex[:, 3] - simplex[:, 4]
    AC = simplex[:, 2] - simplex[:, 4]
    AD = simplex[:, 1] - simplex[:, 4]
    AO = -simplex[:, 4]
    AB_O = sign((AC × AD) ⋅ AO) == sign((AC × AD) ⋅ AB)
    AC_O = sign((AD × AB) ⋅ AO) == sign((AD × AB) ⋅ AC)
    AD_O = sign((AB × AC) ⋅ AO) == sign((AB × AC) ⋅ AD)
    if (AB_O && AC_O && AD_O)
        return psimplex, qsimplex, AO, true, 4
    elseif (!AB_O)
        psimplex = pickcolumns(psimplex, 1, 2, 4)
        qsimplex = pickcolumns(qsimplex, 1, 2, 4)
        return findtriangle(psimplex, qsimplex)
    elseif (!AC_O)
        psimplex = pickcolumns(psimplex, 3, 1, 4)
        qsimplex = pickcolumns(qsimplex, 3, 1, 4)
        return findtriangle(psimplex, qsimplex)
    else
        psimplex = pickcolumns(psimplex, 2, 3, 4)
        qsimplex = pickcolumns(qsimplex, 2, 3, 4)
        return findtriangle(psimplex, qsimplex)
    end
end

"""
    nearestfromsimplex(psimplex, qsimplex, dir2origin)

Compute the nearest points between two simplexes given the
direction to in origin in the Minkowski difference space
"""
function nearestfromsimplex(psimplex::SMatrix{N, M}, qsimplex::SMatrix{N, M}, dir2origin::SVector{N}, sz::Int) where {N, M}
    if sz == 1
        return psimplex[:, 1], qsimplex[:, 1]
    elseif sz == 2
        λ = linecombination(psimplex - qsimplex, dir2origin)
        pclose = reducecolumns(psimplex, λ...)*SVector{2}(λ)
        qclose = reducecolumns(qsimplex, λ...)*SVector{2}(λ)
        return pclose, qclose
    elseif sz == 3
        λ = trianglecombination(psimplex - qsimplex, dir2origin)
        pclose = reducecolumns(psimplex, λ...)*SVector{3}(λ)
        qclose = reducecolumns(qsimplex, λ...)*SVector{3}(λ)
        return pclose, qclose
    end
end

function linecombination(simplex::SMatrix{N}, vec::SVector{N}) where {N}
    AV = vec - simplex[:, 2]
    AB = simplex[:, 1] - simplex[:, 2]
    λ = (AB ⋅ AV)/(AB ⋅ AB)
    if λ < 0
        return 0.0, 1.0 - λ
    else
        return λ, 1.0 - λ
    end
end

function trianglecombination(simplex::SMatrix{N}, vec::SVector{N}) where {N}
    AO = -simplex[:, 3]
    AV = vec - simplex[:, 3]
    AB = simplex[:, 2] - simplex[:, 3]
    AC = simplex[:, 1] - simplex[:, 3]
    BC = simplex[:, 1] - simplex[:, 2]
    T = eltype(simplex)

    if (AC ⋅ AB * BC - AC ⋅ BC * AB) ⋅ AV > zero(T)^3
        if AC ⋅ AV > zero(T)^2
            dir = AO - proj(AC, AV)
            idx = SMatrix{3, 2}(1, 0, 0, 0, 0, 1)
            simplex = pickcolumns(simplex, 1, 3)
            λAC1, λAC2 = linecombination(simplex, -dir)
            return λAC1, 0.0, λAC2
        else
            dir = AO - proj(AB, AV)
            simplex = pickcolumns(simplex, 2, 3)
            λAB1, λAB2 = linecombination(simplex, -dir)
            return 0.0, λAB1, λAB2
        end
    elseif (AB ⋅ BC * AB - AB ⋅ AB * BC) ⋅ AV > zero(T)^3
        dir = AO - proj(AB, AV)
        simplex = pickcolumns(simplex, 2, 3)
        λAB1, λAB2 = linecombination(simplex, -dir)
        return 0.0, λAB1, λAB2
    else
        sABC = [AC AB]
        λ = (sABC' * sABC) \ (sABC' * AV)
        return λ[1], λ[2], 1-sum(λ)
    end
end


text:
module MatrixColorings
if VERSION < v"0.4.0-dev"
    using Docile
end

include("base-fixes.jl")
using Ragged, ArrayViews

import Base: start, next, done, length

export color_jac, recover_jac!, recover_jac, color2seedmatrix, add_delta!, remove_delta!

# this follows "ColPack: Software for Graph Coloring and Related
# Problems in Scientific Computing" GEBREMEDHIN & et al

typealias MatrixLike Any # maybe have a trait here one day
typealias VectorLike Any # maybe have a trait here one day

########
## Misc helper functions
########

# Allocates all vectors inside a Vector{Vector{T}} and element type
# eltype(eltype(vec)).
allocate!{T}(vec::Vector{T}, n::Int) =
    (vec[:] = Vector{T}[ Array(eltype(T), n) for s=1:length(vec)])

#########
## Graphs
#########
# Reinventing the wheel!

# Compact storage for an undirected graph.
# Neighbors of vertex i are adjlist[i]
abstract AGraph
# assumes fields
# v::UnitRange{Int}  # vertices
# adjlist::RaggedArray{Int,Int} # adjacency list, i.e. adjlist[v]
                                # gives connected vertices.
length(g::AGraph) = length(g.v)

immutable Graph <: AGraph
    v::UnitRange{Int}  # vertices
    adjlist::RaggedArray{Int,Int} # adjacency list
    work::BitArray{1} # work array for various algos
    function Graph(v,a)
        @assert size(a,2)==length(v)
        length(v)>0 && @assert v[1]==1
        new(v,a,falses(length(v)))
    end
end
immutable BiGraph <: AGraph
    v::UnitRange{Int}  # vertices
    adjlist::RaggedArray{Int,Int} # adjacency list
    work::BitArray{1} # work array for various algos
    v1::UnitRange{Int} # first partition:
    v2::UnitRange{Int} # second partition
    function BiGraph(v,a,v1,v2)
        @assert size(a,2)==length(v)
        if length(v)>0
            @assert v[1]==1
            @assert v[1]==1
            @assert v1[1]==1
            @assert v1[end]+1==v2[1]
            @assert v2[end]==v[end]
        else
            @assert length(v1)==length(v2)==0
        end
        new(v,a,falses(length(v)),v1,v2)
    end
end

@doc "Iterates over all direct neighbor vertices." ->
immutable Neigh1{G<:AGraph}
    g::G
    v::Int
end
start(n1::Neigh1) = 1 # state is row-index of adjlist
next(n1::Neigh1, state) = (n1.g.adjlist[state,n1.v], state+1)
done(n1::Neigh1, state) = Ragged.collength(n1.g.adjlist,n1.v)+1==state

@doc "Iterates over all 2nd neighbor vertices." ->
immutable Neigh2{G<:AGraph}
    g::G
    v::Int
end
# The problem here is that we don't know we've reached the end until
# we've reached it, thus use the dict.jl skip_deleted trick.

# Type to hold the state:
type Neigh2State{G<:AGraph}
    n1::Neigh1{G}  # the current Neigh1 to iterate over
    n1_state::Int  # the state of n1
    n2_state::Int  # the state encoding which n1 we're on
    item::Int      # the actual item
end
# Function which does all the logic
function _get_next(n2::Neigh2, s::Neigh2State)
    # Iterate over n1, discarding an item if it has been returned
    # already or if it is n2.v.
    visited = n2.g.work
    num_n1 = Ragged.collength(n2.g.adjlist, n2.v)
    @inbounds while visited[s.item] # if item has been visited, try next
        while done(s.n1, s.n1_state) # if n1 is done go onto next first neighbor
            s.n2_state += 1 # onto the next neigh
            if s.n2_state <= num_n1
                s.n1 = Neigh1(n2.g, n2.g.adjlist[s.n2_state, n2.v])
                s.n1_state = start(s.n1)
            else
                s.n2_state = -1  # this indicates completely done,
                                 # i.e. no next item.
                return s
            end
        end
        s.item, s.n1_state = next(s.n1, s.n1_state)
    end
    visited[s.item] = true
    return s
end

# State is (Neigh1, state-of-Neigh1, i in g.adjlist[i,v])
function start(n2::Neigh2)
    if length(n2.g)==0 || Ragged.collength(n2.g.adjlist,n2.v)==0
        # no 1st neighbors
        n1 = Neigh1(n2.g, -1)
        n1_state = -1
        n2_state = -1
        return Neigh2State(n1, n1_state, n2_state, 0)
    else
        fill!(n2.g.work, false) # initialize work array to keep track
                                # of visited vertices.
        n2.g.work[n2.v] = true  #
        item = n2.v
        n1 = Neigh1(n2.g, n2.g.adjlist[1,n2.v])
        n1_state, n2_state = start(n1), 1
        return _get_next(n2, Neigh2State(n1, n1_state, n2_state, item))
    end
end
next(n2::Neigh2, state) = (state.item, _get_next(n2, state))
done(n2::Neigh2, state) = state.n2_state==-1

######
# Colors
######
immutable Color{T}
    c::Vector{T}
    dim::Int
end

############
## Jacobians "partial distance-2 coloring"
############

macro _matrix2bipartie_graph()
    # set-up for function of same name
    esc(quote
    nr,nc = size(A)
    nv = nr+nc
    v = 1:nv
    v1 = 1:nr
    v2 = nr+1:nv
    adj = Array(Vector{Int}, nv)
    allocate!(adj, 0)
    end)
end

@doc """

\"The structure of a Jacobian matrix A is represented by the bipartite
graph Gb(A) = (V₁, V₂, E), where the vertex sets V₁ and V₂ represent
the rows and columns of A, respectively, and each nonzero matrix entry
A_ ij is represented by the edge (r_i , c_j) in E.\"

""" ->
function matrix2bipartie_graph(A::MatrixLike)
    @_matrix2bipartie_graph
    for j=1:nc,i=1:nr
        if A[i,j]!=0
            vv1, vv2 = v1[i], v2[j]
            push!(adj[vv1], vv2)
            push!(adj[vv2], vv1)
        end
    end
    return BiGraph(v, RaggedArray(adj), v1, v2)
end
function matrix2bipartie_graph(A::SparseMatrixCSC)
    # This will include all entries of the sparse matrix, even if they
    # happen to be zero.
    @_matrix2bipartie_graph
    for j=1:nc
        for nz in nzrange(A,j)
            i = rowvals(A)[nz]
            vv1, vv2 = v1[i], v2[j]
            push!(adj[vv1], vv2)
            push!(adj[vv2], vv1)
        end
    end
    return BiGraph(v, RaggedArray(adj), v1, v2)
end

@doc """
I think sorting can improve the coloring, but no idea how to
sort."""->
function sort_vertices!(g::AGraph)
    nothing
    # probably need to return a permeation to reverse the sort.
end

@doc """
\"In a bipartite graph G b = (V 1 , V 2 , E), a partial distance-2
coloring on the vertex set V 1 (or V 2 ) is an assignment of colors to
the vertices in V 1 (or V 2 ) such that a pair of vertices connected
by a path of length exactly two edges receives different colors. The
term partial, which is sometimes omitted when the context is clear, is
used to emphasize that the other vertex set remains uncolored.\"

Some small testing suggest that the performance of this function is
within a factor 10 of COLPACK.
""" ->
function partial_dist2_coloring(g::BiGraph, dim=2)
    # dim==2 column coloring
    vv = dim==1 ? g.v1 : g.v2
#    vv = ordering(g) # TODO, see $5.  Color-based ordering seems best, $5.6
    nvv = length(vv)

    # This is the output.  Possible colors range from 1:nnv, indexed
    # by index into vv.
    color = zeros(Int, nvv)
    # A work array holding vertex numbers, indexed by
    # color. Initialize with a non-vertex number:
    forbiddenColors = zeros(Int,nvv)
    # - forbiddenColors[c] == v indicates that the color c is
    #   impermissible for the vertex v
    # - forbiddenColors[c] != v indicates that c is a candidate color
    #   for the vertex v, regardless of the actual value in
    #   forbiddenColors[c].
    for (i,v) in enumerate(vv)
        for w in Neigh2(g, v) # replace here for other types of coloring
            ii = w-vv[1]+1 # index of w into vv
            @assert w==vv[ii] # TODO remove
            if color[ii]>0
                forbiddenColors[color[ii]] = i
            end
        end
        for (col,fc) in enumerate(forbiddenColors)
            if fc!=i
                color[i] = col
                break
            end
        end
    end
    return Color(color, dim)
end

####
#
####

function color2seedmatrix(c::Color, bi::BiGraph)
    # Returns the seed matrix: B = AS (or B = SA in case of row
    # coloring)
    #
    # S is a sparse matrix as that is useful for recovery
    dim = c.dim
    color = c.c
    nc = maximum(color)
    vv = dim==1 ? bi.v1 : bi.v2 - bi.v2[1]+1
    S = zeros(Int8, length(vv), nc) # TODO: could use Bools here?
#    S = falses(length(vv), nc)
    for (i,c) in enumerate(color)
        S[vv[i],c] = 1
    end
    S = dim==1 ? S' : S
    return sparse(S)
end

######
# Jacobians
######

# Returns a seedmatrix for a Jacobian matrix.  Dim==2 is a column
# coloring (default).
function color_jac(A, dim=2)
    bi = matrix2bipartie_graph(A)
    c = partial_dist2_coloring(bi, dim)
    color2seedmatrix(c, bi)
end

function recover_jac!(jac::SparseMatrixCSC, Bcol, S, j::Int)
    # Recovers part of the Jacobian from one column j of its dense
    # representation B and seed matrix S.  Updates jac in-place.  jac
    # needs to have the exact sparsity pattern which was used to
    # create the coloring!
    #
    # B    = jac  S
    # nxp    nxm  mxp
    #
    # TODO: really not sure whether this is the best implementation.

    njac = nonzeros(jac)
    rowjac = rowvals(jac)
    rowS = rowvals(S)
    for i in nzrange(S,j)
        # some of Bcol (=B[:,j]) belongs into jac[:,row]
        for ii in nzrange(jac,rowS[i])
            njac[ii] = Bcol[rowjac[ii]]
        end
    end
    nothing
end
function recover_jac!(jac::SparseMatrixCSC, B, S)
    # Recovers the Jacobian from its dense representation B and seed
    # matrix S.  Updates jac in-place.  jac needs to have the exact
    # sparsity pattern which was used to create the coloring!
    #
    # B    = jac  S
    # nxp    nxm  mxp
    #
    # TODO: really not sure whether this is the best implementation.
    dim = size(jac,2)==size(S,1) ? 2 : 1  # TODO: this can mess up if n==m==p
    if dim==1
        error("not implemented")
    end
    njac = nonzeros(jac)
    rowjac = rowvals(jac)
    rowS = rowvals(S)
    for j=1:size(S,2) # columns of S and B
        # recover_jac!(jac, B[:,j], S, j) # TODO: this makes a copy...
        for i in nzrange(S,j)
            # some of Bcol (=B[:,j]) belongs into jac[:,row]
            for ii in nzrange(jac,rowS[i])
                njac[ii] = B[rowjac[ii],j]
            end
        end
    end
    nothing
end

function add_delta!(x::VectorLike, S, c, delta)
    # Adds a delta to all indices of x which correspond to color c of
    # seedmatrix S.
    if length(delta)==1
        for i in nzrange(S,c)
            x[rowvals(S)[i]] += delta
        end
    else
        for i in nzrange(S,c)
            ii = rowvals(S)[i]
            x[ii] += delta[ii]
        end
    end
    nothing
end
function remove_delta!(x::VectorLike, S, c, delta)
    # reverses add_delta!
    if length(delta)==1
        for i in nzrange(S,c)
            x[rowvals(S)[i]] -= delta
        end
    else
        for i in nzrange(S,c)
            ii = rowvals(S)[i]
            x[ii] -= delta[ii]
        end
    end
    nothing
end

end # module


text:
if !@isdefined(isnothing)
    isnothing(x) = x === nothing
end

splitlines(str; kwargs...) = readlines(IOBuffer(str); kwargs...)

askwargs(kwargs) = (; kwargs...)
function askwargs(flag::Bool)
    if !flag
        throw(ArgumentError("expect `true`"))
    end
    return NamedTuple()
end

struct LazyTestResult
    label::String
    message::String
    pass::Bool
end

ispass(result::LazyTestResult) = result.pass

# Infix operator wrapping `ispass` so that the failure case is pretty-printed
⊜(result, yes::Bool) = ispass(result)::Bool == yes

# To be shown via `@test` when failed:
function Base.show(io::IO, result::LazyTestResult)
    print(io, "⟪result: ")
    show(io, MIME"text/plain"(), result)
    print(io, "⟫")
end

function Base.show(io::IO, ::MIME"text/plain", result::LazyTestResult)
    if ispass(result)
        printstyled(io, "✔ PASS"; color = :green, bold = true)
    else
        printstyled(io, "😭 FAILED"; color = :red, bold = true)
    end
    println(io, ": ", result.label)
    for line in eachline(IOBuffer(result.message))
        println(io, " "^4, line)
    end
end

function root_project_or_failed_lazytest(pkg::PkgId)
    label = "$pkg"

    srcpath = Base.locate_package(pkg)
    if srcpath === nothing
        return LazyTestResult(
            label,
            """
            Package $pkg does not have a corresponding source file.
            """,
            false,
        )
    end

    pkgpath = dirname(dirname(srcpath))
    root_project_path, found = project_toml_path(pkgpath)
    if !found
        return LazyTestResult(
            label,
            """
            Project.toml file at project directory does not exist:
            $root_project_path
            """,
            false,
        )
    end
    return root_project_path
end

module _TempModule
end

eval_string(code::AbstractString) = include_string(_TempModule, code)

function checked_repr(obj)
    code = repr(obj)
    if !isequal(eval_string(code), obj)
        error("`$repr` is not `repr`-safe")
    end
    return code
end

const stdlibs = try
    Pkg.Types.stdlibs
catch
    try
        # https://github.com/JuliaLang/Pkg.jl/pull/1559
        Pkg.Types.stdlib  # julia < 1.4
    catch
        # https://github.com/JuliaLang/Pkg.jl/pull/696
        Pkg.Types.gather_stdlib_uuids  # julia < 1.1
    end
end

const _project_key_order = ["name", "uuid", "keywords", "license", "desc", "deps", "compat"]
project_key_order(key::String) =
    something(findfirst(x -> x == key, _project_key_order), length(_project_key_order) + 1)

print_project(io, dict) =
    TOML.print(io, dict, sorted = true, by = key -> (project_key_order(key), key))

ensure_exception(e::Exception) = e
ensure_exception(x) = ErrorException(string(x))

"""
    trydiff(label_a => text_a, label_b => text_b) -> string or exception
"""
function trydiff(
    (label_a, text_a)::Pair{<:AbstractString,<:AbstractString},
    (label_b, text_b)::Pair{<:AbstractString,<:AbstractString},
)
    # TODO: use a pure-Julia function
    cmd = `diff --label $label_a --label $label_b -u`
    mktemp() do path_a, io_a
        print(io_a, text_a)
        close(io_a)
        mktemp() do path_b, io_b
            print(io_b, text_b)
            close(io_b)
            try
                return read(ignorestatus(`$cmd $path_a $path_b`), String)
            catch err
                return ensure_exception(err)
            end
        end
    end
end

function format_diff(
    (label_a, text_a)::Pair{<:AbstractString,<:AbstractString},
    (label_b, text_b)::Pair{<:AbstractString,<:AbstractString},
)
    diff = trydiff(label_a => text_a, label_b => text_b)
    diff isa AbstractString && return diff
    # Fallback:
    return """
    *** $label_a ***
    $text_a

    *** $label_b ***
    $text_b
    """
end


text:
using ScoreDrivenModels, Distributions, LinearAlgebra, RecipesBase

using Test, Random, HypothesisTests, DelimitedFiles

include("utils.jl")
include("test_links.jl")
include("test_distributions.jl")

include("test_recursion.jl")
include("test_initial_params.jl")
include("test_diagnostics.jl")
include("test_estimate.jl")
include("test_simulate.jl")
include("test_cross_validation.jl")
include("test_visualization.jl")


text:
module LUMND

    export factorize, solve

    using Printf

    include("util.jl")
    include("metis.jl")
    include("partition.jl")
    include("factorize.jl")
    include("solve.jl")

end


text:
using Printf

"""
initializes x, d, dp arrays
    input:            n problem size
                      p order of differentiation
                      x[n] collocation points

    output: clog[n]   log of diagonal scaling
            s[n]      sign of diagonal scaling
            d[p*n]    initialized digonal coefficients of Dm

    other : a[n]      alpha(x[n])
            dp[p*n*n] work array

    log of function alpha
"""
function init!(n, p, x, a, clog, s, d, dp)

    for i = 1:n
        a[i] = alpha(0, x[i])
    end

    # scaling matrix
    for i = 1:n
        ci = 0.0e0
        si = 1.0e0
        xi = x[i]
        for k = 1:n
            (k == i) && continue
            xik = xi - x[k]
            if xik < 0.0e0
                si = - si
            end
            ci = ci + log(abs(xik))
        end
        clog[i] = ci + a[i]
        s[i] = si
    end

    # initialize d

    for k = 1:n
        xk = x[k]
        for m = 1:p
            i = k + (m - 1) * n
            d[i] = alpha(m, xk)
        end
    end

    return
end

"""
computes diagonal elements of Dm, m = 1[1]p

    input : n      problem size
            p      order of differentiation
            x[n]   collocation points

    output: d[p*n] diagonal elements of Dm
"""
function diag!(n, p, x, d)
    for k = 1:n
        xk = x[k]
        j = k + 1 > n ? 1 : k + 1
        while j != k
            xkj = 1.0e0 / (xk - x[j])
            for m = p:-1:2
                k1 = k + (m - 1) * n
                d[k1] = d[k1] + Float64(m) * xkj * d[k1 - n]
            end
            d[k] = d[k] + xkj
            j = mod(j, n) + 1
        end
    end
    return
end

"""
computes (off-diagonal) elements of Dm, m = 1[1]p

    input : n          problem size
            p          order of differentiation
            x[n]       collocation points
            d[p*n]     diagonal elements of Dm
                       output from subroutine diag

    output: dp[p*n*n]  differentiation matrices Dm

"""
function offdiag!(n, p, x, d, dp)
    for j = 1:n
        j1 = j - 1
        xj = x[j]
        for k = 1:n
            (k == j) && continue
            i = k + j1 * n
            dp[i] = 1.0e0 / (x[k] - xj)
        end
        dp[j + j1 * n] = d[j]
    end

    (p == 1) && return

    n2 = n * n
    for m = 2:p
        m1 = (m - 2) * n
        m2 = m1 + n
        fm = Float64(m)
        for j = 1:n
            j1 = (m1 + j - 1) * n
            j2 = j1 + n2
            xj = x[j]
            for k = 1:n
                (k == j) && continue
                k1 = j1 + k
                k2 = j2 + k
                dp[k2] = fm * (d[k + m1] - dp[k1]) / (x[k] - xj)
            end
            dp[j + j2] = d[j + m2]
        end
    end
    return
end

"""
Given n, p, (x_i, i = 1,n),
      and alpha^(j)(x_i), i = 1,n, j = 0,p-1

  computes D1_m and C such that D_m = C D1_m C , m = 1[1]p
  dp[m-1]*n*n+1,...,m*n*n) contains D1_m
  Cii = s[i]*exp(clog[i]), i = 1[1]n
"""
function  make_scaled_dp()
    #integer nmax, pmax
    nmax = 64
    pmax = 8

    x    = zeros(nmax)
    d    = zeros(nmax * pmax)
    dp   = zeros(pmax * nmax * nmax)
    a    = zeros(nmax)
    clog = zeros(nmax)
    s    = zeros(nmax)

    p = 2
    n = 8

    # collocation points
    for i ∈ 1:n
        x[i] = 10.0e0 * (i-1) / (n-1)
    end

    init!(n, p, x, a, clog, s, d, dp)
    diag!(n, p, x, d)
    offdiag!(n, p, x, d, dp)

    istart = (p - 1) * n * n + 1
    iend = p * n * n

    # to match printing from the Welfert paper
    for i ∈ istart:8:iend
        @printf("%8.3f %8.3f %8.3f %8.3f %8.3f %8.3f %8.3f %8.3f\n", dp[i:i+7]...)
    end

    @printf("\n%8.3f %8.3f %8.3f %8.3f %8.3f %8.3f %8.3f %8.3f\n", s[1:8]...)
    @printf("\n%8.3f %8.3f %8.3f %8.3f %8.3f %8.3f %8.3f %8.3f\n", clog[1:8]...)

end

"""
    if m = 0: log of alpha(x)
    if m > 0: alpha^m(x)/alpha(x)

    Here alpha(x) = 1 / (1 + x) ** n
"""
function alpha(m, x)
    n = 8
    if m == 0
        alpha = - float(n) * log(1.0e0 + x)
    else
        ai = 1.0e0
        for i = n:n + m - 1
            ai = - ai * float(i) / (1.0e0 + x)
        end
        alpha = ai
    end
    return alpha
end


text:
# Subtypes used for the main RiskAdjustedLinearization type
mutable struct RALNonlinearSystem{L <: AbstractRALF, S <: AbstractRALF, V <: AbstractRALF}
    μ::RALF2
    Λ::L
    Σ::S
    ξ::RALF2
    𝒱::V
    ccgf::Function
end

Λ_eltype(m::RALNonlinearSystem{L, S}) where {L, S} = L
Σ_eltype(m::RALNonlinearSystem{L, S}) where {L, S} = S

@inline function update!(m::RALNonlinearSystem{L, S, V}, z::C1, y::C1, Ψ::C2;
                         select::Vector{Symbol} = Symbol[:μ, :ξ, :𝒱]) where {L, S, V <: RALF2,
                                                                             C1 <: AbstractVector{<: Number},
                                                                             C2 <: AbstractMatrix{<: Number}}

    if :μ in select
        m.μ(z, y)
    end

    if :ξ in select
        m.ξ(z, y)
    end

    if :𝒱 in select
        m.𝒱(z, Ψ)
    end

    m
end

@inline function update!(m::RALNonlinearSystem{L, S, V}, z::C1, y::C1, Ψ::C2;
                         select::Vector{Symbol} = Symbol[:μ, :ξ, :𝒱]) where {L, S, V <: RALF4,
                                                                             C1 <: AbstractVector{<: Number}, C2 <: AbstractMatrix{<: Number}}

    if :μ in select
        m.μ(z, y)
    end

    if :ξ in select
        m.ξ(z, y)
    end

    if :𝒱 in select
        m.𝒱(z, y, Ψ, z)
    end

    m
end

mutable struct RALLinearizedSystem{JC5 <: AbstractMatrix{<: Number},
                                   JC6 <: AbstractMatrix{<: Number}, SJC <: AbstractDict{Symbol, NamedTuple}}
    μz::RALF2
    μy::RALF2
    ξz::RALF2
    ξy::RALF2
    J𝒱::Union{RALF2, RALF3}
    Γ₅::JC5
    Γ₆::JC6
    sparse_jac_caches::SJC
end

function RALLinearizedSystem(μz::RALF2, μy::RALF2, ξz::RALF2, ξy::RALF2, J𝒱::AbstractRALF,
                             Γ₅::AbstractMatrix{<: Number}, Γ₆::AbstractMatrix{<: Number})
    RALLinearizedSystem(μz, μy, ξz, ξy, J𝒱, Γ₅, Γ₆, Dict{Symbol, NamedTuple}())
end

@inline function update!(m::RALLinearizedSystem{JC5, JC6}, z::C1, y::C1, Ψ::C2;
                         select::Vector{Symbol} =
                         Symbol[:Γ₁, :Γ₂, :Γ₃, :Γ₄, :JV]) where {#JV <: RALF2,
                                                                 JC5, JC6,
                                                                 C1 <: AbstractVector{<: Number}, C2 <: AbstractMatrix{<: Number}}

    if :Γ₁ in select
        m.μz(z, y)
    end

    if :Γ₂ in select
        m.μy(z, y)
    end

    if :Γ₃ in select
        m.ξz(z, y)
    end

    if :Γ₄ in select
        m.ξy(z, y)
    end

    if :JV in select
        if isa(m.J𝒱, RALF2)
            m.J𝒱(z, Ψ)
        else
            m.J𝒱(z, y, Ψ)
        end
    end

    m
end

abstract type AbstractRiskAdjustedLinearization end

"""
    RiskAdjustedLinearization(μ, Λ, Σ, ξ, Γ₅, Γ₆, ccgf, z, y, Ψ, Nε)
    RiskAdjustedLinearization(nonlinear_system, linearized_system, z, y, Ψ, Nz, Ny, Nε)

Creates a first-order perturbation around the stochastic steady state of a discrete-time dynamic economic model.

The first method is the main constructor most users will want, while the second method is the default constructor.

### Inputs for First Method
- `μ::Function`: expected state transition function
- `ξ::Function`: nonlinear terms of the expectational equations
- `ccgf::Function`: conditional cumulant generating function of the exogenous shocks
- `Λ::Function` or `Λ::AbstractMatrix`: function or matrix mapping endogenous risk into state transition equations
- `Σ::Function` or `Σ::AbstractMatrix`: function or matrix mapping exogenous risk into state transition equations
- `Γ₅::AbstractMatrix{<: Number}`: coefficient matrix on one-period ahead expectation of state variables
- `Γ₆::AbstractMatrix{<: Number}`: coefficient matrix on one-period ahead expectation of jump variables
- `z::AbstractVector{<: Number}`: state variables in stochastic steady state
- `y::AbstractVector{<: Number}`: jump variables in stochastic steady state
- `Ψ::AbstractMatrix{<: Number}`: matrix linking deviations in states to deviations in jumps, i.e. ``y_t - y = \\Psi(z_t - z)``.
- `Nε::Int`: number of exogenous shocks

### Keywords for First Method
- `sss_vector_cache_init::Function = dims -> Vector{T}(undef, dims)`: initializer for the cache of steady state vectors.
- `Λ_cache_init::Function = dims -> Matrix{T}(undef, dims)`:  initializer for the cache of `Λ`
- `Σ_cache_init::Function = dims -> Matrix{T}(undef, dims)`:  initializer for the cache of `Λ`
- `jacobian_cache_init::Function = dims -> Matrix{T}(undef, dims)`: initializer for the cache of the Jacobians of `μ`, `ξ`, and `𝒱 `.
- `jump_dependent_shock_matrices::Bool = false`: if true, `Λ` and `Σ` are treated as `Λ(z, y)` and `Σ(z, y)`
    to allow dependence on jumps.
- `sparse_jacobian::Vector{Symbol} = Symbol[]`: pass the symbols `:μ`, `:ξ`, and/or `:𝒱 ` to declare that
    the Jacobians of these functions are sparse and should be differentiated using sparse methods from SparseDiffTools.jl
- `sparsity::AbstractDict = Dict{Symbol, Mtarix}()`: a dictionary for declaring the
    sparsity patterns of the Jacobians of `μ`, `ξ`, and `𝒱 `. The relevant keys are `:μz`, `:μy`, `:ξz`, `:ξy`, and `:J𝒱 `.
- `colorvec::AbstractDict = Dict{Symbol, Vector{Int}}()`: a dictionary for declaring the
    the matrix coloring vector. The relevant keys are `:μz`, `:μy`, `:ξz`, `:ξy`, and `:J𝒱 `.
- `sparsity_detection::Bool = false`: if true, use SparseDiffTools to determine the sparsity pattern.
    When false (default), the sparsity pattern is estimated by differentiating the Jacobian once
    with `ForwardDiff` and assuming any zeros in the calculated Jacobian are supposed to be zeros.

### Inputs for Second Method
- `nonlinear_system::RALNonlinearSystem`
- `linearized_system::RALLinearizedSystem`
- `z::AbstractVector{<: Number}`: state variables in stochastic steady state
- `y::AbstractVector{<: Number}`: jump variables in stochastic steady state
- `Ψ::AbstractMatrix{<: Number}`: matrix linking deviations in states to deviations in jumps, i.e. ``y_t - y = \\Psi(z_t - z)``.
- `Nz::Int`: number of state variables
- `Ny::Int`: number of jump variables
- `Nε::Int`: number of exogenous shocks
"""
mutable struct RiskAdjustedLinearization{C1 <: AbstractVector{<: Number}, C2 <: AbstractMatrix{<: Number}} <: AbstractRiskAdjustedLinearization

    nonlinear::RALNonlinearSystem
    linearization::RALLinearizedSystem
    z::C1 # Coefficients, TODO: at some point, we may or may not want to make z, y, and Ψ also DiffCache types
    y::C1
    Ψ::C2
    Nz::Int      # Dimensions
    Ny::Int
    Nε::Int
end

# The following constructor is typically the main constructor for most users.
# It will call a lower-level constructor that uses automatic differentiation
# to calculate the Jacobian functions.
# Note that here we pass in the ccgf, rather than 𝒱
function RiskAdjustedLinearization(μ::M, Λ::L, Σ::S, ξ::X, Γ₅::JC5, Γ₆::JC6, ccgf::CF,
                                   z::AbstractVector{T}, y::AbstractVector{T}, Ψ::AbstractMatrix{T},
                                   Nε::Int; sss_vector_cache_init::Function = dims -> Vector{T}(undef, dims),
                                   Λ_cache_init::Function = dims -> Matrix{T}(undef, dims),
                                   Σ_cache_init::Function = dims -> Matrix{T}(undef, dims),
                                   jump_dependent_shock_matrices::Bool = false,
                                   jacobian_cache_init::Function = dims -> Matrix{T}(undef, dims),
                                   sparse_jacobian::Vector{Symbol} = Symbol[],
                                   sparsity::AbstractDict{Symbol, AbstractMatrix} = Dict{Symbol, AbstractMatrix}(),
                                   colorvec::AbstractDict = Dict{Symbol, Vector{Int}}(),
                                   sparsity_detection::Bool = false) where {T <: Number, M <: Function, L, S,
                                                                            X <: Function,
                                                                            JC5 <: AbstractMatrix{<: Number},
                                                                            JC6 <: AbstractMatrix{<: Number},
                                                                            CF <: Function}
    # Get dimensions
    Nz  = length(z)
    Ny  = length(y)
    Nzy = Nz + Ny
    if Nε < 0
        throw(BoundsError("Nε cannot be negative"))
    end

    # Create wrappers enabling caching for μ and ξ
    Nzchunk = ForwardDiff.pickchunksize(Nz)
    Nychunk = ForwardDiff.pickchunksize(Ny)
    _μ = RALF2(μ, z, y, sss_vector_cache_init(Nz), (max(min(Nzchunk, Nychunk), 2), Nzchunk, Nychunk))
    _ξ = RALF2(ξ, z, y, sss_vector_cache_init(Ny), (max(min(Nzchunk, Nychunk), 2), Nzchunk, Nychunk))

    # Apply dispatch on Λ and Σ to figure what they should be
    return RiskAdjustedLinearization(_μ, Λ, Σ, _ξ, Γ₅, Γ₆, ccgf, z, y, Ψ, Nz, Ny, Nε, sss_vector_cache_init = sss_vector_cache_init,
                                     Λ_cache_init = Λ_cache_init,
                                     Σ_cache_init = Σ_cache_init,
                                     jump_dependent_shock_matrices = jump_dependent_shock_matrices,
                                     jacobian_cache_init = jacobian_cache_init,
                                     sparse_jacobian = sparse_jacobian, sparsity = sparsity,
                                     colorvec = colorvec, sparsity_detection = sparsity_detection)
end

# Constructor that uses ForwardDiff to calculate Jacobian functions.
# Users will not typically use this constructor, however, because it requires
# various functions of the RALNonlinearSystem and RALLinearizedSystem to already
# be wrapped with either an RALF1 or RALF2 type.
function RiskAdjustedLinearization(μ::M, Λ::L, Σ::S, ξ::X, Γ₅::JC5, Γ₆::JC6, ccgf::CF,
                                   z::AbstractVector{T}, y::AbstractVector{T}, Ψ::AbstractMatrix{T},
                                   Nz::Int, Ny::Int, Nε::Int; sss_vector_cache_init::Function = dims -> Vector{T}(undef, dims),
                                   jacobian_cache_init::Function = dims -> Matrix{T}(undef, dims),
                                   sparse_jacobian::Vector{Symbol} = Symbol[],
                                   sparsity::AbstractDict{Symbol, AbstractMatrix} = Dict{Symbol, AbstractMatrix}(),
                                   colorvec::AbstractDict = Dict{Symbol, Vector{Int}}(),
                                   sparsity_detection::Bool = false) where {T <: Number, M <: RALF2, L <: RALF1, S <: RALF1,
                                                                            X <: RALF2,
                                                                            JC5 <: AbstractMatrix{<: Number},
                                                                            JC6 <: AbstractMatrix{<: Number},
                                                                            CF <: Function}

    jac_cache = Dict{Symbol, NamedTuple}()

    # Use RALF2 wrapper to create Jacobian functions with caching for μ, ξ.
    # Use the tuple to select the correct Dual cache b/c μ is in place
    if :μ in sparse_jacobian
        μz, μy, jac_cache[:μz], jac_cache[:μy] =
            construct_μ_jacobian_function(μ, z, y;
                                          sparsity_z = haskey(sparsity, :μz) ? sparsity[:μz] : nothing,
                                          sparsity_y = haskey(sparsity, :μy) ? sparsity[:μy] : nothing,
                                          colorvec_z = haskey(sparsity, :μz) ? sparsity[:μz] : nothing,
                                          colorvec_y = haskey(sparsity, :μy) ? sparsity[:μy] : nothing,
                                          sparsity_detection = sparsity_detection)
    else
        μz = RALF2((F, z, y) -> ForwardDiff.jacobian!(F, x -> μ(x, y, (1, 2)), z), z, y,
                   jacobian_cache_init((Nz, Nz)))
        μy = RALF2((F, z, y) -> ForwardDiff.jacobian!(F, x -> μ(z, x, (2, 3)), y), z, y,
                   jacobian_cache_init((Nz, Ny)))
    end

    if :ξ in sparse_jacobian
        ξz, ξy, jac_cache[:ξz], jac_cache[:ξy] =
            construct_ξ_jacobian_function(ξ, z, y;
                                          sparsity_z = haskey(sparsity, :ξz) ? sparsity[:ξz] : nothing,
                                          sparsity_y = haskey(sparsity, :ξy) ? sparsity[:ξy] : nothing,
                                          colorvec_z = haskey(sparsity, :ξz) ? sparsity[:ξz] : nothing,
                                          colorvec_y = haskey(sparsity, :ξy) ? sparsity[:ξy] : nothing,
                                          sparsity_detection = sparsity_detection)
    else
        ξz = RALF2((F, z, y) -> ForwardDiff.jacobian!(F, x -> ξ(x, y, (1, 2)), z), z, y,
                   jacobian_cache_init((Ny, Nz)))
        ξy = RALF2((F, z, y) -> ForwardDiff.jacobian!(F, x -> ξ(z, x, (2, 3)), y), z, y,
                   jacobian_cache_init((Ny, Ny)))
    end

    # Check if we need to compute the left divide or not
    avoid_Λ = isa(get_cache_type(Λ), AbstractMatrix) && all(Λ.cache .== 0.)

    # Check if Σ's cache is sparse, which matters only if Λ is relevant b/c
    # a left-divide of two sparse matrices A and B (i.e. A \ B) will not work
    # b/c the LU factorization algoritm employed needs more structure on A or B.
    sparse_cache_Σ = !avoid_Λ && isa(Σ.cache, AbstractMatrix) ? issparse(Σ.cache) : issparse(Σ.cache.du)

    # Create RALF2 wrappers for 𝒱 and its Jacobian J𝒱
    if applicable(ccgf, Γ₅, z) # Check if ccgf is in place or not
        _𝒱 = function _𝒱_oop(F, z, Ψ)
            Σ0 = Σ(z)
            if size(Σ0) != (Nz, Nε)
                Σ0 = sparse_cache_Σ ? Matrix(reshape(Σ0, Nz, Nε)) : reshape(Σ0, Nz, Nε)
            end
            if avoid_Λ
                F .= ccgf((Γ₅ + Γ₆ * Ψ) * Σ0, z)
            else
                Λ0 = Λ(z)
                if size(Λ0) != (Nz, Ny)
                    Λ0 = reshape(Λ0, Nz, Ny)
                end
                F .= ccgf((Γ₅ + Γ₆ * Ψ) * ((I - (Λ0 * Ψ)) \ Σ0), z)
            end
        end
    else # in place
        _𝒱 = function _𝒱_ip(F, z, Ψ)
            Σ0 = Σ(z)
            if size(Σ0) != (Nz, Nε)
                Σ0 = reshape(Σ0, Nz, Nε)
            end
            if avoid_Λ
                ccgf(F, (Γ₅ + Γ₆ * Ψ) * Σ0, z)
            else
                Λ0 = Λ(z)
                if size(Λ0) != (Nz, Ny)
                    Λ0 = reshape(Λ0, Nz, Ny)
                end
                ccgf(F, (Γ₅ + Γ₆ * Ψ) * ((I - (Λ0 * Ψ)) \ Σ0), z)
            end
        end
    end
    Nzchunk = ForwardDiff.pickchunksize(Nz)
    Nychunk = ForwardDiff.pickchunksize(Ny)
    𝒱 = RALF2((F, z, Ψ) -> _𝒱(F, z, Ψ), z, Ψ, sss_vector_cache_init(Ny), (max(min(Nzchunk, Nychunk), 2), Nzchunk))

    if :𝒱 in sparse_jacobian
        J𝒱, jac_cache[:J𝒱] = construct_𝒱_jacobian_function(𝒱, ccgf, Λ, Σ, Γ₅, Γ₆, z, Ψ;
                                                           sparsity = haskey(sparsity, :J𝒱) ? sparsity[:J𝒱] : nothing,
                                                           colorvec = haskey(colorvec, :J𝒱) ? colorvec[:J𝒱] : nothing,
                                                           sparsity_detection = sparsity_detection)
    else
        _J𝒱(F, z, Ψ) = ForwardDiff.jacobian!(F, x -> 𝒱(x, Ψ, (1, 2)), z)
        J𝒱           = RALF2((F, z, Ψ) -> _J𝒱(F, z, Ψ), z, Ψ, jacobian_cache_init((Ny, Nz)))
    end

    # Form underlying RAL blocks
    nonlinear_system  = RALNonlinearSystem(μ, Λ, Σ, ξ, 𝒱, ccgf)
    linearized_system = RALLinearizedSystem(μz, μy, ξz, ξy, J𝒱, Γ₅, Γ₆, jac_cache)

    return RiskAdjustedLinearization(nonlinear_system, linearized_system, z, y, Ψ, Nz, Ny, Nε)
end

# Handles case where Λ and Σ are RALF2
function RiskAdjustedLinearization(μ::M, Λ::L, Σ::S, ξ::X, Γ₅::JC5, Γ₆::JC6, ccgf::CF,
                                   z::AbstractVector{T}, y::AbstractVector{T}, Ψ::AbstractMatrix{T},
                                   Nz::Int, Ny::Int, Nε::Int; sss_vector_cache_init::Function = dims -> Vector{T}(undef, dims),
                                   jacobian_cache_init::Function = dims -> Matrix{T}(undef, dims),
                                   sparse_jacobian::Vector{Symbol} = Symbol[],
                                   sparsity::AbstractDict{Symbol, AbstractMatrix} = Dict{Symbol, AbstractMatrix}(),
                                   colorvec::AbstractDict = Dict{Symbol, Vector{Int}}(),
                                   sparsity_detection::Bool = false) where {T <: Number, M <: RALF2, L <: RALF2, S <: RALF2,
                                                                            X <: RALF2,
                                                                            JC5 <: AbstractMatrix{<: Number},
                                                                            JC6 <: AbstractMatrix{<: Number},
                                                                            CF <: Function}

    jac_cache = Dict{Symbol, NamedTuple}()

    # Use RALF2 wrapper to create Jacobian functions with caching for μ, ξ.
    # Use the tuple to select the correct Dual cache b/c μ is in place
    if :μ in sparse_jacobian
        μz, μy, jac_cache[:μz], jac_cache[:μy] =
            construct_μ_jacobian_function(μ, z, y;
                                          sparsity_z = haskey(sparsity, :μz) ? sparsity[:μz] : nothing,
                                          sparsity_y = haskey(sparsity, :μy) ? sparsity[:μy] : nothing,
                                          colorvec_z = haskey(sparsity, :μz) ? sparsity[:μz] : nothing,
                                          colorvec_y = haskey(sparsity, :μy) ? sparsity[:μy] : nothing,
                                          sparsity_detection = sparsity_detection)
    else
        μz = RALF2((F, z, y) -> ForwardDiff.jacobian!(F, x -> μ(x, y, (1, 2)), z), z, y,
                   jacobian_cache_init((Nz, Nz)))
        μy = RALF2((F, z, y) -> ForwardDiff.jacobian!(F, x -> μ(z, x, (2, 3)), y), z, y,
                   jacobian_cache_init((Nz, Ny)))
    end

    if :ξ in sparse_jacobian
        ξz, ξy, jac_cache[:ξz], jac_cache[:ξy] =
            construct_ξ_jacobian_function(μ, z, y;
                                          sparsity_z = haskey(sparsity, :ξz) ? sparsity[:ξz] : nothing,
                                          sparsity_y = haskey(sparsity, :ξy) ? sparsity[:ξy] : nothing,
                                          colorvec_z = haskey(sparsity, :ξz) ? sparsity[:ξz] : nothing,
                                          colorvec_y = haskey(sparsity, :ξy) ? sparsity[:ξy] : nothing,
                                          sparsity_detection = sparsity_detection)
    else
        ξz = RALF2((F, z, y) -> ForwardDiff.jacobian!(F, x -> ξ(x, y, (1, 2)), z), z, y,
                   jacobian_cache_init((Ny, Nz)))
        ξy = RALF2((F, z, y) -> ForwardDiff.jacobian!(F, x -> ξ(z, x, (2, 3)), y), z, y,
                   jacobian_cache_init((Ny, Ny)))
    end

    # Check if we need to compute the left divide or not
    avoid_Λ = isa(get_cache_type(Λ), AbstractMatrix) && all(Λ.cache .== 0.)

    # Check if Σ's cache is sparse, which matters only if Λ is relevant b/c
    # a left-divide of two sparse matrices A and B (i.e. A \ B) will not work
    # b/c the LU factorization algoritm employed needs more structure on A or B.
    sparse_cache_Σ = !avoid_Λ && isa(Σ.cache, AbstractMatrix) ? issparse(Σ.cache) : issparse(Σ.cache.du)

    # Create RALF2 wrappers for 𝒱 and its Jacobian J𝒱
    if applicable(ccgf, Γ₅, z) # Check if ccgf is in place or not
        _𝒱 = function _𝒱_oop(F, z, y, Ψ, zₜ)
            yₜ = y + Ψ * (zₜ - z)
            Σ0 = Σ(zₜ, yₜ)
            if size(Σ0) != (Nz, Nε)
                Σ0 = reshape(Σ0, Nz, Nε)
            end
            if avoid_Λ
                F .= ccgf((Γ₅ + Γ₆ * Ψ) * Σ0, zₜ)
            else
                Λ0 = Λ(zₜ, yₜ)
                if size(Λ0) != (Nz, Ny)
                    Λ0 = reshape(Λ0, Nz, Ny)
                end
                F .= ccgf((Γ₅ + Γ₆ * Ψ) * ((I - (Λ0 * Ψ)) \ Σ0), zₜ)
            end
        end
    else # in place
        _𝒱 = function _𝒱_ip(F, z, y, Ψ, zₜ)
            yₜ = y + Ψ * (zₜ - z)
            Σ0 = Σ(zₜ, yₜ)
            if size(Σ0) != (Nz, Nε)
                Σ0 = reshape(Σ0, Nz, Nε)
            end
            if avoid_Λ
                ccgf(F, (Γ₅ + Γ₆ * Ψ) * Σ0, zₜ)
            else
                Λ0 = Λ(zₜ, yₜ)
                if size(Λ0) != (Nz, Ny)
                    Λ0 = reshape(Λ0, Nz, Ny)
                end
                ccgf(F, (Γ₅ + Γ₆ * Ψ) * ((I - (Λ0 * Ψ)) \ Σ0), zₜ)
            end
        end
    end
    Nzchunk = ForwardDiff.pickchunksize(Nz)
    Nychunk = ForwardDiff.pickchunksize(Ny)
    𝒱       = RALF4((F, z, y, Ψ, zₜ) -> _𝒱(F, z, y, Ψ, zₜ), z, y, Ψ, z, sss_vector_cache_init(Ny),
                    (max(min(Nzchunk, Nychunk), 2), Nzchunk))

    if :𝒱 in sparse_jacobian
        J𝒱, jac_cache[:J𝒱] = construct_𝒱_jacobian_function(𝒱, ccgf, Λ, Σ, Γ₅, Γ₆, z, y, Ψ;
                                                           sparsity = haskey(sparsity, :J𝒱) ? sparsity[:J𝒱] : nothing,
                                                           colorvec = haskey(colorvec, :J𝒱) ? colorvec[:J𝒱] : nothing,
                                                           sparsity_detection = sparsity_detection)
    else
        _J𝒱(F, z, y, Ψ) = ForwardDiff.jacobian!(F, zₜ -> 𝒱(z, y, Ψ, zₜ, (4, 2)), z) # use zₜ argument to infer the cache
        J𝒱              = RALF3((F, z, y, Ψ) -> _J𝒱(F, z, y, Ψ), z, y, Ψ, jacobian_cache_init((Ny, Nz)))
    end

    # Form underlying RAL blocks
    nonlinear_system  = RALNonlinearSystem(μ, Λ, Σ, ξ, 𝒱, ccgf)
    linearized_system = RALLinearizedSystem(μz, μy, ξz, ξy, J𝒱, Γ₅, Γ₆, jac_cache)

    return RiskAdjustedLinearization(nonlinear_system, linearized_system, z, y, Ψ, Nz, Ny, Nε)
end

# The following four constructors cover different common cases for the Λ and Σ functions.
function RiskAdjustedLinearization(μ::M, Λ::L, Σ::S, ξ::X, Γ₅::JC5, Γ₆::JC6, ccgf::CF,
                                   z::AbstractVector{T}, y::AbstractVector{T}, Ψ::AbstractMatrix{T},
                                   Nz::Int, Ny::Int, Nε::Int; sss_vector_cache_init::Function = dims -> Vector{T}(undef, dims),
                                   Λ_cache_init::Function = dims -> Matrix{T}(undef, dims),
                                   Σ_cache_init::Function = dims -> Matrix{T}(undef, dims),
                                   jump_dependent_shock_matrices::Bool = false,
                                   jacobian_cache_init::Function = dims -> Matrix{T}(undef, dims),
                                   sparse_jacobian::Vector{Symbol} = Symbol[],
                                   sparsity::AbstractDict = Dict{Symbol, Matrix}(),
                                   colorvec::AbstractDict = Dict{Symbol, Vector{Int}}(),
                                   sparsity_detection::Bool = false) where {T <: Number, M <: RALF2, L <: Function, S <: Function,
                                                                            X <: RALF2,
                                                                            JC5 <: AbstractMatrix{<: Number},
                                                                            JC6 <: AbstractMatrix{<: Number},
                                                                            CF <: Function}

    # Create wrappers enabling caching for Λ and Σ
    Nzchunk = ForwardDiff.pickchunksize(Nz)
    Nychunk = ForwardDiff.pickchunksize(Ny)
    if jump_dependent_shock_matrices
        _Λ = RALF2(Λ, z, y, Λ_cache_init((Nz, Ny)), (max(min(Nzchunk, Nychunk), 2), Nzchunk))
        _Σ = RALF2(Σ, z, y, Σ_cache_init((Nz, Nε)), (max(min(Nzchunk, Nychunk), 2), Nzchunk))
    else
        _Λ = RALF1(Λ, z, Λ_cache_init((Nz, Ny)))
        _Σ = RALF1(Σ, z, Σ_cache_init((Nz, Nε)))
    end

    return RiskAdjustedLinearization(μ, _Λ, _Σ, ξ, Γ₅, Γ₆, ccgf, z, y, Ψ, Nz, Ny, Nε, sss_vector_cache_init = sss_vector_cache_init,
                                     jacobian_cache_init = jacobian_cache_init, sparse_jacobian = sparse_jacobian,
                                     sparsity = sparsity, sparsity_detection = sparsity_detection, colorvec = colorvec)
end

function RiskAdjustedLinearization(μ::M, Λ::L, Σ::S, ξ::X, Γ₅::JC5, Γ₆::JC6, ccgf::CF,
                                   z::AbstractVector{T}, y::AbstractVector{T}, Ψ::AbstractMatrix{T},
                                   Nz::Int, Ny::Int, Nε::Int; sss_vector_cache_init::Function = dims -> Vector{T}(undef, dims),
                                   Λ_cache_init::Function = dims -> Matrix{T}(undef, dims),
                                   Σ_cache_init::Function = dims -> Matrix{T}(undef, dims),
                                   jump_dependent_shock_matrices::Bool = false,
                                   jacobian_cache_init::Function = dims -> Matrix{T}(undef, dims),
                                   sparse_jacobian::Vector{Symbol} = Symbol[],
                                   sparsity::AbstractDict = Dict{Symbol, Matrix}(),
                                   colorvec::AbstractDict = Dict{Symbol, Vector{Int}}(),
                                   sparsity_detection::Bool = false) where {T <: Number, M <: RALF2,
                                                                            L <: AbstractMatrix{<: Number}, S <: Function,
                                                                            X <: RALF2,
                                                                            JC5 <: AbstractMatrix{<: Number},
                                                                            JC6 <: AbstractMatrix{<: Number},
                                                                            CF <: Function}

    # Create wrappers enabling caching for Λ and Σ
    if jump_dependent_shock_matrices
        _Λ = RALF2(Λ)
        _Σ = RALF2(Σ, z, y, Σ_cache_init((Nz, Nε)), (max(min(Nzchunk, Nychunk), 2), Nzchunk))
    else
        _Λ = RALF1(Λ)
        _Σ = RALF1(Σ, z, Σ_cache_init((Nz, Nε)))
    end

    return RiskAdjustedLinearization(μ, _Λ, _Σ, ξ, Γ₅, Γ₆, ccgf, z, y, Ψ, Nz, Ny, Nε, sss_vector_cache_init = sss_vector_cache_init,
                                     jacobian_cache_init = jacobian_cache_init, sparse_jacobian = sparse_jacobian,
                                     sparsity = sparsity, sparsity_detection = sparsity_detection, colorvec = colorvec)
end

function RiskAdjustedLinearization(μ::M, Λ::L, Σ::S, ξ::X, Γ₅::JC5, Γ₆::JC6, ccgf::CF,
                                   z::AbstractVector{T}, y::AbstractVector{T}, Ψ::AbstractMatrix{T},
                                   Nz::Int, Ny::Int, Nε::Int; sss_vector_cache_init::Function = dims -> Vector{T}(undef, dims),
                                   Λ_cache_init::Function = dims -> Matrix{T}(undef, dims),
                                   Σ_cache_init::Function = dims -> Matrix{T}(undef, dims),
                                   jump_dependent_shock_matrices::Bool = false,
                                   jacobian_cache_init::Function = dims -> Matrix{T}(undef, dims),
                                   sparse_jacobian::Vector{Symbol} = Symbol[],
                                   sparsity::AbstractDict = Dict{Symbol, Matrix}(),
                                   colorvec::AbstractDict = Dict{Symbol, Vector{Int}}(),
                                   sparsity_detection::Bool = false) where {T <: Number, M <: RALF2, L <: Function, S <: AbstractMatrix{<: Number},
                                                                            X <: RALF2,
                                                                            JC5 <: AbstractMatrix{<: Number},
                                                                            JC6 <: AbstractMatrix{<: Number},
                                                                            CF <: Function}

    # Create wrappers enabling caching for Λ and Σ
    Nzchunk = ForwardDiff.pickchunksize(Nz)
    Nychunk = ForwardDiff.pickchunksize(Ny)
    if jump_dependent_shock_matrices
        _Λ = RALF2(Λ, z, y, Λ_cache_init((Nz, Ny)), (max(min(Nzchunk, Nychunk), 2), Nzchunk))
        _Σ = RALF2(Σ)
    else
        _Λ = RALF1(Λ, z, Λ_cache_init((Nz, Ny)))
        _Σ = RALF1(Σ)
    end

    return RiskAdjustedLinearization(μ, _Λ, _Σ, ξ, Γ₅, Γ₆, ccgf, z, y, Ψ, Nz, Ny, Nε, sss_vector_cache_init = sss_vector_cache_init,
                                     jacobian_cache_init = jacobian_cache_init, sparse_jacobian = sparse_jacobian,
                                     sparsity = sparsity, sparsity_detection = sparsity_detection, colorvec = colorvec)
end

function RiskAdjustedLinearization(μ::M, Λ::L, Σ::S, ξ::X, Γ₅::JC5, Γ₆::JC6, ccgf::CF,
                                   z::AbstractVector{T}, y::AbstractVector{T}, Ψ::AbstractMatrix{T},
                                   Nz::Int, Ny::Int, Nε::Int; sss_vector_cache_init::Function = dims -> Vector{T}(undef, dims),
                                   Λ_cache_init::Function = dims -> Matrix{T}(undef, dims),
                                   Σ_cache_init::Function = dims -> Matrix{T}(undef, dims),
                                   jump_dependent_shock_matrices::Bool = false,
                                   jacobian_cache_init::Function = dims -> Matrix{T}(undef, dims),
                                   sparse_jacobian::Vector{Symbol} = Symbol[],
                                   sparsity::AbstractDict = Dict{Symbol, Matrix}(),
                                   sparsity_detection::Bool = false) where {T <: Number, M <: RALF2,
                                                                            L <: AbstractMatrix{<: Number}, S <: AbstractMatrix{<: Number},
                                                                            X <: RALF2,
                                                                            JC5 <: AbstractMatrix{<: Number},
                                                                            JC6 <: AbstractMatrix{<: Number},
                                                                            CF <: Function}

    # Create wrappers enabling caching for Λ and Σ
    _Λ = RALF1(Λ)
    _Σ = RALF1(Σ)

    return RiskAdjustedLinearization(μ, _Λ, _Σ, ξ, Γ₅, Γ₆, ccgf, z, y, Ψ, Nz, Ny, Nε, sss_vector_cache_init = sss_vector_cache_init,
                                     jacobian_cache_init = jacobian_cache_init, sparse_jacobian = sparse_jacobian,
                                     sparsity = sparsity, sparsity_detection = sparsity_detection, colorvec = colorvec)
end

## Print statements for RAL objects
function Base.show(io::IO, m::AbstractRiskAdjustedLinearization)
    @printf io "Risk-Adjusted Linearization of an Economic Model\n"
    @printf io "No. of state variables:      %i\n" m.Nz
    @printf io "No. of jump variables:       %i\n" m.Ny
    @printf io "No. of exogenous shocks:     %i\n" m.Nε
end

function Base.show(io::IO, m::RALNonlinearSystem)
    @printf io "RALNonlinearSystem"
end

function Base.show(io::IO, m::RALLinearizedSystem)
    @printf io "RALLinearizedSystem"
end

## Indexing for convenient access to steady state values
@inline function Base.getindex(m::RiskAdjustedLinearization, sym::Symbol)
    if sym in [:μ_sss, :ξ_sss, :𝒱_sss, :Σ_sss, :Λ_sss]
        m.nonlinear[sym]
    elseif sym in [:Γ₁, :Γ₂, :Γ₃, :Γ₄, :Γ₅, :Γ₆, :JV]
        m.linearization[sym]
    else
        throw(KeyError("key $sym not found"))
    end
end

@inline function Base.getindex(m::RALNonlinearSystem, sym::Symbol)
    if sym == :μ_sss
        isnothing(m.μ.cache) ? error("μ is out of place, so its stochastic steady state value is not cached.") : m.μ.cache.du
    elseif sym == :ξ_sss
        isnothing(m.ξ.cache) ? error("ξ is out of place, so its stochastic steady state value is not cached.") : m.ξ.cache.du
    elseif sym == :𝒱_sss
        m.𝒱.cache.du
    elseif sym == :Σ_sss
        if isnothing(m.Σ.cache)
            error("Λ is out of place, so its stochastic steady state value is not cached.")
        elseif isa(m.Σ.cache, DiffCache)
            m.Σ.cache.du
        else
            m.Σ.cache
        end
    elseif sym == :Λ_sss
        if isnothing(m.Λ.cache)
            error("Λ is out of place, so its stochastic steady state value is not cached.")
        elseif isa(m.Λ.cache, DiffCache)
            m.Λ.cache.du
        else
            m.Λ.cache
        end
    else
        throw(KeyError("key $sym not found"))
    end
end

@inline function Base.getindex(m::RALLinearizedSystem, sym::Symbol)
    if sym == :Γ₁
        m.μz.cache.du
    elseif sym == :Γ₂
        m.μy.cache.du
    elseif sym == :Γ₃
        m.ξz.cache.du
    elseif sym == :Γ₄
        m.ξy.cache.du
    elseif sym == :Γ₅
        m.Γ₅
    elseif sym == :Γ₆
        m.Γ₆
    elseif sym == :JV
        m.J𝒱.cache.du
    else
        throw(KeyError("key $sym not found"))
    end
end

## Methods for using RiskAdjustedLinearization
@inline getvalues(m::RiskAdjustedLinearization) = (m.z, m.y, m.Ψ)
@inline getvecvalues(m::RiskAdjustedLinearization) = vcat(m.z, m.y, vec(m.Ψ))
@inline nonlinear_system(m::RiskAdjustedLinearization) = m.nonlinear
@inline linearized_system(m::RiskAdjustedLinearization) = m.linearization

@inline function update!(m::RiskAdjustedLinearization)
    update!(nonlinear_system(m), m.z, m.y, m.Ψ)
    update!(linearized_system(m), m.z, m.y, m.Ψ)

    m
end

@inline function update!(m::RiskAdjustedLinearization, z::C1, y::C1, Ψ::C2;
                         update_cache::Bool = true) where {C1 <: AbstractVector{<: Number}, C2 <: AbstractMatrix{<: Number}}

    # Update values of the affine approximation
    m.z .= z
    m.y .= y
    m.Ψ .= Ψ

    # Update the cached vectors and Jacobians
    if update_cache
        update!(m)
    end

    m
end


text:
# This file was generated, do not modify it. # hide
data = OpenML.load(42178) # data set from OpenML.org
df0 = DataFrames.DataFrame(data)
first(df0, 4)

max_stars_repo_name: alan-turing-institute/MLJTutorials

max_issues_repo_name: alan-turing-institute/MLJTutorials

max_forks_repo_name: alan-turing-institute/MLJTutorials


text:
function distance(distfn, x::Vector{<:Number}, y::Vector{<:Number}, nbins = 200)
    min_x, max_x = extrema(x)
    min_y, max_y = extrema(y)
    min_val = min(min_x, min_y)
    max_val = max(max_x, max_y)
    val_range = range(min_val, stop = max_val, length = nbins)
    hist_x = fit(Histogram, x, val_range)
    hist_y = fit(Histogram, y, val_range)
    hist_x = normalize(hist_x, mode = :density)
    hist_y = normalize(hist_y, mode = :density)
    distance = distfn(hist_x.weights, hist_y.weights)
    return distance
end


text:
setup_tree() = Tree("a",[Tree(x,y) for (x,y) in zip("bcd","xyz")])

@testset "Interpreter testing" begin
    t = Tree("a",[Tree("b",[]), Tree("c",[]),"d"])
    struct Interp1 <: Interpreter end
    @rule a(i::Interp1,tree) = append!(visit_children(i,tree), ["e"])
    @rule b(i::Interp1,tree) = "B"
    @rule c(i::Interp1,tree) = "C"

    @test visit(Interp1(),t) == ["B","C","d","e"]

    struct Interp3 <: Interpreter end
    @rule b(i::Interp3, tree) = "B"
    @rule c(i::Interp3, tree) = "C"

    @test visit(Interp3(),t) == ["B","C","d"]
end

@testset "Transformer testing" begin
    t = Tree("add", [Tree("sub", [Tree("i", ["3"]), Tree("f", ["1.1"])]), Tree("i", ["1"])])

    struct T <: Transformer end

    @inline_rule i(t::T,args...) = Base.parse.(Int64,args...)
    @inline_rule f(t::T,args...) = Base.parse.(Float64,args...)
    @inline_rule sub(t::T,args...) = args[1] - args[2]
    @inline_rule add(t::T,args...) = sum(args)

    @test transform(T(),t) == 2.9
end


text:
# transforms.jl

"""
Compute the Z transform of a sequence, defined as `S(z) = \sum_k s_k z^{-k}`.
"""
function ztransform(s::Sequence, z)
    T = promote_type(eltype(s), eltype(z), eltype(1/z))
    S = zero(T)
    for k in nzrange(s)
        S += s[k] * z^(-k)
    end
    S
end

"""
The Fourier transform of a sequence is defined as `S(ω) = \sum_k s_k e^{-i ω k}`.
It corresponds to the Z transform with `z = e^{i ω}`. The Fourier transform is a
`2π`-periodic function of `ω`.
"""
fouriertransform(s::Sequence, ω) = ztransform(s, exp(im*ω))

"""
The ZTransform type is a lazy wrapper for the `ztransform` function.
"""
immutable ZTransform{S <: Sequence}
    sequence     ::  S
end

sequence(zt::ZTransform) = zt.sequence

(zt::ZTransform)(z) = ztransform(sequence(zt), z)

"""
The FourierTransform type is a lazy wrapper for the `fouriertransform` function.
"""
immutable FourierTransform{S <: Sequence}
    seq     ::  S
end

sequence(ft::FourierTransform) = ft.sequence

(ft::FourierTransform)(ω) = fouriertransform(sequence(ft), ω)


text:
"""
Builds the tree.

Rules of recursion: 1) Believe that it works. 2) Start by checking
for the base case (no further information gain). 3) Prepare for
giant stack traces.
"""
function build_tree(rows::Data)::Node
    # Try partitioning the dataset on each of the unique attribute,
    # calculate the information gain,
    # and return the question that produces the highest gain.
    gain, question = find_best_split(rows)

    # Base case: no further info gain
    # Since we can ask no further questions,
    # we'll return a leaf.
    if gain ≈ 0.0
        return Leaf(rows)
    end
    # If we reach here, we have found a useful feature / value
    # to partition on.
    true_rows, false_rows = partition(rows, question)

    # Recursively build the true branch.
    true_branch = build_tree(true_rows)

    # Recursively build the false branch.
    false_branch = build_tree(false_rows)

    # Return a Question node.
    # This records the best feature / value to ask at this point,
    # as well as the branches to follow
    # depending on the answer.
    Decision(question, true_branch, false_branch)
end


text:
#inspired by nucleus theme for Atom
dark_palette = [
    colorant"#FE4365", # red
    colorant"#eca25c", # orange
    colorant"#3f9778", # green
    colorant"#005D7F" # blue
]
dark_bg = colorant"#363D46"

_themes[:dark] = PlotTheme(
    bg = dark_bg,
    bginside = colorant"#30343B",
    fg = colorant"#ADB2B7",
    fgtext = colorant"#FFFFFF",
    fgguide = colorant"#FFFFFF",
    fglegend = colorant"#FFFFFF",
    legendfontcolor = colorant"#FFFFFF",
    legendtitlefontcolor = colorant"#FFFFFF",
    titlefontcolor = colorant"#FFFFFF",
    palette = expand_palette(dark_bg, dark_palette; lchoices = [57], cchoices = [100]),
    colorgradient = :fire
)


text:
using DifferentialEquations
using Plots
using Distributions
using Random: Random
########### Defines the model parameters for Figure 3 ###########
n_parasites = 100; # number of parasites
K = 100; # carrying capacity
ux = 0.2; # mortality rate of uninfected hosts
# mortality rate of infected hosts (always >= ux)
Random.seed!(6376);
ui = rand(200:1000, n_parasites) / 1000;
# define 3 random constraints per strain: r1, r2, r3
# r1 = virulence of the strain
# r2 = fraction of virulence attributable to fecundity loss
# r3 = fraction of offspring of infected hosts whichare infected
r1 = rand(Float64, n_parasites);
r2 = rand(Float64, n_parasites);
r3 = rand() * r1 # in this case, r3 is limited by r1
bx = 1.0; # birth rate of uninfected individuals
bi = bx .* r3 .* (1 .- r1 .* r2); # birth rate of infected individuals
α = rand(Float64, n_parasites); # cost of vertical transmission
βy = r1 .- (α .* bi) / bx; # horizontal transmission rate
# ensure no values in βy are negative by setting them to 0
for (index, value) in enumerate(βy)
    if value < 0
        βy[index] = 0
    end
end
ei = bx .* (1 .- r3) .* (1 .- (r1 .* r2)); # number of infected offspring from an infected host
# set up the model parameters
parameters = (bx=bx, βy=βy, ei=ei, c=c, K=K, ux=ux, bi=bi, ui=ui);
# define the numbers of individuals in each category
Y = zeros(Float64, n_parasites); # number of parasites of each strain
Y[1] = 1.0; # start with one initial parasite
X0 = 80.0; # number of initial uninfected hosts
new_U = vcat(X0, Y); # merge both together into one array
# define the limits of the simulation
windowstart = 0.0;
windowsize = 1000.0;
windowend = windowstart + windowsize;
N = zeros(Float64, (n_parasites + 1, (n_parasites - 1) * Int(windowsize) + 1));
figure_directory = "Figure3" # output folder for subplots of Figure 3