text:
#jl #! format: off
# One Machine against Infinite Bus (OMIB) simulation with [PowerSimulationsDynamics.jl](https://github.com/NREL-SIIP/PowerSimulationsDynamics.jl)

# **Originally Contributed by**: Rodrigo Henriquez and José Daniel Lara

# # Introduction

# This tutorial will introduce you to the functionality of `PowerSimulationsDynamics`
# for running power system dynamic simulations.

# This tutorial presents a simulation of a two-bus system with an infinite bus
# (represented as a voltage source behind an impedance) at bus 1, and a classic
# machine on bus 2. The perturbation will be the trip of one of the two circuits
# (doubling its resistance and impedance) of the line that connects both buses.

# ## Dependencies
using SIIPExamples #hide
using PowerSimulationsDynamics
PSID = PowerSimulationsDynamics
using PowerSystems
using Sundials
using Plots
gr()

# `PowerSystems` (abbreviated with `PSY`) is used to properly define the data structure and establish an equilibrium
# point initial condition with a power flow routine, while `Sundials` is
# used to solve the problem defined in `PowerSimulationsDynamics`.

# ## Load the system
# _The following command requires that you have executed the
# [dynamic systems data example](https://nbviewer.jupyter.org/github/NREL-SIIP/SIIPExamples.jl/blob/master/notebook/2_PowerSystems_examples/09_loading_dynamic_systems_data.ipynb)
# previously to generate the json file._
file_dir = joinpath(
    dirname(dirname(pathof(SIIPExamples))),
    "script",
    "4_PowerSimulationsDynamics_examples",
    "Data",
)
omib_sys = System(joinpath(file_dir, "omib_sys.json"))

# ## Build the simulation and initialize the problem

# The next step is to create the simulation structure. This will create the indexing
# of our system that will be used to formulate the differential-algebraic system of
# equations. To do so, it is required to specify the perturbation that will occur in
# the system. `PowerSimulationsDynamics` supports three types of perturbations:

# - Network Switch: Change in the Y-bus values.
# - Branch Trip: Disconnects a line from the system.
# - Change in Reference Parameter

# Here, we will use a Branch Trip perturbation, that is modeled by modifying the
# specifying which line we want to trip. In this case we disconnect one of the lines
# that connects BUS 1 and BUS 2, named "BUS 1-BUS 2-i_1".

# With this, we are ready to create our simulation structure:
time_span = (0.0, 30.0)
perturbation_trip = BranchTrip(1.0, "BUS 1-BUS 2-i_1")
sim = PSID.Simulation(PSID.ImplicitModel, omib_sys, pwd(), time_span, perturbation_trip)

# This will automatically initialize the system by running a power flow
# and update `V_ref`, `P_ref` and hence `eq_p` (the internal voltage) to match the
# solution of the power flow. It will also initialize the states in the equilibrium,
# which can be printed with:
print_device_states(sim)

# To examine the calculated initial conditions, we can export them into a dictionary:
x0_init = PSID.get_initial_conditions(sim)

# ## Run the Simulation

# Finally, to run the simulation we simply use:
PSID.execute!(
    sim, #simulation structure
    IDA(), #Sundials DAE Solver
    dtmax = 0.02,
); #Arguments: Maximum timestep allowed

# In some cases, the dynamic time step used for the simulation may fail. In such case, the
# keyword argument `dtmax` can be used to limit the maximum time step allowed for the simulation.

# ## Exploring the solution

# `PowerSimulationsDynamics` has two functions to obtain different
# states of the solution:
#  - `get_state_series(sim, ("generator-102-1", :δ))`: can be used to obtain the solution as
# a tuple of time and the required state. In this case, we are obtaining the rotor angle `:δ`
# of the generator named `"generator-102-1"`.

angle = get_state_series(sim, ("generator-102-1", :δ));
Plots.plot(angle, xlabel = "time", ylabel = "rotor angle [rad]", label = "rotor angle")

# - `get_voltagemag_series(sim, 102)`: can be used to obtain the voltage magnitude as a
# tuple of time and voltage. In this case, we are obtaining the voltage magnitude at bus 102
# (where the generator is located).

volt = get_voltage_magnitude_series(sim, 102);
Plots.plot(volt, xlabel = "time", ylabel = "Voltage [pu]", label = "V_2")

# ## Optional: Small Signal Analysis

# `PowerSimulationsDynamics` uses automatic differentiation to compute the reduced Jacobian
# of the system for the differential states. This can be used to analyze the local stability
# of the linearized system. We need to re-initialize our simulation:

sim2 = PSID.Simulation(PSID.ImplicitModel, omib_sys, pwd(), time_span, perturbation_trip)

small_sig = small_signal_analysis(sim2)

# The `small_sig` result can report the reduced jacobian for ``\delta`` and ``\omega``,

small_sig.reduced_jacobian

# and can also be used to report the eigenvalues of the reduced linearized system:

small_sig.eigenvalues


text:
"""
# Usage

    solve(order, α, T, N, X, i, κ, v, fx0, fgz, f0t, flt, ADV_DIF())

### References

```tex
C.P. Li, R.F. Wu, H.F.  Ding. High-order approximation to Caputo derivatives and Caputo-type advection-diffusion equations (I). Communications in Applied and Industrial Mathematics, 2014, 6(2), e-536: 1-32.  DOI: 10.1685/journal.caim.536.
J.X. Cao, C.P. Li, Y.Q. Chen. High-order approximation to Caputo derivatives and Caputo-type advection-diffusion equations (II). Fractional Calculus and Applied Analysis, 2015, 18(3): 735-761.
H.F. Li, J.X. Cao, C.P. Li. High-order approximation to Caputo derivatives and Caputo-type advection-diffusion equations (III).submitted.
```
"""
struct ADV_DIF <: FractionalDiffEqAlgorithm end

function solve(order::Integer, α, T, N, X, i, κ, v, fx0, fgz, f0t, flt, ::ADV_DIF)
    h = X/i
    gm = T/N
    miu = h^2*gm^(-α)/gamma(1-α)
    u = zeros(i-1, N+1)

    for i2=1:i-1
        u[i2, 1]=fx0(X*i2/i)
    end

    f = zeros(i-1, N)
    for i5=1:N
        for i6=1:i-1
            q1=[T*i5/N X*i6/i α]
            f[i6, i5]=fgz(q1)
        end
    end

    H = zeros(i-1, N)
    for i7 = 1:N
        q1 = [i7/N α]
        q2 = [T*(i7)/N α]
        H[1, i7] = (κ+v*h/2)*f0t(q1)
        H[i-1, i7] = (κ-v*h/2)*flt(q2)
    end

    for n=1:order-2
        g1 = ones(1, n)
        g2=g1*wj(n, n+1, α)
        g = zeros(n, 1)
        for i8=1:n
            g[i8, 1]=g2[1, i8]
        end
        u1=zeros(i-1, n)
        for i9=1:n
            u1[:, i9] = u[:, i9]
        end
        H1 = -miu*u1*g + h^2*f[:, n]+H[:, n]

        A = zeros(i-1, i-1)
        for i10=1:i-1
            A[i10, i10]=miu*g2[1, n+1]+2*κ
            if i10>1
                A[i10, i10-1]=-(κ+v*h/2)
            end
            if i10<i-1
                A[i10, i10+1]=-κ+v*h/2
            end
        end
        u[:, n+1]=zg(A, H1)
    end
    for n=order-1:N
        g1 = ones(1, n)
        g2 = g1*wj(n, order, α)
        g=zeros(n,1)
        for i8=1:n
            g[i8, 1]=g2[1, i8]
        end
        u1=zeros(i-1, n)
        for i9=1:n
            u1[:,i9] = u[:, i9]
        end
        H1=-miu*(u1*g) + h^2*f[:, n] + H[:, n]

        A=zeros(i-1, i-1)
        for i10=1:i-1
            A[i10, i10] = miu*g2[1, n+1] + 2*κ
            if i10 !== 1
                A[i10, i10-1] = -(κ+v*h/2)
            end
            if i10<i-1
                A[i10, i10+1]=-κ+v*h/2
            end
        end
        u[:, n+1] = zg(A, H1)
    end
    return u
end

function wj(n::Int, r::Int, a)
    A = zeros(n, n+1)
    for iw=1:r-2
        for jw=1:iw+1
        A[iw, jw] = w(iw+1-jw, iw+1, iw, n, a)
        end
    end
    for iw=r-1:n
        for jw=iw-r+2:iw+1
        A[iw, jw] = w(iw+1-jw, r, iw, n, a)
        end
    end
    return A
end

function zg(A, H)
    i = length(H)
    r = zeros(i)
    l = zeros(i)
    y = zeros(i)
    U = zeros(i)
    for t=1:i
        if t==1
            r[t] = A[1, 1]
        else
            l[t] = A[t, t-1]/r[t-1]
            r[t] = A[t, t] - l[t]*A[t-1, t]
        end
    end

    for p=1:i
        if p == 1
            y[p] = H[1, 1]
        else
            y[p] = H[p, 1] - l[p]*y[p-1]
        end
    end
    for k=0:i-1
       if k==0
        U[i] = y[i]/r[i]
       else
        U[i-k]=(y[i-k] - A[i-k, i-k+1]*U[i-k+1])/r[i-k]
       end
    end
    return U
end

function w(i::Int, r, j, n, a)
    ar=ones(r-1)
    br=ones(r-1)
    for lj=1:r-2
        jj=r-lj-1
        kj=r-2
        tj=i-1
        aj=collect(Int, 0:tj)
        bj=collect(Int, tj+2:kj+1)
        cj=[aj; bj]
        dj=collect(Int, -1:tj-1)
        ej=collect(Int, tj+1:kj)
        fj=[dj; ej]
        yj = binomial.(cj, jj)
        pj = binomial.(fj, jj)
        sj=1
        tj=1
        for m=1:jj
            sj = sj.*yj[:, m]
            tj = tj.*pj[:, m]
            ar[lj] = sum(sj)
            br[lj] = sum(tj)
        end
    end
    s=0
    for l=1:r-1
        k=1
        for m1=1:l
            k=k/(m1-a)
        end
        k = k*factorial(l)
        k = k*(ar[l]*((n-j)^(l-a))-br[l]*((n-j+1)^(l-a)))
        s = s+k
    end
    s = s*((-1)^(i+1))/(factorial(i)*factorial(r-1-i))
    return s
end

ext: jl


text:
const CACHE = LRU(maxsize = @load_preference("cache_max_size", 1000))

resolve(url) = JSON3.read(String(HTTP.get(url).body))

function cached_resolve(url)
    @debug "url $url in cache: $(haskey(CACHE,url))"

    get!(CACHE, url) do
        @debug "get $url", typeof(url)
        resolve(url)
    end
end

"""
    STAC.set_cache_max_size(cache_max_size::Integer)

Set the maximum number of URLs saved in cache (permanentaly). The default is 1000.
"""
function set_cache_max_size(cache_max_size::Integer)
    @set_preferences!("cache_max_size" => cache_max_size)
    resize!(CACHE; maxsize = cache_max_size)
end

"""
    STAC.empty_cache()

Empty the URL cache.
"""
empty_cache() = empty!(CACHE)


text:
module GeometryPrimitives
using StaticArrays, LinearAlgebra, ChainRulesCore, Zygote, Tullio, LoopVectorization
using Statistics: mean
# using Zygote: @ignore # remove as soon as ChainRules/ChainRulesCore adds replacement
# using ChainRulesCore: @non_differentiable

export Shape, Shape1, Shape2, Shape3
export surfpt_nearby, normal, bounds, translate

abstract type Shape{N,N²,D,T} end # a solid geometric shape in N dimensions (N² = N*N is needed in some shapes, e.g., Box)
const Shape1 = Shape{1,1}
const Shape2 = Shape{2,4}
const Shape3 = Shape{3,9}

Base.ndims(o::Shape{N}) where {N} = N

# The following functions return Any due to the limitations of Julia's dispatch system.
# Therefore, always call them with type assertions.  See
# https://discourse.julialang.org/t/extending-base-in-type-stably/5341/12
# https://github.com/JuliaLang/julia/issues/23210
Base.in(x::AbstractVector{<:Real}, o::Shape{N}) where {N} = SVector{N}(x) in o
surfpt_nearby(x::AbstractVector{<:Real}, o::Shape{N}) where {N} = surfpt_nearby(SVector{N}(x), o)
normal(x::AbstractVector{<:Real}, o::Shape) = surfpt_nearby(x, o)[2]  # outward direction even for x inside o
translate(o::Shape{N}, ∆::AbstractVector{<:Real}) where {N} = translate(o, SVector{N}(∆))

function orthoaxes(n::SVector{3,<:Real})
    u_temp = abs(n[3]) < abs(n[1]) ? SVector(0,0,1) : SVector(1,0,0)
    v = normalize(n × u_temp)
    u = v × n

    return u, v
end

function orthoaxes(n::SVector{N,<:Real}) where {N}
    u_temp = abs(n[3]) < abs(n[1]) ? SVector(0,0,1) : SVector(1,0,0)
    v = normalize(n × u_temp)
    u = v × n

    return u, v
end

SVector{2,T}(x::SVector{3,T}) where T<:Real = SVector{2,T}(x[1],x[2])
SVector{2}(x::SVector{3}) = SVector{2}(x[1],x[2])
Base.in(x::SVector{3,<:Real}, s::Shape2) = (x2=SVector{2}(x[1],x[2]); in(x2,s))

include("box.jl")
include("ellipsoid.jl")
include("sphere.jl")
include("prism/prism.jl")
include("periodize.jl")
include("kdtree.jl")
include("vxlcut.jl")

end # module


text:
mutable struct CartPoleEnv
    gravity::Float32
    masscart::Float32
    masspole::Float32
    total_mass::Float32
    length::Float32  # actually half the pole's length
    polemass_length::Float32
    force_mag::Float32
    τ::Float32   # seconds between state updates
    kinematics_integrator::AbstractString

    # Angle at which to fail the episode
    θ_threshold_radians::Float32
    x_threshold::Float32
    action_space::Discrete
    observation_space::Box
    viewer
    state

    steps_beyond_done
