module SplitDataset

using DiskArrays
using Glob
using HDF5
using JSON3
using Printf
using PaddedBlocks

struct H5SplitDataset{T,N} <: AbstractDiskArray{T, N}
    files::Array{String,1}
    blocks::Blocks{N}
    stackname::String
end

haschunks(a::H5SplitDataset) = DiskArrays.Chunked()
eachchunk(a::H5SplitDataset) = DiskArrays.GridChunks(a, a.blocks.block_size)
Base.size(a::H5SplitDataset) = a.blocks.full_size

"Collects filenames and block dimensions for a HDF5 dataset"
function file_blocks_names(path)
    return h5open(path, "r") do file
        return Blocks(
            Int64.(tuple(collect_attributes(file["stack_metadata"]["full_size"])...)[end:-1:1]),
            Int64.(tuple(collect_attributes(file["stack_metadata"]["block_size"])...)[end:-1:1]),
        ),
        collect_attributes(file["stack_metadata"]["files"])
    end
end

function stack_metadata_from_json(path)
    mddict = open(path, "r") do f
        return JSON3.read(f)
    end
    return Blocks(
        Int64.(Tuple(mddict["shape_full"]))[end:-1:1],
        Int64.(Tuple(mddict["shape_block"]))[end:-1:1],
    )
end

function H5SplitDataset(folder::String, prefix::String = "")
    jsonpath = joinpath(folder, "stack_metadata.json")
    files = sort(glob(prefix * "*.h5", folder))
    if isfile(jsonpath)
        blocks = stack_metadata_from_json(jsonpath)
        filenames = files
    else
        blocks, filenames = file_blocks_names(files[end])
    end
    n_dim = length(blocks.full_size)

    stack_name = h5open(files[1]) do f
        if "stack" in names(f)
            return "stack"
        else
            return "stack_$(n_dim)D"
        end
    end

    el1 = h5read(files[1], stack_name, tuple((1:1 for _ in 1:n_dim)...))
    return H5SplitDataset{eltype(el1),n_dim}(files, blocks, stack_name)
end

function H5SplitDataset(T, path, full_size::NTuple{N, Int}, block_size::NTuple{N, Int}) where {N}
    blocks = Blocks(Int64.(full_size), Int64.(block_size))
    files = [joinpath(path, @sprintf("%04d.h5", i)) for i in 0:(length(blocks)-1)]
    if !isdir(path)
        mkpath(path)
    end
    open(joinpath(path, "stack_metadata.json"), "w") do f
        return JSON3.write(
            f,
            Dict(
                "shape_full" => blocks.full_size[end:-1:1],
                "shape_block" => blocks.block_size[end:-1:1],
            ),
        )
    end
    return H5SplitDataset{T, length(full_size)}(files, blocks, "stack")
end

function H5SplitDataset(
    folder,
    a::AbstractArray{T,N},
    block_size::NTuple{N,Int},
) where {N,T}
    blocks = Blocks(size(a), block_size)
    mkpath(folder)
    files = String[]
    for (i_block, sl) in enumerate(slices(blocks))
        filename = @sprintf("%04d.h5", i_block)
        push!(files, filename)
        savepath = joinpath(folder, filename)
        el1 = h5write(savepath, "stack_$(N)D", a[sl...])
    end
    open(joinpath(folder, "stack_metadata.json"), "w") do f
        return JSON3.write(
            f,
            Dict(
                "shape_full" => blocks.full_size[end:-1:1],
                "shape_block" => blocks.block_size[end:-1:1],
            ),
        )
    end
    return H5SplitDataset(folder)
end

"Collects attributes form a group in a deepdish-saved HDF5 file"
function collect_attributes(hgroup)
    atrnms = names(attrs(hgroup))
    nums = sort([
        begin
            mtch = match(r"i(\d+)", nam)
            parse(mtch.captures[1])
        end for nam in filter(n -> ismatch(r"i(\d+)", n), atrnms)
    ])
    return [read(attrs(hgroup)["i$i"]) for i in nums]
end

function get_matching_slices(f_idx, file_limits, block_limits, block_size)
    source_slices = tuple((
        (ci == s ? si : 1):(ci == e ? ei : bs)
        for
        (ci, s, e, si, ei, bs) in zip(
            f_idx,
            file_limits[1, :],
            file_limits[2, :],
            block_limits[1, :],
            block_limits[2, :],
            block_size
        ))...)

    read_size = tuple((sl.stop - sl.start + 1 for sl in source_slices)...)

    rel_idx = f_idx .- tuple(file_limits[1, :]...) .+ 1
    target_slices = tuple((
        begin
            st = (st_idx == 1 ? 1 : (st_idx - 2) * bs + bs - first_idx + 2)
            st:(st+sz-1)
        end for
        (st_idx, bs, first_idx, sz) in
        zip(rel_idx, block_size, block_limits[1, :], read_size)
    )...)
    return source_slices, target_slices
end

function DiskArrays.readblock!(dset::H5SplitDataset{T, N}, target, i::AbstractUnitRange...) where {T,N}
    idx = i
    file_limits, block_limits = PaddedBlocks.blocks_to_take(dset.blocks, idx)
    s2i = LinearIndices(dset.blocks.blocks_per_dim)
    for f_idx in
        Iterators.product((s:e for (s, e) in zip(file_limits[1, :], file_limits[2, :]))...)
        i_file = s2i[f_idx...]
        source_slices, target_slices = get_matching_slices(f_idx, file_limits, block_limits, dset.blocks.block_size)
        ar = h5read(dset.files[i_file], dset.stackname, source_slices)
        target[target_slices...] .= ar

    end
end

function DiskArrays.writeblock!(dset::H5SplitDataset{T, N}, target, i::AbstractUnitRange...) where {T,N}
    idx = i
    file_limits, block_limits = PaddedBlocks.blocks_to_take(dset.blocks, idx)
    s2i = LinearIndices(dset.blocks.blocks_per_dim)
    for f_idx in
        Iterators.product((s:e for (s, e) in zip(file_limits[1, :], file_limits[2, :]))...)
        i_file = s2i[f_idx...]

        source_slices, target_slices = get_matching_slices(f_idx, file_limits, block_limits, dset.blocks.block_size)

        # TODO handle wrong existing datasets
        h5open(dset.files[i_file], "cw") do fid
            if !exists(fid, dset.stackname)
                h5dset = d_create(fid, dset.stackname, datatype(T),
                dataspace(dset.blocks.block_size...),
                "chunk", HDF5.heuristic_chunk(T, (dset.blocks.block_size)), "blosc", 3)
            else
                h5dset = fid[dset.stackname]
            end

            h5dset[source_slices...] = target[target_slices...]
        end

    end
end

export H5SplitDataset

end


module Whirl
using LinearAlgebra
using StaticArrays
using Statistics
using Plots
import Base

include("num.jl")
include("error.jl")
include("debug.jl")
include("geometry.jl")
include("units.jl")
include("eos.jl")
include("kernel.jl")
include("initial.jl")
include("velocity.jl")
include("mass.jl")
include("energy.jl")
include("pressure.jl")
include("diffusion.jl")
include("boundary.jl")
include("fluid.jl")
include("stepping.jl")
include("simulation.jl")
include("experiments.jl")
include("plotting.jl")

include("test.jl")

end # module


module InheritanceTest
using ConcreteInterfaces.ComputedFieldTypes
using Base.Test

abstract type Bar end

@computed struct Parametric{T} <: Bar
    x::Base.promote_op(+, T, Float64)
end

@computed struct NonParametric <: Bar
    x::Base.promote_op(+, Int, Float64)
end

@computed struct NonParametricNoParent
    x::Base.promote_op(+, Int, Float64)
end

@testset "inheritance" begin
    @test isa(@inferred(Parametric{Int}(1)), Parametric{Int, Float64})
    @test isa(Parametric{Int}(1), Bar)
    @test isa(@inferred(Parametric{Float64}(1.0)), Parametric{Float64, Float64})

    @test isa(@inferred(NonParametric(1)), NonParametric{Float64})
    @test isa(NonParametric(1), Bar)

    @test isa(@inferred(NonParametricNoParent(1)), NonParametricNoParent{Float64})
    @test !isa(NonParametricNoParent(1), Bar)
end
end


# Exports
export Orbit
export interpolate
export Spacecraft
export Location
export Request
export load_requests
export set_collect_duration
export GroundStation
export load_stations
export Opportunity
export Done
export Noop
export Sunpoint
export Collect
export Contact
export SatPlanningProblem
export add_locations
export clear_locations
export clear_opportunities
export clear_all

#########
# Orbit #
#########

"""
Orbit provides a data structure which wraps SatelliteDynamics propagators

Arguments:
- `epc0::Union{Epoch, Real}` Initial Epoch of propagation
- `epcf::Union{Epoch, Real}` Final Epoch of propagation
- `eci0::Array{<:Real, 1}` Initial inertial state at initial Epoch [m; m/s]
- `id::Integer` Identification number for this orbit
- `timestep::Real` Timestep used for propgation output
- `dtmax::Real` Maximum allowable integrator step
- `mass::Real` Satellite mass [kg]
- `area_drag::Real` Wetted wind-facing area [m^2]
- `coef_drag::Real` Coefficient of drag
- `area_srp::Real` Area illuminated by solar radiation pressure [m^2]
- `coef_srp::Real` Coefficient of reflectivity
- `n_grav::Integer` Gravity model degree
- `m_grav::Integer` Gravity model order
- `drag::Bool` Enable drag perturbation in propgation
- `srp::Bool` Enable solar radiation pressure perturbation in propgation
- `moon::Bool` Enable thirdbody moon perturbation in propgation
- `sun::Bool` Enable thirdbody sun perturbation in propgation
- `relativity::Bool` Enable relativity perturbation in propgation

Attributes:
- `id::Int32` Identification number for the specific satellite
- `t::Array{Float64, 1}` Elapsed time of propagator output
- `epc::Array{Epoch, 1}` Epoch of propagator output
- `eci::Array{Float64, 2}` Earth Centered Inertial state of satellite output. State entries alightex with column output.
- `ecef::Array{Float64, 2}` Earth Centered Earth Fixed state of statellite output. State entries alightex with column output.
- `oe::Array{Float64, 2}` Osculating orbital element state of statellite output. State entries alightex with column output.
"""
mutable struct Orbit
    id::Int32
    t::Array{Float64, 1}
    epc::Array{Epoch, 1}
    eci::Array{Float64, 2}
    ecef::Array{Float64, 2}
    oe::Array{Float64, 2}

    function Orbit(epc0::Union{Epoch, Real},  epcf::Union{Epoch, Real}, eci0::Array{<:Real, 1};
                id=0::Integer, timestep=1.0::Real, dtmax=30.0::Real,
                mass=100.0::Real, area_drag=1.0::Real, coef_drag=2.3::Real,
                area_srp=1.0::Real, coef_srp=1.8::Real,
                n_grav=20::Integer, m_grav=20::Integer,
                drag=true::Bool, srp=true::Bool, moon=true::Bool,
                sun=true::Bool, relativity=true::Bool)

        # Initialize State Vector
        orb  = EarthInertialState(epc0, eci0, dt=timestep,
            mass=mass,
            area_drag=area_drag, coef_drag=coef_drag,
            area_srp=area_srp, coef_srp=coef_srp,
            n_grav=n_grav, m_grav=m_grav,
            drag=drag, srp=srp,
            moon=moon, sun=sun,
            relativity=relativity
        )

        # Propagate the orbit
        t, epc, eci = sim!(orb, epcf)

        ecef = zeros(Float64, 6, length(epc))
        oe   = zeros(Float64, 6, length(epc))

        for i in 1:length(t)
            ecef[:, i] = sECItoECEF(epc[i], eci[:, i])
            oe[:, i]   = sCARTtoOSC(eci[:, i], use_degrees=true)
        end

        new(id, t, epc, eci, ecef, oe)
    end
end

function Base.show(io::IO, orb::Orbit)

    s = @sprintf "Orbit(%d)" orb.id

    print(io, s)
end

# Helpfer function
function searchsortednearest(a,x)
    idx = searchsortedfirst(a,x)
    if (idx==1); return idx; end
    if (idx>length(a)); return length(a); end
    if (a[idx]==x); return idx; end
    if (abs(a[idx]-x) < abs(a[idx-1]-x))
       return idx
    else
       return idx-1
    end
 end

"""
Interpolate orbit information to a given `Epoch`.

This interpolation is not particularly accurate, care should be used to determine
that the use-case can handle the interprolation errors introduced.

Arguments:
- `orbit::Orbit` Orbit data to interpolate
- `epc::Union{Epoch, Real}` Epoch to interpolate state output to.

Returns:
- `eci::Arrray{Float64, 1}` Earth intertial state information interpolated to time of Epoch
"""
function interpolate(orbit::Orbit, epc::Union{Epoch, Real})
    # Check validity of input
    if epc < orbit.epc[1]
        throw(ArgumentError("Invaid interpolation time. Cannot interpolate before start of Orbit."))
    end

    if epc > orbit.epc[end]
        throw(ArgumentError("Invaid interpolation time. Cannot interpolate after end of Orbit."))
    end

    # Exit early if epoch is end of interpolation (corner case)
    if epc == orbit.epc[end]
        return orbit.eci[:, end]
    end

    # Get Uppler and lower time bounds
    idx_l = searchsortedfirst(orbit.epc, epc)
    idx_u = idx_l + 1

    # Epoch
    epc_l = orbit.epc[idx_l]
    epc_u = orbit.epc[idx_u]

    # Earth Intertial State
    eci_l = orbit.eci[:, idx_l]
    eci_u = orbit.eci[:, idx_u]

    # Interpolate
    eci = zeros(typeof(eci_l[1]), 6)
    for i in 1:length(eci)
        eci[i] = (eci_u[i] - eci_l[i])/(epc_u - epc_l)*(epc - epc_l) + eci_l[i]
    end

    return eci
