text:
#
println("Aq) Test of the Green(function) expansion.")

if  false
    name          = "Lithium 1s^2 2s ground configuration"
    refConfigs    = [Configuration("[He] 2s")]
    greenSettings = GreenSettings(5, [0, 1, 2], 0.01, true, LevelSelection())
    #
    wa          = Representation(name, Nuclear.Model(8.), Radial.Grid(true), refConfigs,
                                ## GreenExpansion( Atomic.SingleCSFwithoutCI(), Basics.DeExciteSingleElectron(),
                                ## GreenExpansion( Atomic.CoreSpaceCI(), Basics.DeExciteSingleElectron(),
                                    GreenExpansion( AtomicState.DampedSpaceCI(), Basics.DeExciteSingleElectron(),
                                                    [LevelSymmetry(1//2, Basics.plus), LevelSymmetry(3//2, Basics.plus)], 3, greenSettings) )
    println(wa)

    wb = generate(wa, output=true)
    #
elseif true
    # Test example for MS "Approximate atomic Green functions" in molecules (2021)
    name            = "Approximate Green function for neon-like ions."
    refConfigs      = [Configuration("1s^2 2s^2 2p^6")]
    levelSymmetries = [LevelSymmetry(1, Basics.minus)]  ## , LevelSymmetry(2, Basics.minus)]
    greenSettings   = GreenSettings( 50, [0, 1, 2], 0.01, true, LevelSelection())
    greenExpansion  = GreenExpansion(AtomicState.DampedSpaceCI(), Basics.DeExciteSingleElectron(),
    # greenExpansion  = GreenExpansion(AtomicState.DampedSpaceCI(), Basics.DeExciteTwoElectrons(),
    # greenExpansion  = GreenExpansion(AtomicState.SingleCSFwithoutCI(), Basics.DeExciteSingleElectron(),
                                     levelSymmetries, 10, greenSettings)
    #
    rep = Representation(name, Nuclear.Model(18.), Radial.Grid(true), refConfigs, greenExpansion)
    #
    gExp = generate(rep, output=true)
    gChannels = gExp["Green channels"]
    cha1      = gChannels[1]
    @show cha1.gMultiplet.levels[1].J, cha1.gMultiplet.levels[1].energy
    @show cha1.gMultiplet.levels[2].J, cha1.gMultiplet.levels[2].energy
    @show cha1.gMultiplet.levels[3].J, cha1.gMultiplet.levels[3].energy
    Basics.displayLevels(stdout, gChannels)
    #
end


text:
@cache struct TanYam7Cache{uType,rateType,uNoUnitsType,TabType} <: OrdinaryDiffEqMutableCache
  u::uType
  uprev::uType
  fsalfirst::rateType
  k2::rateType
  k3::rateType
  k4::rateType
  k5::rateType
  k6::rateType
  k7::rateType
  k8::rateType
  k9::rateType
  k10::rateType
  utilde::uType
  tmp::uType
  atmp::uNoUnitsType
  k::rateType
  tab::TabType
end

function alg_cache(alg::TanYam7,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Val{true})
  tab = TanYam7ConstantCache(constvalue(uBottomEltypeNoUnits),constvalue(tTypeNoUnits))
  k1 = zero(rate_prototype); k2 = zero(rate_prototype) ; k3 = zero(rate_prototype); k4 = zero(rate_prototype)
  k5 = zero(rate_prototype); k6 = zero(rate_prototype) ; k7 = zero(rate_prototype); k8 = zero(rate_prototype)
  k9 = zero(rate_prototype); k10= zero(rate_prototype) ;
  utilde = zero(u); tmp = zero(u); atmp = similar(u,uEltypeNoUnits); k = zero(rate_prototype)
  TanYam7Cache(u,uprev,k1,k2,k3,k4,k5,k6,k7,k8,k9,k10,utilde,tmp,atmp,k,tab)
end

alg_cache(alg::TanYam7,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Val{false}) = TanYam7ConstantCache(constvalue(uBottomEltypeNoUnits),constvalue(tTypeNoUnits))

@cache struct DP8Cache{uType,rateType,uNoUnitsType,TabType} <: OrdinaryDiffEqMutableCache
  u::uType
  uprev::uType
  k1::rateType
  k2::rateType
  k3::rateType
  k4::rateType
  k5::rateType
  k6::rateType
  k7::rateType
  k8::rateType
  k9::rateType
  k10::rateType
  k11::rateType
  k12::rateType
  k13::rateType
  k14::rateType
  k15::rateType
  k16::rateType
  kupdate::rateType
  udiff::rateType
  bspl::rateType
  dense_tmp3::rateType
  dense_tmp4::rateType
  dense_tmp5::rateType
  dense_tmp6::rateType
  dense_tmp7::rateType
  utilde::uType
  tmp::uType
  atmp::uNoUnitsType
  tab::TabType
end

function alg_cache(alg::DP8,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Val{true})
  k1 = zero(rate_prototype); k2  = zero(rate_prototype); k3  = zero(rate_prototype);  k4 = zero(rate_prototype)
  k5 = zero(rate_prototype); k6  = zero(rate_prototype); k7  = zero(rate_prototype);  k8 = zero(rate_prototype)
  k9 = zero(rate_prototype); k10 = zero(rate_prototype); k11 = zero(rate_prototype); k12 = zero(rate_prototype)
  kupdate = zero(rate_prototype); utilde = zero(u);
  tmp = zero(u); atmp = similar(u,uEltypeNoUnits)
  k13 = zero(rate_prototype)
  k14 = zero(rate_prototype)
  k15 = zero(rate_prototype)
  k16 = zero(rate_prototype)
  udiff = zero(rate_prototype)
  bspl = zero(rate_prototype)
  # dense_tmp1 = udiff
  # dense_tmp2 = bspl
  dense_tmp3 = zero(rate_prototype)
  dense_tmp4 = zero(rate_prototype)
  dense_tmp5 = zero(rate_prototype)
  dense_tmp6 = zero(rate_prototype)
  dense_tmp7 = zero(rate_prototype)
  tab = DP8ConstantCache(constvalue(uBottomEltypeNoUnits),constvalue(tTypeNoUnits))
  DP8Cache(u,uprev,k1,k2,k3,k4,k5,k6,k7,k8,k9,k10,k11,k12,k13,k14,k15,k16,kupdate,
           udiff,bspl,dense_tmp3,dense_tmp4,dense_tmp5,dense_tmp6,dense_tmp7,
           utilde,tmp,atmp,tab)
end

alg_cache(alg::DP8,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Val{false}) = DP8ConstantCache(constvalue(uBottomEltypeNoUnits),constvalue(tTypeNoUnits))

@cache struct TsitPap8Cache{uType,rateType,uNoUnitsType,TabType} <: OrdinaryDiffEqMutableCache
  u::uType
  uprev::uType
  fsalfirst::rateType
  k2::rateType
  k3::rateType
  k4::rateType
  k5::rateType
  k6::rateType
  k7::rateType
  k8::rateType
  k9::rateType
  k10::rateType
  k11::rateType
  k12::rateType
  k13::rateType
  utilde::uType
  tmp::uType
  atmp::uNoUnitsType
  k::rateType
  tab::TabType
end

function alg_cache(alg::TsitPap8,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Val{true})
  tab = TsitPap8ConstantCache(constvalue(uBottomEltypeNoUnits),constvalue(tTypeNoUnits))
  k1 = zero(rate_prototype); k2 = zero(rate_prototype); k3 = zero(rate_prototype); k4 = zero(rate_prototype)
  k5 = zero(rate_prototype); k6 = zero(rate_prototype); k7 = zero(rate_prototype); k8 = zero(rate_prototype)
  k9 = zero(rate_prototype); k10 = zero(rate_prototype); k11 = zero(rate_prototype); k12 = zero(rate_prototype)
  k13 = zero(rate_prototype); utilde = zero(u); k = zero(rate_prototype)
  tmp = zero(u); atmp = similar(u,uEltypeNoUnits);
  TsitPap8Cache(u,uprev,k1,k2,k3,k4,k5,k6,k7,k8,k9,k10,k11,k12,k13,utilde,tmp,atmp,k,tab)
end

alg_cache(alg::TsitPap8,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Val{false}) = TsitPap8ConstantCache(constvalue(uBottomEltypeNoUnits),constvalue(tTypeNoUnits))

@cache struct PFRK87Cache{uType,rateType,uNoUnitsType,TabType} <: OrdinaryDiffEqMutableCache
  u::uType
  uprev::uType
  fsalfirst::rateType
  k2::rateType
  k3::rateType
  k4::rateType
  k5::rateType
  k6::rateType
  k7::rateType
  k8::rateType
  k9::rateType
  k10::rateType
  k11::rateType
  k12::rateType
  k13::rateType
  utilde::uType
  tmp::uType
  atmp::uNoUnitsType
  k::rateType
  tab::TabType
end

function alg_cache(alg::PFRK87,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Val{true})
  tab = PFRK87ConstantCache(constvalue(uBottomEltypeNoUnits),constvalue(tTypeNoUnits))
  k1 = zero(rate_prototype)
  k2 = zero(rate_prototype)
  k3 = zero(rate_prototype)
  k4 = zero(rate_prototype)
  k5 = zero(rate_prototype)
  k6 = zero(rate_prototype)
  k7 = zero(rate_prototype)
  k8 = zero(rate_prototype)
  k9 = zero(rate_prototype)
  k10 = zero(rate_prototype)
  k11 = zero(rate_prototype)
  k12 = zero(rate_prototype)
  k13 = zero(rate_prototype)
  utilde = zero(u)
  k = zero(rate_prototype)
  tmp = zero(u)
  atmp = similar(u,uEltypeNoUnits)
  PFRK87Cache(u,uprev,k1,k2,k3,k4,k5,k6,k7,k8,k9,k10,k11,k12,k13,utilde,tmp,atmp,k,tab)
end

alg_cache(alg::PFRK87,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Val{false}) = PFRK87ConstantCache(constvalue(uBottomEltypeNoUnits),constvalue(tTypeNoUnits))

ext: jl


text:
@testset "show" begin
    @testset "vector" begin
        a = VectorDifference([3], [1], [2], [-1], [3], [4])
        @test sprint(show, MIME("text/plain"), a) == """
        VectorDifference with indices:
         modified: [3]
         added: [1]
         removed: [2]
        and values:
         common: [-1]
         added: [3]
         removed: [4]"""
    end

    @testset "matrix" begin
        a = MatrixDifference(([1], [1]), ([2], [2]), ([3], [3]), [1], [2, 2], [3, 3])
        @test sprint(show, MIME("text/plain"), a) == """
        MatrixDifference with indices:
         modified: ([1], [1])
         added: ([2], [2])
         removed: ([3], [3])
        and values:
         common: [1]
         added: [2, 2]
         removed: [3, 3]"""
    end

    @testset "dictionary" begin
        a = DictDifference(Dict(:x=>1, :y=>[0.0, 1.0]), Dict(), Dict(:z=>3))
        @test sprint(show, MIME("text/plain"), a) == """
        DictDifference with values:
         modified: Dict{Symbol,Any}(:y => [0.0, 1.0],:x => 1)
         added: Dict{Any,Any}()
         removed: Dict(:z => 3)"""
    end

    @testset "named tuple" begin
        a = NamedTupleDifference((x=1, y=[0.0, 1.0]), NamedTuple(), (z=3,))
        @test sprint(show, MIME("text/plain"), a) == """
        NamedTupleDifference with values:
         modified: (x = 1, y = [0.0, 1.0])
         added: NamedTuple()
         removed: (z = 3,)"""
    end

    @testset "set" begin
        a = SetDifference(Set([]), Set([]), Set([]))
        @test sprint(show, MIME("text/plain"), a) == """
        SetDifference with values:
         common: Set{Any}()
         added: Set{Any}()
         removed: Set{Any}()"""
    end
end


text:
import CSV
import DataFrames

function load_data_set()
    df = CSV.read("$(@__DIR__)/coal.csv", DataFrames.DataFrame)
    dates = df[!,1]
    dates = dates .- minimum(dates)
    dates * 365.25 # convert years to days
end


text:
using LinearMaps
using IterativeSolvers

# Newton's algorithm to solve SCF equations
#
# Newton algorithm consists of iterating over density matrices like
#       P <- P + δP, with additional renormalization
# where δP solves ([1])
#       (Ω+K)δP = -[P, [P, H(P)]]
# where (Ω+K) is the constrained Hessian of the energy. It is
# defined as a super operator on the space tangent to the constraints
# manifold at point P_∞, the solution of our problem.
#   - K is the unconstrained Hessian of the energy on the manifold;
#   - Ω represents the influence of the curvature of the manifold :
#         ΩδP = -[P_∞, [H(P_∞), δP]].
#     In practice, we dont have access to P_∞ so we just use the current P.
#     Another way to see Ω is to link it to the four-point independent-particle
#     susceptibility. Indeed, if we define the following extension and
#     retraction operators
#       - E : builds a density matrix Eρ given a density ρ via Eρ(r,r') = δ(r,r')ρ(r),
#       - R : builds a density ρ given a density matrix O via ρ(r) = O(r,r),
#     then we have the relation χ0₂ = Rχ0₄E, where χ0₂ is the two-point
#     independent-particle susceptibility (it returns δρ from a given δV) and
#     χ0₄ = -Ω^{-1} is the four-point independent-particle susceptibility.
#
# For further details :
# [1] Eric Cancès, Gaspard Kemlin, Antoine Levitt. Convergence analysis of
#     direct minimization and self-consistent iterations. SIAM Journal of Matrix
#     Analysis and Applications, 42(1):243–274 (2021)
#
# We presented the algorithm in the density matrix framework, but in practice,
# we implement it with orbitals.
# In this framework, an element on the tangent space at ψ can be written as
#       δP = Σ |ψi><δψi| + hc
# where the δψi are of size Nb and are all orthogonal to the ψj, 1 <= j <= N.
# Therefore we store them in the same kind of array than ψ, with
# δψ[ik][:, i] = δψi for each k-point.
# In this framework :
#   - computing Ωδψ can be done analytically with the formula
#         Ωδψi = H*δψi - Σj λji δψj
#     where λij = <ψi|H|ψj>. This is one way among others to extend the
#     definition of ΩδP = -[P_∞, [H(P_∞), δP]] to any point P on the manifold --
#     we chose this one because it makes Ω self-adjoint for the Frobenius scalar
#     product;
#   - computing Kδψ can be done by constructing the δρ associated with δψ, then
#     using the exchange-correlation kernel to compute the associated δV, then
#     acting with δV on the ψ and projecting on the tangent space.

#  Compute the gradient of the energy, projected on the space tangent to ψ, that
#  is to say H(ψ)*ψ - ψ*λ where λ is the set of Rayleigh coefficients associated
#  to the ψ.
function compute_projected_gradient(basis::PlaneWaveBasis, ψ, occupation)
    ρ = compute_density(basis, ψ, occupation)
    _, H = energy_hamiltonian(basis, ψ, occupation; ρ=ρ)

    [proj_tangent_kpt(H.blocks[ik] * ψk, ψk) for (ik, ψk) in enumerate(ψ)]
end

# Projections on the space tangent to ψ
function proj_tangent_kpt!(δψk, ψk)
    # δψk = δψk - ψk * (ψk'δψk)
    mul!(δψk, ψk, ψk'δψk, -1, 1)
end
proj_tangent_kpt(δψk, ψk) = proj_tangent_kpt!(deepcopy(δψk), ψk)

function proj_tangent(δψ, ψ)
    [proj_tangent_kpt(δψ[ik], ψk) for (ik, ψk) in enumerate(ψ)]
end
function proj_tangent!(δψ, ψ)
    [proj_tangent_kpt!(δψ[ik], ψk) for (ik, ψk) in enumerate(ψ)]
end

"""
    apply_Ω(δψ, ψ, H::Hamiltonian, Λ)

Compute the application of Ω defined at ψ to δψ. H is the Hamiltonian computed
from ψ and Λ is the set of Rayleigh coefficients ψk' * Hk * ψk at each k-point.
"""
function apply_Ω(δψ, ψ, H::Hamiltonian, Λ)
    δψ = proj_tangent(δψ, ψ)
    Ωδψ = [H.blocks[ik] * δψk - δψk * Λ[ik] for (ik, δψk) in enumerate(δψ)]
    proj_tangent!(Ωδψ, ψ)
end

"""
    apply_K(basis::PlaneWaveBasis, δψ, ψ, ρ, occupation)

Compute the application of K defined at ψ to δψ. ρ is the density issued from ψ.
δψ also generates a δρ, computed with `compute_δρ`.
"""
@views function apply_K(basis::PlaneWaveBasis, δψ, ψ, ρ, occupation)
    δψ = proj_tangent(δψ, ψ)
    δρ = compute_δρ(basis, ψ, δψ, occupation)
    δV = apply_kernel(basis, δρ; ρ=ρ)

    Kδψ = map(enumerate(ψ)) do (ik, ψk)
        kpt = basis.kpoints[ik]
        δVψk = similar(ψk)

        for n = 1:size(ψk, 2)
            ψnk_real = G_to_r(basis, kpt, ψk[:, n])
            δVψnk_real = δV[:, :, :, kpt.spin] .* ψnk_real
            δVψk[:, n] = r_to_G(basis, kpt, δVψnk_real)
        end
        δVψk
    end
    # ensure projection onto the tangent space
    proj_tangent!(Kδψ, ψ)
end

"""
    solve_ΩplusK(basis::PlaneWaveBasis{T}, ψ, res, occupation;
                 tol_cg=1e-10, verbose=false) where T

Return δψ where (Ω+K) δψ = rhs
"""
function solve_ΩplusK(basis::PlaneWaveBasis{T}, ψ, rhs, occupation;
                     tol_cg=1e-10, verbose=false) where T
    @assert mpi_nprocs() == 1  # Distributed implementation not yet available
    filled_occ = filled_occupation(basis.model)
    # for now, all orbitals have to be fully occupied -> need to strip them beforehand
    @assert all(all(occ_k .== filled_occ) for occ_k in occupation)

    # compute quantites at the point which define the tangent space
    ρ = compute_density(basis, ψ, occupation)
    _, H = energy_hamiltonian(basis, ψ, occupation; ρ=ρ)

    pack(ψ) = reinterpret_real(pack_ψ(ψ))
    unpack(x) = unpack_ψ(reinterpret_complex(x), size.(ψ))
    unsafe_unpack(x) = unsafe_unpack_ψ(reinterpret_complex(x), size.(ψ))

    # project rhs on the tangent space before starting
    proj_tangent!(rhs, ψ)
    rhs_pack = pack(rhs)

    # preconditioner
    Pks = [PreconditionerTPA(basis, kpt) for kpt in basis.kpoints]
    for ik = 1:length(Pks)
        precondprep!(Pks[ik], ψ[ik])
    end
    function f_ldiv!(x, y)
        δψ = unpack(y)
        proj_tangent!(δψ, ψ)
        Pδψ = [ Pks[ik] \ δψk for (ik, δψk) in enumerate(δψ)]
        proj_tangent!(Pδψ, ψ)
        x .= pack(Pδψ)
    end

    # Rayleigh-coefficients
    Λ = [ψk'Hψk for (ψk, Hψk) in zip(ψ, H * ψ)]

    # mapping of the linear system on the tangent space
    function ΩpK(x)
        δψ = unsafe_unpack(x)
        Kδψ = apply_K(basis, δψ, ψ, ρ, occupation)
        Ωδψ = apply_Ω(δψ, ψ, H, Λ)
        pack(Ωδψ + Kδψ)
    end
    J = LinearMap{T}(ΩpK, size(rhs_pack, 1))

    # solve (Ω+K) δψ = rhs on the tangent space with CG
    δψ, history = cg(J, rhs_pack, Pl=FunctionPreconditioner(f_ldiv!),
                  reltol=0, abstol=tol_cg, verbose=verbose, log=true)

    (; δψ=unpack(δψ), history)
end

# Solves Ω+K using a split algorithm
# With χ04P = -Ω^-1,
# (Ω+K)^-1 = Ω^-1 (1 - K(1+Ω^-1 K)^-1 Ω^-1)
# (Ω+K)^-1 = -χ04P (1 + K(1 - χ04P K)^-1 χ04P)
# (Ω+K)^-1 = -χ04P (1 + E K2P (1 - χ02P K2P)^-1 R χ04P)
# where χ02P = R χ04P E and K2P = R K E
function solve_ΩplusK_split(basis::PlaneWaveBasis{T}, ψ, rhs, occupation;
                            tol_dyson=1e-8, tol_cg=1e-12, verbose=false) where T
    ρ = compute_density(basis, ψ, occupation)
    _, H = energy_hamiltonian(basis, ψ, occupation; ρ=ρ)

    λ = [real.(eigvals(ψk'Hψk)) for (ψk, Hψk) in zip(ψ, H * ψ)]
    occupation, εF = compute_occupation(basis, λ)

    # compute δρ0 (ignoring interactions)
    δψ0 = apply_χ0_4P(H, ψ, occupation, εF, λ, rhs; tol_cg, verbose)
    δρ0 = compute_δρ(basis, ψ, δψ0, occupation)

    pack(δρ) = vec(δρ)
    unpack(δρ) = reshape(δρ, size(ρ))
    # compute total δρ
    function eps_fun(δρ)
        δρ = unpack(δρ)
        δV = apply_kernel(basis, δρ; ρ)
        χ0δV = apply_χ0(H, ψ, εF, λ, δV; tol_cg)
        pack(δρ - χ0δV)
    end
    J = LinearMap{T}(eps_fun, length(pack(δρ0)))
    δρ = unpack(gmres(J, pack(δρ0); reltol=0, abstol=tol_dyson, verbose))
    δV = apply_kernel(basis, δρ; ρ)

    δVψ = [DFTK.RealSpaceMultiplication(basis, kpt, @views δV[:, :, :, kpt.spin]) * ψ[ik]
           for (ik, kpt) in enumerate(basis.kpoints)]
    δψ = apply_χ0_4P(H, ψ, occupation, εF, λ, δVψ; tol_cg, verbose)
    .- (δψ0 .+ δψ)
end

"""
    newton(basis::PlaneWaveBasis{T}; ψ0=nothing,
           tol=1e-6, tol_cg=1e-10, maxiter=20, verbose=false,
           callback=NewtonDefaultCallback(),
           is_converged=NewtonConvergenceDensity(tol))

Newton algorithm. Be careful that the starting point needs to be not too far
from the solution.
"""
function newton(basis::PlaneWaveBasis{T}, ψ0;
                tol=1e-6, tol_cg=1e-10, maxiter=20, verbose=false,
                callback=ScfDefaultCallback(),
                is_converged=ScfConvergenceDensity(tol)) where T

    # setting parameters
    model = basis.model
    @assert model.temperature == 0 # temperature is not yet supported

    # check that there are no virtual orbitals
    filled_occ = filled_occupation(model)
    n_spin = model.n_spin_components
    n_bands = div(model.n_electrons, n_spin * filled_occ, RoundUp)
    @assert n_bands == size(ψ0[1], 2)

    # number of kpoints and occupation
    Nk = length(basis.kpoints)
    occupation = [filled_occ * ones(T, n_bands) for ik = 1:Nk]

    # iterators
    n_iter = 0

    # orbitals, densities and energies to be updated along the iterations
    ψ = deepcopy(ψ0)
    ρ = compute_density(basis, ψ, occupation)
    energies, H = energy_hamiltonian(basis, ψ, occupation; ρ=ρ)
    converged = false

    # perform iterations
    while !converged && n_iter < maxiter
        n_iter += 1

        # compute Newton step and next iteration
        res = compute_projected_gradient(basis, ψ, occupation)
        # solve (Ω+K) δψ = -res so that the Newton step is ψ <- ψ + δψ
        δψ = solve_ΩplusK(basis, ψ, -res, occupation; tol_cg, verbose).δψ
        ψ  = [ortho_qr(ψ[ik] + δψ[ik]) for ik in 1:Nk]

        ρ_next = compute_density(basis, ψ, occupation)
        energies, H = energy_hamiltonian(basis, ψ, occupation; ρ=ρ_next)
        info = (ham=H, basis=basis, converged=converged, stage=:iterate,
                ρin=ρ, ρout=ρ_next, n_iter=n_iter, energies=energies, algorithm="Newton")
        callback(info)

        # update and test convergence
        converged = is_converged(info)
        ρ = ρ_next
    end

    # Rayleigh-Ritz
    eigenvalues = []
    for ik = 1:Nk
        Hψk = H.blocks[ik] * ψ[ik]
        F = eigen(Hermitian(ψ[ik]'Hψk))
        push!(eigenvalues, F.values)
        ψ[ik] .= ψ[ik] * F.vectors
    end

    εF = nothing  # does not necessarily make sense here, as the
                  # Aufbau property might not even be true

    # return results and call callback one last time with final state for clean
    # up
    info = (ham=H, basis=basis, energies=energies, converged=converged,
            ρ=ρ, eigenvalues=eigenvalues, occupation=occupation, εF=εF,
            n_iter=n_iter, ψ=ψ, stage=:finalize, algorithm="Newton")
    callback(info)
    info
end


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

mutable struct IoK8sApiCoreV1DownwardAPIVolumeSource <: SwaggerModel
    defaultMode::Any # spec type: Union{ Nothing, Int32 } # spec name: defaultMode
    items::Any # spec type: Union{ Nothing, Vector{IoK8sApiCoreV1DownwardAPIVolumeFile} } # spec name: items

    function IoK8sApiCoreV1DownwardAPIVolumeSource(;defaultMode=nothing, items=nothing)
        o = new()
        validate_property(IoK8sApiCoreV1DownwardAPIVolumeSource, Symbol("defaultMode"), defaultMode)
        setfield!(o, Symbol("defaultMode"), defaultMode)
        validate_property(IoK8sApiCoreV1DownwardAPIVolumeSource, Symbol("items"), items)
        setfield!(o, Symbol("items"), items)
        o
    end
end # type IoK8sApiCoreV1DownwardAPIVolumeSource

const _property_map_IoK8sApiCoreV1DownwardAPIVolumeSource = Dict{Symbol,Symbol}(Symbol("defaultMode")=>Symbol("defaultMode"), Symbol("items")=>Symbol("items"))
const _property_types_IoK8sApiCoreV1DownwardAPIVolumeSource = Dict{Symbol,String}(Symbol("defaultMode")=>"Int32", Symbol("items")=>"Vector{IoK8sApiCoreV1DownwardAPIVolumeFile}")
Base.propertynames(::Type{ IoK8sApiCoreV1DownwardAPIVolumeSource }) = collect(keys(_property_map_IoK8sApiCoreV1DownwardAPIVolumeSource))
Swagger.property_type(::Type{ IoK8sApiCoreV1DownwardAPIVolumeSource }, name::Symbol) = Union{Nothing,eval(Meta.parse(_property_types_IoK8sApiCoreV1DownwardAPIVolumeSource[name]))}
Swagger.field_name(::Type{ IoK8sApiCoreV1DownwardAPIVolumeSource }, property_name::Symbol) =  _property_map_IoK8sApiCoreV1DownwardAPIVolumeSource[property_name]

function check_required(o::IoK8sApiCoreV1DownwardAPIVolumeSource)
    true
end

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


text:
using JLD2, FileIO
using Test

struct Anint; x::Int; end
struct Bparam{T}; x::T; end

@testset "Inline Union Representation" begin
    fn = mktemp()[1]*".jld2"
    u = Union{Int,Float64}[1, 2.5, 3, 4.5]
    save(fn, "u", u)
    @test all(skipmissing(u) .== skipmissing(load(fn, "u")))

    jldopen(fn, "r") do f
        @test u == f["u"]
        # Test again to verify that the correct object was cached
        # See issue #348
        @test u == f["u"]
    end

    u = Union{Float64, Missing}[1.0, missing, 2.0, 10.5, missing]
    save(fn, "u", u; compress=true)
    @test all(skipmissing(u) .== skipmissing(load(fn, "u")))

    u = Union{Anint, Missing}[Anint(10), missing, missing, Anint(20)]
    save(fn, "u", u)
    @test all(skipmissing(u) .== skipmissing(load(fn, "u")))

    u = Union{Anint, Float64, Int}[Anint(1), 2.5, 3]
    save(fn, "u", u)
    @test u == load(fn, "u")

    u = Union{Anint,Bparam}[Anint(1), Bparam(3)]
    save(fn, "u", u)
    @test u == load(fn, "u")

    u = Bparam[Bparam(1), Bparam(2.0)]
    save(fn, "u", u)
    @test u == load(fn, "u")

    u = Union{Float32, Missing}[rand(5,5);]
    save(fn, "u", u)
    @test u == load(fn, "u")
end


text:
# JDNOTE: This might be merged with the structs in simulation_store
mutable struct SimulationInfo
    number::Int
    name::String
    executions::Int
    execution_count::Int
    caches::Set{CacheKey}
    end_of_interval_step::Int
    chronolgy_dict::Dict{Int, <:FeedForwardChronology}
    requires_rebuild::Bool
    sequence_uuid::Base.UUID
end

struct TimeSeriesCacheKey
    component_uuid::Base.UUID
    time_series_type::Type{<:IS.TimeSeriesData}
    name::String
end

mutable struct ProblemInternal
    optimization_container::OptimizationContainer
    status::BuildStatus
    run_status::RunStatus
    base_conversion::Bool
    output_dir::Union{Nothing, String}
    simulation_info::Union{Nothing, SimulationInfo}
    time_series_cache::Dict{TimeSeriesCacheKey, <:IS.TimeSeriesCache}
    ext::Dict{String, Any}
    console_level::Base.CoreLogging.LogLevel
    file_level::Base.CoreLogging.LogLevel
end

function ProblemInternal(
    optimization_container::OptimizationContainer;
    ext = Dict{String, Any}(),
)
    return ProblemInternal(
        optimization_container,
        BuildStatus.EMPTY,
        RunStatus.READY,
        true,
        nothing,
        nothing,
        Dict{TimeSeriesCacheKey, IS.TimeSeriesCache}(),
        ext,
        Logging.Warn,
        Logging.Info,
    )
end

function configure_logging(internal::ProblemInternal, file_mode)
    return IS.configure_logging(
        console = true,
        console_stream = stderr,
        console_level = internal.console_level,
        file = true,
        filename = joinpath(internal.output_dir, PROBLEM_BUILD_LOG_FILENAME),
        file_level = internal.file_level,
        file_mode = file_mode,
        tracker = nothing,
        set_global = false,
    )
end

"""Default PowerSimulations Operation Problem Type"""
struct GenericOpProblem <: PowerSimulationsOperationsProblem end

"""
    OperationsProblem(::Type{M},
    template::OperationsProblemTemplate,
    sys::PSY.System,
    jump_model::Union{Nothing, JuMP.Model}=nothing;
    kwargs...) where {M<:AbstractOperationsProblem,
                      T<:PM.AbstractPowerFormulation}

This builds the optimization problem of type M with the specific system and template.

# Arguments

- `::Type{M} where M<:AbstractOperationsProblem`: The abstract operation model type
- `template::OperationsProblemTemplate`: The model reference made up of transmission, devices,
                                          branches, and services.
- `sys::PSY.System`: the system created using Power Systems
- `jump_model::Union{Nothing, JuMP.Model}`: Enables passing a custom JuMP model. Use with care

# Output

- `problem::OperationsProblem`: The operation model containing the model type, built JuMP model, Power
Systems system.

# Example

```julia
template = OperationsProblemTemplate(CopperPlatePowerModel, devices, branches, services)
OpModel = OperationsProblem(MockOperationProblem, template, system)
```

# Accepted Key Words

- `optimizer`: The optimizer that will be used in the optimization model.
- `PTDF::PTDF`: Passes the PTDF matrix into the optimization model for StandardPTDFModel networks.
- `horizon::Int`: Manually specify the length of the forecast Horizon
- `warm_start::Bool`: True will use the current operation point in the system to initialize variable values. False initializes all variables to zero. Default is true
- `balance_slack_variables::Bool`: True will add slacks to the system balance constraints
- `services_slack_variables::Bool`: True will add slacks to the services requirement constraints
- `constraint_duals::Vector{Symbol}`: Vector with the duals to query from the optimization model
- `system_to_file::Bool:`: True to create a copy of the system used in the model. Default true.
- `export_pwl_vars::Bool`: True to export all the pwl intermediate variables. It can slow down significantly the solve time. Default to false.
- `allow_fails::Bool`: True to allow the simulation to continue even if the optimization step fails. Use with care, default to false.
- `optimizer_log_print::Bool`: True to print the optimizer solve log. Default to false.
- `direct_mode_optimizer::Bool` True to use the solver in direct mode. Creates a [JuMP.direct_model](https://jump.dev/JuMP.jl/dev/reference/models/#JuMP.direct_model). Default to false.
- `use_parameters::Bool`: True will substitute will implement formulations using ParameterJuMP parameters. Default is false.
- `use_forecast_data::Bool`: If true uses the data in the system forecasts. If false uses the data for current operating point in the system.
- `initial_time::Dates.DateTime`: Initial Time for the model solve
- `time_series_cache_size::Int`: Size in bytes to cache for each time array. Default is 1 MiB. Set to 0 to disable.

"""
mutable struct OperationsProblem{M <: AbstractOperationsProblem}
    template::OperationsProblemTemplate
    sys::PSY.System
    internal::Union{Nothing, ProblemInternal}
    ext::Dict{String, Any}

    function OperationsProblem{M}(
        template::OperationsProblemTemplate,
        sys::PSY.System,
        settings::Settings,
        jump_model::Union{Nothing, JuMP.Model} = nothing,
    ) where {M <: AbstractOperationsProblem}
        internal = ProblemInternal(OptimizationContainer(sys, settings, jump_model))
        new{M}(template, sys, internal, Dict{String, Any}())
    end
end

function OperationsProblem{M}(
    template::OperationsProblemTemplate,
    sys::PSY.System,
    jump_model::Union{Nothing, JuMP.Model} = nothing;
    optimizer = nothing,
    PTDF = nothing,
    horizon = UNSET_HORIZON,
    warm_start = true,
    balance_slack_variables = false,
    services_slack_variables = false,
    constraint_duals = Vector{Symbol}(),
    system_to_file = true,
    export_pwl_vars = false,
    allow_fails = false,
    optimizer_log_print = false,
    direct_mode_optimizer = false,
    use_parameters = false,
    use_forecast_data = true,
    initial_time = UNSET_INI_TIME,
    time_series_cache_size::Int = IS.TIME_SERIES_CACHE_SIZE_BYTES,
) where {M <: AbstractOperationsProblem}
    settings = Settings(
        sys;
        horizon = horizon,
        initial_time = initial_time,
        optimizer = optimizer,
        use_parameters = use_parameters,
        time_series_cache_size = time_series_cache_size,
        warm_start = warm_start,
        balance_slack_variables = balance_slack_variables,
        services_slack_variables = services_slack_variables,
        constraint_duals = constraint_duals,
        system_to_file = system_to_file,
        export_pwl_vars = export_pwl_vars,
        allow_fails = allow_fails,
        PTDF = PTDF,
        optimizer_log_print = optimizer_log_print,
        direct_mode_optimizer = direct_mode_optimizer,
        use_forecast_data = use_forecast_data,
    )
    return OperationsProblem{M}(template, sys, settings, jump_model)
end

"""
    OperationsProblem(::Type{M},
    template::OperationsProblemTemplate,
    sys::PSY.System,
    optimizer::MOI.OptimizerWithAttributes,
    jump_model::Union{Nothing, JuMP.Model}=nothing;
    kwargs...) where {M<:AbstractOperationsProblem}
This builds the optimization problem of type M with the specific system and template
# Arguments
- `::Type{M} where M<:AbstractOperationsProblem`: The abstract operation model type
- `template::OperationsProblemTemplate`: The model reference made up of transmission, devices,
                                          branches, and services.
- `sys::PSY.System`: the system created using Power Systems
- `jump_model::Union{Nothing, JuMP.Model}`: Enables passing a custom JuMP model. Use with care
# Output
- `Stage::OperationsProblem`: The operation model containing the model type, unbuilt JuMP model, Power
Systems system.
# Example
```julia
template = OperationsProblemTemplate(CopperPlatePowerModel, devices, branches, services)
problem = OperationsProblem(MyOpProblemType template, system, optimizer)
```
# Accepted Key Words
- `initial_time::Dates.DateTime`: Initial Time for the model solve
- `PTDF::PTDF`: Passes the PTDF matrix into the optimization model for StandardPTDFModel networks.
- `warm_start::Bool` True will use the current operation point in the system to initialize variable values. False initializes all variables to zero. Default is true
- `balance_slack_variables::Bool` True will add slacks to the system balance constraints
- `services_slack_variables::Bool` True will add slacks to the services requirement constraints
- `export_pwl_vars::Bool` True will write the results of the piece-wise-linear intermediate variables. Slows down the simulation process significantly
- `allow_fails::Bool` True will allow the simulation to continue if the optimizer can't find a solution. Use with care, can lead to unwanted behaviour or results
- `optimizer_log_print::Bool` Uses JuMP.unset_silent() to print the optimizer's log. By default all solvers are set to `MOI.Silent()`
"""
function OperationsProblem(
    ::Type{M},
    template::OperationsProblemTemplate,
    sys::PSY.System,
    jump_model::Union{Nothing, JuMP.Model} = nothing;
    kwargs...,
) where {M <: AbstractOperationsProblem}
    return OperationsProblem{M}(template, sys, jump_model; kwargs...)
end

function OperationsProblem(
    template::OperationsProblemTemplate,
    sys::PSY.System,
    jump_model::Union{Nothing, JuMP.Model} = nothing;
    kwargs...,
)
    return OperationsProblem{GenericOpProblem}(template, sys, jump_model; kwargs...)
end

"""
    OperationsProblem(filename::AbstractString)

Construct an OperationsProblem from a serialized file.

# Arguments
- `filename::AbstractString`: path to serialized file
- `jump_model::Union{Nothing, JuMP.Model}` = nothing: The JuMP model does not get
   serialized. Callers should pass whatever they passed to the original problem.
- `optimizer::Union{Nothing,MOI.OptimizerWithAttributes}` = nothing: The optimizer does
   not get serialized. Callers should pass whatever they passed to the original problem.
"""
function OperationsProblem(
    filename::AbstractString;
    jump_model::Union{Nothing, JuMP.Model} = nothing,
    optimizer::Union{Nothing, MOI.OptimizerWithAttributes} = nothing,
    kwargs...,
)
    return deserialize_problem(
        OperationsProblem,
        filename;
        jump_model = jump_model,
        optimizer = optimizer,
        kwargs...,
    )
end

# Default implementations of getter/setter functions for OperationsProblem.
is_built(problem::OperationsProblem) = problem.internal.status == BuildStatus.BUILT
is_empty(problem::OperationsProblem) = problem.internal.status == BuildStatus.EMPTY
warm_start_enabled(problem::OperationsProblem) =
    get_warm_start(get_optimization_container(problem).settings)
built_for_simulation(problem::OperationsProblem) = get_simulation_info(problem) !== nothing
get_caches(x::OperationsProblem) =
    built_for_simulation(x) ? get_simulation_info(x).caches : nothing
get_constraints(problem::OperationsProblem) =
    get_internal(problem).optimization_container.constraints
get_end_of_interval_step(problem::OperationsProblem) =
    get_simulation_info(problem).end_of_interval_step
get_execution_count(problem::OperationsProblem) =
    get_simulation_info(problem).execution_count
get_executions(problem::OperationsProblem) = get_simulation_info(problem).executions
get_initial_time(problem::OperationsProblem) = get_initial_time(get_settings(problem))
get_horizon(problem::OperationsProblem) = get_horizon(get_settings(problem))
get_internal(problem::OperationsProblem) = problem.internal
get_jump_model(problem::OperationsProblem) =
    get_internal(problem).optimization_container.JuMPmodel
get_name(x::OperationsProblem) = built_for_simulation(x) ? get_simulation_info(x).name : ""

get_optimization_container(problem::OperationsProblem) =
    problem.internal.optimization_container
function get_resolution(problem::OperationsProblem{<:AbstractOperationsProblem})
    resolution = PSY.get_time_series_resolution(get_system(problem))
    return IS.time_period_conversion(resolution)
end
get_problem_base_power(problem::OperationsProblem) = PSY.get_base_power(problem.sys)
get_settings(problem::OperationsProblem) = get_optimization_container(problem).settings
get_solve_timed_log(problem::OperationsProblem) =
    get_optimization_container(problem).solve_timed_log
get_simulation_info(problem::OperationsProblem) = problem.internal.simulation_info
get_simulation_number(problem::OperationsProblem) = problem.internal.simulation_info.number
get_status(problem::OperationsProblem) = problem.internal.status
get_system(problem::OperationsProblem) = problem.sys
get_template(problem::OperationsProblem) = problem.template
get_output_dir(problem::OperationsProblem) = problem.internal.output_dir
get_variables(problem::OperationsProblem) = get_optimization_container(problem).variables
get_parameters(problem::OperationsProblem) = get_optimization_container(problem).parameters
get_duals(problem::OperationsProblem) = get_optimization_container(problem).duals

get_run_status(problem::OperationsProblem) = problem.internal.run_status
set_run_status!(problem::OperationsProblem, status) = problem.internal.run_status = status
get_time_series_cache(problem::OperationsProblem) = problem.internal.time_series_cache
empty_time_series_cache!(x::OperationsProblem) = empty!(get_time_series_cache(x))

function get_time_series_values!(
    time_series_type::Type{<:IS.TimeSeriesData},
    problem::OperationsProblem,
    component,
    name,
    initial_time,
    horizon;
    ignore_scaling_factors = true,
)
    if !use_time_series_cache(get_settings(problem))
        return IS.get_time_series_values(
            time_series_type,
            component,
            name,
            start_time = initial_time,
            len = horizon,
            ignore_scaling_factors = true,
        )
    end

    cache = get_time_series_cache(problem)
    key = TimeSeriesCacheKey(IS.get_uuid(component), time_series_type, name)
    if haskey(cache, key)
        ts_cache = cache[key]
    else
        ts_cache = make_time_series_cache(
            time_series_type,
            component,
            name,
            initial_time,
            horizon,
            ignore_scaling_factors = true,
        )
        cache[key] = ts_cache
    end

    ts = IS.get_time_series_array!(ts_cache, initial_time)
    return TimeSeries.values(ts)
end

function make_time_series_cache(
    time_series_type::Type{T},
    component,
    name,
    initial_time,
    horizon;
    ignore_scaling_factors = true,
) where {T <: IS.TimeSeriesData}
    key = TimeSeriesCacheKey(IS.get_uuid(component), T, name)
    if T <: IS.SingleTimeSeries
        cache = IS.StaticTimeSeriesCache(
            PSY.SingleTimeSeries,
            component,
            name,
            start_time = initial_time,
            ignore_scaling_factors = ignore_scaling_factors,
        )
    elseif T <: IS.Deterministic
        cache = IS.ForecastCache(
            IS.AbstractDeterministic,
            component,
            name,
            start_time = initial_time,
            horizon = horizon,
            ignore_scaling_factors = ignore_scaling_factors,
        )
    elseif T <: IS.Probabilistic
        cache = IS.ForecastCache(
            IS.Probabilistic,
            component,
            name,
            start_time = initial_time,
            horizon = horizon,
            ignore_scaling_factors = ignore_scaling_factors,
        )
    else
        error("not supported yet: $T")
    end

    @debug "Made time series cache for $(summary(component))" name initial_time
    return cache
end

function get_initial_conditions(
    problem::OperationsProblem,
    ic::InitialConditionType,
    device::PSY.Device,
)
    key = ICKey(ic, device)
    return get_initial_conditions(get_optimization_container(problem), key)
end

set_console_level!(problem::OperationsProblem, val) =
    get_internal(problem).console_level = val
set_file_level!(problem::OperationsProblem, val) = get_internal(problem).file_level = val
set_executions!(problem::OperationsProblem, val::Int) =
    problem.internal.simulation_info.executions = val
set_execution_count!(problem::OperationsProblem, val::Int) =
    get_simulation_info(problem).execution_count = val
set_initial_time!(problem::OperationsProblem, val::Dates.DateTime) =
    set_initial_time!(get_settings(problem), val)
set_simulation_info!(problem::OperationsProblem, info::SimulationInfo) =
    problem.internal.simulation_info = info
function set_status!(problem::OperationsProblem, status::BuildStatus)
    problem.internal.status = status
    return
end
set_output_dir!(problem::OperationsProblem, path::AbstractString) =
    get_internal(problem).output_dir = path

function reset!(problem::OperationsProblem{T}) where {T <: AbstractOperationsProblem}
    if built_for_simulation(problem::OperationsProblem)
        set_execution_count!(problem, 0)
    end
    container = OptimizationContainer(get_system(problem), get_settings(problem), nothing)
    problem.internal.optimization_container = container
    empty_time_series_cache!(problem)
    set_status!(problem, BuildStatus.EMPTY)
    return
end

function advance_execution_count!(problem::OperationsProblem)
    info = get_simulation_info(problem)
    info.execution_count += 1
    # Reset execution count at the end of step
    if get_execution_count(problem) == get_executions(problem)
        info.execution_count = 0
    end
    return
end

function build_pre_step!(problem::OperationsProblem)
    TimerOutputs.@timeit BUILD_PROBLEMS_TIMER "Build pre-step" begin
        if !is_empty(problem)
            @info "OptimizationProblem status not BuildStatus.EMPTY. Resetting"
            reset!(problem)
        end
        settings = get_settings(problem)
        if !get_use_parameters(settings)
            set_use_parameters!(settings, built_for_simulation(problem))
            @debug "Set use_parameters = true for use in simulation"
        end
        # Initial time are set here because the information is specified in the
        # Simulation Sequence object and not at the problem creation.
        @info "Initializing Optimization Container"
        optimization_container_init!(
            get_optimization_container(problem),
            get_transmission_model(get_template(problem)),
            get_system(problem),
        )
        set_status!(problem, BuildStatus.IN_PROGRESS)
    end
    return
end

function initialize_simulation_info!(problem::OperationsProblem, ::FeedForwardChronology)
    @assert built_for_simulation(problem)
    system = get_system(problem)
    resolution = get_resolution(problem)
    interval = IS.time_period_conversion(PSY.get_forecast_interval(system))
    end_of_interval_step = Int(interval / resolution)
    get_simulation_info(problem).end_of_interval_step = end_of_interval_step
end

function initialize_simulation_info!(problem::OperationsProblem, ::RecedingHorizon)
    @assert built_for_simulation(problem)
    get_simulation_info(problem).end_of_interval_step = 1
end

function _build!(problem::OperationsProblem{<:AbstractOperationsProblem}, serialize::Bool)
    TimerOutputs.@timeit BUILD_PROBLEMS_TIMER "Problem $(get_name(problem))" begin
        try
            build_pre_step!(problem)
            problem_build!(problem)
            serialize && serialize_problem(problem)
            serialize && serialize_optimization_model(problem)
            set_status!(problem, BuildStatus.BUILT)
            log_values(get_settings(problem))
            !built_for_simulation(problem) && @info "\n$(BUILD_PROBLEMS_TIMER)\n"
        catch e
            set_status!(problem, BuildStatus.FAILED)
            bt = catch_backtrace()
            @error "Operation Problem Build Failed" exception = e, bt
        end
    end
    return get_status(problem)
end

"""Implementation of build for any OperationsProblem"""
function build!(
    problem::OperationsProblem{<:AbstractOperationsProblem};
    output_dir::String,
    console_level = Logging.Error,
    file_level = Logging.Info,
    disable_timer_outputs = false,
    serialize = true,
)
    mkpath(output_dir)
    set_output_dir!(problem, output_dir)
    set_console_level!(problem, console_level)
    set_file_level!(problem, file_level)
    TimerOutputs.reset_timer!(BUILD_PROBLEMS_TIMER)
    disable_timer_outputs && TimerOutputs.disable_timer!(BUILD_PROBLEMS_TIMER)
    logger = configure_logging(problem.internal, "w")
    try
        Logging.with_logger(logger) do
            return _build!(problem, serialize)
        end
    finally
        close(logger)
    end
end

"""
Default implementation of build method for Operational Problems for models conforming with PowerSimulationsOperationsProblem specification. Overload this function to implement a custom build method
"""
function problem_build!(problem::OperationsProblem{<:PowerSimulationsOperationsProblem})
    build_impl!(
        get_optimization_container(problem),
        get_template(problem),
        get_system(problem),
    )
end

serialize_optimization_model(::OperationsProblem) = nothing
serialize_problem(::OperationsProblem) = nothing

function serialize_optimization_model(
    op_problem::OperationsProblem{<:PowerSimulationsOperationsProblem},
)
    name = get_name(op_problem)
    problem_name = isempty(name) ? "OptimizationModel" : "$(name)_OptimizationModel"
    json_file_name = "$problem_name.json"
    json_file_name = joinpath(get_output_dir(op_problem), json_file_name)
    serialize_optimization_model(get_optimization_container(op_problem), json_file_name)
end

struct OperationsProblemSerializationWrapper
    template::OperationsProblemTemplate
    sys::Union{Nothing, String}
    settings::Settings
    op_problem_type::DataType
end

function serialize_problem(
    op_problem::OperationsProblem{<:PowerSimulationsOperationsProblem},
)
    # A PowerSystem cannot be serialized in this format because of how it stores
    # time series data. Use its specialized serialization method instead.
    problem_name = isempty(get_name(op_problem)) ? "OperationProblem" : get_name(op_problem)
    sys_to_file = get_system_to_file(get_settings(op_problem))
    if sys_to_file
        sys = get_system(op_problem)
        sys_filename = joinpath(get_output_dir(op_problem), make_system_filename(sys))
        # Skip serialization if the system is already in the folder
        !ispath(sys_filename) && PSY.to_json(sys, sys_filename)
    else
        sys_filename = nothing
    end
    optimization_container = get_optimization_container(op_problem)
    obj = OperationsProblemSerializationWrapper(
        op_problem.template,
        sys_filename,
        optimization_container.settings_copy,
        typeof(op_problem),
    )
    bin_file_name = "$problem_name.bin"
    bin_file_name = joinpath(get_output_dir(op_problem), bin_file_name)
    Serialization.serialize(bin_file_name, obj)
    @info "Serialized OperationsProblem to" bin_file_name
end

function deserialize_problem(::Type{OperationsProblem}, filename::AbstractString; kwargs...)
    obj = Serialization.deserialize(filename)
    if !(obj isa OperationsProblemSerializationWrapper)
        throw(IS.DataFormatError("deserialized object has incorrect type $(typeof(obj))"))
    end
    sys = get(kwargs, :system, nothing)
    settings = restore_from_copy(obj.settings; optimizer = kwargs[:optimizer])
    if sys === nothing
        if obj.sys === nothing && !settings[:sys_to_file]
            throw(
                IS.DataFormatError(
                    "Operations Problem System was not serialized and a System has not been specified.",
                ),
            )
        elseif !ispath(obj.sys)
            throw(IS.DataFormatError("PowerSystems.System file $(obj.sys) does not exist"))
        end
        sys = PSY.System(obj.sys)
    end

    return obj.op_problem_type(obj.template, sys, kwargs[:jump_model]; settings...)
end

function calculate_aux_variables!(problem::OperationsProblem)
    optimization_container = get_optimization_container(problem)
    system = get_system(problem)
    aux_vars = get_aux_variables(optimization_container)
    for key in keys(aux_vars)
        calculate_aux_variable_value!(optimization_container, key, system)
    end
    return
end

function solve_impl(problem::OperationsProblem; optimizer = nothing)
    if !is_built(problem)
        error(
            "Operations Problem Build status is $(get_status(problem)). Solve can't continue",
        )
    end
    model = get_jump_model(problem)
    if optimizer !== nothing
        JuMP.set_optimizer(model, optimizer)
    end
    if model.moi_backend.state == MOIU.NO_OPTIMIZER
        @error("No Optimizer has been defined, can't solve the operational problem")
        return RunStatus.FAILED
    end
    @assert model.moi_backend.state != MOIU.NO_OPTIMIZER
    status = RunStatus.RUNNING
    timed_log = get_solve_timed_log(problem)
    _, timed_log[:timed_solve_time], timed_log[:solve_bytes_alloc], timed_log[:sec_in_gc] =
        @timed JuMP.optimize!(model)
    model_status = JuMP.primal_status(model)
    if model_status != MOI.FEASIBLE_POINT::MOI.ResultStatusCode
        return RunStatus.FAILED
    else
        calculate_aux_variables!(problem::OperationsProblem)
        status = RunStatus.SUCCESSFUL
    end
    return status
end

"""
Default solve method the operational model for a single instance. Solves problems
    that conform to the requirements of OperationsProblem{<: PowerSimulationsOperationsProblem}
# Arguments
- `op_problem::OperationModel = op_problem`: operation model
# Examples
```julia
results = solve!(OpModel)
```
# Accepted Key Words
- `output_dir::String`: If a file path is provided the results
automatically get written to feather files
- `optimizer::MOI.OptimizerWithAttributes`: The optimizer that is used to solve the model
"""
function solve!(problem::OperationsProblem{<:PowerSimulationsOperationsProblem}; kwargs...)
    status = solve_impl(problem; kwargs...)
    set_run_status!(problem, status)
    return status
end

function write_problem_results!(
    step::Int,
    problem::OperationsProblem{<:PowerSimulationsOperationsProblem},
    start_time::Dates.DateTime,
    store::SimulationStore,
    exports,
)
    stats = OptimizerStats(problem, step)
    write_optimizer_stats!(store, get_name(problem), stats, start_time)
    write_model_results!(store, problem, start_time; exports = exports)
    return
end

function write_problem_results!(
    ::Int,
    ::OperationsProblem{T},
    ::Dates.DateTime,
    ::SimulationStore,
    _,
) where {T <: AbstractOperationsProblem}
    @info "Write results to Store not implemented for problems $T"
    return
end

"""
Default solve method for an operational model used inside of a Simulation. Solves problems that conform to the requirements of OperationsProblem{<: PowerSimulationsOperationsProblem}

# Arguments
- `step::Int`: Simulation Step
- `op_problem::OperationModel`: operation model
- `start_time::Dates.DateTime`: Initial Time of the simulation step in Simulation time.
- `store::SimulationStore`: Simulation output store

# Accepted Key Words
- `exports`: realtime export of output. Use wisely, it can have negative impacts in the simulation times
"""
function solve!(
    step::Int,
    problem::OperationsProblem{<:AbstractOperationsProblem},
    start_time::Dates.DateTime,
    store::SimulationStore;
    exports = nothing,
)
    solve_status = solve!(problem)
    if solve_status == RunStatus.SUCCESSFUL
        write_problem_results!(step, problem, start_time, store, exports)
        advance_execution_count!(problem)
    end

    return solve_status
end

function write_model_results!(store, problem, timestamp; exports = nothing)
    optimization_container = get_optimization_container(problem)
    if exports !== nothing
        export_params = Dict{Symbol, Any}(
            :exports => exports,
            :exports_path => joinpath(exports.path, get_name(problem)),
            :file_type => get_export_file_type(exports),
            :resolution => get_resolution(problem),
            :horizon => get_horizon(get_settings(problem)),
        )
    else
        export_params = nothing
    end

    # This line should only be called if the problem is exporting duals. Otherwise ignore.
    if is_milp(get_optimization_container(problem))
        @warn "Problem $(get_simulation_info(problem).name) is a MILP, duals can't be exported"
    else
        _write_model_dual_results!(
            store,
            optimization_container,
            problem,
            timestamp,
            export_params,
        )
    end

    _write_model_parameter_results!(
        store,
        optimization_container,
        problem,
        timestamp,
        export_params,
    )
    _write_model_variable_results!(
        store,
        optimization_container,
        problem,
        timestamp,
        export_params,
    )
    _write_model_aux_variable_results!(
        store,
        optimization_container,
        problem,
        timestamp,
        export_params,
    )
    return
end

function _write_model_dual_results!(
    store,
    optimization_container,
    problem,
    timestamp,
    exports,
)
    problem_name_str = get_name(problem)
    problem_name = Symbol(problem_name_str)
    if exports !== nothing
        exports_path = joinpath(exports[:exports_path], "duals")
        mkpath(exports_path)
    end

    for name in get_constraint_duals(optimization_container.settings)
        constraint = get_constraint(optimization_container, name)
        write_result!(
            store,
            problem_name,
            STORE_CONTAINER_DUALS,
            name,
            timestamp,
            constraint,
            [name],
        )

        if exports !== nothing &&
           should_export_dual(exports[:exports], timestamp, problem_name_str, name)
            horizon = exports[:horizon]
            resolution = exports[:resolution]
            file_type = exports[:file_type]
            df = axis_array_to_dataframe(constraint, [name])
            time_col = range(timestamp, length = horizon, step = resolution)
            DataFrames.insertcols!(df, 1, :DateTime => time_col)
            export_result(file_type, exports_path, name, timestamp, df)
        end
    end
end

function _write_model_parameter_results!(
    store,
    optimization_container,
    problem,
    timestamp,
    exports,
)
    problem_name_str = get_name(problem)
    problem_name = Symbol(problem_name_str)
    if exports !== nothing
        exports_path = joinpath(exports[:exports_path], "parameters")
        mkpath(exports_path)
    end

    parameters = get_parameters(optimization_container)
    (isnothing(parameters) || isempty(parameters)) && return
    horizon = get_horizon(get_settings(problem))

    for (name, container) in parameters
        !isa(container.update_ref, UpdateRef{<:PSY.Component}) && continue
        param_array = get_parameter_array(container)
        multiplier_array = get_multiplier_array(container)
        @assert_op length(axes(param_array)) == 2
        num_columns = size(param_array)[1]
        data = Array{Float64}(undef, horizon, num_columns)
        for r_ix in param_array.axes[2], (c_ix, name) in enumerate(param_array.axes[1])
            val1 = _jump_value(param_array[name, r_ix])
            val2 = multiplier_array[name, r_ix]
            data[r_ix, c_ix] = val1 * val2
        end

        write_result!(
            store,
            problem_name,
            STORE_CONTAINER_PARAMETERS,
            name,
            timestamp,
            data,
            param_array.axes[1],
        )

        if exports !== nothing &&
           should_export_parameter(exports[:exports], timestamp, problem_name_str, name)
            resolution = exports[:resolution]
            file_type = exports[:file_type]
            df = DataFrames.DataFrame(data, param_array.axes[1])
            time_col = range(timestamp, length = horizon, step = resolution)
            DataFrames.insertcols!(df, 1, :DateTime => time_col)
            export_result(file_type, exports_path, name, timestamp, df)
        end
    end
end

function _write_model_variable_results!(
    store,
    optimization_container,
    problem,
    timestamp,
    exports,
)
    problem_name_str = get_name(problem)
    problem_name = Symbol(problem_name_str)
    if exports !== nothing
        exports_path = joinpath(exports[:exports_path], "variables")
        mkpath(exports_path)
    end

    for (name, variable) in get_variables(optimization_container)
        write_result!(
            store,
            problem_name,
            STORE_CONTAINER_VARIABLES,
            name,
            timestamp,
            variable,
        )

        if exports !== nothing &&
           should_export_variable(exports[:exports], timestamp, problem_name_str, name)
            horizon = exports[:horizon]
            resolution = exports[:resolution]
            file_type = exports[:file_type]
            df = axis_array_to_dataframe(variable)
            time_col = range(timestamp, length = horizon, step = resolution)
            DataFrames.insertcols!(df, 1, :DateTime => time_col)
            export_result(file_type, exports_path, name, timestamp, df)
        end
    end
end

function _write_model_aux_variable_results!(
    store,
    optimization_container,
    problem,
    timestamp,
    exports,
)
    problem_name_str = get_name(problem)
    problem_name = Symbol(problem_name_str)
    if exports !== nothing
        # TODO: Should the export go to a folder aux_variables
        exports_path = joinpath(exports[:exports_path], "variables")
        mkpath(exports_path)
    end

    for (key, variable) in get_aux_variables(optimization_container)
        name = encode_key(key)
        write_result!(
            store,
            problem_name,
            STORE_CONTAINER_VARIABLES,
            name,
            timestamp,
            variable,
        )

        if exports !== nothing &&
           should_export_variable(exports[:exports], timestamp, problem_name_str, name)
            horizon = exports[:horizon]
            resolution = exports[:resolution]
            file_type = exports[:file_type]
            df = axis_array_to_dataframe(variable)
            time_col = range(timestamp, length = horizon, step = resolution)
            DataFrames.insertcols!(df, 1, :DateTime => time_col)
            export_result(file_type, exports_path, name, timestamp, df)
        end
    end
end

# Here because requires the problem to be defined
# This is a method a user defining a custom cache will have to define. This is the definition
# in PSI for the building the TimeStatusChange
function get_initial_cache(cache::AbstractCache, problem::OperationsProblem)
    throw(ArgumentError("Initialization method for cache $(typeof(cache)) not defined"))
end

function get_initial_cache(cache::TimeStatusChange, problem::OperationsProblem)
    ini_cond_on = get_initial_conditions(
        get_optimization_container(problem),
        InitialTimeDurationOn,
        cache.device_type,
    )

    ini_cond_off = get_initial_conditions(
        get_optimization_container(problem),
        InitialTimeDurationOff,
        cache.device_type,
    )

    device_axes = Set((
        PSY.get_name(ic.device) for ic in Iterators.Flatten([ini_cond_on, ini_cond_off])
    ),)
    value_array = JuMP.Containers.DenseAxisArray{Dict{Symbol, Any}}(undef, device_axes)

    for ic in ini_cond_on
        device_name = PSY.get_name(ic.device)
        condition = get_condition(ic)
        status = (condition > 0.0) ? 1.0 : 0.0
        value_array[device_name] = Dict(:count => condition, :status => status)
    end

    for ic in ini_cond_off
        device_name = PSY.get_name(ic.device)
        condition = get_condition(ic)
        status = (condition > 0.0) ? 0.0 : 1.0
        if value_array[device_name][:status] != status
            throw(
                IS.ConflictingInputsError(
                    "Initial Conditions for $(device_name) are not compatible. The values provided are invalid",
                ),
            )
        end
    end

    return value_array
end

function get_initial_cache(cache::StoredEnergy, problem::OperationsProblem)
    ini_cond_level = get_initial_conditions(
        get_optimization_container(problem),
        InitialEnergyLevel,
        cache.device_type,
    )

    device_axes = Set([PSY.get_name(ic.device) for ic in ini_cond_level],)
    value_array = JuMP.Containers.DenseAxisArray{Float64}(undef, device_axes)
    for ic in ini_cond_level
        device_name = PSY.get_name(ic.device)
        condition = get_condition(ic)
        value_array[device_name] = condition
    end
    return value_array
end

function get_timestamps(problem::OperationsProblem)
    start_time = model_initial_time(get_optimization_container(problem))
    resolution = get_resolution(problem)
    horizon = get_horizon(problem)
    return range(start_time, length = horizon, step = resolution)
end

function write_data(problem::OperationsProblem, output_dir::AbstractString; kwargs...)
    write_data(get_optimization_container(problem), output_dir; kwargs...)
    return
end

struct ProblemSerializationWrapper
    template::OperationsProblemTemplate
    sys::String
    settings::Settings
    problem_type::DataType
end

################ Functions to debug optimization models#####################################
""" "Each Tuple corresponds to (con_name, internal_index, moi_index)"""
function get_all_constraint_index(problem::OperationsProblem)
    con_index = Vector{Tuple{Symbol, Int, Int}}()
    optimization_container = get_optimization_container(problem)
    for (key, value) in get_constraints(optimization_container)
        for (idx, constraint) in enumerate(value)
            moi_index = JuMP.optimizer_index(constraint)
            push!(con_index, (key, idx, moi_index.value))
        end
    end
    return con_index
end

""" "Each Tuple corresponds to (con_name, internal_index, moi_index)"""
function get_all_var_index(problem::OperationsProblem)
    var_index = Vector{Tuple{Symbol, Int, Int}}()
    optimization_container = get_optimization_container(problem)
    for (key, value) in get_variables(optimization_container)
        for (idx, variable) in enumerate(value)
            moi_index = JuMP.optimizer_index(variable)
            push!(var_index, (key, idx, moi_index.value))
        end
    end
    return var_index
end

function get_con_index(problem::OperationsProblem, index::Int)
    optimization_container = get_optimization_container(problem)
    constraints = get_constraints(optimization_container)
    for i in get_all_constraint_index(problem::OperationsProblem)
        if i[3] == index
            return constraints[i[1]].data[i[2]]
        end
    end
    @info "Index not found"
    return
end

function get_var_index(problem::OperationsProblem, index::Int)
    optimization_container = get_optimization_container(problem)
    variables = get_variables(optimization_container)
    for i in get_all_var_index(problem::OperationsProblem)
        if i[3] == index
            return variables[i[1]].data[i[2]]
        end
    end
    @info "Index not found"
    return
end


text:
@testset "Simple simulated DAQ device" begin

end