rot(n::Int) = M  -> reduce((M, _) -> rotr90(M), 1:n; init = M)
flip(M::Matrix)::Matrix = M[end:-1:1, :]

const TRANSFORMATIONS = [[rot(i) for i in 0:3];
                         [rot(i) ∘ flip for i in 0:3]]

function parse_input(file::AbstractString)
    tiles = open(file, "r") do io
        tiles = []
        while !eof(io)
            push!(tiles, parse_tile(io))
        end
        return tiles
    end

    tilemap = Dict()
    sidemap = Dict()
    for t in tiles, (tr_id, tr) in enumerate(TRANSFORMATIONS)
        t_rot_id = (id = t.id, tr = tr_id)
        t_rot = tr(t.tile)
        sides = tile_sides(t_rot)
        tilemap[t_rot_id] = (; inside = t_rot[2:end - 1, 2:end - 1], sides...)
        for side in pairs(sides)
            push!(get!(sidemap, side, NamedTuple[]), t_rot_id)
        end
    end
    return (tilemap, sidemap)
end

function parse_tile(io::IO)
    id_line = readline(io)
    id_m = match(r"^Tile ([0-9]+):$", id_line)
    @assert id_m !== nothing id_line
    id = parse(Int, id_m.captures[1])
    tile = hcat([split(readline(io), "") for _ in 1:10]...) |> permutedims
    @assert eof(io) || isempty(readline(io))
    return (; id, tile)
end

function tile_sides(tile::Matrix)
    side(v) = Symbol(join(v, ""))

    return (top    = side(tile[1, 1:end]),
            bottom = side(tile[end, 1:end]),
            left   = side(tile[1:end, 1]),
            right  = side(tile[1:end, end]))
end

up(ci::CartesianIndex) = ci - CartesianIndex(1, 0)
left(ci::CartesianIndex) = ci - CartesianIndex(0, 1)

function arrange(tiles, sides)::Matrix
    w = tiles |> keys .|> first |> unique |> length |> sqrt |> Int
    tileset = (unmatched_tops  = unmatched(tiles, :top),
               matched_bottoms = matched(tiles, :bottom),
               unmatched_lefts = unmatched(tiles, :left),
               matched_rights  = matched(tiles, :right),
               all             = keys(tiles),
               dict            = tiles)
    arr = Matrix{Any}(nothing, w, w)
    ok = arrange!(arr, keys(arr), [], tileset, sides)
    @assert ok
    println("found")
    return arr
end

unmatched(tiles, side::Symbol) = setdiff(keys(tiles), matched(tiles, side))

function matched(tiles, side::Symbol)
    opp_side::Symbol = opposite(side)
    return filter(keys(tiles)) do i
        tiles[i][side] ∈ [t[opp_side] for (j, t) in tiles if j.id != i.id]
    end
end

opposite(side::Symbol)::Symbol = opposite(Val(side))
opposite(::Val{:top}) = :bottom
opposite(::Val{:bottom}) = :top
opposite(::Val{:left}) = :right
opposite(::Val{:right}) = :left

function arrange!(arr::Matrix, cis, used, tileset, sides)::Bool
    isempty(cis) && return true
    (ci, cis...) = cis

    suitable = filter(∉(used) ∘ first, tileset.all)
    if ci[1] == 1
        intersect!(suitable, tileset.unmatched_tops)
    else
        border = tileset.dict[arr[up(ci)]].bottom
        intersect!(suitable, sides[:top => border])
    end
    if ci[1] < size(arr, 1)
        intersect!(suitable, tileset.matched_bottoms)
    end
    if ci[2] == 1
        intersect!(suitable, tileset.unmatched_lefts)
    else
        border = tileset.dict[arr[left(ci)]].right
        intersect!(suitable, sides[:left => border])
    end
    if ci[2] < size(arr, 2)
        intersect!(suitable, tileset.matched_rights)
    end

    for t in suitable
        arr[ci] = t
        arrange!(arr, cis, [used; t.id], tileset, sides) && return true
    end
    print(".")                  # backtrack
    return false
end

q1(arr) = map(t -> t.id, corners(arr)) |> prod

corners(M::Matrix)::Matrix = M[[1, end], [1, end]]

function monster(w::Int)::String
    m = ["..................#."
         "#....##....##....###"
         ".#..#..#..#..#..#..."]
    return join(m, repeat('.', w - length(m[1])))
end

function q2(arr, tiles)::Int
    return count(==('#'), mask_monsters(arr, tiles))
end

function mask_monsters(arr, tiles)::String
    images = assemble(arr, tiles)
    w = size(arr, 2) * size(first(tiles).second.inside, 2)
    m = monster(w)
    mre = Regex(m)
    image = images[only(findall(img -> count(mre, img) > 0, images))]
    mask = findall(==('#'), m)
    vimage = split(image, "")
    for i in findall(mre, image; overlap = true) .|> first
        @. vimage[i + mask - 1] = "O"
    end
    return join(vimage, "")
end

function assemble(arr, tiles)::Vector{String}
    grid = size(first(tiles).second.inside)
    image = fill("", size(arr) .* grid)
    for (i, t_id) in pairs(arr)
        tl = (i - one(i)) * grid + one(i)
        br  = i * grid
        image[tl:br] .= tiles[t_id].inside
    end
    return map(tr -> join(tr(image), ""), TRANSFORMATIONS)
end

Base.:*(ci::CartesianIndex, t::Tuple) = CartesianIndex(ci.I .* t)

let (tiles, sides) = parse_input("day20_test.in"),
    arr = arrange(tiles, sides)

    @assert q1(arr) == 20899048083289
    @assert q2(arr, tiles) == 273
end

(tiles, sides) = parse_input("day20.in")
arr = arrange(tiles, sides)
println("Q1: ", q1(arr))
println("Q2: ", q2(arr, tiles))


module SampleData

using Celeste: Model, DeterministicVI
import Celeste: Synthetic
import Celeste.SDSSIO: RunCamcolField, load_field_images, load_field_catalog,
                       SDSSDataSet, SDSSBackground, SDSSPSFMap

using Distributions
using StaticArrays
import WCS
import FITSIO
import DataFrames

export empty_model_params,
       sample_ce, perturb_params,
       sample_star_fluxes, sample_galaxy_fluxes,
       gen_sample_star_dataset, gen_sample_galaxy_dataset,
       gen_two_body_dataset, gen_three_body_dataset, gen_n_body_dataset,
       make_elbo_args, true_star_init

const sample_star_fluxes = [
    4.451805E+03,1.491065E+03,2.264545E+03,2.027004E+03,1.846822E+04]
const sample_galaxy_fluxes = [
    1.377666E+01, 5.635334E+01, 1.258656E+02,
    1.884264E+02, 2.351820E+02] * 100  # 1x wasn't bright enough

# A world coordinate system where the world and pixel coordinates are the same.
const wcs_id = WCS.WCSTransform(2,
                    cd = Float64[1 0; 0 1],
                    ctype = ["none", "none"],
                    crpix = Float64[1, 1],
                    crval = Float64[1, 1]);

# globals to hold already-loaded test images and catalogs
const DATADIR = joinpath(Pkg.dir("Celeste"), "test", "data")
const DATASET = SDSSDataSet(DATADIR)
const SDSS_FIELD_IMAGES =
    Dict{RunCamcolField,Vector{Image{SDSSBackground,SDSSPSFMap}}}()
const SDSS_FIELD_CATALOGS = Dict{RunCamcolField, Vector{CatalogEntry}}()

# Fetch sdss data for field `rcf` if not already present. See DATADIR/Makefile
# for possible make targets.
function fetch_sdss_data(rcf, make_target)
    # ensure images and catalog are downloaded
    wd = pwd()
    cd(DATADIR)
    make_output = readstring(`make $(make_target) RUN=$(rcf.run) CAMCOL=$(rcf.camcol) FIELD=$(rcf.field)`)

    # only print something if we actually ran a download
    if !startswith(make_output, "make: Nothing to be done for")
        print(make_output)
    end

    cd(wd)
end
fetch_sdss_data(run, camcol, field, make_target) =
    fetch_sdss_data(RunCamcolField(run, camcol, field), make_target)

"""
    get_sdss_images(run, camcol, field) -> Vector{<:Image}

Return lazily-loaded images from the given SDSS field.
"""
function get_sdss_images(run, camcol, field)
    rcf = RunCamcolField(run, camcol, field)
    if !haskey(SDSS_FIELD_IMAGES, rcf)
        fetch_sdss_data(rcf, "all")
        SDSS_FIELD_IMAGES[rcf] = load_field_images(DATASET, rcf)
    end
    return SDSS_FIELD_IMAGES[rcf]
end

"""
    get_sdss_catalog(run, camcol, field) -> Vector{CatalogEntry}

Return lazily-loaded catalog for the given SDSS field.
"""
function get_sdss_catalog(run, camcol, field)
    rcf = RunCamcolField(run, camcol, field)
    if !haskey(SDSS_FIELD_CATALOGS, rcf)
        fetch_sdss_data(rcf, "photoobj")
        SDSS_FIELD_CATALOGS[rcf] = load_field_catalog(DATASET, rcf)
    end
    return SDSS_FIELD_CATALOGS[rcf]
end

"""
Turn a images and vector of catalog entries into elbo arguments
that can be used with Celeste.
"""
function make_elbo_args(images::Vector{<:Image},
                        catalog::Vector{CatalogEntry};
                        active_source=-1,
                        patch_radius_pix::Float64=NaN,
                        include_kl=true)
    patches = Model.get_sky_patches(images,
                                    catalog,
                                    radius_override_pix=patch_radius_pix)
    S = length(catalog)
    active_sources = active_source > 0 ? [active_source] :
                                          S <= 3 ? collect(1:S) : [1,2,3]
    ElboArgs(images, patches, active_sources; include_kl=include_kl)
end

function empty_model_params(S::Int)
    vp = [DeterministicVI.generic_init_source([ 0., 0. ]) for s in 1:S]
    ElboArgs(Image[],
             vp,
             Matrix{ImagePatch}(S, 0),
             collect(1:S))
end

function sample_ce(pos, is_star::Bool)
    CatalogEntry(pos, is_star, sample_star_fluxes, sample_galaxy_fluxes,
        0.1, 0.7, pi/4, 4.0)
end

# for testing away from the truth, where derivatives != 0
function perturb_params(vp)
    for vs in vp
        vs[ids.is_star] = [ 0.4, 0.6 ]
        vs[ids.pos[1]] += .8
        vs[ids.pos[2]] -= .7
        vs[ids.flux_loc] -= log(10)
        vs[ids.flux_scale] *= 25.
        vs[ids.gal_frac_dev] += 0.05
        vs[ids.gal_axis_ratio] += 0.05
        vs[ids.gal_angle] += pi/10
        vs[ids.gal_radius_px] *= 1.2
        vs[ids.color_mean] += 0.5
        vs[ids.color_var] =  1e-1
    end
end

function cropped_sample_images(new_H, new_W)
    srand(1)
    sample_images = get_sdss_images(3900, 6, 269)
    images = deepcopy(sample_images)

    for img in images
        img.H = new_H
        img.W = new_W
        img.pixels = zeros(Float32, new_H, new_W)
        img.nelec_per_nmgy = img.nelec_per_nmgy[1:new_H]
        img.wcs = wcs_id
    end

    return images
end

function gen_sample_star_dataset(; perturb=true)
    images = cropped_sample_images(20, 23)
    catalog = [sample_ce([10.1, 12.2], true),]
    Synthetic.gen_images!(images, catalog)
    ea = make_elbo_args(images, catalog)

    vp = Vector{Float64}[DeterministicVI.catalog_init_source(ce) for ce in catalog]
    if perturb
        perturb_params(vp)
    end

    ea, vp, catalog
end

function gen_sample_galaxy_dataset(; perturb=true, include_kl=true)
    images = cropped_sample_images(20, 23)
    catalog = [sample_ce([8.5, 9.6], false),]
    Synthetic.gen_images!(images, catalog)
    ea = make_elbo_args(images, catalog; include_kl=include_kl)

    vp = Vector{Float64}[DeterministicVI.catalog_init_source(ce) for ce in catalog]
    if perturb
        perturb_params(vp)
    end

    ea, vp, catalog
end

# A small two-body dataset for quick unit testing.  These objects
# will be too close to be identifiable.
function gen_two_body_dataset(; perturb=true)
    images = cropped_sample_images(20, 23)
    catalog = [
        sample_ce([4.5, 3.6], false),
        sample_ce([10.1, 12.1], true)
    ]
    Synthetic.gen_images!(images, catalog)
    ea = make_elbo_args(images, catalog)

    vp = Vector{Float64}[DeterministicVI.catalog_init_source(ce) for ce in catalog]
    if perturb
        perturb_params(vp)
    end

    ea, vp, catalog
end

function gen_three_body_dataset(; perturb=true)
    images = cropped_sample_images(112, 238)
    catalog = [
        sample_ce([4.5, 3.6], false),
        sample_ce([60.1, 82.2], true),
        sample_ce([71.3, 100.4], false),
    ];
    Synthetic.gen_images!(images, catalog);
    ea = make_elbo_args(images, catalog);

    vp = Vector{Float64}[DeterministicVI.catalog_init_source(ce) for ce in catalog]
    if perturb
        perturb_params(vp)
    end

    ea, vp, catalog
end

"""
Generate a large dataset with S randomly placed bodies and non-constant
background.
"""
function gen_n_body_dataset(S::Int; patch_pixel_radius=20., perturb=true)
    images = cropped_sample_images(900, 1000)

    fluxes = [4.451805E+03,1.491065E+03,2.264545E+03,2.027004E+03,1.846822E+04]
    locations = rand(2, S) .* [900, 1000]
    world_locations = WCS.pix_to_world(images[3].wcs, locations)
    catalog = CatalogEntry[CatalogEntry(world_locations[:, s], true,
            fluxes, fluxes, 0.1, 0.7, pi/4, 4.0) for s in 1:S];

    Synthetic.gen_images!(images, catalog);

    # Make non-constant background.
    for b=1:5
        images[b].nelec_per_nmgy = fill(images[b].nelec_per_nmgy[1], images[b].H)
        images[b].sky = SDSSBackground(fill(images[b].sky[1,1], images[b].H, images[b].W),
                                       collect(1:images[b].H), collect(1:images[b].W),
                                       ones(images[b].H))
    end

    ea = make_elbo_args(
        images, catalog, patch_radius_pix=patch_pixel_radius)

    vp = Vector{Float64}[DeterministicVI.catalog_init_source(ce) for ce in catalog]
    if perturb
        perturb_params(vp)
    end

    ea, vp, catalog