end

##############
# Spacecraft #
##############

"""
Object to store spacecraft information

Model Parameters:
- `powergen_sunpoint::Float64` Power generated when sunpointed
- `powergen_image::Float64` Power generated when imaging (should be negative)
- `powergen_downlink::Float64` Power generated when downlinking (should be negative)
- `datagen_backorbit::Float64` Data generated in backorbit
- `datagen_image::Float64` Data generated when imaging
- `datagen_downlink::Float64` Data generated when downlinking (should be negative)
- `slew_rate::Float64` Spacecraft slew rate (degrees per second)
"""
@with_kw mutable struct Spacecraft
    # Core Parameters
    id::Integer
    tle::TLE

    # Resource Model
    powergen_sunpoint::Float64 = 0.0
    powergen_image::Float64 = 0.0
    powergen_downlink::Float64 = 0.0
    datagen_backorbit::Float64 = 0.0
    datagen_image::Float64 = 0.0
    datagen_downlink::Float64 = 0.0
    slew_rate::Float64 = 1.0
end

############
# Location #
############

abstract type Location end

function Base.isequal(ll::Location, rl::Location)
    return (
        (ll.id == rl.id) &&
        (ll.lat == rl.lat) &&
        (ll.lon == rl.lon) &&
        (ll.ecef == rl.ecef) && true
    )
end
Base.:(==)(ll::Location, rl::Location) = Base.isequal(ll, rl)

function Base.show(io::IO, loc::Location)

    s = @sprintf "%s(ID: %s, Lon: %.3f, Lat: %0.3f)" typeof(loc) string(loc.id) loc.lon loc.lat

    print(io, s)
end

###########
# Request #
###########

"""
Request represents a single area/point on Earth to be observed. It represents a location
for which it is possible to collect the entire area in a singel collect

Attributes:
- `lon::Float64` Longitude of request center [deg]
- `lat::Float64` Lattitude of iamge center [deg]
- `ecef::Array{Float64, 1}` Earth-Centered-Earth-Fixed Coordinates of request center [m]
- `look_angle_min::Float64` Minimum look angle of valid collect
- `look_angle_max::Float64` Maximum look angle of valid collect
- `collect_duration` Required time to well of request for a feasible collection
- `require_zero_doppler::Bool` Require that the collect occurs centered at zero Doppler offset
- `reward::Float64` Reward for for request collection
- `id::Union{UUID, Integer}` Unique request identifier
"""
mutable struct Request <: Location
    lon::Float64
    lat::Float64
    ecef::Array{Float64, 1}
    look_angle_min::Float64
    look_angle_max::Float64
    require_zero_doppler::Bool
    collect_duration::Float64
    reward::Float64
    id::Union{UUID, Integer}

    function Request(lon::Real, lat::Real ; id::Union{UUID, Integer}=uuid4(),
                    look_angle_min::Real=5.0, look_angle_max::Real=55.0,
                    require_zero_doppler::Bool=true, collect_duration::Real=1.0,
                    reward::Real=1.0, use_degrees::Bool=true)

        # Compute ECEF coordinates of scene center
        ecef = sGEODtoECEF([lon, lat], use_degrees=use_degrees::Real)

        # Generate new request object
        new(lon, lat, ecef, look_angle_min, look_angle_max,
            require_zero_doppler, collect_duration, reward, id)
    end
end

Base.hash(r::Request, h::UInt) = hash(r.id, hash(r.lon, hash(r.lat, hash(r.look_angle_min, hash(r.look_angle_max, hash(:Request, h))))))

function Base.show(io::IO, req::Request)

    s = @sprintf "Request(ID: %s, Lon: %.3f, Lat: %0.3f, Reward: %.2f, Duration: %.2f)" string(req.id) req.lon req.lat req.reward req.collect_duration

    print(io, s)
end

"""
Loads request data from a JSON file.

THe JSON file is expected to contain an array of JSON objects where each object
has the following format:

```json
{
    "lon": 22,
    "lat": 12.23,
    "reward": 42.0,
    "look_angle_min": 5.0,
    "look_angle_max": 55.0
}
```

Arguments:
- `file::String` Filepath to JSON file encoding request
"""
function load_requests(file::String; id_offset::Integer=0, require_zero_doppler::Bool=true)
    # Parse data file
    data = JSON.parsefile(file)

    requests = Array{Request, 1}(undef, length(data))

    for (i, req) in enumerate(data)
        lon = req["lon"]
        lat = req["lat"]
        look_angle_min = req["look_angle_min"]
        look_angle_max = req["look_angle_max"]
        id  = i + id_offset

        requests[i] = Request(lon, lat, id=id,
            look_angle_min=look_angle_min,
            look_angle_max=look_angle_max
        )
    end

    println("Loaded $(length(requests)) requests.")

    return requests
end

"""
Set the request collection duration for satellites in the constellation.
"""
function set_collect_duration(requests::Array{<:Request, 1}, collect_duration::Real)
    if collect_duration <= 0.0
        throw(ArgumentError("Invalid collection duration."))
    end

    for request in requests
        request.collect_duration = collect_duration
    end
end

#################
# GroundStation #
#################

"""
GroundStation represents a location which the satellite can use to communicate
with the ground.

Attributes:
- `lon::Float64` Longitude of request center [deg]
- `lat::Float64` Lattitude of iamge center [deg]
- `ecef::Array{Float64, 1}` Earth-Centered-Earth-Fixed Coordinates of request center [m]
- `elevation_min::Float64` Minimum look angle of valid collect
- `id::Union{UUID, Integer}` Unique station identifier
"""
mutable struct GroundStation <: Location
    lon::Float64
    lat::Float64
    ecef::Array{Float64, 1}
    elevation_min::Float64
    duration_min::Float64
    id::Union{UUID, Integer}

    function GroundStation(lon::Real, lat::Real ; id::Union{UUID, Integer}=uuid4(),
                    elevation_min::Real=5.0, duration_min::Real=0.0,
                    use_degrees::Bool=true)

        # Compute ECEF coordinates of scene center
        ecef = sGEODtoECEF([lon, lat], use_degrees=use_degrees)

        # Generate new request object
        new(lon, lat, ecef, elevation_min, duration_min, id)
    end
end

Base.hash(s::GroundStation, h::UInt) = hash(s.id, hash(s.lon, hash(s.lat, hash(s.elevation_min, hash(:GroundStation, h)))))

# function Base.isequal(ll::GroundStation, rl::GroundStation)
#     return (
#         (ll.id == rl.id) &&
#         (ll.lat == rl.lat) &&
#         (ll.lon == rl.lon) &&
#         (ll.ecef == rl.ecef) && true
#     )
# end
# Base.:(==)(ll::GroundStation, rl::GroundStation) = Base.isequal(ll, rl)

function Base.show(io::IO, sta::GroundStation)

    s = @sprintf "GroundStation(ID: %s, Lon: %.3f, Lat: %0.3f)" string(sta.id) sta.lon sta.lat

    print(io, s)
end

"""
Loads request data from a JSON file.

THe JSON file is expected to contain an array of JSON objects where each object
has the following format:

```json
{
    "lon": 22,
    "lat": 12.23,
    "elevation_min": 5.0
}
```

Arguments:
- `file::String` Filepath to JSON file encoding request
"""
function load_stations(file::String; id_offset::Integer=0)
    data = JSON.parsefile(file)

    # Initialize array of Requests
    stations = Array{GroundStation, 1}(undef, length(data))

    for (i, sta) in enumerate(data)
        lon = sta["lon"]
        lat = sta["lat"]
        elevation_min = sta["elevation_min"]
        id  = i + id_offset

        stations[i] = GroundStation(lon, lat, id=id, elevation_min=elevation_min)
    end

    println("Loaded $(length(stations)) stations.")

    return stations
end

"""
Set minimum access duration for all stations
"""
function set_station_duration_min(stations::Array{<:GroundStation, 1}, duration_min::Real)
    if duration_min <= 0.0
        throw(ArgumentError("Invalid minimum contact duration duration."))
    end

    for station in stations
        station.duration_min = duration_min
    end
end

###############
# Opportunity #
###############

"""
Opportunity represents an "opportunity" for action for a spacecraft

Common Attributes:
- `id::Integer` Unique Identifier for action
- `orbit::TLE` Orbit object associated with this collect
- `location::Location` Location object on the ground for interaction
- `t_start::Union{Epoch, Real}` Start of acquisition window
- `t_mid::Union{Epoch, Real}` Mid-time of acquisition window
- `t_end::Union{Epoch, Real}` End of possible acquisition window
- `duration::Real` Length of action
"""
abstract type Opportunity end

"""
"""
@with_kw mutable struct Done <: Opportunity
    t_start::Union{Epoch, Real} = Inf
end

"""
No-operation action. Advances state without changing other values
"""
@with_kw mutable struct Noop <: Opportunity
    id::Integer = 0
    t_start::Union{Epoch, Real} = 0.0
end

function Base.show(io::IO, noop::Noop)

    s = @sprintf "Noop(time: %s)" noop.t_start

    print(io, s)
end

@with_kw mutable struct Sunpoint <: Opportunity
    id::Integer = 0
    t_start::Union{Epoch, Real} = 0.0
end

function Base.show(io::IO, sp::Sunpoint)

    s = @sprintf "Sunpoint(time: %s)" sp.t_start

    print(io, s)
end

"""
Opportunity type for a Request collection.
"""
@with_kw mutable struct Collect <: Opportunity
    id::Integer
    spacecraft::Spacecraft
    location::Union{Request, Nothing}
    t_start::Union{Epoch, Real}
    t_mid::Union{Epoch, Real}
    t_end::Union{Epoch, Real}
    duration::Float64
    reward::Float64
    nr::Int

    function Collect(t_start::Union{Epoch, Real}, t_end::Union{Epoch, Real};
        id::Integer=0, spacecraft=nothing::Union{Spacecraft, Nothing},
        location=nothing::Union{Location, Nothing},
        nr::Int=0)

        duration = t_end - t_start
        t_mid    = t_start + duration/2.0
        new(id, spacecraft, location, t_start, t_mid, t_end, duration, location.reward, nr)
    end
end

Base.hash(c::Collect, h::UInt) = hash(c.id, hash(c.spacecraft.id, hash(c.duration, hash(:Collect, h))))

function JSON.lower(col::Collect)
    return Dict(
        "id" => col.id,
        "location" => col.location.id,
        "lon" => col.location.lon,
        "lat" => col.location.lat,
        "t_start" => string(col.t_start),
        "t_end"  => string(col.t_end),
        "duration" => col.duration
    )
end

function Base.isequal(cl::Collect, cr::Collect)
    return (
        (cl.id == cr.id) &&
        ((cl.spacecraft == nothing && cr.spacecraft == nothing) ||
         ((cl.spacecraft != nothing && cr.spacecraft != nothing &&
          (cl.spacecraft.id == cr.spacecraft.id)))) &&
        ((cl.location == nothing && cr.location == nothing) ||
         ((cl.location != nothing && cr.location != nothing &&
          (cl.location.id == cr.location.id)))) &&
        (cl.t_start == cr.t_start) &&
        (cl.t_end == cr.t_end) && true
    )
end
Base.:(==)(cl::Collect, cr::Collect) = Base.isequal(cl, cr)

function Base.show(io::IO, col::Collect)

    if typeof(col.t_start) == Epoch
        s = @sprintf "Collect(ID: %s, Spacecraft %s, Request: %s, t_start: %s, t_end: %s, Reward: %.2f, Remaining: %d)" string(col.id) string(col.spacecraft.id) string(col.location.id) col.t_start col.t_end col.reward col.nr
    else
        s = @sprintf "Collect(ID: %s, Spacecraft %s, Request: %s, t_start: %.3f, t_end: %.3f, Reward: %.2f, Remaining: %d)" string(col.id) string(col.spacecraft.id) string(col.location.id) col.t_start col.t_end col.reward col.nr
    end

    print(io, s)
end

"""
Opportunity type for a ground contact.
"""
mutable struct Contact <: Opportunity
    id::Integer
    spacecraft::Spacecraft
    location::Union{GroundStation, Nothing}
    t_start::Union{Epoch, Real}
    t_mid::Union{Epoch, Real}
    t_end::Union{Epoch, Real}
    duration::Float64
    nr::Int

    function Contact(t_start::Union{Epoch, Real}, t_end::Union{Epoch, Real};
        id::Integer=0, spacecraft=nothing::Union{Spacecraft, Nothing},
        location=nothing::Union{Location, Nothing}, nr::Int=0)

        duration = t_end - t_start
        t_mid    = t_start + duration/2.0
        new(id, spacecraft, location, t_start, t_mid, t_end, duration, nr)
    end
