text:
println("GHCN")

using Pkg
Pkg.activate(".");
Pkg.instantiate();

using ArgParse
using Bijectors
using Clustering
using CSV
using DataFrames
using Dates
using BenchmarkTools
using DrWatson
using FileIO
using Formatting
using GADM
using GeoDatasets
using GHCNData
using Images
using LinearAlgebra
using LineSearches
using Optim
using ParameterHandling
using Plots
using Random
using Stheno
using TemporalGPs
using TexTables
using Zygote

using ParameterHandling: positive, deferred, flatten
using TemporalGPs: approx_posterior_marginals, RegularInTime, Separable
using Zygote: dropgrad

using TemporalGPs:
    approx_posterior_marginals,
    get_time,
    get_space,
    RectilinearGrid,
    Separable

const exp_dir_name = joinpath("global_historical_climatology_network");
const data_dir = joinpath(datadir(), exp_dir_name);
const metadata_dir = joinpath(data_dir, "meta");
const training_output_dir = joinpath(data_dir, "training");
const prediction_output_dir = joinpath(data_dir, "predictions");
const results_dir = joinpath(projectdir(), "results", exp_dir_name);

# Washington + a bit of Canada.
const lon_interval = (-125, -122);
const lat_interval = (47, 49);

# Time intervals to consider.
const time_interval = (Date(1980), Date(2020));
const element = "TMAX"

function make_tmax_data_path(lat_interval, lon_interval, time_interval)
    return joinpath(
        data_dir,
        "tmax_data-$(lat_interval[1])-$(lat_interval[2])" *
            "-$(lon_interval[1])-$(lon_interval[2])" *
            "-$(time_interval[1])-$(time_interval[2])" *
            ".jld2",
    )
end

raw_tmax_data_path = make_tmax_data_path(lat_interval, lon_interval, time_interval);

raw_tmax_data, station_lat_lon_pairs, IDs = if isfile(raw_tmax_data_path)

    @info "DATA ALREADY SERIALISED"

    # If the data has been serialised already, just load it.
    data = FileIO.load(raw_tmax_data_path)

    data["raw_tmax_data"], data["station_lat_lon_pairs"], data["IDs"]
else

    @info "SERIALISING RAW DATA"

    # Otherwise load the raw data, serialise it, and load the serialised version.
    raw_tmax_data, station_lat_lon_pairs, IDs = GHCNData.select_data(
        time_interval, lat_interval, lon_interval, element,
    )

    # Serialise the data.
    FileIO.save(
        raw_tmax_data_path,
        Dict(
            "raw_tmax_data" => raw_tmax_data,
            "station_lat_lon_pairs" => station_lat_lon_pairs,
            "IDs" => IDs,
        ),
    )

    # Load the data and return it.
    data = FileIO.load(raw_tmax_data_path)
    data["raw_tmax_data"], data["station_lat_lon_pairs"], data["IDs"]
end

const station_lats = first.(station_lat_lon_pairs);
const station_lons = last.(station_lat_lon_pairs);

# Get the elevations associated with each station.
station_metadata = load_station_metadata();
elevation_dict = Dict([row.ID => row.ELEVATION for row in eachrow(station_metadata)]);
elevations = [elevation_dict[ID] for ID in IDs];

# TMAX data is in 10ths of a degree C. Remove obvious outliers.
const tmax_data = map(
    x -> (x !== missing && (-30 < x < 60)) ? x : missing,
    raw_tmax_data ./ 10,
);
const tmax_dates = time_interval[1]:Day(1):(time_interval[2]-Day(1))

