using Parameters
using LinearAlgebra
using StaticArrays
using Dierckx
using Calculus

abstract type AbstractDomain end
export AbstractDomain

abstract type AbstractConstantKDomain <: AbstractDomain  end
export AbstractConstantKDomain

abstract type AbstractVariableKDomain <: AbstractDomain end
export AbstractVariableKDomain

@with_kw struct ConstantTPDomain{N<:AbstractPhase,S<:Integer,W<:Real, W2<:Real, I<:Integer, Q<:AbstractArray} <: AbstractConstantKDomain
    phase::N
    interfaces::Array{AbstractInterface,1} = Array{AbstractInterface,1}()
    indexes::Q #assumed to be in ascending order
    constantspeciesinds::Array{S,1}
    T::W
    P::W
    kfs::Array{W,1}
    krevs::Array{W,1}
    efficiencyinds::Array{I,1}
    Gs::Array{W,1}
    rxnarray::Array{UInt16,2}
    mu::W = 0.0
    diffusivity::Array{W,1} = Array{Float64,1}()
    jacobian::Array{W,2} = Array{Float64,2}(undef,(0,0))
    sensitivity::Bool = false
    jacuptodate::MArray{Tuple{1},Bool,1,1}=MVector(false)
    t::MArray{Tuple{1},W2,1,1}=MVector(0.0)
end
function ConstantTPDomain(;phase::E2,interfaces::Array{Q,1}=Array{EmptyInterface,1}(),initialconds::Dict{X,X2},constantspecies::Array{X3,1}=Array{String,1}(),
    sparse::Bool=false,sensitivity::Bool=false) where {E<:Real,E2<:AbstractPhase,Q<:AbstractInterface,W<:Real,X,X2,X3}
    #set conditions and initialconditions
    T = 0.0
    P = 0.0
    if sensitivity
        y0 = zeros(length(phase.species)*(1+length(phase.species)+length(phase.reactions)))
    else
        y0 = zeros(length(phase.species))
    end
    spnames = [x.name for x in phase.species]
    for (key,val) in initialconds
        if key == "T"
            T = val
        elseif key == "P"
            P = val
        else
            ind = findfirst(isequal(key),spnames)
            @assert typeof(ind)<: Integer  "$key not found in species list: $spnames"
            y0[ind] = val
        end
    end

    @assert T != 0.0
    @assert P != 0.0
    ns = y0
    N = sum(ns)

    if length(constantspecies) > 0
        spcnames = getfield.(phase.species,:name)
        constspcinds = [findfirst(isequal(k),spcnames) for k in constantspecies]
    else
        constspcinds = Array{Int64,1}()
    end
    efficiencyinds = [rxn.index for rxn in phase.reactions if typeof(rxn.kinetics)<:AbstractFalloffRate && length(rxn.kinetics.efficiencies) > 0]
    Gs = calcgibbs(phase,T)
    if :solvent in fieldnames(typeof(phase)) && typeof(phase.solvent) != EmptySolvent
        mu = phase.solvent.mu(T)
    else
        mu = 0.0
    end
    if phase.diffusionlimited
        diffs = getfield.(phase.species,:diffusion)(T=T,mu=mu,P=P)
    else
        diffs = Array{typeof(T),1}()
    end
    C = P/(R*T)
    V = N*R*T/P
    kfs,krevs = getkfkrevs(phase=phase,T=T,P=P,C=C,N=N,ns=ns,Gs=Gs,diffs=diffs,V=V)
    if sparse
        jacobian=spzeros(typeof(T),length(phase.species),length(phase.species))
    else
        jacobian=zeros(typeof(T),length(phase.species),length(phase.species))
    end
    rxnarray = getreactionindices(phase)
    return ConstantTPDomain(phase,interfaces,SVector(phase.species[1].index,phase.species[end].index),constspcinds,
        T,P,kfs,krevs,efficiencyinds,Gs,rxnarray,mu,diffs,jacobian,sensitivity,MVector(false),MVector(0.0)), y0
end
export ConstantTPDomain

@with_kw struct ConstantVDomain{N<:AbstractPhase,S<:Integer,W<:Real,W2<:Real,Q<:AbstractArray} <: AbstractVariableKDomain
    phase::N
    interfaces::Array{AbstractInterface,1} = Array{AbstractInterface,1}()
    indexes::Q #assumed to be in ascending order
    constantspeciesinds::Array{S,1}
    V::W
    rxnarray::Array{UInt16,2}
    jacobian::Array{W,2}
    sensitivity::Bool = false
    jacuptodate::MArray{Tuple{1},Bool,1,1}=MVector(false)
    t::MArray{Tuple{1},W2,1,1}=MVector(0.0)
end
function ConstantVDomain(;phase::Z,interfaces::Array{Q,1}=Array{EmptyInterface,1}(),initialconds::Dict{X,E},constantspecies::Array{X2,1}=Array{String,1}(),
    sparse::Bool=false,sensitivity::Bool=false) where {E,X,X2,Z<:IdealGas,Q<:AbstractInterface}

    #set conditions and initialconditions
    T = 0.0
    P = 0.0
    V = 0.0
    ns = zeros(length(phase.species))
    spnames = [x.name for x in phase.species]
    for (key,val) in initialconds
        if key == "T"
            T = val
        elseif key == "P"
            P = val
        elseif key == "V"
            V = val
        else
            ind = findfirst(isequal(key),spnames)
            @assert typeof(ind)<: Integer  "$key not found in species list: $spnames"
            ns[ind] = val
        end
    end
    @assert V != 0.0 || (T != 0.0 && P != 0.0)
    N = sum(ns)
    if V == 0.0
        V = N*R*T/P
    elseif T == 0.0
        T = P*V/(R*N)
    elseif P == 0.0
        P = N*R*T/V
    else
        throw(error("ConstantVDomain overspecified with T,P and V"))
    end
    if sensitivity
        y0 = vcat(ns,T,zeros((length(ns)+1)*(length(ns)+length(phase.reactions))))
    else
        y0 = vcat(ns,T)
    end
    if length(constantspecies) > 0
        spcnames = getfield.(phase.species,:name)
        constspcinds = [findfirst(isequal(k),spcnames) for k in constantspecies]
    else
        constspcinds = Array{Int64,1}()
    end
    if sparse
        jacobian=zeros(typeof(T),length(phase.species)+1,length(phase.species)+1)
    else
        jacobian=zeros(typeof(T),length(phase.species)+1,length(phase.species)+1)
    end
    rxnarray = getreactionindices(phase)
    return ConstantVDomain(phase,interfaces,SVector(phase.species[1].index,phase.species[end].index,phase.species[end].index+1),constspcinds,
    V,rxnarray,jacobian,sensitivity,MVector(false),MVector(0.0)), y0
end
export ConstantVDomain

@with_kw struct ParametrizedTPDomain{N<:AbstractPhase,S<:Integer,W<:Real,W2<:Real,Q<:AbstractArray} <: AbstractVariableKDomain
    phase::N
    interfaces::Array{AbstractInterface,1} = Array{AbstractInterface,1}()
    indexes::Q #assumed to be in ascending order
    constantspeciesinds::Array{S,1}
    T::Function
    P::Function
    rxnarray::Array{UInt16,2}
    jacobian::Array{W,2}
    sensitivity::Bool = false
    jacuptodate::MArray{Tuple{1},Bool,1,1}=MVector(false)
    t::MArray{Tuple{1},W2,1,1}=MVector(0.0)
end
function ParametrizedTPDomain(;phase::Z,interfaces::Array{Q,1}=Array{EmptyInterface,1}(),initialconds::Dict{X,Any},constantspecies::Array{X2,1}=Array{String,1}(),
    sparse::Bool=false,sensitivity::Bool=false) where {X,X2,Z<:IdealGas,Q<:AbstractInterface}

    #set conditions and initialconditions
    T = 0.0
    P = 0.0
    V = 0.0
    ts = 0.0
    ns = zeros(length(phase.species))
    spnames = [x.name for x in phase.species]
    for (key,val) in initialconds
        if key == "T"
            T = val
        elseif key == "P"
            P = val
        elseif key == "V"
            V = val
        elseif key == "ts"
            ts = val
        else
            ind = findfirst(isequal(key),spnames)
            @assert typeof(ind) <: Integer  "$key not found in species list: $spnames"
            ns[ind] = val
        end
    end
    @assert V != 0.0 || (T != 0.0 && P != 0.0)
    if isa(T,AbstractArray)
        q = Spline1D(ts,T;k=3,s=1e-11)
        Tfcn(x::Float64) = q(x)
    elseif isa(T,Function)
        Tfcn = T
    else
        throw(error("ParametrizedTPDomain must take \"T\" as a function or if an array of times for \"ts\" is supplied as an array of volumes"))
    end
    if isa(P,AbstractArray)
        v = Spline1D(ts,P;k=3,s=1e-11)
        Pfcn(x::Float64) = v(x)
    elseif isa(P,Function)
        Pfcn = P
    else
        throw(error("ParametrizedTPDomain must take \"P\" as a function or if an array of times for \"ts\" is supplied as an array of volumes"))
    end

    N = sum(ns)
    V = N*R*Tfcn(0.0)/Pfcn(0.0)

    if sensitivity
        y0 = zeros(length(phase.species)*(length(phase.species)+1+length(phase.reactions)))
    else
        y0 = zeros(length(phase.species))
    end

    y0[phase.species[1].index:phase.species[end].index] = ns

    if length(constantspecies) > 0
        spcnames = getfield.(phase.species,:name)
        constspcinds = [findfirst(isequal(k),spcnames) for k in constantspecies]
    else
        constspcinds = Array{Int64,1}()
    end
    if sparse
        jacobian=zeros(typeof(V),length(phase.species)+1,length(phase.species)+1)
    else
        jacobian=zeros(typeof(V),length(phase.species)+1,length(phase.species)+1)
    end
    rxnarray = getreactionindices(phase)
    return ParametrizedTPDomain(phase,interfaces,SVector(phase.species[1].index,phase.species[end].index),constspcinds,
    Tfcn,Pfcn,rxnarray,jacobian,sensitivity,MVector(false),MVector(0.0)), y0
end
export ParametrizedTPDomain

@with_kw struct ParametrizedVDomain{N<:AbstractPhase,S<:Integer,W<:Real,W2<:Real,Q<:AbstractArray} <: AbstractVariableKDomain
    phase::N
    interfaces::Array{AbstractInterface,1} = Array{AbstractInterface,1}()
    indexes::Q #assumed to be in ascending order
    constantspeciesinds::Array{S,1}
    V::Function
    rxnarray::Array{UInt16,2}
    jacobian::Array{W,2}
    sensitivity::Bool = false
    jacuptodate::MArray{Tuple{1},Bool,1,1}=MVector(false)
    t::MArray{Tuple{1},W2,1,1}=MVector(0.0)
end
function ParametrizedVDomain(;phase::Z,interfaces::Array{Q,1}=Array{EmptyInterface,1}(),initialconds::Dict{X,Any},constantspecies::Array{X2,1}=Array{String,1}(),
    sparse::Bool=false,sensitivity::Bool=false) where {X,X2,E<:Real,Z<:IdealGas,Q<:AbstractInterface}

    #set conditions and initialconditions
    T = 0.0
    P = 0.0
    V = 0.0
    ts = Array{Float64,1}()
    ns = zeros(length(phase.species))
    spnames = [x.name for x in phase.species]
    @assert "V" in keys(initialconds)
    for (key,val) in initialconds
        if key == "T"
            T = val
        elseif key == "P"
            P = val
        elseif key == "V"
            V = val
        elseif key == "ts"
            ts = val
        else
            ind = findfirst(isequal(key),spnames)
            @assert typeof(ind)<: Integer  "$key not found in species list: $spnames"
            ns[ind] = val
        end
    end
    @assert isa(V,Function) || isa(V,AbstractArray)
    if isa(V,AbstractArray)
        q = Spline1D(ts,V;k=3,s=1e-11)
        Vfcn = f(x::Float64) = q(x)
    elseif isa(V,Function)
        Vfcn = V
    else
        throw(error("ParametrizedVDomain must take \"V\" as a function or if an array of times for \"ts\" is supplied as an array of volumes"))
    end
    N = sum(ns)
    if T == 0.0
        T = P*Vfcn(0.0)/(R*N)
    elseif P == 0.0
        P = N*R*T/Vfcn(0.0)
    else
        ns *= (P*Vfcn(0.0)/(R*T))/sum(ns) #automatically scale down moles if pressure specified
    end
    if sensitivity
        y0 = vcat(ns,T,zeros((length(ns)+1)*(length(ns)+length(phase.reactions))))
    else
        y0 = vcat(ns,T)
    end
    if length(constantspecies) > 0
        spcnames = getfield.(phase.species,:name)
        constspcinds = [findfirst(isequal(k),spcnames) for k in constantspecies]
    else
        constspcinds = Array{Int64,1}()
    end
    if sparse
        jacobian=zeros(typeof(T),length(phase.species)+1,length(phase.species)+1)
    else
        jacobian=zeros(typeof(T),length(phase.species)+1,length(phase.species)+1)
    end
    rxnarray = getreactionindices(phase)
    return ParametrizedVDomain(phase,interfaces,SVector(phase.species[1].index,phase.species[end].index,phase.species[end].index+1),constspcinds,
    Vfcn,rxnarray,jacobian,sensitivity,MVector(false),MVector(0.0)), y0