end

Base.hash(c::Contact, h::UInt) = hash(c.id, hash(c.spacecraft.id, hash(c.duration, hash(:Contact, h))))

function JSON.lower(con::Contact)
    return Dict(
        "id" => con.id,
        "location" => con.location.id,
        "lon" => con.location.lon,
        "lat" => con.location.lat,
        "t_start" => string(con.t_start),
        "t_end"  => string(con.t_end),
        "duration" => con.duration
    )
end

function Base.isequal(cl::Contact, cr::Contact)
    return (
        (cl.id == cr.id) &&
        ((cl.spacecraft == nothing && cr.spacecraft == nothing) ||
         ((cl.spacecraft != nothing && cr.spacecraft != nothing &&
          (cl.spacecraft.id == cr.spacecraft.id)))) &&
        ((cl.location == nothing && cr.location == nothing) ||
         ((cl.location != nothing && cr.location != nothing &&
          (cl.location.id == cr.location.id)))) &&
        (cl.t_start == cr.t_start) &&
        (cl.t_end == cr.t_end) && true
    )
end
Base.:(==)(cl::Contact, cr::Contact) = Base.isequal(cl, cr)

function Base.show(io::IO, con::Contact)

    minutes = floor(con.duration/60)
    seconds = con.duration - 60*minutes

    if typeof(con.t_start) == Epoch
        s = @sprintf "Contact(ID: %s,  Spacecraft %s, Station: %s, t_start: %s, t_end: %s, Duration: %02dm %02ds, Remaining: %d)" string(con.id) string(con.spacecraft.id) string(con.location.id) con.t_start con.t_end minutes seconds con.nr
    else
        s = @sprintf "Contact(ID: %s,  Spacecraft %s, Station: %s, t_start: %.3f, t_end: %.3f, Duration: %02dm %02ds, Remaining: %d)" string(con.id) string(con.spacecraft.id) string(con.location.id) con.t_start con.t_end minutes seconds con.nr
    end

    print(io, s)
end

#############
# MDP State #
#############

@with_kw struct SatMDPState
    time::Union{Epoch, Float64}
    last_cdo_action::Opportunity
    last_action::Opportunity
    request_ids::Array{Int, 1} = Int[]
    power::Float64 = 1.0
    data::Float64 = 0.0
end
Base.hash(s::SatMDPState, h::UInt) = hash(s.time, hash(s.last_action, hash(s.power, hash(s.data, hash(s.request_ids, hash(:SatMDPState, h))))))

function Base.show(io::IO, state::SatMDPState)

    s = @sprintf "SatMDPState(Time: %s, Last CDO Action: %s, Action: %s, Requests: %d, Power: %.2f, Data: %.2f)" string(state.time) string(state.last_cdo_action) string(state.last_action) length(state.request_ids) state.power state.data

    print(io, s)
end

function POMDPs.isequal(sl::SatMDPState, sr::SatMDPState)
    return (
        (sl.time == sr.time) &&
        (sl.last_action == sr.last_action) &&
        (sl.last_cdo_action == sr.last_cdo_action) &&
        (sl.power == sr.power) &&
        (sl.data == sr.data) &&
        (sl.request_ids == sr.request_ids)
    )
end

####################
# Planning Problem #
####################

@with_kw mutable struct SatPlanningProblem <: POMDPs.MDP{SatMDPState, Opportunity}
    # General Planning Settings
    t_start::Epoch
    t_end::Epoch

    # Spacecraft Settings
    spacecraft::Array{Spacecraft, 1} = Spacecraft[]
    constraints::Array{Function,1} = Function[]

    # Planning Locations
    locations::Array{<:Location, 1} = Location[]
    stations::Array{GroundStation, 1} = GroundStation[]
    requests::Array{Request, 1} = Request[]

    # Planning Opportunities
    opportunities::Array{<:Opportunity, 1} = Opportunity[]
    contacts::Array{Contact, 1} = Contact[]
    collects::Array{Collect, 1} = Collect[]

    # Lookup Tables
    lt_locations::Dict{Integer, Location} = Dict{Integer, Location}()
    lt_opportunities::Dict{Integer, Opportunity} = Dict{Integer, Opportunity}()
    lt_contacts::Dict{Integer, Contact} = Dict{Integer, Contact}()
    lt_collects::Dict{Integer, Collect} = Dict{Integer, Collect}()
    lt_loc_opps::Dict{Integer, AbstractVector{Integer}} = Dict{Integer, AbstractVector{Integer}}()

    # Action Lookup Table
    actions::AbstractVector{Opportunity} = Opportunity[]
    lt_actions::Dict{Integer, Opportunity} = Dict{Integer, Opportunity}()
    lt_feasible_actions::Dict{Tuple{Integer, Integer}, AbstractVector{Opportunity}} = Dict{Tuple{Opportunity, Opportunity}, AbstractVector{Opportunity}}()

    # Solver Parameters - General
    solve_allow_repeats::Bool = false
    solve_discount::Real   = 0.99
    solve_depth::Int    = 3
    solve_breadth::Int  = 0
    solve_horizon::Real = 90.0

    # # Solver Parameters - MDP
    # mdp_reward_scarcity::Bool = false

    # Solver Parameters - MCTS
    mcts_n_iterations::Int = 50
    mcts_exploration_constant::Real = 1.0
end

"""
Clear problem locations.
"""
function clear_locations(problem::SatPlanningProblem)
    # Reset Locations
    problem.locations = Location[]
    problem.stations = GroundStation[]
    problem.requests = Request[]
end

"""
Clear problem opportunities.
"""
function clear_opportunities(problem::SatPlanningProblem)
    # Reset Opportunities
    problem.opportunities = Opportunity[]
    problem.contacts = Contact[]
    problem.collects = Collect[]

    problem.lt_opportunities = Dict{Integer, Opportunity}()
    problem.lt_contacts      = Dict{Integer, Contact}()
    problem.lt_collects      = Dict{Integer, Collect}()
    problem.lt_loc_opps      = Dict{Integer, Array{Integer, 1}}()

    # Action Lookup Table
    lt_actions = Opportunity[]
    lt_feasible_actions = Dict{Tuple{Opportunity, Opportunity}, AbstractVector{Opportunity}}()
end

"""
Clear all problem opportunities
"""
function clear_all(problem::SatPlanningProblem)
    clear_locations(problem)
    clear_opportunities(problem)
end

"""
Add locations to planning problem.
"""
function add_locations(problem::SatPlanningProblem, locations::Array{<:Location})
    if length(locations) == 0
        @warn "No locations in array. Nothing to add..."
        return
    end

    # TODO: Check for duplicate IDs in requests (assumes already unique)

    # Add location to problem
    push!(problem.locations, locations...)

    # Create separate arrays of contacts and requests
    push!(problem.stations, filter(x -> typeof(x) == GroundStation, locations)...)
    push!(problem.requests, filter(x -> typeof(x) == Request, locations)...)

    # Update location lookup table
    for loc in problem.locations
        problem.lt_locations[loc.id] = loc
    end

    return
end


const SVF = MOI.SingleVariable
const SAF = MOI.ScalarAffineFunction{Float64}

# indices
const VI = MOI.VariableIndex
const CI = MOI.ConstraintIndex

# sets
const BOUNDS = Union{
    MOI.EqualTo{Float64},
    MOI.GreaterThan{Float64},
    MOI.LessThan{Float64},
    MOI.Interval{Float64}
}

const VAR_TYPES = Union{
    MOI.ZeroOne,
    MOI.Integer
}

"""
Optimizer struct
"""
mutable struct Optimizer <: MOI.AbstractOptimizer
    inner::Union{CoM, Nothing}
    variable_info::Vector{Variable}
    # which variable index, (:leq,:geq,:eq,:Int,:Bin), and lower and upper bound
    var_constraints::Vector{Tuple{Int64,Symbol,Int64,Int64}}
    status::MOI.TerminationStatusCode
    options::SolverOptions
end

include("variables.jl")
include("constraints.jl")
include("objective.jl")
include("results.jl")

MOI.get(::Optimizer, ::MOI.SolverName) = "ConstraintSolver"

"""
Optimizer struct constructor
"""
function Optimizer(;options...)
    com = CS.init()
    options = combine_options(options)
    return Optimizer(
        com,
        [],
        [],
        MOI.OPTIMIZE_NOT_CALLED,
        options
    )
end

"""
    MOI.is_empty(model::Optimizer)
"""
function MOI.is_empty(model::Optimizer)
    return isempty(model.variable_info) &&
           isempty(model.var_constraints)
end

"""
    MOI.empty!(model::Optimizer)
"""
function MOI.empty!(model::Optimizer)
    model.inner = CS.init()
    empty!(model.variable_info)
    empty!(model.var_constraints)
    model.status = MOI.OPTIMIZE_NOT_CALLED
    # !important => don't remove the options
end

"""
Copy constructor for the optimizer
"""
MOIU.supports_default_copy_to(model::Optimizer, copy_names::Bool) = !copy_names
function MOI.copy_to(model::Optimizer, src::MOI.ModelLike; kws...)
    return MOI.Utilities.automatic_copy_to(model, src; kws...)
end

"""
    MOI.optimize!(model::Optimizer)
"""
function MOI.optimize!(model::Optimizer)
    # check if every variable has bounds and is an Integer
    check_var_bounds(model)

    set_pvals!(model)
    set_constraint_hashes!(model)

    status = solve!(model)
    set_status!(model, status)
end


module Resampling

import JuliaRL
import Base.get

export
    AbstractExperienceReplay,
    ExperienceReplay,
    WeightedExperienceReplay,
    add!,
    size,
    getindex,
    sample

include("Replay.jl")

export
    GVF,
    FunctionalCumulant,
    StateTerminationDiscount,
    ConstantDiscount,
    PersistentPolicy,
    RandomPolicy,
    UniformRandomPolicy,
    InterpolationPolicy,
    FavoredRandomPolicy

include("GVF.jl")

export Tabular, SparseLayer
include("CustomLayers.jl")

export
    BatchTD,
    IncTD,
    WISBatchTD,
    VTrace,
    IncNormIS,
    WSNormIS,
    BatchSarsa,
    BatchExpectedSarsa,
    update!
include("Learning.jl")

include("Environments.jl")

export MarkovChainUtil
include("exp_util.jl")

# Defunc learning algorithms ported from old codebase for the variance experiments:

include("TabularRL.jl")

end # module


export convertCsvToBasketsExample, dppLearningExample,
  dppLearningBayesianExample, predictionExample,
  predictionForMCMCSamplesExample, predictionMetricsExample

dataDir = "../data/Amazon-baby-registry"

# An example of converting a file in sparse CSV basket format to the data
# format expected by the DPP modules.
function convertCsvToBasketsExample()
  csvBasketDataFileName = "$dataDir/1_100_100_100_apparel_regs.csv"
  trainingBasketsDictFileName = "$dataDir/apparel-regs-training-basketsDict.jls"
  testBasketsDictFileName = "$dataDir/apparel-regs-test-basketsDict.jls"
  trainingSetSizePercent = 0.8

  convertSparseCsvToBaskets(csvBasketDataFileName, trainingBasketsDictFileName,
    testBasketsDictFileName, trainingSetSizePercent)
end

# An example of using stochastic gradient ascent for DPP learning
function dppLearningExample()
  trainingBasketsDictFileName = "$dataDir/apparel-regs-training-basketsDict.jls"
  trainingBasketsDictObjectName = "trainingBasketsDict"
  testBasketsDictFileName = "$dataDir/apparel-regs-test-basketsDict.jls"
  testBasketsDictObjectName = "testBasketsDict"
  learnedModelOutputDirName = "example-SGA-model"
  numItemTraits = 30
  alpha = 1.0

  doDPPLearningSparseVectorData(trainingBasketsDictFileName,
    trainingBasketsDictObjectName, testBasketsDictFileName,
    testBasketsDictObjectName, learnedModelOutputDirName, numItemTraits, alpha)

end

# An example of using stochastic gradient HMC for DPP learning
function dppLearningBayesianExample()
  trainingBasketsDictFileName = "$dataDir/apparel-regs-training-basketsDict.jls"
  trainingBasketsDictObjectName = "trainingBasketsDict"
  testBasketsDictFileName = "$dataDir/apparel-regs-test-basketsDict.jls"
  testBasketsDictObjectName = "testBasketsDict"
  learnedModelOutputDirName = "example-SGHMC-model"
  numItemTraits = 30

  doDPPBayesianLearningSparseVectorData(trainingBasketsDictFileName,
    trainingBasketsDictObjectName, testBasketsDictFileName,
    testBasketsDictObjectName, learnedModelOutputDirName, numItemTraits,
    runStochasticGradientHamiltonianMonteCarloSampler)
end