end

function true_star_init()
    ea, vp, catalog = gen_sample_star_dataset(perturb=false)

    vp[1][ids.is_star] = [ 1.0 - 1e-4, 1e-4 ]
    vp[1][ids.flux_scale] = 1e-4
    vp[1][ids.flux_loc] = log(sample_star_fluxes[3]) - 0.5 * vp[1][ids.flux_scale]
    vp[1][ids.color_var] = 1e-4

    ea, vp, catalog
end

end # End module


"""
The default Updater for a POMCPPlanner.

When a new observation is recieved from the main simulation, a RootUpdater simply updates the
root node for the action decision on the next time step to the child of the current root node
corresponding to the observation. That way, all the information from the tree is preserved.
"""
type RootUpdater{U<:POMDPs.Updater} <: POMDPs.Updater{BeliefNode}
    node_belief_updater::U # updates the belief between nodes if necessary
end

# version with a particle reinvigorator
function update{R<:ParticleReinvigorator,A,O}(updater::RootUpdater{R}, b_old::BeliefNode, a::A, o::O, b=nothing)
    if !haskey(b_old.children[a].children, o)
        new_collection = handle_unseen_observation(updater.node_belief_updater,
                                                   b_old, a, o)
        new_node = ObsNode(o, 0, new_collection, Dict{A,ActNode{A,O,ObsNode{typeof(new_collection),A,O}}}())
        b_old.children[a].children[o] = new_node
    end
    b_new = b_old.children[a].children[o]

    reinvigorate!(b_new.B,
                  updater.node_belief_updater,
                  b_old, a, o)

    return b_new
end

# version with a user-supplied belief updater
function update{A,O}(updater::RootUpdater, b_old::BeliefNode, a::A, o::O, b=nothing)
    if !haskey(b_old.children[a].children, o)
        # if there is no node for the observation, attempt to create one
        new_belief = update(updater.node_belief_updater, b_old.B, a, o)
        new_node = ObsNode(o, 0, new_belief, Dict{A,ActNode{A,O,ObsNode{typeof(new_belief),A,O}}}())
        b_old.children[a].children[o] = new_node
    end

    return b_old.children[a].children[o]
end

updater(policy::POMCPPlanner) = RootUpdater(policy.node_belief_updater)
create_belief(updater::RootUpdater) = RootNode(0, create_belief(updater.node_belief_updater), Dict{Any,ActNode}())
create_belief{R<:ParticleReinvigorator}(updater::RootUpdater{R}) = RootNode(0, nothing, Dict{Any,ActNode}())

initialize_belief(up::RootUpdater, b, new_belief::BeliefNode) = new_belief
initialize_belief(up::RootUpdater, b) = RootNode(0, b, Dict{Any,ActNode}())
initialize_belief(::RootUpdater, n::RootNode) = n

function rand(rng::AbstractRNG, d::BeliefNode, s)
    rand(rng, d.B, s)
end

function rand(rng::AbstractRNG, d::BeliefNode)
    rand(rng, d.B)
end


# License for this file: MIT (expat)
# Copyright 2017-2018, DLR Institute of System Dynamics and Control
#
# This file is part of module
#   Modia3D.Composition (Modia3D/Composition/_module.jl)
#

# Utility function that should not be directly called (only to be called from attach(..)
function attachAndReverseParents(newParent::Object3D, obj::Object3D)::Nothing
   @assert(!(newParent ≡ obj))

   # Save elements of obj
   r_rel  = obj.r_rel
   R_rel  = obj.R_rel
   parent = obj.parent

   # Reverse obj, so that newParent is the new parent
   obj.parent = newParent
   push!(newParent.children, obj)

   # Reverse all parents of obj
   obj1::Object3D = obj
   obj2::Object3D = parent
   while !(parent === obj1)
      parent_r_rel = parent.r_rel
      parent_R_rel = parent.R_rel

      if r_rel ≡ Modia3D.ZeroVector3D
         parent.r_rel = Modia3D.ZeroVector3D
         parent.r_abs = obj1.r_abs
      else
         parent.r_rel = -R_rel*r_rel
      end

      if R_rel ≡ Modia3D.NullRotation
         parent.R_rel = Modia3D.NullRotation
         parent.R_abs = obj1.R_abs
      else
         parent.R_rel = R_rel'
      end

      obj2 = parent
      parent = obj2.parent
      r_rel  = parent_r_rel
      R_rel  = parent_R_rel
      obj2.parent = obj1
      removeChild!(obj2, obj1)
      push!(obj1.children, obj2)
      obj1 = obj2
   end

   # If !isnothing(scene), visit all objs and set scene
#=
   if !isnothing(scene)
      stack = scene.stack
      empty!(stack)
      append!(stack, obj.children)
      while length(stack) > 0
         obj1 = pop!(stack)
         obj1.scene = scene
         append!(stack, obj1.children)
      end
   end
=#
   return nothing
end

"""
    (obj1, obj2, cutJoint) = attach(frame_a, frame_b)

If frame_a and frame_b do not have the same root, they are attached to each other
and cutJoint = false is returned.

If they have the same root, the tree is not modified and cutJoint=true is returned.
"""
function attach(obj1::Object3D, obj2::Object3D)
   root1 = rootObject3D(obj1)
   root2 = rootObject3D(obj2)
   #println("attach: obj1 = ", Modia3D.fullName(obj1), ", root = ", Modia3D.fullName(root1))
   #println("attach: obj2 = ", Modia3D.fullName(obj2), ", root = ", Modia3D.fullName(root2))

   if root1 ≡ root2
      # Compute absolute positions
      # updatePosition!(root1)
      return (obj1,obj2,true)
   end

   if isWorld(root2)
      attachAndReverseParents(obj2, obj1)
      return (obj2,obj1,false)
   else
      attachAndReverseParents(obj1, obj2)
      return (obj1,obj2,false)
   end
end

#=
function connect(obj1::Object3D, obj2::Object3D;
                 r::AbstractVector = Modia3D.ZeroVector3D,
                 R::Union{Frames.RotationMatrix,Nothing} = nothing,
                 q::Union{Frames.Quaternion,Nothing} = nothing,
                 fixed::Bool = true)::Nothing
   if !isnothing(R) && !isnothing(q)
      error("Modia3D.connect(...): either R or q must be nothing but both have a value.")
   end
   if !isnothing(R)
      Modia3D.assertRotationMatrix(R)
   elseif !isnothing(q)
      Modia3D.assertQuaternion(q)
   end

   (parentObject3D, obj, cutJoint) = attach(obj1, obj2)
  # println("... connect, fixed = ", fixed, ", obj1=",Modia3D.fullName(obj1), ", obj2 = ", Modia3D.fullName(obj2), ", obj = ", Modia3D.fullName(obj))

   if cutJoint
      error("Error from Modia3D.Composition.connect(", obj1.name, ",", obj2.name, "):\n",
            "Not yet supported to rigidly connect two objs that have the same root.")
   end

   r_rel = SVector{3,Float64}(r)
   R_rel = !isnothing(R) ? R                   :
   !isnothing(q) ? Modia3D.from_q(q) : Modia3D.NullRotation

   obj.r_rel = obj===obj2 ? r_rel : -R_rel*r_rel
   obj.R_rel = obj===obj2 ? R_rel :  R_rel'

   #r_abs = parent.r_abs + r_rel
   #R_abs = R_rel*parent.R_abs

   if fixed
      obj.joint = fixedJoint
   else
      obj1.hasChildJoint = true
      q_start = !isnothing(R) ? Modia3D.from_R(R) :
      !isnothing(q) ? q                   : Modia3D.NullQuaternion
      q_start = obj===obj2 ? q_start : Modia3D.inverseRotation(q_start)

      obj.joint = FreeMotion(obj, r_start = obj.r_rel, q_start = q_start)
   end
   return nothing
end
=#

function updatePosition!(obj::Object3D)::Nothing
   stack = Object3D[]
   # Push initial children on stack
   append!(stack, obj.children)
   while length(stack) > 0
      obj = pop!(stack)

      # computeKinematics!(obj.joint, obj, Modia3D.KinematicAnalysis, 0.0)

      parent = obj.parent

      if obj.r_rel ≡ Modia3D.ZeroVector3D
         obj.r_abs = parent.r_abs
      else
         obj.r_abs = parent.r_abs + parent.R_abs'*obj.r_rel
      end

      if obj.R_rel ≡ Modia3D.NullRotation
         obj.R_abs = parent.R_abs
      else
         obj.R_abs = obj.R_rel*parent.R_abs
      end

      if length(obj.visualizationFrame) == 1
         obj.visualizationFrame[1].r_abs = obj.r_abs
         obj.visualizationFrame[1].R_abs = obj.R_abs
      end

      append!(stack, obj.children)
   end
   return nothing
end

"""
    R = Rfromrot123(rot123::AbstractVector)

Return rotation matrix `R` from Cardan angles (rotation sequence x-y-z).
`rot123` are the Cardan angles (rotation sequence x-y-z) of rotation from frame `1` to frame `2`.
"""
function Rfromrot123(rot123::AbstractVector)

    (sal, cal) = sincos(rot123[1])
    (sbe, cbe) = sincos(rot123[2])
    (sga, cga) = sincos(rot123[3])
    return @SMatrix[ cbe*cga   sal*sbe*cga+cal*sga  -cal*sbe*cga+sal*sga ;
                    -cbe*sga  -sal*sbe*sga+cal*cga   cal*sbe*sga+sal*cga ;
                     sbe      -sal*cbe               cal*cbe             ]

end

"""
    w = wfromrot123(rot123::AbstractVector, derrot123::AbstractVector)

Return relative rotational velocity Vector3D `w` from frame `1` to frame `2` resolved in frame `2`.

`rot123` are the Cardan angles (rotation sequence x-y-z) of rotation from frame `1` to frame `2`.
`derrot123` are the time derivatives of `rot123`.
"""
function wfromrot123(rot123::AbstractVector, derrot123::AbstractVector)

    # 2_w_12 = (0, 0, gad) + Rga21*[(0, bed, 0) + Rbe21*(ald, 0, 0)]
    #        = (cbe*cga*ald + sga*bed, -cbe*sga*ald + cga*bed, sbe*ald + gad)
    (sbe, cbe) = sincos(rot123[2])
    (sga, cga) = sincos(rot123[3])
    return @SVector[ cbe*cga*derrot123[1] + sga*derrot123[2]                ,
                    -cbe*sga*derrot123[1] + cga*derrot123[2]                ,
                         sbe*derrot123[1]                    + derrot123[3] ]

end

# Next function only for backwards compatibility (do not use for new model)
computeKinematics!(scene::Scene, joint::Modia3D.AbstractJoint, obj::Object3D, analysis::Modia3D.AnalysisType, time::Float64)::Nothing =
    computeKinematics!(scene, [obj], time)

"""
    computeKinematics!(scene::Scene, tree::Vector{Object3D}, time)

Compute position, velocity, acceleration variables of the Object3Ds that are connected
in form of a tree. Variable `tree` contains the Object3Ds in a traversal order (e.g. pre-order traversal).
`tree[1]` is the root object. It is assumed that the kinematic
variables of tree[1].parent have a meaningful value.
"""
function computeKinematics!(scene::Scene, tree::Vector{Object3D}, time)::Nothing
    for obj in tree
        parent    = obj.parent
        jointKind = obj.jointKind

        if jointKind == FixTranslationKind
            obj.r_abs = parent.r_abs + parent.R_abs'*obj.r_rel
            obj.R_abs = parent.R_abs

            obj.v0 = parent.v0 + parent.R_abs'*cross(parent.w, obj.r_rel)
            obj.a0 = parent.a0 + parent.R_abs'*(cross(parent.z, obj.r_rel) + cross(parent.w, cross(parent.w, obj.r_rel)))

            obj.w = parent.w
            obj.z = parent.z

        elseif jointKind == FixKind
            obj.r_abs = parent.r_abs + parent.R_abs'*obj.r_rel
            obj.R_abs = obj.R_rel*parent.R_abs

            obj.v0 = parent.v0 + parent.R_abs'*cross(parent.w, obj.r_rel)
            obj.a0 = parent.a0 + parent.R_abs'*(cross(parent.z, obj.r_rel) + cross(parent.w, cross(parent.w, obj.r_rel)))

            obj.w = obj.R_rel*parent.w
            obj.z = obj.R_rel*parent.z

        elseif jointKind == RevoluteKind
            revolute  = scene.revolute[obj.jointIndex]

            obj.r_abs = parent.r_abs
            obj.R_rel = Frames.rotAxis(revolute.posAxis, revolute.posMovement, revolute.phi)
            obj.R_abs = obj.R_rel*parent.R_abs

            obj.v0 = parent.v0
            obj.a0 = parent.a0

            w_rel = Frames.axisValue(revolute.posAxis, revolute.posMovement, revolute.w)
            z_rel = Frames.axisValue(revolute.posAxis, revolute.posMovement, revolute.a)

            obj.w = obj.R_rel*(parent.w + w_rel)
            obj.z = obj.R_rel*(parent.z + z_rel + cross(parent.w, w_rel))

        elseif jointKind == PrismaticKind
            prismatic = scene.prismatic[obj.jointIndex]

            obj.r_rel = prismatic.eAxis*prismatic.s
            obj.r_abs = parent.r_abs + parent.R_abs'*obj.r_rel
            obj.R_abs = parent.R_abs

            v_rel  = prismatic.eAxis*prismatic.v
            a_rel  = prismatic.eAxis*prismatic.a
            obj.v0 = parent.v0 + parent.R_abs'*(v_rel + cross(parent.w, obj.r_rel))
            obj.a0 = parent.a0 + parent.R_abs'*(a_rel + cross(parent.z, obj.r_rel) + cross(parent.w, v_rel + cross(parent.w, obj.r_rel)) )

            obj.w  = parent.w
            obj.z  = parent.z

        elseif jointKind == AbsoluteFreeMotionKind
            freeMotion  = scene.freeMotion[obj.jointIndex]

            obj.r_rel = freeMotion.r
            obj.R_rel = Rfromrot123(freeMotion.rot)

            obj.r_abs = obj.r_rel
            obj.R_abs = obj.R_rel

            obj.v0 = freeMotion.v
            obj.a0 = freeMotion.a

            obj.w  = freeMotion.w
            obj.z  = freeMotion.z

        elseif jointKind == FreeMotionKind
            freeMotion  = scene.freeMotion[obj.jointIndex]

            obj.r_rel = freeMotion.r
            obj.R_rel = Rfromrot123(freeMotion.rot)

            obj.r_abs = parent.r_abs + parent.R_abs'*obj.r_rel
            obj.R_abs = obj.R_rel*parent.R_abs

            obj.v0 = parent.v0 + parent.R_abs'*(freeMotion.v + cross(parent.w, obj.r_rel))
            obj.a0 = parent.a0 + parent.R_abs'*(freeMotion.a + cross(parent.z, obj.r_rel) +
                                                    cross(parent.w, cross(parent.w, obj.r_rel)))
            obj.w  = parent.R_abs*parent.w + freeMotion.w
            obj.z  = parent.R_abs*parent.z + freeMotion.z

        else
            error("Bug in Modia3D/src/Composition/specifics/specifics.jl (computeKinematics!): jointKind = $jointKind is not known")
        end
    end
    return nothing