# Determine locations at which to make predictions.
const N_base = 250;
const N_lons = N_base * (lon_interval[2] - lon_interval[1]);
const N_lats = N_base * (lat_interval[2] - lat_interval[1]);
const uk_lats_range = range(lat_interval[1], lat_interval[2]; length=N_lons);
const uk_lons_range = range(lon_interval[1], lon_interval[2]; length=N_lats);
const lat_lon_grid = tuple.(uk_lats_range, uk_lons_range');

# Look at some basic statistics of the data.
@show minimum(filter(!ismissing, tmax_data));
@show maximum(filter(!ismissing, tmax_data));
@show prop_missing = count(ismissing, tmax_data) / length(tmax_data);

# Construct the spatial locations of the stations in a GP-friendly format.
# Combine with the elevation data.
get_vector(x::Tuple{A, B} where {A, B}) = [x[1], x[2]]
get_vector(x::Tuple{A, B, C} where {A, B, C}) = [x[1], x[2], x[3]]

lat_lon_mat = hcat(get_vector.(station_lat_lon_pairs)...);
space_inputs_mat = Float64.(vcat(lat_lon_mat, elevations' ./ 100));
const space_inputs = ColVecs(space_inputs_mat);

# Utility functionality for working with `ParameterHandling` a bit more conveniently.
function wrap_objective(f, θ)
    θ_flat, unflatten = flatten(θ)
    function f_flat(θ_flat)
        return f(ParameterHandling.value(unflatten(θ_flat)))
    end
    return f_flat, θ_flat, unflatten
end

function construct_training_outputs(test_idx::Vector{Int}, outputs::AbstractVector)
    train_outputs = deepcopy(outputs)
    train_outputs[test_idx] .= missing
    return test_idx, train_outputs
end

function test_indices_name(start_date::Date, finish_date::Date, Nte::Int)
    return joinpath(data_dir, "$start_date-$finish_date-$Nte.jld2")
end

function maybe_create_test_indices(
    rng::AbstractRNG,
    start_date::Date,
    finish_date::Date,
    Nte::Int,
    outputs::AbstractVector,
)
    fname = test_indices_name(start_date, finish_date, Nte)

    if !isfile(fname)
        @warn "CREATING NEW TEST INDICES FOR DATES $start_date-$finish_date-$Nte."
        present_idx = findall(!ismissing, outputs)
        test_indices = present_idx[randperm(rng, length(present_idx))[1:Nte]]
        FileIO.save(fname, Dict("test_indices" => test_indices))
    else
        @warn "NOT CREATING NEW TEST INDICES AS THEY ALREADY EXIST."
    end
end

function get_test_indices(start_date::Date, finish_date::Date, Nte::Int)
    fname = test_indices_name(start_date, finish_date, Nte)
    return FileIO.load(fname)["test_indices"]
end

# Make it possible to parse strings from the command line into Vector{Int}s.
function ArgParse.parse_item(::Type{Vector{Int}}, x::String)
    @assert x[1] == '[' && x[end] == ']'
    return CSV.read(IOBuffer(x[2:prevind(x, end)]); header=false, transpose=true).Column1
end

function bad_experiment_type_error(experiment_type::String)
    throw(error("Unknown experiment type: $experiment_type."))
end

function sor_data(inputs::RectilinearGrid, outputs::AbstractVector, M::Int)

    # Compute the K-means clustering of the data.
    println(size(get_space(inputs).X))
    cluster_indices = if size(get_space(inputs).X, 2) > M
        println("Using subset of data.")
        Random.seed!(123456)
        clustering_result = kmeans(get_space(inputs).X, M)
        cluster_indices = fill(0, M)
        for (n, a) in enumerate(assignments(clustering_result))
            if cluster_indices[a] == 0
                cluster_indices[a] = n
            end
        end
        cluster_indices
    else
        println("Using all data!")
        collect(1:M)
    end

    # Subset the SoR data.
    inputs_sor = RectilinearGrid(get_space(inputs)[cluster_indices], get_time(inputs))
    outputs_sor = collect(vec(
        reshape(outputs, length(get_space(inputs)), :)[cluster_indices, :],
    ))
    return inputs_sor, outputs_sor
end

# Default behaviour is Matern32 -- Matern52 experiments were added afterwards.
function build_kernel(::Val{:separable}, params::NamedTuple)
    return Separable(
        stretch(EQ(), params.λ_space),
        params.σ^2 * stretch(Matern32(), params.λ_time),
    )
end

function build_kernel(::Val{:sum_separable}, params::Vector{<:NamedTuple})
    length(params) != 2 && throw(error("Expected two sets of params."))
    k1 = build_kernel(Val(:separable), params[1])
    k2 = build_kernel(Val(:separable), params[2])
    return k1 + k2
end

function build_kernel(::Val{:separable_matern_52}, params::NamedTuple)
    return Separable(
        stretch(EQ(), params.λ_space),
        params.σ^2 * stretch(Matern52(), params.λ_time),
    )
end

function build_kernel(::Val{:sum_separable_matern_52}, params::Vector{<:NamedTuple})
    length(params) != 2 && throw(error("Expected two sets of params."))
    k1 = build_kernel(Val(:separable_matern_52), params[1])
    k2 = build_kernel(Val(:separable_matern_52), params[2])
    return k1 + k2
end

function build_kernel(params::NamedTuple{(:kernel_type, :params)})
    return build_kernel(params.kernel_type, params.params)
end

function parse_commandline()
    s = ArgParseSettings()
    @add_arg_table! s begin
        "--run-experiments"
            help = "true / false - run the experiments or not"
            arg_type = Bool
            default = true
        "--M"
            help = "number of pseudo-points"
            arg_type = Int
            default = 60
    end
    return parse_args(s)
end

const args = parse_commandline()

BLAS.set_num_threads(8);

# Need to be able to handle `Val`s for dispatch reasons.
function ParameterHandling.flatten(::Type{T}, ::Val{V}) where {T<:Real, V}
    v = T[]
    unflatten_to_Val(::Vector{T}) = Val{V}()
    return v, unflatten_to_Val
end

ParameterHandling.value(v::Val) = v

#
# Specify experiments to run.
#

const settings_fname = joinpath(metadata_dir, "settings.bson")

let
    mkpath(metadata_dir)

    separable_kernel_info_matern_32 = (
        name = "separable-Matern-32",
        value = (
            kernel_type = Val(:separable),
            params = (
                σ = positive(1.0),
                λ_space = fill(positive(1.0), 3),
                λ_time = positive(1e-2),
            ),
        ),
    )

    sum_separable_kernel_info_matern_32 = (
        name = "sum-separable-Matern32",
        value = (
            kernel_type = Val(:sum_separable),
            params = [
                (
                    σ = positive(0.7),
                    λ_space = fill(positive(1.0), 3),
                    λ_time = positive(1e-2),
                ),
                (
                    σ = positive(0.3),
                    λ_space = fill(positive(5.0), 3),
                    λ_time = positive(1e-1),
                ),
            ],
        ),
    )

    separable_kernel_info_matern_52 = (
        name = "separable-Matern-52",
        value = (
            kernel_type = Val(:separable_matern_52),
            params = (
                σ = positive(1.0),
                λ_space = fill(positive(1.0), 3),
                λ_time = positive(1e-2),
            ),
        ),
    )

    sum_separable_kernel_info_matern_52 = (
        name = "sum-separable-Matern52",
        value = (
            kernel_type = Val(:sum_separable_matern_52),
            params = [
                (
                    σ = positive(0.7),
                    λ_space = fill(positive(1.0), 3),
                    λ_time = positive(1e-2),
                ),
                (
                    σ = positive(0.3),
                    λ_space = fill(positive(5.0), 3),
                    λ_time = positive(1e-1),
                ),
            ],
        ),
    )

    wsave(
        settings_fname,
        :settings => vcat(

            dict_list(Dict(
                # Use the subset of regressors approximation.
                :experiment_type => ["SoD"],

                # Try both kinds of kernel.
                :kernel_info => [sum_separable_kernel_info_matern_52],

                # The total number of spatial locations to use.
                # :M => [1, 2, 5],
                :M => [10, 20, 50],
                # :M => [99],

                # The first date for the training data.
                :start_date => [Date(2000)],
            )),

            dict_list(Dict(
                # Use the new pseudo-point approximation.
                :experiment_type => ["pseudo-point"],

                # The total number of pseudo-points to use per time step.
                # :M => [1, 2, 5],
                :M => [10, 20, 50],
                # :M => [50],

                # Kernels to try, specified using the `ParameterHandling.jl` interface.
                :kernel_info => [sum_separable_kernel_info_matern_52],

                # The first date for the training data.
                :start_date => [Date(2000)],
            ),),
        ),
    )

end

#
# Run experiments. This is split into two sections:
# 1. training: this takes by far the longest amount of time. The entire model is saved as an
#     output from this step.
# 2. prediction: loads the results from step 1 and uses them to make predictions. Not run in
#     the training block because you don't want to have to re-run the training procedure to
#     make slightly different kinds of predictions. Not in the results block because it can
#     take a while, which is less than helpful when tweaking plots.
#

# if args["run-experiments"]

println("Running experiments")

mkpath(training_output_dir)

let
    settings = load(settings_fname)[:settings];

    for (n, setting) in enumerate(settings)

        println("Experiment $n / $(length(settings))")
        display(setting)
        println()

        experiment_type = setting[:experiment_type]

        # Compute the observation times.
        first_date = setting[:start_date]
        first_date_index = only(findall(tmax_dates .== first_date))

        time_inputs_as_dates = tmax_dates[first_date_index:end]
        time_inputs_as_range = RegularSpacing(0.0, 1.0, length(time_inputs_as_dates))
        inputs = RectilinearGrid(space_inputs, time_inputs_as_range)
        outputs_raw = collect(vec(tmax_data[:, first_date_index:end]))

        # If a train-test split hasn't already been created, create one and write it to
        # disk so that it is permanent. SHOULD BE COMMITTED IN THE GIT HISTORY!
        N_obs = length(filter(!ismissing, outputs_raw))
        Nte = Int(floor(0.1 * N_obs))
        rng = MersenneTwister(123456)
        finish_date = time_interval[2]
        maybe_create_test_indices(rng, first_date, finish_date, Nte, outputs_raw)

        # Construct the training data set.
        test_idx = get_test_indices(first_date, finish_date, Nte)
        test_idx, outputs_train_raw = construct_training_outputs(test_idx, outputs_raw)

        # Standardise the training data.
        m = mean(filter(!ismissing, outputs_train_raw))
        s = std(filter(!ismissing, outputs_train_raw))
        outputs_train = (outputs_train_raw .- m) ./ s
        @show typeof(outputs_train)
        println("Total number of training data is $(length(filter(!ismissing, outputs_train))).")
        println("Total number of data is $(length(filter(!ismissing, outputs_raw)))")
        @show mean(filter(!ismissing, outputs_train)), std(filter(!ismissing, outputs_train))

        # Create the objective function, wrapping it in `ParameterHandling.jl` stuff to
        # make it play nicely with Optim.
        objective, flat_initial_parameters, unflatten = if experiment_type == "SoD"

            inputs_sor, outputs_train_sor = sor_data(inputs, outputs_train, setting[:M])

            # Pick one point from each cluster.
            wrap_objective(
                θ -> begin
                    k = build_kernel(θ.k.kernel_type, θ.k.params)
                    f = to_sde(GP(k, GPC()))
                    fx = f(dropgrad(inputs_sor), θ.σ²)
                    return -logpdf(fx, dropgrad(outputs_train_sor))
                end,
                (
                    k = setting[:kernel_info].value,
                    σ² = bounded(0.5, 1e-2, 2.0),
                ),
            )
        elseif experiment_type == "pseudo-point"

            # Compute pseudo-point initialisation. Because Clustering doesn't let you
            # provide a seed to the kmeans++ initialisation, it's necessary to set the
            # global seed.
            Random.seed!(123456)
            clustering_result = kmeans(space_inputs.X, setting[:M])
            Z_r = clustering_result.centers

            wrap_objective(
                θ -> begin
                    k = build_kernel(θ.k.kernel_type, θ.k.params)
                    f = to_sde(GP(k, GPC()))
                    fx = f(dropgrad(inputs), θ.σ²)
                    z_r = ColVecs(θ.Z_r)
                    return -elbo(fx, dropgrad(outputs_train), z_r)
                end,
                (
                    k = setting[:kernel_info].value,
                    σ² = bounded(0.5, 1e-2, 2.0),
                    Z_r = fixed(Z_r),
                ),
            )
        else
            bad_experiment_type_error(experiment_type)
        end

        # Optimise the parameters.
        training_results = Optim.optimize(
            θ -> TemporalGPs.time_ad(Val(:disabled), "objective", objective, θ),
            θ -> Zygote.gradient(objective, θ)[1],
            flat_initial_parameters,
            LBFGS(
                alphaguess = LineSearches.InitialStatic(scaled=true, alpha=1),
                linesearch = LineSearches.BackTracking(),
                m=50,
            ),
            Optim.Options(
                x_tol = 1e-4,
                f_tol = 1e-8,
                g_tol = 1e-1,
                f_calls_limit = 1_000,
                g_calls_limit = 1_000,
                iterations = 1_000,
                show_trace = true,
                extended_trace = true,
                store_trace=true,
                show_every=10,
            );
            inplace=false,
        )

        display(training_results)
        println()

        wsave(
            joinpath(
                training_output_dir,
                savename(
                    if experiment_type == "exact"
                        Dict(
                            :experiment_type => experiment_type,
                            :kernel => setting[:kernel_info].name,
                            :start_date => first_date,
                        )
                    else
                        Dict(
                            :experiment_type => experiment_type,
                            :M => setting[:M],
                            :kernel => setting[:kernel_info].name,
                            :start_date => first_date,
                        )
                    end,
                    "jld2",
                ),
            ),
            Dict(
                "setting" => setting,
                "training_results" => training_results,
                "unflatten" => unflatten,
                "test_idx" => test_idx,
            ),
        )
    end
end

# end

#
# Make predictions at locations that we wish to view them, and compute RMSE on data etc.
#

println("Generating predictions.")

mkpath(prediction_output_dir)

let

    include(joinpath(srcdir(), "elevation_data.jl"))

    trained_models = collect_results(training_output_dir);

    for (n, trained_model_) in enumerate(eachrow(trained_models))

        println("n, $(trained_model_.path)")
        println(n)

        # Defensive copy. Something in here is mutating for some unknown reason.
        trained_model = deepcopy(trained_model_);

        # Construct the inputs.
        first_date = trained_model.setting[:start_date]
        first_date_index = only(findall(tmax_dates .== first_date))
        time_inputs_as_dates = tmax_dates[first_date_index:end]
        time_inputs_as_range = RegularSpacing(0.0, 1.0, length(time_inputs_as_dates))
        inputs = RectilinearGrid(space_inputs, time_inputs_as_range);

        # Reconstruct the training data.
        outputs_raw = collect(vec(tmax_data[:, first_date_index:end]))
        test_idx = trained_model.test_idx
        @show typeof(test_idx)
        _, outputs_train_raw = construct_training_outputs(test_idx, outputs_raw)

        # Standardise the training data.
        m = mean(filter(!ismissing, outputs_train_raw))
        s = std(filter(!ismissing, outputs_train_raw))
        outputs_train = (outputs_train_raw .- m) ./ s
        @show typeof(outputs_train)
        println("Total number of data is $(length(outputs_train)).")

        # Build the model.
        unflatten = trained_model.unflatten
        results = trained_model.training_results
        learned_parameters = ParameterHandling.value(unflatten(results.minimizer))

        k = build_kernel(learned_parameters.k)
        f = to_sde(GP(k, GPC()))

        # Get the elevation for each point in the lat-lon grid and return a new data set.
        lat_lon_el_grid = map(lat_lon_grid) do lat_lon
            lat, lon = lat_lon
            elevation = nearest_point(elevation_data_mat, lat, lon) / 100
            return (lat, lon, elevation)
        end

        # Make predictions at entire grid of points at which we might consider predicting.
        x_r_pr = ColVecs(hcat(get_vector.(lat_lon_el_grid)...))
        x_r = space_inputs
        pos = length(time_inputs_as_range) - 300

        # target_times = [1800, 3600];
        # trace_times = map(x -> x.metadata["time"], results.trace);
        # trace_indices_raw = map(x -> findfirst(trace_times .> x), target_times)
        # trace_indices = map(x -> x === nothing ? 1 : x, trace_indices_raw)

        experiment_type = trained_model.setting[:experiment_type]
        (preds_mat, station_preds) = if experiment_type ∈ ["SoD", "SoR"]

            inputs_sor, outputs_train_sor = sor_data(
                inputs, outputs_train, trained_model.setting[:M],
            )
            inputs_space_sor = get_space(inputs_sor)

            # Compute SoR posterior marginals at requested locations and time.
            # For some reason, my exact inference algorithm is broken at the time of
            # writing, but using the approximate inference algorithm with z_r = x_r yields
            # SoR inference.
            σ² = learned_parameters.σ²
            fx = f(inputs_sor, σ²)
            preds = @time approx_posterior_marginals(
                dtc, fx, outputs_train_sor, inputs_space_sor, x_r_pr, pos,
            )
            preds_mat = reshape(preds, N_lons, :)
            preds_mat = Normal.(mean.(preds_mat), std.(preds_mat) .+ sqrt(σ²))

            # Compute posterior marginals at the training data.
            fx = f(inputs_sor, learned_parameters.σ²)
            station_preds = @time approx_posterior_marginals(
                dtc, fx, outputs_train_sor, inputs_space_sor, x_r,
            )
            station_preds = Normal.(mean.(station_preds), std.(station_preds) .+ sqrt(σ²))

            (preds_mat, station_preds)
        elseif experiment_type == "pseudo-point"
            z_r = ColVecs(learned_parameters.Z_r)
            σ² = learned_parameters.σ²
            fx = f(inputs, σ²)

            # Compute approx posterior marginals at requested locations and time.
            preds = @time approx_posterior_marginals(
                dtc, fx, outputs_train, z_r, x_r_pr, pos,
            )
            preds_mat = reshape(preds, N_lons, :)
            preds_mat = Normal.(mean.(preds_mat), std.(preds_mat) .+ sqrt(σ²))

            # Compute approx. posterior marginals at the training data.
            station_preds = @time approx_posterior_marginals(
                dtc, fx, outputs_train, z_r, x_r,
            )
            station_preds = Normal.(mean.(station_preds), std.(station_preds) .+ sqrt(σ²))

            preds_mat, station_preds
        else
            bad_experiment_type_error(experiment_type)
        end

        # Compute the RSMSEs on the training data.
        train_idx = setdiff(eachindex(outputs_raw), test_idx)
        preds_train = station_preds # Test data is missing from training, so doesn't matter.
        preds_test = station_preds[test_idx]
        outputs_test = (outputs_raw[test_idx] .- m) ./ s
        @show size(preds_train), size(preds_test), size(outputs_test), size(outputs_train)
        smse_train = mean(abs2, filter(!ismissing, outputs_train .- mean.(preds_train)))
        smse_test = mean(abs2, filter(!ismissing, outputs_test .- mean.(preds_test)))

        # Compute the Posterior Predictive Log Probs (PPLP).
        present_train_idx = findall(!ismissing, outputs_train)
        present_test_idx = findall(!ismissing, outputs_test)
        @show size(present_train_idx), size(present_test_idx)
        pplp_train = mean(logpdf.(
            preds_train[present_train_idx],
            outputs_train[present_train_idx],
        ))
        pplp_test = mean(logpdf.(
            preds_test[present_test_idx],
            outputs_test[present_test_idx],
        ))

        # Write the results.
        @time wsave(
            joinpath(
                prediction_output_dir,
                savename(
                    if experiment_type == "exact"
                        Dict(
                            :experiment_type => experiment_type,
                            :kernel => trained_model.setting[:kernel_info].name,
                            :start_date => first_date,
                        )
                    else
                        Dict(
                            :experiment_type => experiment_type,
                            :M => trained_model.setting[:M],
                            :kernel => trained_model.setting[:kernel_info].name,
                            :start_date => first_date,
                        )
                    end,
                    "jld2",
                ),
            ),
            Dict(
                "trained_model" => trained_model,
                "preds_train" => preds_train,
                "preds_test" => preds_test,
                "preds_extrapolate" => preds_mat,
                "x_r_pr" => x_r_pr,
                "lat_lon_pairs" => lat_lon_grid,
                "smse_train" => smse_train,
                "smse_test" => smse_test,
                "pplp_train" => pplp_train,
                "pplp_test" => pplp_test,
            ),
        )
    end

end

#
# Visualise results.
#

# Load shoreline data.
# const shore_data = GeoDatasets.gshhg('f');

struct GridBox{Tlats<:Tuple{<:Real, <:Real}, Tlons<:Tuple{<:Real, <:Real}}
    lats::Tlats
    lons::Tlons
end

const box = GridBox(lat_interval, lon_interval);

function in_box(box::GridBox, lat::Real, lon::Real)
    return box.lats[1] <= lat && box.lats[2] > lat &&
        box.lons[1] <= lon && box.lons[2] > lon
end

let

    results = @time collect_results(prediction_output_dir);

    _, us_states = GADM.get("USA"; children=true);
    state_map = Dict(
        map(
            ((state_name, geom), ) -> state_name => geom,
            zip(us_states.NAME_1, us_states.geom),
        ),
    );

    _, can_states = GADM.get("CAN"; children=true);
    can_state_map = Dict(
        map(
            ((state_name, geom), ) -> state_name => geom,
            zip(can_states.NAME_1, can_states.geom),
        ),
    );

    country_geometries = [
        # GADM.get("AUT"),
        # GADM.get("FIN"),
        # GADM.get("DNK"),
        # GADM.get("ITA"),
        # GADM.get("LIE"),
        # GADM.get("NLD"),
        # GADM.get("NOR"),
        # GADM.get("POL"),
        # GADM.get("SWE"),
        # GADM.get("CHE"),
        # GADM.get("DEU"),
        (geom=state_map["Washington"],),
        (geom=can_state_map["British Columbia"],),
    ];

    # Compute some additional columns that will be used later.
    let
        optim_results = map(x -> x.training_results, results.trained_model)
        results.iters = Optim.iterations.(optim_results)
        results.timings = Optim.time_run.(optim_results)
        results.time_per_iter = results.timings ./ results.iters
        results.trace = getproperty.(optim_results, :trace)
    end;

    gr()

    for result in eachrow(results)

        preds_mat = result.preds_extrapolate
        trained_model = deepcopy(result.trained_model)

        # Construct the inputs.
        first_date = trained_model.setting[:start_date]
        first_date_index = only(findall(tmax_dates .== first_date))
        time_inputs_as_dates = tmax_dates[first_date_index:end]
        time_inputs_as_range = RegularSpacing(0.0, 1.0, length(time_inputs_as_dates))
        inputs = RectilinearGrid(space_inputs, time_inputs_as_range);

        # Get the pseudo-point locations.
        unflatten = trained_model.unflatten
        training_results = trained_model.training_results
        learned_parameters = ParameterHandling.value(unflatten(training_results.minimizer))

        @show learned_parameters.k
        @show result.smse_train, result.smse_test, result.pplp_train, result.pplp_test
        @show training_results.minimum

        experiment_type = trained_model.setting[:experiment_type]
        colour_scheme = :viridis;

        # Plot the posterior mean.
        let
            plt = plot(
                xlims=lon_interval,
                ylims=lat_interval,
                xticks=nothing,
                yticks=nothing,
                legend=:bottomleft,
            )

            heatmap!(
                plt, uk_lons_range, uk_lats_range, mean.(preds_mat);
                label="", c=cgrad(colour_scheme), right_margin=6Plots.mm,
            )

            # Plot national borders in the relevant area of Europe.
            for country in country_geometries
                plot!(plt, country.geom; fillcolor=nothing, label="")
            end

            scatter!(
                plt, getindex.(space_inputs, 2), getindex.(space_inputs, 1);
                label="Stations Locations",
                # color=:blue,
                markersize=5,
                marker=:square,
                color=colorant"#A95AA1",
            )

            if experiment_type == "pseudo-point"
                z_r = ColVecs(learned_parameters.Z_r)
                scatter!(
                    plt, getindex.(z_r, 2), getindex.(z_r, 1);
                    label="z",
                    # color=:red,
                    markersize=4,
                    color=colorant"#F5793A", # Orange
                )
            end

            # for (lons, lats) in shore_data

            #     # Filter out stuff outside of the box.
            #     points_in_box = filter(
            #         ((lat, lon), ) -> in_box(box, lat, lon),
            #         collect(zip(lats, lons)),
            #     )

            #     if length(points_in_box) > 0
            #         lats_f = first.(points_in_box)
            #         lons_f = last.(points_in_box)

            #         # Plot stuff that's in the box.
            #         plot!(plt, lons_f, lats_f; label="", color=:black, alpha=1.0)
            #     end
            # end

            mkpath(results_dir)
            savefig(
                plt,
                joinpath(
                    results_dir,
                    savename(
                        if experiment_type == "SoD"
                            Dict(
                                :experiment_type => experiment_type,
                                :M => trained_model.setting[:M],
                                :kernel => trained_model.setting[:kernel_info].name,
                                :start_date => first_date,
                            )
                        elseif experiment_type == "pseudo-point"
                            Dict(
                                :experiment_type => experiment_type,
                                :M => trained_model.setting[:M],
                                :kernel => trained_model.setting[:kernel_info].name,
                                :start_date => first_date,
                            )
                        else
                            bad_experiment_type_error(experiment_type)
                        end,
                        "posterior_mean.png",
                    ),
                )
            )
        end

        # Plot the posterior standard deviation.
        let
            plt = plot(
                xlims=lon_interval,
                ylims=lat_interval,
                xticks=nothing,
                yticks=nothing,
            )

            heatmap!(
                uk_lons_range, uk_lats_range, std.(preds_mat);
                label="", c=colour_scheme, right_margin=3Plots.mm,
            )

            # Plot national borders in the relevant area of Europe.
            for country in country_geometries
                plot!(plt, country.geom; fillcolor=nothing, label="", linecolor=:black)
            end

            scatter!(
                plt, getindex.(space_inputs, 2), getindex.(space_inputs, 1);
                label="Stations Locations",
                markersize=5,
                marker=:square,
                color=colorant"#A95AA1", # Purple
                # color=:blue,
            )

            if experiment_type == "pseudo-point"
                z_r = ColVecs(learned_parameters.Z_r)
                scatter!(
                    plt, getindex.(z_r, 2), getindex.(z_r, 1);
                    label="z",
                    # color=:red,
                    markersize=4,
                    color=colorant"#F5793A", # Orange
                )
            end

            # for (lons, lats) in shore_data

            #     # Filter out stuff outside of the box.
            #     points_in_box = filter(
            #         ((lat, lon), ) -> in_box(box, lat, lon),
            #         collect(zip(lats, lons)),
            #     )

            #     if length(points_in_box) > 0
            #         lats_f = first.(points_in_box)
            #         lons_f = last.(points_in_box)

            #         # Plot stuff that's in the box.
            #         plot!(plt, lons_f, lats_f; label="", color=:black, alpha=1.0)
            #     end
            # end

            mkpath(results_dir)
            savefig(
                plt,
                joinpath(
                    results_dir,
                    savename(
                        if experiment_type == "SoD"
                            Dict(
                                :experiment_type => experiment_type,
                                :M => trained_model.setting[:M],
                                :kernel => trained_model.setting[:kernel_info].name,
                                :start_date => first_date,
                            )
                        elseif experiment_type == "pseudo-point"
                            Dict(
                                :experiment_type => experiment_type,
                                :M => trained_model.setting[:M],
                                :kernel => trained_model.setting[:kernel_info].name,
                                :start_date => first_date,
                            )
                        else
                            bad_experiment_type_error(experiment_type)
                        end,
                        "posterior_std.png",
                    ),
                )
            )
        end
    end

    #
    # Build TexTable to pretty print the results.
    #

    my_format(x) = format(round(x; sigdigits=3); stripzeros=true)

    name_map = Dict(
        "SoD" => "SoD",
        "exact" => "Exact",
        "pseudo-point" => "P-P",
        # "separable" => "Separable",
        # "sum-separable" => "Sum-Separable",
        "separable-Matern-52" => "Separable",
        "sum-separable-Matern52" => "Sum-Separable",
    );

    # Group everything by date.
    results_df_unsorted = DataFrame(
        start_date = map(x -> x.setting[:start_date], results.trained_model),
        lml_approx = map(
            x -> my_format(x.training_results.minimum),
            results.trained_model,
        ),
        rsmse_test = sqrt.(results.smse_test),
        npplp_test = -results.pplp_test,
        rsmse_test_string = map(my_format, sqrt.(results.smse_test)),
        npplp_test_string = map(my_format, -results.pplp_test),
        row_keys = map(
            x -> begin
                M = x.setting[:M]
                experiment_name = name_map[x.setting[:experiment_type]] * "($M)"
                kernel_name = name_map[x.setting[:kernel_info].name]
                return experiment_name * " - " * kernel_name
            end,
            results.trained_model,
        ),
        M = map(x -> x.setting[:M], results.trained_model),
        experiment_type = map(x -> x.setting[:experiment_type], results.trained_model),
        kernel_name = map(x -> x.setting[:kernel_info].name, results.trained_model),
        run_time = map(x -> Optim.time_run(x.training_results), results.trained_model),
        time_per_iter = results.time_per_iter,
    );

    results_df = sort(results_df_unsorted, :M);

    for group in groupby(results_df, :start_date)

        headers = [
            "RSMSE",
            "NPPLP",
        ]
        col_data = [
            group.rsmse_test_string,
            group.npplp_test_string,
        ]
        cols = [
            TableCol(header, collect(group.row_keys), collect(data)) for
            (header, data) in zip(headers, col_data)
        ]
        results_tex_table = hcat(cols...)

        # Generate paper-ready tables of results, comparing each model.
        start_date = only(unique(group.start_date))
        open(joinpath(results_dir, "performance_results_date=$start_date.txt"), "w") do io
            write(io, to_ascii(results_tex_table))
        end
        open(joinpath(results_dir, "performance_results_date=$start_date.tex"), "w") do io
            write(io, to_tex(results_tex_table))
        end
    end

    pgfplotsx();

    # Wessel's colour pallete that is friendly to those suffering colour-blindness.
    colours = [
        colorant"#4BA6FB", # Modified blue (original was #85C0F9)
        colorant"#A95AA1", # Pink
        colorant"#000000", # Black
        colorant"#F5793A", # Orange
    ]
    # colours = [:black, :red, :blue, :cyan]
    shapes = [:utriangle, :diamond, :square, :circle, :cross]

    # Generate plot. This should be roughly half a page width.
    width_inches = 2.8
    width_px = width_inches * Plots.DPI
    height_px = width_px / 2

    # Build speed-accuracy graphs for each start date in the results.
    for date_group in groupby(results_df, :start_date)
        rsmse_plot = plot(
            size=(width_px, height_px),
            xscale=:log10,
            legend=:topright,
            background_color_legend=RGBA(1, 1, 1, 0),
            foreground_color_legend=RGBA(1, 1, 1, 0),
        )
        npplp_plot = plot(
            size=(width_px, height_px),
            xscale=:log10,
            legend=nothing,
            background_color_legend=RGBA(1, 1, 1, 0),
            foreground_color_legend=RGBA(1, 1, 1, 0),
        )
        start_date = only(unique(date_group.start_date))
        for (n, group) in enumerate(groupby(date_group, [:experiment_type, :kernel_name]))
            experiment_type = only(unique(group.experiment_type))
            kernel_name = only(unique(group.kernel_name))
            label = name_map[experiment_type] * ", " * name_map[kernel_name]
            plot!(
                rsmse_plot, group.run_time, group.rsmse_test;
                xlabel="Run Time (s)",
                ylabel="RSMSE",
                label=label,
                color=colours[n],
                marker=shapes[n],
                markerstrokealpha=0,
                markersize=3,
                markercolor=colours[n],
                tickfontsize=nothing,
                legendfontsize=nothing,
            )
            plot!(
                npplp_plot, group.run_time, group.npplp_test,
                xlabel="Run Time (s)",
                ylabel="NPPLP",
                label=label,
                color=colours[n],
                marker=shapes[n],
                markerstrokealpha=0,
                markersize=3,
                markercolor=colours[n],
                tickfontsize=nothing,
                legendfontsize=nothing,
            )
        end

        savefig(rsmse_plot, joinpath(results_dir, "rsmse-tradeoff_date-$start_date.tikz"))
        savefig(npplp_plot, joinpath(results_dir, "pplp-tradeoff_date-$start_date.tikz"))
    end

    # Build time-per-iteration graphs.
    let
        for date_group in groupby(results_df, :start_date)
            start_date = only(unique(date_group.start_date))
            plt = plot(
                yscale=:log10,
                size=(width_px, height_px),
                legend=:bottomright,
                background_color_legend=RGBA(1, 1, 1, 0),
                foreground_color_legend=RGBA(1, 1, 1, 0),
            )
            for (n, group) in enumerate(groupby(date_group, [:experiment_type, :kernel_name]))
                experiment_type = only(unique(group.experiment_type))
                kernel_name = only(unique(group.kernel_name))
                label =
                plot!(
                    plt, group.M, group.time_per_iter;
                    xlabel="M",
                    ylabel="Time Per Iteration",
                    label=name_map[experiment_type] * ", " * name_map[kernel_name],
                    color=colours[n],
                    marker=shapes[n],
                    markerstrokealpha=0,
                    markersize=2,
                    markercolor=colours[n],
                    tickfontsize=nothing,
                    legendfontsize=6,
                )
            end
            savefig(plt, joinpath(results_dir, "time-per-iteration-$start_date.tikz"))
        end
    end

end


text:
using LightGraphs, GraphPlot

const AbstractLattice = AbstractGraph{Int64}
"""
    Lattice <: AbstractLattice

Abstract supertype for all lattices.
"""
abstract type Lattice <: AbstractLattice end

"""
    NdLattice{N} <: Lattice

Z^`N` lattice type. Generic lattice with `N` directions.
"""
struct NdLattice{N} <: Lattice
    L::SimpleGraph{Int64}
    shape::Tuple{Vararg{Int,N}}

    Vint::Tuple{Vararg{Vector{Int64},N}}
    Vext::Tuple{Vararg{Vector{Int64},N}}

    # Boundary conditions
    pbc::Bool
end

"""
    NdLattice(shape; periodic = false)

Construct an `N`-dimensional cubic lattice of type `NdLattice{N}`.

# Arguments
* `shape`: Tuple of sizes of all dimensions `(nx, ny, ...)`.
"""
function NdLattice(shape::Tuple{Vararg{Int,N}};periodic::Bool = false) where N
    @assert N > 0 "The lattice must be positive-dimensional."
    if N == 1
        return NdLattice{1}(grid(collect(shape);periodic=periodic),shape,([1],),([shape[end]],),periodic)
    else
        L = grid(collect(shape);periodic=periodic);
        lids = LinearIndices(shape)
        symbids = [1;fill(:(:), N-1)]
        Vint = Tuple([Core.eval(CornerSpaceRenorm,Expr(:ref, :($lids), circshift(symbids,d)...))[:] for d in 0:N-1])
        symbids = [:end;fill(:(:), N-1)]
        Vext = Tuple([Core.eval(Main,Expr(:ref, :($lids), circshift(symbids,d)...))[:] for d in 0:N-1])

        return NdLattice{N}(L,shape,Vint,Vext,periodic)
    end
end

Base.show(io::IO, L::NdLattice) = L.pbc ? print(io,join([string(s) for s in L.shape], " x "), " NdLattice with PBC") : print(io,join([string(s) for s in L.shape], " x "), " NdLattice with OBC");
function Base.show(io::IO, ::MIME"text/plain", L::NdLattice)
    print(io,string(typeof(L)),":\n")
    print(io,"  Shape: ",join([string(s) for s in L.shape], " x "),"\n")
    print(io,"  Boundaries:\n")
    for d in 1:length(L.shape)
        print(io,"  d = ",d,":\n")
        print(io,"    Input vertices: ",L.Vint[d],"\n")
        print(io,"    Output vertices: ",L.Vext[d],"\n")
    end

    if L.pbc
        print(io,"  Boundary conditions: periodic")
    else
        print(io,"  Boundary conditions: open")
    end
end
Base.show(io::IO, ::MIME"application/prs.juno.inline", x::NdLattice) = x

# Extend AbstractGraph{Int64} methods to all Lattice subtypes.
@inline Base.eltype(L::Lattice) = Base.eltype(L.L);
@inline LightGraphs.vertices(L::Lattice) = LightGraphs.vertices(L.L);
@inline LightGraphs.edges(L::Lattice) = LightGraphs.edges(L.L);
@inline LightGraphs.nv(L::Lattice) = LightGraphs.nv(L.L);
@inline LightGraphs.ne(L::Lattice) = LightGraphs.ne(L.L);
@inline LightGraphs.has_edge(L::Lattice, s::Int64, d::Int64) = LightGraphs.has_edge(L.L,s,d);
@inline LightGraphs.has_edge(L::Lattice, e::AbstractEdge{Int64}) = LightGraphs.has_edge(L.L,e);
@inline LightGraphs.has_vertex(L::Lattice, v::Int64) = LightGraphs.has_vertex(L.L,v);
@inline LightGraphs.inneighbors(L::Lattice, v::Int64) = LightGraphs.inneighbors(L.L,v);
@inline LightGraphs.outneighbors(L::Lattice, v::Int64) = LightGraphs.outneighbors(L.L,v);
@inline LightGraphs.is_directed(L::Lattice) = false;
@inline LightGraphs.is_directed(T::Type{L}) where {L<:Lattice} = false;

GraphPlot.gplot(L::Lattice; kwargs...) = GraphPlot.gplot(L.L; kwargs...)
GraphPlot.gplot(L::Lattice, locs_x_in::Vector{R}, locs_y_in::Vector{R}; kwargs...) where R<:Real = GraphPlot.gplot(L.L, locs_x_in, locs_y_in; kwargs...)

function rem_boundaries!(L::NdLattice{N}, d::Int) where N
    if L.shape[d] > 2
        eb = [e for e in edges(L) if e.src ∈ L.Vint[d] && e.dst ∈ L.Vext[d]]
        for e in eb
            rem_edge!(L.L,e)
        end
    end
    return L
end

function add_boundaries!(L::NdLattice{N}, d::Int) where N
    for i in 1:length(L.Vext[d])
        add_edge!(L.L,Edge(L.Vext[d][i],L.Vint[d][i]))
    end
    return L
end

"""
    union(L1, L2, d)

Merge two `NdLattice{N}`s along the `d`-th dimension.
"""
function Base.union(L1::NdLattice{N},L2::NdLattice{N},d::Int) where N
    @assert d <= N "Merging direction should be smaller than that of the lattices."
    length(L1.Vext[d]) == length(L2.Vint[d]) || throw(DimensionMismatch("Lattices cannot be merged in this direction."))
    @assert L1.pbc == L2.pbc "Cannot merge a periodic and an open lattices together."
    pbc = L1.pbc
    L1 = deepcopy(L1)
    L2 = deepcopy(L2)
    # Remove boundary edges along the merging dimension if periodic
    if pbc
        rem_boundaries!(L1, d)
        rem_boundaries!(L2, d)
    end

    L = blockdiag(L1.L,L2.L)
    for i in 1:length(L1.Vext[d])
        add_edge!(L,Edge(L1.Vext[d][i],nv(L1)+L2.Vint[d][i]))
    end
    shape = collect(L1.shape);
    shape[d] += L2.shape[d]
    d⊥ = [_d for _d in 1:N if _d!= d]

    Vint = collect(L1.Vint)
    Vext = [nv(L1) .+ vs for vs in L2.Vext]
    for _d in d⊥
        Vint[_d] = [L1.Vint[_d]; nv(L1) .+ L2.Vint[_d]]
        Vext[_d] = [L1.Vext[_d]; nv(L1) .+ L2.Vext[_d]]
    end

    lat = NdLattice{N}(L,Tuple(shape),Tuple(Vint),Tuple(Vext),pbc)
    # Add edges between boundaries if periodic
    if pbc
        add_boundaries!(lat, d)
    end
    return lat
end

"""
    pbc_from_obc(L, d)

Construct an `NdLattice{N}` with periodic boundary conditions from an
`NdLattice{N}` `L` with open boundary conditions.
"""
function pbc_from_obc(Lobc::NdLattice{N}) where N
    @assert !Lobc.pbc "The input lattice already has periodic boundary conditions."
    Lpbc = NdLattice{N}(deepcopy(Lobc.L), deepcopy(Lobc.shape),
                        deepcopy(Lobc.Vint), deepcopy(Lobc.Vext), true)
    for d in 1:N
        add_boundaries!(Lpbc, d)
    end
    return Lpbc
end

"""
    obc_from_pbc(L, d)

Construct an `NdLattice{N}` with open boundary conditions from an `NdLattice{N}`
`L` with periodic boundary conditions.
"""
function obc_from_pbc(Lpbc::NdLattice{N}) where N
    @assert Lpbc.pbc "The input lattice already has open boundary conditions."
    Lobc = NdLattice{N}(deepcopy(Lpbc.L), deepcopy(Lpbc.shape),
                        deepcopy(Lpbc.Vint), deepcopy(Lpbc.Vext), false)
    for d in 1:N
        rem_boundaries!(Lobc, d)
    end
    return Lobc
end

"""
    AbstractSystem

Abstract supertype for all systems.
"""
abstract type AbstractSystem end

"""
    NdSystem <: AbstractSystem

Quantum dissipative system defined on a `NdLattice`.
"""
struct NdSystem{N,L<:NdLattice{N},B<:Basis,
                O1<:AbstractOperator{B,B},O2<:AbstractOperator{B,B},
                O3<:AbstractOperator{B,B},O4<:AbstractOperator{B,B}} <: AbstractSystem
    # Lattice
    lattice::L
    # Bases
    gbasis::B
    # Operators
    H::O1
    trate::Tuple{Vararg{ComplexF64,N}}
    Htint::Tuple{Vararg{Vector{O2},N}}
    Htext::Tuple{Vararg{Vector{O2},N}}
    J::Vector{O3}
    observables::Vector{Dict{String,O4}}
end

"""
    NdSystem(lat, H, trate, lHt, J, obs)

Construct a `NdSystem` from a `NdLattice` and operators defining the model.
# Arguments
* `lat`: `NdLattice{N}`.
* `H`: Hamiltonian of the full system.
* `trate`: tunneling rate or N-`Tuple` of tunneling rates along each dimension.
* `lHt`: local tunneling operator.
* `J`: jump operators of the full system.
* `obs`: Observables to be transformed to the corner space together with the
Liouvillian. Can be a `Dict{String,<local operator type>}` of local operators or
an array of `nv(lat)` `Dict{String,<global operator type>}`s of operators in the
global basis.
"""
function NdSystem(lat::NdLattice{M},H::O1,trate,lHt,J::Vector{O3},obs=missing) where {M,M1,M2,N,S,T<:Number,LB<:Basis,
                                                       B<:CompositeBasis{S,Tuple{Vararg{LB,N}}},
                                                       O1<:AbstractOperator{B,B},
                                                       O3<:AbstractOperator{B,B}}
    gbasis = H.basis_l;
    @assert nv(lat) == length(gbasis.bases)
    @assert H.basis_r == H.basis_l == gbasis
    @assert first(J).basis_l == first(J).basis_r == gbasis

    _trate = begin
        if typeof(trate) <: Number
            Tuple(fill(ComplexF64(trate), M))
        elseif typeof(trate)<:Tuple && eltype(trate)<:Number
            @assert length(trate) == M "Number of coupling rates must match number of dimensions"
            ComplexF64.(trate)
        else
            error("The coupling rate must either be a Number or a Tuple")
        end
    end

    Htint, Htext = begin
        if typeof(lHt) <: AbstractOperator{LB,LB}
            Tuple([embed(gbasis,v,lHt) for v in lat.Vint[d]] for d in 1:M), Tuple([embed(gbasis,v,lHt) for v in lat.Vext[d]] for d in 1:M)
        elseif typeof(lHt)<:Tuple
            @assert length(lHt) == M "Number of coupling operators must match number of dimensions"
            if eltype(lHt)<:AbstractOperator{LB,LB}
                Tuple([embed(gbasis,v,lHt[d]) for v in lat.Vint[d]] for d in 1:M), Tuple([embed(gbasis,v,lHt[d]) for v in lat.Vext[d]] for d in 1:M)
            elseif eltype(lHt)<:Vector{O} where {O<:AbstractOperator{LB,LB}}
                Tuple([embed(gbasis,v,h) for h in lHt[d], v in lat.Vint[d]][:] for d in 1:M), Tuple([embed(gbasis,v,h) for h in lHt[d], v in lat.Vext[d]][:] for d in 1:M)
            else
                error("The tuple of coupling operators must contain either local operators or a vector or local operators")
            end
        else
            error("The coupling operator must be a local operator, a tuple of local operators or a tuple of vectors of local operators")
        end
    end

    _obs = begin
        if ismissing(obs)
             [Dict{String,typeof(H)}() for i in 1:nv(lat)]
         elseif typeof(obs) <: Vector{Dict{String,O}} where {O<:AbstractOperator{B,B}}
             obs
         elseif typeof(obs) <: Dict{String,O} where {O<:AbstractOperator{LB,LB}}
             [Dict(name => embed(gbasis,i,lop) for (name,lop) in obs) for i in 1:nv(lat)]
         elseif typeof(obs) <: Dict{String,O} where {O<:AbstractOperator{B,B}}
             [obs]
         else
             error("If defined, argument obs takes either a Dict{String,O} of local operators or a vector of Dict{String,O} of global operators")
         end
    end
    O4 = eltype(_obs).parameters[2]

    return NdSystem{M,NdLattice{M},B,O1,eltype(first(Htint)),O3,O4}(lat,gbasis,H,_trate,Htint,Htext,J,_obs)
end

function Base.show(io::IO, s::NdSystem)
    s.lattice.pbc ? print(io,join([string(_s) for _s in s.lattice.shape], " x "), " NdSystem with PBC") : print(io,join([string(_s) for _s in s.lattice.shape], " x "), " NdSystem with OBC");
    if typeof(s.gbasis) <: CompositeBasis && all([s.gbasis.bases[i] == s.gbasis.bases[1] for i in 1:length(s.gbasis.bases)])
        print(io, " and local basis: ")
        Base.show(io, first(s.gbasis.bases))
    else
        print(io, " and global basis: ")
        Base.show(io, s.gbasis)
    end
end
function Base.show(io::IO, ::MIME"text/plain", s::NdSystem)
    print(io, "NdSystem:\n")
    print(io,"  Lattice: ")
    Base.show(io,s.lattice)
    if typeof(s.gbasis) <: CompositeBasis && all([s.gbasis.bases[i] == s.gbasis.bases[1] for i in 1:length(s.gbasis.bases)])
        print(io, "\n  Homogeneous local basis: ")
        Base.show(io, first(s.gbasis.bases))
    else
        print(io, "\n  Global basis: ")
        Base.show(io, s.gbasis)
    end
    print(io,"\n  Hamiltonian: ",join([string(_s) for _s in size(s.H)], " x ")," ",typeof(s.H).name,"\n")
    print(io,"  ",length(s.J)," lindblad ops: ",join([string(_s) for _s in size(first(s.J))], " x ")," ",typeof(first(s.J)).name,"\n")
    print(io, "  Tunneling rates: \n")
    for d in 1:length(s.lattice.shape)
        print(io,"    d = ",d,": ",s.trate[d],"\n")
    end
    obs_names = unique(keys.(s.observables))
    if length(obs_names) == 1
        print(io, "  Local observables: ",join([obs for obs in obs_names[1]], ", "))
    end
end
Base.show(io::IO, ::MIME"application/prs.juno.inline", x::NdSystem) = x

GraphPlot.gplot(s::AbstractSystem; kwargs...) = GraphPlot.gplot(s.lattice.L; kwargs...)
GraphPlot.gplot(s::AbstractSystem, locs_x_in::Vector{R}, locs_y_in::Vector{R}; kwargs...) where R<:Real = GraphPlot.gplot(s.lattice.L, locs_x_in, locs_y_in; kwargs...)

"""
    plot_system(s; kwargs...)

Plot the lattice on which `s` is supported, equivalent to `GraphPlot.gplot(s.lattice.L; kwargs...)`.
# Arguments
See documentation of [`GraphPlot.gplot`](@ref).
"""
plot_system(s::AbstractSystem; kwargs...) = GraphPlot.gplot(s.lattice.L; kwargs...)

"""
    plot_system(s, locs_x_in, locs_y_in; kwargs...)

Equivalent to `GraphPlot.gplot(s.lattice.L, locs_x_in, locs_y_in; kwargs...)`.
# Arguments
See documentation of [`GraphPlot.gplot`](@ref).
"""
plot_system(s::AbstractSystem, locs_x_in::Vector{R}, locs_y_in::Vector{R}; kwargs...) where R<:Real = GraphPlot.gplot(s.lattice.L, locs_x_in, locs_y_in; kwargs...)

"""
    CornerBasis <: Basis

Corner basis.
"""
struct CornerBasis <: Basis
    shape::Vector{Int64}
    M::Int64
    local_shapes::Vector{Vector{Int64}}
end

"""
    CornerBasis(M)

Construct a corner basis with corner dimension `M`.
"""
CornerBasis(M::Int64) = CornerBasis([M], M, [[M]])
"""
    CornerBasis(b1, b2, M)

Construct a corner basis with corner dimension `M` from two basis.
"""
CornerBasis(b1::Basis, b2::Basis, M::Int64) = CornerBasis([M], M, [b1.shape, b2.shape])
CornerBasis(b1::CornerBasis, b2::CornerBasis, M::Int64) = CornerBasis([M], M, vcat(b1.local_shapes, b2.local_shapes))

CornerBasis(b::Basis, M::Int64) = CornerBasis([M], M, [b.shape])

import Base: ==
==(b1::CornerBasis, b2::CornerBasis) = (b1.M .== b2.M) &&
                                       (b1.shape == b2.shape) &&
                                       all(b1.local_shapes .== b2.local_shapes)

Base.show(io::IO, cb::CornerBasis) = print(io,"CornerBasis with dimension M = ",cb.M);
Base.show(io::IO, ::MIME"application/prs.juno.inline", x::CornerBasis) = x

ext: jl


text:
module TestCSets
using Test

using Catlab, Catlab.Theories, Catlab.Graphs,
  Catlab.CategoricalAlgebra, Catlab.CategoricalAlgebra.FinSets
using Catlab.Graphs.BasicGraphs: TheoryGraph

# C-set morphisms
#################

# Constructors and accessors.
g, h = Graph(4), Graph(2)
add_edges!(g, [1,2,3], [2,3,4])
add_edges!(h, [1,2], [2,1])
α = CSetTransformation((V=[1,2,1,2], E=[1,2,1]), g, h)
@test components(α) == (V=α[:V], E=α[:E])
@test α[:V] isa FinFunction{Int} && α[:E] isa FinFunction{Int}
@test α[:V](3) == 1
@test α[:E](2) == 2

α′ = CSetTransformation(g, h, V=[1,2,1,2], E=[1,2,1])
@test components(α′) == components(α)
α′′ = CSetTransformation(g, h, V=FinFunction([1,2,1,2]), E=FinFunction([1,2,1]))
@test components(α′′) == components(α)

# Naturality.
@test is_natural(α)
β = CSetTransformation((V=[1,2,1,2], E=[1,1,1]), g, h)
@test !is_natural(β)
β = CSetTransformation((V=[2,1], E=[2,1]), h, h)
@test is_natural(β)

# Category of C-sets.
@test dom(α) === g
@test codom(α) === h
@test compose(α,β) == CSetTransformation((V=α[:V]⋅β[:V], E=α[:E]⋅β[:E]), g, h)
@test force(compose(id(g), α)) == α
@test force(compose(α, id(h))) == α

# Limits
#-------

# Terminal object in Graph: the self-loop.
term = Graph(1)
add_edge!(term, 1, 1)
lim = terminal(Graph)
@test ob(lim) == term
@test delete(lim, g) == CSetTransformation((V=fill(1,4), E=fill(1,3)), g, term)

# Products in Graph: unitality.
lim = product(g, term)
@test ob(lim) == g
@test force(proj1(lim)) == force(id(g))
@test force(proj2(lim)) ==
  CSetTransformation((V=fill(1,4), E=fill(1,3)), g, term)

# Product in Graph: two directed intervals (Reyes et al 2004, p. 48).
I = Graph(2)
add_edge!(I, 1, 2)
prod = ob(product(I, I))
@test nv(prod) == 4
@test ne(prod) == 1
@test src(prod) != tgt(prod)

# Product in Graph: deleting edges by multiplying by an isolated vertex.
g = Graph(4)
add_edges!(g, [1,2,3], [2,3,4])
g0 = ob(product(g, Graph(1)))
@test nv(g0) == nv(g)
@test ne(g0) == 0

# Product in Graph: copying edges by multiplying by the double self-loop.
cycle2 = Graph(1)
add_edges!(cycle2, [1,1], [1,1])
lim = product(g, cycle2)
g2 = ob(lim)
@test nv(g2) == nv(g)
@test ne(g2) == 2*ne(g)
@test src(g2) == repeat(src(g), 2)
@test tgt(g2) == repeat(tgt(g), 2)
α = CSetTransformation((V=[2,3], E=[2]), I, g)
β = CSetTransformation((V=[1,1], E=[2]), I, cycle2)
γ = pair(lim, α, β)
@test force(γ⋅proj1(lim)) == α
@test force(γ⋅proj2(lim)) == β

# Equalizer in Graph from (Reyes et al 2004, p. 50).
g, h = Graph(2), Graph(2)
add_edges!(g, [1,2], [2,1])
add_edges!(h, [1,2,2], [2,1,1])
ϕ = CSetTransformation((V=[1,2], E=[1,2]), g, h)
ψ = CSetTransformation((V=[1,2], E=[1,3]), g, h)
@test is_natural(ϕ) && is_natural(ψ)
eq = equalizer(ϕ, ψ)
@test ob(eq) == I
@test force(incl(eq)[:V]) == FinFunction([1,2])
@test force(incl(eq)[:E]) == FinFunction([1], 2)

# Pullback in Graph from (Reyes et al 2004, p. 53).
# This test exposes an error in the text: using their notation, there should be
# a second edge between the vertices (1,3) and (1,4).
g0, g1, g2 = Graph(2), Graph(3), Graph(2)
add_edges!(g0, [1,1,2], [1,2,2])
add_edges!(g1, [1,2,3], [2,3,3])
add_edges!(g2, [1,2,2], [1,2,2])
ϕ = CSetTransformation((V=[1,2,2], E=[2,3,3]), g1, g0)
ψ = CSetTransformation((V=[1,2], E=[1,3,3]), g2, g0)
@test is_natural(ϕ) && is_natural(ψ)
lim = pullback(ϕ, ψ)
@test nv(ob(lim)) == 3
@test sort!(collect(zip(src(ob(lim)), tgt(ob(lim))))) ==
  [(2,3), (2,3), (3,3), (3,3)]
@test is_natural(proj1(lim)) && is_natural(proj2(lim))

# Colimits
#---------

# Initial object in graph: the empty graph.
colim = initial(Graph)
@test ob(colim) == Graph()
@test create(colim, g) == CSetTransformation((V=Int[], E=Int[]), Graph(), g)

# Coproducts in Graph: unitality.
g = Graph(4)
add_edges!(g, [1,2,3], [2,3,4])
colim = coproduct(g, Graph())
@test ob(colim) == g
@test force(coproj1(colim)) == force(id(g))
@test force(coproj2(colim)) ==
  CSetTransformation((V=Int[], E=Int[]), Graph(), g)

# Coproduct in Graph.
h = Graph(2)
add_edges!(h, [1,2], [2,1])
colim = coproduct(g, h)
coprod = ob(colim)
@test nv(coprod) == 6
@test src(coprod) == [1,2,3,5,6]
@test tgt(coprod) == [2,3,4,6,5]
α = CSetTransformation((V=[1,2,1,2], E=[1,2,1]), g, h)
β = id(h)
γ = copair(colim, α, β)
@test force(coproj1(colim)⋅γ) == α
@test force(coproj2(colim)⋅γ) == force(β)

# Coequalizer in Graph: collapsing a segment to a loop.
g = Graph(2)
add_edge!(g, 1, 2)
α = CSetTransformation((V=[1], E=Int[]), Graph(1), g)
β = CSetTransformation((V=[2], E=Int[]), Graph(1), g)
@test is_natural(α) && is_natural(β)
coeq = coequalizer(α, β)
@test ob(coeq) == ob(terminal(Graph))
@test force(proj(coeq)[:V]) == FinFunction([1,1])
@test force(proj(coeq)[:E]) == FinFunction([1])

# Pushout in Graph from (Reyes et al 2004, p. 59).
α = CSetTransformation((V=[2], E=Int[]), Graph(1), g)
β = CSetTransformation((V=[1], E=Int[]), Graph(1), ob(terminal(Graph)))
@test is_natural(α) && is_natural(β)
colim = pushout(α, β)
@test nv(ob(colim)) == 2
@test src(ob(colim)) == [1,2]
@test tgt(ob(colim)) == [2,2]
@test is_natural(coproj1(colim)) && is_natural(coproj2(colim))

# Attributed C-set morphisms
############################

@present TheoryWeightedGraph <: TheoryGraph begin
  Weight::Data
  weight::Attr(E,Weight)
end
const WeightedGraph = ACSetType(TheoryWeightedGraph, index=[:src,:tgt])

# Constructors and accessors.
g, h = WeightedGraph{Float64}(2), WeightedGraph{Float64}(4)
add_edge!(g, 1, 2, weight=2.0)
add_edges!(h, [1,2,3], [2,3,4], weight=[1.0,2.0,3.0])
α = ACSetTransformation((V=[2,3], E=[2]), g, h)
@test length(components(α)) == 2

# Naturality.
@test is_natural(α)
β = ACSetTransformation((V=[2,1], E=[2]), g, h)
@test !is_natural(β) # Preserves weight but not graph homomorphism
β = ACSetTransformation((V=[1,2], E=[1]), g, h)
@test !is_natural(β) # Graph homomorphism but does not preserve weight

# Colimits
#---------

@present TheoryLabeledGraph <: TheoryGraph begin
  Label::Data
  vlabel::Attr(V,Label)
  elabel::Attr(E,Label)
end
const LabeledGraph = ACSetType(TheoryLabeledGraph, index=[:src,:tgt])

# Initial labeled graph.
@test ob(initial(LabeledGraph{Symbol})) == LabeledGraph{Symbol}()

# Coproduct of labeled graphs.
g = LabeledGraph{Symbol}()
add_vertices!(g, 2, vlabel=[:u,:v])
add_edge!(g, 1, 2, elabel=:e)
h = LabeledGraph{Symbol}()
add_vertex!(h, vlabel=:u)
add_edge!(h, 1, 1, elabel=:f)
coprod = ob(coproduct(g, h))
@test subpart(coprod, :vlabel) == [:u, :v, :u]
@test subpart(coprod, :elabel) == [:e, :f]

# Pushout of labeled graph.
g0 = LabeledGraph{Symbol}()
add_vertex!(g0, vlabel=:u)
α = ACSetTransformation((V=[1], E=Int[]), g0, g)
β = ACSetTransformation((V=[1], E=Int[]), g0, h)
@test is_natural(α) && is_natural(β)
colim = pushout(α, β)
@test src(ob(colim)) == [1,1]
@test tgt(ob(colim)) == [2,1]
@test subpart(ob(colim), :vlabel) == [:u, :v]
@test subpart(ob(colim), :elabel) == [:e, :f]

α′ = ACSetTransformation(V=[2], E=Int[], g0, g)
@test !is_natural(α′) # Vertex labels don't match.
@test_throws ErrorException pushout(α′, β)

end


text:
using PyCall
py"10*10"             #> 100
@pyimport math
math.sin(math.pi / 2) #> 1.0


text:
export periodize

periodize(shp::Shape{N}, A::AbstractMatrix{<:Real}, ∆range::Shape{N}) where {N} = periodize(shp, SMatrix{N,N}(A), ∆range)

function periodize(shp::S,  # shape to periodize
                   A::SMatrix{N,N,<:Real},  # columns: primitive vectors of Bravais lattice
                   ∆range::Shape{N}  # range of translation vectors; boundaries included
				   ) where {N,S<:Shape{N}}
    # R = n₁a₁ + n₂a₂ + n₃a₃ is a translation vector for A = [a₁ a₂ a₃].  Find the ranges of
    # n₁, n₂, n₃ to test the inclusion of R in the Shape ∆range.
    #
    # The minimum and maximum nᵢ's to examine can be found by findfirstg nᵢ's at the corners of
    # the bounding box of ∆range.  This can be easily seen by unitary transformation of aᵢ's
    # into the Cartesian directions.  The transformation transforms the lattice planes into
    # the x-, y-, z-normal planes, and the bounding box into a parallelpiped.  The nᵢ's at
    # the corners of the original bounding box corresponds to the Cartesian indices of the
    # corners of the parallelpiped after transformation, and the minimum and maximum nᵢ's
    # are the Cartesian indices of the corners of the bounding box of this post-transform
    # parallelpiped.  Because the post-transform parallelpiped is enclosed by its bounding
    # box, all the lattice points within the post-transform parallelpiped are also within
    # the bounding box of the parallelpiped, whose corner indices are the minimum and
    # maximum nᵢ's.
    ∆bound = bounds(∆range)  # (SVector{3}, SVector{3}): bounds of translation range
    nmax_fl = SVector(ntuple(k->-Inf, Val(N)))  # [-Inf, -Inf, -Inf] for N = 3
    nmin_fl = SVector(ntuple(k->Inf, Val(N)))  # [Inf, Inf, Inf] for N = 3
    rcart = CartesianIndices(ntuple(k->2, Val(N)))  # (2,2,2) for N = 3
    for icart = rcart  # icart: CartesianIndex (see, e.g., https://julialang.org/blog/2016/02/iteration)
        cnr = map((∆b1,∆b2,i)-> i==1 ? ∆b1 : ∆b2, ∆bound[1], ∆bound[2], SVector(icart.I))  # SVector: corner of ∆bound (note icart.I[k] = 1 or 2)
        n_fl = A \ cnr  # SVector: A * n_fl = corner
        nmax_fl = max.(nmax_fl, n_fl)  # SVector
        nmin_fl = min.(nmin_fl, n_fl)  # SVector
    end

    # Find the integral nᵢ's.
    nmax = ceil.(Int, nmax_fl)  # SVector
    nmin = floor.(Int, nmin_fl)  # SVector

    # For nᵢ's within the range calculated above, check if R = n₁a₁ + n₂a₂ + n₃a₃ is within
    # the Shape ∆range.  If it is, then create a shape transformed by R.
    nshape = prod((nmax.-nmin).+1)  # tentative number of shapes
    shp_array = Vector{S}(undef, nshape)  # preallocation (will be truncated later)
    ishape = 0
    nrange = map((n1,n2)->n1:n2, nmin.data, nmax.data)  # e.g., (1:10, 1:10, 1:10) for N = 3
    for n = CartesianIndices(nrange)  # n: CartesianIndex
		∆ = A * SVector(n.I)  # lattice vector R
		if ∆ ∈ ∆range
			shape = translate(shp, ∆)
			ishape += 1
			shp_array[ishape] = shape
		end
    end

    resize!(shp_array, ishape)  # return resized shp_array
end


text:
 prim[1] =   x[1]
 prim[2] =   x[1]*x[2]
 prim[3] =   x[1]^2
 prim[4] =   x[1]*x[2]*x[3]
 prim[5] =   x[1]^3
 prim[6] =   x[1]*x[2]*x[3]*x[4]
 prim[7] =   x[1]^4
 prim[8] =   x[1]*x[2]*x[5]*x[3]*x[4]
 prim[9] =   x[1]^5
 prim[10] =   x[1]^6

max_stars_repo_name: cortner/ManyBodyIPs

max_issues_repo_name: cortner/ManyBodyIPs

max_forks_repo_name: cortner/ManyBodyIPs


text:
@cache struct KuttaPRK2p5Cache{uType,rateType,TabType} <: OrdinaryDiffEqMutableCache
  u::uType
  uprev::uType
  k::rateType
  k1::rateType
  k2::rateType
  k3::rateType
  k4::rateType
  k5_6::Array{rateType,1}
  tmp::uType
  fsalfirst::rateType
  tab::TabType
end

struct KuttaPRK2p5ConstantCache{T,T2} <: OrdinaryDiffEqConstantCache
  α21::T
  α31::T
  α32::T
  α41::T
  α42::T
  α43::T
  α5_6::Array{T,2}
  β1::T
  β3::T
  β5::T
  β6::T
  c2::T2
  c3::T2
  c4::T2
  c5_6::Array{T2,1}

  function KuttaPRK2p5ConstantCache(T, T2)
    α21 = T(1//3)
    α31 = T(4//25)
    α32 = T(6//25)
    α41 = T(1//4)
    α42 = T(-3)
    α43 = T(15//4)

    α5_6 = Array{T}(undef,(2,4))
    α5_6[1,1] = T(6//81)
    α5_6[1,2]= T(90//81)
    α5_6[1,3]= T(-50//81)
    α5_6[1,4]= T(8//81)
    α5_6[2,1]= T(6//75)
    α5_6[2,2]= T(36//75)
    α5_6[2,3]= T(10//75)
    α5_6[2,4]= T(8//75)

    β1 = T(23//192)
    β3 = T(125//192)
    β5 = T(-81//192)
    β6 = T(125//192)
    c2 = T2(1//3)
    c3 = T2(2//5)
    c4 = T2(1)
    c5_6 = Array{T2}(undef,2)
    c5_6[1] = T2(2//3)
    c5_6[2] = T2(4//5)
    new{T,T2}(α21, α31, α32, α41, α42, α43, α5_6, β1, β3, β5, β6, c2, c3, c4,
     c5_6)
  end
end

function alg_cache(alg::KuttaPRK2p5,u,rate_prototype,::Type{uEltypeNoUnits},::Type{uBottomEltypeNoUnits},::Type{tTypeNoUnits},uprev,uprev2,f,t,dt,reltol,p,calck,::Val{true}) where {uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits}
  tmp = zero(u)
  k = zero(rate_prototype)
  k1 = zero(rate_prototype)
  k2 = zero(rate_prototype)
  k3 = zero(rate_prototype)
  k4 = zero(rate_prototype)
  k5_6 = Array{typeof(rate_prototype)}(undef, 2)
  k5_6[1] = zero(rate_prototype)
  k5_6[2] = zero(rate_prototype)
  fsalfirst = zero(rate_prototype)
  tab = KuttaPRK2p5ConstantCache(constvalue(uBottomEltypeNoUnits), constvalue(tTypeNoUnits))
  KuttaPRK2p5Cache(u,uprev,k,k1,k2,k3,k4,k5_6,tmp,fsalfirst,tab)
end

function alg_cache(alg::KuttaPRK2p5,u,rate_prototype,::Type{uEltypeNoUnits},::Type{uBottomEltypeNoUnits},::Type{tTypeNoUnits},uprev,uprev2,f,t,dt,reltol,p,calck,::Val{false}) where {uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits}
  KuttaPRK2p5ConstantCache(constvalue(uBottomEltypeNoUnits), constvalue(tTypeNoUnits))
end


text:
# Quadrilateral in 2D defined by its 4 vertices.
# NOTE: Quadrilaterals are assumed to be convex. This is because quadrilaterals must be convex
# to be valid finite elements, and we're generating meshes using a finite element mesh generator,
# so this seems reasonable. See link below.
# https://math.stackexchange.com/questions/2430691/jacobian-determinant-for-bi-linear-quadrilaterals
struct Quadrilateral_2D <: Face_2D
    # Counter clockwise order
    points::SVector{4, Point_2D}
end

# Constructors
# -------------------------------------------------------------------------------------------------
Quadrilateral_2D(p₁::Point_2D, p₂::Point_2D, p₃::Point_2D, p₄::Point_2D
                ) = Quadrilateral_2D(SVector(p₁, p₂, p₃, p₄))

# Methods
# -------------------------------------------------------------------------------------------------
function (quad::Quadrilateral_2D)(r::Real, s::Real)
    # See The Visualization Toolkit: An Object-Oriented Approach to 3D Graphics, 4th Edition
    # Chapter 8, Advanced Data Representation, in the interpolation functions section
    rₜ = Float64(r); sₜ = Float64(s)
    return (1 - rₜ)*(1 - sₜ)*quad[1] +
                 rₜ*(1 - sₜ)*quad[2] +
                       rₜ*sₜ*quad[3] +
                 (1 - rₜ)*sₜ*quad[4]
end

function area(quad::Quadrilateral_2D)
    # Using the convex quadrilateral assumption, just return the sum of the areas of the two
    # triangles that partition the quadrilateral. If the convex assumption ever changes, you
    # need to verify that the triangle pair partitions the quadrilateral. Choosing the wrong
    # pair overestimates the area, so just get the areas of both pairs of valid triangles and use
    return sum(area.(triangulate(quad)))
end

function centroid(quad::Quadrilateral_2D)
    tris = triangulate(quad)
    A₁ = area(tris[1])
    A₂ = area(tris[2])
    C₁ = centroid(tris[1])
    C₂ = centroid(tris[2])
    return (A₁*C₁ + A₂*C₂)/(A₁ + A₂)
end

function in(p::Point_2D, quad::Quadrilateral_2D)
    # If the point is to the left of every edge
    #  4<-----3
    #  |      ^
    #  | p    |
    #  |      |
    #  |      |
    #  v----->2
    #  1
    return isleft(p, LineSegment_2D(quad[1], quad[2])) &&
           isleft(p, LineSegment_2D(quad[2], quad[3])) &&
           isleft(p, LineSegment_2D(quad[3], quad[4])) &&
           isleft(p, LineSegment_2D(quad[4], quad[1]))
end

function intersect(l::LineSegment_2D, quad::Quadrilateral_2D)
    # Create the 4 line segments that make up the quadrilateral and intersect each one
    edges = SVector(LineSegment_2D(quad[1], quad[2]),
                    LineSegment_2D(quad[2], quad[3]),
                    LineSegment_2D(quad[3], quad[4]),
                    LineSegment_2D(quad[4], quad[1]))
    ipoints = MVector(Point_2D(), Point_2D(), Point_2D())
    n_ipoints = 0x00000000
    # We need to account for 4 points returned due to vertex intersection
    # The only way we get 4 points though, is if 2 are redundant.
    # So, we return 3, which guarantees all unique points will be returned and the output
    # matched with the triangle output
    for k ∈  1:4
        npoints, point = l ∩ edges[k]
        if npoints === 0x00000001 && n_ipoints !== 0x00000003
            n_ipoints += 0x00000001
            ipoints[n_ipoints] = point
        end
    end
    return n_ipoints, SVector(ipoints)
end

function triangulate(quad::Quadrilateral_2D)
    # Return the two triangles that partition the domain
    A, B, C, D = quad.points
    return SVector(Triangle_2D(A, B, C), Triangle_2D(C, D, A))
end

# Plot
# -------------------------------------------------------------------------------------------------
if enable_visualization
    function convert_arguments(LS::Type{<:LineSegments}, quad::Quadrilateral_2D)
        l₁ = LineSegment_2D(quad[1], quad[2])
        l₂ = LineSegment_2D(quad[2], quad[3])
        l₃ = LineSegment_2D(quad[3], quad[4])
        l₄ = LineSegment_2D(quad[4], quad[1])
        lines = [l₁, l₂, l₃, l₄]
        return convert_arguments(LS, lines)
    end

    function convert_arguments(LS::Type{<:LineSegments}, Q::Vector{Quadrilateral_2D})
        point_sets = [convert_arguments(LS, quad) for quad ∈  Q]
        return convert_arguments(LS, reduce(vcat, [pset[1] for pset ∈  point_sets]))
    end

    function convert_arguments(M::Type{<:Mesh}, quad::Quadrilateral_2D)
        points = [quad[i] for i = 1:4]
        faces = [1 2 3;
                 3 4 1]
        return convert_arguments(M, points, faces)
    end

    function convert_arguments(M::Type{<:Mesh}, Q::Vector{Quadrilateral_2D})
        points = reduce(vcat, [[quad[i] for i = 1:4] for quad ∈  Q])
        faces = zeros(Int64, 2*length(Q), 3)
        j = 0
        for i in 1:2:2*length(Q)
            faces[i    , :] = [1 2 3] .+ j
            faces[i + 1, :] = [3 4 1] .+ j
            j += 4
        end
        return convert_arguments(M, points, faces)
    end
end

max_stars_repo_name: KyleVaughn/MOCNeutronTransport

max_issues_repo_name: KyleVaughn/MOCNeutronTransport

max_forks_repo_name: KyleVaughn/MOCNeutronTransport


text:
# recursive.jl -- Part of 'MultinomialSeries',
# a package for The Julia Programming Language
# Copyright (C) 2020-2021 Markus J. Weber
# License: See accompanying file 'LICENSE.md'

#=---- E X P E R I M E N T A L ----=#

"""
    map_multinomial(fn, m, n)

dims : m = dimension, > 0
order: n = order of application, ≥ 0
k = 1...m
Returns nothing.

Usage:
    map_multinomial(3,3) do k, Cₖₙ
        # This function is called for every combination of k
        # and the associated multinomial coefficient Cₖₙ.
        @show k, Cₖₙ
    end
"""
function map_multinomial(
            fn,
            dims::Integer, order::Integer,
            k::Union{Nothing,Array}=nothing
        )

    # Might quickly overflow !
    n_choose_k(n,k) = factorial(n) ÷ (factorial(k) * factorial(n-k))
        # Check validity of iterations
        #@info "number terms" n_choose_k(order+dims-1, dims-1) dims^order

    if k == nothing
        @boundscheck dims  > 0 || throw(ArgumentError("Multinomial requires positive dimension (>0). Given: $dims"))
        @boundscheck order ≥ 0 || throw(ArgumentError("Multinomial requires non-negative order (≥0). Given: $order"))
        # initial setup
        k = zeros(Int, dims)
    end

    if dims == 1
        k[1] = order
        Cₖₙ = factorial(sum(k)) ÷ prod(factorial.(k))       # Might quickly overflow !
        fn(k, Cₖₙ)
    else
        for kₘ in 0:order
            k[dims] = kₘ
            multinomial_generator(fn, dims-1, order-kₘ, k)
        end
    end

    nothing
end

# end of file


text:
Pkg.add("Ipopt")
#Pkg.add("KNITRO")
Pkg.add("JuMP")

Pkg.clone("https://github.com/JuliaLang/MATLAB.jl.git")

# install http://brew.sh
# brew tap homebrew/science
# brew install mumps [--with-scotch5] [--with-openblas]  # use scotch/openblas at your option
Pkg.clone("https://github.com/JuliaOptimizers/MUMPS.jl.git")
Pkg.build("MUMPS")

using MPI, MUMPS
let
  MPI.Init();
  #mumps = Mumps{Float64}(mumps_definite);

  A = sprand(10, 10, .2) + speye(10); rhs = rand(10);
  x = solve(A, rhs);  # Mumps object is created and destroyed
  norm(x - A \ rhs) / norm(x)
  MPI.Finalize();
end

using MPI
x=1+1

import MUMPS # needs to be intialized before Ipopt don't understand why
MPI.Init();
cntl = MUMPS.default_icntl;
cntl[4] = 1;
var_mumps = MUMPS.Mumps{Float64}(MUMPS.mumps_unsymmetric, cntl, MUMPS.default_cntl64	);  # Real, general unsymmetric
A = sparse(rand(4,4)); rhs = rand(4);       # Happens on all cores
MUMPS.associate_matrix(var_mumps, A);
MUMPS.factorize(var_mumps);
MUMPS.associate_rhs(var_mumps, rhs);
MUMPS.solve(var_mumps);
x_sol = MUMPS.get_solution(var_mumps);
println(var_mumps.infog[12])
MUMPS.finalize(var_mumps);
MPI.Finalize();

Pkg.update()

Pkg.clone("https://github.com/ohinder/Everest.git")
Pkg.clone("https://github.com/ohinder/my_advanced_timer.git")
Pkg.update()

using Everest

max_forks_repo_name: ohinder/Everest


text:
# This file is a part of ChangesOfVariables.jl, licensed under the MIT License (MIT).

"""
    ChangesOfVariables.test_with_logabsdet_jacobian(f, x, getjacobian; compare = isapprox, kwargs...)

Test if [`with_logabsdet_jacobian(f, x)`](@ref) is implemented correctly.

Checks if the result of `with_logabsdet_jacobian(f, x)` is approximately
equal to `(f(x), logabsdet(getjacobian(f, x)))`

So the test uses `getjacobian(f, x)` to calculate a reference Jacobian for
`f` at `x`. Passing `ForwardDiff.jabobian`, `Zygote.jacobian` or similar as
the `getjacobian` function will do fine in most cases. If input and output
of `f` are real scalar values, use `ForwardDiff.derivative`.

Note that the result of `getjacobian(f, x)` must be a real-valued matrix
or a real scalar, so you may need to use a custom `getjacobian` function
that transforms the shape of `x` and `f(x)` internally, in conjunction
with automatic differentiation.

`kwargs...` are forwarded to `compare`.
"""
function test_with_logabsdet_jacobian(f, x, getjacobian; compare = isapprox, kwargs...)
    @testset "test_with_logabsdet_jacobian: $f with input $x" begin
        ref_y, test_type_inference = try
            @inferred(f(x)), true
        catch err
            f(x), false
        end

        y, ladj = if test_type_inference
            @inferred with_logabsdet_jacobian(f, x)
        else
            with_logabsdet_jacobian(f, x)
        end

        ref_ladj = _generalized_logabsdet(getjacobian(f, x))[1]

        @test compare(y, ref_y; kwargs...)
        @test compare(ladj, ref_ladj; kwargs...)
    end
    return nothing
end

_generalized_logabsdet(A) = logabsdet(A)
_generalized_logabsdet(x::Real) = log(abs(x))


text:
using Test
using LinearAlgebra
using Ferrite
using Ferrite: cellid, getcoordinates, CellIterator

using TopOpt
using TopOpt.TopOptProblems: boundingbox, nnodespercell, getgeomorder, getmetadata, getdh, getE, getdim
using TrussTopOpt.TrussTopOptProblems
using TrussTopOpt.TrussTopOptProblems: buckling
using TrussTopOpt.TrussTopOptProblems: getA, default_quad_order
using IterativeSolvers
using Arpack
using Crayons.Box

problem_json = ["buckling_2d_nodal_instab.json", "buckling_2d_global_instab.json", "buckling_2d_debug.json"]
ins_dir = joinpath(@__DIR__, "instances", "fea_examples");

# @testset "Buckling problem solve - $(problem_json[i])" for i in 1:length(problem_json)
    i = 3
    file_name = problem_json[i]
    problem_file = joinpath(ins_dir, file_name)

    node_points, elements, mats, crosssecs, fixities, load_cases = parse_truss_json(problem_file);
    ndim, nnodes, ncells = length(node_points[1]), length(node_points), length(elements)
    loads = load_cases["0"]

    problem = TrussProblem(Val{:Linear}, node_points, elements, loads, fixities, mats, crosssecs);

    # @test getdim(problem) == ndim
    # @test Ferrite.getncells(problem) == ncells
    # @test getE(problem) == [m.E for m in mats]
    # @test problem.black == problem.white == falses(ncells)
    # @test problem.force == loads
    # @test problem.varind == 1:ncells
    # grid = problem.ch.dh.grid
    # @test length(grid.cells) == ncells
    # @test getgeomorder(problem) == 1
    # @test nnodespercell(problem) == 2

    solver = FEASolver(Displacement, Direct, problem)
    # call solver to trigger assemble!
    solver()
    @show solver.u
    scene, layout = draw_truss_problem(problem; u=solver.u,
        default_load_scale=0.2, default_support_scale=0.2, default_arrow_size=0.03)

    try
        global K, Kσ = buckling(problem, solver.globalinfo, solver.elementinfo)
    catch err
        println(RED_BG("ERROR: "))
        println(err)
        if isa(err, SingularException)
            println(RED_FG("Linear elasticity solve failed due to mechanism."))
            K_l = solver.globalinfo.K
            # TODO: use sparse eigen
            # λ, ϕ = eigs(K.data, nev = 2, which=:SM);
            # @show λ
            # @show ϕ

            # ? is this ordered by eigen values' magnitude?
            local F = eigen(Array(K_l))
            @show F
            @assert abs(F.values[1] - 0.0) < eps()

            # * draw eigen mode
            scene, layout = draw_truss_problem(problem; u=F.vectors[:,1])
        end
        throw(err)
    end
    # the linear elastic model must be solved successfully to proceed on buckling calculation
    println(GREEN_FG("Linear elasticity solved."))

    ch = problem.ch
    C = zeros(size(K, 1), length(ch.prescribed_dofs))
    setindex!.(Ref(C), 1, ch.prescribed_dofs, 1:length(ch.prescribed_dofs))
    # C = hcat(C, solver.u)

    # Find the maximum eigenvalue of system Kσ x = 1/λ K x
    # The reason we do it this way is because K is guaranteed to be positive definite while Kσ is not and the LOBPCG algorithm to find generalised eigenvalues requires the matrix on the RHS to be positive definite.
    # https://julialinearalgebra.github.io/IterativeSolvers.jl/stable/eigenproblems/lobpcg/
    c = norm(K)
    # r = lobpcg(Kσ.data, K.data, true, 1; C=C)
    r = lobpcg(Kσ.data, K.data, true, 1)
    @show r

    # # Minimum eigenvalue of the system K x + λ Kσ x = 0
    @show λ = -1/r.λ[1]
    @show v = r.X

    scene, layout = draw_truss_problem(problem; u=v, default_exagg=1.0, exagg_range=10.0,
        default_load_scale=0.2, default_support_scale=0.2, default_arrow_size=0.03)

    # F = eigen(Array(Kσ), Array(K))
    # F = eigen(Array(K), Array(Kσ))
    # @show - 1 ./ F.values

    if 0 < λ && λ < 1
        # 0 < λ_min < 1 means that the equilibrium equation with load
        # λ*f is not solvable for some 0 < λ < 1 and the structure under load
        # f is not stable

        # s = 1.0
        Ks = K .+ λ .* Kσ

        eigen_vals, eigen_vecs = eigs(Ks, nev = 2, which=:SM);
        @show eigen_vals
        @show eigen_vecs
        local F = eigen(Array(Ks))
        @show F

        # cholKs = cholesky(Symmetric(Ks))
        # u_b = cholKs \ solver.globalinfo.f
        # println(GREEN_FG("Buckling problem solved."))

        ## TODO plot analysis result with
        # scene, layout = draw_truss_problem(problem; u=u_b)
        # println("End")
    end

    # TODO eigenmode to show global instability mode
    # Find a test case for verifying the geometric stiffness matrix
    # ? what does MMA need? Inner opt uses MMA or SIMP?
    #  - where to plug in the gradient of the barrier function?
        # function (v::Volume{T})(x, grad = nothing) where {T}
    #  - how to compute ∂K/∂x_e and ∂f/∂x_e?
            # d = ForwardDiff.Dual{T}(x[varind[i]], one(T))
            # if PENALTY_BEFORE_INTERPOLATION
            #     p = density(penalty(d), xmin)
            # else
            #     p = penalty(density(d, xmin))
            # end
            # grad[varind[i]] = -p.partials[1] * cell_comp[i] # Ke
    #  adjoint method: ∂u/∂x_e = K^{-1} * (∂f/∂x_e - ∂K/∂x_e * u)
    # ∂Kσ/∂x_e for all e

    # TODO finite difference to verify the gradient
    # TODO verify the gradient for some analytical problems
    # TODO "manual" interior point loop, adjusting the c value every iter

# end # end test set

# =#


text:
### A Pluto.jl notebook ###
# v0.12.20

using Markdown
using InteractiveUtils

# ╔═╡ f3e35f66-6993-11eb-2b3f-7915f3382768
begin
	using Plots
	using Random
	using TimeSeries
	using Dates
	using Statistics: mean
end

# ╔═╡ 7df5623a-698a-11eb-0291-9b9899764848
md"""
# Introduction to Time Series I

> Author: Edwin Bedolla
>
> Date: Original, 4th April 2020. This version, 7th February 2021.

This document gathers some main concepts in time series analysis as well as some examples written
for the [Julia programming language](https://julialang.org).

First, we import every module that will be used in this document.
"""

# ╔═╡ 059e0846-6994-11eb-0a0a-b1fd93c64ba3
gr();

# ╔═╡ 1d604200-6994-11eb-1b2e-1bb27449c881
md"""
## Definitions

> - A **time series** is a collection of **random variables** indexed in a *ordered set* representing time periods.

> - **Stochastic processes** are a family of _indexed_ random variables $Z(\omega, t)$
>   where $\omega$ belongs to a **sample space** and $t$ belongs to an **index set**.

> - For a given $\omega$, $Z(\omega, t)$ as a function of $t$ is called a sample function or **realisation**.

> - The population that consists of all possible realisations is called the **ensemble**.

From these definitions we can see that a *time series* is actually a *stochastic process*.
We will now turn to see some simple examples of time series and how to program them in `Julia`.

"""

# ╔═╡ 2a51d53c-6994-11eb-3834-0d546ce983c9
md"""
## Examples

### 1. **White noise**

One of the basic time series is **white noise**, which is a time series
generated from uncorrelated variables, which are most of the time *normally distributed*.

"""

# ╔═╡ 5681a222-6994-11eb-1ce0-370b0c741e80
# Create a range of time for a year, spaced evenly every 10 hours
dates = DateTime(2018, 1, 1, 1):Dates.Hour(10):DateTime(2018, 12, 31, 24);

# ╔═╡ 5fd0b4bc-6994-11eb-2423-7b82cb249212
# Build a TimeSeries object with the specified time range and white noise
ts = TimeArray(dates, randn(length(dates)));

# ╔═╡ 722194d8-6994-11eb-0345-25a122d42c0d
plot(ts, leg=false)

# ╔═╡ 85ccb094-6994-11eb-1690-c38c7c43b7c4
md"""
This collection of random variables $\{x_t\}$ has the following properties:

```math
\mu_x = 0
```

```math
\sigma^2_x = 1
```

They are *independently and identically distributed* such that

```math
x_t \sim i.i.d. \mathcal{N}(0, 1)
```

We can see that this time series is very *noisy*, with big peaks all over the place;
if we wanted to do some meaningful analysis on it we might have a difficult time.
Instead we can apply a very simple **smoothing** technique named the **moving average**.

"""

# ╔═╡ b23b3fce-6994-11eb-28ea-e354d2ce8fbf
md"""
### 2. **Moving average**

The **moving average** is an actual time series in itself defined as

```math
v_t = \frac{1}{N} \sum_{i=0}^{N-1} w_{t-i}
```

which means that the *moving average* takes as input the neighboring values
in the past
and future time periods and evaluates them, obtaining the **realization** as an arithmetic
mean. The following example applies a moving average to a white noise time series.

"""

# ╔═╡ ebf6d020-6994-11eb-0c7e-455eb2805265
# Create a range of time for a year, spaced evenly every 50 hours
dates1 = DateTime(2018, 1, 1, 1):Dates.Hour(50):DateTime(2018, 12, 31, 24);

# ╔═╡ f6858432-6994-11eb-32ef-979ff0ffbada
# Build a TimeSeries object with the specified time range and white noise
ts1 = TimeArray(dates1, randn(length(dates1)));

# ╔═╡ 25706cee-6995-11eb-361f-efd74443d4df
# Moving average using 3 values
moving_average = moving(mean, ts1, 3);

# ╔═╡ 2c84f072-6995-11eb-35b2-f76c27eeb3ce
# And plot both of them superimposed
begin
	plot(ts, label = "White noise")
	plot!(moving_average, label = "Moving average")
end

# ╔═╡ 3a0e0968-6995-11eb-175f-612026c9c3d8
md"""
In this case we used the following 3-valued moving average

```math
v_t = \frac{1}{3} \left( w_{t-1} + w_{t} + w_{t+1} \right)
```

We can see that most of the peaks are gone, but we can still observe the main structure
of the original time series.

"""

# ╔═╡ 61dbad7e-6995-11eb-1a7f-8f683553d132
md"""
### 3. **Random walks with drift**

The analysis of *trends* is one of the prime examples of time series, where there is
a need to explore and understand what has been happening throughout several time periods.
For example, one might need to understand how the global temperature has been rising
since 1997 until today.

To model this type of time series we have the following

```math
x_t = \delta t + \sum_{j=1}^{t} w_j
```

where, as before, ``w_j`` is a white noise time series, with ``\sigma_w = 1``;
``\delta`` is the so-called *drift* coefficient that makes the trend in the time
series much more steep. In the special case that we have ``\delta = 0`` we would
have what's called a simple *random walk*.

We can see an example of a random walk *with drift* and *without drift* below.

"""

# ╔═╡ 8265a574-6995-11eb-333c-73dc1f662455
# Create a range of time for a year, spaced evenly every 50 hours
dates2 = DateTime(2018, 1, 1, 1):Dates.Hour(50):DateTime(2018, 12, 31, 24);

# ╔═╡ 8dcfcf46-6995-11eb-15f0-ffb4087992f5
# Build a TimeSeries object with the specified time range and white noise
wt = randn(length(dates2));

# ╔═╡ 93ae4a78-6995-11eb-17f5-a758fa6d7160
# Define a drift value
δ = 0.2;

# ╔═╡ 973d5198-6995-11eb-136f-1d2bf3820068
# We then build the random walk with drift
ts_drift = TimeArray(dates2, cumsum(wt .+ δ));

# ╔═╡ 9d400022-6995-11eb-2064-fd73bfa9a5b2
# And without drift
ts_nodrift = TimeArray(dates2, cumsum(wt));

# ╔═╡ a502c678-6995-11eb-1f1e-752a083ca379
# And plot both of them superimposed
begin
	plot(ts_drift, label = "Random walk with drift")
	plot!(ts_nodrift, label = "Random walk without drift")
end

# ╔═╡ b5292362-6995-11eb-06b8-5f2deab48e73
md"""
Here we see a *trend* in the time series because every time period in the future
the *realization* is always higher than in past values.We might expect that this
time series is always growing. We will see that *trends* are very useful to perform
some basic exploratory analysis on the data itself in further documents.

"""

# ╔═╡ Cell order:
# ╟─7df5623a-698a-11eb-0291-9b9899764848
# ╠═f3e35f66-6993-11eb-2b3f-7915f3382768
# ╠═059e0846-6994-11eb-0a0a-b1fd93c64ba3
# ╟─1d604200-6994-11eb-1b2e-1bb27449c881
# ╟─2a51d53c-6994-11eb-3834-0d546ce983c9
# ╠═5681a222-6994-11eb-1ce0-370b0c741e80
# ╠═5fd0b4bc-6994-11eb-2423-7b82cb249212
# ╠═722194d8-6994-11eb-0345-25a122d42c0d
# ╟─85ccb094-6994-11eb-1690-c38c7c43b7c4
# ╟─b23b3fce-6994-11eb-28ea-e354d2ce8fbf
# ╠═ebf6d020-6994-11eb-0c7e-455eb2805265
# ╠═f6858432-6994-11eb-32ef-979ff0ffbada
# ╠═25706cee-6995-11eb-361f-efd74443d4df
# ╠═2c84f072-6995-11eb-35b2-f76c27eeb3ce
# ╟─3a0e0968-6995-11eb-175f-612026c9c3d8
# ╟─61dbad7e-6995-11eb-1a7f-8f683553d132
# ╠═8265a574-6995-11eb-333c-73dc1f662455
# ╠═8dcfcf46-6995-11eb-15f0-ffb4087992f5
# ╠═93ae4a78-6995-11eb-17f5-a758fa6d7160
# ╠═973d5198-6995-11eb-136f-1d2bf3820068
# ╠═9d400022-6995-11eb-2064-fd73bfa9a5b2
# ╠═a502c678-6995-11eb-1f1e-752a083ca379
# ╟─b5292362-6995-11eb-06b8-5f2deab48e73

max_issues_repo_name: BioLogs/TimeSeriesConcepts


text:
# old quad and adjoint replaced by inplace accumulation versions.
# function quad(A::AbstractArray{<:Number,3}, x)
#     return map(j -> dot(x, view(A, j, :, :), x), 1:size(A, 1))
# end
# # quadratic form pullback
# function quad_pb(Δres::AbstractVector, A::AbstractArray{<:Number,3}, x::AbstractVector)
#     ΔA = similar(A)
#     Δx = zeros(length(x))
#     tmp = x * x'
#     for i in 1:size(A, 1)
#         ΔA[i, :, :] .= tmp .* Δres[i]
#         Δx += (A[i, :, :] + A[i, :, :]') * x .* Δres[i]
#     end
#     return ΔA, Δx
# end

# y += quad(A, x)
# The quad_muladd! uses on a vector of matrices for A
function quad_muladd!(y, A, x)
    @inbounds for j in 1:size(A, 1)
        @views y[j] += dot(x, A[j], x)
    end
    return y
end

# inplace version with accumulation and using the cache of A[i] + A[i]', etc.
function quad_muladd_pb!(ΔA_vec, Δx, Δres, A_vec_sum, x)
    tmp = x * x'  # could add in a temp here
    @inbounds for (i, A_sum) in enumerate(A_vec_sum)  # @views @inbounds  ADD
        ΔA_vec[i] .+= tmp .* Δres[i]
        Δx .+= A_sum * x .* Δres[i]
    end
    return nothing
end

# Temporary.  Eventually, move to sciml NoiseProcess with better rng support/etc.
get_concrete_noise(prob, noise, B, T) = noise # maybe do a promotion to an AbstractVectorOfVector type
get_concrete_noise(prob, noise::Nothing, B, T) = randn(eltype(B), size(B, 2), T) # default is unit Gaussian
get_concrete_noise(prob, noise::UnivariateDistribution, B, T) = rand(noise, size(B, 2), T) # iid

# Utility functions to conditionally check size if not-nothing
maybe_check_size(m::AbstractMatrix, index, val) = (size(m, index) == val)
maybe_check_size(m::Nothing, index, val) = true


text:
module BatteriesBatterySwitch

using ..Ahorn, Maple

@mapdef Entity "batteries/battery_switch" BatterySwitch(x::Integer, y::Integer, rightSide::Bool=false, persistent::Bool=false, horizontal::Bool=true, alwaysFlag::Bool=false)

const placements = Ahorn.PlacementDict()

directions = Dict{String, Tuple{Bool, String, Bool}}(
    "up" => (false, "ceiling", false),
    "down" => (false, "ceiling", true),
    "left" => (true, "rightSide", false),
    "right" => (true, "rightSide", true),
)

for (dir, data) in directions
    key = "Battery Switch ($(uppercasefirst(dir)))"
    horiz, datakey, val = data
    placements[key] = Ahorn.EntityPlacement(
        BatterySwitch,
        "rectangle",
        Dict{String, Any}(
            "horizontal" => horiz,
            datakey => val
        )
    )
end

function Ahorn.selection(entity::BatterySwitch)
    x, y = Ahorn.position(entity)
    if get(entity.data, "horizontal", true)
        left = get(entity.data, "rightSide", false)

        if left
            return Ahorn.Rectangle(x, y - 1, 10, 16)

        else
            return Ahorn.Rectangle(x - 2, y, 10, 16)
        end
    else
        ceiling = get(entity.data, "ceiling", false)

        if ceiling
            return Ahorn.Rectangle(x, y, 16, 12)

        else
            return Ahorn.Rectangle(x, y - 4, 16, 12)
        end
    end
end

function Ahorn.render(ctx::Ahorn.Cairo.CairoContext, entity::BatterySwitch, room::Maple.Room)
    texture = "batteries/battery_switch/insert8"
    if get(entity.data, "horizontal", true)
        right = get(entity.data, "rightSide", false)

        if right
            Ahorn.drawSprite(ctx, texture, 30, 28, rot=pi)

        else
            Ahorn.drawSprite(ctx, texture, -2, 8)
        end
    else
        ceiling = get(entity.data, "ceiling", false)

        if ceiling
            Ahorn.drawSprite(ctx, texture, 8, 30, rot=-pi / 2)

        else
            Ahorn.drawSprite(ctx, texture, 28, -2, rot=pi / 2)
        end
    end
end

end

max_stars_repo_name: PenguinOwl/Batteries

max_issues_repo_name: PenguinOwl/Batteries

max_forks_repo_name: PenguinOwl/Batteries


text:
module SQLdf

using SQLite, DataFrames, RCall

import DataFrames: DataFrame

export sqldf, @sqldf, DataFrame, setRDB, @rput

include("sqldf_SQLite.jl")
include("sqldf_R.jl")
include("qtables.jl")

sqldf = sqldf_SQLite

function setRDB(rdb::String)

    @assert rdb in ["SQLite","R"] """Posible values for RDB are "SQLite" and "R" """

    if rdb == "R"
        if !rcopy(R""" "sqldf" %in% installed.packages() """)
            @error "'sqldf' needs to be installed in R. \nMore information at https://juliainterop.github.io/RCall.jl/latest/installation.html"
        else
            R"library(proto); library(gsubfn); library(RSQLite); library(sqldf)"
        end

        global sqldf = sqldf_R
    end

    if rdb == "SQLite"
         global sqldf = sqldf_SQLite
    end

end

macro sqldf(query)
    return :( sqldf($(esc(query))) )
end

"""
Execute SQL queries on julia DataFrame and Tables
"""
SQLdf

"""
Package: SQLdf

    sqldf(query::String)::DataFrame

Execute a SQL query on julia DataFrames

By default sqldf uses SQLite however it can also connect to R/sqldf via setRDB("R").

When using R/sqldf only DataFrame types are accepted and columns in the DataFrame must
have a type other than Any. In order to  work with dates expressions like
\"""select strftime("%Y", datetime_column, "unixepoch") as year from T\""" may be used.

When using SQLite all types implementing Tables interface will be accepted, however sqldf
will still return and DataFrame.

# Arguments
`query`: SQL query

# Returns
Julia DataFrame with query results

# Examples
```julia-repl
julia> T = DataFrame(C1 = [1,2], C2 = ["a","b"])

julia> query = \"""
               select *
               from T
               where C2 = "a"
               \""";

julia> sqldf(query)
1×2 DataFrame
 Row │ C1     C2
     │ Int64  String
─────┼───────────────
   1 │     1  a

julia> @sqldf query
1×2 DataFrame
 Row │ C1     C2
     │ Int64  String
─────┼───────────────
   1 │     1  a

```
"""
sqldf

end


text:
using Zygote
using Zygote: @adjoint, @nograd

@adjoint function free_energy(β::Real, H::AbstractBlock{N}, sampler, circuit::AbstractBlock{N}, samples) where N
    free_energy(β, H, sampler, circuit, samples), function (adjy)
        adjθ = grad_θ(H, circuit, samples) * adjy
        adjsampler = grad_sampler(β, H, sampler, circuit, samples) .*adjy
        return (nothing, nothing, adjsampler, adjθ, nothing)
    end
end

@nograd update!, dispatch!, get_input_reg, free_energy_local