end

include("vis/cartpole.jl")

function CartPoleEnv()
    gravity = 98f-1
    masscart = 1f0
    masspole = 1f-1
    total_mass = masspole + masscart
    length = 5f-1 # actually half the pole's length
    polemass_length = (masspole * length)
    force_mag = 1f1
    τ = 2f-2  # seconds between state updates
    kinematics_integrator = "euler"

    # Angle at which to fail the episode
    θ_threshold_radians = Float32(12 * 2 * π / 360)
    x_threshold = 24f-1

    # Angle limit set to 2θ_threshold_radians so failing observation is still within bounds
    high = [
        2x_threshold,
        maxintfloat(Float32),
        2θ_threshold_radians,
        maxintfloat(Float32)]

    action_space = Discrete(2)
    observation_space = Box(-high, high, Float32)

    viewer = nothing
    state = nothing

    steps_beyond_done = nothing
    CartPoleEnv(
        gravity, masscart, masspole, total_mass, length, polemass_length,
        force_mag, τ, kinematics_integrator, θ_threshold_radians, x_threshold,
        action_space, observation_space, viewer, state, steps_beyond_done)
end

function step!(env::CartPoleEnv, action)
    @assert action ∈ env.action_space "$action in ($(env.action_space)) invalid"
    state = env.state
    x, ẋ, θ, θ̇  = state[1:1], state[2:2], state[3:3], state[4:4]
    #force = action == 2 ? env.force_mag : -env.force_mag
    force = action .* env.force_mag  # action is +1 or -1
    cosθ = cos.(θ)
    sinθ = sin.(θ)
    temp = (force .+ env.polemass_length * θ̇  .^ 2 .* sinθ) / env.total_mass
    θacc = (env.gravity*sinθ .- cosθ.*temp) ./
           (env.length * (4f0/3 .- env.masspole * cosθ .^ 2 / env.total_mass))
    xacc  = temp .- env.polemass_length * θacc .* cosθ / env.total_mass
    if env.kinematics_integrator == "euler"
        x_ = x  .+ env.τ * ẋ
        ẋ_ = ẋ  .+ env.τ * xacc
        θ_ = θ  .+ env.τ * θ̇
        θ̇_ = θ̇  .+ env.τ * θacc
    else # semi-implicit euler
        ẋ_ = ẋ  .+ env.τ * xacc
        x_ = x  .+ env.τ * ẋ_
        θ̇_ = θ̇  .+ env.τ * θacc
        θ_ = θ  .+ env.τ * θ̇_
    end

    env.state = vcat(x_, ẋ_, θ_, θ̇_)
    done =  !(all(vcat(-env.x_threshold .≤ x_ .≤ env.x_threshold,
            -env.θ_threshold_radians .≤ θ_ .≤ env.θ_threshold_radians)))

    if !done
        reward = 1f0
    elseif env.steps_beyond_done === nothing
        # Pole just fell!
        env.steps_beyond_done = 0
        reward = 1f0
    else
        if env.steps_beyond_done == 0
            @warn "You are calling 'step!()' even though this environment has already returned done = true. You should always call 'reset()' once you receive 'done = true' -- any further steps are undefined behavior."
        end
        env.steps_beyond_done += 1
        reward = 0f0
    end

    return env.state, reward, done, Dict()
end

function reset!(env::CartPoleEnv)
    env.state = param(rand(Float32, 4) * 1f-1 .- 5f-2)

    if isdefined(Main, :CuArrays)
        env.state = env.state |> gpu
    end

    env.steps_beyond_done = nothing

    return env.state
end

show(io::IO, env::CartPoleEnv) = print(io, "CartPoleEnv")


text:
using Documenter
using Unfold
using DocStringExtensions
using Plots
gr()
#unicodeplots()

makedocs(sitename="Unfold.jl",
        #root = joinpath(dirname(pathof(Unfold)), "..", "docs"),
        #prettyurls = get(ENV, "CI", nothing) == "true",
        pages = [
            "index.md",
            "Tutorials"=>[
                "Running these tutorials" => "tutorials/installation.md",
                "Mass Univariate" =>"tutorials/lm_mu.md",
                "Mass Univariate Mixed Model" =>"tutorials/lmm_mu.md",
                "LM Overlap correction" =>"tutorials/lm_overlap.md",
                "(t.b.d.) LMM + Overlap correction" =>"tutorials/lmm_overlap.md",
            ],
            "HowTo"=>[
		      #"Overlap: Different events"=>"HowTo/multiple_events.md",
                      "Custom Solvers / StandardErrors / B2B"=>"HowTo/custom_solvers.md"  ],
            "Explanations"=>["Temporal Basisfunctions" => "./explanations/basisfunctions.md"],
            "Reference"=>["Types" => "references/types.md",
            "Functions" => "references/functions.md"],

        ])

deploydocs(; repo = "github.com/unfoldtoolbox/Unfold.jl", push_preview = true,        devbranch = "main")


text: Base.show(io::IO, x::Infinite) = print(io, isposinf(x) ? "∞" : "-∞")


text:
using AbstractGPs
using TemporalGPs

# Utilising TemporalGPs.jl to work with AbstractGPs requires minimal modification to the
# GP objects from AbstractGPs that you are used to. The primary differences are
# 1. RegularSpacing is a useful type. It's basically a `range` that's hacked together to
#    work nicely with Zygote.jl. At some point, it will hopefully disappear.
# 2. Call `to_sde` on your AbstractGP object to say "use TemporalGPs.jl to do inference".

# This is an example of a very, very noise regression problem.
# You would probably be better of using a pseudo-point approximation for this particular
# data set, but longer time-series tend to lend themselves less well to nice plots.

# Load up the separable kernel from TemporalGPs.
using TemporalGPs: RegularSpacing

# Build a GP as per usual, and wrap it inside a TemporalGPs.jl object.
f_raw = GP(Matern52Kernel());
f = to_sde(f_raw, SArrayStorage(Float64));

# Specify a collection of inputs. Must be increasing.
T = 1_000_000;
x = RegularSpacing(0.0, 1e-4, T);

# Generate some noisy synthetic data from the GP.
σ²_noise = 5.0;
y = rand(f(x, σ²_noise));

# Construct the posterior as per usual.
f_post = posterior(f(x, σ²_noise), y);

# Specify some locations at which to make predictions.
T_pr = 1_200_000;
x_pr = RegularSpacing(0.0, 1e-4, T_pr);

# Compute the exact posterior marginals at `x_pr`.
f_post_marginals = marginals(f_post(x_pr));
m_post_marginals = mean.(f_post_marginals);
σ_post_marginals = std.(f_post_marginals);

# Generate a few posterior samples. Not fantastically-well optimised at present.
f_post_samples = [rand(f_post(x_pr)) for _ in 1:5];

# Visualise the posterior. The if block is just to prevent it running in CI.
if get(ENV, "TESTING", "FALSE") == "FALSE"
    using Plots
    plt = plot();
    scatter!(plt, x, y; label="", markersize=0.1, alpha=0.1);
    plot!(plt, f_post(x_pr); ribbon_scale=3.0, label="");
    plot!(x_pr, f_post_samples; color=:red, label="");
    savefig(plt, "posterior.png");
end


text:
@testset "mul" begin
    @testset "Dense" begin
        M = GBMatrix(rand(-10.0:0.05:10.0, 10, 10))
        Y = GBMatrix(rand(-10.0:0.05:10.0, 10))
        N = GBMatrix(rand(-10.0:0.05:10.0, 10, 11))
        @testset "+.*" begin
            test_frule(mul, M, Y)
            test_frule(mul, M, Y, Semirings.PLUS_TIMES)
            test_rrule(mul, M, Y)
            test_rrule(mul, M, Y, Semirings.PLUS_TIMES)

            test_frule(mul, M, N)
            test_frule(mul, M, N, Semirings.PLUS_TIMES)
            test_rrule(mul, M, N)
            test_rrule(mul, M, N, Semirings.PLUS_TIMES)
        end

        @testset "+.÷" begin
            test_rrule(mul, M, Y, Semirings.PLUS_DIV)
            test_rrule(mul, M, N, Semirings.PLUS_DIV)
        end

        @testset "+.+" begin
            test_frule(mul, M, Y, Semirings.PLUS_PLUS)
            test_frule(mul, M, N, Semirings.PLUS_PLUS)
            test_rrule(mul, M, Y, Semirings.PLUS_PLUS)
            test_rrule(mul, M, N, Semirings.PLUS_PLUS)
        end

        @testset "+.-" begin
            test_frule(mul, M, Y, Semirings.PLUS_MINUS)
            test_frule(mul, M, N, Semirings.PLUS_MINUS)
            test_rrule(mul, M, Y, Semirings.PLUS_MINUS)
            test_rrule(mul, M, N, Semirings.PLUS_MINUS)
        end
    end

    @testset "Sparse" begin
        M = GBMatrix(sprand(100, 10, 0.25))
        Y = GBMatrix(sprand(10, 0.1))
        N = GBMatrix(sprand(10, 75, 0.05))
        @testset "+.*" begin
            test_frule(mul, M, Y)
            test_frule(mul, M, Y, Semirings.PLUS_TIMES)
            test_rrule(mul, M, Y)
            test_rrule(mul, M, Y, Semirings.PLUS_TIMES)

            test_frule(mul, M, N)
            test_frule(mul, M, N, Semirings.PLUS_TIMES)
            test_rrule(mul, M, N)
            test_rrule(mul, M, N, Semirings.PLUS_TIMES)
        end

        @testset "+.÷" begin
            test_rrule(mul, M, Y, Semirings.PLUS_DIV)
            test_rrule(mul, M, N, Semirings.PLUS_DIV)
        end

        @testset "+.+" begin
            test_frule(mul, M, Y, Semirings.PLUS_PLUS)
            test_frule(mul, M, N, Semirings.PLUS_PLUS)
            test_rrule(mul, M, Y, Semirings.PLUS_PLUS)
            test_rrule(mul, M, N, Semirings.PLUS_PLUS)
        end
        @testset "+.-" begin
            test_frule(mul, M, Y, Semirings.PLUS_MINUS)
            test_frule(mul, M, N, Semirings.PLUS_MINUS)
            test_rrule(mul, M, Y, Semirings.PLUS_MINUS)
            test_rrule(mul, M, N, Semirings.PLUS_MINUS)
        end
    end
end


text:
module Record
using MLStyle
using MLStyle.MatchCore
using MLStyle.MatchImpl
using MLStyle.AbstractPatterns
using MLStyle.AbstractPatterns.BasicPatterns
using MLStyle.Qualification

export @as_record, record_def

@nospecialize
function P_partial_struct_decons(t, partial_fields, ps, prepr::AbstractString = "$t")
    function tcons(_...)
        t
    end

    comp = PComp(prepr, tcons;)
    function extract(sub::Any, i::Int, ::Any, ::Any)
        :($sub.$(partial_fields[i]))
    end
    decons(comp, extract, ps)
end

function _compile_record_pattern(
    t::Type,
    self::Function,
    type_params::Any,
    type_args::Any,
    args::Any,
)
    isempty(type_params) || return begin
        call = Expr(:call, t, args...)
        ann = Expr(:curly, t, type_args...)
        self(Where(call, ann, type_params))
    end
    all_field_names = fieldnames(t)
    partial_field_names = Symbol[]
    patterns = Function[]

    @switch args begin
        @case [Expr(:parameters, kwargs...), args...]
        @case let kwargs = []
        end
    end

    n_args = length(args)
    if all(Meta.isexpr(arg, :kw) for arg in args)
        for arg in args
            field_name = arg.args[1]
            field_name in all_field_names || error("$t has no field $field_name.")
            push!(partial_field_names, field_name)
            push!(patterns, self(arg.args[2]))
        end
    elseif length(all_field_names) === n_args
        append!(patterns, map(self, args))
        append!(partial_field_names, all_field_names)
    elseif n_args === 1 && args[1] === :_
    elseif n_args !== 0
        error(
            "count of positional fields should be 0 or the same as the fields($all_field_names)",
        )
    end
    for e in kwargs
        @switch e begin
            @case ::Symbol
            e in all_field_names ||
                error("unknown field name $e for $t when field punnning.")
            push!(partial_field_names, e)
            push!(patterns, P_capture(e))
            continue
            @case Expr(:kw, key::Symbol, value)
            key in all_field_names ||
                error("unknown field name $key for $t when field punnning.")
            push!(partial_field_names, key)
            push!(patterns, and([P_capture(key), self(value)]))
            continue
            @case _

            error("unknown sub-pattern $e in $t.")
        end
    end

    ret = P_partial_struct_decons(t, partial_field_names, patterns)
    isempty(type_args) && return ret
    and([self(Expr(:(::), Expr(:curly, t, type_args...))), ret])
end

function record_def(@nospecialize(Struct), line::LineNumberNode, ::Module)
    quote
        $line
        function $MLStyle.pattern_uncall(
            t::Type{<:$Struct},
            self::Function,
            type_params::Any,
            type_args::Any,
            args::Any,
        )
            $line
            $_compile_record_pattern(t, self, type_params, type_args, args)
        end
    end
end

function as_record(@nospecialize(n), line::LineNumberNode, __module__::Module)
    @switch n begin
        @case :(struct $hd{$(_...)}
                  $(_...)
              end) ||
              :(struct $hd{$(_...)} <: $_
                  $(_...)
              end) ||
              :(struct $hd <: $_
                  $(_...)
              end) ||
              :(struct $hd
                  $(_...)
              end)
        return Expr(:block, n, record_def(hd, line, __module__))
        @case _
        return record_def(n, line, __module__)
    end
end

macro as_record(qualifier, n)
    deprecate_qualifier_macro(qualifier, __source__)
    esc(as_record(n, __source__, __module__))
end

macro as_record(n)
    esc(as_record(n, __source__, __module__))
end
@specialize
end


text:
fname = :ST_TROp
c = Combi(hessian_operator,PDataST,solve_modelST_TR,preprocessST,decreaseGen,TparamsST())
include("Template.jl")