# An example of computing predictions using a DPP model learned by stochastic
# gradient ascent
function predictionExample()
  testBasketsDictFileName = "$dataDir/apparel-regs-test-basketsDict.jls"
  testBasketsDictObjectName = "testBasketsDict"
  numItemTraits = 30
  alpha = 1.0
  learnedDPPParamsFileName = "example-SGA-model/learnedDPPParamsMatrix-k$numItemTraits-lambdaPop$alpha.jls"
  resultsForTestInstancesDictFileName =
    "example-SGA-model/resultsForTestInstancesDict-k$numItemTraits-lambdaPop$alpha-Amazon-apparel-regs.jls"

  computePredictionsSparseVectorData(testBasketsDictFileName,
    testBasketsDictObjectName, learnedDPPParamsFileName,
    resultsForTestInstancesDictFileName)
end

# An example of computing predictions using a DPP model learned by stochastic
# gradient HMC
function predictionForMCMCSamplesExample()
  testBasketsDictFileName = "$dataDir/apparel-regs-test-basketsDict.jld"
  testBasketsDictObjectName = "testBasketsDict"
  numItemTraits = 30
  sampleLag = 1
  collectedMCMCSamplesDirPathName =
    "example-SGHMC-model/learnedDPPParams-k$numItemTraits-SGHMC-collectedSamples"
  resultsForTestInstancesDictFileName =
    "example-SGHMC-model/resultsForTestInstancesDict-k$numItemTraits-collectedSamples-SGHMC-Amazon-apparel-regs.jld"

  computePredictionsForMCMCSamples(testBasketsDictFileName,
    testBasketsDictObjectName, collectedMCMCSamplesDirPathName, sampleLag,
    resultsForTestInstancesDictFileName)
end

# An example of computing prediction metrics using DPP model prediction results
function predictionMetricsExample()
  trainingBasketsDictFileName = "$dataDir/apparel-regs-training-basketsDict.jls"
  trainingBasketsDictObjectName = "trainingBasketsDict"
  testBasketsDictFileName = "$dataDir/apparel-regs-test-basketsDict.jls"
  testBasketsDictObjectName = "testBasketsDict"
  numItemTraits = 30
  alpha = 1.0
  resultsForTestInstancesDictFileName =
    "example-SGA-model/resultsForTestInstancesDict-k$numItemTraits-lambdaPop$alpha-Amazon-apparel-regs.jls"
  # resultsForTestInstancesDictFileName =
  #   "example-SGHMC-model/resultsForTestInstancesDict-k$numItemTraits-collectedSamples-SGHMC-Amazon-apparel-regs.jls"

  computePredictionMetricsSparseVectorData(trainingBasketsDictFileName,
    trainingBasketsDictObjectName, testBasketsDictFileName,
    testBasketsDictObjectName, resultsForTestInstancesDictFileName)
end

# Optimization-based learning pipeline:
# convertCsvToBasketsExample()
# dppLearningExample()
# predictionExample()
# predictionMetricsExample()

# Bayesian learning pipeline:
# dppLearningBayesianExample()
# predictionForMCMCSamplesExample()
# predictionMetricsExample()


"""
    static_analysis(assembly; kwargs...)

Perform a static analysis of the system of nonlinear beams contained in
`assembly`. Return the resulting system and a flag indicating whether the
iteration procedure converged.

# Keyword Arguments
- `prescribed_conditions = Dict{Int,PrescribedConditions{Float64}}()`:
       A dictionary with keys corresponding to the points at
       which prescribed conditions are applied and elements of type
       [`PrescribedConditions`](@ref) which describe the prescribed conditions
       at those points.  If time varying, this input may be provided as a
       function of time.
- `distributed_loads = Dict{Int,DistributedLoads{Float64}}()`: A dictionary
       with keys corresponding to the elements to which distributed loads are
       applied and elements of type [`DistributedLoads`](@ref) which describe
       the distributed loads at those points.  If time varying, this input may
       be provided as a function of time.
 - `linear = false`: Set to `true` for a linear analysis
 - `linearization_state`: Linearization state variables.  Defaults to zeros.
 - `update_linearization_state`: Flag indicating whether to update the linearization state
    variables for a linear analysis with the instantaneous state variables.
 - `method = :newton`: Method (as defined in NLsolve) to solve nonlinear system of equations
 - `linesearch = LineSearches.BackTracking(maxstep=1e6)`: Line search used to solve nonlinear system of equations
 - `ftol = 1e-9`: tolerance for solving nonlinear system of equations
 - `iterations = 1000`: maximum iterations for solving the nonlinear system of equations
 - `tvec = 0`: Time vector/value. May be used in conjunction with time varying
     prescribed conditions and distributed loads to gradually increase
    displacements/loads.
 - `reset_state = true`: Flag indicating whether the state variables should be
    reset prior to performing the analysis.  This keyword argument is only valid
    for the pre-allocated version of this function.
"""
function static_analysis(assembly;
    prescribed_conditions=Dict{Int,PrescribedConditions{Float64}}(),
    distributed_loads=Dict{Int,DistributedLoads{Float64}}(),
    linear=false,
    linearization_state=nothing,
    update_linearization_state=false,
    method=:newton,
    linesearch=LineSearches.BackTracking(maxstep=1e6),
    ftol=1e-9,
    iterations=1000,
    tvec=0.0,
    )

    static = true

    pc = typeof(prescribed_conditions) <: AbstractDict ? prescribed_conditions : prescribed_conditions(tvec[1])

    system = System(assembly, static; prescribed_points=keys(pc))

    return static_analysis!(system, assembly;
        prescribed_conditions=prescribed_conditions,
        distributed_loads=distributed_loads,
        linear=linear,
        linearization_state=linearization_state,
        update_linearization_state=update_linearization_state,
        method=method,
        linesearch=linesearch,
        ftol=ftol,
        iterations=iterations,
        tvec=tvec,
        reset_state=false)
end

"""
    static_analysis!(system, assembly; kwargs...)

Pre-allocated version of `static_analysis`.
"""
function static_analysis!(system, assembly;
    prescribed_conditions=Dict{Int,PrescribedConditions{Float64}}(),
    distributed_loads=Dict{Int,DistributedLoads{Float64}}(),
    linear=false,
    linearization_state=nothing,
    update_linearization_state=false,
    method=:newton,
    linesearch=LineSearches.BackTracking(maxstep=1e6),
    ftol=1e-9,
    iterations=1000,
    tvec=0.0,
    reset_state=true)

    # check to make sure system is static
    @assert system.static == true

    # reset state, if specified
    if reset_state
        reset_state!(system)
    end

    # unpack pre-allocated storage and pointers
    x = system.x
    F = system.r
    J = system.K
    force_scaling = system.force_scaling
    mass_scaling = system.mass_scaling
    irow_point = system.irow_point
    irow_elem = system.irow_elem
    irow_elem1 = system.irow_elem1
    irow_elem2 = system.irow_elem2
    icol_point = system.icol_point
    icol_elem = system.icol_elem

    converged = true
    for t in tvec

        # update stored time
        system.t = t

        # current parameters
        pcond = typeof(prescribed_conditions) <: AbstractDict ?
            prescribed_conditions : prescribed_conditions(t)
        dload = typeof(distributed_loads) <: AbstractDict ?
            distributed_loads : distributed_loads(t)

        # solve the system of equations
        f! = (F, x) -> system_residual!(F, x, assembly, pcond, dload, force_scaling,
            mass_scaling, irow_point, irow_elem, irow_elem1, irow_elem2, icol_point, icol_elem)

        j! = (J, x) -> system_jacobian!(J, x, assembly, pcond, dload, force_scaling,
            mass_scaling, irow_point, irow_elem, irow_elem1, irow_elem2, icol_point, icol_elem)

        if linear
            # linear analysis
            if !update_linearization_state
                if isnothing(linearization_state)
                    x .= 0
                else
                    x .= linearization_state
                end
            end
            f!(F, x)
            j!(J, x)
            x .-= safe_lu(J) \ F
        else
            # nonlinear analysis
            df = NLsolve.OnceDifferentiable(f!, j!, x, F, J)

            result = NLsolve.nlsolve(df, x,
                linsolve=(x, A, b) -> ldiv!(x, safe_lu(A), b),
                method=method,
                linesearch=linesearch,
                ftol=ftol,
                iterations=iterations)

            # update the solution
            x .= result.zero
            J .= df.DF

            # update convergence flag
            converged = result.f_converged
        end
    end

    return system, converged
end

"""
    steady_state_analysis(assembly; kwargs...)

Perform a steady-state analysis for the system of nonlinear beams contained in
`assembly`.  Return the resulting system and a flag indicating whether the
iteration procedure converged.

# Keyword Arguments
 - `prescribed_conditions = Dict{Int,PrescribedConditions{Float64}}()`:
        A dictionary with keys corresponding to the points at
        which prescribed conditions are applied and elements of type
        [`PrescribedConditions`](@ref) which describe the prescribed conditions
        at those points.  If time varying, this input may be provided as a
        function of time.
 - `distributed_loads = Dict{Int,DistributedLoads{Float64}}()`: A dictionary
        with keys corresponding to the elements to which distributed loads are
        applied and elements of type [`DistributedLoads`](@ref) which describe
        the distributed loads at those points.  If time varying, this input may
        be provided as a function of time.
 - `linear = false`: Set to `true` for a linear analysis
 - `linearization_state`: Linearization state variables.  Defaults to zeros.
 - `update_linearization_state`: Flag indicating whether to update the linearization state
    variables for a linear analysis with the current state variables.
 - `method = :newton`: Method (as defined in NLsolve) to solve nonlinear system of equations
 - `linesearch = LineSearches.LineSearches.BackTracking(maxstep=1e6)`: Line search used to solve nonlinear system of equations
 - `ftol = 1e-9`: tolerance for solving nonlinear system of equations
 - `iterations = 1000`: maximum iterations for solving the nonlinear system of equations
 - `origin = zeros(3)`: Global frame origin vector. If time varying, this input
        may be provided as a function of time.
 - `linear_velocity = zeros(3)`: Global frame linear velocity vector. If time
        varying, this vector may be provided as a function of time.
 - `angular_velocity = zeros(3)`: Global frame angular velocity vector. If time
        varying, this vector may be provided as a function of time.
 - `tvec = 0.0`: Time vector/value. May be used in conjunction with time varying
    prescribed conditions, distributed loads, and global motion to gradually
    increase displacements/loads.
 - `reset_state = true`: Flag indicating whether the state variables should be
    reset prior to performing the analysis.  This keyword argument is only valid
    for the pre-allocated version of this function.
"""
function steady_state_analysis(assembly;
    prescribed_conditions=Dict{Int,PrescribedConditions{Float64}}(),
    distributed_loads=Dict{Int,DistributedLoads{Float64}}(),
    linear=false,
    linearization_state=nothing,
    update_linearization_state=false,
    method=:newton,
    linesearch=LineSearches.BackTracking(maxstep=1e6),
    ftol=1e-9,
    iterations=1000,
    origin=(@SVector zeros(3)),
    linear_velocity=(@SVector zeros(3)),
    angular_velocity=(@SVector zeros(3)),
    tvec=0.0,
    )

    static = false

    pc = typeof(prescribed_conditions) <: AbstractDict ? prescribed_conditions : prescribed_conditions(tvec[1])

    system = System(assembly, static; prescribed_points=keys(pc))

    return steady_state_analysis!(system, assembly;
        prescribed_conditions=prescribed_conditions,
        distributed_loads=distributed_loads,
        linear=linear,
        linearization_state=linearization_state,
        update_linearization_state=update_linearization_state,
        method=method,
        linesearch=linesearch,
        ftol=ftol,
        iterations=iterations,
        origin=origin,
        linear_velocity=linear_velocity,
        angular_velocity=angular_velocity,
        tvec=tvec,
        reset_state=false,
        )
end