end
export ParametrizedVDomain
@with_kw struct ConstantTVDomain{N<:AbstractPhase,S<:Integer,W<:Real, W2<:Real, I<:Integer, Q<:AbstractArray} <: AbstractConstantKDomain
    phase::N
    interfaces::Array{AbstractInterface,1} = Array{AbstractInterface,1}()
    indexes::Q #assumed to be in ascending order
    constantspeciesinds::Array{S,1}
    T::W
    V::W
    kfs::Array{W,1}
    krevs::Array{W,1}
    efficiencyinds::Array{I,1}
    Gs::Array{W,1}
    rxnarray::Array{UInt16,2}
    mu::W = 0.0
    diffusivity::Array{W,1} = Array{Float64,1}()
    jacobian::Array{W,2} = Array{Float64,2}(undef,(0,0))
    sensitivity::Bool = false
    jacuptodate::MArray{Tuple{1},Bool,1,1}=MVector(false)
    t::MArray{Tuple{1},W2,1,1}=MVector(0.0)
end
function ConstantTVDomain(;phase::Z,interfaces::Array{Q,1}=Array{EmptyInterface,1}(),initialconds::Dict{X,E},constantspecies::Array{X2,1}=Array{String,1}(),
    sparse=false,sensitivity=false) where {E,X,X2, Z<:AbstractPhase,Q<:AbstractInterface,W<:Real}
    #set conditions and initialconditions
    T = 0.0
    V = 0.0
    P = 1.0e9
    if sensitivity
        y0 = zeros(length(phase.species)*(length(phase.species)+1+length(phase.reactions)))
    else
        y0 = zeros(length(phase.species))
    end
    spnames = [x.name for x in phase.species]
    for (key,val) in initialconds
        if key == "T"
            T = val
        elseif key == "P"
            P = val
        elseif key == "V"
            V = val
        else
            ind = findfirst(isequal(key),spnames)
            @assert typeof(ind)<: Integer  "$key not found in species list: $spnames"
            y0[ind] = val
        end
    end
    @assert T != 0.0
    @assert V != 0.0
    ns = y0
    N = sum(ns)

    if length(constantspecies) > 0
        spcnames = getfield.(phase.species,:name)
        constspcinds = [findfirst(isequal(k),spcnames) for k in constantspecies]
    else
        constspcinds = Array{Int64,1}()
    end
    efficiencyinds = [rxn.index for rxn in phase.reactions if typeof(rxn.kinetics)<:AbstractFalloffRate && length(rxn.kinetics.efficiencies) > 0]
    Gs = calcgibbs(phase,T)
    if :solvent in fieldnames(typeof(phase)) && typeof(phase.solvent) != EmptySolvent
        mu = phase.solvent.mu(T)
    else
        mu = 0.0
    end
    if phase.diffusionlimited
        diffs = [x(T=T,mu=mu,P=P) for x in getfield.(phase.species,:diffusion)]
    else
        diffs = Array{Float64,1}()
    end
    P = 1.0e9  #essentiallly assuming this is a liquid
    C = N/V
    kfs,krevs = getkfkrevs(phase=phase,T=T,P=P,C=C,N=N,ns=ns,Gs=Gs,diffs=diffs,V=V)
    if sparse
        jacobian=zeros(typeof(T),length(phase.species),length(phase.species))
    else
        jacobian=zeros(typeof(T),length(phase.species),length(phase.species))
    end
    rxnarray = getreactionindices(phase)
    return ConstantTVDomain(phase,interfaces,SVector(phase.species[1].index,phase.species[end].index),constspcinds,
        T,V,kfs,krevs,efficiencyinds,Gs,rxnarray,mu,diffs,jacobian,sensitivity,MVector(false),MVector(0.0)), y0
end
export ConstantTVDomain

@with_kw struct ParametrizedTConstantVDomain{N<:AbstractPhase,S<:Integer,W<:Real,W2<:Real,Q<:AbstractArray} <: AbstractVariableKDomain
    phase::N
    interfaces::Array{AbstractInterface,1} = Array{AbstractInterface,1}()
    indexes::Q #assumed to be in ascending order
    constantspeciesinds::Array{S,1}
    T::Function
    V::W
    rxnarray::Array{UInt16,2}
    jacobian::Array{W,2}
    sensitivity::Bool = false
    jacuptodate::MArray{Tuple{1},Bool,1,1}=MVector(false)
    t::MArray{Tuple{1},W2,1,1}=MVector(0.0)
end
function ParametrizedTConstantVDomain(;phase::IdealDiluteSolution,interfaces::Array{Q,1}=Array{EmptyInterface,1}(),initialconds::Dict{X,X3},constantspecies::Array{X2,1}=Array{String,1}(),
    sparse::Bool=false,sensitivity::Bool=false) where {X,X2,X3,Q<:AbstractInterface}
    #set conditions and initialconditions
    T = 0.0
    P = 0.0
    V = 0.0
    ts = Array{Float64,1}()
    ns = zeros(length(phase.species))
    spnames = [x.name for x in phase.species]
    @assert "V" in keys(initialconds)
    for (key,val) in initialconds
        if key == "T"
            T = val
        elseif key == "P"
            P = val
        elseif key == "V"
            V = val
        elseif key == "ts"
            ts = val
        else
            ind = findfirst(isequal(key),spnames)
            @assert typeof(ind)<: Integer  "$key not found in species list: $spnames"
            ns[ind] = val
        end
    end
    if isa(T,AbstractArray)
        q = Spline1D(ts,T;k=3,s=1e-11)
        Tfcn = f(x::Float64) = q(x)
    elseif isa(T,Function)
        Tfcn = T
    else
        throw(error("ParametrizedTConstantVDomain must take \"T\" as a function or if an array of times for \"ts\" is supplied as an array of volumes"))
    end
    N = sum(ns)
    if P == 0.0
        P = 1e8
    else
        throw(error("ParametrizedTConstantVDomain cannot specify P"))
    end
    if sensitivity
        y0 = zeros(length(phase.species)*(length(phase.species)+1+length(phase.reactions)))
    else
        y0 = zeros(length(phase.species))
    end
    y0[phase.species[1].index:phase.species[end].index] = ns
    if length(constantspecies) > 0
        spcnames = getfield.(phase.species,:name)
        constspcinds = [findfirst(isequal(k),spcnames) for k in constantspecies]
    else
        constspcinds = Array{Int64,1}()
    end
    if sparse
        jacobian=zeros(typeof(V),length(phase.species)+1,length(phase.species)+1)
    else
        jacobian=zeros(typeof(V),length(phase.species)+1,length(phase.species)+1)
    end
    rxnarray = getreactionindices(phase)
    return ParametrizedTConstantVDomain(phase,interfaces,SVector(phase.species[1].index,phase.species[end].index),constspcinds,
    Tfcn,V,rxnarray,jacobian,sensitivity,MVector(false),MVector(0.0)), y0
end
export ParametrizedTConstantVDomain

@inline function calcthermo(d::ConstantTPDomain{W,Y},y::J,t::Q) where {W<:IdealGas,Y<:Integer,J<:AbstractArray{Float64,1},Q<:Float64}
    if t != d.t[1]
        d.t[1] = t
        d.jacuptodate[1] = false
    end
    ns = y[d.indexes[1]:d.indexes[2]]
    N = sum(ns)
    V = N*d.T*R/d.P
    cs = ns./V
    C = N/V
    for ind in d.efficiencyinds #efficiency related rates may have changed
        d.kfs[ind],d.krevs[ind] = getkfkrev(d.phase.reactions[ind],d.phase,d.T,d.P,C,N,ns,d.Gs,d.diffusivity,V)
    end
    return ns,cs,d.T,d.P,V,C,N,d.mu,d.kfs,d.krevs,[],[],[],[],0.0
end

@inline function calcthermo(d::ConstantTPDomain{W,Y},y::J,t::Q) where {W<:IdealGas,Y<:Integer,J<:AbstractArray,Q<:Real}
    if t != d.t[1]
        if isa(t,Float64)
            d.t[1] = t
        end
        d.jacuptodate[1] = false
    end
    ns = y[d.indexes[1]:d.indexes[2]]
    N = sum(ns)
    V = N*d.T*R/d.P
    cs = ns./V
    C = N/V
    kfs = convert(typeof(y),copy(d.kfs))
    krevs = convert(typeof(y),copy(d.krevs))
    for ind in d.efficiencyinds #efficiency related rates may have changed
        kfs[ind],krevs[ind] = getkfkrev(d.phase.reactions[ind],d.phase,d.T,d.P,C,N,ns,d.Gs,d.diffusivity,V)
    end
    return ns,cs,d.T,d.P,V,C,N,d.mu,kfs,krevs,[],[],[],[],0.0
end

@inline function calcthermo(d::ConstantVDomain{W,Y},y::J,t::Q) where {W<:IdealGas,Y<:Integer,J<:AbstractArray,Q<:Real}
    if t != d.t[1]
        d.t[1] = t
        d.jacuptodate[1] = false
    end
    ns = y[d.indexes[1]:d.indexes[2]]
    T = y[d.indexes[3]]
    N = sum(ns)
    cs = ns./d.V
    C = N/d.V
    P = C*R*T
    Gs = zeros(length(d.phase.species))
    Us = zeros(length(d.phase.species))
    Cvave = 0.0
    @simd for i = 1:length(d.phase.species)
        @inbounds cpdivR,hdivRT,sdivR = calcHSCpdless(d.phase.species[i].thermo,T)
        @fastmath @inbounds Gs[i] = (hdivRT-sdivR)*R*T
        @fastmath @inbounds Us[i] = (hdivRT-1.0)*R*T
        @fastmath @inbounds Cvave += cpdivR*ns[i]
    end
    @fastmath Cvave *= R/N
    @fastmath Cvave -= R
    if d.phase.diffusionlimited
        diffs = getfield.(d.phase.species,:diffusion)(T=T,mu=mu,P=P)
    else
        diffs = Array{Float64,1}()
    end
    kfs,krevs = getkfkrevs(phase=d.phase,T=T,P=P,C=C,N=N,ns=ns,Gs=Gs,diffs=diffs,V=d.V)
    return ns,cs,T,P,d.V,C,N,0.0,kfs,krevs,[],Us,Gs,diffs,Cvave
end