end

"""
    computeKinematics_for_leq_mode_pos!(scene::Scene, tree::Vector{Object3D}, time)

Compute accelerations that are only a function of qdd, but not of q and qd.
of the Object3Ds that are connected in form of a tree.
"""
function computeKinematics_for_leq_mode_pos!(scene::Scene, tree::Vector{Object3D}, time)::Nothing
    for obj in tree
        parent    = obj.parent
        jointKind = obj.jointKind

        if jointKind == FixTranslationKind
            obj.a0 = parent.a0 + parent.R_abs'*(cross(parent.z, obj.r_rel))
            obj.z  = parent.z

        elseif jointKind == FixKind
            obj.a0 = parent.a0 + parent.R_abs'*(cross(parent.z, obj.r_rel))
            obj.z  = obj.R_rel*parent.z

        elseif jointKind == RevoluteKind
            revolute = scene.revolute[obj.jointIndex]

            obj.a0 = parent.a0
            z_rel  = Frames.axisValue(revolute.posAxis, revolute.posMovement, revolute.a)
            obj.z  = obj.R_rel*(parent.z + z_rel)

        elseif jointKind == PrismaticKind
            prismatic = scene.prismatic[obj.jointIndex]

            eAxis  = prismatic.posMovement ? prismatic.posAxis : -prismatic.posAxis
            a_rel  = prismatic.eAxis*prismatic.a
            obj.a0 = parent.a0 + parent.R_abs'*(a_rel + cross(parent.z, obj.r_rel))
            obj.z  = parent.z

        elseif jointKind == AbsoluteFreeMotionKind
            freeMotion = scene.freeMotion[obj.jointIndex]

            obj.a0 = freeMotion.a
            obj.z  = freeMotion.z

        elseif jointKind == FreeMotionKind
            freeMotion = scene.freeMotion[obj.jointIndex]

            obj.a0 = parent.a0 + parent.R_abs'*(freeMotion.a + cross(parent.z, obj.r_rel))
            obj.z  = parent.R_abs*parent.z + freeMotion.z

        else
            error("Bug in Modia3D/src/Composition/specifics/specifics.jl (computeKinematics_for_leq_mode_pos!): jointKind = $jointKind is not known")
        end
    end
    return nothing
end

"""
    computeForcesTorquesAndResiduals!(scene::Scene, tree::Vector{Object3D}, time)

Compute forces/torques and residuals in a backward recursion from tree[end] to tree[1].
Variable `tree` contains the Object3Ds in a traversal order (e.g. pre-order traversal).
It is assumed that all force/torque variables are initialized (e.g. to zero), including
tree[1].parent.
"""
function computeForcesTorquesAndResiduals!(scene::Scene, tree::Vector{Object3D}, time)::Nothing
    for i = length(tree):-1:1
        obj       = tree[i]
        parent    = obj.parent
        jointKind = obj.jointKind

        if jointKind == FixTranslationKind
            parent.f += obj.f
            parent.t += obj.t + cross(obj.r_rel, obj.f)

        elseif jointKind == FixKind
            parent.f += obj.R_rel'*obj.f
            parent.t += obj.R_rel'*obj.t + cross(obj.r_rel, obj.f)

        elseif jointKind == RevoluteKind
            revolute  = scene.revolute[obj.jointIndex]
            parent.f += obj.R_rel'*obj.f
            parent.t += obj.R_rel'*obj.t
            obj.f     = -obj.f
            obj.t     = -obj.t
            revolute.residue = -revolute.tau + (revolute.posMovement ? obj.t[revolute.posAxis] : -obj.t[revolute.posAxis])

        elseif jointKind == PrismaticKind
            prismatic = scene.prismatic[obj.jointIndex]
            parent.f += obj.f
            parent.t += obj.t + cross(obj.r_rel, obj.f)
            obj.f     = -obj.f
            obj.t     = -obj.t
            prismatic.residue = -prismatic.f + dot(prismatic.eAxis,obj.f)

        elseif jointKind == AbsoluteFreeMotionKind || jointKind == FreeMotionKind
            freeMotion           = scene.freeMotion[obj.jointIndex]
            freeMotion.residue_f = obj.f
            freeMotion.residue_t = obj.t

        else
            error("Bug in Modia3D/src/Composition/specifics/specifics.jl (computeForcesTorquesAndResiduals!): jointKind = $jointKind is not known")
        end
    end
    return nothing
end

"""
    setJointVariables!(scene::Scene, objects::Vector{Object3D}, args, qdd, startIndex, ndof)

Copy specific variables into the corresponding Object3Ds.
"""
function setJointVariables!(scene::Scene, objects::Vector{Object3D}, args, qdd, startIndex, ndof)::Nothing
    for (i,obj) in enumerate(objects)
        jointKind = obj.jointKind
        args_i    = args[i]
        beg       = startIndex[i]

        if jointKind == RevoluteKind
            @assert(ndof[i] == 1)
            revolute     = scene.revolute[obj.jointIndex]
            revolute.phi = args_i[1]
            revolute.w   = args_i[2]
            revolute.a   = qdd[beg]
            revolute.tau = args_i[3]

        elseif jointKind == PrismaticKind
            @assert(ndof[i] == 1)
            prismatic   = scene.prismatic[obj.jointIndex]
            prismatic.s = args_i[1]
            prismatic.v = args_i[2]
            prismatic.a = qdd[beg]
            prismatic.f = args_i[3]

        elseif jointKind == AbsoluteFreeMotionKind || jointKind == FreeMotionKind
            @assert(ndof[i] == 6)
            freeMotion     = scene.freeMotion[obj.jointIndex]
            freeMotion.r   = SVector{3,Float64}(args_i[1])
            freeMotion.rot = SVector{3,Float64}(args_i[2])
            freeMotion.v   = SVector{3,Float64}(args_i[3])
            freeMotion.w   = SVector{3,Float64}(args_i[4])
            freeMotion.a   = qdd[beg+0:beg+2]
            freeMotion.z   = qdd[beg+3:beg+5]

        else
           error("Bug in Modia3D/src/Composition/specifics/specifics.jl (setJointVariables!): jointKind = $jointKind is not known")
        end
    end
    return nothing
end

"""
    getJointResiduals_for_leq_mode_0!(scene::Scene, objects::Vector{Object3D}, residuals, startIndex::Int, ndof::Int, cache_h)

Copy specific variables into their objects for leq_mode = 0.
"""
function getJointResiduals_for_leq_mode_0!(scene::Scene, objects::Vector{Object3D}, residuals, startIndex::Vector{Int}, ndof::Vector{Int}, cache_h)::Nothing
    for (i,obj) in enumerate(objects)
        jointKind = obj.jointKind
        beg       = startIndex[i]

        if jointKind == RevoluteKind
            @assert(ndof[i] == 1)
            revolute       = scene.revolute[obj.jointIndex]
            cache_h[  beg] = revolute.residue + revolute.tau
            residuals[beg] = revolute.residue

        elseif jointKind == PrismaticKind
            @assert(ndof[i] == 1)
            prismatic      = scene.prismatic[obj.jointIndex]
            cache_h[  beg] = prismatic.residue + prismatic.f
            residuals[beg] = prismatic.residue

        elseif jointKind == AbsoluteFreeMotionKind || jointKind == FreeMotionKind
            @assert(ndof[i] == 6)
            freeMotion             = scene.freeMotion[obj.jointIndex]
            cache_h[  beg+0:beg+2] = freeMotion.residue_f
            cache_h[  beg+3:beg+5] = freeMotion.residue_t

            residuals[beg+0:beg+2] = freeMotion.residue_f
            residuals[beg+3:beg+5] = freeMotion.residue_t

        else
           error("Bug in Modia3D/src/Composition/specifics/specifics.jl (getJointResiduals_for_leq_mode_0!): jointKind = $jointKind is not known")
        end
    end
    return nothing
end

"""
    getJointResiduals_for_leq_mode_pos!(scene::Scene, objects::Vector{Object3D}, residuals, startIndex::Int, ndof::Int, cache_h)

Copy specific variables into their objects for leq_mode > 0.
"""
function getJointResiduals_for_leq_mode_pos!(scene::Scene, objects::Vector{Object3D}, residuals, startIndex::Vector{Int}, ndof::Vector{Int}, cache_h)::Nothing
    for (i,obj) in enumerate(objects)
        jointKind = obj.jointKind
        beg       = startIndex[i]

        if jointKind == RevoluteKind
            @assert(ndof[i] == 1)
            residuals[beg] = scene.revolute[obj.jointIndex].residue + cache_h[beg]

        elseif jointKind == PrismaticKind
            @assert(ndof[i] == 1)
            residuals[beg] = scene.prismatic[obj.jointIndex].residue + cache_h[beg]

        elseif jointKind == AbsoluteFreeMotionKind || jointKind == FreeMotionKind
            @assert(ndof[i] == 6)
            freeMotion             = scene.freeMotion[obj.jointIndex]
            residuals[beg+0:beg+2] = freeMotion.residue_f + cache_h[beg+0:beg+2]
            residuals[beg+3:beg+5] = freeMotion.residue_t + cache_h[beg+3:beg+5]

        else
           error("Bug in Modia3D/src/Composition/specifics/specifics.jl (getJointResiduals_for_leq_mode_0!): jointKind = $jointKind is not known")
        end
    end
    return nothing
end

# For backwards compatibility (do not use for new models)
function setAngle!(revolute::Revolute, phi::Float64)
   obj          = revolute.obj2
   revolute.phi = phi
   obj.R_rel    = Frames.rotAxis(revolute.posAxis, revolute.posMovement, phi)
   obj.R_abs    = obj.R_rel*obj.parent.R_abs
   return revolute
end

# For backwards compatibility (do not use for new models)
function setDistance!(prismatic::Prismatic, s::Float64)
   obj         = prismatic.obj2
   prismatic.s = s
   obj.r_rel   = prismatic.eAxis*s
   obj.r_abs   = parent.r_abs + parent.R_abs'*obj.r_rel
   return prismatic
end


"""
    TrajectoryOptimization
Primary module for setting up and solving trajectory optimization problems.
"""
module TrajectoryOptimization

using ForwardDiff
using DocStringExtensions
using Interpolations
using RecipesBase
using LinearAlgebra
using Statistics
using Random
using SparseArrays
using StaticArrays
using Logging
using Formatting
using BenchmarkTools
using Parameters
using Rotations
using MathOptInterface
using Quaternions
using UnsafeArrays

const MOI = MathOptInterface
const MAX_ELEM = 170

import Base.copy

export
    Dynamics,
    Problems,
    Controllers

export
    state_dim,
    control_dim

# Primary types
export
    Model,
    QuadraticCost,
    LQRCost,
    LQRObjective,
    GenericCost,
    Trajectory

export
    Problem,
    iLQRSolver,
    iLQRSolverOptions,
    AugmentedLagrangianSolver,
    AugmentedLagrangianSolverOptions,
    AugmentedLagrangianProblem,
    ALTROSolverOptions,
    DIRCOLSolver,
    DIRCOLSolverOptions,
    ProjectedNewtonSolver,
    ProjectedNewtonSolverOptions,
    SequentialNewtonSolver,
    Discrete,
    Continuous,
    Constraint,
    BoundConstraint,
    Equality,
    Inequality,
    Objective,
    Constraints

export
    rk3,
    rk4,
    midpoint,
    add_constraints!,
    goal_constraint,
    initial_controls!,
    initial_state!,
    circle_constraint,
    sphere_constraint

# Primary methods
export
    solve,
    solve!,
    rollout!,
    rollout,
    forwardpass!,
    backwardpass!,
    cost,
    max_violation,
    max_violation_direct,
    infeasible_control,
    line_trajectory,
    evaluate!,
    jacobian!

export
    get_sizes,
    num_constraints,
    get_num_constraints,
    get_num_controls,
    init_results,
    to_array,
    get_N,
    to_dvecs,
    quat2rot,
    sphere_constraint,
    circle_constraint,
    plot_trajectory!,
    plot_vertical_lines!,
    convergence_rate,
    plot_obstacles,
    evals,
    reset,
    reset_evals,
    final_time,
    total_time,
    count_constraints,
    inequalities,
    equalities,
    bounds,
    labels,
    terminal,
    stage,
    interp_rows

# Static methods
export
    convertProblem