"""
    steady_state_analysis!(system, assembly; kwargs...)

Pre-allocated version of `steady_state_analysis`.
"""
function steady_state_analysis!(system, assembly;
    prescribed_conditions=Dict{Int,PrescribedConditions{Float64}}(),
    distributed_loads=Dict{Int,DistributedLoads{Float64}}(),
    linear=false,
    linearization_state=nothing,
    update_linearization_state=false,
    method=:newton,
    linesearch=LineSearches.BackTracking(maxstep=1e6),
    ftol=1e-9,
    iterations=1000,
    origin=(@SVector zeros(3)),
    linear_velocity=(@SVector zeros(3)),
    angular_velocity=(@SVector zeros(3)),
    tvec=0.0,
    reset_state=true,
    )

    # check to make sure the simulation is dynamic
    @assert system.static == false

    # reset state, if specified
    if reset_state
        reset_state!(system)
    end

    # unpack pointers to pre-allocated storage
    x = system.x
    F = system.r
    J = system.K
    force_scaling = system.force_scaling
    mass_scaling = system.mass_scaling
    irow_point = system.irow_point
    irow_elem = system.irow_elem
    irow_elem1 = system.irow_elem1
    irow_elem2 = system.irow_elem2
    icol_point = system.icol_point
    icol_elem = system.icol_elem

    converged = true
    for t in tvec

        # update stored time
        system.t = t

        # current parameters
        pcond = typeof(prescribed_conditions) <: AbstractDict ?
            prescribed_conditions : prescribed_conditions(t)
        dload = typeof(distributed_loads) <: AbstractDict ?
            distributed_loads : distributed_loads(t)
        x0 = typeof(origin) <: AbstractVector ? SVector{3}(origin) : SVector{3}(origin(t))
        v0 = typeof(linear_velocity) <: AbstractVector ? SVector{3}(linear_velocity) : SVector{3}(linear_velocity(t))
        ω0 = typeof(angular_velocity) <: AbstractVector ? SVector{3}(angular_velocity) : SVector{3}(angular_velocity(t))

        f! = (F, x) -> system_residual!(F, x, assembly, pcond,
            dload, force_scaling, mass_scaling, irow_point, irow_elem, irow_elem1,
            irow_elem2, icol_point, icol_elem, x0, v0, ω0)

        j! = (J, x) -> system_jacobian!(J, x, assembly, pcond,
            dload, force_scaling, mass_scaling, irow_point, irow_elem, irow_elem1,
            irow_elem2, icol_point, icol_elem, x0, v0, ω0)

        # solve the system of equations
        if linear
            # linear analysis
            if !update_linearization_state
                if isnothing(linearization_state)
                    x .= 0
                else
                    x .= linearization_state
                end
            end
            f!(F, x)
            j!(J, x)
            x .-= safe_lu(J) \ F
        else
            # nonlinear analysis
            df = NLsolve.OnceDifferentiable(f!, j!, x, F, J)

            result = NLsolve.nlsolve(df, x,
                linsolve=(x, A, b) -> ldiv!(x, safe_lu(A), b),
                method=method,
                linesearch=linesearch,
                ftol=ftol,
                iterations=iterations)

            # update the solution
            x .= result.zero
            J .= df.DF

            # update the convergence flag
            convergence = result.f_converged
        end
    end

    return system, converged
end

"""
    eigenvalue_analysis(assembly; kwargs...)

Compute the eigenvalues and eigenvectors of the system of nonlinear beams
contained in `assembly`.  Return the modified system, eigenvalues, eigenvectors,
and a convergence flag indicating whether the corresponding steady-state analysis
converged.

# Keyword Arguments
 - `prescribed_conditions = Dict{Int,PrescribedConditions{Float64}}()`:
        A dictionary with keys corresponding to the points at
        which prescribed conditions are applied and elements of type
        [`PrescribedConditions`](@ref) which describe the prescribed conditions
        at those points.  If time varying, this input may be provided as a
        function of time.
 - `distributed_loads = Dict{Int,DistributedLoads{Float64}}()`: A dictionary
        with keys corresponding to the elements to which distributed loads are
        applied and elements of type [`DistributedLoads`](@ref) which describe
        the distributed loads at those points.  If time varying, this input may
        be provided as a function of time.
 - `linear = false`: Set to `true` for a linear analysis
 - `linearization_state`: Linearization state variables.  Defaults to zeros.
 - `update_linearization_state`: Flag indicating whether to update the linearization state
    variables for a linear analysis with the current state variables.
 - `method = :newton`: Method (as defined in NLsolve) to solve nonlinear system of equations
 - `linesearch = LineSearches.LineSearches.BackTracking(maxstep=1e6)`: Line search used to solve nonlinear system of equations
 - `ftol = 1e-9`: tolerance for solving nonlinear system of equations
 - `iterations = 1000`: maximum iterations for solving the nonlinear system of equations
 - `reset_state = true`: Flag indicating whether the state variables should be
    reset prior to performing the steady-state analysis.  This keyword argument
    is only valid for the pre-allocated version of this function.
 - `find_steady_state = reset_state && !linear`: Flag indicating whether the
    steady state solution should be found prior to performing the eigenvalue analysis.
 - `origin = zeros(3)`: Global frame origin.
    If time varying, this vector may be provided as a function of time.
 - `linear_velocity = zeros(3)`: Global frame linear velocity vector.
    If time varying, this vector may be provided as a function of time.
    May be provided either as a constant or as a function of time.
 - `angular_velocity = zeros(3)`: Global frame angular velocity vector.
    If time varying, this vector may be provided as a function of time.
 - `tvec`: Time vector. May be used in conjunction with time varying
    prescribed conditions, distributed loads, and global motion to gradually
    increase displacements/loads during the steady-state analysis.
 - `nev = 6`: Number of eigenvalues to compute
"""
function eigenvalue_analysis(assembly;
    prescribed_conditions=Dict{Int,PrescribedConditions{Float64}}(),
    distributed_loads=Dict{Int,DistributedLoads{Float64}}(),
    method=:newton,
    linear=false,
    linearization_state=nothing,
    update_linearization_state=false,
    linesearch=LineSearches.BackTracking(maxstep=1e6),
    ftol=1e-9,
    iterations=1000,
    find_steady_state=!linear,
    origin=(@SVector zeros(3)),
    linear_velocity=(@SVector zeros(3)),
    angular_velocity=(@SVector zeros(3)),
    tvec=0.0,
    nev=6
    )

    static = false

    pc = typeof(prescribed_conditions) <: AbstractDict ? prescribed_conditions : prescribed_conditions(tvec[1])

    system = System(assembly, static; prescribed_points=keys(pc))

    return eigenvalue_analysis!(system, assembly;
        prescribed_conditions=prescribed_conditions,
        distributed_loads=distributed_loads,
        linear=linear,
        linearization_state=linearization_state,
        update_linearization_state=update_linearization_state,
        method=method,
        linesearch=linesearch,
        ftol=ftol,
        iterations=iterations,
        reset_state=false,
        find_steady_state=find_steady_state,
        origin=origin,
        linear_velocity=linear_velocity,
        angular_velocity=angular_velocity,
        tvec=tvec,
        nev=nev,
        )
end

"""
    eigenvalue_analysis!(system, assembly; kwargs...)

Pre-allocated version of `eigenvalue_analysis`.  Uses the state variables stored in
`system` as an initial guess for iterating to find the steady state solution.
"""
function eigenvalue_analysis!(system, assembly;
    prescribed_conditions=Dict{Int,PrescribedConditions{Float64}}(),
    distributed_loads=Dict{Int,DistributedLoads{Float64}}(),
    linear=false,
    linearization_state=nothing,
    update_linearization_state=false,
    method=:newton,
    linesearch=LineSearches.BackTracking(maxstep=1e6),
    ftol=1e-9,
    iterations=1000,
    reset_state=true,
    find_steady_state=!linear && reset_state,
    origin=(@SVector zeros(3)),
    linear_velocity=(@SVector zeros(3)),
    angular_velocity=(@SVector zeros(3)),
    tvec=0.0,
    nev=6,
    )

    if reset_state
        reset_state!(system)
    end

    # perform steady state analysis (if nonlinear)
    if find_steady_state
        system, converged = steady_state_analysis!(system, assembly;
            prescribed_conditions=prescribed_conditions,
            distributed_loads=distributed_loads,
            linear=linear,
            linearization_state=linearization_state,
            update_linearization_state=update_linearization_state,
            method=method,
            linesearch=linesearch,
            ftol=ftol,
            iterations=iterations,
            origin=origin,
            linear_velocity=linear_velocity,
            angular_velocity=angular_velocity,
            tvec=tvec,
            )
    else
        # set linearization state variables
        if linear && !update_linearization_state
            if isnothing(linearization_state)
                system.x .= 0
            else
                system.x .= linearization_state
            end
        end
        # converged by default
        converged = true
    end

    # unpack state vector, stiffness, and mass matrices
    x = system.x # populated during steady state solution
    K = system.K # needs to be updated
    M = system.M # still needs to be populated

    # unpack scaling parameter
    force_scaling = system.force_scaling
    mass_scaling = system.mass_scaling

    # also unpack system indices
    irow_point = system.irow_point
    irow_elem = system.irow_elem
    irow_elem1 = system.irow_elem1
    irow_elem2 = system.irow_elem2
    icol_point = system.icol_point
    icol_elem = system.icol_elem

    # current time
    t = system.t

    # current parameters
    pcond = typeof(prescribed_conditions) <: AbstractDict ?
        prescribed_conditions : prescribed_conditions(t)
    dload = typeof(distributed_loads) <: AbstractDict ?
        distributed_loads : distributed_loads(t)
    x0 = typeof(origin) <: AbstractVector ? SVector{3}(origin) : SVector{3}(origin(t))
    v0 = typeof(linear_velocity) <: AbstractVector ? SVector{3}(linear_velocity) : SVector{3}(linear_velocity(t))
    ω0 = typeof(angular_velocity) <: AbstractVector ? SVector{3}(angular_velocity) : SVector{3}(angular_velocity(t))

    # solve for the system stiffness matrix
    K = system_jacobian!(K, x, assembly, pcond, dload, force_scaling, mass_scaling,
        irow_point, irow_elem, irow_elem1, irow_elem2, icol_point, icol_elem, x0, v0, ω0)

    # solve for the system mass matrix
    M = system_mass_matrix!(M, x, assembly, force_scaling, mass_scaling, irow_point, irow_elem,
        irow_elem1, irow_elem2, icol_point, icol_elem)

    # construct linear map
    T = eltype(system)
    nx = length(x)
    Kfact = safe_lu(K)
    f! = (b, x) -> ldiv!(b, Kfact, M * x)
    fc! = (b, x) -> mul!(b, M', Kfact' \ x)
    A = LinearMap{T}(f!, fc!, nx, nx; ismutating=true)

    # compute eigenvalues and eigenvectors
    λ, V = partialeigen(partialschur(A; nev=min(nx, nev), which=LM())[1])

    # sort eigenvalues by magnitude
    perm = sortperm(λ, by=(λ) -> (abs(λ), imag(λ)), rev=true)
    λ .= λ[perm]
    V .= V[:,perm]

    # eigenvalues are actually -1/λ, no modification necessary for eigenvectors
    λ .= -1 ./ λ

    return system, λ, V, converged
end

"""
    initial_condition_analysis(assembly, t0; kwargs...)

Perform an analysis to obtain a consistent set of initial conditions.  Return the
final system with the new initial conditions.

# Keyword Arguments
 - `prescribed_conditions: A dictionary with keys corresponding to the points at
        which prescribed conditions are applied and elements of type
        [`PrescribedConditions`](@ref) which describe the prescribed conditions
        at those points.  If time varying, this input may be provided as a
        function of time.
 - `distributed_loads: A dictionary with keys corresponding to the elements to
        which distributed loads are applied and elements of type
        [`DistributedLoads`](@ref) which describe the distributed loads at those
        points.  If time varying, this input may be provided as a function of
        time.
 - `linear = false`: Set to `true` for a linear analysis
 - `linearization_state`: Linearization state variables.  Defaults to zeros.
 - `method = :newton`: Method (as defined in NLsolve) to solve nonlinear system of equations
 - `linesearch = LineSearches.LineSearches.BackTracking(maxstep=1e6)`: Line search used to solve nonlinear system of equations
 - `ftol = 1e-9`: tolerance for solving nonlinear system of equations
 - `iterations = 1000`: maximum iterations for solving the nonlinear system of equations
 - `reset_state = true`: Flag indicating whether the state variables should be
    reset prior to performing the analysis.  This keyword argument is only valid
    for the pre-allocated version of this function.
 - `origin = zeros(3)`: Global frame origin.
    If time varying, this vector may be provided as a function of time.
 - `linear_velocity = zeros(3)`: Global frame linear velocity vector.
    If time varying, this vector may be provided as a function of time.
    May be provided either as a constant or as a function of time.
 - `angular_velocity = zeros(3)`: Global frame angular velocity vector.
    If time varying, this vector may be provided as a function of time.
 - `u0=fill(zeros(3), length(assembly.elements))`: Initial displacment of each beam element,
 - `theta0=fill(zeros(3), length(assembly.elements))`: Initial angular displacement of each beam element,
 - `udot0=fill(zeros(3), length(assembly.elements))`: Initial time derivative with respect to `u`
 - `thetadot0=fill(zeros(3), length(assembly.elements))`: Initial time derivative with respect to `theta`
 - `save=1:length(tvec)`: Steps at which to save the time history
"""
function initial_condition_analysis(assembly, t0;
    prescribed_conditions=Dict{Int,PrescribedConditions{Float64}}(),
    distributed_loads=Dict{Int,DistributedLoads{Float64}}(),
    linear=false,
    linearization_state=nothing,
    method=:newton,
    linesearch=LineSearches.BackTracking(maxstep=1e6),
    ftol=1e-9,
    iterations=1000,
    origin=(@SVector zeros(3)),
    linear_velocity=(@SVector zeros(3)),
    angular_velocity=(@SVector zeros(3)),
    u0=fill((@SVector zeros(3)), length(assembly.elements)),
    theta0=fill((@SVector zeros(3)), length(assembly.elements)),
    udot0=fill((@SVector zeros(3)), length(assembly.elements)),
    thetadot0=fill((@SVector zeros(3)), length(assembly.elements)),
    )

    static = false

    pc = typeof(prescribed_conditions) <: AbstractDict ? prescribed_conditions : prescribed_conditions(t0)

    system = System(assembly, static; prescribed_points=keys(pc))

    return initial_condition_analysis!(system, assembly, t0;
        prescribed_conditions=prescribed_conditions,
        distributed_loads=distributed_loads,
        linear=linear,
        linearization_state=linearization_state,
        method=method,
        linesearch=linesearch,
        ftol=ftol,
        iterations=iterations,
        reset_state=false,
        origin=origin,
        linear_velocity=linear_velocity,
        angular_velocity=angular_velocity,
        u0=u0,
        theta0=theta0,
        udot0=udot0,
        thetadot0=thetadot0,
        )