text:
using Blink, WebIO

t() = scatter(y=rand(10))
p = plot([t(), t(), t(), t()])
w = Blink.Window()
body!(w, p.scope)
sleep(3.0)  # make sure we give time for svg to render

# hook up testing observables
on(p.scope["__gd_contents"]) do x end

function update_data!()
    p.scope["__get_gd_contents"][] = "data"
    sleep(1)
end

function update_layout!()
    p.scope["__get_gd_contents"][] = "layout"
    sleep(1)
end

@testset "api methods" begin

    @testset "to_image" begin
        @test to_image(p)[1:21] == "data:image/png;base64"
        @test to_image(p, format="svg")[1:18] == "data:image/svg+xml"
        @test to_image(p, format="jpeg")[1:22] =="data:image/jpeg;base64"
        @test to_image(p, format="webp")[1:22] =="data:image/webp;base64"
        @test_throws ErrorException to_image(p, format="pdf")
        @test to_image(p, imageDataOnly=true, format="svg")[1:4] == "<svg"
    end

    @testset "restyle" begin
        p.scope["__gd_contents"][] = Dict() # reset
        restyle!(p, marker_color="pink", 1)
        update_data!()

        restyle!(p, marker_symbol="star")
        update_data!()
        for i in 1:4
            @test p.scope["__gd_contents"][][i]["marker"]["symbol"] == "star"
        end
    end

    @testset "relayout" begin
        relayout!(p, title="This is a title")
        update_layout!()
        @test p.scope["__gd_contents"][]["title"] == "This is a title"
    end

    @testset "update" begin
        update!(p, marker_size=10, layout=Layout(xaxis_title="This is x"))
        update_layout!()
        @test p.scope["__gd_contents"][]["xaxis"]["title"] == "This is x"

        update_data!()
        for i in 1:4
            @test p.scope["__gd_contents"][][i]["marker"]["size"] == 10
        end
    end

    @testset "addtraces" begin
        update_data!()
        @test length(p.scope["__gd_contents"][]) == 4

        new_trace = scatter(x=1:10, y=(1:10.0).^2, name="New trace")
        addtraces!(p, new_trace)
        update_data!()

        @test length(p.scope["__gd_contents"][]) == 5
        @test all(p.scope["__gd_contents"][][5]["x"] .== 1:10)
        @test all(p.scope["__gd_contents"][][5]["y"] .== (1:10).^2)
        @test p.scope["__gd_contents"][][5]["name"] == "New trace"
    end

    @testset "movetraces" begin
        update_data!()
        @test p.scope["__gd_contents"][][5]["name"] == "New trace"

        movetraces!(p, [5, 4], [4, 5])
        update_data!()
        @test p.scope["__gd_contents"][][4]["name"]  == "New trace"
    end

    @testset "deletetraces" begin
        update_data!()
        @test length(p.scope["__gd_contents"][]) == 5
        deletetraces!(p, 4)

        update_data!()
        @test length(p.scope["__gd_contents"][]) == 4
        @test !haskey(p.scope["__gd_contents"][][4], "name")
    end

    @testset "extendtraces" begin
        extendtraces!(p, 1, -1; y=[2, 2.5, 3.0])
        update_data!()
        @test length(p.scope["__gd_contents"][][1]["y"]) == 13
        @test all(p.scope["__gd_contents"][][1]["y"][end-2:end] .== [2, 2.5, 3.0])
    end

    @testset "prependtraces" begin
        prependtraces!(p, 1, -1; y=[3, 4, 5])
        update_data!()
        @test length(p.scope["__gd_contents"][][1]["y"]) == 16
        @test all(p.scope["__gd_contents"][][1]["y"][1:3] .== [3, 4, 5])
    end
end


text:
using DataFrames, JuMP, Ipopt, DelimitedFiles;

include("reducedIO.jl")

# evaluating variables and parameters
#A = ones(numSectors)
# compensation of employees
#LCurrent = zeros(numSectors)
# consumption?
#CCurrent = zeros(numSectors)
# output
#QCurrent = zeros(numSectors)
#
#MCurrentT = zeros(numSectors)
#
#XCurrentT = zeros(numSectors)
# new capital
capFlowsReceivable = zeros(numSectors)
# labour share
alpha = zeros(numSectors)
# intermediate share
mu = zeros(numSectors)
# consumption share
xi = zeros(numSectors)
# matrix of intermediate flows as shares
gammaM = zeros(numSectors,numSectors)
# matrix of capital investment as shares
gammaX = zeros(numSectors,numSectors)
# intermediate inputs
#MCurrent = zeros(numSectors,numSectors)
# capital flows - sector to sector
#XCurrent = zeros(numSectors,numSectors)
# was called M, =reducedIO[numSectors+3,i]/(reducedIO[numSectors+1,i]+reducedIO[numSectors+3,i])
#KNext = zeros(numSectors)

epsM = 0.1*0.98
epsQ = 1*0.98
epsD = 1*0.98
epsX = 1*0.98
epsLS = 1*0.98
deltaC = ones(numSectors)*0.0098
# next ones are from aldrich
beta = 0.984*0.98
deltaK = 0.01*0.98

#calculating from data
for i in 1:numSectors
    #QCurrent[i] = reducedIO[numSectors+8,i]
    #LCurrent[i] = reducedIO[numSectors+2,i]
    #CCurrent[i] = (reducedIO[i,numSectors+2]+reducedIO[i,numSectors+3])
    alpha[i] = reducedIO[numSectors+2,i]/(reducedIO[numSectors+2,i]+reducedIO[numSectors+3,i])+0.001
    mu[i] = reducedIO[numSectors+3,i]/(reducedIO[numSectors+1,i]+reducedIO[numSectors+3,i])+0.001
    xi[i] = (reducedIO[i,numSectors+2] + reducedIO[i,numSectors+3])/(reducedIO[numSectors+1,numSectors+2] + reducedIO[numSectors+1,numSectors+2])+0.001
    capFlowsReceivable[i] = sum(reducedCapFlows[:,i])+0.001
    for j in 1:numSectors
        gammaM[i,j] = reducedIO[j,i]/(reducedIO[j,numSectors+1])+0.001
        gammaX[i,j] = reducedCapFlows[j,i]/sum(reducedCapFlows[:,i])+0.001
        #MCurrent[i,j] = reducedIO[i,j]
        #XCurrent[i,j] = reducedCapFlows[i,j]
    end
end

A = ones(numSectors).*0.98;

modBasic = Model(Ipopt.Optimizer)
@variable(modBasic, CCurrent[1:numSectors]>=0.01)
@variable(modBasic, LCurrent[1:numSectors]>=0.01)
@variable(modBasic, KCurrent[1:numSectors]>=0.01)
@variable(modBasic, KNext[1:numSectors]>=0.01)
@variable(modBasic, KNextNext[1:numSectors]>=0.01)
@variable(modBasic, MCurrent[1:numSectors,1:numSectors]>=0.01)
@variable(modBasic, XCurrent[1:numSectors,1:numSectors]>=0.01)
@variable(modBasic, QCurrent[1:numSectors]>=0.01)
@variable(modBasic, MCurrentT[1:numSectors]>=0.01)
@variable(modBasic, XCurrentT[1:numSectors]>=0.01)
@NLobjective(modBasic, Max, beta*(log(sum(xi[j]^(1/epsD)*(deltaC[j]*CCurrent[j])^((0.0001+epsD-1)/epsD) for j in 1:numSectors)^((0.0001+epsD)/(epsD*0.999-1)))
-epsLS/(0.999*epsLS+1)+sum(LCurrent[i] for i in 1:numSectors)^((epsLS+1)/0.999*epsLS)));
# Market clearing constraint
@NLconstraint(modBasic, sum(QCurrent[j]-CCurrent[j] - sum(MCurrent[i,j]+XCurrent[i,j] for i in 1:numSectors) for j in 1:numSectors) == 0);
@NLconstraint(modBasic, sum(XCurrentT[j]+(1-deltaK)*KCurrent[j]-KNext[j] for  j in 1:numSectors) == 0);
@NLconstraint(modBasic, sum((1-deltaK)*KNext[j]-KNextNext[j] for j in 1:numSectors) == 0);
for j in 1:numSectors;
    #@NLconstraint(modBasic, XCurrentT[j]==sum(gammaX[i,j]^(1/(0.0001+epsX))*XCurrent[i,j]^((0.0001+epsX-1)/epsX) for i in 1:numSectors)^((0.0001+epsX)/(0.0001+epsX-1)));
    #@NLconstraint(modBasic, MCurrentT[j]==sum(gammaM[i,j]^(1/epsM)*MCurrent[i,j]^((0.0001+epsM-1)/epsM) for i in 1:numSectors)^((0.0001+epsM)/(epsM-1)));
    #impose Q constraint from appendix F.1
    @NLconstraint(modBasic, A[j]*((1-mu[j])^(1/(0.0001+epsQ))*((KCurrent[j]/(0.0001+alpha[j]))^(0.0001+alpha[j])*(LCurrent[j]/(1-alpha[j]))^(1-alpha[j]))^((0.0001+epsQ-1)/epsQ)+mu[j]^(1/(0.0001+epsQ))*MCurrentT[j]^((0.0001+epsQ-1)/epsQ))^(epsQ/(0.99*epsQ-1)) == QCurrent[j]);
end
optimize!(modBasic)
#=
lagrange = E0*sum(beta*(log(sum(xi[j]^(1/epsD)*(deltaC[j]*C[j])^((epsD-1)/epsD))^(epsD/(epsD*0.999-1)))
-epsLS/(epsLS+1)+sum(LCurrent[j])^((epsLS+1)/epsLS)
+sum(P[j]*(QCurrent[j]-CCurrent[j] - sum(MCurrent[j,i]+XCurrent[j,i])))))
=#

#=
constraints:

sum(QCurrent[j]-CCurrent[j] - sum(MCurrent[j,i]+XCurrent[j,i])) == 0

sum(XCurrentT[j]+(1-deltaK)*KCurrent[j]-Kt+1[j]) == 0

