text:
module targetfigures

using ..matrixcore
using ..fresnelltools
using ..analyticalmaterials.hs_Palm

export scann_singleparameter, scan_plasmon_singleparameter, scan_plasmon_dualparameter, scan_plasmon_dualthicknesses, make_layered_tm_system, make_layered_θ_tm_system, make_d3_system, make_d4_system, make_d2_dualparameter_system, make_hc_system

function scann_singleparameter(system::Function, scanrange::Vector{<:Number}, postprocessing::Function=v -> abs(v[2, 1])^2)::Tuple{Vector{<:Number}, Vector}
    # This function provides a unified interface for scanning structures for different design parameters.
    # It akes a parameter "system", a function returning (for instance) a total scattering matrix, a scans it for all elements in the scanrange and a postprocessing expression.
    # Returns the scanrange and the output data.
    # The default postprocessing stage takes a scattering matrix and returns the power reflection coefficient.

    output = Vector(undef, length(scanrange))

    for i in 1:length(scanrange)
        instance = system(scanrange[i])
        output[i] = postprocessing(instance)
    end

    scanrange, output
end

function plasmon_minima(system::Function, λs::Vector{Float64}, predipps::Int64=1)::Tuple{Float64, Float64}
    # This function finds the plasmon minima, as well as the relevant frequency at which it appears.
    # It takes as its arguments a function returning a scattering matrix as a function of a wavelength, a range of wavelengths and a number of predipps.
    # The predipps number is the number of local maxima the function can encounter before searching for a global minima.
    # This is useful for situations in which the plasmon dip is weaker than the high-frequency (low wavelength) absorption of the material.
    # Note however that, if set too high, the algorithm might measure a local minima pas the plasmon dip instead.
    # The function returns the plasmon minima power reflection and the wavelength at which it was found.

    minλ = λs[1]
    minr = abs(system(minλ)[2, 1])^2
    initialrise = predipps
    thisr = 1

    for λ in λs
        lastr = thisr
        thisr = abs(system(λ)[2, 1])^2

        if (initialrise >= 1) && (thisr < lastr)
            initialrise -= 1
            minr = thisr
            minλ = λ
        end

        if thisr < minr
            minr = thisr
            minλ = λ
        end
    end

    minr, minλ
end

function plasmon_halfwidth(system::Function, λs::Vector{Float64}, predipps::Int64=1, widthratio::Float64=0.5)::Tuple{Float64, Float64, Float64, Float64, Float64}
    # This function finds the plasmon minima dip half peak width (half can be changed by widthratio parameter).
    # It opperates similarly to the plasmon_minima function.

    λleftmax = 0.0

    λleft = 0.0
    λright = 0.0
    λmin = 0.0

    minr = 1.0
    maxr = 0.0

    initialrise = predipps
    thisr = abs(system(λs[1])[2, 1])^2

    for λ in λs
        lastr = thisr
        thisr = abs(system(λ)[2, 1])^2

        if initialrise >= 1
            if thisr < lastr
                initialrise -= 1
                minr = thisr
                λmin = λ
            end

            if thisr > maxr
                maxr = thisr
               λleftmax = λ
            end
        end

        if thisr < minr
            minr = thisr
            λmin = λ
        end

    end

    rthreshold = minr + ((maxr - minr) * widthratio)

    for λ in λs
        if λ >= λleftmax
            thisr = abs(system(λ)[2, 1])^2

            if (thisr > rthreshold) && (λ <= λmin)
                λleft = λ
            end

            if (thisr < rthreshold) && (λ >= λmin)
                λright = λ
            end

            if (thisr > rthreshold) && (λ >= λmin)
                break
            end

        end
    end

    width = λright - λleft

    width, maxr - minr, λleft, λright, λmin
end

function scan_minima(xs::Vector{<:Number}, ys::Vector{<:Number})::Tuple{Number, Number}
    # An optimized function for finding the minimum of a curve.

    @assert(length(xs) == length(ys))
    @assert(length(xs) > 0)

    @inbounds begin
        xmin = xs[1]
        ymin = ys[1]

        for i in 1:length(xs)
            if ys[i] < ymin
                xmin = xs[i]
                ymin = ys[i]
            end
        end
    end

    xmin, ymin
end

function scan_maxima(xs::Vector{<:Number}, ys::Vector{<:Number})::Tuple{Number, Number}
    # An optimized function for finding the minimum of a curve.

    @assert(length(xs) == length(ys))
    @assert(length(xs) > 0)

    @inbounds begin
        xmax = xs[1]
        ymax = ys[1]

        for i in 1:length(xs)
            if ys[i] > ymax
                xmax = xs[i]
                ymax = ys[i]
            end
        end
    end

    xmax, ymax
end

function scan_plasmon_singleparameter(system::Function, scanrange::Vector{<:Number}, λs::Vector{Float64}, predipps::Int64=1)::Tuple{Vector{<:Number}, Vector{Float64}, Vector{Float64}, Vector{Float64}, Vector{Float64}}
    # A tool for running a single parameter plasmon scan.
    # Takes the following arguments:
    #   system: A function of the scanparameter returning a function on the form λ -> S.
    #   scanrange: A range of scan parameters.
    #   λs: A range of wavelengths.
    #   predipps: The number of local maxima in the reflection spectrum to scan through before finding the plasmon minima.
    # Returns the scanrange, the minimum reflection coefficients and the relevant wavelengths.
    # Functions similarly to the single thickness scann one-off written in demos/thickness.jl.

    plasmonfinder(subsystem) = plasmon_minima(subsystem, λs, predipps)
    plasmonwidthfinder(subsystem) = plasmon_halfwidth(subsystem, λs, predipps)

    _, plasmon_minima_output = scann_singleparameter(system, scanrange, plasmonfinder)
    _, plasmon_width_output = scann_singleparameter(system, scanrange, plasmonwidthfinder)

    scanrange, [x[1] for x in plasmon_minima_output], [x[2] for x in plasmon_minima_output], [z[1] for z in plasmon_width_output], [z[2] for z in plasmon_width_output]
end

function scan_plasmon_dualparameter(system::Function, parameter1s::Vector{<:Number}, parameter2s::Vector{<:Number}, λs::Vector{Float64}, predipps::Int64=1)::Tuple{Vector{<:Number}, Vector, Vector}
    # A tool for scanning for plasmon peaks across two parameters.
    # The function runs the single parameter scan on parameter 2 for each of parameter 1, it then plots the optimal parameter 2 as a function of parameter 1.

    minRP2s = Vector(undef, length(parameter1s))
    minλP2s = Vector(undef, length(parameter1s))

    for i in 1:length(parameter1s)
        p1 = parameter1s[i]

        _, reflections, wavelengths, _, _ = scan_plasmon_singleparameter(system(p1), parameter2s, λs, predipps)
        minRP2, _ = scan_minima(parameter2s, reflections)
        minλP2, _ = scan_minima(parameter2s, wavelengths) # TODO this might be a suboptimal target function, try to find something better based on the plateaus observed.

        minRP2s[i] = minRP2
        minλP2s[i] = minλP2
    end

    parameter1s, minRP2s, minλP2s
end

function scan_plasmon_dualthicknesses(system::Function, d1min, d1max, d1step, d2min, d2max, d2step, λs::Vector{Float64}, predipps::Int64=1, layertolerance::Float64=0.5)::Tuple{Vector{<:Number}, Vector, Vector, Vector, Vector}
    # A tool for scanning for plasmon peaks across two parameters.
    # The function runs the single parameter scan on parameter 2 for each of parameter 1, it then plots the optimal parameter 2 as a function of parameter 1.

    parameter1s = [x for x in d1min:d1step:d1max]

    minRP2s = Vector(undef, length(parameter1s))
    minλP2s = Vector(undef, length(parameter1s))
    minFWHMP2s = Vector(undef, length(parameter1s))
    maxΔRP2s = Vector(undef, length(parameter1s))

    for i in 1:length(parameter1s)
        p1 = parameter1s[i]

        parameter2s = [x for x in d2min:d2step:min(d2max, p1 * layertolerance)]

        _, reflections, wavelengths, peakwidths, peakdeltas = scan_plasmon_singleparameter(system(p1), parameter2s, λs, predipps)
        minRP2, _ = scan_minima(parameter2s, reflections)
        minλP2, _ = scan_minima(parameter2s, wavelengths) # TODO this might be a suboptimal target function, try to find something better based on the plataus observed.
        minFWHMP2, _ = scan_minima(parameter2s, peakwidths)
        maxΔRP2, _ = scan_maxima(parameter2s, peakdeltas)

        minRP2s[i] = minRP2
        minλP2s[i] = minλP2
        minFWHMP2s[i] = minFWHMP2
        maxΔRP2s[i] = maxΔRP2
    end

    parameter1s, minRP2s, minλP2s, minFWHMP2s, maxΔRP2s
end

function make_layered_tm_system(ns::Vector{Function}, ds::Vector{Float64}, θi::Float64)::Function
    # A tool for generating λ -> S functions for arbitrary layered systems.

    @assert(length(ns) == length(ds))

    function S(λ)
        θw = θi

        elements = Vector(undef, 2 * length(ns) - 2)

        for i in 1:(length(ns) - 1)
            bulk = FresnellSlab(ns[i](λ), 2*π/λ, ds[i], θw)
            _, border, θw = FresnellBoundrary(ns[i](λ), ns[i + 1](λ), θw)

            elements[2*i - 1] = bulk
            elements[2*i] = border
        end

        CascadeScattering(elements)
    end

    S
end

function make_layered_θ_tm_system(ns::Vector{Function}, ds::Vector{Float64}, λ::Float64)::Function
    # A tool for generating θ -> S functions for arbitrary layered systems.

    @assert(length(ns) == length(ds))

    function S(θi)
        θw = θi

        elements = Vector(undef, 2 * length(ns) - 2)

        for i in 1:(length(ns) - 1)
            bulk = FresnellSlab(ns[i](λ), 2*π/λ, ds[i], θw)
            _, border, θw = FresnellBoundrary(ns[i](λ), ns[i + 1](λ), θw)

            elements[2*i - 1] = bulk
            elements[2*i] = border
        end

        CascadeScattering(elements)
    end

    S
end

function make_d3_system(n1::Function, n2::Function, n3::Function, n4::Function, n5::Function, d1::Float64, d2::Float64, d4::Float64, d5::Float64, θi::Float64)::Function
    function tm_singleparameter_system(d3::Float64)::Function
        make_layered_tm_system([n1, n2, n3, n4, n5], [d1, d2, d3, d4, d5], θi)
    end

    tm_singleparameter_system
end
function make_d4_system(n1::Function, n2::Function, n3::Function, n4::Function, n5::Function, d1::Float64, d2::Float64, d3::Float64, d5::Float64, θi::Float64)::Function
    function tm_singleparameter_system(d4::Float64)::Function
        make_layered_tm_system([n1, n2, n3, n4, n5], [d1, d2, d3, d4, d5], θi)
    end

    tm_singleparameter_system
end
function make_d2_dualparameter_system(n1::Function, n2::Function, n3::Function, n4::Function, n5::Function, d1::Float64, d4::Float64, d5::Float64, θi::Float64)::Function
    function tm_dualparameter_system(d2::Float64)::Function
        make_d3_system(n1, n2, n3, n4, n5, d1, d2, d4, d5, θi)
    end

    tm_dualparameter_system
end
function make_hc_system(n1::Function, n2::Function, n3::Function, n4::Function, n5::Function, d1::Float64, d2::Float64, d3::Float64, d4::Float64, d5::Float64, θi::Float64, Pd_c::Float64)::Function
    function tm_singleparameter_system(H_c::Float64)::Function

        function modified_n3(λ::Float64)::Number
            √(h(λ, H_c, Pd_c) * n3(λ)^2)
        end

        make_layered_tm_system([n1, n2, modified_n3, n4, n5], [d1, d2, d3, d4, d5], θi)
    end

    tm_singleparameter_system
end

end # targetfigures


text:
#GracePlot demonstration 6: Building a template
#-------------------------------------------------------------------------------

using GracePlot
import GracePlot.Meter

#==Input data
===============================================================================#
x=collect(-10:0.1:10)

#==Configuration
===============================================================================#
pdefaults = defaults(linewidth=1, linestyle=:solid)
dataline = line(width=4)
frameline = line(width=2.5)
refsize = 20e-2 #Reference: 20cm canvas
(wpage, hpage) = (1.6, 1) #Relative dimensions
(nrows, ncols) = (2,2)
Δxlegend = .17

#Offsets from ideal graph boundary to allow sapce for tick labels, etc:
(Δxmin, Δxmax) = (.09, .02)
(Δymin, Δymax) = (.065, .055)

#==Position calculations
===============================================================================#
function graph_position(row, col) #0-based row/col
	const (w, h) = ((wpage-Δxlegend)/ncols, hpage/nrows) #Graph  w/h
	const (xstart, ystart) = (w*col, hpage-h*(row+1))

	return limits(
		xmin = xstart+Δxmin, xmax = xstart+w-Δxmax,
		ymin = ystart+Δymin, ymax = ystart+h-Δymax
	)
end

function legend_position(row, col)
	gpos = graph_position(row, col)
	return (gpos.xmax, gpos.ymax)