end

"""
    initial_condition_analysis!(system, assembly, t0; kwargs...)

Pre-allocated version of `initial_condition_analysis`.
"""
function initial_condition_analysis!(system, assembly, t0;
    prescribed_conditions=Dict{Int,PrescribedConditions{Float64}}(),
    distributed_loads=Dict{Int,DistributedLoads{Float64}}(),
    linear=false,
    linearization_state=nothing,
    method=:newton,
    linesearch=LineSearches.BackTracking(maxstep=1e6),
    ftol=1e-9,
    iterations=1000,
    reset_state=true,
    origin=(@SVector zeros(3)),
    linear_velocity=(@SVector zeros(3)),
    angular_velocity=(@SVector zeros(3)),
    u0=fill((@SVector zeros(3)), length(assembly.elements)),
    theta0=fill((@SVector zeros(3)), length(assembly.elements)),
    udot0=fill((@SVector zeros(3)), length(assembly.elements)),
    thetadot0=fill((@SVector zeros(3)), length(assembly.elements)),
    )

    # check to make sure the simulation is dynamic
    @assert system.static == false

    if reset_state
        reset_state!(system)
    end

    # unpack pre-allocated storage and pointers for system
    x = system.x
    F = system.r
    J = system.K
    force_scaling = system.force_scaling
    mass_scaling = system.mass_scaling
    irow_point = system.irow_point
    irow_elem = system.irow_elem
    irow_elem1 = system.irow_elem1
    irow_elem2 = system.irow_elem2
    icol_point = system.icol_point
    icol_elem = system.icol_elem
    udot = system.udot
    θdot = system.θdot
    Pdot = system.Pdot
    Hdot = system.Hdot

    nelem = length(assembly.elements)

    # set current time step
    system.t = t0

    # set current parameters
    pcond = typeof(prescribed_conditions) <: AbstractDict ?
        prescribed_conditions : prescribed_conditions(t0)
    dload = typeof(distributed_loads) <: AbstractDict ?
        distributed_loads : distributed_loads(t0)
    x0 = typeof(origin) <: AbstractVector ? SVector{3}(origin) : SVector{3}(origin(t0))
    v0 = typeof(linear_velocity) <: AbstractVector ? SVector{3}(linear_velocity) : SVector{3}(linear_velocity(t0))
    ω0 = typeof(angular_velocity) <: AbstractVector ? SVector{3}(angular_velocity) : SVector{3}(angular_velocity(t0))

    # construct residual and jacobian functions
    f! = (F, x) -> system_residual!(F, x, assembly, pcond, dload, force_scaling,
        mass_scaling, irow_point, irow_elem, irow_elem1, irow_elem2, icol_point, icol_elem,
        x0, v0, ω0, u0, theta0, udot0, thetadot0)

    j! = (J, x) -> system_jacobian!(J, x, assembly, pcond, dload, force_scaling,
        mass_scaling, irow_point, irow_elem, irow_elem1, irow_elem2, icol_point, icol_elem,
        x0, v0, ω0, u0, theta0, udot0, thetadot0)

    # solve system of equations
    if linear
        # linear analysis
        if isnothing(linearization_state)
            x .= 0
        else
            x .= linearization_state
        end
        f!(F, x)
        j!(J, x)
        x .-= safe_lu(J) \ F
    else
        # nonlinear analysis
        df = OnceDifferentiable(f!, j!, x, F, J)

        result = NLsolve.nlsolve(df, x,
            linsolve=(x, A, b) -> ldiv!(x, safe_lu(A), b),
            method=method,
            linesearch=linesearch,
            ftol=ftol,
            iterations=iterations)

        x .= result.zero
        J .= df.DF
    end

    # get convergence flag
    converged = result.f_converged

    # save states and state rates
    for ielem = 1:nelem
        icol = icol_elem[ielem]
        # extract rotation parameters for this beam element
        C = get_C(theta0[ielem])
        Cab = assembly.elements[ielem].Cab
        CtCab = C' * Cab
        # save states and state rates
        udot[ielem] = udot0[ielem]
        θdot[ielem] = thetadot0[ielem]
        Pdot[ielem] = CtCab' * SVector(x[icol], x[icol + 1], x[icol + 2]) .* mass_scaling
        Hdot[ielem] = CtCab' * SVector(x[icol + 3], x[icol + 4], x[icol + 5]) .* mass_scaling
# restore original state vector
        x[icol:icol + 2] .= u0[ielem]
        x[icol + 3:icol + 5] .= theta0[ielem]
    end

    return system, converged
end

"""
    time_domain_analysis(assembly, tvec; kwargs...)

Perform a time-domain analysis for the system of nonlinear beams contained in
`assembly` using the time vector `tvec`.  Return the final system, a post-processed
solution history, and a convergence flag indicating whether the iterations
converged for each time step.

# Keyword Arguments
 - `prescribed_conditions: A dictionary with keys corresponding to the points at
        which prescribed conditions are applied and elements of type
        [`PrescribedConditions`](@ref) which describe the prescribed conditions
        at those points.  If time varying, this input may be provided as a
        function of time.
 - `distributed_loads: A dictionary with keys corresponding to the elements to
        which distributed loads are applied and elements of type
        [`DistributedLoads`](@ref) which describe the distributed loads at those
        points.  If time varying, this input may be provided as a function of
        time.
 - `linear = false`: Set to `true` for a linear analysis
 - `linearization_state`: Linearization state variables.  Defaults to zeros.
 - `update_linearization_state`: Flag indicating whether to update the linearization state
    variables for a linear analysis with the current state variables.
 - `method = :newton`: Method (as defined in NLsolve) to solve nonlinear system of equations
 - `linesearch = LineSearches.LineSearches.BackTracking(maxstep=1e6)`: Line search used to solve nonlinear system of equations
 - `ftol = 1e-9`: tolerance for solving nonlinear system of equations
 - `iterations = 1000`: maximum iterations for solving the nonlinear system of equations
 - `reset_state = true`: Flag indicating whether the state variables should be
    reset prior to performing the analysis.  This keyword argument is only valid
    for the pre-allocated version of this function.
 - `initialize = true`: Flag indicating whether a consistent set of initial
    conditions should be found using [`initial_condition_analysis`](@ref). If
    `false`, the keyword arguments `u0`, `theta0`, `udot0` and `thetadot0` will
    be ignored and the system state vector will be used as the initial state
    variables.
 - `origin`: Global frame origin vector. If time varying, this input
        may be provided as a function of time.
 - `linear_velocity`: Global frame linear velocity vector. If time
        varying, this vector may be provided as a function of time.
 - `angular_velocity`: Global frame angular velocity vector. If time
        varying, this vector may be provided as a function of time.
 - `u0=fill(zeros(3), length(assembly.elements))`: Initial displacment of each beam element,
 - `theta0=fill(zeros(3), length(assembly.elements))`: Initial angular displacement of each beam element,
 - `udot0=fill(zeros(3), length(assembly.elements))`: Initial time derivative with respect to `u`
 - `thetadot0=fill(zeros(3), length(assembly.elements))`: Initial time derivative with respect to `theta`
 - `save=1:length(tvec)`: Steps at which to save the time history
"""
function time_domain_analysis(assembly, tvec;
    prescribed_conditions=Dict{Int,PrescribedConditions{Float64}}(),
    distributed_loads=Dict{Int,DistributedLoads{Float64}}(),
    linear=false,
    linearization_state=nothing,
    update_linearization_state=false,
    method=:newton,
    linesearch=LineSearches.BackTracking(maxstep=1e6),
    ftol=1e-9,
    iterations=1000,
    initialize=true,
    origin=(@SVector zeros(3)),
    linear_velocity=(@SVector zeros(3)),
    angular_velocity=(@SVector zeros(3)),
    u0=fill((@SVector zeros(3)), length(assembly.elements)),
    theta0=fill((@SVector zeros(3)), length(assembly.elements)),
    udot0=fill((@SVector zeros(3)), length(assembly.elements)),
    thetadot0=fill((@SVector zeros(3)), length(assembly.elements)),
    save=1:length(tvec)
    )

    static = false

    pc = typeof(prescribed_conditions) <: AbstractDict ? prescribed_conditions : prescribed_conditions(tvec[1])

    system = System(assembly, static; prescribed_points=keys(pc))

    return time_domain_analysis!(system, assembly, tvec;
        prescribed_conditions=prescribed_conditions,
        distributed_loads=distributed_loads,
        linear=linear,
        linearization_state=linearization_state,
        update_linearization_state=update_linearization_state,
        method=method,
        linesearch=linesearch,
        ftol=ftol,
        iterations=iterations,
        reset_state=false,
        initialize=initialize,
        origin=origin,
        linear_velocity=linear_velocity,
        angular_velocity=angular_velocity,
        u0=u0,
        theta0=theta0,
        udot0=udot0,
        thetadot0=thetadot0,
        save=save,
        )
end

"""
    time_domain_analysis!(system, assembly, tvec; kwargs...)

Pre-allocated version of `time_domain_analysis`.
"""
function time_domain_analysis!(system, assembly, tvec;
    prescribed_conditions=Dict{Int,PrescribedConditions{Float64}}(),
    distributed_loads=Dict{Int,DistributedLoads{Float64}}(),
    linear=false,
    linearization_state=nothing,
    update_linearization_state=false,
    method=:newton,
    linesearch=LineSearches.BackTracking(maxstep=1e6),
    ftol=1e-9,
    iterations=1000,
    reset_state=true,
    initialize=true,
    origin=(@SVector zeros(3)),
    linear_velocity=(@SVector zeros(3)),
    angular_velocity=(@SVector zeros(3)),
    u0=fill((@SVector zeros(3)), length(assembly.elements)),
    theta0=fill((@SVector zeros(3)), length(assembly.elements)),
    udot0=fill((@SVector zeros(3)), length(assembly.elements)),
    thetadot0=fill((@SVector zeros(3)), length(assembly.elements)),
    save=1:length(tvec)
    )

    # check to make sure the simulation is dynamic
    @assert system.static == false

    if reset_state
        reset_state!(system)
    end

    # perform initial condition analysis
    if initialize
        system, converged = initial_condition_analysis!(system, assembly, tvec[1];
            prescribed_conditions=prescribed_conditions,
            distributed_loads=distributed_loads,
            linear=linear,
            linearization_state=linearization_state,
            method=method,
            linesearch=linesearch,
            ftol=ftol,
            iterations=iterations,
            reset_state=false,
            origin=origin,
            linear_velocity=linear_velocity,
            angular_velocity=angular_velocity,
            u0=u0,
            theta0=theta0,
            udot0=udot0,
            thetadot0=thetadot0,
            )
    else
        # converged by default
        converged = true
    end

    # unpack pre-allocated storage and pointers for system
    x = system.x
    F = system.r
    J = system.K
    force_scaling = system.force_scaling
    mass_scaling = system.mass_scaling
    irow_point = system.irow_point
    irow_elem = system.irow_elem
    irow_elem1 = system.irow_elem1
    irow_elem2 = system.irow_elem2
    icol_point = system.icol_point
    icol_elem = system.icol_elem
    udot = system.udot
    θdot = system.θdot
    Pdot = system.Pdot
    Hdot = system.Hdot

    # number of beam elements
    nelem = length(assembly.elements)

    # initialize storage for each time step
    isave = 1
    history = Vector{AssemblyState{eltype(system)}}(undef, length(save))

    # add initial state to the solution history
    if isave in save
        pcond = typeof(prescribed_conditions) <: AbstractDict ?
            prescribed_conditions : prescribed_conditions(tvec[1])
        history[isave] = AssemblyState(system, assembly, prescribed_conditions=pcond)
        isave += 1
    end

    # --- Begin Time Domain Simulation --- #

    for it = 2:length(tvec)

        # update current time
        system.t = tvec[it]

        # current time step size
        dt = tvec[it] - tvec[it - 1]

        # current parameters
        pcond = typeof(prescribed_conditions) <: AbstractDict ?
            prescribed_conditions : prescribed_conditions(tvec[it])
        dload = typeof(distributed_loads) <: AbstractDict ?
            distributed_loads : distributed_loads(tvec[it])
        x0 = typeof(origin) <: AbstractVector ? SVector{3}(origin) : SVector{3}(origin(tvec[it]))
        v0 = typeof(linear_velocity) <: AbstractVector ? SVector{3}(linear_velocity) : SVector{3}(linear_velocity(tvec[it]))
        ω0 = typeof(angular_velocity) <: AbstractVector ? SVector{3}(angular_velocity) : SVector{3}(angular_velocity(tvec[it]))

        # current initialization parameters
        for ielem = 1:nelem
            icol = icol_elem[ielem]
            # get beam element states
            u = SVector(x[icol], x[icol + 1], x[icol + 2])
            θ = SVector(x[icol + 3], x[icol + 4], x[icol + 5])
            P = SVector(x[icol + 12], x[icol + 13], x[icol + 14]) .* mass_scaling
            H = SVector(x[icol + 15], x[icol + 16], x[icol + 17]) .* mass_scaling
            # extract rotation parameters
            C = get_C(θ)
            Cab = assembly.elements[ielem].Cab
            CtCab = C' * Cab
            # store `udot_init` in `udot`
            udot[ielem] = 2 / dt * u + udot[ielem]
            # store `θdot_init` in `θdot`
            θdot[ielem] = 2 / dt * θ + θdot[ielem]
            # store `CtCabPdot_init` in `Pdot`
            Pdot[ielem] = 2 / dt * CtCab * P + CtCab * Pdot[ielem]
            # store `CtCabHdot_init` in `Hdot`
            Hdot[ielem] = 2 / dt * CtCab * H + CtCab * Hdot[ielem]
        end

        # solve for the state variables at the next time step
        f! = (F, x) -> system_residual!(F, x, assembly, pcond, dload, force_scaling,
            mass_scaling, irow_point, irow_elem, irow_elem1, irow_elem2, icol_point, icol_elem,
            x0, v0, ω0, udot, θdot, Pdot, Hdot, dt)

        j! = (J, x) -> system_jacobian!(J, x, assembly, pcond, dload, force_scaling,
            mass_scaling, irow_point, irow_elem, irow_elem1, irow_elem2, icol_point, icol_elem,
            x0, v0, ω0, udot, θdot, Pdot, Hdot, dt)

        # solve system of equations
        if linear
            # linear analysis
            if !update_linearization_state
                if isnothing(linearization_state)
                    x .= 0
                else
                    x .= linearization_state
                end
            end
            f!(F, x)
            j!(J, x)
            x .-= safe_lu(J) \ F
        else
            df = OnceDifferentiable(f!, j!, x, F, J)

            result = NLsolve.nlsolve(df, x,
                linsolve=(x, A, b) -> ldiv!(x, safe_lu(A), b),
                method=method,
                linesearch=linesearch,
                ftol=ftol,
                iterations=iterations)

            x .= result.zero
            J .= df.DF
        end

        # set new state rates
        for ielem = 1:nelem
            icol = icol_elem[ielem]
            # get beam element states
            u = SVector(x[icol], x[icol + 1], x[icol + 2])
            θ = SVector(x[icol + 3], x[icol + 4], x[icol + 5])
            P = SVector(x[icol + 12], x[icol + 13], x[icol + 14]) .* mass_scaling
            H = SVector(x[icol + 15], x[icol + 16], x[icol + 17]) .* mass_scaling
            # extract rotation parameters
            C = get_C(θ)
            Cab = assembly.elements[ielem].Cab
            CtCab = C' * Cab
            # save state rates
            udot[ielem] = 2 / dt * u - udot[ielem]
            θdot[ielem] = 2 / dt * θ - θdot[ielem]
            Pdot[ielem] = 2 / dt * P - CtCab' * Pdot[ielem]
            Hdot[ielem] = 2 / dt * H - CtCab' * Hdot[ielem]
        end

        # add state to history
        if it in save
            history[isave] = AssemblyState(system, assembly, prescribed_conditions=prescribed_conditions)
            isave += 1
        end

        # stop early if unconverged
        if !linear && !result.f_converged
            converged = false
            break
        end

    end

    return system, history, converged
