text:
module Diagnostics

using NCDatasets
using Statistics
using LinearAlgebra
using EnsembleKalmanProcesses
using EnsembleKalmanProcesses.ParameterDistributions
import EnsembleKalmanProcesses: construct_sigma_ensemble, construct_mean, construct_cov
import EnsembleKalmanProcesses: construct_successful_mean, construct_successful_cov

using ..ReferenceModels
using ..ReferenceStats
using ..HelperFuncs
const NC = NCDatasets

export io_dictionary_ensemble, io_dictionary_reference, io_dictionary_metrics
export io_dictionary_particle_state, io_dictionary_particle_eval
export io_dictionary_val_metrics, io_dictionary_val_particle_eval
export io_dictionary_val_reference, io_dictionary_prior

"""
    io_dictionary_reference()
    io_dictionary_reference(
        ref_stats::ReferenceStatistics,
        ref_models::Vector{ReferenceModel},
        write_full_stats::Bool = true,
    )

Dictionary of diagnostics pertaining to the `ReferenceModel`s and `ReferenceStatistics` that define the inverse problem.

Elements:

 - `Gamma` :: Covariance matrix in the inverse problem latent space (regularized low-dimensional encoding).
 - `Gamma_full` :: Covariance matrix of normalized observed variables in full space (possibly ill-conditioned). Only written to file if `write_full_stats` is true.
 - `Gamma_full_diag` :: Diagonal of `Gamma_full`, useful when `Gamma_full` is not written to file.
 - `y` :: Observations in the inverse problem latent space (low-dimensional encoding).
 - `y_full` :: Normalized observations in full space.
 - `P_pca` :: PCA projection matrix from full space to low-dimensional latent space.
 - `num_vars` :: Maximum number of observed fields (not dimensions) per `ReferenceModel`.
 - `var_dof` :: Maximum number of degrees of freedom of each field per `ReferenceModel`.
 - `config_pca_dim` :: Dimensionality of the latent space associated with each `ReferenceModel`.
 - `config_name` :: Name of each `ReferenceModel` used to construct the inverse problem.
 - `config_z_obs` :: Vertical locations of the observations of each `ReferenceModel`.
 - `norm_factor` :: Pooled variance used to normalize each field of each `ReferenceModel`.
"""
function io_dictionary_reference()
    io_dict = Dict(
        "Gamma" => (; dims = ("out", "out"), group = "reference", type = Float64),
        "Gamma_full" => (; dims = ("out_full", "out_full"), group = "reference", type = Float64),
        "Gamma_full_diag" => (; dims = ("out_full",), group = "reference", type = Float64),
        "y" => (; dims = ("out",), group = "reference", type = Float64),
        "y_full" => (; dims = ("out_full",), group = "reference", type = Float64),
        "P_pca" => (; dims = ("out_full", "out"), group = "reference", type = Float64),
        "num_vars" => (; dims = ("config",), group = "reference", type = Int16),
        "var_dof" => (; dims = ("config",), group = "reference", type = Int16),
        "config_pca_dim" => (; dims = ("config",), group = "reference", type = Int16),
        "config_name" => (; dims = ("config",), group = "reference", type = String),
        "config_z_obs" => (; dims = ("config", "dof"), group = "reference", type = Float64),
        "norm_factor" => (; dims = ("config", "config_field"), group = "reference", type = Float64),
    )
    return io_dict
end
function io_dictionary_reference(
    ref_stats::ReferenceStatistics,
    ref_models::Vector{ReferenceModel},
    write_full_stats::Bool = true,
)
    orig_dict = io_dictionary_reference()
    d_full = full_length(ref_stats)
    d = pca_length(ref_stats)
    num_vars = [length(norm_scale) for norm_scale in ref_stats.norm_vec]
    var_dof = Int.([size(P_pca, 1) for P_pca in ref_stats.pca_vec] ./ num_vars)
    config_pca_dim = [size(P_pca, 2) for P_pca in ref_stats.pca_vec]
    config_name = [
        rm.case_name == "LES_driven_SCM" ? join(split(basename(rm.y_dir), ".")[2:end], "_") : rm.case_name for
        rm in ref_models
    ]

    config_z_obs = zeros(length(ref_models), maximum(var_dof))
    for (i, rm) in enumerate(ref_models)
        z_obs = get_z_obs(rm)
        config_z_obs[i, 1:length(z_obs)] = z_obs
    end

    P_pca_full = zeros(d_full, d)
    idx_row = 1
    idx_col = 1
    for P_pca in ref_stats.pca_vec
        rows, cols = size(P_pca)
        P_pca_full[idx_row:(idx_row + rows - 1), idx_col:(idx_col + cols - 1)] = P_pca
        idx_row += rows
        idx_col += cols
    end
    io_dict = Dict(
        "Gamma" => Base.setindex(orig_dict["Gamma"], ref_stats.Γ, :field),
        "Gamma_full_diag" => Base.setindex(orig_dict["Gamma_full_diag"], Array(diag(ref_stats.Γ_full)), :field),
        "y" => Base.setindex(orig_dict["y"], ref_stats.y, :field),
        "y_full" => Base.setindex(orig_dict["y_full"], ref_stats.y_full, :field),
        "P_pca" => Base.setindex(orig_dict["P_pca"], P_pca_full, :field),
        "num_vars" => Base.setindex(orig_dict["num_vars"], num_vars, :field),
        "var_dof" => Base.setindex(orig_dict["var_dof"], var_dof, :field),
        "config_pca_dim" => Base.setindex(orig_dict["config_pca_dim"], config_pca_dim, :field),
        "config_name" => Base.setindex(orig_dict["config_name"], config_name, :field),
        "config_z_obs" => Base.setindex(orig_dict["config_z_obs"], config_z_obs, :field),
    )
    max_num_fields = maximum([length(norm_vec) for norm_vec in ref_stats.norm_vec])
    norm_factor = zeros(length(ref_stats.norm_vec), max_num_fields)
    for (i, norm_vec) in enumerate(ref_stats.norm_vec)
        num_fields = length(norm_vec)
        norm_factor[i, 1:num_fields] = norm_vec
    end
    io_dict["norm_factor"] = Base.setindex(orig_dict["norm_factor"], norm_factor, :field)

    if write_full_stats
        io_dict["Gamma_full"] = Base.setindex(orig_dict["Gamma_full"], Array(ref_stats.Γ_full), :field)
    end
    return io_dict
end