@inline function calcthermo(d::ParametrizedVDomain{W,Y},y::J,t::Q) where {W<:IdealGas,Y<:Integer,J<:AbstractArray,Q<:Real}
    if t != d.t[1]
        d.t[1] = t
        d.jacuptodate[1] = false
    end
    V = d.V(t)
    ns = y[d.indexes[1]:d.indexes[2]]
    T = y[d.indexes[3]]
    N = sum(ns)
    cs = ns./V
    C = N/V
    P = C*R*T
    Gs = zeros(length(d.phase.species))
    Us = zeros(length(d.phase.species))
    Cvave = 0.0
    @simd for i = 1:length(d.phase.species)
        @inbounds cpdivR,hdivRT,sdivR = calcHSCpdless(d.phase.species[i].thermo,T)
        @fastmath @inbounds Gs[i] = (hdivRT-sdivR)*R*T
        @fastmath @inbounds Us[i] = (hdivRT-1.0)*R*T
        @fastmath @inbounds Cvave += cpdivR*ns[i]
    end
    @fastmath Cvave *= R/N
    @fastmath Cvave -= R
    if d.phase.diffusionlimited
        diffs = getfield.(d.phase.species,:diffusion)(T=T,mu=mu,P=P)
    else
        diffs = Array{Float64,1}()
    end
    kfs,krevs = getkfkrevs(phase=d.phase,T=T,P=P,C=C,N=N,ns=ns,Gs=Gs,diffs=diffs,V=V)
    return ns,cs,T,P,V,C,N,0.0,kfs,krevs,[],Us,Gs,diffs,Cvave
end

@inline function calcthermo(d::ParametrizedTConstantVDomain{W,Y},y::J,t::Q) where {W<:IdealDiluteSolution,Y<:Integer,J<:AbstractArray,Q<:Real}
    if t != d.t[1]
        d.t[1] = t
        d.jacuptodate[1] = false
    end
    V = d.V
    T = d.T(t)
    ns = y[d.indexes[1]:d.indexes[2]]
    N = sum(ns)
    cs = ns./V
    C = N/V
    P = C*R*T
    Gs = zeros(length(d.phase.species))
    mu = d.phase.solvent.mu(T)
    for i = 1:length(d.phase.species)
        @inbounds cpdivR,hdivRT,sdivR = calcHSCpdless(d.phase.species[i].thermo,T)
        @fastmath @inbounds Gs[i] = (hdivRT-sdivR)*R*T
    end
    if d.phase.diffusionlimited
        diffs = [x(T=T,mu=mu,P=P) for x in getfield.(d.phase.species,:diffusion)]
    else
        diffs = Array{Float64,1}()
    end
    kfs,krevs = getkfkrevs(phase=d.phase,T=T,P=P,C=C,N=N,ns=ns,Gs=Gs,diffs=diffs,V=V)
    return ns,cs,T,P,V,C,N,0.0,kfs,krevs,[],[],Gs,diffs,0.0
end

@inline function calcthermo(d::ParametrizedTPDomain{W,Y},y::J,t::Q) where {W<:IdealGas,Y<:Integer,J<:AbstractArray,Q<:Real}
    if t != d.t[1]
        d.t[1] = t
        d.jacuptodate[1] = false
    end
    T = d.T(t)
    @assert T < 10000.0
    P = d.P(t)
    ns = y[d.indexes[1]:d.indexes[2]]
    N = sum(ns)
    V = N*T*R/P
    cs = ns./V
    C = N/V
    P = C*R*T
    Gs = zeros(length(d.phase.species))
    Us = zeros(length(d.phase.species))
    Cvave = 0.0
    @simd for i = 1:length(d.phase.species)
        @inbounds cpdivR,hdivRT,sdivR = calcHSCpdless(d.phase.species[i].thermo,T)
        @fastmath @inbounds Gs[i] = (hdivRT-sdivR)*R*T
    end
    @fastmath Cvave *= R/N
    @fastmath Cvave -= R
    if d.phase.diffusionlimited
        diffs = getfield.(d.phase.species,:diffusion)(T=T,mu=mu,P=P)
    else
        diffs = Array{Float64,1}()
    end
    kfs,krevs = getkfkrevs(phase=d.phase,T=T,P=P,C=C,N=N,ns=ns,Gs=Gs,diffs=diffs,V=V)
    return ns,cs,T,P,V,C,N,0.0,kfs,krevs,[],[],Gs,diffs,0.0
end

@inline function calcthermo(d::ConstantTVDomain{W,Y},y::J,t::Q) where {W<:IdealDiluteSolution,Y<:Integer,J<:AbstractArray,Q<:Real}
    if t != d.t[1]
        d.t[1] = t
        d.jacuptodate[1] = false
    end
    ns = y[d.indexes[1]:d.indexes[2]]
    N = sum(ns)
    cs = ns./d.V
    C = N/d.V
    P = 1.0e9
    return ns,cs,d.T,P,d.V,C,N,d.mu,d.kfs,d.krevs,[],[],[],[],0.0
end
export calcthermo

@inline function calcdomainderivatives!(d::Q,dydt::Array{Z7,1};t::Z10,T::Z4,P::Z9,Us::Array{Z,1},V::Z2,C::Z3,ns::Array{Z5,1},N::Z6,Cvave::Z8) where {Q<:AbstractDomain,Z10,Z9,Z8<:Real,Z7<:Real,W<:IdealGas,Y<:Integer,Z6,Z,Z2,Z3,Z4,Z5<:Real}
    for ind in d.constantspeciesinds #make dydt zero for constant species
        @inbounds dydt[ind] = 0.0
    end
end

@inline function calcdomainderivatives!(d::ConstantVDomain{W,Y},dydt::Array{K,1};t::Z10,T::Z4,P::Z9,Us::Array{Z,1},V::Z2,C::Z3,ns::Array{Z5,1},N::Z6,Cvave::Z7) where {Z10,Z9,W<:IdealGas,Z7<:Real,K<:Real,Y<:Integer,Z6,Z,Z2,Z3,Z4,Z5<:Real}
    @views @fastmath @inbounds dydt[d.indexes[3]] = -dot(Us,dydt[d.indexes[1]:d.indexes[2]])/(N*Cvave) #divide by V to cancel ωV to ω
    for ind in d.constantspeciesinds #make dydt zero for constant species
        @inbounds dydt[ind] = 0.0
    end
end

@inline function calcdomainderivatives!(d::ParametrizedVDomain{W,Y},dydt::Array{K,1};t::Z10,T::Z4,P::Z9,Us::Array{Z,1},V::Z2,C::Z3,ns::Array{Z5,1},N::Z6,Cvave::Z7) where {Z10,Z9,W<:IdealGas,Z7<:Real,K<:Real,Y<:Integer,Z6,Z,Z2,Z3,Z4,Z5<:Real}
    @views @fastmath @inbounds dydt[d.indexes[3]] = (-dot(Us,dydt[d.indexes[1]:d.indexes[2]])-P*Calculus.derivative(d.V,t))/(N*Cvave) #divide by V to cancel ωV to ω

    for ind in d.constantspeciesinds #make dydt zero for constant species
        @inbounds dydt[ind] = 0.0
    end
end
export calcdomainderivatives!

function getreactionindices(ig::Q) where {Q<:AbstractPhase}
    arr = zeros(UInt16,(6,length(ig.reactions)))
    for (i,rxn) in enumerate(ig.reactions)
        arr[1:length(rxn.reactantinds),i] = rxn.reactantinds
        arr[4:length(rxn.productinds)+3,i] = rxn.productinds
    end
    return arr
end
export getreactionindices


isdefined(Base, :__precompile__) && __precompile__()
module PowerSimulations

#################################################################################
# Exports

# Base Models
export Stage
export Simulation
export OperationsProblem
export OperationsProblemTemplate
export InitialCondition
export SimulationSequence

#Network Relevant Exports
export StandardPTDFModel
export CopperPlatePowerModel
export AreaBalancePowerModel

######## Device Models ########
export DeviceModel
export FixedOutput
######## Service Models ########
export ServiceModel
export RangeReserve
export StepwiseCostReserve
export PIDSmoothACE
export GroupReserve
######## Branch Models ########
export StaticLine
export StaticLineBounds
export StaticLineUnbounded
export StaticTransformer
export StaticTransformerBounds
export StaticTransformerUnbounded
export FlowMonitoredLine
export HVDCLossless
export HVDCDispatch
export HVDCUnbounded
#export VoltageSourceDC
######## Load Models ########
export StaticPowerLoad
export InterruptiblePowerLoad
export DispatchablePowerLoad
######## Renewable Formulations ########
export RenewableFullDispatch
export RenewableConstantPowerFactor
######## Hydro Formulations ########
export HydroDispatchRunOfRiver
export HydroDispatchReservoirBudget
export HydroDispatchReservoirStorage
export HydroCommitmentRunOfRiver
export HydroCommitmentReservoirBudget
export HydroCommitmentReservoirStorage
export HydroDispatchPumpedStorage
export HydroDispatchPumpedStoragewReservation
######## Renewable Formulations ########
export BookKeeping
export BookKeepingwReservation
######## Thermal Formulations ########
export ThermalStandardUnitCommitment
export ThermalBasicUnitCommitment
export ThermalDispatch
export ThermalRampLimited
export ThermalDispatchNoMin
export ThermalMultiStartUnitCommitment

###### Regulation Device Formulation #######
export DeviceLimitedRegulation
export ReserveLimitedRegulation

# feedforward chrons
export RecedingHorizon
export Synchronize
export Consecutive
export FullHorizon
export Range

# feedforward models
export UpperBoundFF
export SemiContinuousFF
export RangeFF
export IntegralLimitFF
export ParameterFF

# InitialConditions chrons
export InterStageChronology
export IntraStageChronology

# Initial Conditions Quantities
export DevicePower
export DeviceStatus
export TimeDurationON
export TimeDurationOFF
export EnergyLevel

# cache_models
export TimeStatusChange
export StoredEnergy

#operation_models
export GenericOpProblem
export UnitCommitmentProblem
export EconomicDispatchProblem
#export OptimalPowerFlow

# Functions
## Construction Exports
export construct_device!
export construct_network!
## Op Model Exports
export solve!
export get_initial_conditions
export set_transmission_model!
export set_devices_template!
export set_branches_template!
export set_services_template!
export set_device_model!
export set_branch_model!
export set_device_model!
export set_model!
## Sim Model Exports
export build!
export execute!
export make_references
## Template Exports
export template_economic_dispatch
export template_unit_commitment
export template_agc_reserve_deployment
export EconomicDispatchProblem
export UnitCommitmentProblem
export AGCReserveDeployment
export run_economic_dispatch
export run_unit_commitment
## Results interfaces
export get_duals

## Utils Exports
export SimulationResultsReference
export write_results
export check_file_integrity
export load_results
export load_operation_results
export load_simulation_results
export write_to_CSV
export get_all_constraint_index
export get_all_var_index
export get_con_index
export get_var_index
export get_result_variable
export get_variable_names
export configure_logging
export show_recorder_events
export list_simulation_events
export show_simulation_events

# Variables / Parameters
export ACTIVE_POWER
export ENERGY
export ENERGY_BUDGET
export FLOW_ACTIVE_POWER
export ON
export REACTIVE_POWER
export ACTIVE_POWER_IN
export ACTIVE_POWER_OUT
export RESERVE
export SERVICE_REQUIREMENT
export START
export STOP
export THETA
export VM
export INFLOW
export SPILLAGE
export SLACK_UP
export SLACK_DN

# Constraints
export ACTIVE
export ACTIVE_RANGE
export ACTIVE_RANGE_LB
export ACTIVE_RANGE_UB
export COMMITMENT
export DURATION
export DURATION_DOWN
export DURATION_UP
export ENERGY_CAPACITY
export ENERGY_LIMIT
export FEEDFORWARD
export FEEDFORWARD_UB
export FEEDFORWARD_BIN
export FEEDFORWARD_INTEGRAL_LIMIT
export FLOW_LIMIT
export FLOW_LIMIT_FROM_TO
export FLOW_LIMIT_TO_FROM
export FLOW_REACTIVE_POWER_FROM_TO
export FLOW_REACTIVE_POWER_TO_FROM
export FLOW_ACTIVE_POWER_FROM_TO
export FLOW_ACTIVE_POWER_TO_FROM
export FLOW_ACTIVE_POWER
export FLOW_REACTIVE_POWER
export INPUT_POWER_RANGE
export OUTPUT_POWER_RANGE
export RAMP
export RAMP_DOWN
export RAMP_UP
export RATE_LIMIT
export RATE_LIMIT_FT
export RATE_LIMIT_TF
export REACTIVE
export REACTIVE_RANGE
export REQUIREMENT
export INFLOW_RANGE