FOC's:
P_t-1 - beta*P*(1-deltaC[j])=beta*xi[j]^(1/epsD)*deltaC[j]^((epsD-1)/epsD)*
CCurrent[j]^(-1/epsD)*(sum(xi[i]^(1/epsD)*(deltaC[i]*CCurrent[i])^((epsD-1)/eqsD))^(-1)

P[i]/P[j] = A[j]^((epsQ-1)/epsQ)*(QCurrent[j]*mu[j]/MCurrentOne[j])^(1/epsQ)*(MCurrentOne[j]*gammaM[j])^(1/epsM)

=#

max_stars_repo_name: uq-aibe/aim

max_issues_repo_name: uq-aibe/aim

max_forks_repo_name: uq-aibe/aim


text:
#!/usr/bin/env julia

# Test for Wrapper for GDIP library libGDIP.so

##################################################
# Imports
##################################################

using Printf

# import PyPlot in version 2.8.0+
import Pkg
try
    using PyPlot
catch
    Pkg.add("PyPlot")
    using PyPlot
end
if PyPlot.version < v"2.8.0"
    error("Please upgrade PyPlot to version 2.8.0+")
    @assert false
end
Circle = matplotlib.patches.Circle

# import wrapper for the GDIP solution
include("GDIP.jl")

##################################################
# Settings
##################################################

# Turning radius
turning_radius = 1.0
# Step of the theta2 angle in the visualization
step = 2 * pi / 100
# Wait time for the animation
wait_time = 0.01
# Save figures into a directory "images"
save_figures = false
if length(ARGS) > 0
    if "-save" == ARGS[1]
        save_figures = true
    end
end
if save_figures
    mkpath("images")
end
# Sensing radius in case of GDIP solution
sensing_radius_gdip = 0.5

# End points
point1 = [0.0, 0.0]
point2 = [5.0, 0.0]

# First angle interval (the second is changed)
interval1 = [1.0, 0.5]

# Enable only specific parts
DUB = false
DIP = false
GDIP = true

#matplotlib.rcParams['text.usetex'] = true

##################################################
# Functions
##################################################

function plot_points(points; specs = "b", zorder = 1)
    x_val = [x[1] for x in points]
    y_val = [x[2] for x in points]
    return plt.plot(x_val, y_val, specs, zorder = zorder)
end

function plot_interval(point, interval, sensing_radius)
    ax = plt.gca()
    circle = Circle(point, sensing_radius, facecolor="yellow", edgecolor="orange", linewidth=1, alpha=0.2)
    ax.add_patch(circle)

    for angle in [interval[1], interval[1] + interval[2]]
        last = point + turning_radius * [cos(angle), sin(angle)]
        points = [point, last]
        plot_points(points, specs = "y-", zorder = 1)
    end
end

function plot_arrow(px, py, angle; alpha = 1)
    rad = 0.7 * turning_radius
    dx = rad * cos(angle)
    dy = rad * sin(angle)
    plt.arrow(px, py, dx, dy, head_width=0.05, head_length=0.1, fc="k", ec="k", zorder=2, alpha=alpha)
    plt.plot([px, px + dx], [py, py + dy], "k", visible=false)
end

function plot_figure(samples, point1, point2, title; sensing_radius = 0, interval1 = nothing, interval2 = nothing)
    # Clear and basic settings
    plt.clf()
    plt.axis("equal")
    plt.title(title)
    plt.xlabel("x [-]")
    plt.ylabel("y [-]")
    plt.ylim([-3, 3])
    plt.xlim([-1.5, 6.5])
    plt.grid(alpha=0.2)

    # Plot regions based on sensing radius
    if interval1 != nothing
        plot_interval(point1, interval1, sensing_radius)
    end
    if interval2 != nothing
        plot_interval(point2, interval2, sensing_radius)
    end

    # Plot samples of the Dubins maneuver
    ps = plot_points(samples)
    plt.legend([ps[1]], ["Dubins maneuver"])

    # First arrow
    p1 = samples[1]
    p2 = samples[2]
    angle = atan(p2[2]-p1[2], p2[1]-p1[1])
    plot_arrow(p1[1], p1[2], angle)
    plot_arrow(point1[1], point1[2], angle, alpha = 0.3)
    plt.plot(p1[1], p1[2], "k.")

    # Second arrow
    p1 = samples[length(samples)]
    p2 = samples[length(samples)-1]
    angle = atan(p1[2]-p2[2], p1[1]-p2[1])
    plot_arrow(p1[1], p1[2], angle)
    plot_arrow(point2[1], point2[2], angle, alpha = 0.3)
    plt.plot(p1[1], p1[2], "k.")

    if ! save_figures
        plt.pause(wait_time)
    end

    if save_figures
        filename = @sprintf("images/%05d.png", image_idx)
        #filename = "images/{:05d}.png".format(image_idx)
        global image_idx = image_idx + 1
        print(filename * "\n")
        plt.savefig(filename)
    end
end

image_idx = 0

##################################################
# Basic Dubins maneuver
##################################################
if DUB
    for t2 in 0:step:(2*pi)
        start = [point1[1], point1[2], 1.0]
        final = [point2[1], point2[2], t2]
        gdip_init_dubins_maneuver(start, final, 1.0)
        samples = gdip_sample_many(step_size)

        plot_figure(samples, point1, point2, "Dubins maneuver")
    end
end

##################################################
# Dubins Interval Problem (DIP)
##################################################
if DIP
    sensing_radius = 0.0

    for t2 in 0:step:(2*pi)
        interval2 = [2 + t2, 0.5]
        gdip_init_dip(point1, interval1, point2, interval2, turning_radius)
        samples = gdip_sample_many(0.1)

        plot_figure(samples, point1, point2, "Dubins Interval Problem (DIP)",
            sensing_radius = sensing_radius, interval1 = interval1, interval2 = interval2)

    end
end

##################################################
# Generalized Dubins Interval Problem (GDIP)
##################################################
if GDIP
    sensing_radius = sensing_radius_gdip

    for t2 in 0:step:(2*pi)
        interval2 = [2 + t2, 0.5]

        gdip_init_gdip(point1, interval1, sensing_radius, point2, interval2, sensing_radius, turning_radius)
        samples = gdip_sample_many(0.1)

        plot_figure(samples, point1, point2, "Dubins Interval Problem (DIP)",
            sensing_radius = sensing_radius, interval1 = interval1, interval2 = interval2)

    end
end

max_stars_repo_name: petvana/gdip

max_issues_repo_name: petvana/gdip

max_forks_repo_name: petvana/gdip


text:
const SIZEOF_HDR1 = Int32(348)
const SIZEOF_HDR2 = Int32(540)

const NP1_MAGIC = (0x6e,0x2b,0x31,0x00)
const NI1_MAGIC = (0x6e,0x69,0x31,0x00)
const NP2_MAGIC = (0x6e,0x2b,0x32,0x00,0x0d,0x0a,0x1a,0x0a)
const NI2_MAGIC = (0x6e,0x69,0x32,0x00,0x0d,0x0a,0x1a,0x0a)

function to_eltype_error(@nospecialize(x))
    throw("NIfTI parse erorr: $x does not have a corresponding eltype ")
end

eltype_to_int16(::Type{T}) where {T} = error("Unsupported data type $T")

for T in (Int16, Int32)
    symint = Symbol(lowercase(string(T)))
    @eval begin
        #= TODO handle intent codes
        @inline function to_intent(x::Int16)
            if i === $T(1001)
                return Estimate
            elseif i === $T(1002)
                return NeuroLabel
            elseif i === $T(1003)
                return NeuroNames
            elseif i === $T(1004)
                return SMatrix
            elseif i === $T(1005)
                return Symmetric
            elseif i === $T(1006)
                return DisplacementVector
            elseif i === $T(1007)
                return SVector
            elseif i === $T(1008)
                return Point
            elseif i === $T(1009)
                return Triangle
            elseif i === $T(1010)
                return Quat
            elseif i === $T(1011)
                return Dimensionless
            elseif i === $T(2001)
                return TimeSeries
            elseif i === $T(2002)
                return Vector{Point}
            elseif i === $T(2003)
                return GiftiRGB
            elseif i === $T(2004)
                return GiftiRGBA
            elseif i === $T(2005)
                return Polygon
            elseif i === $T(2006)
               return FSLDisplacementVector # TODO
            elseif i === $T(2007)
                return FSLCubicSplineCoefficient  # TODO
            elseif i === $T(2008)
                return FSLDCTCoefficients  # TODO
            elseif i === $T(2009)
                return FSLQuadraticSplineCoefficients  # TODO
            elseif i === $T(2016)
                return FSLTopupCubicSplineCoefficients  # TODO
            elseif i === $T(2017)
                return TopupQuadraticSplineCoefficients # TODO
            elseif i === $T(2018)
                return TopupField  # TODO
            elseif i == $T(3000)
                return ConnUnkown
            elseif i == $T(3001)
                return ConnDense
            elseif i == $T(3002)
                return ConnDenseSeries
            elseif i == $T(3003)
                return ConnParcels
            elseif i == $T(3004)
                return ConnParcelSeries
            elseif i == $T(3006)
                return ConnDenseScalar
            elseif i == $T(3007)
                return ConnDenseLabel
            elseif i == $T(3008)
                return ConnParcelScalr
            elseif i == $T(3009)
                return ConnParcelDense
            elseif i == $T(3010)
                return ConnDenseParcel
            elseif i == $T(3013)
                return ConnPPSc
            else
                return UnkownIntent
            end
        end
        =#

        @inline function $(Symbol(:xform_to_, symint))(x::Symbol)
            if x === :ScannerSpace
                return $T(1)
            elseif x === :AnatomicalSpace
                return $T(2)
            elseif x === :TalairachSpace
                return $T(3)
            elseif x === :MNI152Space
                return $T(4)
            elseif x === :OtherTemplate
                return $T(5)
            else
                return $T(0)
            end
        end

        @inline function to_xform(i::$T)
            if i === $T(1)
                return :ScannerSpace
            elseif i === $T(2)
                return :AnatomicalSpace
            elseif i === $T(3)
                return :TalairachSpace
            elseif i === $T(4)
                return :MNI152Space
            elseif i === $T(5)
                return :OtherTemplate
            else
                return :UnkownSpace
            end
        end

        @inline function to_eltype(x::$T)
            if x === $T(1)
                return Bool
            elseif x === $T(2)
                return UInt8
            elseif x === $T(4)
                return Int16
            elseif x === $T(8)
                return Int32
            elseif x === $T(16)
                return Float32
            elseif x === $T(32)
                return ComplexF32
            elseif x === $T(64)
                return Float64
            elseif x === $T(256)
                return Int8
            elseif x === $T(512)
                return UInt16
            elseif x === $T(768)
                return UInt32
            elseif x === $T(1024)
                return Int64
            elseif x === $T(1280)
                return UInt64
            elseif x === $T(1792)
                return ComplexF64
            else
                to_eltype_error(x)
            end
        end

        $(Symbol(:eltype_to_, symint))(::Type{Bool}) = $T(1)
        $(Symbol(:eltype_to_, symint))(::Type{UInt8}) = $T(2)
        $(Symbol(:eltype_to_, symint))(::Type{Int16}) = $T(4)
        $(Symbol(:eltype_to_, symint))(::Type{Int32}) = $T(8)
        $(Symbol(:eltype_to_, symint))(::Type{Float32}) = $T(16)
        $(Symbol(:eltype_to_, symint))(::Type{ComplexF32}) = $T(32)
        $(Symbol(:eltype_to_, symint))(::Type{Float64}) = $T(64)
        $(Symbol(:eltype_to_, symint))(::Type{Int8}) = $T(256)
        $(Symbol(:eltype_to_, symint))(::Type{UInt16}) = $T(512)
        $(Symbol(:eltype_to_, symint))(::Type{UInt32}) = $T(768)
        $(Symbol(:eltype_to_, symint))(::Type{Int64}) = $T(1024)
        $(Symbol(:eltype_to_, symint))(::Type{UInt64}) = $T(1280)
        $(Symbol(:eltype_to_, symint))(::Type{ComplexF64}) = $T(1792)
    end
end

to_dimensions(d::NTuple{8,Int16}) = ntuple(i ->Int(@inbounds(getfield(d, i + 1))), first(d))
to_dimensions(d::NTuple{8,Int}) = ntuple(i ->@inbounds(getfield(d, i + 1)), first(d))

function hdr_to_img(file::AbstractString)
    volume_name = replace(file, r"\.\w+(\.gz)?$" => "")*".img"
    if isfile(volume_name)
        return volume_name
    else
        volume_name *= ".gz"
        if isfile(volume_name)
            return volume_name
        else
            error("NIfTI file is dual file storage, but $volume_name does not exist")
        end
    end
end

function string_tuple(x::String, n::Int)
    a = codeunits(x)
    padding = zeros(UInt8, n-length(a))
    (a..., padding...)
end
string_tuple(x::AbstractString) = string_tuple(bytestring(x))

# Conversion factors to mm/ms
# http://nifti.nimh.nih.gov/nifti-1/documentation/nifti1fields/nifti1fields_pages/xyzt_units.html
const SPATIAL_UNIT_MULTIPLIERS = [
    1000,   # 1 => NIfTI_UNITS_METER
    1,      # 2 => NIfTI_UNITS_MM
    0.001   # 3 => NIfTI_UNITS_MICRON
]
const TIME_UNIT_MULTIPLIERS = [
    1000,   # NIfTI_UNITS_SEC
    1,      # NIfTI_UNITS_MSEC
    0.001,  # NIfTI_UNITS_USEC
    1,      # NIfTI_UNITS_HZ
    1,      # NIfTI_UNITS_PPM
    1       # NIfTI_UNITS_RADS
]

# Gets dim to be used in header
function to_dim_i16(x::NTuple{N}) where {N}
    return ntuple(Val(8)) do i
        if i === 1
            Int16(N)
        elseif i > (N + 1)
            zero(Int16)
        else
            @inbounds(getfield(x, i - 1))
        end
    end
end

#=
to_spatial_units(x::Int8) = _to_spatial_units(x & 0x07)
@inline function _to_spatial_units(x::UInt8)
    if x === UInt8(1)
        return m  # meter
    elseif x === UInt8(2)
        return mm  # millimeters
    elseif x === UInt8(3)
        return μm  # microns
    else
        return Unitful.FreeUnits{(),NoDims,nothing}()
    end
end

to_time_units(x::Int8) = _to_time_units(x & 0x38)
@inline function _to_time_units(x::Int8)
    if x === Int8(8)
        return s  # second
    elseif x === Int8(16)
        return ms  # milliseconds
    elseif x === Int8(24)
        return μs # microseconds
    elseif x === Int8(32)
        return Hz  # hertz
    #(UInt8(40), u"ppm"),  # parts per million
    elseif x === Int8(48)
        return rad/s  # radians per second
    else
        return Unitful.FreeUnits{(),NoDims,nothing}()
    end
end

units_to_int8(::typeof(Unitful.s)) = Int8(1)
units_to_int8(::typeof(Unitful.ms)) = Int8(16)
units_to_int8(::typeof(Unitful.μs)) = Int8(24)
units_to_int8(::typeof(Unitful.Hz)) = Int8(32)
units_to_int8(::typeof((Unitful.rad/Unitful.s))) = Int8(48)
units_to_int8(::Unitful.FreeUnits{(), NoDims, nothing}) = Int8(0)

function space_time_to_xyzt(ss, tt)
    return (units_to_int8(ss) & 0x07) | (units_to_int8(tt) & 0x38)
end
=#


text:
module Architectures

export AbstractArchitecture, AbstractMultiArchitecture
export CPU, GPU
export device, device_event, architecture, array_type, arch_array

using CUDA
using KernelAbstractions
using CUDAKernels
using Adapt
using OffsetArrays

"""
    AbstractArchitecture

Abstract supertype for architectures supported by Oceananigans.
"""
abstract type AbstractArchitecture end

"""
    AbstractMultiArchitecture

Abstract supertype for Distributed architectures supported by Oceananigans.
"""
abstract type AbstractMultiArchitecture <: AbstractArchitecture end

"""
    CPU <: AbstractArchitecture

Run Oceananigans on one CPU node. Uses multiple threads if the environment
variable `JULIA_NUM_THREADS` is set.
"""
struct CPU <: AbstractArchitecture end

"""
    GPU <: AbstractArchitecture

Run Oceananigans on a single NVIDIA CUDA GPU.
"""
struct GPU <: AbstractArchitecture end

#####
##### These methods are extended in Distributed.jl
#####

device(::CPU) = KernelAbstractions.CPU()
device(::GPU) = CUDAKernels.CUDADevice()

architecture() = nothing
architecture(::Number) = nothing
architecture(::Array) = CPU()
architecture(::CuArray) = GPU()

"""
    child_architecture(arch)

Return `arch`itecture of child processes.
On single-process, non-distributed systems, return `arch`.
"""
child_architecture(arch) = arch

array_type(::CPU) = Array
array_type(::GPU) = CuArray

arch_array(::CPU, a::Array)   = a
arch_array(::CPU, a::CuArray) = Array(a)
arch_array(::GPU, a::Array)   = CuArray(a)
arch_array(::GPU, a::CuArray) = a

arch_array(arch, a::AbstractRange) = a
arch_array(arch, a::OffsetArray) = OffsetArray(arch_array(arch, a.parent), a.offsets...)
arch_array(arch, ::Nothing) = nothing
arch_array(arch, a::Number) = a

device_event(arch) = Event(device(arch))

end


text:
export iauLd
"""
Apply light deflection by a solar-system body, as part of
transforming coordinate direction into natural direction.

This function is part of the International Astronomical Union's
SOFA (Standards of Fundamental Astronomy) software collection.

Status:  support function.

Given:
   bm     double     mass of the gravitating body (solar masses)
   p      double[3]  direction from observer to source (unit vector)
   q      double[3]  direction from body to source (unit vector)
   e      double[3]  direction from body to observer (unit vector)
   em     double     distance from body to observer (au)
   dlim   double     deflection limiter (Note 4)

Returned:
   p1     double[3]  observer to deflected source (unit vector)

Notes:

   1. The algorithm is based on Expr. (70) in Klioner (2003) and
   Expr. (7.63) in the Explanatory Supplement (Urban & Seidelmann
   2013), with some rearrangement to minimize the effects of machine
   precision.

   2. The mass parameter bm can, as required, be adjusted in order to
   allow for such effects as quadrupole field.

   3. The barycentric position of the deflecting body should ideally
   correspond to the time of closest approach of the light ray to
   the body.

   4. The deflection limiter parameter dlim is phi^2/2, where phi is
   the angular separation (in radians) between source and body at
   which limiting is applied.  As phi shrinks below the chosen
   threshold, the deflection is artificially reduced, reaching zero
   for phi = 0.

   5. The returned vector p1 is not normalized, but the consequential
   departure from unit magnitude is always negligible.

   6. The arguments p and p1 can be the same array.

   7. To accumulate total light deflection taking into account the
   contributions from several bodies, call the present function for
   each body in succession, in decreasing order of distance from the
   observer.

   8. For efficiency, validation is omitted.  The supplied vectors must
   be of unit magnitude, and the deflection limiter non-zero and
   positive.

References:

   Urban, S. & Seidelmann, P. K. (eds), Explanatory Supplement to
   the Astronomical Almanac, 3rd ed., University Science Books
   (2013).

   Klioner, Sergei A., "A practical relativistic model for micro-
   arcsecond astrometry in space", Astr. J. 125, 1580-1597 (2003).

Called:
   iauPdp       scalar product of two p-vectors
   iauPxp       vector product of two p-vectors

This revision:   2013 October 9

SOFA release 2018-01-30

Copyright (C) 2018 IAU SOFA Board.  See notes at end.
"""

# void iauLd(double bm, double p[3], double q[3], double e[3],
#            double em, double dlim, double p1[3])

function iauLd(rh::Real, p::AbstractVector{<:Real}, q::AbstractVector{<:Real},
               e::AbstractVector{<:Real}, em::Real, dlim::Real)
   p1 = zeros(Float64, 3)

   ccall((:iauLd, libsofa_c), Cvoid,
         (Cdouble, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble},
         Cdouble, Cdouble, Ptr{Cdouble}),
          convert(Float64, rh), convert(Array{Float64, 1}, p),
          convert(Array{Float64, 1}, q), convert(Array{Float64, 1}, e),
          convert(Float64, em), convert(Float64, dlim), p1)

   return SVector{3}(p1)