include("utils.jl")
include("rotations.jl")
include("logger.jl")
include("knotpoint.jl")
include("expansions.jl")
include("model.jl")
include("costfunctions.jl")
include("objective.jl")
include("solver_opts.jl")
include("solvers.jl")
include("abstract_constraint.jl")
include("constraints.jl")
include("dynamics_constraints.jl")
include("integration.jl")
include("dynamics.jl")

include("cost.jl")
include("static_methods.jl")
include("constraint_vals.jl")
include("constraint_sets.jl")
include("problem.jl")
# include("solvers/silqr/silqr_solver.jl")
# include("solvers/silqr/silqr_methods.jl")
include("solvers/ilqr/ilqr.jl")
include("solvers/ilqr/ilqr_solve.jl")
include("solvers/ilqr/backwardpass.jl")
include("solvers/ilqr/rollout.jl")
include("solvers/augmented_lagrangian/sal_solver.jl")
include("solvers/augmented_lagrangian/sal_methods.jl")
include("solvers/direct/primals.jl")
include("solvers/direct/pn.jl")
include("solvers/direct/pn_methods.jl")
include("solvers/altro/altro_solver.jl")

include("solvers/direct/moi.jl")
include("solvers/direct/copy_blocks.jl")
include("solvers/direct/direct_constraints.jl")

include("problems.jl")
include("controllers.jl")

write_ipopt_options()
Logging.global_logger(default_logger(true))
end


"""
Interior point algorithms for conic models.
"""
module Solvers

using DocStringExtensions
using Printf
using LinearAlgebra
using SparseArrays
import SuiteSparse
import LinearMaps
import IterativeSolvers
using Test
import Base.convert
import Hypatia.Cones
import Hypatia.Models
import Hypatia.SparseNonSymCache
import Hypatia.SparseSymCache
import Hypatia.diag_min
import Hypatia.update_fact
import Hypatia.inv_prod
import Hypatia.free_memory
import Hypatia.int_type
import Hypatia.nonsymm_fact_copy!
import Hypatia.symm_fact_copy!
import Hypatia.posdef_fact_copy!
import Hypatia.outer_prod!

const RealOrNothing = Union{Real, Nothing}

include("point.jl")

# solver termination status codes
@enum Status begin
    NotLoaded
    Loaded
    SolveCalled
    Optimal
    PrimalInfeasible
    DualInfeasible
    IllPosed
    PrimalInconsistent
    DualInconsistent
    SlowProgress
    IterationLimit
    TimeLimit
    NumericalFailure
    UnknownStatus
end

convert(::Type{String}, status::Status) = string(status)

abstract type Stepper{T <: Real} end

abstract type SystemSolver{T <: Real} end

"""
$(TYPEDEF)

Hypatia's interior point solver type. See source for options and defaults.
"""
mutable struct Solver{T <: Real}
    # main options
    verbose::Bool
    iter_limit::Int
    time_limit::Float64
    tol_rel_opt::T
    tol_abs_opt::T
    tol_feas::T
    tol_infeas::T
    tol_illposed::T
    tol_slow::T
    preprocess::Bool
    reduce::Bool
    rescale::Bool
    init_use_indirect::Bool
    init_tol_qr::T
    stepper::Stepper{T}
    syssolver::SystemSolver{T}

    # current status of the solver object and info
    status::Status
    solve_time::Float64
    num_iters::Int

    # helpful performance metrics for major subprocedures
    time_rescale::Float64 # rescale affine data
    time_initx::Float64 # preprocess dual equalities and finding initial x point
    time_inity::Float64 # preprocess primal equalities and finding initial y point
    time_unproc::Float64 # unprocess final point
    time_loadsys::Float64 # initialize/load system solver
    time_upsys::Float64 # update LHS and factorization etc for directions solving
    time_upfact::Float64 # update inner factorization for directions solving only
    time_uprhs::Float64 # update RHSs for directions
    time_getdir::Float64 # solve for directions
    time_search::Float64 # searches for alpha parameters

    # model and preprocessed model data
    orig_model::Models.Model{T}
    model::Models.Model{T}
    x_keep_idxs::AbstractVector{Int}
    y_keep_idxs::AbstractVector{Int}
    Ap_R::UpperTriangular{T, <:AbstractMatrix{T}}
    Ap_Q::Union{UniformScaling, AbstractMatrix{T}}
    reduce_cQ1
    reduce_Rpib0
    reduce_GQ1
    reduce_Ap_R
    reduce_Ap_Q
    reduce_y_keep_idxs
    reduce_row_piv_inv

    # current iterate
    point::Point{T}
    mu::T

    # result (solution) point
    result::Point{T}

    # residuals
    x_residual::Vector{T}
    y_residual::Vector{T}
    z_residual::Vector{T}
    tau_residual::T
    x_norm_res_t::T
    y_norm_res_t::T
    z_norm_res_t::T
    x_norm_res::T
    y_norm_res::T
    z_norm_res::T

    # direction solving helpers
    res_norm_cutoff::T # desired worst-case residual norm (inf) on direction
    max_ref_steps::Int # maximum number of iterative refinement steps

    # convergence parameters
    primal_obj_t::T
    dual_obj_t::T
    primal_obj::T
    dual_obj::T
    gap::T
    x_feas::T
    y_feas::T
    z_feas::T
    tau_feas::T

    # termination condition helpers
    x_conv_tol::T
    y_conv_tol::T
    z_conv_tol::T
    # TODO count how many slow with counter field and use option for max slow count
    prev_is_slow::Bool
    prev2_is_slow::Bool
    worst_dir_res::T

    # data scaling
    used_rescaling::Bool
    b_scale::Vector{T}
    c_scale::Vector{T}
    h_scale::Vector{T}

    function Solver{T}(;
        verbose::Bool = true,
        iter_limit::Int = 1000,
        time_limit::Real = Inf,
        tol_rel_opt::RealOrNothing = nothing,
        tol_abs_opt::RealOrNothing = nothing,
        tol_feas::RealOrNothing = nothing,
        tol_infeas::RealOrNothing = nothing,
        tol_illposed::RealOrNothing = nothing,
        default_tol_power::RealOrNothing = nothing,
        default_tol_relax::RealOrNothing = nothing,
        tol_slow::Real = 1e-3,
        preprocess::Bool = true,
        reduce::Bool = true,
        rescale::Bool = true,
        init_use_indirect::Bool = false,
        init_tol_qr::Real = 1000 * eps(T),
        stepper::Stepper{T} = default_stepper(T),
        syssolver::SystemSolver{T} = default_syssolver(T),
        ) where {T <: Real}
        if isa(syssolver, QRCholSystemSolver{T})
            @assert preprocess # require preprocessing for QRCholSystemSolver # TODO only need primal eq preprocessing or reduction
        end
        if reduce
            @assert preprocess # cannot use reduction without preprocessing # TODO only need primal eq preprocessing
        end
        # @assert !(init_use_indirect && preprocess) # cannot use preprocessing and indirect methods for initial point

        if isnothing(default_tol_power)
            default_tol_power = (T <: LinearAlgebra.BlasReal ? 0.5 : 0.4)
        end
        default_tol_power = T(default_tol_power)
        default_tol_loose = eps(T) ^ default_tol_power
        default_tol_tight = eps(T) ^ (T(1.5) * default_tol_power)
        if !isnothing(default_tol_relax)
            default_tol_loose *= T(default_tol_relax)
            default_tol_tight *= T(default_tol_relax)
        end
        if isnothing(tol_rel_opt)
            tol_rel_opt = default_tol_loose
        end
        if isnothing(tol_abs_opt)
            tol_abs_opt = default_tol_tight
        end
        if isnothing(tol_feas)
            tol_feas = default_tol_loose
        end
        if isnothing(tol_infeas)
            tol_infeas = default_tol_tight
        end
        if isnothing(tol_illposed)
            tol_illposed = default_tol_tight / 100
        end
        @assert min(tol_rel_opt, tol_abs_opt, tol_feas,
            tol_infeas, tol_illposed, tol_slow) >= 0

        solver = new{T}()

        solver.verbose = verbose
        solver.iter_limit = iter_limit
        solver.time_limit = time_limit
        solver.tol_rel_opt = tol_rel_opt
        solver.tol_abs_opt = tol_abs_opt
        solver.tol_feas = tol_feas
        solver.tol_infeas = tol_infeas
        solver.tol_illposed = tol_illposed
        solver.tol_slow = tol_slow
        solver.preprocess = preprocess
        solver.reduce = reduce
        solver.rescale = rescale
        solver.init_use_indirect = init_use_indirect
        solver.init_tol_qr = init_tol_qr
        solver.stepper = stepper
        solver.syssolver = syssolver
        solver.status = NotLoaded

        return solver
    end
end

default_stepper(T) = CombinedStepper{T}()
default_syssolver(T) = QRCholDenseSystemSolver{T}()

function solve(solver::Solver{T}) where {T <: Real}
    @assert solver.status == Loaded
    solver.status = SolveCalled
    start_time = time()
    solver.num_iters = 0
    solver.solve_time = NaN

    solver.time_rescale = 0
    solver.time_initx = 0
    solver.time_inity = 0
    solver.time_unproc = 0
    solver.time_loadsys = 0
    solver.time_upsys = 0
    solver.time_upfact = 0
    solver.time_uprhs = 0
    solver.time_getdir = 0
    solver.time_search = 0

    solver.res_norm_cutoff = 0
    solver.max_ref_steps = 5

    solver.x_norm_res_t = NaN
    solver.y_norm_res_t = NaN
    solver.z_norm_res_t = NaN
    solver.x_norm_res = NaN
    solver.y_norm_res = NaN
    solver.z_norm_res = NaN

    solver.primal_obj_t = NaN
    solver.dual_obj_t = NaN
    solver.primal_obj = NaN
    solver.dual_obj = NaN
    solver.gap = NaN
    solver.x_feas = NaN
    solver.y_feas = NaN
    solver.z_feas = NaN
    solver.tau_feas = NaN

    # preprocess and find initial point
    orig_model = solver.orig_model
    result = solver.result = Point(orig_model)
    # copy original model to solver.model, which may be modified
    model = solver.model = Models.Model{T}(
        orig_model.c, orig_model.A, orig_model.b, orig_model.G, orig_model.h,
        orig_model.cones, obj_offset = orig_model.obj_offset)
    (init_z, init_s) = initialize_cone_point(solver.orig_model)

    solver.time_rescale = @elapsed solver.used_rescaling = rescale_data(solver)

    if solver.reduce
        # TODO don't find point / unnecessary stuff before reduce
        solver.time_inity = @elapsed init_y = find_initial_y(solver, init_z, true)
        solver.time_initx = @elapsed init_x = find_initial_x(solver, init_s)
    else
        solver.time_initx = @elapsed init_x = find_initial_x(solver, init_s)
        solver.time_inity = @elapsed init_y = find_initial_y(solver, init_z, false)
    end

    if solver.status == SolveCalled
        point = solver.point = Point(model)
        point.x .= init_x
        point.y .= init_y
        point.z .= init_z
        point.s .= init_s
        point.tau[] = one(T)
        point.kap[] = one(T)
        calc_mu(solver)
        if isnan(solver.mu) || abs(one(T) - solver.mu) > sqrt(eps(T))
            @warn("initial mu is $(solver.mu) but should be 1 (this could " *
                "indicate a problem with cone barrier oracles)")
        end
        Cones.load_point.(model.cones, point.primal_views)
        Cones.load_dual_point.(model.cones, point.dual_views)

        # setup iteration helpers
        solver.x_residual = zero(model.c)
        solver.y_residual = zero(model.b)
        solver.z_residual = zero(model.h)
        solver.tau_residual = 0

        solver.x_conv_tol = inv(1 + norm(model.c, Inf))
        solver.y_conv_tol = inv(1 + norm(model.b, Inf))
        solver.z_conv_tol = inv(1 + norm(model.h, Inf))
        solver.prev_is_slow = false
        solver.prev2_is_slow = false
        solver.worst_dir_res = 0

        stepper = solver.stepper
        load(stepper, solver)
        solver.time_loadsys = @elapsed load(solver.syssolver, solver)

        solver.verbose && print_header(stepper, solver)
        flush(stdout)

        # iterate from initial point
        while true
            improv = calc_convergence_params(solver)

            if solver.verbose
                print_iteration(stepper, solver)
                flush(stdout)
            end

            check_convergence(solver) && break

            if solver.num_iters == solver.iter_limit
                solver.verbose && println("iteration limit reached; terminating")
                solver.status = IterationLimit
                break
            end

            if time() - start_time >= solver.time_limit
                solver.verbose && println("time limit reached; terminating")
                solver.status = TimeLimit
                break
            end

            if expect_improvement(stepper)
                if improv < solver.tol_slow
                    if solver.prev_is_slow && solver.prev2_is_slow
                        if solver.verbose
                            println("slow progress in consecutive " *
                                "iterations; terminating")
                        end
                        solver.status = SlowProgress
                        break
                    else
                        solver.prev2_is_slow = solver.prev_is_slow
                        solver.prev_is_slow = true
                    end
                else
                    solver.prev2_is_slow = solver.prev_is_slow
                    solver.prev_is_slow = false
                end
            end

            solver.res_norm_cutoff = T(1e-4) * max(solver.x_norm_res,
                solver.y_norm_res, solver.z_norm_res, solver.tau_feas)
            solver.worst_dir_res = 0

            step(stepper, solver) || break
            flush(stdout)
            calc_mu(solver)

            if min(point.tau[], point.kap[], solver.mu) <= 0
                @warn("numerical failure: tau is $(point.tau[]), " *
                    "kappa is $(point.kap[]), mu is $(solver.mu); terminating")
                solver.status = NumericalFailure
                break
            end

            flush(stdout)
            solver.num_iters += 1
        end

        # finalize result point
        solver.time_unproc = @elapsed postprocess(solver)
    end

    solver.solve_time = time() - start_time

    # free memory used by some system solvers
    free_memory(solver.syssolver)

    if solver.verbose
        println("\nstatus is $(solver.status) after $(solver.num_iters) " *
            "iterations and $(trunc(solver.solve_time, digits=3)) seconds\n")
    end
    flush(stdout)

    return solver
end

