"""
    StateVars{names,griddvars,TU,TV,TI,DF,DG}

Generic container for holding discretized state arrays from symbolic variables (`Var` types).
The `uproto` field represetns a "prototype" of the prognostic state array which should fully
define the state of the system at any given point.
"""
struct StateVars{names,griddvars,TU,TV,TI,DF,DG}
    uproto::TU # prognostic state vector prototype
    vars::NamedTuple{names,TV} # variable metadata
    diag::NamedTuple{names,DF} # off-grid diagnostic variables
    griddiag::NamedTuple{griddvars,DG} # on-grid diagnostic variables
    layeridx::NamedTuple{names,TI}
end

"""
    StateVars(vars::NamedTuple, D::Numerics.AbstractDiscretization, cachector::Function, arrayproto::Type{A}=Vector) where {A<:AbstractVector}

Constructs a `StateVars` container from `vars` and the given discretization `D`. The function `cachector` must accept argments `f(name::Symbol, x::AbstractVector)` where
`name` is the name of the state variable and `x` is the array generated by the discretization.
"""
function StateVars(
    @nospecialize(vars::NamedTuple),
    @nospecialize(D::Numerics.AbstractDiscretization),
    zs::NTuple,
    ::Type{TCache},
    ::Type{TArray}=Vector;
    cache_kwargs...
) where {TCache,TArray<:AbstractVector}
    _flatten(vars) = Flatten.flatten(vars, Flatten.flattenable, Var)
    # select each variable type
    progvars = map(group -> filter(isprognostic, group), vars)
    diagvars = map(group -> filter(isdiagnostic, group), vars)
    algvars = map(group -> filter(isalgebraic, group), vars)
    # create variables for time delta variables (divergence/residual)
    dpvars = map(group -> map(CryoGrid.DVar, filter(var -> isalgebraic(var) || isprognostic(var), group)), vars)
    gridprogvars = Tuple(unique(filter(isongrid, tuplejoin(_flatten(progvars), _flatten(algvars)))))
    freeprogvars = map(group -> filter(!isongrid, group), progvars)
    vartypes = map(vartype, tuplejoin(gridprogvars, _flatten(freeprogvars)))
    @assert all(map(==(first(vartypes)), vartypes)) "All prognostic variables must have same data type"
    uproto = prognosticstate(TArray, D, freeprogvars, gridprogvars)
    # build off-grid (i.e. "free") diagnostic state vectors
    freediagvars = map(group -> filter(!isongrid, group), diagvars)
    freediagstate = map(group -> (;map(v -> varname(v) => TCache(varname(v), instantiate(v, D, TArray); cache_kwargs...), group)...), freediagvars)
    # build on-grid diagnostic state vectors;
    griddiagvars = Tuple(unique(filter(isongrid, _flatten(diagvars))))
    griddiagstate = map(v -> varname(v) => TCache(varname(v), instantiate(v, D, TArray); cache_kwargs...), griddiagvars)
    # join prognostic variables with delta and flux variables, then build nested named tuples in each group with varnames as keys
    allvars = map(vars -> NamedTuple{map(varname, vars)}(vars), map(tuplejoin, vars, dpvars))
    # get layer indices
    zs = tuple(zs..., D[end], D[end])
    layeridx = map((k, z1, z2) -> k => bounds(D[z1..z2]), keys(allvars), zs[1:end-1], zs[2:end])
    return StateVars(uproto, allvars, (;freediagstate...), (;griddiagstate...), (;layeridx...))
end

@generated function getvar(::Val{name}, vs::StateVars{layers,griddvars}, u, du=nothing) where {name,layers,griddvars}
    pax = ComponentArrays.indexmap(first(ComponentArrays.getaxes(u))) # get prognostic variable index map (name -> indices)
    dnames = map(n -> CryoGrid.dname(n), keys(pax)) # get names of delta/derivative variables
    # case 1) variable is diagnostic and lives on the grid
    if name ∈ griddvars
        quote
            return retrieve(vs.griddiag.$name, u)
        end
    # case 2) variable is prognostic
    elseif name ∈ keys(pax)
        quote
            return u.$name
        end
    # case 3) variable is a prognostic derivative or residual
    elseif du != Nothing && name ∈ dnanes
        i = findfirst(n -> n == name, dnames)::Int
        quote
            return du.$(keys(pax)[i])
        end
    # case 4) no variables match the given name
    else
        :(return nothing)
    end
end

function getvars(vs::StateVars{layers,gridvars,TU}, u::ComponentVector, du::ComponentVector, vals::Union{Symbol,<:Pair{Symbol}}...) where {layers,gridvars,T,A,pax,TU<:ComponentVector{T,A,Tuple{Axis{pax}}}}
    # case 1: grid variable (no layer specified)
    isprognostic(name::Symbol) = name ∈ keys(pax)
    # case 2: non-grid variable on specific layer; ignore here, defer until handled below
    isprognostic(other) = false
    symbols(name::Symbol) = tuple(name)
    symbols(names::NTuple{N,Symbol}) where N = names
    # map over non-prognostic variables, selecting variables from cache
    vars = map(filter(!isprognostic, vals)) do val # map over given variable names, ignoring prognostic variables
        # in case val is a differential var (will be nothing otherwise)
        dvar_ind = findfirst(n -> val == CryoGrid.dname(n), keys(pax))
        if !isnothing(dvar_ind)
            val => du[keys(pax)[dvar_ind]]    
        elseif val ∈ gridvars
            val => getvar(Val{val}(), vs, u, du)
        elseif isa(val, Pair)
            layername = val[1]
            # handle either a single variable name or multiple, also filtering out prognostic variables
            layervars = filter(!isprognostic, symbols(val[2]))
            layername => (;map(n -> n => retrieve(getproperty(vs.diag[layername], n)), layervars)...)
        else
            error("no state variable named $val defined")
        end
    end
    return (;vars...)
end

"""
    build_mass_matrix(u::ComponentVector, states::StateVars)

Constructs a mass matrix `M⋅∂u∂t = f(u)` suitable for the prognostic state vector `u` based on the
defined variable types.
"""
function build_mass_matrix(states::StateVars)
    M_diag = similar(states.uproto)
    M_idxmap = ComponentArrays.indexmap(getaxes(M_diag)[1])
    allvars = Flatten.flatten(states.vars, Flatten.flattenable, Var)
    progvars = map(varname, filter(isprognostic, allvars))
    algvars = map(varname, filter(isalgebraic, allvars))
    for name in keys(M_idxmap)
        M_diag_var = @view M_diag[name]
        if isa(M_diag_var, ComponentArray)
            for layervar in keys(M_diag_var)
                M_diag_sub_var = @view M_diag_var[layervar]
                if layervar ∈ progvars
                    M_diag_sub_var .= one(eltype(M_diag))
                elseif layervar ∈ algvars
                    M_diag_sub_var .= zero(eltype(M_diag))
                end
            end
        else
            if name ∈ progvars
                M_diag_var .= one(eltype(M_diag))
            elseif name ∈ algvars
                M_diag_var .= zero(eltype(M_diag))
            end
        end
    end
    # if no algebraic variables are present, use identity matrix
    num_algebraic = length(M_diag) - sum(M_diag)
    M = num_algebraic > 0 ? Diagonal(M_diag) : I
end