end


text:
"Information about a hardware thread"
struct HWThread
    threadId::Int
    coreId::Int
    packageId::Int
    apicId::Int
    dieId::Int
    inCpuSet::Int
end

# function Base.show(io::IO, hwt::HWThread)
#     print("HWThread()")
# end

"Information about a cache level"
struct CacheLevel
    level::Int
    type::Symbol
    associativity::Int
    sets::Int
    lineSize::Int
    size::Int
    threads::Int
    inclusive::Int
end

# function Base.show(io::IO, clvl::CacheLevel)
#     print(io, "CacheLevel($(clvl.level), ...)")
# end

"CPU topology information"
struct CpuTopology
    numHWThreads::Int
    activeHWThreads::Int
    numSockets::Int
    numDies::Int
    numCoresPerSocket::Int
    numThreadsPerCore::Int
    numCacheLevels::Int
    threadPool::Vector{HWThread}
    cacheLevels::Vector{CacheLevel}
end

# struct Likwid_Configuration
#     configFileName::String
#     topologyCfgFileName::String
#     daemonPath::String
#     groupPath::String
#     daemonMode::AccessMode
#     maxNumThreads::Cint
#     maxNumNodes::Cint
# end

"CPU information"
struct CpuInfo
    family::Int
    model::Int
    stepping::Int
    vendor::Int
    part::Int
    clock::Int
    turbo::Bool
    osname::String
    name::String
    short_name::String
    features::String
    isIntel::Bool
    architecture::String
    supportUncore::Bool
    supportClientmem::Bool
    featureFlags::UInt
    perf_version::Int
    perf_num_ctr::Int
    perf_width_ctr::Int
    perf_num_fixed_ctr::Int
end

"Information about a NUMA node"
struct NumaNode
    id::Int
    totalMemory::Int # kB
    freeMemory::Int # kB
    numberOfProcessors::Int
    processors::Vector{Int}
    numberOfDistances::Int
    distances::Vector{Int}
end

Base.show(io::IO, nn::NumaNode) = print(io, "NumaNode()")

"CPU topology information"
struct NumaTopology
    numberOfNodes::Int
    nodes::Vector{NumaNode}
end

"An affinity domain"
struct AffinityDomain
    tag::String
    numberOfProcessors::Int
    numberOfCores::Int
    processorList::Vector{Int}
end

"Information about the affinity domains"
struct AffinityDomains
    numberOfSocketDomains::Int
    numberOfNumaDomains::Int
    numberOfProcessorsPerSocket::Int
    numberOfCacheDomains::Int
    numberOfCoresPerCache::Int
    numberOfProcessorsPerCache::Int
    numberOfAffinityDomains::Int
    domains::Vector{AffinityDomain}
end

"Turbo boost information"
struct TurboBoost
    numSteps::Int
    steps::Vector{Float64}
end

Base.show(io::IO, t::TurboBoost) = print(io, "TurboBoost()")

# @cenum PowerType::UInt32 begin
#     PKG = 0
#     PP0 = 1
#     PP1 = 2
#     DRAM = 3
#     PLATFORM = 4
# end

"Power domain information"
struct PowerDomain
    id::Int
    type::LibLikwid.PowerType
    supportFlags::UInt32
    energyUnit::Float64
    tdp::Float64
    minPower::Float64
    maxPower::Float64
    maxTimeWindow::Float64
    supportInfo::Bool
    supportStatus::Bool
    supportPerf::Bool
    supportPolicy::Bool
    supportLimit::Bool
end

Base.show(io::IO, pd::PowerDomain) = print(io, "PowerDomain($(pd.type), ...)")

"Power information"
struct PowerInfo
    baseFrequency::Float64
    minFrequency::Float64
    turbo::TurboBoost
    hasRAPL::Bool
    powerUnit::Float64
    timeUnit::Float64
    uncoreMinFreq::Float64
    uncoreMaxFreq::Float64
    perfBias::Int
    domains::NTuple{5,PowerDomain}
end

# struct PowerData
#     domain::Int
#     before::UInt32
#     after::UInt32
# end

struct Likwid_Configuration
    configFileName::String
    topologyCfgFileName::String
    daemonPath::String
    groupPath::String
    daemonMode::LibLikwid.AccessMode
    maxNumThreads::Int
    maxNumNodes::Int
end

"Essential information about a performance group"
struct GroupInfoCompact
    name::String
    shortinfo::String
    longinfo::String
end

Base.show(io::IO, gi::GroupInfoCompact) = print(io, "$(gi.name) => $(gi.shortinfo)")

"Detailed information about a GPU device"
struct GpuDevice
    devid::Int
    numaNode::Int
    name::String
    short_name::String
    mem::Int
    compute_capability_major::Int
    compute_capability_minor::Int
    maxThreadsPerBlock::Int
    maxThreadsDim::NTuple{3,Int}
    maxGridSize::NTuple{3,Int}
    sharedMemPerBlock::Int
    totalConstantMemory::Int
    simdWidth::Int
    memPitch::Int
    regsPerBlock::Int
    clockRatekHz::Int
    textureAlign::Int
    surfaceAlign::Int
    l2Size::Int
    memClockRatekHz::Int
    pciBus::Int
    pciDev::Int
    pciDom::Int
    maxBlockRegs::Int
    numMultiProcs::Int
    maxThreadPerMultiProc::Int
    memBusWidth::Int
    unifiedAddrSpace::Bool
    ecc::Bool
    asyncEngines::Int
    mapHostMem::Bool
    integrated::Bool
end

Base.show(io::IO, d::GpuDevice) = print(io, "GpuDevice($(d.name), ...)")

"Topology information of GPUs"
struct GpuTopology
    numDevices::Int
    devices::Vector{GpuDevice}
end

# SHOW
const SHOW_TYPES = Union{
    CpuTopology,
    CpuInfo,
    NumaTopology,
    NumaNode,
    AffinityDomain,
    AffinityDomains,
    PowerInfo,
    PowerDomain,
    TurboBoost,
    Likwid_Configuration,
    GroupInfoCompact,
    GpuTopology,
    GpuDevice,
}

function Base.show(io::IO, mime::MIME{Symbol("text/plain")}, x::SHOW_TYPES)
    summary(io, x)
    println(io)
    T = typeof(x)
    nfields = length(fieldnames(T))
    for (i, field) in enumerate(fieldnames(T))
        char = i == nfields ? "└" : "├"
        xfield = getproperty(x, field)
        if (xfield isa AbstractVector || xfield isa NTuple) &&
           !(T in (NumaNode, AffinityDomain, TurboBoost, GpuDevice))
            print(io, char, " ", field, ": ... (", length(xfield), " elements)")
        elseif field == :totalMemory || field == :freeMemory
            print(io, char, " ", field, ": ", round(xfield / 1024 / 1024; digits=2), " GB")
        elseif T == GpuDevice && field == :mem
            print(
                io,
                char,
                " ",
                field,
                ": ",
                round(xfield / 1024 / 1024 / 1024; digits=2),
                " GB",
            )
        elseif field in (:baseFrequency, :minFrequency, :uncoreMinFreq, :uncoreMaxFreq)
            print(io, char, " ", field, ": ", xfield, " MHz")
        else
            print(io, char, " ", field, ": ", xfield)
        end
        i !== nfields && println(io)
    end
end

# function Base.show(io::IO, mime::MIME{Symbol("text/plain")}, x::LibLikwid.TimerData)
#     summary(io, x); println(io)
#     T = typeof(x)
#     nfields = length(fieldnames(T))
#     for (i,field) in enumerate(fieldnames(T))
#         char = i == nfields ? "└" : "├"
#         if getproperty(x, field) isa AbstractVector && !(T in (NumaNode, AffinityDomain))
#             print(io, char, " ", field, ": ... (", length(getproperty(x, field)), " elements)")
#         else
#             print(io, char, " ", field, ": ", getproperty(x, field))
#         end
#         i !== nfields && println(io)
#     end
# end

function Base.show(io::IO, x::LibLikwid.TimerData)
    return print(
        io, "TimerData(cycles start: $(x.start.int64), cycles stop: $(x.stop.int64))"
    )
end
Base.show(io::IO, x::LibLikwid.TscCounter) = print(io, x.int64, " (TscCounter)")


text:
include("run_sim.jl")

T = 1000000
p_E_range = 0.5
p_H_range = 0.002
λ_H_range = 0.2:0.2:(15-0.2)
λ_E_range = 5
#transitions_list = [tr_chains_d1, tr_chains_d50,tr_prioH, tr_prioE]
transitions_list = [tr_prioH, tr_prioE]

tic()
table1 = run_multiple_sim(λ_E_range, λ_H_range, p_E_range, p_H_range, transitions_list);
toc()
writetable("csv/fig_3.1.csv", table1)

T = 1000000
p_E_range = 0.5
p_H_range = 0.002
λ_H_range = 0.2:0.2:(15-0.2)
λ_E_range = 5
#transitions_list = [tr_chains_d1, tr_chains_d50,tr_prioH, tr_prioE]
transitions_list = [tr_prioH, tr_prioE]

tic()
table2 = run_multiple_sim(λ_E_range, λ_H_range, p_E_range, p_H_range, transitions_list, w_E = true);
toc()
writetable("csv/fig_3.2.csv", table2)
#
# T = 1000000
# p_E_range = 0.5
# p_H_range = 0.01
# λ_H_range = 0.1:0.1:(15-0.1)
# λ_E_range = 5
# transitions_list = [tr_prioH, tr_prioE, tr_random_prio]
#
# tic()
# table1 = run_multiple_sim(λ_E_range, λ_H_range, p_E_range, p_H_range, transitions_list);
# toc()
#
# writetable("csv/fig_3.1.csv", table1)
#
# T = 1000000
# p_E_range = 0.5
# p_H_range = 0.02
# λ_H_range = 0.5:0.5:(15-0.5)
# λ_E_range = 5
# transitions_list = [tr_chains_d1, tr_chains_d10,
#                     tr_prioH, tr_prioE]
# tic()
# table2 = run_multiple_sim(λ_E_range, λ_H_range, p_E_range, p_H_range, transitions_list, w_E = true);
# toc()
# writetable("csv/fig_3.2.csv", table2)
#
# T = 1000000
# p_E_range = 0.5
# p_H_range = 0.002
# λ_H_range = 0.2:0.2:(15-0.2)
# λ_E_range = 5
# #transitions_list = [tr_chains_d1, tr_chains_d50,tr_prioH, tr_prioE]
# transitions_list = [tr_prioH, tr_prioE]
#
# tic()
# table3 = run_multiple_sim(λ_E_range, λ_H_range, p_E_range, p_H_range, transitions_list, w_E = true);
# toc()
# writetable("csv/fig_3.3.csv", table3)


text:
# Non-conditional negative log-likelihood.
function negloglik(::Val{S}, kfun, params, pts, vals, w1) where{S}
  if !iszero(S)
    updatebuf_avx!(w1, Val(S), pts, pts, kfun, params)
  else
    updatebuf!(w1, pts, pts, kfun, params)
  end
  K   = cholesky!(Symmetric(w1))
  tmp = K.U'\vals # alloc 1, but this function only gets called once per nll.
  (logdet(K), sum(x->x^2, tmp))
end