"""
    io_dictionary_val_reference()
    io_dictionary_val_reference(
        ref_stats::ReferenceStatistics,
        ref_models::Vector{ReferenceModel},
        write_full_stats::Bool = true,
    )

Dictionary of diagnostics pertaining to the `ReferenceModel`s and `ReferenceStatistics` in the validation set.

Elements:

 - `Gamma_val` :: Covariance matrix in latent space, using the same truncation as for the training set.
 - `Gamma_full_val` :: Covariance matrix of normalized observed variables in full space. Only written to file if `write_full_stats` is true.
 - `Gamma_full_diag_val` :: Diagonal of `Gamma_full_val`, useful when `Gamma_full_val` is not written to file.
 - `y_val` :: Observations in latent space, for observed fields in the validation set.
 - `y_full_val` :: Normalized observations in full space, for the validation set.
 - `P_pca_val` :: PCA projection matrix from full space to low-dimensional latent space, for the validation set.
 - `num_vars_val` :: Maximum number of observed fields (not dimensions) per validation `ReferenceModel`.
 - `var_dof_val` :: Maximum number of degrees of freedom of each field per validation `ReferenceModel`.
 - `config_pca_dim_val` :: Dimensionality of the latent space associated with each validation `ReferenceModel`.
 - `config_name_val` :: Name of each `ReferenceModel` in the validation set.
 - `config_z_obs_val` :: Vertical locations of the observations of each validation `ReferenceModel`.
 - `norm_factor_val` :: Pooled variance used to normalize each field of each validation `ReferenceModel`.
"""
function io_dictionary_val_reference()
    io_dict = Dict(
        "Gamma_val" => (; dims = ("out_val", "out_val"), group = "reference", type = Float64),
        "Gamma_full_val" => (; dims = ("out_full_val", "out_full_val"), group = "reference", type = Float64),
        "Gamma_full_diag_val" => (; dims = ("out_full_val",), group = "reference", type = Float64),
        "y_val" => (; dims = ("out_val",), group = "reference", type = Float64),
        "y_full_val" => (; dims = ("out_full_val",), group = "reference", type = Float64),
        "P_pca_val" => (; dims = ("out_full_val", "out_val"), group = "reference", type = Float64),
        "num_vars_val" => (; dims = ("config_val",), group = "reference", type = Int16),
        "var_dof_val" => (; dims = ("config_val",), group = "reference", type = Int16),
        "config_pca_dim_val" => (; dims = ("config_val",), group = "reference", type = Int16),
        "config_name_val" => (; dims = ("config_val",), group = "reference", type = String),
        "config_z_obs_val" => (; dims = ("config_val", "dof_val"), group = "reference", type = Float64),
        "norm_factor_val" => (; dims = ("config_val", "config_field_val"), group = "reference", type = Float64),
    )
    return io_dict
end
function io_dictionary_val_reference(
    ref_stats::ReferenceStatistics,
    ref_models::Vector{ReferenceModel},
    write_full_stats::Bool = true,
)
    orig_dict = io_dictionary_val_reference()
    d_full = full_length(ref_stats)
    d = pca_length(ref_stats)
    num_vars = [length(norm_scale) for norm_scale in ref_stats.norm_vec]
    var_dof = Int.([size(P_pca, 1) for P_pca in ref_stats.pca_vec] ./ num_vars)
    config_pca_dim = [size(P_pca, 2) for P_pca in ref_stats.pca_vec]
    config_name = [
        rm.case_name == "LES_driven_SCM" ? join(split(basename(rm.y_dir), ".")[2:end], "_") : rm.case_name for
        rm in ref_models
    ]

    config_z_obs = zeros(length(ref_models), maximum(var_dof))
    for (i, rm) in enumerate(ref_models)
        z_obs = get_z_obs(rm)
        config_z_obs[i, 1:length(z_obs)] = z_obs
    end

    P_pca_full = zeros(d_full, d)
    idx_row = 1
    idx_col = 1
    for P_pca in ref_stats.pca_vec
        rows, cols = size(P_pca)
        P_pca_full[idx_row:(idx_row + rows - 1), idx_col:(idx_col + cols - 1)] = P_pca
        idx_row += rows
        idx_col += cols
    end
    io_dict = Dict(
        "Gamma_val" => Base.setindex(orig_dict["Gamma_val"], ref_stats.Γ, :field),
        "Gamma_full_diag_val" =>
            Base.setindex(orig_dict["Gamma_full_diag_val"], Array(diag(ref_stats.Γ_full)), :field),
        "y_val" => Base.setindex(orig_dict["y_val"], ref_stats.y, :field),
        "y_full_val" => Base.setindex(orig_dict["y_full_val"], ref_stats.y_full, :field),
        "P_pca_val" => Base.setindex(orig_dict["P_pca_val"], P_pca_full, :field),
        "num_vars_val" => Base.setindex(orig_dict["num_vars_val"], num_vars, :field),
        "var_dof_val" => Base.setindex(orig_dict["var_dof_val"], var_dof, :field),
        "config_pca_dim_val" => Base.setindex(orig_dict["config_pca_dim_val"], config_pca_dim, :field),
        "config_name_val" => Base.setindex(orig_dict["config_name_val"], config_name, :field),
        "config_z_obs_val" => Base.setindex(orig_dict["config_z_obs_val"], config_z_obs, :field),
    )
    max_num_fields = maximum([length(norm_vec) for norm_vec in ref_stats.norm_vec])
    norm_factor = zeros(length(ref_stats.norm_vec), max_num_fields)
    for (i, norm_vec) in enumerate(ref_stats.norm_vec)
        num_fields = length(norm_vec)
        norm_factor[i, 1:num_fields] = norm_vec
    end
    io_dict["norm_factor_val"] = Base.setindex(orig_dict["norm_factor_val"], norm_factor, :field)

    if write_full_stats
        io_dict["Gamma_full_val"] = Base.setindex(orig_dict["Gamma_full_val"], Array(ref_stats.Γ_full), :field)
    end
    return io_dict
end

"""
    io_dictionary_prior()
    io_dictionary_prior(priors::ParameterDistribution)

Parameter prior diagnostics dictionary.

Elements:

 - `u_mean_prior` :: Prior mean in unconstrained parameter space.
 - `phi_mean_prior` :: Prior mean in constrained parameter space.
 - `u_var_prior` :: Diagonal of the prior covariance in unconstrained space.
 - `phi_low_unc_prior` :: Lower uncertainty bound (μ-1σ_prior) of prior in constrained space.
 - `phi_upp_unc_prior` :: Upper uncertainty bound (μ+1σ_prior) of prior in constrained space.
 - `phi_low_std_prior` :: Lower standard bound (μ-1) of prior in constrained space. Useful measure of minimum allowed values for bounded parameters.
 - `phi_upp_std_prior` :: Upper standard bound (μ+1) of prior in constrained space. Useful measure of maximum allowed values for bounded parameters.
"""
function io_dictionary_prior()
    io_dict = Dict(
        "u_mean_prior" => (; dims = ("param",), group = "prior", type = Float64),
        "phi_mean_prior" => (; dims = ("param",), group = "prior", type = Float64),
        "u_var_prior" => (; dims = ("param",), group = "prior", type = Float64),
        "phi_low_unc_prior" => (; dims = ("param",), group = "prior", type = Float64),
        "phi_upp_unc_prior" => (; dims = ("param",), group = "prior", type = Float64),
        "phi_low_std_prior" => (; dims = ("param",), group = "prior", type = Float64),
        "phi_upp_std_prior" => (; dims = ("param",), group = "prior", type = Float64),
    )
    return io_dict