#################################################################################
# Imports
import Logging
import Serialization
#Modeling Imports
import JuMP
# so that users do not need to import JuMP to use a solver with PowerModels
import JuMP: optimizer_with_attributes
export optimizer_with_attributes
import MathOptInterface
import ParameterJuMP
import LinearAlgebra
import PowerSystems
import InfrastructureSystems
# so that users have access to IS.Results interfaces
import InfrastructureSystems:
    get_variables, get_total_cost, get_optimizer_log, write_results, get_timestamp, get_name
export get_name
export get_model_base_power
export get_variables
export get_dual_values
export get_total_cost
export get_optimizer_log
export get_timestamp
export write_results
import PowerModels
import TimerOutputs

#TimeStamp Management Imports
import Dates
import TimeSeries

#I/O Imports
import DataFrames
import Feather
import JSON
import CSV
import SHA

include("core/definitions.jl")

################################################################################
# Includes

include("logging.jl")
include("utils.jl")

#Models and constructors
include("core/results.jl")
include("core/abstract_types.jl")
include("core/aux_structs.jl")

include("services_models/services_model.jl")
include("devices_models/device_model.jl")
include("network_models/networks.jl")

include("core/parameters.jl")
include("core/variables.jl")
include("core/constraints.jl")
include("core/cache.jl")
include("core/initial_condition_types.jl")
include("core/initial_condition.jl")
include("core/initial_conditions.jl")
include("core/operations_problem_template.jl")
include("core/psi_settings.jl")
include("core/psi_container.jl")
include("core/update_initial_conditions.jl")
include("core/operations_problem_results.jl")
include("core/operations_problem.jl")
include("core/simulation_stages.jl")
include("core/simulation_sequence.jl")
include("core/simulation.jl")
include("devices_models/devices/common.jl")
include("core/feedforward.jl")
include("core/simulation_results.jl")
include("core/recorder_events.jl")

#Device Modeling components
include("devices_models/devices/common/device_range_constraints.jl")
include("devices_models/devices/common/nodal_expression.jl")
include("devices_models/devices/renewable_generation.jl")
include("devices_models/devices/thermal_generation.jl")
include("devices_models/devices/electric_loads.jl")
include("devices_models/devices/AC_branches.jl")
include("devices_models/devices/DC_branches.jl")
include("devices_models/devices/storage.jl")
include("devices_models/devices/hydro_generation.jl")
include("devices_models/devices/regulation_device.jl")

#Services Models
include("services_models/agc.jl")
include("services_models/reserves.jl")
include("services_models/GroupReserve.jl")
include("services_models/service_slacks.jl")
include("services_models/services_constructor.jl")

#Network models
include("network_models/copperplate_model.jl")
include("network_models/powermodels_interface.jl")
include("network_models/ptdf_model.jl")
include("network_models/network_slack_variables.jl")
include("network_models/area_balance_model.jl")

#Device constructors
include("devices_models/device_constructors/common/constructor_validations.jl")
include("devices_models/device_constructors/thermalgeneration_constructor.jl")
include("devices_models/device_constructors/hydrogeneration_constructor.jl")
include("devices_models/device_constructors/branch_constructor.jl")
include("devices_models/device_constructors/renewablegeneration_constructor.jl")
include("devices_models/device_constructors/load_constructor.jl")
include("devices_models/device_constructors/storage_constructor.jl")
include("devices_models/device_constructors/regulationdevice_constructor.jl")

#Network constructors
include("network_models/network_constructor.jl")

#Templates
include("operations_problems_templates.jl")
# Printing
include("printing.jl")

end


include("../core/blocks.jl")
include("../../configs.jl")

"""
Upsampling and Pyramid Network for the RetinaFace.
"""
struct FPN o6; o2; o3; o4; o5; merge4; merge3; merge2; scale_cnt; end

function FPN(;in_size=2048, dtype=Array{Float32s}, scale_cnt=5)
    return FPN(
        ConvBn(3, 3, in_size, 256, init=xavier, stride=2, bias=false, padding=1, dtype=dtype),

        ConvBn(1, 1, 256, 256, f=relu, bias=false, dtype=dtype),
        ConvBn(1, 1, 512, 256, f=relu, bias=false, dtype=dtype),
        ConvBn(1, 1, 1024, 256, f=relu, bias=false, dtype=dtype),
        ConvBn(1, 1, 2048, 256, f=relu, bias=false, dtype=dtype),

        ConvBn(3, 3, 256, 256, f=relu, bias=false, dtype=dtype, padding=1),
        ConvBn(3, 3, 256, 256, f=relu, bias=false, dtype=dtype, padding=1),
        ConvBn(3, 3, 256, 256, f=relu, bias=false, dtype=dtype, padding=1),
        scale_cnt
    )
end

function (fpn::FPN)(xs; train=true)
    c2, c3, c4, c5 = xs
    # extra conv layer
    if fpn.scale_cnt == 5
        p6 = fpn.o6(c5, train=train)
        # refucing the dim. to 256 and upsampling
        p5 = fpn.o5(c5, train=train)
        p4 = fpn.o4(c4, train=train) + unpool(p5)
        p4 = fpn.merge4(p4, train=train)
        p3 = fpn.o3(c3, train=train) + unpool(p4)
        p3 = fpn.merge3(p3, train=train)
        p2 = fpn.o2(c2, train=train) + unpool(p3)
        p2 = fpn.merge2(p2, train=train)
        return [p2, p3, p4, p5, p6]

    elseif fpn.scale_cnt == 3
        p5 = fpn.o5(c5, train=train)
        p4 = fpn.o4(c4, train=train) + unpool(p5)
        p4 = fpn.merge4(p4, train=train)
        p3 = fpn.o3(c3, train=train) + unpool(p4)
        p3 = fpn.merge3(p3, train=train)
        return [p3, p4, p5]
    end
    return nothing
end


using MuZero
using Test

@testset "MuZero.jl" begin
    @test true
end

include("../diagnostics/tac/game.jl")
include("../diagnostics/tac/inference_tests.jl")
dir = "$(@__DIR__)/output"
tests = [
    get_policyvalue_tests(Tac),
    get_reward_tests(Tac),
]
config = getconfig(Tac)
config.save_interval = 1
config.checkpoint_interval = 1
config.batch_size = 10
config.num_filters = 3
config.representation_resnet_layers = 1
config.dynamics_resnet_layers = 1
config.reward_fc_layers = 1
config.policy_fc_layers = 1
config.value_fc_layers = 1

@time @testset "Basic learner save/load" begin
    l1 = Learner(Tac; config=config, dir=dir, tests=tests)
    @test get_steps(l1) == 0
    save!(l1)
    l2 = Learner(Tac; config=config, uuid=l1.uuid, dir=dir, tests=tests)
    @test get_steps(l2) == 0
    @test l1.train.saved_at == l2.train.saved_at
    @test l1.training_started == l2.training_started
    @test l1.stats.losses == l2.stats.losses
    @test l1.stats.accuracies == l2.stats.accuracies
    @test l1.stats.step_count == l2.stats.step_count
end

@time @testset "Basic network training" begin
    l = Learner(Tac; config=config, dir=dir, tests=tests)
    selfplay!(l; limit=10)
    @test length(l.stats.losses) == 0
    @test length(l.stats.accuracies) == 0
    fit!(l, steps=1)
    @test length(l.stats.accuracies) == 2
    @test length(l.stats.losses) == 2
    @test get_steps(l) == 1
    fit!(l, steps=1)
    @test length(l.stats.accuracies) == 3
    @test length(l.stats.losses) == 3
    @test get_steps(l) == 2
end

@time @testset "Training/test networks save/load" begin
    l1 = Learner(Tac; config=config, dir=dir, tests=tests)
    selfplay!(l1; limit=10)
    fit!(l1, steps=1)
    batch_tensor = randtrainingbatch(l1) |> tensor
    train_net_initial_out1 = initial_inference(l1.train, batch_tensor.images) |> NetworkOutput
    test_net_initial_out1 = initial_inference(l1.test, batch_tensor.images) |> NetworkOutput
    loss1, accuracy1 = get_loss(l1), get_accuracy(l1)
    save!(l1)
    l2 = Learner(Tac; config=config, uuid=l1.uuid, dir=dir, tests=tests)
    train_net_initial_out2 = initial_inference(l2.train, batch_tensor.images) |> NetworkOutput
    test_net_initial_out2 = initial_inference(l2.test, batch_tensor.images) |> NetworkOutput
    loss2, accuracy2 = get_loss(l2), get_accuracy(l2)
    @test train_net_initial_out1 == train_net_initial_out2
    @test test_net_initial_out1 == test_net_initial_out2
    @test accuracy1 == accuracy2
    @test loss1 == loss2
end


export SimulationBuilder, set!, build

mutable struct SimulationBuilder{N}
    unit_system::UnitSystem
    kernel::Kernel{N}
    eos::EquationOfState
    mass_distribution::InitialMassDistribution{N}
    velocity_distribution::InitialVelocityDistribution{N}
    energy_distribution::InitialEnergyDistribution{N}
    mass_component_type::Tuple{Type,Any}
    energy_component_type::Tuple{Type,Any}
    diffusion_component_type::Tuple{Type,Any}
    boundary_component_type::Tuple{Type,Any}
    stepper_type::Tuple{Type,Any}

    function SimulationBuilder(
        mass_distribution::InitialMassDistribution{N},
        energy_distribution::InitialEnergyDistribution{N};
        unit_system::UnitSystem = SI,
        kernel::Kernel{N} = TabulatedKernel(CubicSplineKernel{N}()),
        eos::EquationOfState = AdiabaticIdealGasEOS(1.4, NORMALIZED),
        velocity_distribution::InitialVelocityDistribution{N} = StaticVelocityDistribution{
            N,
        }(),
    ) where {N}
        new{N}(
            unit_system,
            kernel,
            eos,
            mass_distribution,
            velocity_distribution,
            energy_distribution,
            (AdaptiveMass, []),
            (StandardEnergy, []),
            (StandardDiffusion, []),
            (NoBoundaries{N}, []),
            (EulerStepper, []),
        )
    end
end

function Base.show(io::IO, builder::SimulationBuilder{N}) where {N}
    print(
        io,
        """
        SimulationBuilder{
            $(N)D
            $(typeof(builder.mass_distribution))
            $(typeof(builder.velocity_distribution))
            $(typeof(builder.energy_distribution))
            $(typeof(builder.eos))
            $(typeof(builder.kernel))
            $(builder.mass_component_type[1])
            $(builder.energy_component_type[1])
            $(builder.diffusion_component_type[1])
            $(builder.boundary_component_type[1])
            $(builder.stepper_type[1])
        }""",
    )
end

(Base.:+)(builder::SimulationBuilder, property) = set!(builder, property)
(Base.:+)(
    builder::SimulationBuilder,
    (property, kwargs)::Tuple{<:Any,NamedTuple},
) = set!(builder, property, kwargs...)

function set!(builder::SimulationBuilder{N}, unit_system::UnitSystem) where {N}
    builder.unit_system = unit_system
    builder
end

function set!(builder::SimulationBuilder{N}, kernel::Kernel{N}) where {N}
    builder.kernel = kernel
    builder
end

function set!(builder::SimulationBuilder{N}, eos::EquationOfState) where {N}
    builder.eos = eos
    builder
end

function set!(
    builder::SimulationBuilder{N},
    mass_distribution::InitialMassDistribution{N},
) where {N}
    builder.mass_distribution = mass_distribution
    builder
end

function set!(
    builder::SimulationBuilder{N},
    velocity_distribution::InitialVelocityDistribution{N},
) where {N}
    builder.velocity_distribution = velocity_distribution
    builder
end

function set!(
    builder::SimulationBuilder{N},
    energy_distribution::InitialEnergyDistribution{N},
) where {N}
    builder.energy_distribution = energy_distribution
    builder
end

function set!(
    builder::SimulationBuilder{N},
    mass_component_type::Type{<:MassComponent};
    kwargs...,
) where {N}
    builder.mass_component_type = (mass_component_type, kwargs)
    builder
end

function set!(
    builder::SimulationBuilder{N},
    energy_component_type::Type{<:EnergyComponent};
    kwargs...,
) where {N}
    builder.energy_component_type = (energy_component_type, kwargs)
    builder