end


import Base:
  # types
  Array,
  CartesianIndices,
  Vector,
  NTuple,
  Tuple,
  # symbols
  +,
  -,
  *,
  ^,
  /,
  ==,
  <,
  >,
  !,
  # functions
  adjoint,
  allunique,
  axes,
  complex,
  convert,
  copy,
  copyto!,
  deepcopy,
  deleteat!,
  eachindex,
  eltype,
  fill!,
  filter,
  filter!,
  findall,
  findfirst,
  getindex,
  hash,
  intersect,
  intersect!,
  isapprox,
  isassigned,
  isempty,
  isless,
  iterate,
  keys,
  lastindex,
  length,
  map,
  map!,
  ndims,
  permutedims,
  promote_rule,
  push!,
  resize!,
  setdiff,
  setdiff!,
  setindex!,
  show,
  similar,
  size,
  summary,
  truncate,
  zero,
  # macros
  @propagate_inbounds

import Base.Broadcast:
  # types
  AbstractArrayStyle,
  Broadcasted,
  BroadcastStyle,
  DefaultArrayStyle,
  Style,
  # functions
  _broadcast_getindex,
  broadcasted,
  broadcastable,
  instantiate

import ITensors.ContractionSequenceOptimization:
  contraction_cost, optimal_contraction_sequence

import HDF5: read, write

import LinearAlgebra:
  axpby!,
  axpy!,
  dot,
  eigen,
  exp,
  factorize,
  lmul!,
  mul!,
  norm,
  normalize!,
  qr,
  rmul!,
  svd,
  tr

using ITensors.NDTensors:
  EmptyNumber,
  _Tuple,
  _NTuple,
  blas_get_num_threads,
  eachblock,
  eachdiagblock,
  fill!!,
  randn!!,
  timer

import ITensors.NDTensors:
  # Modules
  Strided, # to control threading
  # Types
  AliasStyle,
  AllowAlias,
  NeverAlias,
  # Methods
  array,
  blockdim,
  blockoffsets,
  contract,
  dense,
  denseblocks,
  dim,
  dims,
  disable_tblis,
  eachnzblock,
  enable_tblis,
  ind,
  inds,
  insertblock!,
  insert_diag_blocks!,
  matrix,
  maxdim,
  mindim,
  nblocks,
  nnz,
  nnzblocks,
  nzblock,
  nzblocks,
  outer,
  permuteblocks,
  polar,
  scale!,
  setblockdim!,
  setinds,
  setstorage,
  sim,
  storage,
  sum,
  tensor,
  truncate!,
  using_tblis,
  vector,
  # Deprecated
  addblock!,
  store

import Random: randn!

using SerializedElementArrays: SerializedElementVector

const DiskVector{T} = SerializedElementVector{T}

import SerializedElementArrays: disk


abstract type AbstractSherlog <: AbstractFloat end

struct Sherlog64{T<:AbstractFloat,i} <: AbstractSherlog    # T is the bitpattern type for logging, i the logbook id
    val::Float64                                           # the value is always Float64
end

# conversions back from Sherlog64
Base.Float64(x::Sherlog64) = x.val
Base.Float32(x::Sherlog64) = Float32(x.val)
Base.Float16(x::Sherlog64) = Float16(x.val)
Base.Int64(x::Sherlog64) = Int64(x.val)
Base.Int32(x::Sherlog64) = Int32(x.val)
Base.Int16(x::Sherlog64) = Int16(x.val)

# otherwise conversion to itself is ambigious
Sherlog64{T,i}(x::Sherlog64) where {T<:AbstractFloat,i} = Sherlog64{T,i}(x.val)
Sherlog64(x::Sherlog64{T,i}) where {T<:AbstractFloat,i} = Sherlog64{T,i}(x.val)

# generator functions from Float, Int and Bool
Sherlog64(x::AbstractFloat) = Sherlog64{Float16,1}(x)   # use Float16 and logbook #1 as standard
Sherlog64{T}(x::AbstractFloat) where T = Sherlog64{T,1}(x)      # use logbook #1 if no logbook provided
Sherlog64(x::Integer) = Sherlog64{Float16,1}(x)           # same for integers
Sherlog64{T}(x::Integer) where T = Sherlog64{T,1}(x)
Sherlog64{T}(x::Bool) where T = if x Sherlog64{T,1}(1) else Sherlog64{T,1}(0) end

Base.promote_rule(::Type{Int64},::Type{Sherlog64{T,i}}) where {T,i} = Sherlog64{T,i}
Base.promote_rule(::Type{Int32},::Type{Sherlog64{T,i}}) where {T,i} = Sherlog64{T,i}
Base.promote_rule(::Type{Int16},::Type{Sherlog64{T,i}}) where {T,i} = Sherlog64{T,i}

Base.promote_rule(::Type{Float64},::Type{Sherlog64{T,i}}) where {T,i} = Sherlog64{T,i}
Base.promote_rule(::Type{Float32},::Type{Sherlog64{T,i}}) where {T,i} = Sherlog64{T,i}
Base.promote_rule(::Type{Float16},::Type{Sherlog64{T,i}}) where {T,i} = Sherlog64{T,i}

Base.bitstring(x::Sherlog64) = bitstring(x.val)
Base.show(io::IO,x::Sherlog64) = print(io,"Sherlog64(",string(x.val),")")

Base.eps(::Type{Sherlog64{T,i}}) where {T,i} = eps(Float64)
Base.eps(x::Sherlog64) = eps(x.val)

Base.typemin(::Sherlog64{T,i}) where {T,i} = Sherlog64{T,i}(typemin(Float64))
Base.typemax(::Sherlog64{T,i}) where {T,i} = Sherlog64{T,i}(typemax(Float64))
Base.floatmin(::Sherlog64{T,i}) where {T,i} = Sherlog64{T,i}(floatmin(Float64))
Base.floatmax(::Sherlog64{T,i}) where {T,i} = Sherlog64{T,i}(floatmax(Float64))
Base.precision(::Sherlog64{T,i}) where {T,i} = Sherlog64{T,i}(precision(Float64))

function +(x::Sherlog64{T,i},y::Sherlog64{T,i}) where {T,i}
    r = x.val + y.val
    log_it(T,r,i)
    return Sherlog64{T,i}(r)
end

function -(x::Sherlog64{T,i},y::Sherlog64{T,i}) where {T,i}
    r = x.val - y.val
    log_it(T,r,i)
    return Sherlog64{T,i}(r)
end

function *(x::Sherlog64{T,i},y::Sherlog64{T,i}) where {T,i}
    r = x.val * y.val
    log_it(T,r,i)
    return Sherlog64{T,i}(r)
end

function /(x::Sherlog64{T,i},y::Sherlog64{T,i}) where {T,i}
    r = x.val / y.val
    log_it(T,r,i)
    return Sherlog64{T,i}(r)
end

function ^(x::Sherlog64{T,i},y::Sherlog64{T,i}) where {T,i}
    r = x.val ^ y.val
    log_it(T,r,i)
    return Sherlog64{T,i}(r)
end

for O in ( :(-), :(+),
           :sign,
           :prevfloat, :nextfloat,
           :round, :trunc, :ceil, :floor,
           :inv, :abs, :sqrt, :cbrt,
           :exp, :expm1, :exp2, :exp10,
           :log, :log1p, :log2, :log10,
           :rad2deg, :deg2rad, :mod2pi, :rem2pi,
           :sin, :cos, :tan, :csc, :sec, :cot,
           :asin, :acos, :atan, :acsc, :asec, :acot,
           :sinh, :cosh, :tanh, :csch, :sech, :coth,
           :asinh, :acosh, :atanh, :acsch, :asech, :acoth,
           :sinc, :sinpi, :cospi,
           :sind, :cosd, :tand, :cscd, :secd, :cotd,
           :asind, :acosd, :atand, :acscd, :asecd, :acotd
          )
    @eval begin
        function Base.$O(x::Sherlog64{T,i}) where {T,i}
            r = $O(x.val)
            log_it(T,r,i)
            return Sherlog64{T,i}(r)
        end
    end
end

for O in ( :(<), :(<=))
    @eval begin
        Base.$O(x::Sherlog64, y::Sherlog64) = $O(x.val, y.val)
    end
end


using QuantumOptics: tensor, CompositeBasis
using .PhysicalConstants: ħ, c

export Trap, get_basis, Efield_from_pi_time, Efield_from_pi_time!,
       transition_frequency, set_gradient!, Efield_from_rabi_frequency,
       Efield_from_rabi_frequency!, global_beam!, get_η

#############################################################################################
# an ion trap with a linear ion chain configuration
#############################################################################################