end
function io_dictionary_prior(priors::ParameterDistribution)
    orig_dict = io_dictionary_prior()
    u_mean = mean(priors)
    u_var = var(priors)
    # The estimator of the mean is valid in unconstrained space, so we must transform the mean.
    ϕ_mean = transform_unconstrained_to_constrained(priors, u_mean)
    # Transform prior uncertainty bands to constrained space
    u_low = u_mean .- sqrt.(u_var)
    u_upp = u_mean .+ sqrt.(u_var)
    ϕ_low = transform_unconstrained_to_constrained(priors, u_low)
    ϕ_upp = transform_unconstrained_to_constrained(priors, u_upp)
    # Transform standard uncertainty bands (1σ in unconstrained space) to constrained space
    ϕ_low_std = transform_unconstrained_to_constrained(priors, u_mean .- 1.0)
    ϕ_upp_std = transform_unconstrained_to_constrained(priors, u_mean .+ 1.0)

    io_dict = Dict(
        "u_mean_prior" => Base.setindex(orig_dict["u_mean_prior"], u_mean, :field),
        "phi_mean_prior" => Base.setindex(orig_dict["phi_mean_prior"], ϕ_mean, :field),
        "u_var_prior" => Base.setindex(orig_dict["u_var_prior"], u_var, :field),
        "phi_low_unc_prior" => Base.setindex(orig_dict["phi_low_unc_prior"], ϕ_low, :field),
        "phi_upp_unc_prior" => Base.setindex(orig_dict["phi_upp_unc_prior"], ϕ_upp, :field),
        "phi_low_std_prior" => Base.setindex(orig_dict["phi_low_std_prior"], ϕ_low_std, :field),
        "phi_upp_std_prior" => Base.setindex(orig_dict["phi_upp_std_prior"], ϕ_upp_std, :field),
    )
    return io_dict
end

"""
    io_dictionary_metrics()
    io_dictionary_metrics(ekp::EnsembleKalmanProcess, mse_full::Vector{FT}) where {FT <: Real}

Scalar metrics dictionary.

Evaluations of the data-model mismatch in inverse problem (i.e., latent) space are denoted `loss`.
Errors computed in normalized physical (i.e., full) space are denoted `mse_full`. Differences between
these two metrics include:
 - Covariance matrix defining the inner product (covariance weighting in `loss` vs L2 norm in `mse_full`),
 - Treatment of trailing eigenvalues (truncation and regularization vs considering all eigenmodes).
 - The `loss` includes the L2 penalty term, `mse_full` does not.

Elements:

 - `loss_mean_g` :: `(ḡ - y)'Γ_inv(ḡ - y)`. This is the ensemble mean loss seen by the Kalman inversion process.
 - `loss_mean` :: Ensemble mean of `(g - y)'Γ_inv(g - y)`.
 - `loss_min` :: Ensemble min of `(g - y)'Γ_inv(g - y)`.
 - `loss_max` :: Ensemble max of `(g - y)'Γ_inv(g - y)`.
 - `loss_var` :: Variance estimate of `(g - y)'Γ_inv(g - y)`, empirical (EKI/EKS) or quadrature (UKI).
 - `loss_nn_mean` :: `(g_nn - y)'Γ_inv(nn - y)`, where `g_nn` is the forward model output at the particle closest to the mean in parameter space.
 - `mse_full_mean` :: Ensemble mean of MSE(`g_full`, `y_full`).
 - `mse_full_min` :: Ensemble min of MSE(`g_full`, `y_full`).
 - `mse_full_max` :: Ensemble max of MSE(`g_full`, `y_full`).
 - `mse_full_var` :: Variance estimate of MSE(`g_full`, `y_full`), empirical (EKI/EKS) or quadrature (UKI).
 - `mse_full_nn_mean` :: MSE(`g_full`, `y_full`) of particle closest to the mean in parameter space. The mean in parameter space is the solution to the particle-based inversion.
 - `failures` :: Number of particle failures per iteration. If the calibration is run with the "high_loss" failure handler, this diagnostic will not capture the failures due to masking.
 - `nn_mean_index` :: Particle index of the nearest neighbor to the ensemble mean in parameter space. This index is used to construct `..._nn_mean` metrics.
"""
function io_dictionary_metrics()
    io_dict = Dict(
        "loss_mean_g" => (; dims = ("iteration",), group = "metrics", type = Float64),
        "loss_mean" => (; dims = ("iteration",), group = "metrics", type = Float64),
        "loss_min" => (; dims = ("iteration",), group = "metrics", type = Float64),
        "loss_max" => (; dims = ("iteration",), group = "metrics", type = Float64),
        "loss_var" => (; dims = ("iteration",), group = "metrics", type = Float64),
        "loss_nn_mean" => (; dims = ("iteration",), group = "metrics", type = Float64),
        "mse_full_mean" => (; dims = ("iteration",), group = "metrics", type = Float64),
        "mse_full_min" => (; dims = ("iteration",), group = "metrics", type = Float64),
        "mse_full_max" => (; dims = ("iteration",), group = "metrics", type = Float64),
        "mse_full_var" => (; dims = ("iteration",), group = "metrics", type = Float64),
        "mse_full_nn_mean" => (; dims = ("iteration",), group = "metrics", type = Float64),
        "failures" => (; dims = ("iteration",), group = "metrics", type = Int16),
        "nn_mean_index" => (; dims = ("iteration",), group = "metrics", type = Int16),
    )
    return io_dict