end

function set!(
    builder::SimulationBuilder{N},
    diffusion_component_type::Type{<:DiffusionComponent};
    kwargs...,
) where {N}
    builder.diffusion_component_type = (diffusion_component_type, kwargs)
    builder
end

function set!(
    builder::SimulationBuilder{N},
    boundary_component_type::Type{<:BoundaryComponent{N}};
    kwargs...,
) where {N}
    builder.boundary_component_type = (boundary_component_type, kwargs)
    builder
end

function set!(
    builder::SimulationBuilder{N},
    stepper_type::Type{<:Stepper{N}};
    kwargs...,
) where {N}
    builder.stepper_type = (stepper_type, kwargs)
    builder
end

function build(builder::SimulationBuilder{N}) where {N}
    mass_component, positions = builder.mass_component_type[1](
        builder.kernel,
        builder.mass_distribution;
        builder.mass_component_type[2]...,
    )
    mass_densities = get_mass_densities(mass_component)
    energy_component = builder.energy_component_type[1](
        positions,
        mass_densities,
        builder.energy_distribution,
        builder.eos;
        builder.energy_component_type[2]...,
    )
    diffusion_component =
        builder.diffusion_component_type[1](builder.diffusion_component_type[2]...)
    velocities =
        Velocities(positions, mass_densities, builder.velocity_distribution)
    boundary_component =
        builder.boundary_component_type[1](builder.boundary_component_type[2]...)
    stepper = builder.stepper_type[1](
        positions,
        velocities,
        mass_component,
        energy_component,
        diffusion_component,
        boundary_component,
        builder.eos;
        builder.stepper_type[2]...,
    )
    stepper
end


using ModernGL, GeometryTypes, GLAbstraction, GLWindow, Images, FileIO

# Load our textures. See "downloads.jl" to get the images.
kitten = load(Pkg.dir("GLAbstraction", "tutorials", "images", "kitten.png"))
puppy  = load(Pkg.dir("GLAbstraction", "tutorials", "images", "puppy.png"))

windowhints = [(GLFW.DEPTH_BITS, 32), (GLFW.STENCIL_BITS, 8)]
window = create_glcontext("Depth and stencils 2",
                          resolution=(600,600),
                          windowhints=windowhints)

vao = glGenVertexArrays()
glBindVertexArray(vao)

# The cube. This could be more efficiently represented using indexes,
# but the tutorial doesn't do it that way.
vertex_positions = Vec3f0[
    # The cube
    (-0.5f0, -0.5f0, -0.5f0),
    ( 0.5f0, -0.5f0, -0.5f0),
    ( 0.5f0,  0.5f0, -0.5f0),
    ( 0.5f0,  0.5f0, -0.5f0),
    (-0.5f0,  0.5f0, -0.5f0),
    (-0.5f0, -0.5f0, -0.5f0),

    (-0.5f0, -0.5f0,  0.5f0),
    ( 0.5f0, -0.5f0,  0.5f0),
    ( 0.5f0,  0.5f0,  0.5f0),
    ( 0.5f0,  0.5f0,  0.5f0),
    (-0.5f0,  0.5f0,  0.5f0),
    (-0.5f0, -0.5f0,  0.5f0),

    (-0.5f0,  0.5f0,  0.5f0),
    (-0.5f0,  0.5f0, -0.5f0),
    (-0.5f0, -0.5f0, -0.5f0),
    (-0.5f0, -0.5f0, -0.5f0),
    (-0.5f0, -0.5f0,  0.5f0),
    (-0.5f0,  0.5f0,  0.5f0),

    ( 0.5f0,  0.5f0,  0.5f0),
    ( 0.5f0,  0.5f0, -0.5f0),
    ( 0.5f0, -0.5f0, -0.5f0),
    ( 0.5f0, -0.5f0, -0.5f0),
    ( 0.5f0, -0.5f0,  0.5f0),
    ( 0.5f0,  0.5f0,  0.5f0),

    (-0.5f0, -0.5f0, -0.5f0),
    ( 0.5f0, -0.5f0, -0.5f0),
    ( 0.5f0, -0.5f0,  0.5f0),
    ( 0.5f0, -0.5f0,  0.5f0),
    (-0.5f0, -0.5f0,  0.5f0),
    (-0.5f0, -0.5f0, -0.5f0),

    (-0.5f0,  0.5f0, -0.5f0),
    ( 0.5f0,  0.5f0, -0.5f0),
    ( 0.5f0,  0.5f0,  0.5f0),
    ( 0.5f0,  0.5f0,  0.5f0),
    (-0.5f0,  0.5f0,  0.5f0),
    (-0.5f0,  0.5f0, -0.5f0)]

floor_positions = Vec3f0[
    # The floor
    (-1.0f0, -1.0f0, -0.5f0),
    ( 1.0f0, -1.0f0, -0.5f0),
    ( 1.0f0,  1.0f0, -0.5f0),
    ( 1.0f0,  1.0f0, -0.5f0),
    (-1.0f0,  1.0f0, -0.5f0),
    (-1.0f0, -1.0f0, -0.5f0)
]

vertex_texcoords = Vec2f0[
                          # The cube
                          (0.0f0, 0.0f0),
                          (1.0f0, 0.0f0),
                          (1.0f0, 1.0f0),
                          (1.0f0, 1.0f0),
                          (0.0f0, 1.0f0),
                          (0.0f0, 0.0f0),

                          (0.0f0, 0.0f0),
                          (1.0f0, 0.0f0),
                          (1.0f0, 1.0f0),
                          (1.0f0, 1.0f0),
                          (0.0f0, 1.0f0),
                          (0.0f0, 0.0f0),

                          (1.0f0, 0.0f0),
                          (1.0f0, 1.0f0),
                          (0.0f0, 1.0f0),
                          (0.0f0, 1.0f0),
                          (0.0f0, 0.0f0),
                          (1.0f0, 0.0f0),

                          (1.0f0, 0.0f0),
                          (1.0f0, 1.0f0),
                          (0.0f0, 1.0f0),
                          (0.0f0, 1.0f0),
                          (0.0f0, 0.0f0),
                          (1.0f0, 0.0f0),

                          (0.0f0, 1.0f0),
                          (1.0f0, 1.0f0),
                          (1.0f0, 0.0f0),
                          (1.0f0, 0.0f0),
                          (0.0f0, 0.0f0),
                          (0.0f0, 1.0f0),

                          (0.0f0, 1.0f0),
                          (1.0f0, 1.0f0),
                          (1.0f0, 0.0f0),
                          (1.0f0, 0.0f0),
                          (0.0f0, 0.0f0),
                          (0.0f0, 1.0f0)]

floor_texcoords = Vec2f0[
                          # The floor
                          (0.0f0, 0.0f0),
                          (1.0f0, 0.0f0),
                          (1.0f0, 1.0f0),
                          (1.0f0, 1.0f0),
                          (0.0f0, 1.0f0),
                          (0.0f0, 0.0f0)]

vertex_colors = fill(Vec3f0(1,1,1), 36)
floor_colors = fill(Vec3f0(0,0,0),6)

vertex_shader = vert"""
#version 150

in vec3 position;
in vec3 color;
in vec2 texcoord;

out vec3 Color;
out vec2 Texcoord;

uniform vec3 overrideColor;
uniform mat4 model;
uniform mat4 view;
uniform mat4 proj;

void main()
{
    Color = overrideColor * color;
    Texcoord = texcoord;
    gl_Position = proj * view * model * vec4(position, 1.0);
}
"""
fragment_shader = load(joinpath(dirname(@__FILE__), "shaders", "puppykitten_color.frag"))

model1 = eye(Mat{4,4,Float32})
model2 = translationmatrix_z(-1f0) * scalematrix(Vec3f0(1,1,-1))
view = lookat(Vec3f0((2.5, 2.5, 2)), Vec3f0((0, 0, 0)), Vec3f0((0, 0, 1)))
proj = perspectiveprojection(Float32, 45, 600/600, 1, 10)

## Now render the distinct objects. Rather than always using std_renderobject,
## here we control the settings manually.
# The cube
bufferdict_cube = Dict(:position=>GLBuffer(vertex_positions),
                       :texcoord=>GLBuffer(vertex_texcoords),
                       :color=>GLBuffer(vertex_colors),
                       :texKitten=>Texture(data(kitten)),
                       :texPuppy=>Texture(data(puppy)),
                       :overrideColor=>Vec3f0((1,1,1)),
                       :model=>model1,
                       :view=>view,
                       :proj=>proj)

ro_cube = std_renderobject(bufferdict_cube,
                           LazyShader(vertex_shader, fragment_shader))
prerender!(ro_cube,
    glDisable, GL_STENCIL_TEST
)
# The floor. This is drawn without writing to the depth buffer, but we
# write stencil values.
bufferdict_floor = Dict(:position=>GLBuffer(floor_positions),
                        :texcoord=>GLBuffer(floor_texcoords),
                        :color=>GLBuffer(floor_colors),
                        :texKitten=>Texture(data(kitten)), # with different shaders, wouldn't need these here
                        :texPuppy=>Texture(data(puppy)),
                        :overrideColor=>Vec3f0((1,1,1)),
                        :model=>model1,
                        :view=>view,
                        :proj=>proj)

ro_floor = RenderObject(bufferdict_floor,
                        LazyShader(vertex_shader, fragment_shader))
prerender!(ro_floor,
           glDepthMask, GL_FALSE,                  # don't write to depth buffer
           glEnable, GL_STENCIL_TEST,              # use stencils
           glStencilMask, 0xff,                    # do write to stencil buffer
           glStencilFunc, GL_ALWAYS, 1, 0xff,      # all pass
           glStencilOp, GL_KEEP, GL_KEEP, GL_REPLACE,  # replace stencil value
           glClear, GL_STENCIL_BUFFER_BIT)         # start with empty buffer
postrender!(ro_floor,
            render, ro_floor.vertexarray, GL_TRIANGLES)

# The cube reflection
bufferdict_refl = Dict(:position=>GLBuffer(vertex_positions),
                       :texcoord=>GLBuffer(vertex_texcoords),
                       :color=>GLBuffer(vertex_colors),
                       :texKitten=>Texture(data(kitten)),
                       :texPuppy=>Texture(data(puppy)),
                       :overrideColor=>Vec3f0((0.3,0.3,0.3)),
                       :model=>model2,
                       :view=>view,
                       :proj=>proj)

ro_refl = RenderObject(bufferdict_refl,
                       LazyShader(vertex_shader, fragment_shader))
prerender!(ro_refl,
           glStencilFunc, GL_EQUAL, 1, 0xff,
           glStencilMask, 0x00)
postrender!(ro_refl,
            render, ro_refl.vertexarray, GL_TRIANGLES)

glClearColor(1,1,1,1) # make the background white, so we can see the floor
glClearStencil(0)     # clear the stencil buffer with 0

while !GLFW.WindowShouldClose(window)
    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
    render([ro_cube, ro_floor, ro_refl])

    GLFW.SwapBuffers(window)
    GLFW.PollEvents()
    if GLFW.GetKey(window, GLFW.KEY_ESCAPE) == GLFW.PRESS
        GLFW.SetWindowShouldClose(window, true)
    end
end

GLFW.DestroyWindow(window)  # needed if you're running this from the REPL


# MIT License
#
# Copyright (c) 2018 Martin Biel
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""
    DeterministicEquivalent

Deterministic equivalent memory structure. Stochastic program is stored as one large optimization problem. Supported by any standard `AbstractOptimizer`.