function calc_mu(solver::Solver{T}) where {T <: Real}
    point = solver.point
    solver.mu = (dot(point.z, point.s) + point.tau[] * point.kap[]) /
        (solver.model.nu + 1)
    return solver.mu
end

function calc_convergence_params(solver::Solver{T}) where {T <: Real}
    model = solver.model
    point = solver.point
    tau = point.tau[]

    # x_residual = -A'*y - G'*z - c*tau
    x_residual = solver.x_residual
    mul!(x_residual, model.G', point.z)
    mul!(x_residual, model.A', point.y, true, true)
    solver.x_norm_res_t = norm(x_residual, Inf)
    @. x_residual += model.c * tau
    solver.x_norm_res = norm(x_residual, Inf) / tau
    @. x_residual *= -1
    x_feas = solver.x_norm_res * solver.x_conv_tol

    # y_residual = A*x - b*tau
    y_residual = solver.y_residual
    mul!(y_residual, model.A, point.x)
    solver.y_norm_res_t = norm(y_residual, Inf)
    @. y_residual -= model.b * tau
    solver.y_norm_res = norm(y_residual, Inf) / tau
    y_feas = solver.y_norm_res * solver.y_conv_tol

    # z_residual = s + G*x - h*tau
    z_residual = solver.z_residual
    mul!(z_residual, model.G, point.x)
    @. z_residual += point.s
    solver.z_norm_res_t = norm(z_residual, Inf)
    @. z_residual -= model.h * tau
    solver.z_norm_res = norm(z_residual, Inf) / tau
    z_feas = solver.z_norm_res * solver.z_conv_tol

    # tau_residual = c'*x + b'*y + h'*z + kap
    solver.primal_obj_t = dot(model.c, point.x)
    solver.dual_obj_t = -dot(model.b, point.y) - dot(model.h, point.z)
    solver.tau_residual = solver.primal_obj_t - solver.dual_obj_t + point.kap[]
    tau_feas = abs(solver.tau_residual)

    # check improvement
    improv = zero(T)
    for (curr, prev) in ((x_feas, solver.x_feas), (y_feas, solver.y_feas),
        (z_feas, solver.z_feas), (tau_feas, solver.tau_feas))
        if isnan(prev) || isnan(curr)
            continue
        end
        improv = max(improv, (prev - curr) / (abs(prev) + eps(T)))
    end
    solver.x_feas = x_feas
    solver.y_feas = y_feas
    solver.z_feas = z_feas
    solver.tau_feas = tau_feas

    # gap
    solver.primal_obj = solver.primal_obj_t / tau + model.obj_offset
    solver.dual_obj = solver.dual_obj_t / tau + model.obj_offset
    solver.gap = dot(point.z, point.s)

    return improv
end

function check_convergence(solver::Solver{T}) where {T <: Real}
    # check convergence criteria
    # TODO nearly primal or dual infeasible or nearly optimal cases?
    tau = solver.point.tau[]
    primal_obj_t = solver.primal_obj_t
    dual_obj_t = solver.dual_obj_t

    is_feas = (max(solver.x_feas, solver.y_feas, solver.z_feas) <= solver.tol_feas)
    is_abs_opt = (solver.gap <= solver.tol_abs_opt)
    is_rel_opt = (min(solver.gap / tau, abs(primal_obj_t - dual_obj_t)) <=
        solver.tol_rel_opt * max(tau, min(abs(primal_obj_t), abs(dual_obj_t))))
    if is_feas && (is_abs_opt || is_rel_opt)
        solver.verbose && println("optimal solution found; terminating")
        solver.status = Optimal
        return true
    end

    if dual_obj_t > eps(T) && solver.x_norm_res_t <= solver.tol_infeas * dual_obj_t
        solver.verbose && println("primal infeasibility detected; terminating")
        solver.status = PrimalInfeasible
        solver.primal_obj = primal_obj_t
        solver.dual_obj = dual_obj_t
        return true
    end

    if primal_obj_t < -eps(T) && max(solver.y_norm_res_t, solver.z_norm_res_t) <=
        solver.tol_infeas * -primal_obj_t
        solver.verbose && println("dual infeasibility detected; terminating")
        solver.status = DualInfeasible
        solver.primal_obj = primal_obj_t
        solver.dual_obj = dual_obj_t
        return true
    end

    # TODO experiment with ill-posedness check
    if solver.mu <= solver.tol_illposed && tau <=
        solver.tol_illposed * min(one(T), solver.point.kap[])
        solver.verbose && println("ill-posedness detected; terminating")
        solver.status = IllPosed
        return true
    end

    return false
end

function initialize_cone_point(model::Models.Model{T}) where {T <: Real}
    init_z = zeros(T, model.q)
    init_s = zeros(T, model.q)

    for (cone, idxs) in zip(model.cones, model.cone_idxs)
        Cones.setup_data!(cone)
        primal_k = view(Cones.use_dual_barrier(cone) ? init_z : init_s, idxs)
        dual_k = view(Cones.use_dual_barrier(cone) ? init_s : init_z, idxs)
        Cones.set_initial_point!(primal_k, cone)
        Cones.load_point(cone, primal_k)
        @assert Cones.is_feas(cone)
        g = Cones.grad(cone)
        @. dual_k = -g
        Cones.load_dual_point(cone, dual_k)
        @assert Cones.is_dual_feas(cone)
    end

    return (init_z, init_s)
end

get_status(solver::Solver) = solver.status
get_solve_time(solver::Solver) = solver.solve_time
get_num_iters(solver::Solver) = solver.num_iters

get_primal_obj(solver::Solver) = solver.primal_obj
get_dual_obj(solver::Solver) = solver.dual_obj

get_s(solver::Solver) = copy(solver.result.s)
get_z(solver::Solver) = copy(solver.result.z)
get_x(solver::Solver) = copy(solver.result.x)
get_y(solver::Solver) = copy(solver.result.y)

get_tau(solver::Solver) = solver.point.tau[]
get_kappa(solver::Solver) = solver.point.kap[]
get_mu(solver::Solver) = solver.mu

function load(solver::Solver{T}, model::Models.Model{T}) where {T <: Real}
    # @assert solver.status == NotLoaded # TODO maybe want a reset function that just keeps options
    solver.orig_model = model
    solver.status = Loaded
    return solver
end

include("process.jl")

include("search.jl")

include("steppers/common.jl")

include("systemsolvers/common.jl")

# release memory used by sparse system solvers
free_memory(::SystemSolver) = nothing
free_memory(syssolver::Union{NaiveSparseSystemSolver, SymIndefSparseSystemSolver}) =
    free_memory(syssolver.fact_cache)

# verbose helpers
function print_header(stepper::Stepper, solver::Solver)
    println()
    @printf("%5s %12s %12s |%9s ", "iter", "p_obj", "d_obj", "abs_gap")
    if iszero(solver.model.p)
        @printf("%9s %9s ", "x_feas", "z_feas")
    else
        @printf("%9s %9s %9s ", "x_feas", "y_feas", "z_feas")
    end
    @printf("|%9s %9s %9s |%8s %8s ", "tau", "kap", "mu", "dir_res", "prox")

    print_header_more(stepper, solver)
    println()
    return
end
print_header_more(stepper::Stepper, solver::Solver) = nothing

function print_iteration(stepper::Stepper, solver::Solver)
    @printf("%5d %12.4e %12.4e |%9.2e ", solver.num_iters, solver.primal_obj,
        solver.dual_obj, solver.gap)
    if iszero(solver.model.p)
        @printf("%9.2e %9.2e ", solver.x_feas, solver.z_feas)
    else
        @printf("%9.2e %9.2e %9.2e ", solver.x_feas, solver.y_feas, solver.z_feas)
    end
    @printf("|%9.2e %9.2e %9.2e |", solver.point.tau[], solver.point.kap[],
        solver.mu)

    if !iszero(solver.num_iters)
        @printf("%8.1e %8.1e ", solver.worst_dir_res, stepper.searcher.prox)
        print_iteration_more(stepper, solver)
    end
    println()
    return
end
print_iteration_more(stepper::Stepper, solver::Solver) = nothing

end


#=
Copyright (c) 2018-2022 Chris Coey, Lea Kapelevich, and contributors

This Julia package Hypatia.jl is released under the MIT license; see LICENSE
file in the root directory or at https://github.com/chriscoey/Hypatia.jl
=#

insts = OrderedDict()
insts["minimal"] = [
    ((3, 2, true, 0),),
    ((3, 2, false, 0),),
    ((3, 2, true, 10),),
    ((3, 2, false, 10),),
    ]
insts["fast"] = [
    ((5, 3, true, 0),),
    ((5, 3, false, 0),),
    ((5, 3, true, 10),),
    ((5, 3, false, 10),),
    ((30, 10, true, 0),),
    ((30, 10, false, 0), (default_tol_relax = 1000,)),
    ((30, 10, true, 10),),
    ((30, 10, false, 10),),
    ]
insts["various"] = [
    ((10, 3, true, 10),),
    ((10, 3, false, 10),),
    ((20, 6, true, 10),),
    ((20, 6, false, 10),),
    ((40, 12, true, 10),),
    ((40, 12, false, 10),),
    ]
return (SparsePCANative, insts)


using YaoLocations
using MLStyle
using YaoHIR
using Test

module TestIntrinsic
using YaoHIR: @intrinsic

@intrinsic X
@intrinsic R(theta::T) where {T <: Real}
@intrinsic SWAP

end

@test YaoHIR.routine_name(TestIntrinsic.X) == :X
@test YaoHIR.routine_name(TestIntrinsic.R(1.0)) == :R
@test TestIntrinsic.R(1.0).theta == 1.0

display(YaoHIR.X)
display(YaoHIR.Rx(1.0))
display(YaoHIR.Operation(YaoHIR.X, 2.0))

circ = Chain(
    Gate(YaoHIR.X, Locations(1)),
    Core.SSAValue(1),
    Ctrl(Gate(Core.SSAValue(1), Locations(3)), CtrlLocations(2))
)

print(circ)

@test YaoHIR.leaves(circ) == [Gate(YaoHIR.X, Locations(1)),
    Core.SSAValue(1),
    Ctrl(Gate(Core.SSAValue(1), Locations(3)), CtrlLocations(2))
]


# ------------ Choice sites ------------ #

@inline function (ctx::ParameterBackpropagateContext)(call::typeof(rand),
                                                      addr::T,
                                                      d::Distribution{K}) where {T <: Address, K}
    s = get_top(ctx.tr, addr).val
    ctx.weight += logpdf(d, s)
    return s
end

@inline function (ctx::ChoiceBackpropagateContext)(call::typeof(rand),
                                                   addr::T,
                                                   d::Distribution{K}) where {T <: Address, K}
    s = get_top(ctx.tr, addr).val
    ctx.weight += logpdf(d, s)
    return s
end

# ------------ Learnable ------------ #

@inline function (ctx::ParameterBackpropagateContext)(fn::typeof(learnable), addr::Address)
    return read_parameter(ctx, addr)
end

@inline function (ctx::ChoiceBackpropagateContext)(fn::typeof(learnable), addr::Address)
    return read_parameter(ctx, addr)
end

# ------------ Call sites ------------ #

@inline function (ctx::ParameterBackpropagateContext)(c::typeof(rand),
                                                      addr::T,
                                                      call::Function,
                                                      args...) where T <: Address
    cl = get_sub(ctx.tr, addr)
    param_grads = Gradients()
    ps = get_sub(ctx.initial_params, addr)
    ret = simulate_call_pullback(ps, param_grads, cl, args)
    ctx.param_grads.tree[addr] = param_grads
    return ret
end

@inline function (ctx::ChoiceBackpropagateContext)(c::typeof(rand),
                                                   addr::T,
                                                   call::Function,
                                                   args...) where T <: Address
    cl = get_sub(ctx.tr, addr)
    choice_grads = Gradients()
    ps = get_sub(ctx.initial_params, addr)
    ret = simulate_choice_pullback(ps, choice_grads, get_sub(ctx.select, addr), cl, args)
    ctx.choice_grads.tree[addr] = choice_grads
    return ret
end


##################################################################################
#### DEFINTION OF NEW VARIABLES FOR STORAGE INVESTMENTS ACCODING TO FlexPlan MODEL
##################################################################################

function variable_absorbed_energy(pm::_PM.AbstractPowerModel; nw::Int=_PM.nw_id_default, bounded::Bool = true, report::Bool=true)
    e_abs = _PM.var(pm, nw)[:e_abs] = JuMP.@variable(pm.model,
    [i in _PM.ids(pm, nw, :storage_bounded_absorption)], base_name="$(nw)_e_abs",
    start = 0)

    if bounded
        for (s, storage) in _PM.ref(pm, nw, :storage_bounded_absorption)
            JuMP.set_lower_bound(e_abs[s],  0)
            JuMP.set_upper_bound(e_abs[s],  storage["max_energy_absorption"])
        end
    end

    report && _PM.sol_component_value(pm, nw, :storage, :e_abs, _PM.ids(pm, nw, :storage_bounded_absorption), e_abs)
end

function variable_absorbed_energy_ne(pm::_PM.AbstractPowerModel; nw::Int=_PM.nw_id_default, bounded::Bool = true, report::Bool=true)
    e_abs = _PM.var(pm, nw)[:e_abs_ne] = JuMP.@variable(pm.model,
    [i in _PM.ids(pm, nw, :ne_storage_bounded_absorption)], base_name="$(nw)_e_abs_ne",
    start = 0)

    if bounded
        for (s, storage) in _PM.ref(pm, nw, :ne_storage_bounded_absorption)
            JuMP.set_lower_bound(e_abs[s],  0)
            JuMP.set_upper_bound(e_abs[s],  storage["max_energy_absorption"])
        end
    end

    report && _PM.sol_component_value(pm, nw, :ne_storage, :e_abs_ne, _PM.ids(pm, nw, :ne_storage_bounded_absorption), e_abs)
end

function variable_storage_power_ne(pm::_PM.AbstractPowerModel; kwargs...)
    variable_storage_power_real_ne(pm; kwargs...)
    variable_storage_power_imaginary_ne(pm; kwargs...)
    variable_storage_power_control_imaginary_ne(pm; kwargs...)
    variable_storage_current_ne(pm; kwargs...)
    variable_storage_energy_ne(pm; kwargs...)
    variable_storage_charge_ne(pm; kwargs...)
    variable_storage_discharge_ne(pm; kwargs...)
    variable_storage_indicator(pm; kwargs..., relax=true)
    variable_storage_investment(pm; kwargs...)
end

function variable_storage_power_real_ne(pm::_PM.AbstractPowerModel; nw::Int=_PM.nw_id_default, bounded::Bool=true, report::Bool=true)
    ps = _PM.var(pm, nw)[:ps_ne] = JuMP.@variable(pm.model,
        [i in _PM.ids(pm, nw, :ne_storage)], base_name="$(nw)_ps_ne",
        start = _PM.comp_start_value(_PM.ref(pm, nw, :ne_storage, i), "ps_start")
    )

    if bounded
        inj_lb, inj_ub = _PM.ref_calc_storage_injection_bounds(_PM.ref(pm, nw, :ne_storage), _PM.ref(pm, nw, :bus))

        for i in _PM.ids(pm, nw, :ne_storage)
            if !isinf(inj_lb[i])
                JuMP.set_lower_bound(ps[i], inj_lb[i])
            end
            if !isinf(inj_ub[i])
                JuMP.set_upper_bound(ps[i], inj_ub[i])
            end
        end
    end

    report && _PM.sol_component_value(pm, nw, :ne_storage, :ps_ne, _PM.ids(pm, nw, :ne_storage), ps)
end

function variable_storage_power_imaginary_ne(pm::_PM.AbstractPowerModel; nw::Int=_PM.nw_id_default, bounded::Bool=true, report::Bool=true)
    qs = _PM.var(pm, nw)[:qs_ne] = JuMP.@variable(pm.model,
        [i in _PM.ids(pm, nw, :ne_storage)], base_name="$(nw)_qs_ne",
        start = _PM.comp_start_value(_PM.ref(pm, nw, :ne_storage, i), "qs_start")
    )

    if bounded
        inj_lb, inj_ub = _PM.ref_calc_storage_injection_bounds(_PM.ref(pm, nw, :ne_storage), _PM.ref(pm, nw, :bus))

        for (i, storage) in _PM.ref(pm, nw, :ne_storage)
            JuMP.set_lower_bound(qs[i], max(inj_lb[i], storage["qmin"]))
            JuMP.set_upper_bound(qs[i], min(inj_ub[i], storage["qmax"]))
        end
    end

    report && _PM.sol_component_value(pm, nw, :ne_storage, :qs_ne, _PM.ids(pm, nw, :ne_storage), qs)
end

"apo models ignore reactive power flows"
function variable_storage_power_imaginary_ne(pm::_PM.AbstractActivePowerModel; nw::Int=_PM.nw_id_default, bounded::Bool=true, report::Bool=true)
    report && _IM.sol_component_fixed(pm, _PM.pm_it_sym, nw, :ne_storage, :qs_ne, _PM.ids(pm, nw, :ne_storage), NaN)
end

function variable_storage_power_control_imaginary_ne(pm::_PM.AbstractPowerModel; nw::Int=_PM.nw_id_default, bounded::Bool=true, report::Bool=true)
    qsc = _PM.var(pm, nw)[:qsc_ne] = JuMP.@variable(pm.model,
        [i in _PM.ids(pm, nw, :ne_storage)], base_name="$(nw)_qsc_ne",
        start = _PM.comp_start_value(_PM.ref(pm, nw, :ne_storage, i), "qsc_start")
    )

    if bounded
        inj_lb, inj_ub = _PM.ref_calc_storage_injection_bounds(_PM.ref(pm, nw, :ne_storage), _PM.ref(pm, nw, :bus))

        for (i,storage) in _PM.ref(pm, nw, :ne_storage)

            if !isinf(inj_lb[i]) || haskey(storage, "qmin")
                JuMP.set_lower_bound(qsc[i], max(inj_lb[i], get(storage, "qmin", -Inf)))
            end
            if !isinf(inj_ub[i]) || haskey(storage, "qmax")
                JuMP.set_upper_bound(qsc[i], min(inj_ub[i], get(storage, "qmax",  Inf)))
            end
        end
    end

    report && _PM.sol_component_value(pm, nw, :ne_storage, :qsc_ne, _PM.ids(pm, nw, :ne_storage), qsc)
end

"apo models ignore reactive power flows"
function variable_storage_power_control_imaginary_ne(pm::_PM.AbstractActivePowerModel; nw::Int=_PM.nw_id_default, bounded::Bool=true, report::Bool=true)
    report && _IM.sol_component_fixed(pm, _PM.pm_it_sym, nw, :ne_storage, :qsc_ne, _PM.ids(pm, nw, :ne_storage), NaN)
end

"do nothing by default but some formulations require this"
function variable_storage_current_ne(pm::_PM.AbstractPowerModel; nw::Int=_PM.nw_id_default, bounded::Bool=true, report::Bool=true)
end

function variable_storage_energy_ne(pm::_PM.AbstractPowerModel; nw::Int=_PM.nw_id_default, bounded::Bool=true, report::Bool=true)
    se = _PM.var(pm, nw)[:se_ne] = JuMP.@variable(pm.model,
        [i in _PM.ids(pm, nw, :ne_storage)], base_name="$(nw)_se_ne",
        start = _PM.comp_start_value(_PM.ref(pm, nw, :ne_storage, i), "se_start", 1)
    )

    if bounded
        for (i, storage) in _PM.ref(pm, nw, :ne_storage)
            JuMP.set_lower_bound(se[i], 0)
            JuMP.set_upper_bound(se[i], storage["energy_rating"])
        end
    end

    report && _PM.sol_component_value(pm, nw, :ne_storage, :se_ne, _PM.ids(pm, nw, :ne_storage), se)
end

function variable_storage_charge_ne(pm::_PM.AbstractPowerModel; nw::Int=_PM.nw_id_default, bounded::Bool=true, report::Bool=true)
    sc = _PM.var(pm, nw)[:sc_ne] = JuMP.@variable(pm.model,
        [i in _PM.ids(pm, nw, :ne_storage)], base_name="$(nw)_sc_ne",
        start = _PM.comp_start_value(_PM.ref(pm, nw, :ne_storage, i), "sc_start", 1)
    )

    if bounded
        for (i, storage) in _PM.ref(pm, nw, :ne_storage)
            JuMP.set_lower_bound(sc[i], 0)
            JuMP.set_upper_bound(sc[i], storage["charge_rating"])
        end
    end

    report && _PM.sol_component_value(pm, nw, :ne_storage, :sc_ne, _PM.ids(pm, nw, :ne_storage), sc)
end

function variable_storage_discharge_ne(pm::_PM.AbstractPowerModel; nw::Int=_PM.nw_id_default, bounded::Bool=true, report::Bool=true)
    sd = _PM.var(pm, nw)[:sd_ne] = JuMP.@variable(pm.model,
        [i in _PM.ids(pm, nw, :ne_storage)], base_name="$(nw)_sd_ne",
        start = _PM.comp_start_value(_PM.ref(pm, nw, :ne_storage, i), "sd_start", 1)
    )

    if bounded
        for (i, storage) in _PM.ref(pm, nw, :ne_storage)
            JuMP.set_lower_bound(sd[i], 0)
            JuMP.set_upper_bound(sd[i], storage["discharge_rating"])
        end
    end

    report && _PM.sol_component_value(pm, nw, :ne_storage, :sd_ne, _PM.ids(pm, nw, :ne_storage), sd)
end

function variable_storage_indicator(pm::_PM.AbstractPowerModel; nw::Int=_PM.nw_id_default, relax::Bool=false, report::Bool=true)
    first_n = first_id(pm, nw, :hour, :scenario)
    if nw == first_n
        if !relax
            z = _PM.var(pm, nw)[:z_strg_ne] = JuMP.@variable(pm.model,
                [i in _PM.ids(pm, nw, :ne_storage)], base_name="$(nw)_z_strg_ne",
                binary = true,
                start = 0
            )
        else
            z = _PM.var(pm, nw)[:z_strg_ne] = JuMP.@variable(pm.model,
                [i in _PM.ids(pm, nw, :ne_storage)], base_name="$(nw)_z_strg_ne",
                lower_bound = 0,
                upper_bound = 1,
                start = 0
            )
        end
    else
        z = _PM.var(pm, nw)[:z_strg_ne] = _PM.var(pm, first_n)[:z_strg_ne]
    end
    report && _PM.sol_component_value(pm, nw, :ne_storage, :isbuilt, _PM.ids(pm, nw, :ne_storage), z)
end

function variable_storage_investment(pm::_PM.AbstractPowerModel; nw::Int=_PM.nw_id_default, relax::Bool=false, report::Bool=true)
    first_n = first_id(pm, nw, :hour, :scenario)
    if nw == first_n
        if !relax
            investment = _PM.var(pm, nw)[:z_strg_ne_investment] = JuMP.@variable(pm.model,
                [i in _PM.ids(pm, nw, :ne_storage)], base_name="$(nw)_z_strg_ne_investment",
                binary = true,
                start = 0
            )
        else
            investment = _PM.var(pm, nw)[:z_strg_ne_investment] = JuMP.@variable(pm.model,
                [i in _PM.ids(pm, nw, :ne_storage)], base_name="$(nw)_z_strg_ne_investment",
                lower_bound = 0,
                upper_bound = 1,
                start = 0
            )
        end
    else
        investment = _PM.var(pm, nw)[:z_strg_ne_investment] = _PM.var(pm, first_n)[:z_strg_ne_investment]
    end
    report && _PM.sol_component_value(pm, nw, :ne_storage, :investment, _PM.ids(pm, nw, :ne_storage), investment)
end

# ####################################################
# Constraint Templates: They are used to do all data manipuations and return a function with the same name,
# this way the constraint itself only containts the mathematical formulation
# ###################################################

function constraint_storage_thermal_limit_ne(pm::_PM.AbstractPowerModel, i::Int; nw::Int=_PM.nw_id_default)
    storage = _PM.ref(pm, nw, :ne_storage, i)
    constraint_storage_thermal_limit_ne(pm, nw, i, storage["thermal_rating"])
end

function constraint_storage_losses_ne(pm::_PM.AbstractPowerModel, i::Int; nw::Int=_PM.nw_id_default)
    storage = _PM.ref(pm, nw, :ne_storage, i)

    constraint_storage_losses_ne(pm, nw, i, storage["storage_bus"], storage["r"], storage["x"], storage["p_loss"], storage["q_loss"])
end

function constraint_storage_bounds_ne(pm::_PM.AbstractPowerModel, i::Int; nw::Int=_PM.nw_id_default)
    constraint_storage_bounds_ne(pm, nw, i)
end

function constraint_storage_state(pm::_PM.AbstractPowerModel, i::Int; nw::Int=_PM.nw_id_default)
    storage = _PM.ref(pm, nw, :storage, i)

    if haskey(_PM.ref(pm, nw), :time_elapsed)
        time_elapsed = _PM.ref(pm, nw, :time_elapsed)
    else
        Memento.warn(_LOGGER, "network data should specify time_elapsed, using 1.0 as a default")
        time_elapsed = 1.0
    end
    constraint_storage_state_initial(pm, nw, i, storage["energy"], storage["charge_efficiency"], storage["discharge_efficiency"], storage["stationary_energy_inflow"], storage["stationary_energy_outflow"], storage["self_discharge_rate"], time_elapsed)
end

function constraint_storage_state_ne(pm::_PM.AbstractPowerModel, i::Int; nw::Int=_PM.nw_id_default)
    storage = _PM.ref(pm, nw, :ne_storage, i)

    if haskey(_PM.ref(pm, nw), :time_elapsed)
        time_elapsed = _PM.ref(pm, nw, :time_elapsed)
    else
        Memento.warn(_LOGGER, "network data should specify time_elapsed, using 1.0 as a default")
        time_elapsed = 1.0
    end
    constraint_storage_state_initial_ne(pm, nw, i, storage["energy"], storage["charge_efficiency"], storage["discharge_efficiency"], storage["stationary_energy_inflow"], storage["stationary_energy_outflow"], storage["self_discharge_rate"], time_elapsed)
end

function constraint_storage_state(pm::_PM.AbstractPowerModel, i::Int, nw_1::Int, nw_2::Int)
    storage = _PM.ref(pm, nw_2, :storage, i)

    if haskey(_PM.ref(pm, nw_2), :time_elapsed)
        time_elapsed = _PM.ref(pm, nw_2, :time_elapsed)
    else
        Memento.warn(_LOGGER, "network $(nw_2) should specify time_elapsed, using 1.0 as a default")
        time_elapsed = 1.0
    end

    if haskey(_PM.ref(pm, nw_1, :storage), i)
        constraint_storage_state(pm, nw_1, nw_2, i, storage["charge_efficiency"], storage["discharge_efficiency"], storage["stationary_energy_inflow"], storage["stationary_energy_outflow"], storage["self_discharge_rate"], time_elapsed)
    else
        # if the storage device has status=0 in nw_1, then the stored energy variable will not exist. Initialize storage from data model instead.
        Memento.warn(_LOGGER, "storage component $(i) was not found in network $(nw_1) while building constraint_storage_state between networks $(nw_1) and $(nw_2). Using the energy value from the storage component in network $(nw_2) instead")
        constraint_storage_state_initial(pm, nw_2, i, storage["energy"], storage["charge_efficiency"], storage["discharge_efficiency"], storage["stationary_energy_inflow"], storage["stationary_energy_outflow"], storage["self_discharge_rate"], time_elapsed)
    end
end

function constraint_storage_state_ne(pm::_PM.AbstractPowerModel, i::Int, nw_1::Int, nw_2::Int)
    storage = _PM.ref(pm, nw_2, :ne_storage, i)

    if haskey(_PM.ref(pm, nw_2), :time_elapsed)
        time_elapsed = _PM.ref(pm, nw_2, :time_elapsed)
    else
        Memento.warn(_LOGGER, "network $(nw_2) should specify time_elapsed, using 1.0 as a default")
        time_elapsed = 1.0
    end

    if haskey(_PM.ref(pm, nw_1, :ne_storage), i)
        constraint_storage_state_ne(pm, nw_1, nw_2, i, storage["charge_efficiency"], storage["discharge_efficiency"], storage["stationary_energy_inflow"], storage["stationary_energy_outflow"], storage["self_discharge_rate"], time_elapsed)
    else
        # if the storage device has status=0 in nw_1, then the stored energy variable will not exist. Initialize storage from data model instead.
        Memento.warn(_LOGGER, "storage component $(i) was not found in network $(nw_1) while building constraint_storage_state between networks $(nw_1) and $(nw_2). Using the energy value from the storage component in network $(nw_2) instead")
        constraint_storage_state_initial_ne(pm, nw_2, i, storage["energy"], storage["charge_efficiency"], storage["discharge_efficiency"], storage["stationary_energy_inflow"], storage["stationary_energy_outflow"], storage["self_discharge_rate"], time_elapsed)
    end
end

function constraint_storage_state_final(pm::_PM.AbstractPowerModel, i::Int; nw::Int=_PM.nw_id_default)
    storage = _PM.ref(pm, nw, :storage, i)
    constraint_storage_state_final(pm, nw, i, storage["energy"])
end

function constraint_storage_state_final_ne(pm::_PM.AbstractPowerModel, i::Int; nw::Int=_PM.nw_id_default)
    storage = _PM.ref(pm, nw, :ne_storage, i)
    constraint_storage_state_final_ne(pm, nw, i, storage["energy"])
end

function constraint_storage_excl_slack(pm::_PM.AbstractPowerModel, i::Int; nw::Int=_PM.nw_id_default)
    constraint_storage_excl_slack(pm, nw, i)
end

function constraint_storage_excl_slack_ne(pm::_PM.AbstractPowerModel, i::Int; nw::Int=_PM.nw_id_default)
    constraint_storage_excl_slack_ne(pm, nw, i)
end

function constraint_maximum_absorption(pm::_PM.AbstractPowerModel, i::Int; nw::Int=_PM.nw_id_default)
    storage = _PM.ref(pm, nw, :storage, i)

    if haskey(_PM.ref(pm, nw), :time_elapsed)
        time_elapsed = _PM.ref(pm, nw, :time_elapsed)
    else
        Memento.warn(_LOGGER, "network data should specify time_elapsed, using 1.0 as a default")
        time_elapsed = 1.0
    end
    constraint_maximum_absorption_initial(pm, nw, i, time_elapsed)
end

function constraint_maximum_absorption_ne(pm::_PM.AbstractPowerModel, i::Int; nw::Int=_PM.nw_id_default)
    storage = _PM.ref(pm, nw, :ne_storage, i)

    if haskey(_PM.ref(pm, nw), :time_elapsed)
        time_elapsed = _PM.ref(pm, nw, :time_elapsed)
    else
        Memento.warn(_LOGGER, "network data should specify time_elapsed, using 1.0 as a default")
        time_elapsed = 1.0
    end
    constraint_maximum_absorption_initial_ne(pm, nw, i, time_elapsed)
end

function  constraint_maximum_absorption(pm::_PM.AbstractPowerModel, i::Int, nw_1::Int, nw_2::Int)
    storage = _PM.ref(pm, nw_2, :storage, i)

    if haskey(_PM.ref(pm, nw_2), :time_elapsed)
        time_elapsed = _PM.ref(pm, nw_2, :time_elapsed)
    else
        Memento.warn(_LOGGER, "network $(nw_2) should specify time_elapsed, using 1.0 as a default")
        time_elapsed = 1.0
    end

    if haskey(_PM.ref(pm, nw_1, :storage), i)
        constraint_maximum_absorption(pm, nw_1, nw_2, i, time_elapsed)
    else
        # if the storage device has status=0 in nw_1, then the stored energy variable will not exist. Initialize storage from data model instead.
        Memento.warn(_LOGGER, "storage component $(i) was not found in network $(nw_1) while building constraint_storage_state between networks $(nw_1) and $(nw_2). Using the energy value from the storage component in network $(nw_2) instead")
        constraint_maximum_absorption_initial(pm, nw_2, i, time_elapsed)
    end
end

function  constraint_maximum_absorption_ne(pm::_PM.AbstractPowerModel, i::Int, nw_1::Int, nw_2::Int)
    storage = _PM.ref(pm, nw_2, :ne_storage, i)

    if haskey(_PM.ref(pm, nw_2), :time_elapsed)
        time_elapsed = _PM.ref(pm, nw_2, :time_elapsed)
    else
        Memento.warn(_LOGGER, "network $(nw_2) should specify time_elapsed, using 1.0 as a default")
        time_elapsed = 1.0
    end

    if haskey(_PM.ref(pm, nw_1, :ne_storage), i)
        constraint_maximum_absorption_ne(pm, nw_1, nw_2, i, time_elapsed)
    else
        # if the storage device has status=0 in nw_1, then the stored energy variable will not exist. Initialize storage from data model instead.
        Memento.warn(_LOGGER, "storage component $(i) was not found in network $(nw_1) while building constraint_storage_state between networks $(nw_1) and $(nw_2). Using the energy value from the storage component in network $(nw_2) instead")
        constraint_maximum_absorption_initial_ne(pm, nw_2, i, time_elapsed)
    end
end

function constraint_ne_storage_activation(pm::_PM.AbstractPowerModel, i::Int, prev_nws::Vector{Int}, nw::Int)
    investment_horizon = [nw]
    lifetime = _PM.ref(pm, nw, :ne_storage, i, "lifetime")
    for n in Iterators.reverse(prev_nws[max(end-lifetime+2,1):end])
        i in _PM.ids(pm, n, :ne_storage) ? push!(investment_horizon, n) : break
    end
    constraint_ne_storage_activation(pm, nw, i, investment_horizon)
end

####################################################
############### Constraints
###################################################

function _PM.constraint_storage_thermal_limit(pm::BFARadPowerModel, n::Int, i, rating)
    ps = _PM.var(pm, n, :ps, i)
    qs = _PM.var(pm, n, :qs, i)

    c_perp = cos(π/8) # ~0.92
    c_diag = sin(π/8) + cos(π/8) # == cos(π/8) * sqrt(2), ~1.31

    JuMP.@constraint(pm.model, -c_perp*rating <= ps      <= c_perp*rating)
    JuMP.@constraint(pm.model, -c_perp*rating <=      qs <= c_perp*rating)
    JuMP.@constraint(pm.model, -c_diag*rating <= ps + qs <= c_diag*rating)
    JuMP.@constraint(pm.model, -c_diag*rating <= ps - qs <= c_diag*rating)
end

function constraint_storage_thermal_limit_ne(pm::_PM.AbstractActivePowerModel, n::Int, i, rating)
    ps = _PM.var(pm, n, :ps_ne, i)

    JuMP.lower_bound(ps) < -rating && JuMP.set_lower_bound(ps, -rating)
    JuMP.upper_bound(ps) >  rating && JuMP.set_upper_bound(ps,  rating)
end

function constraint_storage_thermal_limit_ne(pm::BFARadPowerModel, n::Int, i, rating)
    ps = _PM.var(pm, n, :ps_ne, i)
    qs = _PM.var(pm, n, :qs_ne, i)

    c_perp = cos(π/8) # ~0.92
    c_diag = sin(π/8) + cos(π/8) # == cos(π/8) * sqrt(2), ~1.31

    JuMP.@constraint(pm.model, -c_perp*rating <= ps      <= c_perp*rating)
    JuMP.@constraint(pm.model, -c_perp*rating <=      qs <= c_perp*rating)
    JuMP.@constraint(pm.model, -c_diag*rating <= ps + qs <= c_diag*rating)
    JuMP.@constraint(pm.model, -c_diag*rating <= ps - qs <= c_diag*rating)
end

function constraint_storage_losses_ne(pm::_PM.AbstractAPLossLessModels, n::Int, i, bus, r, x, p_loss, q_loss)
    ps = _PM.var(pm, n, :ps_ne, i)
    sc = _PM.var(pm, n, :sc_ne, i)
    sd = _PM.var(pm, n, :sd_ne, i)

    JuMP.@constraint(pm.model, ps + (sd - sc) == p_loss)
end

"Neglects the active and reactive loss terms associated with the squared current magnitude."
function constraint_storage_losses_ne(pm::_PM.AbstractBFAModel, n::Int, i, bus, r, x, p_loss, q_loss)
    ps  = _PM.var(pm, n, :ps_ne, i)
    qs  = _PM.var(pm, n, :qs_ne, i)
    sc  = _PM.var(pm, n, :sc_ne, i)
    sd  = _PM.var(pm, n, :sd_ne, i)
    qsc = _PM.var(pm, n, :qsc_ne, i)

    JuMP.@constraint(pm.model, ps + (sd - sc) == p_loss)
    JuMP.@constraint(pm.model, qs == qsc + q_loss)
end

function constraint_storage_state_initial(pm::_PM.AbstractPowerModel, n::Int, i::Int, energy, charge_eff, discharge_eff, inflow, outflow, self_discharge_rate, time_elapsed)
    sc = _PM.var(pm, n, :sc, i)
    sd = _PM.var(pm, n, :sd, i)
    se = _PM.var(pm, n, :se, i)

    JuMP.@constraint(pm.model, se == ((1-self_discharge_rate)^time_elapsed)*energy + time_elapsed*(charge_eff*sc - sd/discharge_eff + inflow - outflow))
end

function constraint_storage_state_initial_ne(pm::_PM.AbstractPowerModel, n::Int, i::Int, energy, charge_eff, discharge_eff, inflow, outflow, self_discharge_rate, time_elapsed)
    sc = _PM.var(pm, n, :sc_ne, i)
    sd = _PM.var(pm, n, :sd_ne, i)
    se = _PM.var(pm, n, :se_ne, i)
    z = _PM.var(pm, n, :z_strg_ne, i)

    JuMP.@constraint(pm.model, se == ((1-self_discharge_rate)^time_elapsed)*energy*z + time_elapsed*(charge_eff*sc - sd/discharge_eff + inflow * z - outflow * z))
end

function constraint_storage_state(pm::_PM.AbstractPowerModel, n_1::Int, n_2::Int, i::Int, charge_eff, discharge_eff, inflow, outflow, self_discharge_rate, time_elapsed)
    sc_2 = _PM.var(pm, n_2, :sc, i)
    sd_2 = _PM.var(pm, n_2, :sd, i)
    se_2 = _PM.var(pm, n_2, :se, i)
    se_1 = _PM.var(pm, n_1, :se, i)

    JuMP.@constraint(pm.model, se_2 == ((1-self_discharge_rate)^time_elapsed)*se_1 + time_elapsed*(charge_eff*sc_2 - sd_2/discharge_eff + inflow - outflow))
end

function constraint_storage_state_ne(pm::_PM.AbstractPowerModel, n_1::Int, n_2::Int, i::Int, charge_eff, discharge_eff, inflow, outflow, self_discharge_rate, time_elapsed)
    sc_2 = _PM.var(pm, n_2, :sc_ne, i)
    sd_2 = _PM.var(pm, n_2, :sd_ne, i)
    se_2 = _PM.var(pm, n_2, :se_ne, i)
    se_1 = _PM.var(pm, n_1, :se_ne, i)
    z = _PM.var(pm, n_2, :z_strg_ne, i)

    JuMP.@constraint(pm.model, se_2 == ((1-self_discharge_rate)^time_elapsed)*se_1 + time_elapsed*(charge_eff*sc_2 - sd_2/discharge_eff + inflow * z - outflow * z))
end

function constraint_storage_state_final(pm::_PM.AbstractPowerModel, n::Int, i::Int, energy)
    se = _PM.var(pm, n, :se, i)

    JuMP.@constraint(pm.model, se >= energy)
end

function constraint_storage_state_final_ne(pm::_PM.AbstractPowerModel, n::Int, i::Int, energy)
    se = _PM.var(pm, n, :se_ne, i)
    z = _PM.var(pm, n, :z_strg_ne, i)

    JuMP.@constraint(pm.model, se >= energy * z)
end

function constraint_maximum_absorption_initial(pm::_PM.AbstractPowerModel, n::Int, i::Int, time_elapsed)
    sc = _PM.var(pm, n, :sc, i)
    e_abs = _PM.var(pm, n, :e_abs, i)

    JuMP.@constraint(pm.model, e_abs == time_elapsed * sc)
end

function constraint_maximum_absorption_initial_ne(pm::_PM.AbstractPowerModel, n::Int, i::Int, time_elapsed)
    sc = _PM.var(pm, n, :sc_ne, i)
    e_abs = _PM.var(pm, n, :e_abs_ne, i)

    JuMP.@constraint(pm.model, e_abs == time_elapsed * sc)
end

function constraint_maximum_absorption(pm::_PM.AbstractPowerModel, n_1::Int, n_2::Int, i::Int, time_elapsed)
    sc_2 = _PM.var(pm, n_2, :sc, i)
    e_abs_2 = _PM.var(pm, n_2, :e_abs, i)
    e_abs_1 = _PM.var(pm, n_1, :e_abs, i)

    JuMP.@constraint(pm.model, e_abs_2 - e_abs_1 == time_elapsed * sc_2)
end

function constraint_maximum_absorption_ne(pm::_PM.AbstractPowerModel, n_1::Int, n_2::Int, i::Int, time_elapsed)
    sc_2 = _PM.var(pm, n_2, :sc_ne, i)
    e_abs_2 = _PM.var(pm, n_2, :e_abs_ne, i)
    e_abs_1 = _PM.var(pm, n_1, :e_abs_ne, i)

    JuMP.@constraint(pm.model, e_abs_2 - e_abs_1 == time_elapsed * sc_2)
end

function constraint_storage_excl_slack(pm::_PM.AbstractPowerModel, n::Int, i::Int)
    sc = _PM.var(pm, n, :sc, i)
    sd = _PM.var(pm, n, :sd, i)
    sc_max = JuMP.upper_bound(sc)
    sd_max = JuMP.upper_bound(sd)
    s_bound = max(sc_max, sd_max)

    JuMP.@constraint(pm.model, sc + sd <= s_bound)
end

function constraint_storage_excl_slack_ne(pm::_PM.AbstractPowerModel, n::Int, i::Int)
    sc = _PM.var(pm, n, :sc_ne, i)
    sd = _PM.var(pm, n, :sd_ne, i)
    sc_max = JuMP.upper_bound(sc)
    sd_max = JuMP.upper_bound(sd)
    s_bound = max(sc_max, sd_max)

    JuMP.@constraint(pm.model, sc + sd <= s_bound)
end

function constraint_storage_bounds_ne(pm::_PM.AbstractPowerModel, n::Int, i::Int)
    se = _PM.var(pm, n, :se_ne, i)
    sc = _PM.var(pm, n, :sc_ne, i)
    sd = _PM.var(pm, n, :sd_ne, i)
    ps = _PM.var(pm, n, :ps_ne, i)
    qs = _PM.var(pm, n, :qs_ne, i)
    z = _PM.var(pm, n, :z_strg_ne, i)

    se_min = JuMP.lower_bound(se)
    se_max = JuMP.upper_bound(se)
    sc_min = JuMP.lower_bound(sc)
    sc_max = JuMP.upper_bound(sc)
    sd_min = JuMP.lower_bound(sd)
    sd_max = JuMP.upper_bound(sd)
    ps_min = JuMP.lower_bound(ps)
    ps_max = JuMP.upper_bound(ps)
    qs_min = JuMP.lower_bound(qs)
    qs_max = JuMP.upper_bound(qs)

    JuMP.@constraint(pm.model, se  <= se_max * z)
    JuMP.@constraint(pm.model, se  >= se_min * z)
    JuMP.@constraint(pm.model, sc  <= sc_max * z)
    JuMP.@constraint(pm.model, sc  >= sc_min * z)
    JuMP.@constraint(pm.model, sd  <= sd_max * z)
    JuMP.@constraint(pm.model, sd  >= sd_min * z)
    JuMP.@constraint(pm.model, ps  <= ps_max * z)
    JuMP.@constraint(pm.model, ps  >= ps_min * z)
    JuMP.@constraint(pm.model, qs  <= qs_max * z)
    JuMP.@constraint(pm.model, qs  >= qs_min * z)
end

function constraint_storage_bounds_ne(pm::_PM.AbstractActivePowerModel, n::Int, i::Int)
    se = _PM.var(pm, n, :se_ne, i)
    sc = _PM.var(pm, n, :sc_ne, i)
    sd = _PM.var(pm, n, :sd_ne, i)
    ps = _PM.var(pm, n, :ps_ne, i)
    z = _PM.var(pm, n, :z_strg_ne, i)

    se_min = JuMP.lower_bound(se)
    se_max = JuMP.upper_bound(se)
    sc_min = JuMP.lower_bound(sc)
    sc_max = JuMP.upper_bound(sc)
    sd_min = JuMP.lower_bound(sd)
    sd_max = JuMP.upper_bound(sd)
    ps_min = JuMP.lower_bound(ps)
    ps_max = JuMP.upper_bound(ps)

    JuMP.@constraint(pm.model, se  <= se_max * z)
    JuMP.@constraint(pm.model, se  >= se_min * z)
    JuMP.@constraint(pm.model, sc  <= sc_max * z)
    JuMP.@constraint(pm.model, sc  >= sc_min * z)
    JuMP.@constraint(pm.model, sd  <= sd_max * z)
    JuMP.@constraint(pm.model, sd  >= sd_min * z)
    JuMP.@constraint(pm.model, ps  <= ps_max * z)
    JuMP.@constraint(pm.model, ps  >= ps_min * z)
end

function constraint_ne_storage_activation(pm::_PM.AbstractPowerModel, n::Int, i::Int, horizon::Vector{Int})
    indicator = _PM.var(pm, n, :z_strg_ne, i)
    investments = _PM.var.(Ref(pm), horizon, :z_strg_ne_investment, i)

    JuMP.@constraint(pm.model, indicator == sum(investments))
end

# Force investment decisions in two different nws to be the same.
function constraint_ne_storage_investment_same(pm::_PM.AbstractPowerModel, n_1::Int, n_2::Int, i::Int)
    investment_1 = _PM.var(pm, n_1, :z_strg_ne_investment, i)
    investment_2 = _PM.var(pm, n_2, :z_strg_ne_investment, i)

    JuMP.@constraint(pm.model, investment_1 == investment_2)
end


export ClassField

#TODO: create an interface using characters

mutable struct BadPrime <: Exception
  p
end

function Base.show(io::IO, E::BadPrime)
  if isdefined(E, :p)
    println("Bad prime $(E.p) encountered")
  else
    println("Unknown bad prime encountered")
  end
end

mutable struct ClassField_pp{S, T}
  rayclassgroupmap::S
  quotientmap::T
  a::FacElem{nf_elem, AnticNumberField}#Generator of the Kummer Extension

  sup::Array{NfOrdIdl, 1} # the support of a - if known
  sup_known::Bool

  K::NfRel{nf_elem} # the target with the roots of unity
  A::NfRel{nf_elem} # the target
  o::Int # the degree of K - note, in general this is a divisor of the degree of A
  defect::Int # div(degree(A), degree(K)) = div(degree(A), o)
  pe::NfRelElem{nf_elem} #The image of the generator of A in K
  AutG::Vector{NfRelToNfRelMor{nf_elem, nf_elem}}
  AutR::fmpz_mat
  bigK::KummerExt
  degree::Int # The degree of the relative extension we are searching for.
              # In other words, the order of the codomain of quotientmap

  function ClassField_pp{S, T}() where {S, T}
    z = new{S, T}()
    z.degree = -1
    return z
  end
end

function Base.show(io::IO, C::ClassField_pp)
  println(IOContext(io, :compact => true), "Cyclic class field of degree $(degree(C)) defined modulo $(defining_modulus(C))")
  if isdefined(C, :a)
    println(io, "Kummer generator ", C.a)
  end
  if isdefined(C, :K)
    println(io, "Kummer extension: ", C.K)
  end
  if isdefined(C, :A)
    println(io, "Defining  polynomial ", C.A.pol)
  end
end

mutable struct ClassField{S, T}
  rayclassgroupmap::S#Union{MapRayClassGrp{GrpAbFinGen}, MapClassGrp{GrpAbFinGen}}
  quotientmap::T#GrpAbFinGenMap

  conductor::Tuple{NfOrdIdl, Array{InfPlc, 1}}
  relative_discriminant::Dict{NfOrdIdl, Int}
  absolute_discriminant::Dict{fmpz,Int}
  cyc::Array{ClassField_pp{S, T}, 1}
  A::NfRel_ns{nf_elem}
  AbsAutGrpA::Vector{NfRel_nsToNfRel_nsMor{nf_elem}} #The generators for the absolute automorphism
                                                     #group of A
  degree::Int # The degree of the relative extension we are searching for.
              # In other words, the order of the codomain of quotientmap

  function ClassField{S, T}() where {S, T}
    z = new{S, T}()
    z.degree = -1
    return z
  end
end

function Base.show(io::IO, CF::ClassField)
  print(IOContext(io, :compact => true), "Class field defined mod ",
                   defining_modulus(CF), " of structure ",
                   codomain(CF.quotientmap))
end

###############################################################################
#
#  Modulus function
#
###############################################################################

@doc Markdown.doc"""
    defining_modulus(CF::ClassField)
The modulus, ie. an ideal the the set of real places, used to create the
class field.
"""
function defining_modulus(CF::ClassField)
  return _modulus(CF.rayclassgroupmap)
end

function defining_modulus(CF::ClassField_pp)
  return _modulus(CF.rayclassgroupmap)
end

function _modulus(mq::MapRayClassGrp)
  return mq.defining_modulus::Tuple{NfOrdIdl, Vector{InfPlc}}
end

function _modulus(mq::MapClassGrp)
  return (ideal(order(codomain(mq)), 1), InfPlc[])::Tuple{NfOrdIdl, Vector{InfPlc}}
end

###############################################################################
#
#  Base ring and Base field
#
###############################################################################

@doc Markdown.doc"""
    base_ring(A::ClassField)
The maximal order of the field that $A$ is defined over.
"""
function base_ring(A::ClassField)
  return order(codomain(A.rayclassgroupmap))
end

@doc Markdown.doc"""
    base_field(A::ClassField)
The number field that $A$ is defined over.
"""
function base_field(A::ClassField)
  return number_field(base_ring(A))
end

function base_ring(A::ClassField_pp)
  return order(codomain(A.rayclassgroupmap))
end

function base_field(A::ClassField_pp)
  return number_field(base_ring(A))
end

###############################################################################
#
#  Degree
#
###############################################################################

@doc Markdown.doc"""
    degree(A::ClassField)
The degree of $A$ over its base field, ie. the size of the defining ideal group.
"""
function degree(A::ClassField)
  if A.degree == -1
    A.degree = Int(order(codomain(A.quotientmap)))
  end
  return A.degree
end

function degree(A::ClassField_pp{S, T}) where {S, T}
  if A.degree == -1
    A.degree = Int(order(codomain(A.quotientmap)))
  end
  return A.degree
end

###############################################################################
#
#  Exponent
#
###############################################################################

@doc Markdown.doc"""
    exponent(A::ClassField)
The exponent of $A$ over its base field, i.e. the exponent of the Galois
group of the extension.
"""
function exponent(A::ClassField{S, T}) where {S, T}
  return Int(exponent(codomain(A.quotientmap)))
end
###############################################################################
#
#  Compositum
#
###############################################################################

@doc Markdown.doc"""
    compositum(a::ClassField, b::ClassField) -> ClassField
> The compositum of $a$ and $b$ as a (formal) class field.
"""
function compositum(a::ClassField, b::ClassField)
  @assert base_ring(a) == base_ring(b)
  c = lcm(defining_modulus(a)[1], defining_modulus(b)[1])
  d = lcm(degree(a), degree(b))
  c_inf = union(defining_modulus(a)[2], defining_modulus(b)[2])
  r, mr = ray_class_group(c, c_inf, n_quo = Int(d))
  C = ray_class_field(mr)
  @assert domain(C.rayclassgroupmap) == r
  h = norm_group_map(C, [a,b])
  U = intersect(kernel(h[1])[1], kernel(h[2])[1])
  q, mq = quo(codomain(C.quotientmap), U)
  return ray_class_field(mr, GrpAbFinGenMap(C.quotientmap * mq))
end

@doc Markdown.doc"""
  *(A::ClassField, B::ClassField) -> ClassField
The compositum of $a$ and $b$ as a (formal) class field.
"""
*(a::ClassField, b::ClassField) = compositum(a, b)

###############################################################################
#
#  Intersection
#
###############################################################################

@doc Markdown.doc"""
    intersect(a::ClassField, b::ClassField) -> ClassField
The intersect of $a$ and $b$ as a class field.
"""
function Base.intersect(a::ClassField, b::ClassField)
  @assert base_ring(a) == base_ring(b)
  c = lcm(defining_modulus(a)[1], defining_modulus(b)[1])
  c_inf = union(defining_modulus(a)[2], defining_modulus(b)[2])
  d = lcm(degree(a), degree(b))

  r, mr = ray_class_group(c, c_inf, n_quo = Int(d))
  C = ray_class_field(mr)
  h = norm_group_map(C, [a,b])
  U = kernel(h[1])[1] + kernel(h[2])[1]
  q, mq = quo(codomain(C.quotientmap), U)
  return ray_class_field(mr, GrpAbFinGenMap(C.quotientmap * mq))
end

###############################################################################
#
#  Subfield
#
###############################################################################

@doc Markdown.doc"""
    issubfield(a::ClassField, b::ClassField) -> Bool
Determines of $a$ is a subfield of $b$.
"""
function issubfield(a::ClassField, b::ClassField)
  @assert base_ring(a) == base_ring(b)
  c = lcm(defining_modulus(a)[1], defining_modulus(b)[1])
  c_inf = union(defining_modulus(a)[2], defining_modulus(b)[2])
  d = lcm(degree(a), degree(b))

  r, mr = ray_class_group(c, c_inf, n_quo = Int(d))
  C = ray_class_field(mr)
  h = norm_group_map(C, [a,b])
  return issubset(kernel(h[2])[1], kernel(h[1])[1])
end

###############################################################################
#
#  Equality
#
###############################################################################

@doc Markdown.doc"""
    ==(a::ClassField, b::ClassField)
Tests if $a$ and $b$ are equal.
"""
function ==(a::ClassField, b::ClassField)
  @assert base_ring(a) == base_ring(b)
  mq1 = a.quotientmap
  mq2 = b.quotientmap
  if !isisomorphic(codomain(mq1), codomain(mq2))
    return false
  end
  expo = Int(exponent(codomain(mq1)))
  c = lcm(defining_modulus(a)[1], defining_modulus(b)[1])
  c_inf = union(defining_modulus(a)[2], defining_modulus(b)[2])
  r, mr = ray_class_group(c, c_inf, n_quo = expo)
  C = ray_class_field(mr)
  @assert defining_modulus(C) == (c, c_inf)
  h = norm_group_map(C, [a,b])
  return iseq(kernel(h[2])[1], kernel(h[1])[1])
end

###############################################################################
#
#  IsCyclic
#
###############################################################################

@doc Markdown.doc"""
    iscyclic(C::ClassField)
Tests if the (relative) automorphism group of $C$ is cyclic (by checking
the defining ideal group).
"""
function iscyclic(C::ClassField)
  mp = C.quotientmap
  return iscyclic(codomain(mp))
end

###############################################################################
#
#  Maximal p-subfield
#
###############################################################################

@doc Markdown.doc"""
    maximal_p_subfield(C::ClassField, p::Int)
Returns the class field corresponding to the maximal subextension of
prime power degree
"""
function maximal_p_subfield(C::ClassField, p::Int)
  mg = C.quotientmap
  v = valuation(degree(C), p)
  q, mq = quo(codomain(mg), p^v, false)
  s, ms = snf(q)
  mg1 = GrpAbFinGenMap(mg*mq*ms)
  return ray_class_field(C.rayclassgroupmap, mg1)
end