# A negloglik where a mean and covariance matrix are provided.
function negloglik(K::Cholesky, mu, vals)
  mu .-= vals     # note now that mu is not the mean anymore.
  ldiv!(K.U', mu) # "mu" is now actually K.L\(vals - mu)
  (logdet(K), sum(x->x^2, mu))
end

# Conditional negative log-likelihood. Organized in such a way that you
# accurately compute the mean and covariance of the conditioned component and
# then pass those values to negloglik.
function cond_negloglik(::Val{S}, kfun, params, pts, vals,
                        cond_pts, cond_vals, w1, w2, w3) where{S}
  Z = promote_type(eltype(params), eltype(vals))
  # Update the buffers:
  if !iszero(S)
    updatebuf_avx!(w1, Val(S), cond_pts, cond_pts, kfun, params, skipltri=true)
    updatebuf_avx!(w2, Val(S), cond_pts, pts, kfun, params)
    updatebuf_avx!(w3, Val(S), pts, pts, kfun, params, skipltri=false)
  else
    updatebuf!(w1, cond_pts, cond_pts, kfun, params, skipltri=true)
    updatebuf!(w2, cond_pts, pts, kfun, params)
    updatebuf!(w3, pts, pts, kfun, params, skipltri=false)
  end
  # Rename the buffers/factorize the first one:
  K_cond_cond = cholesky!(Symmetric(w1)) # cov between cond-pts and cond-pts
  K_pts_condt = w2                       # cov between cond-pts and pts (transp)
  K_pts_pts   = w3                       # cov between pts and pts
  # conditional mean, the one remaining allocation in this function.
  mu = K_pts_condt'*(K_cond_cond\cond_vals)
  # Conditional covariance, reusing buffers liberally (see ?mul! in a REPL):
  ldiv!(K_cond_cond.U', K_pts_condt)
  mul!(K_pts_pts, adjoint(K_pts_condt), K_pts_condt, -one(Z), one(Z))
  sig = cholesky!(Symmetric(K_pts_pts))
  # Regular negloglik with computed conditional covariance and mean:
  negloglik(sig, mu, vals)
end

function nll(V::AbstractVecchiaConfig{H,D,F}, params::AbstractVector{T};
             memoize=false, execmode=ThreadedEx()) where{H,D,F,T}
  kernel = memoize ? MemoizedKernel(V.kernel, V.pts[1][1], params) : V.kernel
  Z      = promote_type(H,T)
  chsz   = V.chunksize
  ccsz   = chsz*V.blockrank
  out_logdet = zero(Z)
  out_qforms = zero(Z)
  scalar = (V isa ScalarVecchiaConfig) ? Val(D) : Val(0)
  @floop execmode for j in 1:length(V.condix)
    # Allocate work buffers the correct way:
    @init workcc = Array{Z}(undef, ccsz, ccsz)
    @init workcp = Array{Z}(undef, ccsz, chsz)
    @init workpp = Array{Z}(undef, chsz, chsz)
    # Get the likelihood points, data, and buffer for K(pts, pts) sorted out:
    pts  = V.pts[j]
    dat  = V.data[j]
    if scalar != Val(0)
      buf3 = view(workpp, 1:div(length(pts), D),  1:div(length(pts), D))
    else
      buf3 = view(workpp, 1:length(pts),  1:length(pts))
    end
    # if j==1, return just the regular negloglik for those points:
    if isone(j)
      (ldj, qfj) = negloglik(scalar, kernel, params, pts, dat, buf3)
    else
      # Now that we know the conditioning set isn't empty, get the conditioning
      # points sorted out:
      #
      # TODO (cg 2022/04/21 10:39): these are probably among the more
      # problematic allocations remaining. At some point it would be good to
      # write a more manual version that again uses pre-allocated work buffers
      # for this.
      idxs = V.condix[j]
      cpts = reduce(vcat, V.pts[idxs])
      cdat = reduce(vcat, V.data[idxs])
      # Get the buffers loaded up and properly shaped:
      if scalar != Val(0)
        buf1 = view(workcc, 1:div(length(cpts), D), 1:div(length(cpts), D))
        buf2 = view(workcp, 1:div(length(cpts), D), 1:div(length(pts),  D))
      else
        buf1 = view(workcc, 1:length(cpts), 1:length(cpts))
        buf2 = view(workcp, 1:length(cpts), 1:length(pts))
      end
      # Evaluate the conditional nll:
      (ldj, qfj) = cond_negloglik(scalar, kernel, params, pts, dat,
                                  cpts, cdat, buf1, buf2, buf3)
    end
    @reduce(out_logdet += ldj)
    @reduce(out_qforms += qfj)
  end
  # Return the final value, where all the rescaling has been saved to the end
  # (in the sense that the logdets need to be scaled by the number of samples,
  # and the whole thing needs to be divided by two).
  (out_logdet*size(V.data[1],2) + out_qforms)/2
end

# for simple debugging and testing.
function exact_nll(V::VecchiaConfig{H,D,F}, params::Vector{T}) where{H,D,T,F}
  pts = reduce(vcat, V.pts)
  dat = reduce(vcat, V.data)
  buf = Array{T}(undef, length(pts), length(pts))
  (ld, qf) = negloglik(Val(0), V.kernel, params, pts, dat, buf)
  0.5*(ld + qf)
end


text:
#! format: off
requires_initialization(::AbstractHydroFormulation) = false
requires_initialization(::AbstractHydroUnitCommitment) = true

get_variable_multiplier(_, ::Type{<:PSY.HydroGen}, ::AbstractHydroFormulation) = 1.0
get_expression_type_for_reserve(::ActivePowerReserveVariable, ::Type{<:PSY.HydroGen}, ::Type{<:PSY.Reserve{PSY.ReserveUp}}) = ActivePowerRangeExpressionUB
get_expression_type_for_reserve(::ActivePowerReserveVariable, ::Type{<:PSY.HydroGen}, ::Type{<:PSY.Reserve{PSY.ReserveDown}}) = ActivePowerRangeExpressionLB
get_expression_type_for_reserve(::ActivePowerReserveVariable, ::Type{<:PSY.HydroPumpedStorage}, ::Type{<:PSY.Reserve{PSY.ReserveUp}}) = ReserveRangeExpressionUB
get_expression_type_for_reserve(::ActivePowerReserveVariable, ::Type{<:PSY.HydroPumpedStorage}, ::Type{<:PSY.Reserve{PSY.ReserveDown}}) = ReserveRangeExpressionLB

########################### ActivePowerVariable, HydroGen #################################
get_variable_binary(::ActivePowerVariable, ::Type{<:PSY.HydroGen}, ::AbstractHydroFormulation) = false
get_variable_warm_start_value(::ActivePowerVariable, d::PSY.HydroGen, ::AbstractHydroFormulation) = PSY.get_active_power(d)
get_variable_lower_bound(::ActivePowerVariable, d::PSY.HydroGen, ::AbstractHydroFormulation) = PSY.get_active_power_limits(d).min
get_variable_lower_bound(::ActivePowerVariable, d::PSY.HydroGen, ::AbstractHydroUnitCommitment) = 0.0
get_variable_upper_bound(::ActivePowerVariable, d::PSY.HydroGen, ::AbstractHydroFormulation) = PSY.get_active_power_limits(d).max

############## ActivePowerVariable, HydroDispatchRunOfRiver ####################
get_variable_lower_bound(::ActivePowerVariable, d::PSY.HydroGen, ::HydroDispatchRunOfRiver) = 0.0

############## ReactivePowerVariable, HydroGen ####################
get_variable_binary(::ReactivePowerVariable, ::Type{<:PSY.HydroGen}, ::AbstractHydroFormulation) = false
get_variable_warm_start_value(::ReactivePowerVariable, d::PSY.HydroGen, ::AbstractHydroFormulation) = PSY.get_active_power(d)
get_variable_lower_bound(::ReactivePowerVariable, d::PSY.HydroGen, ::AbstractHydroFormulation) = PSY.get_active_power_limits(d).min
get_variable_upper_bound(::ReactivePowerVariable, d::PSY.HydroGen, ::AbstractHydroFormulation) = PSY.get_active_power_limits(d).max

############## EnergyVariable, HydroGen ####################
get_variable_binary(::EnergyVariable, ::Type{<:PSY.HydroGen}, ::AbstractHydroFormulation) = false
get_variable_warm_start_value(pv::EnergyVariable, d::PSY.HydroGen, ::AbstractHydroFormulation) = PSY.get_initial_storage(d)
get_variable_lower_bound(::EnergyVariable, d::PSY.HydroGen, ::AbstractHydroFormulation) = 0.0
get_variable_upper_bound(::EnergyVariable, d::PSY.HydroGen, ::AbstractHydroFormulation) = PSY.get_storage_capacity(d)

########################### EnergyVariableUp, HydroGen #################################
get_variable_binary(::EnergyVariableUp, ::Type{<:PSY.HydroGen}, ::AbstractHydroFormulation) = false
get_variable_warm_start_value(pv::EnergyVariableUp, d::PSY.HydroGen, ::AbstractHydroFormulation) = PSY.get_initial_storage(d).up
get_variable_lower_bound(::EnergyVariableUp, d::PSY.HydroGen, ::AbstractHydroFormulation) = 0.0
get_variable_upper_bound(::EnergyVariableUp, d::PSY.HydroGen, ::AbstractHydroFormulation) = PSY.get_storage_capacity(d).up

########################### EnergyVariableDown, HydroGen #################################
get_variable_binary(::EnergyVariableDown, ::Type{<:PSY.HydroGen}, ::AbstractHydroFormulation) = false
get_variable_warm_start_value(::EnergyVariableDown, d::PSY.HydroGen, ::AbstractHydroFormulation) = PSY.get_initial_storage(d).down
get_variable_lower_bound(::EnergyVariableDown, d::PSY.HydroGen, ::AbstractHydroFormulation) = 0.0
get_variable_upper_bound(::EnergyVariableDown, d::PSY.HydroGen, ::AbstractHydroFormulation) = PSY.get_storage_capacity(d).down

########################### ActivePowerInVariable, HydroGen #################################
get_variable_binary(::ActivePowerInVariable, ::Type{<:PSY.HydroGen}, ::AbstractHydroFormulation) = false
get_variable_lower_bound(::ActivePowerInVariable, d::PSY.HydroGen, ::AbstractHydroFormulation) = 0.0
get_variable_upper_bound(::ActivePowerInVariable, d::PSY.HydroGen, ::AbstractHydroFormulation) = nothing
get_variable_multiplier(::ActivePowerInVariable, d::Type{<:PSY.HydroGen}, ::AbstractHydroFormulation) = -1.0

########################### ActivePowerOutVariable, HydroGen #################################
get_variable_binary(::ActivePowerOutVariable, ::Type{<:PSY.HydroGen}, ::AbstractHydroFormulation) = false
get_variable_lower_bound(::ActivePowerOutVariable, d::PSY.HydroGen, ::AbstractHydroFormulation) = 0.0
get_variable_upper_bound(::ActivePowerOutVariable, d::PSY.HydroGen, ::AbstractHydroFormulation) = nothing
get_variable_multiplier(::ActivePowerOutVariable, d::Type{<:PSY.HydroGen}, ::AbstractHydroFormulation) = 1.0

############## OnVariable, HydroGen ####################
get_variable_binary(::OnVariable, ::Type{<:PSY.HydroGen}, ::AbstractHydroFormulation) = true
get_variable_warm_start_value(::OnVariable, d::PSY.HydroGen, ::AbstractHydroFormulation) = PSY.get_active_power(d) > 0 ? 1.0 : 0.0

############## WaterSpillageVariable, HydroGen ####################
get_variable_binary(::WaterSpillageVariable, ::Type{<:PSY.HydroGen}, ::AbstractHydroFormulation) = false
get_variable_lower_bound(::WaterSpillageVariable, d::PSY.HydroGen, ::AbstractHydroFormulation) = 0.0

############## ReservationVariable, HydroGen ####################
get_variable_binary(::ReservationVariable, ::Type{<:PSY.HydroGen}, ::AbstractHydroFormulation) = true
get_variable_binary(::ReservationVariable, ::Type{<:PSY.HydroPumpedStorage}, ::AbstractHydroFormulation) = true

############## EnergyShortageVariable, HydroGen ####################
get_variable_binary(::EnergyShortageVariable, ::Type{<:PSY.HydroGen}, ::AbstractHydroFormulation) = false
get_variable_lower_bound(::EnergyShortageVariable, d::PSY.HydroGen, ::AbstractHydroFormulation) = 0.0
get_variable_upper_bound(::EnergyShortageVariable, d::PSY.HydroGen, ::AbstractHydroFormulation) = PSY.get_storage_capacity(d)
get_variable_upper_bound(::EnergyShortageVariable, d::PSY.HydroPumpedStorage, ::AbstractHydroFormulation) = PSY.get_storage_capacity(d).up
############## EnergySurplusVariable, HydroGen ####################
get_variable_binary(::EnergySurplusVariable, ::Type{<:PSY.HydroGen}, ::AbstractHydroFormulation) = false
get_variable_upper_bound(::EnergySurplusVariable, d::PSY.HydroGen, ::AbstractHydroFormulation) = 0.0
get_variable_lower_bound(::EnergySurplusVariable, d::PSY.HydroGen, ::AbstractHydroFormulation) = - PSY.get_storage_capacity(d)
get_variable_lower_bound(::EnergySurplusVariable, d::PSY.HydroPumpedStorage, ::AbstractHydroFormulation) = - PSY.get_storage_capacity(d).up
########################### Parameter related set functions ################################
get_multiplier_value(::EnergyBudgetTimeSeriesParameter, d::PSY.HydroGen, ::AbstractHydroFormulation) = PSY.get_storage_capacity(d)
get_multiplier_value(::EnergyTargetTimeSeriesParameter, d::PSY.HydroGen, ::AbstractHydroFormulation) = PSY.get_storage_capacity(d)
get_multiplier_value(::InflowTimeSeriesParameter, d::PSY.HydroGen, ::AbstractHydroFormulation) = PSY.get_inflow(d) * PSY.get_conversion_factor(d)
get_multiplier_value(::OutflowTimeSeriesParameter, d::PSY.HydroGen, ::AbstractHydroFormulation) = PSY.get_outflow(d) * PSY.get_conversion_factor(d)
get_multiplier_value(::TimeSeriesParameter, d::PSY.HydroGen, ::AbstractHydroFormulation) = PSY.get_max_active_power(d)
get_multiplier_value(::TimeSeriesParameter, d::PSY.HydroGen, ::FixedOutput) = PSY.get_max_active_power(d)

get_parameter_multiplier(::VariableValueParameter, d::PSY.HydroGen, ::AbstractHydroFormulation) = 1.0
get_initial_parameter_value(::VariableValueParameter, d::PSY.HydroGen, ::AbstractHydroFormulation) = 1.0
get_expression_multiplier(::OnStatusParameter, ::ActivePowerRangeExpressionUB, d::PSY.HydroGen, ::AbstractHydroFormulation) = PSY.get_active_power_limits(d).max
get_expression_multiplier(::OnStatusParameter, ::ActivePowerRangeExpressionLB, d::PSY.HydroGen, ::AbstractHydroFormulation) = PSY.get_active_power_limits(d).min

#################### Initial Conditions for models ###############
initial_condition_default(::DeviceStatus, d::PSY.HydroGen, ::AbstractHydroFormulation) = PSY.get_status(d)
initial_condition_variable(::DeviceStatus, d::PSY.HydroGen, ::AbstractHydroFormulation) = OnVariable()
initial_condition_default(::DevicePower, d::PSY.HydroGen, ::AbstractHydroFormulation) = PSY.get_active_power(d)
initial_condition_variable(::DevicePower, d::PSY.HydroGen, ::AbstractHydroFormulation) = ActivePowerVariable()
initial_condition_default(::InitialEnergyLevel, d::PSY.HydroGen, ::AbstractHydroFormulation) = PSY.get_initial_storage(d)
initial_condition_variable(::InitialEnergyLevel, d::PSY.HydroGen, ::AbstractHydroFormulation) = EnergyVariable()
initial_condition_default(::InitialEnergyLevelUp, d::PSY.HydroGen, ::AbstractHydroFormulation) = PSY.get_initial_storage(d).up
initial_condition_variable(::InitialEnergyLevelUp, d::PSY.HydroGen, ::AbstractHydroFormulation) = EnergyVariableUp()
initial_condition_default(::InitialEnergyLevelDown, d::PSY.HydroGen, ::AbstractHydroFormulation) = PSY.get_initial_storage(d).down
initial_condition_variable(::InitialEnergyLevelDown, d::PSY.HydroGen, ::AbstractHydroFormulation) = EnergyVariableDown()
initial_condition_default(::InitialTimeDurationOn, d::PSY.HydroGen, ::AbstractHydroFormulation) = PSY.get_status(d) ? PSY.get_time_at_status(d) :  0.0
initial_condition_variable(::InitialTimeDurationOn, d::PSY.HydroGen, ::AbstractHydroFormulation) = OnVariable()
initial_condition_default(::InitialTimeDurationOff, d::PSY.HydroGen, ::AbstractHydroFormulation) = PSY.get_status(d) ? 0.0 : PSY.get_time_at_status(d)
initial_condition_variable(::InitialTimeDurationOff, d::PSY.HydroGen, ::AbstractHydroFormulation) = OnVariable()

########################Objective Function##################################################
proportional_cost(cost::Nothing, ::PSY.HydroGen, ::ActivePowerVariable, ::AbstractHydroFormulation)=0.0
proportional_cost(cost::PSY.OperationalCost, ::OnVariable, ::PSY.HydroGen, ::AbstractHydroFormulation)=PSY.get_fixed(cost)
proportional_cost(cost::PSY.StorageManagementCost, ::EnergySurplusVariable, ::PSY.HydroGen, ::AbstractHydroFormulation)=PSY.get_energy_surplus_cost(cost)
proportional_cost(cost::PSY.StorageManagementCost, ::EnergyShortageVariable, ::PSY.HydroGen, ::AbstractHydroFormulation)=PSY.get_energy_shortage_cost(cost)

objective_function_multiplier(::ActivePowerVariable, ::AbstractHydroFormulation)=OBJECTIVE_FUNCTION_POSITIVE
objective_function_multiplier(::ActivePowerOutVariable, ::AbstractHydroFormulation)=OBJECTIVE_FUNCTION_POSITIVE
objective_function_multiplier(::OnVariable, ::AbstractHydroFormulation)=OBJECTIVE_FUNCTION_POSITIVE
objective_function_multiplier(::EnergySurplusVariable, ::AbstractHydroFormulation)=OBJECTIVE_FUNCTION_NEGATIVE
objective_function_multiplier(::EnergyShortageVariable, ::AbstractHydroFormulation)=OBJECTIVE_FUNCTION_POSITIVE

sos_status(::PSY.HydroGen, ::AbstractHydroFormulation)=SOSStatusVariable.NO_VARIABLE
sos_status(::PSY.HydroGen, ::AbstractHydroUnitCommitment)=SOSStatusVariable.VARIABLE

variable_cost(::Nothing, ::ActivePowerVariable, ::PSY.HydroGen, ::AbstractHydroFormulation)=0.0
variable_cost(cost::PSY.OperationalCost, ::ActivePowerVariable, ::PSY.HydroGen, ::AbstractHydroFormulation)=PSY.get_variable(cost)
variable_cost(cost::PSY.OperationalCost, ::ActivePowerOutVariable, ::PSY.HydroGen, ::AbstractHydroFormulation)=PSY.get_variable(cost)

#! format: on

function get_initial_conditions_device_model(
    ::OperationModel,
    model::DeviceModel{T, <:AbstractHydroFormulation},
) where {T <: PSY.HydroEnergyReservoir}
    return model
end

function get_initial_conditions_device_model(
    ::OperationModel,
    ::DeviceModel{T, <:AbstractHydroFormulation},
) where {T <: PSY.HydroDispatch}
    return DeviceModel(PSY.HydroDispatch, HydroDispatchRunOfRiver)
end

function get_initial_conditions_device_model(
    ::OperationModel,
    ::DeviceModel{T, <:AbstractHydroFormulation},
) where {T <: PSY.HydroPumpedStorage}
    return DeviceModel(PSY.HydroPumpedStorage, HydroDispatchPumpedStorage)
end

function get_default_time_series_names(
    ::Type{<:PSY.HydroGen},
    ::Type{<:Union{FixedOutput, HydroDispatchRunOfRiver, HydroCommitmentRunOfRiver}},
)
    return Dict{Type{<:TimeSeriesParameter}, String}(
        ActivePowerTimeSeriesParameter => "max_active_power",
        ReactivePowerTimeSeriesParameter => "max_active_power",
    )
end

function get_default_time_series_names(
    ::Type{PSY.HydroEnergyReservoir},
    ::Type{<:Union{HydroCommitmentReservoirBudget, HydroDispatchReservoirBudget}},
)
    return Dict{Type{<:TimeSeriesParameter}, String}(
        EnergyBudgetTimeSeriesParameter => "hydro_budget",
    )
end

function get_default_time_series_names(
    ::Type{PSY.HydroEnergyReservoir},
    ::Type{<:Union{HydroDispatchReservoirStorage, HydroCommitmentReservoirStorage}},
)
    return Dict{Type{<:TimeSeriesParameter}, String}(
        EnergyTargetTimeSeriesParameter => "storage_target",
        InflowTimeSeriesParameter => "inflow",
    )
end

function get_default_time_series_names(
    ::Type{PSY.HydroPumpedStorage},
    ::Type{<:HydroDispatchPumpedStorage},
)
    return Dict{Type{<:TimeSeriesParameter}, String}(
        InflowTimeSeriesParameter => "inflow",
        OutflowTimeSeriesParameter => "outflow",
    )
end

function get_default_attributes(
    ::Type{T},
    ::Type{D},
) where {T <: PSY.HydroGen, D <: Union{FixedOutput, AbstractHydroFormulation}}
    return Dict{String, Any}("reservation" => false)
end

function get_default_attributes(
    ::Type{PSY.HydroPumpedStorage},
    ::Type{HydroDispatchPumpedStorage},
)
    return Dict{String, Any}("reservation" => true)
end

"""
Time series constraints
"""
function add_constraints!(
    container::OptimizationContainer,
    T::Type{ActivePowerVariableLimitsConstraint},
    U::Type{<:Union{VariableType, ExpressionType}},
    devices::IS.FlattenIteratorWrapper{V},
    model::DeviceModel{V, W},
    X::Type{<:PM.AbstractPowerModel},
) where {V <: PSY.HydroGen, W <: HydroDispatchRunOfRiver}
    if !has_semicontinuous_feedforward(model, U)
        add_range_constraints!(container, T, U, devices, model, X)
    end
    add_parameterized_upper_bound_range_constraints(
        container,
        ActivePowerVariableTimeSeriesLimitsConstraint,
        U,
        ActivePowerTimeSeriesParameter,
        devices,
        model,
        X,
    )
    return
end

function add_constraints!(
    container::OptimizationContainer,
    T::Type{ActivePowerVariableLimitsConstraint},
    U::Type{<:RangeConstraintLBExpressions},
    devices::IS.FlattenIteratorWrapper{V},
    model::DeviceModel{V, W},
    X::Type{<:PM.AbstractPowerModel},
) where {V <: PSY.HydroGen, W <: HydroDispatchRunOfRiver}
    if !has_semicontinuous_feedforward(model, U)
        add_range_constraints!(container, T, U, devices, model, X)
    end
    return
end

"""
Add semicontinuous range constraints for Hydro Unit Commitment formulation
"""
function add_constraints!(
    container::OptimizationContainer,
    T::Type{ActivePowerVariableLimitsConstraint},
    U::Type{<:Union{VariableType, <:RangeConstraintLBExpressions}},
    devices::IS.FlattenIteratorWrapper{V},
    model::DeviceModel{V, W},
    X::Type{<:PM.AbstractPowerModel},
) where {V <: PSY.HydroGen, W <: HydroCommitmentRunOfRiver}
    add_semicontinuous_range_constraints!(container, T, U, devices, model, X)
    return
end

function add_constraints!(
    container::OptimizationContainer,
    T::Type{ActivePowerVariableLimitsConstraint},
    U::Type{<:Union{VariableType, ExpressionType}},
    devices::IS.FlattenIteratorWrapper{V},
    model::DeviceModel{V, W},
    X::Type{<:PM.AbstractPowerModel},
) where {V <: PSY.HydroGen, W <: HydroCommitmentRunOfRiver}
    add_semicontinuous_range_constraints!(container, T, U, devices, model, X)
    add_parameterized_upper_bound_range_constraints(
        container,
        ActivePowerVariableTimeSeriesLimitsConstraint,
        U,
        ActivePowerTimeSeriesParameter,
        devices,
        model,
        X,
    )
    return
end

"""
Min and max reactive Power Variable limits
"""
function get_min_max_limits(
    x::PSY.HydroGen,
    ::Type{<:ReactivePowerVariableLimitsConstraint},
    ::Type{<:AbstractHydroFormulation},
)
    return PSY.get_reactive_power_limits(x)
end

"""
Min and max active Power Variable limits
"""
function get_min_max_limits(
    x::PSY.HydroGen,
    ::Type{<:ActivePowerVariableLimitsConstraint},
    ::Type{<:AbstractHydroFormulation},
)
    return PSY.get_active_power_limits(x)
end

function get_min_max_limits(
    x::PSY.HydroGen,
    ::Type{<:ActivePowerVariableLimitsConstraint},
    ::Type{HydroDispatchRunOfRiver},
)
    return (min=0.0, max=PSY.get_max_active_power(x))
end

"""
Add power variable limits constraints for hydro unit commitment formulation
"""
function add_constraints!(
    container::OptimizationContainer,
    T::Type{<:PowerVariableLimitsConstraint},
    U::Type{<:Union{VariableType, ExpressionType}},
    devices::IS.FlattenIteratorWrapper{V},
    model::DeviceModel{V, W},
    X::Type{<:PM.AbstractPowerModel},
) where {V <: PSY.HydroGen, W <: AbstractHydroUnitCommitment}
    add_semicontinuous_range_constraints!(container, T, U, devices, model, X)
    return
end

"""
Add power variable limits constraints for hydro dispatch formulation
"""
function add_constraints!(
    container::OptimizationContainer,
    T::Type{<:PowerVariableLimitsConstraint},
    U::Type{<:Union{VariableType, ExpressionType}},
    devices::IS.FlattenIteratorWrapper{V},
    model::DeviceModel{V, W},
    X::Type{<:PM.AbstractPowerModel},
) where {V <: PSY.HydroGen, W <: AbstractHydroDispatchFormulation}
    if !has_semicontinuous_feedforward(model, U)
        add_range_constraints!(container, T, U, devices, model, X)
    end
    return
end

"""
Add input power variable limits constraints for hydro dispatch formulation
"""
function add_constraints!(
    container::OptimizationContainer,
    T::Type{InputActivePowerVariableLimitsConstraint},
    U::Type{<:Union{VariableType, ExpressionType}},
    devices::IS.FlattenIteratorWrapper{V},
    model::DeviceModel{V, W},
    X::Type{<:PM.AbstractPowerModel},
) where {V <: PSY.HydroPumpedStorage, W <: AbstractHydroReservoirFormulation}
    if get_attribute(model, "reservation")
        add_reserve_range_constraints!(container, T, U, devices, model, X)
    else
        if !has_semicontinuous_feedforward(model, U)
            add_range_constraints!(container, T, U, devices, model, X)
        end
    end
    return
end

"""
Add output power variable limits constraints for hydro dispatch formulation
"""
function add_constraints!(
    container::OptimizationContainer,
    T::Type{<:PowerVariableLimitsConstraint},
    U::Type{<:Union{VariableType, ExpressionType}},
    devices::IS.FlattenIteratorWrapper{V},
    model::DeviceModel{V, W},
    X::Type{<:PM.AbstractPowerModel},
) where {V <: PSY.HydroPumpedStorage, W <: AbstractHydroReservoirFormulation}
    if get_attribute(model, "reservation")
        add_reserve_range_constraints!(container, T, U, devices, model, X)
    else
        if !has_semicontinuous_feedforward(model, U)
            add_range_constraints!(container, T, U, devices, model, X)
        end
    end
    return
end

"""
Min and max output active power variable limits for hydro dispatch pumped storage
"""
function get_min_max_limits(
    x::PSY.HydroGen,
    ::Type{<:OutputActivePowerVariableLimitsConstraint},
    ::Type{HydroDispatchPumpedStorage},
)
    return PSY.get_active_power_limits(x)
end

"""
Min and max input active power variable limits for hydro dispatch pumped storage
"""
function get_min_max_limits(
    x::PSY.HydroGen,
    ::Type{<:InputActivePowerVariableLimitsConstraint},
    ::Type{HydroDispatchPumpedStorage},
)
    return PSY.get_active_power_limits_pump(x)
end

######################## Energy balance constraints ############################

"""
This function defines the constraints for the water level (or state of charge)
for the Hydro Reservoir.
"""
function add_constraints!(
    container::OptimizationContainer,
    ::Type{EnergyBalanceConstraint},
    devices::IS.FlattenIteratorWrapper{V},
    model::DeviceModel{V, W},
    ::Type{X},
) where {
    V <: PSY.HydroEnergyReservoir,
    W <: AbstractHydroFormulation,
    X <: PM.AbstractPowerModel,
}
    time_steps = get_time_steps(container)
    resolution = get_resolution(container)
    fraction_of_hour = Dates.value(Dates.Minute(resolution)) / MINUTES_IN_HOUR
    names = [PSY.get_name(x) for x in devices]
    initial_conditions = get_initial_condition(container, InitialEnergyLevel(), V)
    energy_var = get_variable(container, EnergyVariable(), V)
    power_var = get_variable(container, ActivePowerVariable(), V)
    spillage_var = get_variable(container, WaterSpillageVariable(), V)

    constraint = add_constraints_container!(
        container,
        EnergyBalanceConstraint(),
        V,
        names,
        time_steps,
    )
    param = get_parameter_array(container, InflowTimeSeriesParameter(), V)
    multiplier = get_parameter_multiplier_array(container, InflowTimeSeriesParameter(), V)

    for ic in initial_conditions
        device = get_component(ic)
        name = PSY.get_name(device)
        constraint[name, 1] = JuMP.@constraint(
            container.JuMPmodel,
            energy_var[name, 1] ==
            get_value(ic) - power_var[name, 1] * fraction_of_hour -
            spillage_var[name, 1] * fraction_of_hour +
            param[name, 1] * multiplier[name, 1]
        )

        for t in time_steps[2:end]
            constraint[name, t] = JuMP.@constraint(
                container.JuMPmodel,
                energy_var[name, t] ==
                energy_var[name, t - 1] + param[name, t] * multiplier[name, t] -
                power_var[name, t] * fraction_of_hour -
                spillage_var[name, t] * fraction_of_hour
            )
        end
    end
    return
end

"""
This function defines the constraints for the water level (or state of charge)
for the HydroPumpedStorage.
"""
function add_constraints!(
    container::OptimizationContainer,
    ::Type{EnergyCapacityUpConstraint},
    devices::IS.FlattenIteratorWrapper{V},
    model::DeviceModel{V, W},
    ::Type{X},
) where {
    V <: PSY.HydroPumpedStorage,
    W <: AbstractHydroFormulation,
    X <: PM.AbstractPowerModel,
}
    time_steps = get_time_steps(container)
    resolution = get_resolution(container)
    fraction_of_hour = Dates.value(Dates.Minute(resolution)) / MINUTES_IN_HOUR
    names = [PSY.get_name(x) for x in devices]
    initial_conditions = get_initial_condition(container, InitialEnergyLevelUp(), V)

    energy_var = get_variable(container, EnergyVariableUp(), V)
    powerin_var = get_variable(container, ActivePowerInVariable(), V)
    powerout_var = get_variable(container, ActivePowerOutVariable(), V)
    spillage_var = get_variable(container, WaterSpillageVariable(), V)

    constraint = add_constraints_container!(
        container,
        EnergyCapacityUpConstraint(),
        V,
        names,
        time_steps,
    )
    param = get_parameter_array(container, InflowTimeSeriesParameter(), V)
    multiplier = get_parameter_multiplier_array(container, InflowTimeSeriesParameter(), V)

    for ic in initial_conditions
        device = get_component(ic)
        efficiency = PSY.get_pump_efficiency(device)
        name = PSY.get_name(device)
        constraint[name, 1] = JuMP.@constraint(
            container.JuMPmodel,
            energy_var[name, 1] ==
            get_value(ic) +
            (
                powerin_var[name, 1] * efficiency - spillage_var[name, 1] -
                powerout_var[name, 1]
            ) * fraction_of_hour +
            param[name, 1] * multiplier[name, 1]
        )

        for t in time_steps[2:end]
            constraint[name, t] = JuMP.@constraint(
                container.JuMPmodel,
                energy_var[name, t] ==
                energy_var[name, t - 1] +
                param[name, t] * multiplier[name, t] +
                (powerin_var[name, 1] - powerout_var[name, t] - spillage_var[name, t]) *
                fraction_of_hour
            )
        end
    end
    return
end

"""
Add energy capacity down constraints for hydro pumped storage
"""
function add_constraints!(
    container::OptimizationContainer,
    ::Type{EnergyCapacityDownConstraint},
    devices::IS.FlattenIteratorWrapper{V},
    model::DeviceModel{V, W},
    ::Type{X},
) where {
    V <: PSY.HydroPumpedStorage,
    W <: AbstractHydroFormulation,
    X <: PM.AbstractPowerModel,
}
    time_steps = get_time_steps(container)
    resolution = get_resolution(container)
    fraction_of_hour = Dates.value(Dates.Minute(resolution)) / MINUTES_IN_HOUR
    names = [PSY.get_name(x) for x in devices]
    initial_conditions = get_initial_condition(container, InitialEnergyLevelDown(), V)

    energy_var = get_variable(container, EnergyVariableDown(), V)
    powerin_var = get_variable(container, ActivePowerInVariable(), V)
    powerout_var = get_variable(container, ActivePowerOutVariable(), V)
    spillage_var = get_variable(container, WaterSpillageVariable(), V)

    constraint = add_constraints_container!(
        container,
        EnergyCapacityDownConstraint(),
        V,
        names,
        time_steps,
    )

    param = get_parameter_array(container, OutflowTimeSeriesParameter(), V)
    multiplier = get_parameter_multiplier_array(container, OutflowTimeSeriesParameter(), V)

    for ic in initial_conditions
        device = get_component(ic)
        efficiency = PSY.get_pump_efficiency(device)
        name = PSY.get_name(device)
        constraint[name, 1] = JuMP.@constraint(
            container.JuMPmodel,
            energy_var[name, 1] ==
            get_value(ic) -
            (
                spillage_var[name, 1] + powerout_var[name, 1] -
                powerin_var[name, 1] / efficiency
            ) * fraction_of_hour - param[name, 1] * multiplier[name, 1]
        )

        for t in time_steps[2:end]
            constraint[name, t] = JuMP.@constraint(
                container.JuMPmodel,
                energy_var[name, t] ==
                energy_var[name, t - 1] - param[name, t] * multiplier[name, t] +
                (
                    powerout_var[name, 1] - powerin_var[name, t] / efficiency +
                    spillage_var[name, t]
                ) * fraction_of_hour
            )
        end
    end
    return
end

"""
Add energy target constraints for hydro gen
"""
function add_constraints!(
    container::OptimizationContainer,
    ::Type{EnergyTargetConstraint},
    devices::IS.FlattenIteratorWrapper{V},
    model::DeviceModel{V, W},
    ::Type{X},
) where {V <: PSY.HydroGen, W <: AbstractHydroFormulation, X <: PM.AbstractPowerModel}
    time_steps = get_time_steps(container)
    resolution = get_resolution(container)
    set_name = [PSY.get_name(d) for d in devices]
    constraint = add_constraints_container!(
        container,
        EnergyTargetConstraint(),
        V,
        set_name,
        time_steps,
    )

    e_var = get_variable(container, EnergyVariable(), V)
    shortage_var = get_variable(container, EnergyShortageVariable(), V)
    surplus_var = get_variable(container, EnergySurplusVariable(), V)
    param = get_parameter_array(container, EnergyTargetTimeSeriesParameter(), V)
    multiplier =
        get_parameter_multiplier_array(container, EnergyTargetTimeSeriesParameter(), V)

    for d in devices
        name = PSY.get_name(d)
        cost_data = PSY.get_operation_cost(d)
        if isa(cost_data, PSY.StorageManagementCost)
            shortage_cost = PSY.get_energy_shortage_cost(cost_data)
        else
            @debug "Data for device $name doesn't contain shortage costs"
            shortage_cost = 0.0
        end

        if shortage_cost == 0.0
            @warn(
                "Device $name has energy shortage cost set to 0.0, as a result the model will turnoff the EnergyShortageVariable to avoid infeasible/unbounded problem."
            )
            JuMP.delete_upper_bound.(shortage_var[name, :])
            JuMP.set_upper_bound.(shortage_var[name, :], 0.0)
        end
        for t in time_steps
            constraint[name, t] = JuMP.@constraint(
                container.JuMPmodel,
                e_var[name, t] + shortage_var[name, t] + surplus_var[name, t] ==
                multiplier[name, t] * param[name, t]
            )
        end
    end
    return
end

##################################### Water/Energy Budget Constraint ############################
"""
This function define the budget constraint for the
active power budget formulation.

`` sum(P[t]) <= Budget ``
"""

function add_constraints!(
    container::OptimizationContainer,
    ::Type{EnergyBudgetConstraint},
    devices::IS.FlattenIteratorWrapper{V},
    model::DeviceModel{V, W},
    ::Type{X},
) where {V <: PSY.HydroGen, W <: AbstractHydroFormulation, X <: PM.AbstractPowerModel}
    time_steps = get_time_steps(container)
    resolution = get_resolution(container)
    set_name = [PSY.get_name(d) for d in devices]
    constraint =
        add_constraints_container!(container, EnergyBudgetConstraint(), V, set_name)

    variable_out = get_variable(container, ActivePowerVariable(), V)
    param = get_parameter_array(container, EnergyBudgetTimeSeriesParameter(), V)
    multiplier =
        get_parameter_multiplier_array(container, EnergyBudgetTimeSeriesParameter(), V)

    for d in devices
        name = PSY.get_name(d)
        constraint[name] = JuMP.@constraint(
            container.JuMPmodel,
            sum([variable_out[name, t] for t in time_steps]) <= sum([multiplier[name, t] * param[name, t] for t in time_steps])
        )
    end
    return
end

##################################### Auxillary Variables ############################

function calculate_aux_variable_value!(
    container::OptimizationContainer,
    ::AuxVarKey{EnergyOutput, T},
    system::PSY.System,
) where {T <: PSY.HydroGen}
    devices = get_available_components(T, system)
    time_steps = get_time_steps(container)

    p_variable_results = get_variable(container, ActivePowerVariable(), T)
    aux_variable_container = get_aux_variable(container, EnergyOutput(), T)
    for d in devices, t in time_steps
        name = PSY.get_name(d)
        min = PSY.get_active_power_limits(d).min
        aux_variable_container[name, t] = jump_value(p_variable_results[name, t])
    end

    return
end

##################################### Hydro generation cost ############################
function objective_function!(
    container::OptimizationContainer,
    devices::IS.FlattenIteratorWrapper{T},
    ::DeviceModel{T, U},
    ::Type{<:PM.AbstractPowerModel},
) where {T <: PSY.HydroGen, U <: AbstractHydroUnitCommitment}
    add_variable_cost!(container, ActivePowerVariable(), devices, U())
    add_proportional_cost!(container, OnVariable(), devices, U())
    return
end

function objective_function!(
    container::OptimizationContainer,
    devices::IS.FlattenIteratorWrapper{PSY.HydroPumpedStorage},
    ::DeviceModel{PSY.HydroPumpedStorage, T},
    ::Type{<:PM.AbstractPowerModel},
) where {T <: HydroDispatchPumpedStorage}
    add_variable_cost!(container, ActivePowerOutVariable(), devices, T())
    return
end

function objective_function!(
    container::OptimizationContainer,
    devices::IS.FlattenIteratorWrapper{T},
    ::DeviceModel{T, U},
    ::Type{<:PM.AbstractPowerModel},
) where {T <: PSY.HydroGen, U <: AbstractHydroDispatchFormulation}
    add_variable_cost!(container, ActivePowerVariable(), devices, U())
    return
end

function objective_function!(
    container::OptimizationContainer,
    devices::IS.FlattenIteratorWrapper{T},
    ::DeviceModel{T, U},
    ::Type{<:PM.AbstractPowerModel},
) where {
    T <: PSY.HydroPumpedStorage,
    U <: Union{HydroDispatchReservoirStorage, HydroDispatchReservoirBudget},
}
    add_variable_cost!(container, ActivePowerOutVariable(), devices, U())
    add_proportional_cost!(container, EnergySurplusVariable(), devices, U())
    add_proportional_cost!(container, EnergyShortageVariable(), devices, U())
    return
end

function objective_function!(
    container::OptimizationContainer,
    devices::IS.FlattenIteratorWrapper{T},
    ::DeviceModel{T, U},
    ::Type{<:PM.AbstractPowerModel},
) where {T <: PSY.HydroEnergyReservoir, U <: HydroDispatchReservoirStorage}
    add_variable_cost!(container, ActivePowerVariable(), devices, U())
    add_proportional_cost!(container, EnergySurplusVariable(), devices, U())
    add_proportional_cost!(container, EnergyShortageVariable(), devices, U())
    return
end

function objective_function!(
    container::OptimizationContainer,
    devices::IS.FlattenIteratorWrapper{T},
    ::DeviceModel{T, U},
    ::Type{<:PM.AbstractPowerModel},
) where {T <: PSY.HydroEnergyReservoir, U <: HydroCommitmentReservoirStorage}
    add_variable_cost!(container, ActivePowerVariable(), devices, U())
    add_proportional_cost!(container, EnergySurplusVariable(), devices, U())
    add_proportional_cost!(container, EnergyShortageVariable(), devices, U())
    return
end