"""
struct DeterministicEquivalent{N, M, S <: NTuple{M, Scenarios}} <: AbstractStochasticStructure{N}
    decisions::Decisions{N}
    scenarios::S
    model::JuMP.Model
    proxy::NTuple{N,JuMP.Model}

    function DeterministicEquivalent(decisions::Decisions{N}, scenarios::NTuple{M, Scenarios}) where {N, M}
        M == N - 1 || error("Inconsistent number of stages $N and number of scenario types $M")
        proxy = ntuple(Val{N}()) do _
            Model()
        end
        S = typeof(scenarios)
        return new{N,M,S}(decisions, scenarios, Model(), proxy)
    end
end

function StochasticStructure(decisions::Decisions{N}, scenario_types::ScenarioTypes{M}, ::Deterministic) where {N, M}
    scenarios = ntuple(Val(M)) do i
        Vector{scenario_types[i]}()
    end
    return DeterministicEquivalent(decisions, scenarios)
end

function StochasticStructure(decisions::Decisions{N}, scenarios::NTuple{M, Vector{<:AbstractScenario}}, ::Deterministic) where {N, M}
    return DeterministicEquivalent(decisions, scenarios)
end

# Base overloads #
# ========================== #
function Base.print(io::IO, structure::DeterministicEquivalent)
    print(io, "Deterministic equivalent problem\n")
    print(io, structure.model)
end
# ========================== #

# MOI #
# ========================== #
function MOI.get(structure::DeterministicEquivalent, attr::MOI.AbstractModelAttribute)
    return MOI.get(backend(structure.model), attr)
end
function MOI.get(structure::DeterministicEquivalent, attr::MOI.AbstractVariableAttribute, index::MOI.VariableIndex)
    return MOI.get(backend(structure.model), attr, index)
end
function MOI.get(structure::DeterministicEquivalent, attr::MOI.AbstractConstraintAttribute, ci::CI)
    return MOI.get(backend(structure.model), attr, ci)
end
function MOI.get(structure::DeterministicEquivalent, attr::MOI.AbstractConstraintAttribute, ci::CI{F,S}) where {F <: SingleDecision, S}
    con_ref = ConstraintRef(structure.model, ci)
    return MOI.get(structure.model, attr, con_ref)
end
function MOI.get(structure::DeterministicEquivalent{N}, attr::ScenarioDependentModelAttribute) where N
    n = num_scenarios(structure, attr.stage)
    1 <= attr.scenario_index <= n || error("Scenario index $attr.scenario_index not in range 1 to $n.")
    if attr.attr isa MOI.ObjectiveFunction
        return get_stage_objective(structure.decisions, attr.stage, attr.scenario_index)[2]
    elseif attr.attr isa MOI.ObjectiveFunctionType
        return typeof(get_stage_objective(structure.decisions, attr.stage, attr.scenario_index)[2])
    elseif attr.attr isa MOI.ObjectiveSense
        return get_stage_objective(structure.decisions, attr.stage, attr.scenario_index)[1]
    elseif attr.attr isa MOI.ObjectiveValue || attr.attr isa MOI.DualObjectiveValue
        return MOIU.eval_variables(get_stage_objective(structure.decisions, attr.stage, attr.scenario_index)[2]) do idx
            return MOI.get(backend(structure.model), MOI.VariablePrimal(), idx)
        end
    else
        # Most attributes are shared with the deterministic equivalent
        return MOI.get(backend(structure.model), attr.attr)
    end
end
function MOI.get(structure::DeterministicEquivalent, attr::ScenarioDependentVariableAttribute, index::MOI.VariableIndex)
    n = num_scenarios(structure, attr.stage)
    1 <= attr.scenario_index <= n || error("Scenario index $attr.scenario_index not in range 1 to $n.")
    mapped_vi = mapped_index(structure, index, attr.scenario_index)
    return MOI.get(backend(structure.model), attr.attr, mapped_vi)
end
function MOI.get(structure::DeterministicEquivalent, attr::ScenarioDependentConstraintAttribute, ci::CI)
    n = num_scenarios(structure, attr.stage)
    1 <= attr.scenario_index <= n || error("Scenario index $attr.scenario_index not in range 1 to $n.")
    mapped_ci = mapped_index(structure, ci, attr.scenario_index)
    return MOI.get(backend(structure.model), attr.attr, mapped_ci)
end
function MOI.get(structure::DeterministicEquivalent, attr::ScenarioDependentConstraintAttribute, ci::CI{F,S}) where {F <: SingleDecision, S}
    n = num_scenarios(structure, attr.stage)
    1 <= attr.scenario_index <= n || error("Scenario index $attr.scenario_index not in range 1 to $n.")
    mapped_vi = mapped_index(structure, MOI.VariableIndex(ci.value), attr.scenario_index)
    mapped_ci = CI{F,S}(mapped_vi.value)
    con_ref = ConstraintRef(structure.model, mapped_ci)
    return MOI.get(structure.model, attr.attr, con_ref)
end
function MOI.get(structure::DeterministicEquivalent, attr::ScenarioDependentConstraintAttribute, ci::CI{F,S}) where {F <: MOI.SingleVariable, S}
    n = num_scenarios(structure, attr.stage)
    1 <= attr.scenario_index <= n || error("Scenario index $attr.scenario_index not in range 1 to $n.")
    mapped_vi = mapped_index(structure, MOI.VariableIndex(ci.value), attr.scenario_index)
    mapped_ci = CI{F,S}(mapped_vi.value)
    return MOI.get(backend(structure.model), attr.attr, mapped_ci)
end

function MOI.set(structure::DeterministicEquivalent, attr::MOI.Silent, flag)
    MOI.set(backend(structure.model), attr, flag)
    return nothing
end
function MOI.set(structure::DeterministicEquivalent{N}, attr::MOI.AbstractModelAttribute, value) where N
    if attr isa MOI.ObjectiveFunction
        # Get full objective+sense
        dep_obj = copy(value)
        obj_sense = objective_sense(structure.model)
        # Update first-stage objective
        set_stage_objective!(structure.decisions, 1, 1, obj_sense, value)
        # Update main objective
        for (i, sub_objective) in enumerate(structure.decisions.stage_objectives[2])
            (sense, func) = sub_objective
            if obj_sense == sense
                dep_obj += probability(structure, 2, i) * func
            else
                dep_obj -= probability(structure, 2, i) * func
            end
        end
        MOI.set(backend(structure.model), attr, dep_obj)
    elseif attr isa MOI.ObjectiveSense
        # Get full objective+sense
        prev_sense, dep_obj = get_stage_objective(structure.decisions, 1, 1)
        # Update first-stage objective
        set_stage_objective!(structure.decisions, 1, 1, value, dep_obj)
        # Update main objective (if necessary)
        if value != prev_sense
            for (i, sub_objective) in enumerate(structure.decisions.stage_objectives[2])
                (sense, func) = sub_objective
                if value == sense
                    dep_obj += probability(structure, 2, i) * func
                else
                    dep_obj -= probability(structure, 2, i) * func
                end
            end
            MOI.set(backend(structure.model), MOI.ObjectiveFunction{typeof(dep_obj)}(), dep_obj)
        end
        MOI.set(backend(structure.model), MOI.ObjectiveSense(), value)
    else
        # Most attributes are shared with the deterministic equivalent
        MOI.set(backend(structure.model), attr, value)
    end
    return nothing
end
function MOI.set(structure::DeterministicEquivalent, attr::MOI.AbstractVariableAttribute,
                 index::MOI.VariableIndex, value)
    MOI.set(backend(structure.model), attr, index, value)
    return nothing
end
function MOI.set(structure::DeterministicEquivalent, attr::MOI.AbstractConstraintAttribute,
                 ci::MOI.ConstraintIndex, value)
    con_ref = ConstraintRef(structure.model, mapped_ci)
    MOI.set(structure.model, attr.attr, con_ref, value)
    return nothing
end
function MOI.set(structure::DeterministicEquivalent{N}, attr::ScenarioDependentModelAttribute, value) where N
    if attr.attr isa MOI.ObjectiveFunction
        # Get full objective+sense
        obj_sense = objective_sense(structure.model)
        dep_obj = objective_function(structure.model)
        # Update subobjective
        (sub_sense, prev_func) = get_stage_objective(structure.model, attr.stage, attr.scenario_index, Val{N}())
        set_stage_objective!(structure.decisions, attr.stage, attr.scenario_index, sub_sense, value)
        sub_obj = jump_function(structure.model, value)
        # Update main objective
        if obj_sense == sub_sense
            dep_obj += probability(structure, attr.stage, attr.scenario_index) * (sub_obj - prev_func)
        else
            dep_obj -= probability(structure, attr.stage, attr.scenario_index) * (sub_obj - prev_func)
        end
        set_objective_function(structure.model, dep_obj)
    elseif attr.attr isa MOI.ObjectiveSense
        # Get current
        (prev_sense, func) = get_stage_objective(structure.model, attr.stage, attr.scenario_index, Val{N}())
        if value == prev_sense
            # Nothing to do
            return nothing
        end
        # Get full objective+sense
        obj_sense = objective_sense(structure.model)
        dep_obj = objective_function(structure.model)
        # Update subobjective sense
        set_stage_objective!(structure.decisions, attr.stage, attr.scenario_index, value, moi_function(func))
        # Update main objective
        if value == obj_sense
            dep_obj += 2 * probability(structure, attr.stage, attr.scenario_index) * func
        else
            dep_obj -= 2 * probability(structure, attr.stage, attr.scenario_index) * func
        end
        set_objective_function(structure.model, dep_obj)
    else
        # Most attributes are shared with the deterministic equivalent
        MOI.set(backend(structure.model), attr.attr, value)
    end
    return nothing
end
function MOI.set(structure::DeterministicEquivalent, attr::ScenarioDependentVariableAttribute,
                 index::MOI.VariableIndex, value)
    n = num_scenarios(structure, attr.stage)
    1 <= attr.scenario_index <= n || error("Scenario index $attr.scenario_index not in range 1 to $n.")
    mapped_vi = mapped_index(structure, index, attr.scenario_index)
    MOI.set(backend(structure.model), attr.attr, mapped_vi, value)
    return nothing
end
function MOI.set(structure::DeterministicEquivalent, attr::ScenarioDependentConstraintAttribute,
                 ci::CI, value)
    n = num_scenarios(structure, attr.stage)
    1 <= attr.scenario_index <= n || error("Scenario index $attr.scenario_index not in range 1 to $n.")
    mapped_ci = mapped_index(structure, ci, attr.scenario_index)
    MOI.set(backend(structure.model), attr.attr, mapped_ci, value)
    return nothing
end
function MOI.set(structure::DeterministicEquivalent, attr::ScenarioDependentConstraintAttribute,
                 ci::CI{F,S}, value) where {F <: SingleDecision, S}
    n = num_scenarios(structure, attr.stage)
    1 <= attr.scenario_index <= n || error("Scenario index $attr.scenario_index not in range 1 to $n.")
    mapped_vi = mapped_index(structure, MOI.VariableIndex(ci.value), attr.scenario_index)
    mapped_ci = CI{F,S}(mapped_vi.value)
    con_ref = ConstraintRef(structure.model, mapped_ci)
    MOI.set(structure.model, attr.attr, con_ref, value)
    return nothing
end
function MOI.set(structure::DeterministicEquivalent, attr::ScenarioDependentConstraintAttribute,
                 ci::CI{F,S}, value) where {F <: MOI.SingleVariable, S}
    n = num_scenarios(structure, attr.stage)
    1 <= attr.scenario_index <= n || error("Scenario index $attr.scenario_index not in range 1 to $n.")
    mapped_vi = mapped_index(structure, MOI.VariableIndex(ci.value), attr.scenario_index)
    mapped_ci = CI{F,S}(mapped_vi.value)
    MOI.set(backend(structure.model), attr.attr, mapped_ci, value)
    return nothing
end

function MOI.is_valid(structure::DeterministicEquivalent, index::MOI.VariableIndex, stage::Integer)
    stage == 1 || error("No scenario index specified.")
    return MOI.is_valid(backend(structure.model), index)
end
function MOI.is_valid(structure::DeterministicEquivalent, index::MOI.VariableIndex, stage::Integer, scenario_index::Integer)
    stage > 1 || error("There are no scenarios in the first stage.")
    n = num_scenarios(structure, stage)
    1 <= scenario_index <= n || error("Scenario index $scenario_index not in range 1 to $n.")
    mapped_vi = mapped_index(structure, index, scenario_index)
    return MOI.is_valid(backend(structure.model), mapped_vi)
end
function MOI.is_valid(structure::DeterministicEquivalent, ci::MOI.ConstraintIndex{F,S}, stage::Integer) where {F, S}
    stage == 1 || error("No scenario index specified.")
    return MOI.is_valid(backend(structure.model), ci)
end
function MOI.is_valid(structure::DeterministicEquivalent, ci::MOI.ConstraintIndex{F,S}, stage::Integer, scenario_index::Integer) where {F, S}
    stage > 1 || error("There are no scenarios in the first stage.")
    n = num_scenarios(structure, stage)
    1 <= scenario_index <= n || error("Scenario index $scenario_index not in range 1 to $n.")
    mapped_ci = mapped_index(structure, ci, scenario_index)
    return MOI.is_valid(backend(structure.model), mapped_ci)
end
function MOI.delete(structure::DeterministicEquivalent, indices::Vector{MOI.VariableIndex}, stage::Integer)
    stage == 1 || error("No scenario index specified.")
    JuMP.delete(structure.model, DecisionRef.(structure.model, indices))
    return nothing
end
function MOI.delete(structure::DeterministicEquivalent{N}, indices::Vector{MOI.VariableIndex}, stage::Integer, scenario_index::Integer) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    stage > 1 || error("There are no scenarios in the first stage.")
    n = num_scenarios(structure, stage)
    1 <= scenario_index <= n || error("Scenario index $scenario_index not in range 1 to $n.")
    mapped_indices = map(indices) do index
        return mapped_index(structure, index, scenario_index)
    end
    JuMP.delete(structure.model, DecisionRef.(structure.model, mapped_indices))
    return nothing
end
function MOI.delete(structure::DeterministicEquivalent, ci::MOI.ConstraintIndex, stage::Integer)
    stage == 1 || error("No scenario index specified.")
    MOI.delete(backend(structure.model), ci)
    return nothing
end
function MOI.delete(structure::DeterministicEquivalent, cis::Vector{<:MOI.ConstraintIndex}, stage::Integer)
    stage == 1 || error("No scenario index specified.")
    MOI.delete(backend(structure.model), cis)
    return nothing
end
function MOI.delete(structure::DeterministicEquivalent{N}, ci::MOI.ConstraintIndex, stage::Integer, scenario_index::Integer) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    stage > 1 || error("There are no scenarios in the first stage.")
    n = num_scenarios(structure, stage)
    1 <= scenario_index <= n || error("Scenario index $scenario_index not in range 1 to $n.")
    mapped_ci = mapped_index(structure, ci, scenario_index)
    MOI.delete(backend(structure.model), mapped_ci)
    return nothing
end
function MOI.delete(structure::DeterministicEquivalent{N}, cis::Vector{<:MOI.ConstraintIndex}, stage::Integer, scenario_index::Integer) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    stage > 1 || error("There are no scenarios in the first stage.")
    n = num_scenarios(structure, stage)
    1 <= scenario_index <= n || error("Scenario index $scenario_index not in range 1 to $n.")
    mapped_cis = map(cis) do ci
        return mapped_index(structure, ci, scenario_index)
    end
    MOI.delete(backend(structure.model), mapped_cis)
    return nothing
end

# JuMP #
# ========================== #
function decision_dispatch(decision_function::Function,
                           structure::DeterministicEquivalent{N},
                           index::MOI.VariableIndex,
                           stage::Integer,
                           args...) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    stage == 1 || error("No scenario index specified.")
    dref = DecisionRef(structure.model, index)
    return decision_function(dref, args...)
end
function scenario_decision_dispatch(decision_function::Function,
                                    structure::DeterministicEquivalent{N},
                                    index::MOI.VariableIndex,
                                    stage::Integer,
                                    scenario_index::Integer,
                                    args...) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    stage > 1 || error("There are no scenarios in the first stage.")
    n = num_scenarios(structure, stage)
    1 <= scenario_index <= n || error("Scenario index $scenario_index not in range 1 to $n.")
    mapped_vi = mapped_index(structure, index, scenario_index)
    dref = DecisionRef(structure.model, mapped_vi)
    return decision_function(dref, args...)
end
function decision_dispatch!(decision_function!::Function,
                            structure::DeterministicEquivalent{N},
                            index::MOI.VariableIndex,
                            stage::Integer,
                            args...) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    stage == 1 || error("No scenario index specified.")
    dref = DecisionRef(structure.model, index)
    decision_function!(dref, args...)
    return nothing
end
function scenario_decision_dispatch!(decision_function!::Function,
                                     structure::DeterministicEquivalent{N},
                                     index::MOI.VariableIndex,
                                     stage::Integer,
                                     scenario_index::Integer,
                                     args...) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    stage > 1 || error("There are no scenarios in the first stage.")
    n = num_scenarios(structure, stage)
    1 <= scenario_index <= n || error("Scenario index $scenario_index not in range 1 to $n.")
    mapped_vi = mapped_index(structure, index, scenario_index)
    dref = DecisionRef(structure.model, mapped_vi)
    decision_function!(dref, args...)
    return nothing
end

function JuMP.fix(structure::DeterministicEquivalent{N}, index::MOI.VariableIndex, stage::Integer, val::Number) where N
    dref = DecisionRef(structure.model, index)
    fix(dref, val)
    return nothing
end
function JuMP.fix(structure::DeterministicEquivalent{N}, index::MOI.VariableIndex, stage::Integer, scenario_index::Integer, val::Number) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    n = num_scenarios(structure, stage)
    1 <= scenario_index <= n || error("Scenario index $scenario_index not in range 1 to $n.")
    # Fix mapped decision
    mapped_vi = mapped_index(structure, index, scenario_index)
    dref = DecisionRef(structure.model, mapped_vi)
    fix(dref, val)
    return nothing
end
function JuMP.unfix(structure::DeterministicEquivalent{N}, index::MOI.VariableIndex, stage::Integer) where N
    dref = DecisionRef(structure.model, index)
    unfix(dref)
    return nothing
end
function JuMP.unfix(structure::DeterministicEquivalent{N}, index::MOI.VariableIndex, stage::Integer, scenario_index::Integer) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    n = num_scenarios(structure, stage)
    1 <= scenario_index <= n || error("Scenario index $scenario_index not in range 1 to $n.")
    # Unfix mapped decision
    mapped_vi = mapped_index(structure, index, scenario_index)
    dref = DecisionRef(structure.model, mapped_vi)
    unfix(dref)
    return nothing
end

function JuMP.set_objective_sense(structure::DeterministicEquivalent, stage::Integer, sense::MOI.OptimizationSense)
    if stage == 1
        # Changing the first-stage sense modifies the whole objective as usual
        MOI.set(structure, MOI.ObjectiveSense(), sense)
    else
        # Every sub-objective in the given stage should be changed
        for scenario_index in 1:num_scenarios(structure, stage)
            attr = ScenarioDependentModelAttribute(stage, scenario_index, MOI.ObjectiveSense())
            MOI.set(structure, attr, sense)
        end
    end
    return nothing
end

function JuMP.objective_function_type(structure::DeterministicEquivalent)
    return jump_function_type(structure.model,
                              MOI.get(structure, MOI.ObjectiveFunctionType()))
end
function JuMP.objective_function_type(structure::DeterministicEquivalent{N}, stage::Integer, scenario_index::Integer) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    stage > 1 || error("There are no scenarios in the first stage.")
    n = num_scenarios(structure, stage)
    1 <= scenario_index <= n || error("Scenario index $scenario_index not in range 1 to $n.")
    attr = ScenarioDependentModelAttribute(stage, scenario_index, MOI.ObjectiveFunctionType())
    return jump_function_type(structure.model,
                              MOI.get(structure, attr))
end

function JuMP.objective_function(structure::DeterministicEquivalent, FunType::Type{<:AbstractJuMPScalar})
    MOIFunType = moi_function_type(FunType)
    func = MOI.get(structure,
                   MOI.ObjectiveFunction{MOIFunType}())::MOIFunType
    return JuMP.jump_function(structure, 1, func)
end
function JuMP.objective_function(structure::DeterministicEquivalent{N}, stage::Integer, FunType::Type{<:AbstractJuMPScalar}) where N
    if stage == 1
        (sense, obj::FunType) = get_stage_objective(structure.model, 1, Val{N}())
        return obj
    else
        return objective_function(structure.proxy, FunType)
    end
end
function JuMP.objective_function(structure::DeterministicEquivalent{N},
                                 stage::Integer,
                                 scenario_index::Integer,
                                 FunType::Type{<:AbstractJuMPScalar}) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    stage > 1 || error("There are no scenarios in the first stage.")
    n = num_scenarios(structure, stage)
    1 <= scenario_index <= n || error("Scenario index $scenario_index not in range 1 to $n.")
    MOIFunType = moi_function_type(FunType)
    attr = ScenarioDependentModelAttribute(stage, scenario_index, MOI.ObjectiveFunction{MOIFunType}())
    func = MOI.get(structure, attr)::MOIFunType
    return JuMP.jump_function(structure, stage, scenario_index, func)
end

function JuMP._moi_optimizer_index(structure::DeterministicEquivalent, index::VI)
    return decision_index(backend(structure.model), index)
end
function JuMP._moi_optimizer_index(structure::DeterministicEquivalent, index::VI, scenario_index::Integer)
    mapped_vi = mapped_index(structure, index, scenario_index)
    return decision_index(backend(structure.model), mapped_vi)
end
function JuMP._moi_optimizer_index(structure::DeterministicEquivalent, ci::CI)
    return decision_index(backend(structure.model), ci)
end
function JuMP._moi_optimizer_index(structure::DeterministicEquivalent, ci::CI{F,S}) where {F <: SingleDecision, S}
    inner = mapped_constraint(structure.decisions, ci)
    inner.value == 0 && error("Constraint $ci not properly mapped.")
    return decision_index(backend(structure.model), inner)
end
function JuMP._moi_optimizer_index(structure::DeterministicEquivalent, ci::CI, scenario_index::Integer)
    mapped_ci = mapped_index(structure, ci, scenario_index)
    return decision_index(backend(structure.model), mapped_ci)
end
function JuMP._moi_optimizer_index(structure::DeterministicEquivalent, ci::CI{F,S}, scenario_index::Integer) where {F <: SingleDecision, S}
    mapped_vi = mapped_index(structure, MOI.VariableIndex(ci.value), scenario_index)
    mapped_ci = CI{F,S}(mapped_vi.value)
    inner = mapped_constraint(structure.decisions, mapped_ci)
    inner.value == 0 && error("Constraint $mapped_ci not properly mapped.")
    return decision_index(backend(structure.model), inner)
end

function JuMP.set_objective_coefficient(structure::DeterministicEquivalent{N}, index::VI, var_stage::Integer, stage::Integer, coeff::Real) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    var_stage <= stage || error("Can only modify coefficient in current stage of decision or subsequent stages from where decision is taken.")
    if var_stage == 1 && stage == 1
        # Use temporary model to apply modification
        obj = get_stage_objective(structure.model, 1, Val{N}())[2]
        moi_obj = moi_function(obj)
        m = Model()
        MOI.set(backend(m), MOI.ObjectiveFunction{typeof(moi_obj)}(), moi_obj)
        dref = DecisionRef(m, index)
        set_objective_coefficient(m, dref, coeff)
        obj = objective_function(m)
        F = moi_function_type(typeof(obj))
        # Modify full objective
        MOI.set(structure, MOI.ObjectiveFunction{F}(), moi_function(obj))
    elseif (var_stage == 1 && stage > 1) || var_stage > 1
        for scenario_index in 1:num_scenarios(structure, stage)
            # Use temporary model to apply modification
            obj = get_stage_objective(structure.model, stage, scenario_index, Val{N}())[2]
            moi_obj = moi_function(obj)
            m = Model()
            MOI.set(backend(m), MOI.ObjectiveFunction{typeof(moi_obj)}(), moi_obj)
            dref = DecisionRef(m, index)
            set_objective_coefficient(m, dref, coeff)
            obj = objective_function(m)
            attr = ScenarioDependentModelAttribute(stage, scenario_index, MOI.ObjectiveFunction{typeof(obj)}())
            MOI.set(structure, attr, obj)
        end
    end
    return nothing
end
function JuMP.set_objective_coefficient(structure::DeterministicEquivalent{N}, index::VI, var_stage::Integer, stage::Integer, scenario_index::Integer, coeff::Real) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    stage > 1 || error("There are no scenarios in the first stage.")
    n = num_scenarios(structure, stage)
    1 <= scenario_index <= n || error("Scenario index $scenario_index not in range 1 to $n.")
    if var_stage == 1
        # Use temporary model to apply modification
        obj = get_stage_objective(structure.model, stage, scenario_index, Val{N}())[2]
        moi_obj = moi_function(obj)
        m = Model()
        MOI.set(backend(m), MOI.ObjectiveFunction{typeof(moi_obj)}(), moi_obj)
        dref = DecisionRef(m, index)
        set_objective_coefficient(m, dref, coeff)
        obj = objective_function(m)
        F = moi_function_type(typeof(obj))
        attr = ScenarioDependentModelAttribute(stage, scenario_index, MOI.ObjectiveFunction{F}())
        MOI.set(structure, attr, moi_function(obj))
    else
        # Use temporary model to apply modification
        sense, obj = get_stage_objective(structure.model, stage, scenario_index, Val{N}())
        moi_obj = moi_function(obj)
        m = Model()
        MOI.set(backend(m), MOI.ObjectiveFunction{typeof(moi_obj)}(), moi_obj)
        mapped_vi = mapped_index(structure, index, scenario_index)
        dref = DecisionRef(m, mapped_vi)
        set_objective_coefficient(m, dref, coeff)
        obj = objective_function(m)
        set_stage_objective!(structure.decisions, stage, scenario_index, sense, moi_function(obj))
        # Set coefficient of mapped second-stage variable
        dref = DecisionRef(structure.model, mapped_vi)
        set_objective_coefficient(structure.model, dref, probability(structure, stage, scenario_index) * coeff)
    end
    return nothing
end

function JuMP.set_normalized_coefficient(structure::DeterministicEquivalent,
                                         ci::CI{F,S},
                                         index::VI,
                                         value) where {T, F <: Union{AffineDecisionFunction{T}, QuadraticDecisionFunction{T}}, S}
    MOI.modify(backend(structure.model), ci,
               DecisionCoefficientChange(index, convert(T, value)))
    return nothing
end
function JuMP.set_normalized_coefficient(structure::DeterministicEquivalent{N},
                                         ci::CI{F,S},
                                         index::VI,
                                         var_stage::Integer,
                                         stage::Integer,
                                         scenario_index::Integer,
                                         value) where {N, T, F <: Union{AffineDecisionFunction{T}, QuadraticDecisionFunction{T}}, S}
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    stage > 1 || error("There are no scenarios in the first stage.")
    n = num_scenarios(structure, stage)
    1 <= scenario_index <= n || error("Scenario index $scenario_index not in range 1 to $n.")
    mapped_vi = if var_stage == 1
        mapped_vi = index
    else
        mapped_vi = mapped_index(structure, index, scenario_index)
    end
    mapped_ci = mapped_index(structure, ci, scenario_index)
    MOI.modify(backend(structure.model), mapped_ci,
               DecisionCoefficientChange(mapped_vi, convert(T, value)))
    return nothing
end

function JuMP.normalized_coefficient(structure::DeterministicEquivalent{N},
                                     ci::CI{F,S},
                                     index::VI,
                                     var_stage::Integer,
                                     stage::Integer,
                                     scenario_index::Integer) where {N, T, F <: Union{AffineDecisionFunction{T}, QuadraticDecisionFunction{T}}, S}
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    stage > 1 || error("There are no scenarios in the first stage.")
    n = num_scenarios(structure, stage)
    1 <= scenario_index <= n || error("Scenario index $scenario_index not in range 1 to $n.")
    mapped_ci = mapped_index(structure, ci, scenario_index)
    f = MOI.get(structure, MOI.ConstraintFunction(), mapped_ci)::F
    dref = if var_stage == 1
        dref = DecisionRef(structure.model, index)
    else
        mapped_vi = mapped_index(structure, index, scenario_index)
        dref = DecisionRef(structure.model, mapped_vi)
    end
    return JuMP._affine_coefficient(jump_function(structure.model, f), dref)
end

function JuMP.set_normalized_rhs(structure::DeterministicEquivalent,
                                 ci::CI{F,S},
                                 value) where {T,
                                               F <: Union{AffineDecisionFunction{T}, QuadraticDecisionFunction{T}},
                                               S <: MOIU.ScalarLinearSet{T}}
    MOI.set(backend(structure.model), MOI.ConstraintSet(), ci,
            S(convert(T, value)))
    return nothing
end
function JuMP.set_normalized_rhs(structure::DeterministicEquivalent{N},
                                 ci::CI{F,S},
                                 stage::Integer,
                                 scenario_index::Integer,
                                 value) where {N,
                                               T,
                                               F <: Union{AffineDecisionFunction{T}, QuadraticDecisionFunction{T}},
                                               S <: MOIU.ScalarLinearSet{T}}
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    stage > 1 || error("There are no scenarios in the first stage.")
    n = num_scenarios(structure, stage)
    1 <= scenario_index <= n || error("Scenario index $scenario_index not in range 1 to $n.")
    mapped_ci = mapped_index(structure, ci, scenario_index)
    MOI.set(backend(structure.model), MOI.ConstraintSet(), mapped_ci,
            S(convert(T, value)))
    return nothing
end

function DecisionRef(structure::DeterministicEquivalent, index::VI)
    return DecisionRef(structure.model, index)
end
function DecisionRef(structure::DeterministicEquivalent, index::VI, stage::Integer, scenario_index::Integer)
    mapped_vi = mapped_index(structure, index, scenario_index)
    return DecisionRef(structure.model, mapped_vi)
end
function DecisionRef(structure::DeterministicEquivalent{N}, index::VI, at_stage::Integer, stage::Integer, scenario_index::Integer) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    at_stage > 1 || error("There are no scenarios in the first stage.")
    n = num_scenarios(structure, at_stage)
    1 <= scenario_index <= n || error("Scenario index $scenario_index not in range 1 to $n.")
    return DecisionRef(structure.model, index)
end

function JuMP.jump_function(structure::DeterministicEquivalent{N},
                            stage::Integer,
                            f::MOI.AbstractFunction) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    if stage == 1
        return JuMP.jump_function(structure.model, f)
    else
        return JuMP.jump_function(structure.proxy[stage], f)
    end
end
function JuMP.jump_function(structure::DeterministicEquivalent{N},
                            stage::Integer,
                            scenario_index::Integer,
                            f::MOI.AbstractFunction) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    return JuMP.jump_function(structure.model, f)
end

function JuMP.relax_integrality(structure::DeterministicEquivalent)
    unrelax = relax_decision_integrality(structure.model)
    return unrelax
end

# Getters #
# ========================== #
function structure_name(structure::DeterministicEquivalent)
    return "Deterministic equivalent"
end
function scenario_types(structure::DeterministicEquivalent{N}) where N
    return ntuple(Val{N-1}()) do i
        eltype(structure.scenarios[i])
    end
end
function proxy(structure::DeterministicEquivalent{N}, stage::Integer) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    return structure.proxy[stage]
end
function decision(structure::DeterministicEquivalent{N}, index::MOI.VariableIndex, stage::Integer) where N
    stage == 1 || error("No scenario index specified.")
    return decision(structure.decisions, stage, index)
end
function decision(structure::DeterministicEquivalent{N}, index::MOI.VariableIndex, stage::Integer, scenario_index::Integer) where N
    stage > 1 || error("There are no scenarios in the first stage.")
    n = num_scenarios(structure, stage)
    1 <= scenario_index <= n || error("Scenario index $scenario_index not in range 1 to $n.")
    mapped_vi = mapped_index(structure, index, scenario_index)
    return decision(structure.decisions, stage, mapped_vi)
end
function scenario(structure::DeterministicEquivalent{N}, stage::Integer, scenario_index::Integer) where N
    return structure.scenarios[stage-1][scenario_index]
end
function scenarios(structure::DeterministicEquivalent{N}, stage::Integer) where N
    return structure.scenarios[stage-1]
end
function subproblem(structure::DeterministicEquivalent, stage::Integer, scenario_index::Integer)
    error("The determinstic equivalent is not decomposed into subproblems.")
end
function subproblems(structure::DeterministicEquivalent, stage::Integer)
    error("The determinstic equivalent is not decomposed into subproblems.")
end
function num_subproblems(structure::DeterministicEquivalent, stage::Integer)
    return 0
end
function deferred(structure::DeterministicEquivalent)
    return num_variables(structure.model) == 0
end
# ========================== #

# Setters
# ========================== #
function update_known_decisions!(structure::DeterministicEquivalent)
    update_known_decisions!(structure.model)
    return nothing
end
function update_known_decisions!(structure::DeterministicEquivalent, stage::Integer, scenario_index::Integer)
    stage > 1 || error("There are no scenarios in the first stage.")
    n = num_scenarios(structure, stage)
    1 <= scenario_index <= n || error("Scenario index $scenario_index not in range 1 to $n.")
    update_known_decisions!(structure.model)
    return nothing
end

function add_scenario!(structure::DeterministicEquivalent, stage::Integer, scenario::AbstractScenario)
    push!(scenarios(structure, stage), scenario)
    return nothing
end
function add_worker_scenario!(structure::DeterministicEquivalent, stage::Integer, scenario::AbstractScenario, w::Integer)
    add_scenario!(structure, scenario, stage)
    return nothing
end
function add_scenario!(scenariogenerator::Function, structure::DeterministicEquivalent, stage::Integer)
    add_scenario!(structure, stage, scenariogenerator())
    return nothing
end
function add_worker_scenario!(scenariogenerator::Function, structure::DeterministicEquivalent, stage::Integer, w::Integer)
    add_scenario!(scenariogenerator, structure, stage)
    return nothing
end
function add_scenarios!(structure::DeterministicEquivalent, stage::Integer, _scenarios::Vector{<:AbstractScenario})
    append!(scenarios(structure, stage), _scenarios)
    return nothing
end
function add_worker_scenarios!(structure::DeterministicEquivalent, stage::Integer, scenarios::Vector{<:AbstractScenario}, w::Integer)
    add_scenarios!(structure, scenarios, stasge)
    return nothing
end
function add_scenarios!(scenariogenerator::Function, structure::DeterministicEquivalent, stage::Integer, n::Integer)
    for i = 1:n
        add_scenario!(structure, stage) do
            return scenariogenerator()
        end
    end
    return nothing
end
function add_worker_scenarios!(scenariogenerator::Function, structure::DeterministicEquivalent, stage::Integer, n::Integer, w::Integer)
    add_scenarios!(scenariogenerator, structure, n, stage)
    return nothing
end
function sample!(structure::DeterministicEquivalent, stage::Integer, sampler::AbstractSampler, n::Integer)
    sample!(scenarios(structure, stage), sampler, n)
    return nothing
end
# ========================== #

# Indices
# ========================== #
function mapped_index(structure::DeterministicEquivalent{2}, index::MOI.VariableIndex, scenario_index::Integer)
    # The initial number of first-stage decisions is always given by
    num_first_stage_decisions = MOI.get(structure.proxy[1], MOI.NumberOfConstraints{MOI.SingleVariable,SingleDecisionSet{Float64}}())
    # Calculate offset from first-stage auxilliary variables (first-stage decisions are included in second-stage proxy, so deduct them)
    first_stage_offset = MOI.get(structure.proxy[1], MOI.NumberOfVariables()) - num_first_stage_decisions
    # Calculate offset from extra counts of first-stage decisions from second-stage proxy
    first_stage_decision_offset = -(scenario_index - 1) * num_first_stage_decisions
    # Calculate offset from second-stage variables
    scenario_offset = (scenario_index - 1) * MOI.get(structure.proxy[2], MOI.NumberOfVariables())
    return MOI.VariableIndex(index.value + first_stage_offset + first_stage_decision_offset + scenario_offset)
end
function mapped_index(structure::DeterministicEquivalent{2}, ci::CI{F,S}, scenario_index::Integer) where {F,S}
    first_stage_offset = MOI.get(structure.proxy[1], MOI.NumberOfConstraints{F,S}())
    scenario_offset = (scenario_index - 1) * MOI.get(structure.proxy[2], MOI.NumberOfConstraints{F,S}())
    return CI{F,S}(ci.value + first_stage_offset + scenario_offset)
end