end
function io_dictionary_metrics(ekp::EnsembleKalmanProcess, mse_full::Vector{FT}) where {FT <: Real}
    orig_dict = io_dictionary_metrics()

    # Get failures
    failures = length(filter(isnan, mse_full))

    # Get nearest_to_mean point
    nn_mean = get_mean_nearest_neighbor(ekp)

    # Failure-safe variance
    mse_full_var = get_metric_var(ekp, mse_full)
    # Get mse at nearest_to_mean point
    mse_full_nn_mean = mse_full[nn_mean]

    # Get loss (latent space)
    loss = compute_ensemble_loss(ekp)
    # Failure-safe variance
    loss_var = get_metric_var(ekp, loss)
    # Get loss at nearest_to_mean point
    loss_nn_mean = loss[nn_mean]

    # Filter NaNs
    loss_filt = filter(!isnan, loss)
    mse_filt = filter(!isnan, mse_full)

    io_dict = Dict(
        "loss_mean_g" => Base.setindex(orig_dict["loss_mean_g"], get_error(ekp)[end], :field),
        "loss_mean" => Base.setindex(orig_dict["loss_mean"], mean(loss_filt), :field),
        "loss_min" => Base.setindex(orig_dict["loss_min"], minimum(loss_filt), :field),
        "loss_max" => Base.setindex(orig_dict["loss_max"], maximum(loss_filt), :field),
        "loss_var" => Base.setindex(orig_dict["loss_var"], loss_var, :field),
        "loss_nn_mean" => Base.setindex(orig_dict["loss_nn_mean"], loss_nn_mean, :field),
        "mse_full_mean" => Base.setindex(orig_dict["mse_full_mean"], mean(mse_filt), :field),
        "mse_full_min" => Base.setindex(orig_dict["mse_full_min"], minimum(mse_filt), :field),
        "mse_full_max" => Base.setindex(orig_dict["mse_full_max"], maximum(mse_filt), :field),
        "mse_full_var" => Base.setindex(orig_dict["mse_full_var"], mse_full_var, :field),
        "mse_full_nn_mean" => Base.setindex(orig_dict["mse_full_nn_mean"], mse_full_nn_mean, :field),
        "failures" => Base.setindex(orig_dict["failures"], failures, :field),
        "nn_mean_index" => Base.setindex(orig_dict["nn_mean_index"], nn_mean, :field),
    )
    return io_dict
end

"""
    io_dictionary_val_metrics()
    io_dictionary_val_metrics(ekp::EnsembleKalmanProcess, mse_full::Vector{FT}) where {FT <: Real}

Dictionary of scalar validation metrics.

Evaluations of the data-model mismatch in inverse problem (i.e., latent) space are denoted `loss`.
Errors computed in normalized physical (i.e., full) space are denoted `mse_full`. Differences between
these two metrics include:
 - Covariance matrix defining the inner product (covariance weighting in `loss` vs L2 norm in `mse_full`),
 - Treatment of trailing eigenvalues (truncation and regularization vs considering all eigenmodes).
 - The `loss` includes the L2 penalty term, `mse_full` does not.

Elements:

 - `val_loss_mean` :: Ensemble mean of validation `(g - y)'Γ_inv(g - y)`.
 - `val_loss_min` :: Ensemble min of validation `(g - y)'Γ_inv(g - y)`.
 - `val_loss_max` :: Ensemble max of validation `(g - y)'Γ_inv(g - y)`.
 - `val_loss_var` :: Variance estimate of validation `(g - y)'Γ_inv(g - y)`, empirical (EKI/EKS) or quadrature (UKI).
 - `val_loss_nn_mean` :: Validation `(g_nn - y)'Γ_inv(nn - y)`, where `g_nn` is the validation forward model output at the particle closest to the mean in parameter space.
 - `val_mse_full_mean` :: Ensemble mean of MSE(`g_full_val`, `y_full_val`).
 - `val_mse_full_min` :: Ensemble min of MSE(`g_full_val`, `y_full_val`).
 - `val_mse_full_max` :: Ensemble max of MSE(`g_full_val`, `y_full_val`).
 - `val_mse_full_var` :: Variance estimate of MSE(`g_full_val`, `y_full_val`), empirical (EKI/EKS) or quadrature (UKI).
 - `val_mse_full_nn_mean` :: MSE(`g_full_val`, `y_full_val`) of particle closest to the mean in parameter space. The mean in parameter space is the solution to the particle-based inversion.
"""
function io_dictionary_val_metrics()
    io_dict = Dict(
        "val_loss_mean" => (; dims = ("iteration",), group = "metrics", type = Float64),
        "val_loss_min" => (; dims = ("iteration",), group = "metrics", type = Float64),
        "val_loss_max" => (; dims = ("iteration",), group = "metrics", type = Float64),
        "val_loss_var" => (; dims = ("iteration",), group = "metrics", type = Float64),
        "val_loss_nn_mean" => (; dims = ("iteration",), group = "metrics", type = Float64),
        "val_mse_full_mean" => (; dims = ("iteration",), group = "metrics", type = Float64),
        "val_mse_full_min" => (; dims = ("iteration",), group = "metrics", type = Float64),
        "val_mse_full_max" => (; dims = ("iteration",), group = "metrics", type = Float64),
        "val_mse_full_var" => (; dims = ("iteration",), group = "metrics", type = Float64),
        "val_mse_full_nn_mean" => (; dims = ("iteration",), group = "metrics", type = Float64),
    )
    return io_dict
end
function io_dictionary_val_metrics(
    ekp::EnsembleKalmanProcess,
    val_ref_stats::ReferenceStatistics,
    g_val::Matrix{FT},
    val_mse_full::Vector{FT},
) where {FT <: Real}
    orig_dict = io_dictionary_val_metrics()

    # Get nearest_to_mean point
    nn_mean = get_mean_nearest_neighbor(ekp)

    # Failure-safe variance
    val_mse_full_var = get_metric_var(ekp, val_mse_full)
    # Get mse at nearest_to_mean point
    val_mse_full_nn_mean = val_mse_full[nn_mean]

    # Get loss (latent space), augmenting val_ref_stats if necessary
    d_aug = size(g_val, 1)
    d = length(val_ref_stats.y)
    if d_aug > d
        y_val = zeros(d_aug)
        y_val[1:d] = val_ref_stats.y
        y_val[(d + 1):d_aug] = ekp.obs_mean[(end - d_aug + d + 1):end]
        Γ_θ = ekp.obs_noise_cov[(end - d_aug + d + 1):end, (end - d_aug + d + 1):end]
        Γ_val = cat([val_ref_stats.Γ, Γ_θ]..., dims = (1, 2))
    else
        y_val = val_ref_stats.y
        Γ_val = val_ref_stats.Γ
    end

    val_loss = compute_ensemble_loss(g_val, y_val, Γ_val)
    # Failure-safe variance
    val_loss_var = get_metric_var(ekp, val_loss)
    # Get loss at nearest_to_mean point
    val_loss_nn_mean = val_loss[nn_mean]

    # Filter NaNs
    val_loss_filt = filter(!isnan, val_loss)
    val_mse_filt = filter(!isnan, val_mse_full)

    io_dict = Dict(
        "val_loss_mean" => Base.setindex(orig_dict["val_loss_mean"], mean(val_loss_filt), :field),
        "val_loss_min" => Base.setindex(orig_dict["val_loss_min"], minimum(val_loss_filt), :field),
        "val_loss_max" => Base.setindex(orig_dict["val_loss_max"], maximum(val_loss_filt), :field),
        "val_loss_var" => Base.setindex(orig_dict["val_loss_var"], val_loss_var, :field),
        "val_loss_nn_mean" => Base.setindex(orig_dict["val_loss_nn_mean"], val_loss_nn_mean, :field),
        "val_mse_full_mean" => Base.setindex(orig_dict["val_mse_full_mean"], mean(val_mse_filt), :field),
        "val_mse_full_min" => Base.setindex(orig_dict["val_mse_full_min"], minimum(val_mse_filt), :field),
        "val_mse_full_max" => Base.setindex(orig_dict["val_mse_full_max"], maximum(val_mse_filt), :field),
        "val_mse_full_var" => Base.setindex(orig_dict["val_mse_full_var"], val_mse_full_var, :field),
        "val_mse_full_nn_mean" => Base.setindex(orig_dict["val_mse_full_nn_mean"], val_mse_full_nn_mean, :field),
    )
    return io_dict