end

#==Generate plot
===============================================================================#
plot = GracePlot.new(pdefaults, fixedcanvas=true)
	set(plot, canvas(Meter(refsize*wpage), Meter(refsize*hpage))) #Force canvas size
	w = get(plot, :wview); h = get(plot, :hview)
	@show w, h #Confirm canvas size

#Add subplots:
for gidx in 1:4
g = add(plot, subtitle="Subplot $gidx")
	row = div(gidx-1, ncols) #0-based
	col = mod(gidx-1, ncols) #0-based
	set(g, frameline=frameline, view=graph_position(row, col))
	set(g, legend(display=false))

	for i in 1:10
		y = i*x/10
		ds = add(g, x, y, dataline, id="Dataset $i")
	end

	autofit(g)
end

#Position legend:
g = graph(plot, 1) #Get second graph
set(g, legend(loc=legend_position(0,1), loctype=:view, charsize=.75))
set(g, legend(display=true))

#Finalize:
redraw(plot)

#Last line

ext: jl


text:
#///////////////////////////////////////
#// File Name: cost_test.jl
#// Author: Haruki Nishimura (hnishimura@stanford.edu)
#// Date Created: 2021/01/15
#// Description: Test code for src/cost.jl
#///////////////////////////////////////

using ForwardDiff
using LinearAlgebra
using RobotOS

@testset "Cost Test" begin
target_point = [0.0, 0.0];
wp1 = WayPoint2D(target_point, Time(0.0));
wp2 = WayPoint2D(target_point, Time(1.0));
target_trajectory = Trajectory2D([wp1, wp2]);
Cep = Matrix(1.0I, 2, 2);
Cu = Matrix(1.0I, 2, 2);
β_pos = 0.6;
β_col = 0.4;
α_col = 100.0;
λ_col = 1.0;
σ_risk = 1.0;
param = CostParameter(Cep, Cu, β_pos, α_col, β_col, λ_col, σ_risk);

p1 = [3., 4.];
p2 = [7., 8.];
v1 = [1., 1.];
u1 = [2., 2.];
e1 = RobotState([p1; v1]);

# Instantaneous Cost
# # Position
cip = instant_position_cost(e1, target_trajectory, param);
@test cip ≈ 1/2*(p1 - target_point)'*Cep*(p1 - target_point);
# # Control
ciu = instant_control_cost(u1, param);
@test ciu ≈ 1/2*u1'*Cu*u1;
# # Collision
cic = instant_collision_cost(e1, p2, param);
@test cic ≈ α_col*exp(-1/(2*λ_col)*norm(p1 - p2)^2);

# Terminal Cost
# # Position
ctp = terminal_position_cost(e1, target_trajectory, param);
@test ctp ≈ param.β_pos*cip;
# # Collision
ctc = terminal_collision_cost(e1, p2, param);
@test ctc ≈ param.β_col*cic;

# Instantaneous Cost Gradient
# # Position
function instant_position_cost_vector(evec::Vector{<:Real},
                                      t::Time,
                                      target_trajectory::Trajectory2D,
                                      param::CostParameter)
    ep_target = get_position(target_trajectory, t)
    ep = evec[1:2];
    return 1/2*(ep - ep_target)'*param.Cep*(ep - ep_target);
end
function test_instant_position_cost_gradient(e_state::RobotState,
                                             target_trajectory::Trajectory2D,
                                             param::CostParameter)
    return ForwardDiff.gradient(x -> instant_position_cost_vector(x, e1.t,
                                target_trajectory, param), e1.x)
end
gip = instant_position_cost_gradient(e1, target_trajectory, param);
@test gip ≈ test_instant_position_cost_gradient(e1, target_trajectory, param);
# # Collision
function instant_collision_cost_vector(evec::Vector{<:Real},
                                       apvec::Vector{<:Real},
                                       param::CostParameter)
    ep = evec[1:2]
    return param.α_col*exp(-1/(2*param.λ_col)*norm(ep - apvec)^2);
end
function test_instant_collision_cost_gradient(e_state::RobotState,
                                              apvec::Vector{<:Real},
                                              param::CostParameter)
    return ForwardDiff.gradient(x -> instant_collision_cost_vector(x,
                                apvec, param), e_state.x)
end
gic = instant_collision_cost_gradient(e1, p2, param);
@test gic ≈ test_instant_collision_cost_gradient(e1, p2, param);

# Terminal Cost Gradient
# # Position
function test_terminal_position_cost_gradient(e_state::RobotState,
                                             target_trajectory::Trajectory2D,
                                             param::CostParameter)
    return ForwardDiff.gradient(x -> param.β_pos*instant_position_cost_vector(x, e1.t,
                                                 target_trajectory, param), e1.x)
end
gtp = terminal_position_cost_gradient(e1, target_trajectory, param);
@test gtp ≈ test_terminal_position_cost_gradient(e1, target_trajectory, param);
# # Collision
function test_terminal_collision_cost_gradient(e_state::RobotState,
                                               apvec::Vector{<:Real},
                                               param::CostParameter)
    return ForwardDiff.gradient(x -> param.β_col*instant_collision_cost_vector(x,
                                                 apvec, param), e_state.x)
end
gtc = terminal_collision_cost_gradient(e1, p2, param);
@test gtc ≈ test_terminal_collision_cost_gradient(e1, p2, param);

end


text:
module TestIteration
	using Base.Test
	using DataArrays
	using DataFrames

	dv = @data([1, 2, NA])
	dm = DataArray([1 2; 3 4])
	dt = DataArray(zeros(2, 2, 2))

	df = DataFrame(A = 1:2, B = 2:3)

	for el in dv
	    @assert ndims(el) == 0
	end

	for el in dm
	    @assert ndims(el) == 0
	end

	for el in dt
	    @assert ndims(el) == 0
	end

	for row in EachRow(df)
	    @assert isa(row, DataFrame)
	end

	for col in EachCol(df)
	    @assert isa(col, AbstractDataVector)
	end

	@assert isequal(map(x -> minimum(array(x)), EachRow(df)), {1,2})
	@assert isequal(map(minimum, EachCol(df)), DataFrame(A = [1], B = [2]))

	# @test_fail for x in df; end # Raises an error
end


text:
@doc raw"""
```
DiscreteShapeBound(iterations::Int = 2000, smoothness::Int = 1000, albedo::Real = Inf, illumination_direction::Vector{T} where T <: Real = [Inf, Inf, Inf])
```
Same as [`DiscreteShape`](@ref)
except it has its initial conditions bound by the image as per bellow where
``E`` is the brightness of the image:
```math
\begin{gathered}
Z_{i,j}=\begin{cases}-100E_{i,j} &\text{if } E_{i,j}>0.75\\0 &\text{otherwise}\\
\end{cases}\\\\p,q=\nabla E
\end{gathered}
```
# Example
Compute the heightmap for a synthetic image generated by `generate_surface`.
```julia
using Images, Makie, ShapeFromShading

#generate synthetic image
img = generate_surface(SynthSphere(50), 1.0, [0.2,0,0.9])

#calculate the heightmap (using 500 iterations)
discreteShapeBound = DiscreteShapeBound(iterations = 500, 1.0, [0.2,0,0.9])
Z,p,q = discreteShapeBound(img)

#normalize to maximum of 1 (not necessary but makes displaying easier)
Z = Z./maximum(Z)

#display using Makie (Note: Makie can often take several minutes first time)
r = 0.0:0.1:2
surface(r, r, Z)
```
# Reference
1. S. Elhabian, "Hands on Shape from Shading", Computer Vision and Image Processing, 2008.
"""
function (algorithm::DiscreteShapeBound)(img::AbstractArray)
    ρ = algorithm.albedo
    I = algorithm.illumination_direction
    smoothness = algorithm.smoothness
    iterations = algorithm.iterations

    if ρ == Inf || I == [Inf, Inf, Inf]
        ρ,I,σ,τ = estimate_img_properties(img)
    end
    E = Array{Float64}(img)

    #initialize variables
    p = zeros(Complex{Float64},axes(E))
    q = zeros(Complex{Float64},axes(E))
    R = zeros(Complex{Float64},axes(E))
    Z = zeros(axes(E))

    #apply bounding conditions
    for i in CartesianIndices(E)
        if E[i] > 0.75
            Z[i] = -100*E[i]
        else
            Z[i] = 0.0
        end
    end
    q, p = Array{Complex{Float64}}.(imgradients(Z, KernelFactors.sobel, "replicate"))
    return solve_EulerLagrange(ρ, I, iterations, p, q, R, smoothness, E, Z)
end


text:
module RingSignature

#=
Ring signature scheme:

- Uniform and has no center
- Set-up free, no prearranged groups of users
- Signer-ambiguous

# References

- Rivest R.L., Shamir A., Tauman Y. (2001) How to Leak a Secret.
=#

include("util.jl")
include("rsa.jl")

using SHA

using RingSignature.Util

export RingContext, ring_sign, ring_verify

#=
The example code from Wikipedia uses a hash function in place of a symmetric encryption function.

Now that the hash is not reversible, you can't really "solve" the ring equation -- it's more of a
"draw the circles around the arrow after you shoot it" approach. That is, you compute `E_k(u)` for
a randomly genarated `u`, tweak the XOR term, and finally "close the ring" at the signer.