"""
    Trap(;
            configuration::LinearChain, B::Real=0, Bhat::NamedTuple{(:x,:y,:z)=ẑ, ∇B::Real=0,
             δB::Union{Real,Function}=0, lasers::Vector{Laser}
    )

Information necessary to describe the Hamiltonian for a collection of ions in a linear chain
interacting with laser light.
**user-defined fields**
* `configuration<:LinearChain`
* `B`: A real value describing the mean magnitude of the B-field [Tesla].
* `Bhat::NamedTuple{(:x,:y,:z)}`: Describes the direction of the B-field (defaults to ẑ).
* `∇B`: Magnitude of the B-field gradient. We assume that the gradient always points along the
        z-direction. [Tesla / meter]
* `δB::Function`: Time-dependence of the B-field [Tesla]
* `lasers::Array{<:Laser}`: For each laser in the array, the pointing field should contain
        an array of `Tuple{Int,Real}`. The first element specifies the index of an ion
        in the `ions` field that the laser interacts with. The second element specifies a
        scaling factor for the strength of that interaction (to be used, e.g., for
        modeling cross-talk).
**derived fields**
* `_cnst_δB::Bool`: A Boolean flag signifying whether or not `δB` is a constant function.
* `basis<:CompositeBasis`: The basis for describing the combined system, ions + vibrational
        modes. If constructing the Hamiltonian explictly (with [`hamiltonian`](@ref)), then
        the ordering of the basis is set, by convention, as
        ``ion₁ ⊗ ion₂ ⊗ ... ⊗ ion_N ⊗ mode₁ ⊗ mode₂ ⊗ ... ⊗ mode_N``, where the ion bases are
        ordered according to the order in `T.configuration.ions` and the vibrational modes
        are ordered according to the order in
        `[T.configuration.vibrational_modes.x, T.configuration.vibrational_modes.y,
        T.configuration.vibrational_modes.z]`.
    E.g. for:

    ```
    chain = LinearChain(ions=[C1, C2], com_frequencies=(x=2e6,y=2e6,z=1e6),
    selected_modes=(x=[1, 2], y=[], z=[1]))
    ```

    The ordering of the basis would be

    `C1.basis ⊗ C2.basis ⊗ chain.vibrational_modes.x[1].basis
    ⊗ chain.vibrational_modes.x[2].basis ⊗ chain.vibrational_modes.z[1].basis`

    Otherwise, the ordering is according to the form of the initial state used in the solver.
"""
mutable struct Trap
    configuration::LinearChain
    B::Real
    Bhat::NamedTuple{(:x,:y,:z)}
    ∇B::Real
    δB::Function
    lasers::Array{<:Laser}
    basis::CompositeBasis
    _cnst_δB::Bool
    function Trap(;
            configuration::LinearChain, B=0, Bhat=ẑ, ∇B=0, δB::TδB=0, lasers=Laser[]
        ) where {TδB}
        warn = nothing
        for i in 1:length(lasers)
            if length(lasers[i].pointing) == 0
                for n in eachindex(configuration.ions)
                    push!(lasers[i].pointing, (n, 1.0))
                end
            end
            for j in i+1:length(lasers)
                if lasers[j] ≡ lasers[i]
                    lasers[j] = copy(lasers[i])
                    if isnothing(warn)
                        warn = "Some lasers point to the same thing. Making copies."
                        @warn warn
                    end
                end
            end
        end
        @assert isapprox(norm(Bhat), 1, rtol=1e-6) "!(|$Bhat| = 1)"
        for (li, l) in enumerate(lasers), p in l.pointing
            @assert p[1] <= length(configuration.ions) (
                """lasers[$li] points at configuration.ions[$(p[1])], but there are only
                 $(length(configuration.ions)) ions."""
            )
        end
        if TδB <: Number
            _cnst_δB = true
            δBt(t) = δB
        else
            _cnst_δB = false
            δBt = δB
        end
        basis = tensor(
            configuration.ions...,
            configuration.vibrational_modes.x...,
            configuration.vibrational_modes.y...,
            configuration.vibrational_modes.z...,
        )
        new(configuration, B, Bhat, ∇B, δBt, lasers, basis, _cnst_δB)
    end
end

function Base.setproperty!(T::Trap, s::Symbol, v::Tv) where{Tv}
    if s == :δB
        if Tv <: Number
            _cnst_δB = true
            vt(t) = v
        else
            _cnst_δB = false
            vt = v
        end
        Core.setproperty!(T, s, vt)
        Core.setproperty!(T, :_cnst_δB, _cnst_δB)
    elseif s == :basis || s == :_cnst_δB
        return
    elseif s == :configuration
        Core.setproperty!(T, s, v)
        basis = tensor(
            v.ions..., v.vibrational_modes.x..., v.vibrational_modes.y...,
            v.vibrational_modes.z...,
        )
        Core.setproperty!(T, :basis, basis)
    else
        Core.setproperty!(T, s, v)
    end
end

Base.show(io::IO, T::Trap) = print(io, "Trap")  # suppress long output

#############################################################################################
# general functions
#############################################################################################

"""
    get_basis(T::Trap)
Returns the composite basis describing the Hilbert space for `T`.
"""
function get_basis(T::Trap)::CompositeBasis
    tensor(
            T.configuration.ions...,
            T.configuration.vibrational_modes.x...,
            T.configuration.vibrational_modes.y...,
            T.configuration.vibrational_modes.z...,
        )
end

"""
    global_beam!(T::Trap, laser::Laser)
Set `laser` to shine with full intensity on all ions in `Trap`.
"""
function global_beam!(T::Trap, laser::Laser)
    for n in eachindex(T.configuration.ions)
        push!(laser.pointing, (n, 1.0))
    end
end

"""
    Efield_from_pi_time(
        pi_time::Real, Bhat::NamedTuple{(:x,:y,:z)}, laser::Laser, ion::Ion,
        transition::Union{Tuple{String,String},Vector{<:String}}
    )
Compute the E-field needed to get a certain `pi_time` with a certain `laser`-`ion`
`transition`.

Alternatively, one may use
```
Efield_from_pi_time(
            pi_time::Real, T::Trap, laser_index::Int, ion_index::Int,
            transition::Union{Tuple{String,String},Vector{<:String}}
        )
```
which is the same as
`Efield_from_pi_time(pi_time, T.Bhat, T.lasers[laser_index], T.configuration.ions[ion_index],
transition)`
"""
function Efield_from_pi_time(
    pi_time::Real, Bhat::NamedTuple{(:x,:y,:z)}, laser::Laser, ion::Ion,
    transition::Union{Tuple{String,String},Vector{<:String}}
)
    p = laser.pointing
    (γ, ϕ) = map(x -> rad2deg(acos(ndot(Bhat, x))), [laser.ϵ, laser.k])
    s_indx = findall(x -> x[1] == ion.number, p)
    @assert length(s_indx) > 0 "This laser doesn't shine on this ion"
    s = laser.pointing[s_indx[1]][2]
    Ω = s * ion.selected_matrix_elements[tuple(transition...)](1, γ, ϕ)
    if Ω < 1e-15
        # even when coupling strength is zero, numerical error causes it to be finite
        # (on order 1e-16), this is a band-aid to prevent users from unknowingly setting
        # the E-field to something absurd (like 1e20 V/m)
        return Inf
    end
    1 / (2Ω * pi_time)
end

function Efield_from_pi_time(
        pi_time::Real, T::Trap, laser_index::Int, ion_index::Int,
        transition::Union{Tuple{String,String},Vector{<:String}}
    )
    Efield_from_pi_time(
            pi_time, T.Bhat, T.lasers[laser_index], T.configuration.ions[ion_index],
            transition
        )
end

"""
    Efield_from_pi_time!(
            pi_time::Real, Bhat::NamedTuple{(:x,:y,:z)}, laser::Laser, ion::Ion,
            transition::Union{Tuple{String,String},Vector{<:String}}
        )
Same as [`Efield_from_pi_time`](@ref), but updates `laser[:E]` in-place.
"""
function Efield_from_pi_time!(
        pi_time::Real, Bhat::NamedTuple{(:x,:y,:z)}, laser::Laser, ion::Ion,
        transition::Union{Tuple{String,String},Vector{<:String}}
    )
    Efield::Float64 = Efield_from_pi_time(pi_time, Bhat, laser, ion, transition)
    laser.E = t -> Efield
end

function Efield_from_pi_time!(
        pi_time::Real, T::Trap, laser_index::Int, ion_index::Int,
        transition::Union{Tuple{String,String},Vector{<:String}}
    )
    Efield::Float64 = Efield_from_pi_time(pi_time, T, laser_index, ion_index, transition)
    T.lasers[laser_index].E = t -> Efield
end

"""
    Efield_from_rabi_frequency(
        Ω::Real, Bhat::NamedTuple{(:x,:y,:z)}, laser::Laser, ion::Ion,
        transition::Union{Tuple{String,String},Vector{<:String}}
    )
Compute the Efield needed to get a certain rabi frequency `Ω` with a certain `laser`-`ion`
`transition`.

Alternatively, one may use
```
Efield_from_rabi_frequency(
            Ω::Real, T::Trap, laser_index::Int, ion_index::Int,
            transition::Union{Tuple{String,String},Vector{<:String}}
        )
```
which is the same as
`Efield_from_rabi_frequency(pi_time, T.Bhat, T.lasers[laser_index],
T.configuration.ions[ion_index], transition)`
"""
function Efield_from_rabi_frequency(
    Ω::Real, Bhat::NamedTuple{(:x,:y,:z)}, laser::Laser, ion::Ion,
    transition::Union{Tuple{String,String},Vector{<:String}})
    Efield_from_pi_time(1 / 2Ω, Bhat, laser, ion, transition)
end

function Efield_from_rabi_frequency(
        Ω::Real, T::Trap, laser_index::Int, ion_index::Int,
        transition::Union{Tuple{String,String},Vector{<:String}}
    )
    Efield_from_rabi_frequency(
            Ω, T.Bhat, T.lasers[laser_index], T.configuration.ions[ion_index],
            transition
        )
end

"""
    Efield_from_rabi_frequency!(
        Ω::Real, Bhat::NamedTuple{(:x,:y,:z)}, laser::Laser, ion::Ion,
        transition::Union{Tuple{String,String},Vector{<:String}}
    )
Same as [`Efield_from_rabi_frequency`](@ref), but updates `laser[:E]` in-place.
"""
function Efield_from_rabi_frequency!(
        Ω::Real, Bhat::NamedTuple{(:x,:y,:z)}, laser::Laser, ion::Ion,
        transition::Union{Tuple{String,String},Vector{<:String}}
    )
    Efield::Float64 = Efield_from_rabi_frequency(Ω, Bhat, laser, ion, transition)
    laser.E = t -> Efield
end

function Efield_from_rabi_frequency!(
        Ω::Real, T::Trap, laser_index::Int, ion_index::Int,
        transition::Union{Tuple{String,String},Vector{<:String}}
    )
    Efield::Float64 = Efield_from_rabi_frequency(Ω, T, laser_index, ion_index, transition)
    T.lasers[laser_index].E = t -> Efield
end

"""
    transition_frequency(
        B::Real, ion::Ion, transition::Union{Tuple{String,String},Vector{<:String}}
    )
Compute the transition frequency of the `ion`'s selected transition under Bfield `B`.
Alternatively, one may use:
```
transition_frequency(
        T::Trap, ion::Ion, transition::Union{Tuple{String,String},Vector{<:String}}
    )
```
which is the same as `transition_frequency(T.B, ion, transition)` or
```
transition_frequency(
        T::Trap, ion_index::Int, transition::Union{Tuple{String,String},Vector{<:String}}
    )
```
which is the same as `transition_frequency(T.B, T.configuration.ions[ion_index], transition)`.
"""
function transition_frequency(
        B::Real, ion::Ion, transition::Union{Tuple{String,String},Vector{<:String}}
    )
    diff([map(x -> zeeman_shift(B, ion.selected_level_structure[x]), transition)...])[1]
end

function transition_frequency(
        T::Trap, ion::Ion, transition::Union{Tuple{String,String},Vector{<:String}}
    )
    B = T.B + T.∇B * ion.position
    transition_frequency(B, ion, transition)
end

function transition_frequency(
        T::Trap, ion_index::Int, transition::Union{Tuple{String,String},Vector{<:String}}
    )
    ion = T.configuration.ions[ion_index]
    B = T.B + T.∇B * ion.position
    transition_frequency(B, ion, transition)
end

"""
    set_gradient!(
            T::Trap, ion_indxs::Tuple{Int,Int}, transition::Tuple{String,String}, f::Real
        )
Sets the Bfield gradient in place to achieve a detuning `f` between the `transition` of two
ions, which are assumed to be of the same species. `ion_indxs` refer to the
ordering of the ions in the chain.
"""
function set_gradient!(
        T::Trap, ion_indxs::Tuple{Int,Int}, transition::Tuple{String,String}, f::Real
    )
    ion1 = T.configuration.ions[ion_indxs[1]]
    ion2 = T.configuration.ions[ion_indxs[2]]
    separation = abs(ion1.position - ion2.position)
    E1, E2 = map(x -> zeeman_shift(1, ion1.selected_level_structure[x]), transition)
    T.∇B = f / (abs(E2 - E1) * separation)
end

# In QunatumOptics.jl, this method will return true whenever the shapes of b1 and b2 match,
# but we'd like to distinguish, i.e., between Ion ⊗ mode1 ⊗ mode2 and Ion ⊗ mode2 ⊗ mode1
# when mode1.N == mode2.N but mode1.axis ≠ mode2.axis.
function (QuantumOptics.:(==)(b1::T, b2::T)
    where {T<:CompositeBasis{<:Vector{Int},<:Tuple{Vararg{<:IonSimBasis}}}})
N = length(b1.bases)
if N ≠ length(b2.bases)
    return false
end
for i in 1:N
    if !(b1.bases[i] == b2.bases[i])
        return false
    end
end
return true
end

"""
    get_η(V::VibrationalMode, L::Laser, I::Ion)
The Lamb-Dicke parameter:
``|k|cos(\\theta)\\sqrt{\\frac{\\hbar}{2m\\nu}}``
for a given vibrational mode, ion and laser.
"""
function get_η(V::VibrationalMode, L::Laser, I::Ion; scaled=false)
    @fastmath begin
        k = 2π / L.λ
        scaled ? ν = 1 : ν = V.ν
        x0 = √(ħ / (2 * I.mass * 2π * ν))
        cosθ = ndot(L.k, V.axis)
        k * x0 * cosθ * V.mode_structure[I.number]
    end
end