end

"""
    io_dictionary_particle_state()
    io_dictionary_particle_state(ekp::EnsembleKalmanProcess, priors::ParameterDistribution)

Dictionary of particle-wise parameter diagnostics, not involving forward model evaluations.

Elements:

 - `u`   :: Parameter ensemble in unconstrained (inverse problem) space.
 - `phi` :: Parameter ensemble in constrained (physical) space.
"""
function io_dictionary_particle_state()
    io_dict = Dict(
        "u" => (; dims = ("particle", "param", "iteration"), group = "particle_diags", type = Float64),
        "phi" => (; dims = ("particle", "param", "iteration"), group = "particle_diags", type = Float64),
    )
    return io_dict
end
function io_dictionary_particle_state(ekp::EnsembleKalmanProcess, priors::ParameterDistribution)
    orig_dict = io_dictionary_particle_state()
    u = get_u_final(ekp)
    ϕ = transform_unconstrained_to_constrained(priors, u)
    io_dict =
        Dict("u" => Base.setindex(orig_dict["u"], u', :field), "phi" => Base.setindex(orig_dict["phi"], ϕ', :field))
    return io_dict
end

"""
    io_dictionary_particle_eval()
    io_dictionary_particle_eval(
        ekp::EnsembleKalmanProcess,
        g_full::Matrix{FT},
        mse_full::Vector{FT},
        d::IT,
        d_full::IT,
        batch_indices::Vector{IT},
    ) where {FT <: Real, IT <: Integer}

Dictionary of particle-wise diagnostics involving forward model evaluations.

Elements:

 - `g` :: Forward model evaluation in inverse problem space.
 - `g_full` :: Forward model evaluation in primitive output space, normalized using the pooled field covariance.
 - `mse_full` :: Particle-wise evaluation of MSE(`g_full`, `y_full`).
 - `batch_indices` :: Indices of `ReferenceModel`s evaluated per iteration.
"""
function io_dictionary_particle_eval()
    io_dict = Dict(
        "g" => (; dims = ("particle", "out_aug", "iteration"), group = "particle_diags", type = Float64),
        "g_full" => (; dims = ("particle", "out_full", "iteration"), group = "particle_diags", type = Float64),
        "mse_full" => (; dims = ("particle", "iteration"), group = "particle_diags", type = Float64),
        "batch_indices" => (; dims = ("batch_index", "iteration"), group = "particle_diags", type = Int16),
    )
    return io_dict