See "Rivest's ring signatures with hashes instead of symmetric encryption"
(<https://crypto.stackexchange.com/q/52608>) for a detailed analysis.
=#

struct RingContext
    n::Int
    upper::BigInt
    rsa_list::Vector{RSA}

    function RingContext(n::Int, key_size::Int)
        upper = big(2) ^ (key_size - 1)
        rsa_list = [RSA(key_size) for _ in 1:n]
        new(n, upper, rsa_list)
    end
end

# Computes `E_k(m)`.
function encrypt(ring_ctx::RingContext, k::Vector{UInt8}, m::BigInt)::BigInt
    m = convert(Vector{UInt8}, m)
    convert(BigInt, sha1(vcat(m, k)))  # Manually digest twice to eliminate concatenation
end

# Computes `y = g_i(x)`.
#
# The original paper uses an extended trap-door function. However, as the hash function is no longer
# strictly equiprobable, it seems OK to use the vanilla trap-door, right?
function trapdoor(ring_ctx::RingContext, i::Int, x::BigInt)::BigInt
    rsa_encrypt(ring_ctx.rsa_list[i], x)
end

# Computes `x = g_s^{-1}(y)`. Only can be called by the signer.
function backdoor(ring_ctx::RingContext, signer::Int, y::BigInt)::BigInt
    rsa_decrypt(ring_ctx.rsa_list[signer], y)
end

function ring_sign(
    ring_ctx::RingContext,
    message::Vector{UInt8},
    signer::Int,
)::Tuple{BigInt, Vector{BigInt}}
    k = sha1(message)

    n = ring_ctx.n
    x = zeros(BigInt, n)

    u = rand(big(0):ring_ctx.upper)
    iv = v = encrypt(ring_ctx, k, u)
    for i in vcat((signer + 1):n, 1:(signer - 1))
        x[i] = rand(big(0):ring_ctx.upper)
        y = trapdoor(ring_ctx, i, x[i])
        v = encrypt(ring_ctx, k, v ⊻ y)
        if i == n
            iv = v
        end
    end
    x[signer] = backdoor(ring_ctx, signer, v ⊻ u)
    (iv, x)
end

function ring_verify(
    ring_ctx::RingContext,
    message::Vector{UInt8},
    sig::Tuple{BigInt, Vector{BigInt}}
)::Bool
    k = sha1(message)
    (v, x) = sig
    n = ring_ctx.n
    y = [trapdoor(ring_ctx, i, x[i]) for i in 1:n]
    v_expected = v
    for i in 1:n
        v = encrypt(ring_ctx, k, v ⊻ y[i])
    end
    v == v_expected
end

end # module

max_stars_repo_name: gyk/TrivialSolutions

max_issues_repo_name: gyk/TrivialSolutions

max_forks_repo_name: gyk/TrivialSolutions


text:
# this is an example script to run the package...
using Pkg
Pkg.activate(".")
using JuliaTutor
menu()


text: include("Traversal.jl")


text:
module Model

using Random
using SparseArrays

struct Factorization{T <: AbstractFloat}
    user_embeddings:: Matrix{T}
    item_embeddings:: Matrix{T}
    item_bias:: Vector{T}

    function Factorization{T}(n_users::Int,
                              n_items::Int,
                              n_components::Int) where {T<:AbstractFloat}
        return new{T}(
            zeros(T, n_components, n_users),
            zeros(T, n_components, n_items),
            zeros(T, n_items))
    end
end

n_users(model::Factorization) = size(model.user_embeddings, 2)
n_items(model::Factorization) = size(model.item_embeddings, 2)
n_components(model::Factorization) = size(model.user_embeddings, 1)

function(model::Factorization)(user::Int, item::Int)
    s = model.item_bias[item]
    for i=1:n_components(model)
        s += model.item_embeddings[i, item] * model.user_embeddings[i, user]
    end
    return s
end

warp_constant(model::Factorization{T}, n) where {T<:AbstractFloat} =
    convert(T, log((n_items(model) - 1) / n))

function warp_embeddings_grad(model::Factorization,
                              user::Int,
                              item_p::Int,
                              item_n::Int,
                              i::Int)
    return (model.item_embeddings[i, item_n] - model.item_embeddings[i, item_p],
            -model.user_embeddings[i, user],
            +model.user_embeddings[i, user])

end

abstract type Optimizer{T, N} end

struct SGD{T <: AbstractFloat, N} <: Optimizer{T, N}
    θ::Array{T, N}
    ϵ::T
    SGD(θ::Array{T}, ϵ::AbstractFloat) where {T <: AbstractFloat, N} =
        new{T, N}(θ, convert(T, ϵ))
end

function step!(optimizer::SGD{T, N},
               g::T,
               i::Vararg{Int, N}) where {T <: AbstractFloat, N}
    optimizer.θ[i...] -= optimizer.ϵ * g
end

struct Adadelta{T <: AbstractFloat, N} <: Optimizer{T, N}
    θ::Array{T, N}
    Eg²::Array{T, N}
    EΔ²::Array{T, N}
    ϵ::T
    η::T
    function Adadelta(θ::Array{T, N},
                      ϵ::AbstractFloat,
                      η::AbstractFloat) where {T <: AbstractFloat, N}
        return new{T, N}(θ,
                         zero(θ),
                         zero(θ),
                         convert(T, ϵ),
                         convert(T, η))
    end
end

function step!(optimizer::Adadelta{T, N},
               g::T,
               i::Vararg{Int, N}) where {T<:AbstractFloat, N}
    EΔ²i = optimizer.EΔ²[i...]
    Eg²i = optimizer.Eg²[i...]
    ϵ = optimizer.ϵ
    η = optimizer.η

    Eg²i = η*Eg²i + (1 - η)*g^2
    Δi = -sqrt((EΔ²i + ϵ)/(Eg²i + ϵ))*g
    EΔ²i = η*EΔ²i + (1 - η)*Δi^2

    optimizer.θ[i...] += Δi
    optimizer.EΔ²[i...] = EΔ²i
    optimizer.Eg²[i...] = Eg²i
end

struct Adagrad{T <: AbstractFloat, N} <: Optimizer{T, N}
    θ::Array{T, N}
    G::Array{T, N}
    ϵ::T
    η::T
    function Adagrad(θ::Array{T},
                     ϵ::AbstractFloat,
                     η::AbstractFloat) where {T <: AbstractFloat, N}
        return new{T, N}(θ,
                         zero(θ),
                         convert(T, ϵ),
                         convert(T, η))
    end
end

function step!(optimizer::Adagrad{T, N},
               g::T,
               i::Vararg{Int, N}) where {T<:AbstractFloat, N}
    Gi = optimizer.G[i...]
    ϵ = optimizer.ϵ
    η = optimizer.η

    Gi += g^2
    Δi = -η/sqrt(Gi + ϵ)*g

    optimizer.θ[i...] += Δi
    optimizer.G[i...] = Gi
end

abstract type FactorizationOptimizer{T <: AbstractFloat} end

macro optimizertype(OptimizerType, params...)
    structname = Symbol("Factorization", OptimizerType)
    return quote
        struct $structname{T} <: FactorizationOptimizer{T}
            user_embeddings::$OptimizerType{T, 2}
            item_embeddings::$OptimizerType{T, 2}
            item_bias::$OptimizerType{T, 1}

            function $structname(model::Factorization{T},
                                 $(params...)) where {T}
                return new{T}(
                    $OptimizerType(model.user_embeddings, $(params...)),
                    $OptimizerType(model.item_embeddings, $(params...)),
                    $OptimizerType(model.item_bias, $(params...)))
            end
        end
    end
end

@optimizertype(SGD, ϵ)
@optimizertype(Adadelta, ϵ, η)
@optimizertype(Adagrad, ϵ, η)

function warpstep!(model::Factorization{T},
                   optimizer::FactorizationOptimizer{T},
                   user::Int,
                   item_p::Int,
                   item_n::Int,
                   n::Int) where {T}
    C = warp_constant(model, n)
    for i=1:n_components(model)
        gu, gp, gn = warp_embeddings_grad(model, user, item_p, item_n, i)
        step!(optimizer.user_embeddings, C*gu, i, user)
        step!(optimizer.item_embeddings, C*gp, i, item_p)
        step!(optimizer.item_embeddings, C*gn, i, item_n)
    end
    step!(optimizer.item_bias, convert(T, -1), item_p)
    step!(optimizer.item_bias, convert(T, +1), item_n)
end

struct Interactions{S}
    sparse::SparseMatrixCSC{S, Int}
    users::Vector{Int}

    function Interactions(interactions::SparseMatrixCSC{S, Int}) where {S}
        n_interactions = length(interactions.nzval)
        users = Array{Int}(undef, length(interactions.nzval))
        j = 1
        for i=1:length(interactions.colptr)-1
            for j=interactions.colptr[i]:(interactions.colptr[i+1] - 1)
                users[j] = i
            end
        end
        return new{S}(interactions, users)
    end
end

function init!(interactions::Interactions, model::Factorization)
    interactions = interactions.sparse

    fill!(model.item_bias, 0)
    for i=1:length(interactions.nzval)
        model.item_bias[interactions.rowval[i]] += 1
    end
    model.item_bias ./= sum(model.item_bias)

    scale = 1.0 / n_components(model)
    rand!(model.user_embeddings)
    rand!(model.item_embeddings)
    model.user_embeddings .-= 0.5
    model.item_embeddings .-= 0.5
    model.user_embeddings .*= scale
    model.item_embeddings .*= scale
end

function is_positive(interactions::SparseMatrixCSC{S, Int},
                     user::Int,
                     item::Int) where {S}
    for i=interactions.colptr[user]:interactions.colptr[user+1] - 1
        if item == interactions.rowval[i]
            return true
        end
    end
    return false
end

function find_inversion(model::Factorization,
                        user::Int,
                        item::Int,
                        interactions::SparseMatrixCSC{S, Int},
                        max_sample::Int) where {S}
    n = n_items(model)
    s = model(user, item)
    for i=1:max_sample
        j = abs(rand(typeof(item)) % n) + 1
        t = model(user, j)
        if t > (s - 1) && !is_positive(interactions, user, j)
            return (i, j)
        end
    end
    return (max_sample, 0)
end

function epoch!(model::Factorization, optimizer::FactorizationOptimizer, interactions::Interactions, max_sample::Int)
    max_sample = min(n_items(model) - 1, max_sample)
    n_interactions = length(interactions.sparse.nzval)
    @Threads.threads for i in randperm(n_interactions)
        user = interactions.users[i]
        item_p = interactions.sparse.rowval[i]
        n, item_n = find_inversion(model, user, item_p, interactions.sparse, max_sample)
        if item_n > 0
            warpstep!(model, optimizer, user, item_p, item_n, n)
        end
    end
end

export Factorization
export n_users
export n_items
export n_components
export init!
export zeros

export FactorizationSGD
export FactorizationAdadelta
export FactorizationAdagrad

export Interactions
export is_positive
export epoch!

end

max_stars_repo_name: plafl/julia-cython-comp

max_issues_repo_name: plafl/julia-cython-comp

max_forks_repo_name: plafl/julia-cython-comp


text:
using SortFilters
using Test
using Random
using Statistics

len = 100
Random.seed!(1)
T = Int
x = rand(T(-10):T(10), len)
window = 9

mx_ref = similar(x)
mn_ref = similar(x)
med_ref = similar(x, Float64)
p20_ref = similar(x, Float64)
p25_p75_ref = similar(x, NTuple{2,Float64})

for i = 1:len
    chunk = x[max(1, i-window+1) : i]
    mx_ref[i] = maximum(chunk)
    mn_ref[i] = minimum(chunk)
    med_ref[i] = median(chunk)
    p20_ref[i] = quantile(chunk, 0.20)
    p25_p75_ref[i] = quantile(chunk, (0.25, 0.75))
end

@testset "reset filter state and run" begin

f = MovSortFilter{T}(window)
mx = run(f, x, 1.0)
reset!(f)
mx_ = run(f, x, 100)

reset!(f)
mn = run(f, x, 0.0)
reset!(f)
mn_ = run(f, x, 0)

reset!(f)
med = run(f, x, 0.5)
reset!(f)
med_ = run(f, x, 50)

reset!(f)
p25_p75 = run(f, x, (0.25, 0.75))
reset!(f)
p25_p75_ = run(f, x, (25, 75))

@test mx == mx_ == mx_ref
@test mn == mn_ == mn_ref
@test med[window:end] == med_[window:end] == med_ref[window:end]
@test p25_p75[window:end] == p25_p75_[window:end] == p25_p75_ref[window:end]

reset!(f)
p20 = run(f, x, 0.20)
@test p20[window:end] == p20_ref[window:end] # should match nearest index quantiles

end

@testset "maintain filter state and run!" begin

mx 		 = similar(mx_ref, Float64)
mx_ 	 = similar(mx_ref)
mn 		 = similar(mn_ref, Float64)
mn_ 	 = similar(mn_ref)
med 	 = similar(med_ref, Float64)
med_ 	 = similar(med_ref)
p25_p75  = similar(p25_p75_ref, NTuple{2,Float64})
p25_p75_ = similar(p25_p75_ref)
p20      = similar(p20_ref, Float64)

part1 = x->view(x, 1 : len÷2)
part2 = x->view(x, len÷2+1 : len)
part1(x)

f = MovSortFilter{T}(window)
run!(part1(mx), f, part1(x), 1.0)
run!(part2(mx), f, part2(x), 1.0)
reset!(f)
run!(part1(mx_), f, part1(x), 100)
run!(part2(mx_), f, part2(x), 100)

reset!(f)
run!(part1(mn), f, part1(x), 0.0)
run!(part2(mn), f, part2(x), 0.0)
reset!(f)
run!(part1(mn_), f, part1(x), 0)
run!(part2(mn_), f, part2(x), 0)

reset!(f)
run!(part1(med), f, part1(x), 0.5)
run!(part2(med), f, part2(x), 0.5)
reset!(f)
run!(part1(med_), f, part1(x), 50)
run!(part2(med_), f, part2(x), 50)

reset!(f)
run!(part1(p25_p75), f, part1(x), (0.25, 0.75))
run!(part2(p25_p75), f, part2(x), (0.25, 0.75))
reset!(f)
run!(part1(p25_p75_), f, part1(x), (25, 75))
run!(part2(p25_p75_), f, part2(x), (25, 75))

@test mx == mx_ == mx_ref
@test mn == mn_ == mn_ref
@test med[window:end] == med_[window:end] == med_ref[window:end]
@test p25_p75[window:end] == p25_p75_[window:end] == p25_p75_ref[window:end]

reset!(f)
run!(part1(p20), f, part1(x), 0.20)
run!(part2(p20), f, part2(x), 0.20)
@test p20[window:end] == p20_ref[window:end] # should match nearest index quantiles

end


text:
#TODO: use vendored FunctionWrappers for now: https://github.com/JuliaLang/julia/issues/40187
include("../../lib/FunctionWrappers/FunctionWrappers.jl")
using .FunctionWrappers: FunctionWrapper
struct Call{V,F<:FunctionWrapper} <: State{V}
    value::F
end

Call(; unit, _value, _type, _calltype, _...) = begin
    V = valuetype(_type, value(unit))
    F = _calltype
    Call{V,F}(_value)
end

(c::Call)(a...) = value(c)(a...)

constructortags(::Val{:Call}) = (:unit,)

#HACK: showing s.value could trigger StackOverflowError
Base.show(io::IO, s::Call) = print(io, "<call>")

gencallargtype(t) = isnothing(t) ? :Float64 : esc(t)

updatetags!(d, ::Val{:Call}; kwargs, _...) = begin
    #FIXME: lower duplicate efforts in genvartype()
    N = d[:_type]
    U = d[:unit]
    V = @q $C.valuetype($N, $U)
    extract(a) = let k, t, u
        @capture(a, k_::t_(u_) | k_::t_ | k_(u_) | k_)
        @q $C.valuetype($(gencallargtype(t)), $u)
    end
    F = @q FunctionWrapper{$V, Tuple{$(extract.(kwargs)...)}}
    d[:_calltype] = F
end

genvartype(v::VarInfo, ::Val{:Call}; V, _...) = begin
    extract(a) = let k, t, u
        @capture(a, k_::t_(u_) | k_::t_ | k_(u_) | k_)
        @q $C.valuetype($(gencallargtype(t)), $u)
    end
    F = @q FunctionWrapper{$V, Tuple{$(extract.(v.kwargs)...)}}
    @q Call{$V,$F}
end

gendefault(v::VarInfo, ::Val{:Call}) = genfunc(v)

genupdate(v::VarInfo, ::Val{:Call}, ::MainStep; kw...) = genvalue(v)


text:
using Compat.Test

@testset "Neural Network" begin
    include("NN/NN.jl")
end

@testset "Tensor Network" begin
    include("TN/TN.jl")
end

@testset "Quantum Circuit" begin
    include("QuCircuit/QuCircuit.jl")
end


text:
"""
    NetworkData{S<:Real, T<:NetworkParams}
"""
mutable struct NetworkData{S<:Real, T<:NetworkParams} <: AbstractData
    nPort::Integer
    nPoint::Integer
    ports::Vector{Port}
    frequency::Vector{S}
    params::Vector{T}
    function NetworkData(ports::Vector{Port}, frequency::AbstractVector{S},
        params::Vector{T}) where {S<:Real, T<:NetworkParams}
        nPoint = length(frequency)
        nPort = length(ports)
        if (length(params) != nPoint) | (length(frequency) != nPoint)
            error("NetworkData Error: the number of data points doesn't match with number of frequency points")
        end
        if ~all(n -> (nPort == params[n].nPort), 1:nPoint)
            error("NetworkData Error: the number of ports in params doesn't match with `nPort`")
        end

        if T<:TwoPortParams
            if is_uniform(ports) == false
                error("TwoPortParams Error: two-port parameters are defined only for uniform port impedances")
            end
        end
        new{S, T}(nPort, nPoint, ports, frequency, params)
    end
end

is_uniform(ntwk::NetworkData{S, T}) where {S<:Real, T<:NetworkParams} =
    is_uniform(ntwk.ports)
is_two_port(ntwk::NetworkData{S, T}) where {S<:Real, T<:NetworkParams} =
    (ntwk.nPort == 2)

check_frequency_identical(ntwkA::NetworkData, ntwkB::NetworkData,
    ntwkC::NetworkData...) = begin
    ntwks = [ntwkA, ntwkB, ntwkC...]
    freqs = [ntwk.frequency for ntwk in ntwks]
    # foldl doesn't seem to work reliably here
    # foldl(==, freqs)
    all(n -> (freqs[1] == freqs[n]), 1:length(freqs))
    end

check_port_impedance_identical(ntwkA::NetworkData{S1, T1}, k,
    ntwkB::NetworkData{S2, T2}, l) where {S1<:Real, S2<:Real, T1<:NetworkParams,
    T2<:NetworkParams} = (ntwkA.ports[k].Z0 == ntwkB.ports[k].Z0)

"""
    NetworkData(frequency, params; port_impedance=50.0)
Convenience constructor for creating a `NetworkData` object of uniform port
impedances.
"""
NetworkData(frequency::AbstractVector{S}, params::Vector{T};
    Z0=50.0) where {S<:Real, T<:NetworkParams} =
    NetworkData([Port(Z0) for n in 1:params[1].nPort], frequency, params)

"""
    show(io::IO, D::NetworkData{T}) where {T<:NetworkParams}
Pretty-printing of `NetworkData`.
"""
function show(io::IO, D::NetworkData{S, T}) where {S<:Real, T<:NetworkParams}
    write(io, "$(D.nPort)-port $(typeof(D)):\n")
    write(io, "  # datapoints = $(D.nPoint)\n")
    write(io, "  Port Informaton:\n")
    for (n, p) in enumerate(D.ports)
        write(io, "    Port $(n) → (index = $(p.index), Z0 = Impedance($(p.Z0.data)))\n")
    end
end

impedances(D::NetworkData) = impedances(D.ports)

"""
`getindex` methods for `NetworkData{T<:NetworkParams}`. The first input argument
in the square bracket must be a 2-integer tuple (i, j) denoting which element of
parameter to get. The second input argument is either a scalar integer, Range,
integer vector, mask, or colon(:) that determines which datapoint to choose.
i.e., `D[(i, j), n]` returns the n-th data of parameter Tᵢⱼ
where T∈{S,Y,Z,ABCD,...}.
"""
getindex(D::NetworkData{S, T}, I1::Tuple{Integer, Integer},
    I2::Integer) where {S<:Real, T<:NetworkParams} = D.params[I2].data[I1...]
getindex(D::NetworkData{S, T}, I1::Tuple{Integer, Integer},
    I2::Range) where {S<:Real, T<:NetworkParams} = [D.params[n].data[I1...] for n in I2]
getindex(D::NetworkData{S, T}, I1::Tuple{Integer, Integer},
    I2::Vector) where {S<:Real, T<:NetworkParams} =
    [D.params[n].data[I1...] for n in I2]
getindex(D::NetworkData{S, T}, I1::Tuple{Integer, Integer},
    I2::Vector{Bool}) where {S<:Real, T<:NetworkParams} =
    (length(D.params) == length(I2))?
    [D.params[n].data[I1...] for n in Base.LogicalIndex(I2)]:
    error("Length of the mask different from lenghth of the array attemped to access")
getindex(D::NetworkData{S, T}, I1::Tuple{Integer, Integer},
    I2::BitArray{1}) where {S<:Real, T<:NetworkParams} =
    getindex(D, I1, collect(I2))
getindex(D::NetworkData{S, T}, I1::Tuple{Integer, Integer},
    I2::Colon) where {S<:Real, T<:NetworkParams} = getindex(D, I1, 1:length(D.params))
getindex(D::NetworkData{S, T}, I1::Tuple{Integer, Integer}) where {S<:Real, T<:NetworkParams} =
    getindex(D, I1, :)
# setindex!: TODO?
#


text:
using JLisp
using Test

macro test_sexp(ex, expected)
    quote
        let ex = JLisp.sexp($(QuoteNode(ex)); flatten=true)
            @test ex == $expected
        end
    end
end

@testset "jlisp" begin
    @testset "S-expression parsing" begin
        @test_sexp {max 1 2} Expr(:call, :max, 1, 2)
        @test_sexp {+ 1 2} Expr(:call, :+, 1, 2)
        @test_sexp {:> 1 2} Expr(:call, :>, 1, 2)
        @test_sexp {:def x 1} Expr(:(=), :x, 1)
        @test_sexp {:const x 1} Expr(:const, Expr(:(=), :x, 1))
        @test_sexp {:if true 1 2} Expr(:if, :true, 1, 2)
        @test_sexp {:for {x xs} x} Expr(:for, Expr(:block, Expr(:(=), :x, :xs)), :x)
        @test_sexp {:for {{x xs}} x} Expr(:for, Expr(:block, Expr(:(=), :x, :xs)), :x)
        @test_sexp {:while true fun} Expr(:while, :true, :fun)
        @test_sexp {:module Foo 1} Expr(:module, true, :Foo, Expr(:block, 1))
        @test_sexp {:baremodule Foo 1} Expr(:module, false, :Foo, Expr(:block, 1))
        @test_sexp {:quote 1 2 3} Expr(:quote, Expr(:block, 1, 2, 3))
        @test_sexp {:begin 1 2 3} Expr(:block, 1, 2, 3)
        @test_sexp {:abstract Foo} Expr(:abstract, :Foo)
        @test_sexp {:abstract Foo <: Real} Expr(:abstract, Expr(:<:, :Foo, :Real))
        @test_sexp {:primitive Foo 8} Expr(:primitive, :Foo, 8)
        @test_sexp {:primitive Foo <: Real 8} Expr(:primitive, Expr(:<:, :Foo, :Real), 8)
        @test_sexp {:mutable Foo} Expr(:struct, true, :Foo, Expr(:block))
        @test_sexp {:return} Expr(:return, nothing)
        @test_sexp {:return 1} Expr(:return, 1)
        @test_sexp {:|| foo bar} Expr(:||, :foo, :bar)

        @test_sexp {:function foo {x :kw y=1 z::Bool} x} Expr(
            :function,
            Expr(
                :call,
                :foo,
                Expr(:parameters, Expr(:kw, :y, 1), Expr(:(::), :z, :Bool)),
                :x,
            ),
            Expr(:block, :x),
        )

        @test_sexp {:macro foo {x} x} Expr(
            :macro,
            Expr(:call, :foo, :x),
            Expr(:block, :x),
        )

        @test_sexp {:struct Foo a b::Int} Expr(
            :struct,
            false,
            :Foo,
            Expr(:block, :a, Expr(:(::), :b, :Int)),
        )

        @test_sexp {:let {{x 1} {y 2}} 1} Expr(
            :let,
            Expr(:block, Expr(:(=), :x, 1), Expr(:(=), :y, 2)),
            Expr(:block, 1),
        )

        @test_sexp {:cond {foo bar} {{foo bar} baz}} Expr(
            :if,
            :foo,
            Expr(:block, :bar),
            Expr(:elseif, Expr(:call, :foo, :bar), Expr(:block, :baz)),
        )

        @test_sexp {:try {foo} {:catch ex {rethrow}} {:finally {bar}}} Expr(
            :try,
            Expr(:block, Expr(:call, :foo)),
            :ex,
            Expr(:block, Expr(:call, :rethrow)),
            Expr(:block, Expr(:call, :bar)),
        )

        @test_sexp {:try {foo} {:catch {rethrow}} {:finally {bar}}} Expr(
            :try,
            Expr(:block, Expr(:call, :foo)),
            false,
            Expr(:block, Expr(:call, :rethrow)),
            Expr(:block, Expr(:call, :bar)),
        )

        ex = JLisp.sexp(:({@show {max 1 2}}); flatten=true)
        @test ex.head === :macrocall && first(ex.args) === Symbol("@show")
        @test last(ex.args) == Expr(:call, :max, 1, 2)
    end

    @testset "String parsing" begin
        code = """
            {:let {x 0}
                {:while {:< x 10}
                    {:+= x 1}}
                {:if {:!= x 10}
                    {error "Something is broken!"}}}
            """
        ex = JLisp.parse(code)
        expected = Expr(
            :block,
            Expr(
                :let,
                Expr(:block, Expr(:(=), :x, 0)),
                Expr(
                    :block,
                    Expr(
                        :while,
                        Expr(:call, :<, :x, 10),
                        Expr(:+=, :x, 1)
                    ),
                    Expr(
                        :if,
                        Expr(:call, :!=, :x, 10),
                        Expr(:call, :error, "Something is broken!"),
                    ),
                ),
            ),
        )
        @test ex == expected
    end
end

max_stars_repo_name: christopher-dG/jlisp

max_issues_repo_name: christopher-dG/jlisp

max_forks_repo_name: christopher-dG/jlisp


text:
"""
```
solve(m::AbstractDSGEModel)
```

Driver to compute the model solution and augment transition matrices.

### Inputs

- `m`: the model object

## Keyword Arguments

- `regime_switching::Bool`: true if the state space system features regime switching
- `regimes::Union{Int, Vector{Int}, UnitRange{Int}}`: specifies the specific regime to solve for.

### Outputs
 - TTT, RRR, and CCC matrices of the state transition equation:
```
S_t = TTT*S_{t-1} + RRR*ϵ_t + CCC
```
"""
function solve(m::AbstractDSGEModel{T}; regime_switching::Bool = false,
               gensys_regimes::Vector{UnitRange{Int64}} = UnitRange{Int64}[1:1],
               gensys2_regimes::Vector{UnitRange{Int64}} = Vector{UnitRange{Int}}(undef, 0),
               regimes::Vector{Int} = Int[1],
               verbose::Symbol = :high) where {T <: Real}

    uncertain_altpolicy = haskey(get_settings(m), :uncertain_altpolicy) && get_setting(m, :uncertain_altpolicy)
    if regime_switching
        return solve_regime_switching(m; gensys_regimes = gensys_regimes,
                                      gensys2_regimes = gensys2_regimes, regimes = regimes,
                                      uncertain_altpolicy = uncertain_altpolicy,
                                      verbose = verbose)
    else
        # do NOT delete this Boolean for apply_altpolicy; it is needed for
        # the scenarios code, which continues to use :alternative_policy as a Setting,
        # AND for implementing imperfect awareness via uncertain_altpolicy
        apply_altpolicy = haskey(get_settings(m), :alternative_policy) && get_setting(m, :alternative_policy).solve != solve
        if get_setting(m, :solution_method) == :gensys
            if !apply_altpolicy

                # Get equilibrium condition matrices
                Γ0, Γ1, C, Ψ, Π = eqcond(m)

                # Solve model
                TTT_gensys, CCC_gensys, RRR_gensys, eu = gensys(Γ0, Γ1, C, Ψ, Π, 1+1e-6, verbose = verbose)

                # Check for LAPACK exception, existence and uniqueness
                if eu[1] != 1 || eu[2] != 1
                    throw(GensysError())
                end

                TTT_gensys = real(TTT_gensys)
                RRR_gensys = real(RRR_gensys)
                CCC_gensys = real(CCC_gensys)

                # Augment states
                TTT, RRR, CCC = augment_states(m, TTT_gensys, RRR_gensys, CCC_gensys)
            else
                # do NOT delete this else block; it is needed for
                # the scenarios code, which continues to use :alternative_policy as a Setting
                # and the original method for implementing alternative policies
                # (when DSGE.jl was created in 2015)
                # AND for implementing imperfect awareness via uncertain_altpolicy

                # Change the policy rule
                TTT, RRR, CCC = get_setting(m, :alternative_policy).solve(m)
            end

            if uncertain_altpolicy && apply_altpolicy
                weights = get_setting(m, :regime_eqcond_info)[1].weights
                altpols = get_setting(m, :alternative_policies)

                TTT_gensys, RRR_gensys, CCC_gensys =
                    gensys_uncertain_altpol(m, weights, altpols; TTT = TTT[inds, inds])

                # Augment states
                TTT, RRR, CCC = augment_states(m, TTT_gensys, RRR_gensys, CCC_gensys)
            end
        elseif get_setting(m, :solution_method) == :klein
            TTT_jump, TTT_state = klein(m)

            # Transition
            TTT, RRR = klein_transition_matrices(m, TTT_state, TTT_jump)
            CCC = zeros(n_model_states(m))
        end

        return TTT, RRR, CCC
    end
end

"""
solve(m::PoolModel)
```

Driver to compute the model solution when using the PoolModel type

### Inputs

- `m`: the PoolModel object

### Outputs
- Φ: transition function
- F_ϵ: distribution of structural shock
- F_λ: prior on the initial λ_0

"""
function solve(m::PoolModel)
    return transition(m)
end

"""
```
GensysError <: Exception
```
A `GensysError` is thrown when:

1. Gensys does not give existence and uniqueness, or
2. A LAPACK error was thrown while computing the Schur decomposition of Γ0 and Γ1

If a `GensysError`is thrown during Metropolis-Hastings, it is caught by
`posterior`.  `posterior` then returns a value of `-Inf`, which
Metropolis-Hastings always rejects.

### Fields

* `msg::String`: Info message. Default = \"Error in Gensys\"
"""
mutable struct GensysError <: Exception
    msg::String
end
GensysError() = GensysError("Error in Gensys")
Base.showerror(io::IO, ex::GensysError) = print(io, ex.msg)

"""
```
KleinError <: Exception
```
A `KleinError` is thrown when:
1. A LAPACK error was thrown while computing the pseudo-inverse of U22*U22'

If a `KleinError`is thrown during Metropolis-Hastings, it is caught by
`posterior`.  `posterior` then returns a value of `-Inf`, which
Metropolis-Hastings always rejects.

### Fields

* `msg::String`: Info message. Default = \"Error in Klein\"
"""
mutable struct KleinError <: Exception
    msg::String
end
KleinError() = KleinError("Error in Klein")
Base.showerror(io::IO, ex::KleinError) = print(io, ex.msg)

"""
```
solve_regime_switching(m::AbstractDSGEModel)

solve_one_regime(m::AbstractDSGEModel)

solve_non_gensys2_regimes!(m::AbstractDSGEModel, Γ0s::Vector{Matrix{S}}, Γ01::Vector{Matrix{S}},
    Cs::Vector{Vector{S}}, Ψs::Vector{Matrix{S}}, Πs::Vector{Matrix{S}};
    TTTs::Vector{Matrix{S}}, RRRs::Vector{Matrix{S}}, CCCs::Vector{Vector{S}};
    regimes::Vector{Int} = Int[1], uncertain_altpolicy::Bool = false,
    altpolicy_solve::Function = solve, verbose::Symbol = :high)

solve_gensys2!(m::AbstractDSGEModel, Γ0s::Vector{Matrix{S}}, Γ01::Vector{Matrix{S}},
    Cs::Vector{Vector{S}}, Ψs::Vector{Matrix{S}}, Πs::Vector{Matrix{S}},
    TTT_gensys_final::AbstractMatrix{S}, RRR_gensys_final::AbstractMatrix{S},
    CCC_gensys_final::AbstractMatrix{S},
    TTTs::Vector{Matrix{S}}, RRRs::Vector{Matrix{S}}, CCCs::Vector{Vector{S}};
    fcast_gensys2_regimes::Vector{Int} = Int[1], verbose::Symbol = :high)
```

calculates the reduced form state transition matrices in the case of regime switching.
These functions are intended to be internal functions hidden from the user but are separate
from the definition of the main `solve` function to ensure the main function is
comprehensible.
"""
function solve_regime_switching(m::AbstractDSGEModel{T};
                                uncertain_altpolicy::Bool = false,
                                gensys_regimes::Vector{UnitRange{Int64}} = UnitRange{Int64}[1:1],
                                gensys2_regimes::Vector{UnitRange{Int64}} = Vector{UnitRange{Int64}}(undef, 0),
                                regimes::Vector{Int} = Int[1],
                                verbose::Symbol = :high) where {T <: Real}

    altpolicy_solve = alternative_policy(m).solve
    gensys2 = haskey(get_settings(m), :gensys2) && get_setting(m, :gensys2)
    uncertain_temporary_altpolicy = haskey(get_settings(m), :uncertain_temporary_altpolicy) && get_setting(m, :uncertain_temporary_altpolicy)
    replace_eqcond = haskey(get_settings(m), :replace_eqcond) && get_setting(m, :replace_eqcond)

    if get_setting(m, :solution_method) == :gensys
        if length(regimes) == 1 # Calculate the solution to a specific regime
            return solve_one_regime(m; regime = regimes[1],
                                    uncertain_altpolicy = uncertain_altpolicy, verbose = verbose)
        else # Calculate the reduced-form state space matrices for all regimes
            Γ0s = Vector{Matrix{Float64}}(undef, length(regimes))
            Γ1s = Vector{Matrix{Float64}}(undef, length(regimes))
            Cs = Vector{Vector{Float64}}(undef, length(regimes))
            Ψs = Vector{Matrix{Float64}}(undef, length(regimes))
            Πs = Vector{Matrix{Float64}}(undef, length(regimes))

            do_replace_eqcond = haskey(get_settings(m), :replace_eqcond) && get_setting(m, :replace_eqcond) &&
                haskey(get_settings(m), :regime_eqcond_info)
            has_eqcond_reg_map = haskey(get_settings(m), :identical_eqcond_regimes)
            for reg in regimes
                if has_eqcond_reg_map && haskey(get_setting(m, :identical_eqcond_regimes), reg)
                    reg_to_copy = get_setting(m, :identical_eqcond_regimes)[reg]
                    Γ0s[reg] = copy(Γ0s[reg_to_copy])
                    Γ1s[reg] = copy(Γ1s[reg_to_copy])
                    Cs[reg]  = copy( Cs[reg_to_copy])
                    Ψs[reg]  = copy( Ψs[reg_to_copy])
                    Πs[reg]  = copy( Πs[reg_to_copy])
                else
                    Γ0s[reg], Γ1s[reg], Cs[reg], Ψs[reg], Πs[reg] = do_replace_eqcond && haskey(get_setting(m, :regime_eqcond_info), reg) ?
                        get_setting(m, :regime_eqcond_info)[reg].alternative_policy.eqcond(m, reg) : eqcond(m, reg)
                end
            end

            TTTs = Vector{Matrix{Float64}}(undef, length(regimes))
            RRRs = Vector{Matrix{Float64}}(undef, length(regimes))
            CCCs = Vector{Vector{Float64}}(undef, length(regimes))

            # Solve model for gensys regimes
            for reg_range in gensys_regimes
                solve_non_gensys2_regimes!(m, Γ0s, Γ1s, Cs, Ψs, Πs, TTTs, RRRs, CCCs;
                                           regimes = collect(reg_range),
                                           uncertain_altpolicy = uncertain_altpolicy,
                                           verbose = verbose)
            end

            # Solve for gensys2 regimes
            if gensys2
                for reg_range in gensys2_regimes
                    solve_gensys2!(m, Γ0s, Γ1s, Cs, Ψs, Πs,
                                   TTTs, RRRs, CCCs; gensys2_regimes = collect(reg_range),
                                   uncertain_altpolicy = uncertain_altpolicy,
                                   uncertain_temporary_altpolicy = uncertain_temporary_altpolicy,
                                   verbose = verbose)
                end
            end

            return TTTs, RRRs, CCCs
        end
    else
        error("Regime switching has not been implemented for other solution methods.")
    end
end

function solve_one_regime(m::AbstractDSGEModel{T}; regime::Int = 1, uncertain_altpolicy::Bool = false,
                          verbose::Symbol = :high) where {T <: Real}

    # do NOT delete the second part of the Boolean for apply_altpolicy; it is needed for
    # the scenarios code, which continues to use :alternative_policy as a Setting
    apply_altpolicy = haskey(get_settings(m), :regime_eqcond_info) ||
        (haskey(get_settings(m), :alternative_policy) && get_setting(m, :alternative_policy).key != :historical)
    altpolicy_solve = alternative_policy(m).solve

    if altpolicy_solve == solve || !apply_altpolicy

        # Get equilibrium condition matrices
        Γ0, Γ1, C, Ψ, Π  = eqcond(m, regime)

        # Solve model
        TTT_gensys, CCC_gensys, RRR_gensys, eu = gensys(Γ0, Γ1, C, Ψ, Π, 1+1e-6, verbose = verbose)

        # Check for LAPACK exception, existence and uniqueness
        if eu[1] != 1 || eu[2] != 1
            throw(GensysError())
        end

        TTT_gensys = real(TTT_gensys)
        RRR_gensys = real(RRR_gensys)
        CCC_gensys = real(CCC_gensys)

        # Augment states
        TTT, RRR, CCC = augment_states(m, TTT_gensys, RRR_gensys, CCC_gensys; regime_switching = true,
                                       reg = regime)
    else
        TTT, RRR, CCC = altpolicy_solve(m; regime_switching = true, regimes = Int[regime])
    end

    if uncertain_altpolicy && apply_altpolicy
        # Since we only have one regime in solve_one_regime, no time-varying credibility
        weights = get_setting(m, :regime_eqcond_info)[regime].weights
        altpols = get_setting(m, :alternative_policies)
        inds = 1:n_states(m)
        TTT_gensys, RRR_gensys, CCC_gensys = gensys_uncertain_altpol(m, weights, altpols; TTT = TTT[inds, inds],
                                                                     regimes = Int[regime])

        # Augment states
        TTT, RRR, CCC = augment_states(m, TTT_gensys, RRR_gensys, CCC_gensys; regime_switching = true,
                                       reg = regime)
    end

    return TTT, RRR, CCC
end

function solve_non_gensys2_regimes!(m::AbstractDSGEModel, Γ0s::Vector{Matrix{S}}, Γ1s::Vector{Matrix{S}},
                                    Cs::Vector{Vector{S}}, Ψs::Vector{Matrix{S}}, Πs::Vector{Matrix{S}},
                                    TTTs::Vector{Matrix{S}}, RRRs::Vector{Matrix{S}}, CCCs::Vector{Vector{S}};
                                    regimes::Vector{Int} = Int[1], uncertain_altpolicy::Bool = false,
                                    verbose::Symbol = :high) where {S <: Real}

    # TODO: the commented step below should be unnecessary?
    # uncertain_altpol = haskey(get_settings(m), :uncertain_altpolicy) && get_setting(m, :uncertain_altpolicy)
    uncertain_altpol = uncertain_altpolicy
    has_perf_cred_same_reg = haskey(get_settings(m), :perfect_credibility_identical_transitions)
    if has_perf_cred_same_reg
        perf_cred_same_reg = get_setting(m, :perfect_credibility_identical_transitions)
    end
    if uncertain_altpol
        altpols = get_setting(m, :alternative_policies)
    end

    for reg in regimes
        # Check if, in the case of perfect credibility, whether or not the current regime's gensys solution
        # is identical to another regime's that has already been computed
        if !uncertain_altpol && has_perf_cred_same_reg && haskey(perf_cred_same_reg, reg) && perf_cred_same_reg[reg] != reg
            TTTs[reg] = TTTs[perf_cred_same_reg[reg]]
            RRRs[reg] = RRRs[perf_cred_same_reg[reg]]
            CCCs[reg] = CCCs[perf_cred_same_reg[reg]]

            continue # if so, skip the remainder of this loop
        end

        TTT_gensys, CCC_gensys, RRR_gensys, eu =
            gensys(Γ0s[reg], Γ1s[reg], Cs[reg], Ψs[reg], Πs[reg],
                   1+1e-6, verbose = verbose)

        # Check for LAPACK exception, existence and uniqueness
        if eu[1] != 1 || eu[2] != 1
            throw(GensysError("Error in Gensys, Regime $reg"))
        end

        if uncertain_altpolicy && haskey(get_settings(m), :regime_eqcond_info) && haskey(get_setting(m, :regime_eqcond_info), reg)
            weights = get_setting(m, :regime_eqcond_info)[reg].weights
            if !isempty(weights)
                # Time-varying credibility weights for the regime reg
                altpols = get_setting(m, :alternative_policies)
                TTT_gensys, RRR_gensys, CCC_gensys =
                    gensys_uncertain_altpol(m, weights, altpols; TTT = TTT_gensys,
                                            regime_switching = true, regimes = Int[reg])
            end
        end

        TTT_gensys = real(TTT_gensys)
        RRR_gensys = real(RRR_gensys)
        CCC_gensys = real(CCC_gensys)

        # Populate the TTTs, etc., for regime `reg`
        TTTs[reg], RRRs[reg], CCCs[reg] =
            augment_states(m, TTT_gensys, RRR_gensys, CCC_gensys, reg = reg)
    end

    return TTTs, RRRs, CCCs
end

function solve_gensys2!(m::AbstractDSGEModel, Γ0s::Vector{Matrix{S}}, Γ1s::Vector{Matrix{S}},
                        Cs::Vector{Vector{S}}, Ψs::Vector{Matrix{S}}, Πs::Vector{Matrix{S}},
                        TTTs::Vector{Matrix{S}}, RRRs::Vector{Matrix{S}}, CCCs::Vector{Vector{S}};
                        gensys2_regimes::Vector{Int} = Vector{Int}(undef, 0),
                        uncertain_altpolicy::Bool = false, uncertain_temporary_altpolicy::Bool = false,
                        verbose::Symbol = :high) where {S <: Real}

    # Solve for the final regime of the alternative rule
    altpolicy_solve = alternative_policy(m).solve
    TTT_final, RRR_final, CCC_final = altpolicy_solve(m; regime_switching = true,
                                                      regimes = Int[last(gensys2_regimes)])

    n_endo = length(m.endogenous_states)
    TTT_final = TTT_final[1:n_endo, 1:n_endo] # make sure the non-augmented version
    RRR_final = RRR_final[1:n_endo, :]        # is returned
    CCC_final = CCC_final[1:n_endo]

    TTT_final = real(TTT_final)
    RRR_final = real(RRR_final)
    CCC_final = real(CCC_final)

    # if we're using an uncertain alternative policy, we have to compute the
    # weighted final transition matrix
    if uncertain_altpolicy
        weights = get_setting(m, :regime_eqcond_info)[last(gensys2_regimes)].weights
        altpols = get_setting(m, :alternative_policies)

        TTT_final_weighted, RRR_final_weighted, CCC_final_weighted =
            gensys_uncertain_altpol(m, weights, altpols; TTT = TTT_final,
                                    regimes = Int[last(gensys2_regimes)])

        TTT_final_weighted = real(TTT_final_weighted) # need to define this as different
        RRR_final_weighted = real(RRR_final_weighted) # from TTT_final, which is intended to
        CCC_final_weighted = real(CCC_final_weighted) # be the perfect credibility version
    end

    # Populate TTTs, RRRs, CCCs matrices
    if uncertain_temporary_altpolicy

        # Setup
        ffreg = first(gensys2_regimes) + 1
        altpols = get_setting(m, :alternative_policies)
        weights = Vector{Vector{Float64}}(undef, length(gensys2_regimes) - 1)
        for (i, reg) in enumerate(gensys2_regimes[2]:gensys2_regimes[end])
            weights[i] = get_setting(m, :regime_eqcond_info)[reg].weights
        end
        if length(altpols) == 1
            if altpols[1].solve == solve && haskey(get_settings(m), :uncertain_altpolicy) && get_setting(m, :uncertain_altpolicy)
                revert_unc_altpol = true # Need to make sure uncertain_altpolicy is off if we're using the default solve function
                m <= Setting(:uncertain_altpolicy, false) # since this altpols[1].solve should solve for a perfectly credible policy
            else
                revert_unc_altpol = false
            end
            Talt, _, Calt = altpols[1].solve(m)
            Talt = Talt[1:n_endo,1:n_endo]
            Calt = Calt[1:n_endo]
            if revert_unc_altpol
                m <= Setting(:uncertain_altpolicy, true)
            end
        else
            @assert isa.(altpols, AltPolicy) "All policies in `get_setting(m, :alternative_policies)` should have type AltPolicy."

            Talt = Vector{Matrix{Float64}}(undef, length(altpols)) # can call a faster version
            Calt = Vector{Vector{Float64}}(undef, length(altpols))

            if altpols[1].solve == solve && haskey(get_settings(m), :uncertain_altpolicy) && get_setting(m, :uncertain_altpolicy)
                revert_unc_altpol = true # Need to make sure uncertain_altpolicy is off if we're using the default solve function
                m <= Setting(:uncertain_altpolicy, false) # since this altpols[1].solve should solve for a perfectly credible policy
            else
                revert_unc_altpol = false
            end

            for i in 1:length(altpols)
                Talt[i], _, Calt[i] = altpols[i].solve(m)
                Talt[i] = Talt[i][1:n_endo,1:n_endo]
                Calt[i] = Calt[i][1:n_endo]
            end

            if revert_unc_altpol
                m <= Setting(:uncertain_altpolicy, true)
            end
        end

        # Calculate gensys2 matrices under belief that the desired lift-off policy will occur
        # TODO: generalize to having multiple distinct sets of regimes which are gensys2 regimes
        Tcal, Rcal, Ccal = gensys2(m, Γ0s[gensys2_regimes], Γ1s[gensys2_regimes],
                                   Cs[gensys2_regimes], Ψs[gensys2_regimes], Πs[gensys2_regimes],
                                   TTT_final, RRR_final, CCC_final,
                                   length(gensys2_regimes) - 1)
        Tcal[end] = TTT_final
        Rcal[end] = RRR_final
        Ccal[end] = CCC_final

        # Now calculate transition matrices under an uncertain ZLB
        # Γ0_til, etc., are eqcond matrices implementing ZLB, i.e. zero_rate_rule
        # It is assumed that there is no time variation in the equilibrium conditions
        # for the ZLB, so we can just use ffreg to identify the correct set of
        # equilibrium conditions.
        Γ0_til, Γ1_til, Γ2_til, C_til, Ψ_til =
            gensys_to_predictable_form(Γ0s[ffreg], Γ1s[ffreg], Cs[ffreg], Ψs[ffreg], Πs[ffreg];
                                       use_sparse = (haskey(get_settings(m), :gensys2_sparse_matrices) &&
                                                     get_setting(m, :gensys2_sparse_matrices)))

        ng2  = length(Tcal) - 1 # number of gensys2 regimes
        nzlb = haskey(get_settings(m), :temporary_altpolicy_length) ? get_setting(m, :temporary_altpolicy_length) : ng2

        # Use Tcal, Rcal, & Ccal from 2 as inputs b/c use t + 1 matrix, not t
        # Then, if nzlb = 1, Tcal should have length 2, and you only need the lift-off matrix
        Tcal[1:(1 + nzlb)], Rcal[1:(1 + nzlb)], Ccal[1:(1 + nzlb)] =
            gensys2_uncertain_altpol(weights, Talt, Calt,
                                     Tcal[2:(1 + nzlb)], Rcal[2:(1 + nzlb)], Ccal[2:(1 + nzlb)],
                                     Γ0_til, Γ1_til, Γ2_til, C_til, Ψ_til)

        if uncertain_altpolicy
            Tcal[end] = TTT_final_weighted
            Rcal[end] = RRR_final_weighted
            Ccal[end] = CCC_final_weighted
        else
            Tcal[end] = TTT_final
            Rcal[end] = RRR_final
            Ccal[end] = CCC_final
        end

    else
        Tcal, Rcal, Ccal = gensys2(m, Γ0s[gensys2_regimes], Γ1s[gensys2_regimes],
                                   Cs[gensys2_regimes], Ψs[gensys2_regimes], Πs[gensys2_regimes],
                                   TTT_final, RRR_final, CCC_final,
                                   length(gensys2_regimes) - 1)

        if uncertain_altpolicy
            Tcal[end] = TTT_final_weighted
            Rcal[end] = RRR_final_weighted
            Ccal[end] = CCC_final_weighted
        else
            Tcal[end] = TTT_final
            Rcal[end] = RRR_final
            Ccal[end] = CCC_final
        end
    end

    # only need to populate regimes during which a temporary altpolicy holds
    populate_reg = gensys2_regimes[2:end]
    for (i, reg) in enumerate(populate_reg)
        TTTs[reg], RRRs[reg], CCCs[reg] = augment_states(m, Tcal[i], Rcal[i], Ccal[i], reg = reg)
    end

    return TTTs, RRRs, CCCs
end

"""
```
solve_uncertain_altpolicy(m::AbstractDSGEModel{T},
                          system_perfect_cred_totpolicies::Vector{Union{System, RegimeSwitchingSystem}},
                          is_altpol::Vector{Bool};
                          gensys_regimes::Vector{UnitRange{Int64}} = UnitRange{Int64}[1:1],
                          gensys2_regimes::Vector{UnitRange{Int64}} = Vector{UnitRange{Int}}(undef, 0),
                          regimes::Vector{Int} = Int[1], verbose::Symbol = :high) where {T <: Real}
```

computes the regime-switching state space transition matrices when agents have imperfect awareness
about the policy being implemented and some of these policies may be temporary policies
(and thus would be specified as a `MultiPeriodAltPolicy`).
"""
function solve_uncertain_altpolicy(m::AbstractDSGEModel{T},
                                   system_perfect_cred_totpolicies::Vector{Union{System, RegimeSwitchingSystem}},
                                   is_altpol::Vector{Bool};
                                   gensys_regimes::Vector{UnitRange{Int64}} = UnitRange{Int64}[1:1],
                                   gensys2_regimes::Vector{UnitRange{Int64}} = Vector{UnitRange{Int}}(undef, 0),
                                   regimes::Vector{Int} = Int[1], verbose::Symbol = :high) where {T <: Real}

    Γ0s = Vector{Matrix{Float64}}(undef, length(regimes))
    Γ1s = Vector{Matrix{Float64}}(undef, length(regimes))
    Cs = Vector{Vector{Float64}}(undef, length(regimes))
    Ψs = Vector{Matrix{Float64}}(undef, length(regimes))
    Πs = Vector{Matrix{Float64}}(undef, length(regimes))

    @assert haskey(get_settings(m), :regime_eqcond_info) "The setting :regime_eqcond_info must exist"

    has_eqcond_reg_map = haskey(get_settings(m), :identical_eqcond_regimes)
    for reg in regimes
        if has_eqcond_reg_map && haskey(get_setting(m, :identical_eqcond_regimes), reg)
            reg_to_copy = get_setting(m, :identical_eqcond_regimes)[reg]
            Γ0s[reg] = copy(Γ0s[reg_to_copy])
            Γ1s[reg] = copy(Γ1s[reg_to_copy])
            Cs[reg]  = copy( Cs[reg_to_copy])
            Ψs[reg]  = copy( Ψs[reg_to_copy])
            Πs[reg]  = copy( Πs[reg_to_copy])
        else
            Γ0s[reg], Γ1s[reg], Cs[reg], Ψs[reg], Πs[reg] = haskey(get_setting(m, :regime_eqcond_info), reg) ?
                get_setting(m, :regime_eqcond_info)[reg].alternative_policy.eqcond(m, reg) : eqcond(m, reg)
        end
    end

    TTTs = Vector{Matrix{Float64}}(undef, length(regimes))
    RRRs = Vector{Matrix{Float64}}(undef, length(regimes))
    CCCs = Vector{Vector{Float64}}(undef, length(regimes))

    for reg_range in gensys_regimes
        solve_non_gensys2_regimes!(m, Γ0s, Γ1s, Cs, Ψs, Πs, TTTs, RRRs, CCCs,
                                   system_perfect_cred_totpolicies, is_altpol;
                                   regimes = collect(reg_range),
                                   verbose = verbose)
    end

    for reg_range in gensys2_regimes
        solve_gensys2!(m, Γ0s, Γ1s, Cs, Ψs, Πs,
                       TTTs, RRRs, CCCs, system_perfect_cred_totpolicies, is_altpol;
                       gensys2_regimes = collect(reg_range),
                       verbose = verbose)
    end

    return TTTs, RRRs, CCCs
end

function solve_non_gensys2_regimes!(m::AbstractDSGEModel, Γ0s::Vector{Matrix{S}}, Γ1s::Vector{Matrix{S}},
                                    Cs::Vector{Vector{S}}, Ψs::Vector{Matrix{S}}, Πs::Vector{Matrix{S}},
                                    TTTs::Vector{Matrix{S}}, RRRs::Vector{Matrix{S}}, CCCs::Vector{Vector{S}},
                                    system_perfect_cred_totpolicies::Vector{Union{System, RegimeSwitchingSystem}},
                                    is_altpol::Vector{Bool}; regimes::Vector{Int} = Int[1],
                                    verbose::Symbol = :high) where {S <: Real}

    for reg in regimes
        if haskey(get_settings(m), :regime_eqcond_info) && haskey(get_setting(m, :regime_eqcond_info), reg) &&
            !isempty(get_setting(m, :regime_eqcond_info)[reg].weights)

            weights = get_setting(m, :regime_eqcond_info)[reg].weights

            # Time-varying credibility weights for the regime reg
            TTT_gensys, RRR_gensys, CCC_gensys =
                gensys_uncertain_altpol(m, weights, reg, system_perfect_cred_totpolicies, is_altpol)
        else
            TTT_gensys, CCC_gensys, RRR_gensys, eu =
                gensys(Γ0s[reg], Γ1s[reg], Cs[reg], Ψs[reg], Πs[reg],
                       1+1e-6, verbose = verbose)

            # Check for LAPACK exception, existence and uniqueness
            if eu[1] != 1 || eu[2] != 1
                throw(GensysError("Error in Gensys, Regime $reg"))
            end
        end

        TTT_gensys = real(TTT_gensys)
        RRR_gensys = real(RRR_gensys)
        CCC_gensys = real(CCC_gensys)

        # Populate the TTTs, etc., for regime `reg`
        TTTs[reg], RRRs[reg], CCCs[reg] =
            augment_states(m, TTT_gensys, RRR_gensys, CCC_gensys, reg = reg)
    end

    return TTTs, RRRs, CCCs
end

function solve_gensys2!(m::AbstractDSGEModel, Γ0s::Vector{Matrix{S}}, Γ1s::Vector{Matrix{S}},
                        Cs::Vector{Vector{S}}, Ψs::Vector{Matrix{S}}, Πs::Vector{Matrix{S}},
                        TTTs::Vector{Matrix{S}}, RRRs::Vector{Matrix{S}}, CCCs::Vector{Vector{S}},
                        system_perfect_cred_totpolicies::Vector{Union{System, RegimeSwitchingSystem}},
                        is_altpol::Vector{Bool};
                        gensys2_regimes::Vector{Int} = Vector{Int}(undef, 0),
                        verbose::Symbol = :high) where {S <: Real}

    # Solve for the final regime of the alternative rule
    TTT_final = system_perfect_cred_totpolicies[1][gensys2_regimes[end], :TTT] # policy 1 is the perfectly credible version of the implemented policy
    RRR_final = system_perfect_cred_totpolicies[1][gensys2_regimes[end], :RRR]
    CCC_final = system_perfect_cred_totpolicies[1][gensys2_regimes[end], :CCC]

    # if we're using an uncertain alternative policy, we have to compute the
    # weighted final transition matrix
    weights = get_setting(m, :regime_eqcond_info)[last(gensys2_regimes)].weights
    TTT_final_weighted, RRR_final_weighted, CCC_final_weighted =
        gensys_uncertain_altpol(m, weights, last(gensys2_regimes), system_perfect_cred_totpolicies, is_altpol)

    TTT_final_weighted = real(TTT_final_weighted) # need to define this as different
    RRR_final_weighted = real(RRR_final_weighted) # from TTT_final, which is intended to
    CCC_final_weighted = real(CCC_final_weighted) # be the perfect credibility version

    # Now calculate transition matrices under an uncertain ZLB
    # Γ0_til, etc., are eqcond matrices implementing ZLB, i.e. zero_rate_rule
    # It is assumed that there is no time variation in the equilibrium conditions
    # for the ZLB, so we can just use ffreg to identify the correct set of
    # equilibrium conditions.
    ffreg = first(gensys2_regimes) + 1
    Γ0_til, Γ1_til, Γ2_til, C_til, Ψ_til =
        gensys_to_predictable_form(Γ0s[ffreg], Γ1s[ffreg], Cs[ffreg], Ψs[ffreg], Πs[ffreg];
                                   use_sparse = (haskey(get_settings(m), :gensys2_sparse_matrices) &&
                                                 get_setting(m, :gensys2_sparse_matrices)))

    # Use Tcal, Rcal, & Ccal from 2 as inputs b/c use t + 1 matrix, not t
    # Then, if nzlb = 1, Tcal should have length 2, and you only need the lift-off matrix
    weights = Vector{Vector{Float64}}(undef, length(gensys2_regimes) - 1) # prep time varying weights
    for (i, reg) in enumerate(gensys2_regimes[2]:gensys2_regimes[end])
        weights[i] = get_setting(m, :regime_eqcond_info)[reg].weights
    end
    Tcal, Rcal, Ccal =
        gensys2_uncertain_altpol(weights, gensys2_regimes, 1:n_states(m), system_perfect_cred_totpolicies, is_altpol,
                                 Γ0_til, Γ1_til, Γ2_til, C_til, Ψ_til)

    ng2  = length(Tcal) - 1 # number of gensys2 regimes
    nzlb = haskey(get_settings(m), :temporary_altpolicy_length) ? get_setting(m, :temporary_altpolicy_length) : ng2

    @assert nzlb == ng2 "The setting :temporary_altpolicy_length does not match the number of gensys2 regimes."

    Tcal[end] = TTT_final_weighted
    Rcal[end] = RRR_final_weighted
    Ccal[end] = CCC_final_weighted

    # only need to populate regimes during which a temporary altpolicy holds
    populate_reg = gensys2_regimes[2:end]
    for (i, reg) in enumerate(populate_reg)
        TTTs[reg], RRRs[reg], CCCs[reg] = augment_states(m, Tcal[i], Rcal[i], Ccal[i], reg = reg)
    end

    return TTTs, RRRs, CCCs
end

function solve_non_gensys2_regimes!(m::AbstractDSGEModel, Γ0s::Vector{Matrix{S}}, Γ1s::Vector{Matrix{S}},
                                    Cs::Vector{Vector{S}}, Ψs::Vector{Matrix{S}}, Πs::Vector{Matrix{S}},
                                    TTTs::Vector{Matrix{S}}, RRRs::Vector{Matrix{S}}, CCCs::Vector{Vector{S}},
                                    system_perfect_cred_totpolicies::Vector{Union{System, RegimeSwitchingSystem}};
                                    regimes::Vector{Int} = Int[1], verbose::Symbol = :high) where {S <: Real}

    for reg in regimes
        if haskey(get_settings(m), :regime_eqcond_info) && haskey(get_setting(m, :regime_eqcond_info), reg) &&
            !isempty(get_setting(m, :regime_eqcond_info)[reg].weights)

            weights = get_setting(m, :regime_eqcond_info)[reg].weights

            # Time-varying credibility weights for the regime reg
            TTT_gensys, RRR_gensys, CCC_gensys =
                gensys_uncertain_altpol(m, weights, reg, system_perfect_cred_totpolicies, trues(length(weights) - 1))
        else
            TTT_gensys, CCC_gensys, RRR_gensys, eu =
                gensys(Γ0s[reg], Γ1s[reg], Cs[reg], Ψs[reg], Πs[reg],
                       1+1e-6, verbose = verbose)

            # Check for LAPACK exception, existence and uniqueness
            if eu[1] != 1 || eu[2] != 1
                throw(GensysError("Error in Gensys, Regime $reg"))
            end
        end

        TTT_gensys = real(TTT_gensys)
        RRR_gensys = real(RRR_gensys)
        CCC_gensys = real(CCC_gensys)

        # Populate the TTTs, etc., for regime `reg`
        TTTs[reg], RRRs[reg], CCCs[reg] =
            augment_states(m, TTT_gensys, RRR_gensys, CCC_gensys, reg = reg)
    end

    return TTTs, RRRs, CCCs
end


text:
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Description
# ==============================================================================
#
#   This file contains functions to create printing recipes.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

# Constants that identify some actions inside the column receipes. They **must
# be** lower than zero.
const _LEFT_LINE = -1
const _COLUMN_LINE = -2
const _RIGHT_LINE = -3

# Constants that identify some actions inside the row receipes. They **must be**
# a `NTuple{4, Int}` with the first index lower than 0.
const _ROW_LINE = (-1, 0, 0, 0)
const _CONTINUATION_LINE = (-2, 0, 0, 0)

# Create the printing recipes for the rows and columns of the table. This
# function also computed how many rows and columns are omitted in the printing.
function _create_printing_recipe(
    display::Display,
    header_num_rows::Int,
    num_filtered_rows::Int,
    num_filtered_cols::Int,
    num_printed_rows::Int,
    num_printed_cols::Int,
    num_lines_in_row::Vector{Int},
    cols_width::Vector{Int},
    id_rows::Vector{Int},
    hlines::Vector{Int},
    vlines::Vector{Int},
    Δdisplay_lines::Int,
    Δc::Int,
    # Configurations
    crop::Symbol,
    noheader::Bool,
    show_omitted_cell_summary::Bool,
    vcrop_mode::Symbol
)
    # Column printing recipe
    # ==========================================================================

    col_printing_recipe   = Vector{Int}(undef, 0)
    data_horizontal_limit = 0
    fully_printed_cols    = 0

    # Verify if the output must be cropped.
    if crop == :both || crop == :horizontal
        data_horizontal_limit = display.size[2] - 1
    end

    # The variable stores the current line length to check how many columns we
    # can print.
    line_length = 0

    if 0 ∈ vlines
        line_length += 1
        push!(col_printing_recipe, _LEFT_LINE)
    end

    @inbounds for i = 1:num_printed_cols
        # Space before the next row.
        line_length += 1

        # Width of the column.
        line_length += cols_width[i]
        push!(col_printing_recipe, i)

        Δ = data_horizontal_limit > 0 ? data_horizontal_limit - line_length : 1
        Δ ≤ 0 && break

        fully_printed_cols += 1

        # Space after the row printing.
        line_length += 1

        if i ∈ vlines
            if i != num_printed_cols
                push!(col_printing_recipe, _COLUMN_LINE)
            else
                push!(col_printing_recipe, _RIGHT_LINE)
            end
            line_length += 1
        end
    end

    # Compute the number of omitted columns.
    num_omitted_cols = clamp(
        num_filtered_cols - (fully_printed_cols - Δc),
        0,
        num_filtered_cols
    )

    # Row printing recipe
    # ==========================================================================

    # The row printing recipe can be a symbol, which indicates a type of line
    # that must be drawn, or a tuple of three integers that describe:
    #
    #     1. The index of the row that must be printed in `data_str`.
    #     2. The index of the row that must be printed in `data`.
    #     3. How many lines must be printed.
    #     4. The initial printing line.

    row_printing_recipe = Vector{NTuple{4, Int}}(undef, 0)
    data_vertical_limit = 0
    fully_printed_rows  = 0

    # Compute how many additional lines we need to print the cropping
    # information.
    crop_space = show_omitted_cell_summary ? 2 : 1

    # Verify if the output must be vertically cropped.
    data_vcropped = false

    if (display.size[1] > 0) && ( (crop == :both) || (crop == :vertical) )
        # Compute the number of lines required to print the table. Notice that
        # we need to omit the last line since it is always printed.
        total_hlines = count(x -> 0 ≤ x < (!noheader + num_printed_rows), hlines)
        total_table_lines  = !noheader ? header_num_rows : 0
        total_table_lines += total_hlines + sum(num_lines_in_row)

        # Check if we can print all the table lines in the available space.
        data_vertical_limit = display.size[1] - Δdisplay_lines - crop_space
        Δ = data_vertical_limit - total_table_lines

        # Given the additional space we have when printing the continuation
        # line and possibly the summary information, we need to check if
        # the remaining lines fit on it.
        data_vcropped = true

        if Δ > -crop_space
            data_vcropped = false
        else
            if show_omitted_cell_summary && (num_omitted_cols == 0) && (Δ+2 ≥ 0)
                data_vcropped = false
            elseif (Δ+1 ≥ 0)
                data_vcropped = false
            end
        end
    end

    # If data is not cropped, then we just need to add all the columns to the
    # recipe.
    if !data_vcropped
        fully_printed_rows = num_printed_rows

        @inbounds for i = 1:num_printed_rows
            # Index of this row in `data`.
            ir = id_rows[i]

            push!(row_printing_recipe, (i, ir, num_lines_in_row[i], 1))

            if (i != num_printed_rows) && ((i+!noheader) ∈ hlines)
                push!(row_printing_recipe, _ROW_LINE)
            end
        end
    else
        # Compute the required size for the header.
        header_length  = 0 ∈ hlines ? 1 : 0

        if !noheader
            header_length += header_num_rows
            header_length += 1 ∈ hlines ? 1 : 0
        end

        # This variable stores the number of printed lines to verify how many we
        # can print.
        printed_lines = header_length

        fully_printed_rows = 0

        if vcrop_mode == :bottom
            @inbounds for i = 1:num_printed_rows
                num_lines_row_i = num_lines_in_row[i]

                # Index of this row in `data`.
                ir = id_rows[i]

                # This variable contains the number of available lines we have
                # to print the data.
                Δ = data_vertical_limit - (num_lines_row_i + printed_lines)

                # Verify if the entire row can be printed.
                if Δ < 0
                    remaining_rows = data_vertical_limit - printed_lines

                    if remaining_rows > 0
                        push!(row_printing_recipe, (i, ir, remaining_rows, 1))
                    end

                    push!(row_printing_recipe, _CONTINUATION_LINE)
                    break
                else
                    push!(row_printing_recipe, (i, ir, num_lines_row_i, 1))
                    printed_lines += num_lines_row_i
                    fully_printed_rows += 1

                    # Check if we have space for drawing the row line.
                    Δ = data_vertical_limit - printed_lines

                    if (Δ < 0)
                        push!(row_printing_recipe, _CONTINUATION_LINE)
                        break
                    else
                        if (Δ > 0) && (i != num_printed_rows) && ((i+!noheader) ∈ hlines)
                            push!(row_printing_recipe, _ROW_LINE)
                            printed_lines += 1
                        end
                    end
                end
            end

        elseif vcrop_mode == :middle
            # In this mode, we will fill the printing using two passes. In the
            # first, we compute the top of the table and, in the second, the
            # bottom of the table. In the latter, it is easier to fill from the
            # end to the beginning of the table. Hence, we create a temporary
            # vector of the recipe that will be reversed and merged into the
            # main one.
            row_printing_recipe_end = Vector{NTuple{4, Int}}(undef, 0)

            # Compute the amount of space left to print the table data.
            table_data_space = data_vertical_limit - header_length

            Δs = cld(table_data_space, 2)
            Δr = cld(num_printed_rows, 2)

            printed_lines = 0

            # Fill the top of the table.
            @inbounds for i = 1:Δr
                num_lines_row_i = num_lines_in_row[i]

                # Index of this row in `data`.
                ir = id_rows[i]

                # This variable contains the number of available lines we have
                # to print the data.
                Δ = Δs - (num_lines_row_i + printed_lines)

                # Verify if the entire row can be printed.
                if Δ < 0
                    remaining_rows = Δs - printed_lines

                    if remaining_rows > 0
                        push!(row_printing_recipe, (i, ir, remaining_rows, 1))
                    end

                    break
                else
                    push!(row_printing_recipe, (i, ir, num_lines_row_i, 1))
                    printed_lines += num_lines_row_i
                    fully_printed_rows += 1

                    # Check if we have space for drawing the row line.
                    Δ = Δs - printed_lines

                    if (Δ < 0)
                        break
                    else
                        if (Δ > 0) && (i != num_printed_rows) && ((i+!noheader) ∈ hlines)
                            push!(row_printing_recipe, _ROW_LINE)
                            printed_lines += 1
                        end
                    end
                end
            end

            # In this case, we know that the data is cropped. Hence, we print
            # the continuation line in the middle.
            push!(row_printing_recipe, _CONTINUATION_LINE)

            Δs = table_data_space - Δs

            printed_lines = 0

            @inbounds for i = num_printed_rows:-1:Δr+1
                num_lines_row_i = num_lines_in_row[i]

                # Index of this row in `data`.
                ir = id_rows[num_filtered_rows - (num_printed_rows - i)]

                # This variable contains the number of available lines we have
                # to print the data.
                Δ = Δs - (num_lines_row_i + printed_lines)

                # Verify if the entire row can be printed.
                if Δ < 0
                    remaining_rows = Δs - printed_lines

                    # This row is cropped and we are in the bottom part. Hence,
                    # we need to print the lower part of the row.
                    if remaining_rows > 0
                        push!(
                            row_printing_recipe_end,
                            (i, ir, remaining_rows, num_lines_row_i - remaining_rows + 1)
                        )
                    end

                    break
                else
                    push!(row_printing_recipe_end, (i, ir, num_lines_row_i, 1))
                    printed_lines += num_lines_row_i
                    fully_printed_rows += 1

                    # Check if we have space for drawing the row line.
                    Δ = Δs - printed_lines

                    if (Δ < 0)
                        break
                    else
                        if (Δ > 0) && ((i-1+!noheader) ∈ hlines)
                            push!(row_printing_recipe_end, _ROW_LINE)
                            printed_lines += 1
                        end
                    end
                end
            end

            # Merge the top and bottom parts.
            row_printing_recipe = vcat(
                row_printing_recipe,
                reverse(row_printing_recipe_end)
            )
        else
            error("Unknown vertical crop mode.")
        end
    end

    num_omitted_rows = num_filtered_rows - fully_printed_rows

    return row_printing_recipe, col_printing_recipe, num_omitted_rows, num_omitted_cols
end


text:
#
# types.jl --
#
# Definitions of types.
#
#-------------------------------------------------------------------------------
#
# This file is part of "AndorCameras.jl" released under the MIT license.
#
# Copyright (C) 2017-2021, Éric Thiébaut.
#

"""

`AndorError` is an exception sub-type to report errors occuring in calls to the
Sofware Develement Kit (SDK) of Andor cameras.

"""
struct AndorError <: Exception
    func::Symbol
    code::Cint
    AndorError(func::Symbol, code::Integer) = new(func, code)
    AndorError(func::AbstractString, code::Integer) = new(Symbol(func), code)
end

"""

`AndorCameraModel` is used to wrap a constant identifier to quickly identify a
specific camera model.

"""
struct AndorCameraModel
    id::Int
end
Base.:(==)(a::AndorCameraModel, b::AndorCameraModel) = a.id == b.id
const _UNKNOWN_MODEL  = AndorCameraModel(0)
const _SIM_CAM_MODEL  = AndorCameraModel(1)
const _ZYLA_MODEL     = AndorCameraModel(2)
const _ZYLA_USB_MODEL = AndorCameraModel(3)

# Camera structure.
mutable struct Camera <: ScientificCamera
    state::Int
    bufs::Vector{Vector{UInt8}}        # buffers for the frame grabber
    lastimg::Array{T,2} where {T}      # last image
    bytesperline::Int
    clockfrequency::Int
    model::AndorCameraModel
    handle::AT.Handle
    mono12packed::Bool
    Camera() = new(0,
                   Vector{Vector{UInt8}}(undef, 0),
                   Array{UInt8,2}(undef, 0, 0),
                   0, 0, _UNKNOWN_MODEL, AT.HANDLE_UNINITIALISED, false)
end

const AndorCamera = Camera

"""

Concrete sub-types of `AbstractFeature` are simple wrappers around
`WideString` instances to allow for type-stability.

"""
abstract type AbstractFeature end

for T in (:CommandFeature, :BooleanFeature, :EnumeratedFeature,
          :IntegerFeature, :FloatingPointFeature, :StringFeature,
          :BooleanOrEnumeratedFeature, )
    @eval struct $T <: AbstractFeature
        name::WideString
    end
end


text:
using ClimateMachine

if parse(Bool, get(ENV, "CLIMATEMACHINE_PLOT_EDMF_COMPARISON", "false"))
    using Plots
end

using OrderedCollections
using Test
using NCDatasets
using Dierckx
using PrettyTables
using Printf
using ClimateMachine.ArtifactWrappers

# Get PyCLES_output dataset folder:
#! format: off
PyCLES_output_dataset = ArtifactWrapper(
    @__DIR__,
    isempty(get(ENV, "CI", "")),
    "PyCLES_output",
    ArtifactFile[
    # ArtifactFile(url = "https://caltech.box.com/shared/static/johlutwhohvr66wn38cdo7a6rluvz708.nc", filename = "Rico.nc",),
    ArtifactFile(url = "https://caltech.box.com/shared/static/zraeiftuzlgmykzhppqwrym2upqsiwyb.nc", filename = "Gabls.nc",),
    # ArtifactFile(url = "https://caltech.box.com/shared/static/toyvhbwmow3nz5bfa145m5fmcb2qbfuz.nc", filename = "DYCOMS_RF01.nc",),
    # ArtifactFile(url = "https://caltech.box.com/shared/static/ivo4751camlph6u3k68ftmb1dl4z7uox.nc", filename = "TRMM_LBA.nc",),
    # ArtifactFile(url = "https://caltech.box.com/shared/static/4osqp0jpt4cny8fq2ukimgfnyi787vsy.nc", filename = "ARM_SGP.nc",),
    ArtifactFile(url = "https://caltech.box.com/shared/static/jci8l11qetlioab4cxf5myr1r492prk6.nc", filename = "Bomex.nc",),
    # ArtifactFile(url = "https://caltech.box.com/shared/static/pzuu6ii99by2s356ij69v5cb615200jq.nc", filename = "Soares.nc",),
    # ArtifactFile(url = "https://caltech.box.com/shared/static/7upt639siyc2umon8gs6qsjiqavof5cq.nc", filename = "Nieuwstadt.nc",),
    ],
)
PyCLES_output_dataset_path = get_data_folder(PyCLES_output_dataset)
#! format: on

include("variable_map.jl")

function compute_mse(
    grid,
    bl,
    time_cm,
    dons_arr,
    ds,
    experiment,
    best_mse,
    t_compare,
    plot_dir = nothing,
)
    mse = Dict()

    # Ensure domain matches:
    z_les = ds["z_half"][:]
    z_cm = get_z(grid; rm_dupes = true)
    @info "Z extent for LES vs CLIMA:"
    @show extrema(z_cm)
    @show extrema(z_les)

    time_les = ds["t"][:]

    # Find the nearest matching final time:
    t_cmp = min(time_cm[end], time_les[end])

    # Accidentally running a short simulation
    # could improve MSE. So, let's test that
    # we run for at least t_compare. We should
    # increase this as we can reach higher CFL.
    @test t_cmp >= t_compare

    # Ensure z_cm and dons_arr fields are consistent lengths:
    @test length(z_cm) == length(dons_arr[1][first(keys(dons_arr[1]))])

    data_cm = Dict()
    dons_cont = Dict()
    cm_variables = []
    computed_mse = []
    table_best_mse = []
    mse_reductions = []
    pycles_variables = []
    data_scales = []
    pycles_weight = []
    for (ftc) in keys(best_mse)
        # Only compare fields defined for var_map
        tup = var_map(ftc)
        tup == nothing && continue

        # Unpack the data
        LES_var = tup[1]
        facts = tup[2]
        push!(cm_variables, ftc)
        push!(pycles_variables, LES_var)
        data_ds = ds.group["profiles"]
        data_les = data_ds[LES_var][:]

        # Scale the data for comparison
        ρ = data_ds["rho"][:]
        a_up = data_ds["updraft_fraction"][:]
        a_en = 1 .- data_ds["updraft_fraction"][:]
        ρa_up = ρ .* a_up
        ρa_en = ρ .* a_en
        ρa = occursin("updraft", ftc) in ftc ? ρa_up : ρa_en
        if :a in facts && :ρ in facts
            data_les .*= ρa
            push!(pycles_weight, "ρa")
        elseif :ρ in facts
            data_les .*= ρ
            push!(pycles_weight, "ρ")
        else
            push!(pycles_weight, "1")
        end

        # Interpolate data
        steady_data = length(size(data_les)) == 1
        if steady_data
            data_les_cont = Spline1D(z_les, data_les)
        else # unsteady data
            data_les_cont = Spline2D(time_les, z_les, data_les')
        end
        data_cm_arr_ = [
            dons_arr[i][ftc][i_z]
            for i in 1:length(time_cm), i_z in 1:length(z_cm)
        ]
        data_cm_arr = reshape(data_cm_arr_, (length(time_cm), length(z_cm)))
        data_cm[ftc] = Spline2D(time_cm, z_cm, data_cm_arr)

        # Compute data scale
        data_scale = sum(abs.(data_les)) / length(data_les)
        push!(data_scales, data_scale)

        # Plot comparison
        if plot_dir ≠ nothing
            p = plot()
            if steady_data
                data_les_cont_mapped = map(z -> data_les_cont(z), z_cm)
            else
                data_les_cont_mapped = map(z -> data_les_cont(t_cmp, z), z_cm)
            end
            plot!(
                data_les_cont_mapped,
                z_cm ./ 10^3,
                xlabel = ftc,
                ylabel = "z [km]",
                label = "PyCLES",
            )
            plot!(
                map(z -> data_cm[ftc](t_cmp, z), z_cm),
                z_cm ./ 10^3,
                xlabel = ftc,
                ylabel = "z [km]",
                label = "CM",
            )
            mkpath(plot_dir)
            ftc_name = replace(ftc, "." => "_")
            savefig(joinpath(plot_dir, "$ftc_name.png"))
        end

        # Compute mean squared error (mse)
        if steady_data
            mse_single_var = sum(map(z_cm) do z
                (data_les_cont(z) - data_cm[ftc](t_cmp, z))^2
            end)
        else
            mse_single_var = sum(map(z_cm) do z
                (data_les_cont(t_cmp, z) - data_cm[ftc](t_cmp, z))^2
            end)
        end
        # Normalize by data scale
        mse[ftc] = mse_single_var / data_scale^2

        push!(mse_reductions, (best_mse[ftc] - mse[ftc]) / best_mse[ftc] * 100)
        push!(computed_mse, mse[ftc])
        push!(table_best_mse, best_mse[ftc])
    end

    # Tabulate output
    header = [
        "Variable" "Variable" "Weight" "Data scale" "MSE" "MSE" "MSE"
        "ClimateMachine (EDMF)" "PyCLES" "PyCLES" "" "Computed" "Best" "Reduction (%)"
    ]
    table_data = hcat(
        cm_variables,
        pycles_variables,
        pycles_weight,
        data_scales,
        computed_mse,
        table_best_mse,
        mse_reductions,
    )

    @info @sprintf(
        "Experiment comparison: %s at time t=%s\n",
        experiment,
        t_cmp
    )
    hl_worsened_mse = Highlighter(
        (data, i, j) -> !sufficient_mse(data[i, 5], data[i, 6]) && j == 5,
        crayon"red bold",
    )
    hl_worsened_mse_reduction = Highlighter(
        (data, i, j) -> !sufficient_mse(data[i, 5], data[i, 6]) && j == 7,
        crayon"red bold",
    )
    hl_improved_mse = Highlighter(
        (data, i, j) -> sufficient_mse(data[i, 5], data[i, 6]) && j == 7,
        crayon"green bold",
    )
    pretty_table(
        table_data,
        header,
        formatters = ft_printf("%.16e", 5:6),
        header_crayon = crayon"yellow bold",
        subheader_crayon = crayon"green bold",
        highlighters = (
            hl_worsened_mse,
            hl_improved_mse,
            hl_worsened_mse_reduction,
        ),
        crop = :none,
    )

    return mse
end

sufficient_mse(computed_mse, best_mse) = computed_mse <= best_mse + sqrt(eps())

function test_mse(computed_mse, best_mse, key)
    mse_not_regressed = sufficient_mse(computed_mse[key], best_mse[key])
    @test mse_not_regressed
    mse_not_regressed || @show key
end

function dons(diag_vs_z)
    return Dict(map(keys(first(diag_vs_z))) do k
        string(k) => [getproperty(ca, k) for ca in diag_vs_z]
    end)
end

get_dons_arr(diag_arr) = [dons(diag_vs_z) for diag_vs_z in diag_arr]

dons_arr = get_dons_arr(diag_arr)