end
function io_dictionary_particle_eval(
    ekp::EnsembleKalmanProcess,
    g_full::Matrix{FT},
    mse_full::Vector{FT},
    d::IT,
    d_full::IT,
    batch_indices::Vector{IT},
) where {FT <: Real, IT <: Integer}
    orig_dict = io_dictionary_particle_eval()

    g_aug = get_g_final(ekp)
    d_batch, N_ens = size(g_aug)
    # Fill "g" array with zeros and modify leading rows with possibly batched `g`
    g_filled = zeros(d, N_ens)
    g_filled[1:d_batch, :] = g_aug
    # Fill "g_full" array with zeros and modify leading rows with possibly batched `g`
    d_full_batch = size(g_full, 1)
    g_full_filled = zeros(d_full, N_ens)
    g_full_filled[1:d_full_batch, :] = g_full

    io_dict = Dict(
        "g" => Base.setindex(orig_dict["g"], g_filled', :field), # Avoid params in augmented state
        "g_full" => Base.setindex(orig_dict["g_full"], g_full_filled', :field),
        "mse_full" => Base.setindex(orig_dict["mse_full"], mse_full, :field),
        "batch_indices" => Base.setindex(orig_dict["batch_indices"], batch_indices, :field),
    )
    return io_dict
end

"""
    io_dictionary_val_particle_eval()
    io_dictionary_val_particle_eval(
        g::Matrix{FT},
        g_full::Matrix{FT},
        mse_full::Vector{FT},
        d::IT,
        d_full::IT,
        batch_indices::Vector{IT},
    ) where {FT <: Real, IT <: Integer}

Dictionary of particle-wise validation diagnostics involving forward model evaluations.

Elements:

 - `val_g` :: Validation forward model evaluation in reduced space.
 - `val_g_full` :: Validation forward model evaluation in primitive output space, normalized using the pooled field covariance.
 - `val_mse_full` :: Particle-wise evaluation of MSE(`val_g_full`, `val_y_full`).
 - `val_batch_indices` :: Indices of validation `ReferenceModel`s evaluated per iteration.
"""
function io_dictionary_val_particle_eval()
    io_dict = Dict(
        "val_g" => (; dims = ("particle", "out_aug_val", "iteration"), group = "particle_diags", type = Float64),
        "val_g_full" =>
            (; dims = ("particle", "out_full_val", "iteration"), group = "particle_diags", type = Float64),
        "val_mse_full" => (; dims = ("particle", "iteration"), group = "particle_diags", type = Float64),
        "val_batch_indices" => (; dims = ("batch_index_val", "iteration"), group = "particle_diags", type = Int16),
    )
    return io_dict
end
function io_dictionary_val_particle_eval(
    g::Matrix{FT},
    g_full::Matrix{FT},
    mse_full::Vector{FT},
    d_aug::IT,
    d_full::IT,
    batch_indices::Vector{IT},
) where {FT <: Real, IT <: Integer}
    orig_dict = io_dictionary_val_particle_eval()

    d_batch, N_ens = size(g)
    # Fill "g" array with zeros and modify leading rows with possibly batched `g`
    g_filled = zeros(d_aug, N_ens)
    g_filled[1:d_batch, :] = g
    # Fill "g_full" array with zeros and modify leading rows with possibly batched `g`
    d_full_batch = size(g_full, 1)
    g_full_filled = zeros(d_full, N_ens)
    g_full_filled[1:d_full_batch, :] = g_full

    io_dict = Dict(
        "val_g" => Base.setindex(orig_dict["val_g"], g_filled', :field),
        "val_g_full" => Base.setindex(orig_dict["val_g_full"], g_full_filled', :field),
        "val_mse_full" => Base.setindex(orig_dict["val_mse_full"], mse_full, :field),
        "val_batch_indices" => Base.setindex(orig_dict["val_batch_indices"], batch_indices, :field),
    )
    return io_dict
end

"""
    io_dictionary_ensemble()
    io_dictionary_ensemble(ekp::EnsembleKalmanProcess, priors::ParameterDistribution)

Dictionary of ensemble parameter diagnostics.

Elements:

 - `u_mean` :: Ensemble mean parameter in unconstrained (inverse problem) space.
 - `phi_mean` :: Ensemble mean parameter in constrained (physical) space.
 - `u_cov` :: Sample parameter covariance in unconstrained (inverse problem) space.
 - `phi_cov` :: Sample parameter covariance in constrained (physical) space.
 - `phi_low_unc` :: Lower uncertainty bound (μ-1σ) of the parameter value in constrained (physical) space.
 - `phi_upp_unc` :: Upper uncertainty bound (μ+1σ) of the parameter value in constrained (physical) space.
"""
function io_dictionary_ensemble()
    io_dict = Dict(
        "u_mean" => (; dims = ("param", "iteration"), group = "ensemble_diags", type = Float64),
        "phi_mean" => (; dims = ("param", "iteration"), group = "ensemble_diags", type = Float64),
        "u_cov" => (; dims = ("param", "param", "iteration"), group = "ensemble_diags", type = Float64),
        "phi_cov" => (; dims = ("param", "param", "iteration"), group = "ensemble_diags", type = Float64),
        "phi_low_unc" => (; dims = ("param", "iteration"), group = "ensemble_diags", type = Float64),
        "phi_upp_unc" => (; dims = ("param", "iteration"), group = "ensemble_diags", type = Float64),
    )
    return io_dict
end
function io_dictionary_ensemble(ekp::EnsembleKalmanProcess, priors::ParameterDistribution)
    orig_dict = io_dictionary_ensemble()
    u_mean = get_u_mean(ekp)
    u_cov = get_u_cov(ekp)
    # The estimator of the mean is valid in unconstrained space, so we must transform the mean.
    ϕ_mean = transform_unconstrained_to_constrained(priors, u_mean)
    # Transform uncertainty bands to constrained space
    u_low = u_mean .- sqrt.(diag(u_cov))
    u_upp = u_mean .+ sqrt.(diag(u_cov))
    ϕ_low = transform_unconstrained_to_constrained(priors, u_low)
    ϕ_upp = transform_unconstrained_to_constrained(priors, u_upp)
    # The covariance of ϕ is not the transformed covariance, this is just a linear approximator.
    ϕ_cov = get_ϕ_cov(ekp, priors)
    io_dict = Dict(
        "u_mean" => Base.setindex(orig_dict["u_mean"], u_mean, :field),
        "phi_mean" => Base.setindex(orig_dict["phi_mean"], ϕ_mean, :field),
        "u_cov" => Base.setindex(orig_dict["u_cov"], u_cov, :field),
        "phi_cov" => Base.setindex(orig_dict["phi_cov"], ϕ_cov, :field),
        "phi_low_unc" => Base.setindex(orig_dict["phi_low_unc"], ϕ_low, :field),
        "phi_upp_unc" => Base.setindex(orig_dict["phi_upp_unc"], ϕ_upp, :field),
    )
    return io_dict
end

function get_u_mean(ekp::EnsembleKalmanProcess)
    if isa(ekp.process, Unscented)
        return get_u_mean_final(ekp)
    else
        u = get_u_final(ekp)
        return vcat(mean(u, dims = 2)...)
    end
end

function get_u_cov(ekp::EnsembleKalmanProcess)
    if isa(ekp.process, Unscented)
        return deepcopy(ekp.process.uu_cov[end])
    else
        u = get_u_final(ekp)
        return cov(u, dims = 2)
    end
end

"""
    get_ϕ_cov(ekp::EnsembleKalmanProcess, priors::ParameterDistribution)

Get the last parameter covariance estimate in constrained (physical) space.

For ensemble methods, the covariance of the transformed parameters is returned.
For unscented methods, the covariance is computed through a quadrature on the
transformed quadrature points. The covariance of the transformed parameters
returned here is equal to the transformed covariance only under a first order
Taylor approximation, which is consistent with other approximations underlying the
calibration method.

Inputs:

 - `ekp`    :: The EnsembleKalmanProcess.
 - `priors` :: The priors defining transformations between constrained and unconstrained space.

Outputs:

 - The parameter covariance in constrained space.

"""
function get_ϕ_cov(ekp::EnsembleKalmanProcess, priors::ParameterDistribution)
    if isa(ekp.process, Unscented)
        u_mean = get_u_mean_final(ekp)
        u_cov = deepcopy(ekp.process.uu_cov[end])
        u_p = construct_sigma_ensemble(ekp.process, u_mean, u_cov)
        ϕ_p = transform_unconstrained_to_constrained(priors, u_p)
        ϕ_p_mean = construct_mean(ekp, ϕ_p)
        return construct_cov(ekp, ϕ_p, ϕ_p_mean)
    else
        u = get_u_final(ekp)
        ϕ = transform_unconstrained_to_constrained(priors, u)
        return cov(ϕ, dims = 2)
    end
end

"""
    get_metric_var(ekp::EnsembleKalmanProcess, metric::Vector{FT}) where {FT <: Real}

Computes the ensemble variance of a scalar metric.

For ensemble methods, the sample variance of the metric is returned. For unscented methods,
the variance is computed through a quadrature. Ensemble members where the metric is `NaN`
are filtered out of the computation.

Inputs:

 - `ekp`    :: The EnsembleKalmanProcess.
 - `metric` :: A vector containing the value of the metric for each ensemble member.

Outputs:

 - The ensemble variance of `metric`.

"""
function get_metric_var(ekp::EnsembleKalmanProcess, metric::Vector{FT}) where {FT <: Real}
    if isa(ekp.process, Unscented)
        if any(isnan.(metric))
            succ_ens = [i for i = 1:length(metric) if !isnan(metric[i])]
            metric_mean = construct_successful_mean(ekp, metric, succ_ens)
            return construct_successful_cov(ekp, metric, metric_mean, succ_ens)
        else
            metric_mean = construct_mean(ekp, metric)
            return construct_cov(ekp, metric, metric_mean)
        end
    else
        return var(filter(!isnan, metric))
    end
end

"""
    get_mean_nearest_neighbor(ekp::EnsembleKalmanProcess)

Returns the index of the nearest neighbor to the ensemble mean parameter, in unconstrained space.
"""
function get_mean_nearest_neighbor(ekp::EnsembleKalmanProcess)
    u = get_u_final(ekp)
    u_mean = mean(u, dims = 2)
    return argmin(vcat(sum((u .- u_mean) .^ 2, dims = 1)...))
end

"""
    compute_ensemble_loss(
        g::AbstractMatrix{FT},
        y::AbstractVector{FT},
        Γ::Union{AbstractMatrix{FT}, UniformScaling{FT}},
    ) where {FT <: Real}
    compute_ensemble_loss(ekp::EnsembleKalmanProcess)

Computes the covariance-weighted error `(g - y)'Γ_inv(g - y)` for each ensemble member.
"""
function compute_ensemble_loss(
    g::AbstractMatrix{FT},
    y::AbstractVector{FT},
    Γ::Union{AbstractMatrix{FT}, UniformScaling{FT}},
) where {FT <: Real}
    diff = g .- y # [d, N_ens]
    loss = zeros(size(g, 2)) # [N_ens, 1]
    for i in 1:size(g, 2)
        Γ_inv_diff = Γ \ diff[:, i] # [d, 1]
        loss[i] = dot(diff[:, i], Γ_inv_diff)
    end
    return loss
end
function compute_ensemble_loss(ekp::EnsembleKalmanProcess)
    g = get_g_final(ekp)
    return compute_ensemble_loss(g, ekp.obs_mean, ekp.obs_noise_cov)
end

end # module


text:
using OrdinaryDiffEq, Test, Statistics
using DiffEqProblemLibrary.ODEProblemLibrary: importodeproblems; importodeproblems()
import DiffEqProblemLibrary.ODEProblemLibrary: prob_ode_2Dlinear
prob = deepcopy(prob_ode_2Dlinear)
prob2 = ODEProblem(prob.f,prob.u0,(1.0,0.0),1.01)

sol = solve(prob2,DP5(),dt=-1/4,tstops=[0.5])
@test sol.t == [1.0,.75,.5,0]

sol = solve(prob2,RK4(),dt=-1/4, adaptive=false)
@test sol.t == [1.00,.75,.5,.25,0]

sol = solve(prob2,RK4(),dt=-1/4,tstops=[0.5,0.33], adaptive=false)
@test sol.t ≈ [1.00,.75,.5,.33,0.08,0]

ext: jl


text:
# Tests for the PDict type.
# Author: Noah C. Benson <n@nben.net>

function compare_test(
    p::DIMM, s::DMUT, ks::Vector{K}, vs::Vector{V}, n::Int
) where {K,V, DIMM <: AbstractDict{K,V}, DMUT <: AbstractDict{K,V}}
    pd = [:set, :delete, :get]
    for i in 1:n
        q = rand(pd)
        if q == :set
            k = rand(ks)
            v = rand(vs)
            p = Air.push(p, k => v)
            push!(s, k => v)
        elseif q == :delete && length(s) > 0
            k = rand(keys(s))
            p = Air.delete(p, k)
            delete!(s, k)
        else
            k = rand(ks)
        end
        @test isequal(p, s)
        @test length(p) == length(s)
        @test get(p, k, nothing) == get(s, k, nothing)
        @test (==)((k => get(p, k, nothing)) in p,
                   (k => get(s, k, nothing)) in s)
    end
end
@testset "PDict" begin
    # Next, using the above function
    syms = Symbol[:a, :b, :c, :d, :e, :f, :g, :h, :i, :j]
    nums = Real[10, 20, 30, 40, 50, 60, 70, 80, 90, 100]
    n = 100
    @testset "PIdDict" begin
        compare_test(Air.PIdDict{Symbol,Real}(),
                     Base.IdDict{Symbol,Real}(),
                     syms, nums, n)
        compare_test(Air.PIdDict{Symbol,Real}(:b=>20, :d=>40, :e=>50),
                     Base.IdDict{Symbol,Real}(:b=>20, :d=>40, :e=>50),
                     syms, nums, n)
    end
    @testset "PDict" begin
        compare_test(Air.PDict{Symbol,Real}(),
                     Base.Dict{Symbol,Real}(),
                     syms, nums, n)
        compare_test(Air.PDict{Symbol,Real}(:b=>20, :d=>40, :e=>50),
                     Base.Dict{Symbol,Real}(:b=>20, :d=>40, :e=>50),
                     syms, nums, n)
    end
end


text:
using ContinuumArrays, LinearAlgebra, FastTransforms, QuasiArrays, Test
import ContinuumArrays: Basis, Weight, Map, LazyQuasiArrayStyle, TransformFactorization, ExpansionLayout

"""
This is a simple implementation of Chebyshev for testing. Use ClassicalOrthogonalPolynomials
for the real implementation.
"""
struct Chebyshev <: Basis{Float64}
    n::Int
end

struct ChebyshevWeight <: Weight{Float64} end

Base.:(==)(::Chebyshev, ::Chebyshev) = true
Base.:(==)(::ChebyshevWeight, ::ChebyshevWeight) = true
Base.axes(T::Chebyshev) = (Inclusion(-1..1), Base.OneTo(T.n))
ContinuumArrays.grid(T::Chebyshev) = chebyshevpoints(Float64, T.n, Val(1))
Base.axes(T::ChebyshevWeight) = (Inclusion(-1..1),)

Base.getindex(::Chebyshev, x::Float64, n::Int) = cos((n-1)*acos(x))
Base.getindex(::ChebyshevWeight, x::Float64) = 1/sqrt(1-x^2)
Base.getindex(w::ChebyshevWeight, ::Inclusion) = w # TODO: make automatic

LinearAlgebra.factorize(L::Chebyshev) =
    TransformFactorization(grid(L), plan_chebyshevtransform(Array{Float64}(undef, size(L,2))))
LinearAlgebra.factorize(L::Chebyshev, n) =
    TransformFactorization(grid(L), plan_chebyshevtransform(Array{Float64}(undef, size(L,2),n),1))

# This is wrong but just for tests
QuasiArrays.layout_broadcasted(::Tuple{ExpansionLayout,Any}, ::typeof(*), a::ApplyQuasiVector{<:Any,typeof(*),<:Tuple{Chebyshev,Any}}, b::Chebyshev) = b * Matrix(I, 5, 5)

struct QuadraticMap{T} <: Map{T} end
struct InvQuadraticMap{T} <: Map{T} end

QuadraticMap() = QuadraticMap{Float64}()
InvQuadraticMap() = InvQuadraticMap{Float64}()

Base.getindex(::QuadraticMap, r::Number) = 2r^2-1
Base.axes(::QuadraticMap{T}) where T = (Inclusion(0..1),)
Base.axes(::InvQuadraticMap{T}) where T = (Inclusion(-1..1),)
Base.getindex(d::InvQuadraticMap, x::Number) = sqrt((x+1)/2)
ContinuumArrays.invmap(::QuadraticMap{T}) where T = InvQuadraticMap{T}()
ContinuumArrays.invmap(::InvQuadraticMap{T}) where T = QuadraticMap{T}()

@testset "Chebyshev" begin
    T = Chebyshev(5)
    w = ChebyshevWeight()
    wT = w .* T
    x = axes(T,1)

    @testset "basics" begin
        F = factorize(T)
        g = grid(F)
        @test T \ exp.(x) == F \ exp.(x) == F \ exp.(g) == chebyshevtransform(exp.(g), Val(1))
        @test all(checkpoints(T) .∈ Ref(axes(T,1)))

        @test T \ [exp.(x) cos.(x)] ≈ [F \ exp.(x) F \ cos.(x)]
    end

    @testset "Weighted" begin
        @test MemoryLayout(w) isa WeightLayout
        @test MemoryLayout(w[Inclusion(0..1)]) isa WeightLayout

        @test wT == wT

        wT2 = w .* T[:,2:4]
        wT3 = wT[:,2:4]
        @test MemoryLayout(wT) isa WeightedBasisLayout
        @test MemoryLayout(wT2) isa WeightedBasisLayout
        @test MemoryLayout(wT3) isa SubWeightedBasisLayout
        @test grid(wT) == grid(wT2) == grid(wT3) == grid(T)

        @test ContinuumArrays.unweighted(wT) ≡ T
        @test ContinuumArrays.unweighted(wT2) ≡ T[:,2:4]
        @test ContinuumArrays.unweighted(wT3) ≡ T[:,2:4]

        @test ContinuumArrays.weight(wT) ≡ ContinuumArrays.weight(wT2) ≡ ContinuumArrays.weight(wT3) ≡ w

        @test wT \ @.(exp(x) / sqrt(1-x^2)) ≈ T \ exp.(x)

        @test wT \ w ≈ [1; zeros(4)]

        @test (x .* wT)[0.1,:] ≈ 0.1 * wT[0.1,:]

        a = wT / wT \ @.(exp(x) / sqrt(1-x^2))
        @test wT \ (a .* T) == I # fake multiplication matrix
    end
    @testset "Mapped" begin
        y = affine(0..1, x)

        @test summary(T[y,:]) == "Chebyshev affine mapped to 0..1"
        @test MemoryLayout(wT[y,:]) isa MappedWeightedBasisLayout
        @test MemoryLayout(w[y] .* T[y,:]) isa MappedWeightedBasisLayout
        @test wT[y,:][[0.1,0.2],1:5] == (w[y] .* T[y,:])[[0.1,0.2],1:5] == (w .* T[:,1:5])[y,:][[0.1,0.2],:]
        @test MemoryLayout(wT[y,1:3]) isa MappedWeightedBasisLayout
        @test wT[y,1:3][[0.1,0.2],1:2] == wT[y[[0.1,0.2]],1:2]

        @testset "QuadraticMap" begin
            m = QuadraticMap()
            mi = InvQuadraticMap()
            @test 0.1 ∈ m
            @test -0.1 ∈ m
            @test 2 ∉ m
            @test 0.1 ∈ mi
            @test -0.1 ∉ mi

            @test m[findfirst(isequal(0.1), m)] ≈ 0.1
            @test m[findlast(isequal(0.1), m)] ≈ 0.1
            @test m[findall(isequal(0.1), m)] ≈ [0.1]

            @test m[Inclusion(0..1)] ≡ m
            @test_throws BoundsError m[Inclusion(-1..1)]
            T = Chebyshev(5)
            M = T[m,:]
            @test MemoryLayout(M) isa MappedBasisLayout
            @test MemoryLayout(M[:,1:3]) isa MappedBasisLayout
            @test M[0.1,:] ≈ T[2*0.1^2-1,:]
            x = axes(M,1)
            @test x == Inclusion(0..1)
            @test M \ exp.(x) ≈ T \ exp.(sqrt.((axes(T,1) .+ 1)/2))
        end
    end

    @testset "Broadcasted" begin
        T = Chebyshev(5)
        F = factorize(T)
        x = axes(T,1)
        a = 1 .+ x .+ x.^2
        # The following are wrong, just testing dispatch
        @test T \ (a .* T) == I
        @test T \ (a .* (T * (T \ a))) isa Vector
        f = exp.(x) .* a # another broadcast layout
        @test T \ f == F \ f

        ã = T * (T \ a)
        @test T \ (ã .* ã) ≈ [1.5,1,0.5,0,0]
    end
end


text:
@testset "Testing Handle" begin
    # We test these operations on a `msr` register
    handle = CounterTools.Handle("/dev/cpu/0/msr")

    # Read the initial state of the CHA 0 control register
    location = CounterTools.IndexZero(0xE01)
    seek(handle, location)

    # Test reading and unsafe_reading
    initial_state = read(handle, UInt64)
    @test initial_state == unsafe_read(handle, UInt64, location)

    # Temporarily write a new value
    event = CounterTools.UncoreSelectRegister(event = 0x50, umask = 0x01)
    seek(handle, location)
    write(handle, event)

    @test read(handle, UInt64) == CounterTools.value(event)
    buffer = Vector{UInt8}(undef, sizeof(UInt64))
    @test unsafe_read(handle, UInt64, location; buffer = buffer) == CounterTools.value(event)

    # Write back the initial state
    seek(handle, location)
    write(handle, initial_state)
    @test read(handle, UInt64) == initial_state
    @test unsafe_read(handle, UInt64, location; buffer = buffer) == initial_state
end