"""
    TimeSeries{D, Ω, I, T, TT} <: AbstractDiagnostic

A diagnostic for collecting and storing time series.
"""
struct TimeSeries{D, Ω, I, T, TT} <: AbstractDiagnostic
    diagnostic :: D
     frequency :: Ω
      interval :: I
          data :: T
          time :: Vector{TT}
end

"""
    TimeSeries(diagnostic, model; frequency=nothing, interval=nothing)

A `TimeSeries` `Diagnostic` that records a time series of `diagnostic(model)`.

Example
=======
```julia
julia> model = IncompressibleModel(grid=RegularCartesianGrid(size=(16, 16, 16), length=(1, 1, 1)));

julia> max_u = TimeSeries(FieldMaximum(abs, model.velocities.u), model; frequency=1)

julia> model.diagnostics[:max_u] = max_u; data(model.velocities.u) .= π; time_step!(model, Nt=3, Δt=1e-16)

julia> max_u.data
3-element Array{Float64,1}:
 3.141592653589793
 3.1415926025389127
 3.1415925323439517
```
"""
function TimeSeries(diagnostic, model; frequency=nothing, interval=nothing)
    TD = typeof(diagnostic(model))
    TT = typeof(model.clock.time)
    return TimeSeries(diagnostic, frequency, interval, TD[], TT[])
end

@inline (time_series::TimeSeries)(model) = time_series.diagnostic(model)

function run_diagnostic(model, diag::TimeSeries)
    push!(diag.data, diag(model))
    push!(diag.time, model.clock.time)
    return nothing
end

"""
    TimeSeries(diagnostics::NamedTuple, model; frequency=nothing, interval=nothing)

A `TimeSeries` `Diagnostic` that records a `NamedTuple` of time series of
`diag(model)` for each `diag` in `diagnostics`.

Example
=======
```julia
julia> model = IncompressibleModel(grid=RegularCartesianGrid(size=(16, 16, 16), length=(1, 1, 1))); Δt = 1.0;

julia> cfl = TimeSeries((adv=AdvectiveCFL(Δt), diff=DiffusiveCFL(Δt)), model; frequency=1);

julia> model.diagnostics[:cfl] = cfl; time_step!(model, Nt=3, Δt=Δt)

julia> cfl.data
(adv = [0.0, 0.0, 0.0, 0.0], diff = [0.0002688, 0.0002688, 0.0002688, 0.0002688])

julia> cfl.adv
4-element Array{Float64,1}:
 0.0
 0.0
 0.0
 0.0
```
"""
function TimeSeries(diagnostics::NamedTuple, model; frequency=nothing, interval=nothing)
    TT = typeof(model.clock.time)
    TDs = Tuple(typeof(diag(model)) for diag in diagnostics)
    data = NamedTuple{propertynames(diagnostics)}(Tuple(T[] for T in TDs))
    return TimeSeries(diagnostics, frequency, interval, data, TT[])
end

function (time_series::TimeSeries{<:NamedTuple})(model)
    names = propertynames(time_series.diagnostic)
    return NamedTuple{names}(Tuple(diag(model) for diag in time_series.diagnostics))
end

function run_diagnostic(model, diag::TimeSeries{<:NamedTuple})
    ntuple(Val(length(diag.diagnostic))) do i
        Base.@_inline_meta
        push!(diag.data[i], diag.diagnostic[i](model))
    end
    push!(diag.time, model.clock.time)
    return nothing
end

Base.getproperty(ts::TimeSeries{<:NamedTuple}, name::Symbol) = get_time_series_field(ts, name)

"Returns a field of time series, or a field of `time_series.data` when possible."
function get_time_series_field(time_series, name)
    if name ∈ propertynames(time_series)
        return getfield(time_series, name)
    else
        return getproperty(time_series.data, name)
    end
end


function serialize(vis::CoreVisualizer, queue::CommandQueue)
    utime = time_ns()
    delete_cmds = Dict{String, Any}[]
    setgeometry_cmds = Dict{String, Any}[]
    settransform_cmds = Dict{String, Any}[]
    for path in queue.delete
        push!(delete_cmds, Dict("path" => path))
    end
    for path in queue.setgeometry
        visdata = vis.tree[path].data
        push!(setgeometry_cmds, serialize(path, visdata.geometries))
    end
    for path in queue.settransform
        visdata = vis.tree[path].data
        tform = serialize(visdata.transform)
        push!(settransform_cmds,
              Dict{String, Any}("path" => path,
                                "transform" => tform
              )
        )
    end
    data = Dict{String, Any}(
        "utime" => utime,
        "delete" => delete_cmds,
        "setgeometry" => setgeometry_cmds,
        "settransform" => settransform_cmds
    )
end

function serialize(path::AbstractVector, geomdatas::Vector{GeometryData})
    params = serialize.(geomdatas)
    if length(params) == 1
        Dict("path" => serialize(path),
             "geometry" => params[1])
    else
        Dict("path" => serialize(path),
             "geometries" => params)
    end
end

function serialize(geomdata::GeometryData)
    params = serialize(geomdata.geometry)
    params["color"] = serialize(geomdata.color)
    transform = compose(geomdata.transform,
                        intrinsic_transform(geomdata.geometry))
    if transform != IdentityTransformation()
        params["transform"] = serialize(transform)
    end
    params
end

intrinsic_transform(g) = IdentityTransformation()
intrinsic_transform(geomdata::GeometryData) = intrinsic_transform(geomdata.geometry)
intrinsic_transform(g::HyperRectangle) = Translation(center(g)...)
intrinsic_transform(g::HyperSphere) = Translation(center(g)...)
intrinsic_transform(g::HyperEllipsoid) = Translation(center(g)...)
intrinsic_transform(g::HyperCube) = Translation(center(g)...)
function intrinsic_transform(g::Cylinder{3})
    # DrakeVisualizer wants a cylinder to lie along the z axis
    R = rotation_between(SVector(0, 0, 1), g.extremity)
    Translation(center(g)) ∘ LinearMap(R)
end

serialize(color::Colorant) = (red(color),
                              green(color),
                              blue(color),
                              alpha(color))
serialize(p::Path) = string.(p)
serialize(v::Vector) = v
serialize(v::Vec) = convert(Vector, v)
serialize(v::Point) = convert(Vector, v)
serialize(v::StaticArray) = convert(Vector, v)
serialize(face::Face{N, T}) where {N, T} =
  raw.(convert(Face{N, GeometryTypes.OffsetInteger{-1, Int}}, face))
serialize(g::HyperRectangle) = Dict("type" => "box", "lengths" => serialize(widths(g)))
serialize(g::HyperSphere) = Dict("type" => "sphere", "radius" => radius(g))
serialize(g::HyperEllipsoid) = Dict("type" => "ellipsoid", "radii" => serialize(radii(g)))
serialize(g::Cylinder{3}) = Dict("type" => "cylinder",
                                      "length" => norm(g.extremity),
                                      "radius" => radius(g))
serialize(g::HyperCube) = Dict("type" => "box", "lengths" => widths(g))
serialize(g::GeometryPrimitive) = serialize(GLNormalMesh(g))
serialize(g::Triad) = Dict{String, Any}("type" => "triad",
                                        "scale" => g.scale,
                                        "tube" => g.tube)

function serialize(g::AbstractMesh)
    Dict("type" => "mesh_data",
         "vertices" => serialize.(vertices(g)),
         "faces" => serialize.(faces(g)))
end

function serialize(f::MeshFile)
    Dict("type" => "mesh_file",
         "filename" => f.filename,
         "scale" => SVector(1.0, 1.0, 1.0))
end

function serialize(g::PointCloud)
    params = Dict("type" => "pointcloud",
                  "points" => serialize.(g.points),
                  "channels" => Dict{String, Any}())
    for (channel, values) in g.channels
        params["channels"][string(channel)] = serialize.(values)
    end
    params
end

function serialize(g::PolyLine)
    params = Dict("type" => "line",
                  "points" => serialize.(g.points),
                  "radius" => g.radius,
                  "closed" => g.closed
        )
    if g.start_head !== nothing
        params["start_head"] = true
        params["head_radius"] = g.start_head.radius
        params["head_length"] = g.start_head.length
    end
    if g.end_head !== nothing
        params["end_head"] = true
        params["head_radius"] = g.end_head.radius
        params["head_length"] = g.end_head.length
    end
    params
end

function serialize(tform::Transformation)
    Dict{String, Vector{Float64}}("translation" => translation(tform),
                      "quaternion" => quaternion(tform))
end

quaternion(::IdentityTransformation) = SVector(1., 0, 0, 0)
quaternion(tform::AbstractAffineMap) = quaternion(transform_deriv(tform, SVector(0., 0, 0)))
quaternion(matrix::UniformScaling) = quaternion(IdentityTransformation())
quaternion(matrix::AbstractMatrix) = quaternion(Quat(matrix))
quaternion(quat::Quat) = SVector(quat.w, quat.x, quat.y, quat.z)

translation(tform::Transformation) = tform(SVector(0., 0, 0))


include("../intcode.jl")

tape = Tape(parse_data(readline(open("input.txt"))))

const NORTH = 1
const SOUTH = 2
const WEST = 3
const EAST = 4

const dirs = Dict(
    NORTH => [0, +1],
    SOUTH => [0, -1],
    WEST => [-1, 0],
    EAST => [+1, 0],
)

const UNEXPLORED = -1
const WALL = 0
const EMPTY = 1
const OXYGEN = 2

const chars = Dict(
    UNEXPLORED => ' ',
    WALL => '#',
    EMPTY => '.',
    OXYGEN => 'x',
)

function squeeze_map(areamap, pos)
    is_explored = areamap .!= UNEXPLORED
    exploredX = (1:size(areamap, 1))[maximum(is_explored, dims=2)[:]]
    exploredY = (1:size(areamap, 2))[maximum(is_explored, dims=1)[:]]
    lowX, hiX = minimum(exploredX), maximum(exploredX)
    lowY, hiY = minimum(exploredY), maximum(exploredY)
    shifted_pos = pos - [lowX-1, lowY-1]
    areamap[lowX:hiX, lowY:hiY], shifted_pos
end

function plot_map(areamap, pos)
    for y=size(areamap, 2):-1:1
        for x=1:size(areamap, 1)
            if [x, y] == pos
                ch = (areamap[x, y] == OXYGEN) ? '@' : 'O'
            else
                ch = chars[areamap[x, y]]
            end
            print(ch)
        end
        println()
    end
end

function read_dir()
    valid_input = map(string, collect(keys(dirs)))
    while !((input = readline()) in valid_input) end
    parse(Int, input)
end

function steer!(tape)
    areamap = fill(-1, 1001, 1001)
    pos = [501, 501]
    areamap[pos...] = 1
    while true
        out = take!(tape.outch)
        if isa(out, AwaitingInput)
            print("Direction (N=1, S=2, W=3, E=4): ")
            direction = read_dir()
            put!(tape.inch, direction)
            output = take!(tape.outch)
            dir_vec = dirs[direction]
            newpos = pos + dir_vec
            areamap[newpos...] = output
            if output != WALL
                pos = newpos
            end

            plot_map(squeeze_map(areamap, pos)...)

        elseif isa(out, CodeFinished)
            println("FINISHED")
            return
        else
            println("OUTPUT: ", out)
        end
    end
end


@testset "Space" begin

@testset "1D grids" begin
    a = GridSpace((5, ))
    @test size(a) == (5, )
    @test typeof(a.s) <: Array{Array{Int64,1},1}
    @test a.metric == :chebyshev

    b = GridSpace((5, ), periodic = false)
    @test size(b) == (5, )
    @test typeof(b.s) <: Array{Array{Int64,1},1}
    @test b.metric == :chebyshev

    c = GridSpace((3, ), metric = :euclidean)
    @test size(c) == (3, )
    @test typeof(c.s) <: Array{Array{Int64,1},1}
    @test c.metric == :euclidean

    d = GridSpace((3, ), metric = :euclidean, periodic = false)
    @test size(d) == (3, )
    @test typeof(d.s) <: Array{Array{Int64,1},1}
    @test d.metric == :euclidean
end

@testset "2D grids" begin
    a = GridSpace((2, 3))
    @test size(a) == (2, 3)
    @test typeof(a.s) <: Array{Array{Int64,1},2}
    @test a.metric == :chebyshev

    b = GridSpace((2, 3), periodic = false)
    @test size(b) == (2, 3)
    @test typeof(b.s) <: Array{Array{Int64,1},2}
    @test b.metric == :chebyshev

    c = GridSpace((3, 3), metric = :euclidean)
    @test size(c) == (3, 3)
    @test typeof(c.s) <: Array{Array{Int64,1},2}
    @test c.metric == :euclidean

    d = GridSpace((3, 4), metric = :euclidean, periodic = false)
    @test size(d) == (3, 4)
    @test typeof(d.s) <: Array{Array{Int64,1},2}
    @test d.metric == :euclidean
end

@testset "3D grids" begin
    a = GridSpace((2, 3, 4))
    @test size(a) == (2, 3, 4)
    @test typeof(a.s) <: Array{Array{Int64,1},3}
    @test a.metric == :chebyshev

    b = GridSpace((2, 3, 7), periodic = false)
    @test size(b) == (2, 3, 7)
    @test typeof(b.s) <: Array{Array{Int64,1},3}
    @test b.metric == :chebyshev

    c = GridSpace((3, 3, 3), metric = :euclidean)
    @test size(c) == (3, 3, 3)
    @test typeof(c.s) <: Array{Array{Int64,1},3}
    @test c.metric == :euclidean

    d = GridSpace((3, 4, 2), metric = :euclidean, periodic = false)
    @test size(d) == (3, 4, 2)
    @test typeof(d.s) <: Array{Array{Int64,1},3}
    @test d.metric == :euclidean
end

@testset "Positions" begin
    space = GridSpace((3, 3))
    model = ABM(Agent1, space)
    empty = collect(empty_positions(model))
    @test length(empty) > 0
    for n in [1, 5, 6, 9, 2, 3, 4]
        add_agent!(empty[n], model)
    end
    # only positions (1,3) and (2,3) should be empty
    @test random_empty(model) ∈ [(1,3), (2,3)]
    pos_map = [(1, 1)  (1, 2)  (1, 3)
               (2, 1)  (2, 2)  (2, 3)
               (3, 1)  (3, 2)  (3, 3)]
    @test collect(positions(model)) == pos_map
    random_positions = positions(model, :random)
    @test all(n ∈ pos_map for n in random_positions)
    @test positions(model, :population) == [pos_map[i] for i in [1, 2, 3, 4, 5, 6, 9, 7, 8]]
    @test length(ids_in_position(5, model)) > length(ids_in_position(7, model))
    @test_throws ErrorException positions(model, :notreal)
end

@testset "Nearby Agents" begin
    undirected = ABM(Agent5, GraphSpace(path_graph(5)))
    @test nearby_positions(3, undirected) == [2, 4]
    @test nearby_positions(1, undirected) == [2]
    @test nearby_positions(5, undirected) == [4]
    @test nearby_positions(3, undirected; neighbor_type = :out) ==
    nearby_positions(3, undirected; neighbor_type = :in) ==
    nearby_positions(3, undirected; neighbor_type = :all) ==
    nearby_positions(3, undirected; neighbor_type = :default)
    add_agent!(1, undirected, rand())
    add_agent!(2, undirected, rand())
    add_agent!(3, undirected, rand())
    # We expect id 2 to be included for a grid based search
    @test sort!(collect(nearby_ids(2, undirected))) == [1, 2, 3]
    # But to be excluded if we are looking around it.
    @test sort!(collect(nearby_ids(undirected[2], undirected))) == [1, 3]

    directed = ABM(Agent5, GraphSpace(path_digraph(5)))
    @test nearby_positions(3, directed) == [4]
    @test nearby_positions(1, directed) == [2]
    @test nearby_positions(5, directed) == []
    @test nearby_positions(3, directed; neighbor_type = :default) ==
          nearby_positions(3, directed)
    @test nearby_positions(3, directed; neighbor_type = :in) == [2]
    @test nearby_positions(3, directed; neighbor_type = :out) == [4]
    @test sort!(nearby_positions(3, directed; neighbor_type = :all)) == [2, 4]
    add_agent!(1, directed, rand())
    add_agent!(2, directed, rand())
    add_agent!(3, directed, rand())
    @test sort!(nearby_ids(2, directed)) == [2, 3]
    @test sort!(nearby_ids(2, directed; neighbor_type=:in)) == [1, 2]
    @test sort!(nearby_ids(2, directed; neighbor_type=:all)) == [1, 2, 3]
    @test collect(nearby_ids(directed[2], directed)) == [3]
    @test collect(nearby_ids(directed[2], directed; neighbor_type=:in)) == [1]
    @test sort!(collect(nearby_ids(directed[2], directed; neighbor_type=:all))) == [1, 3]

    gridspace = ABM(Agent3, GridSpace((3, 3); metric = :euclidean, periodic = false))
    @test collect(nearby_positions((2, 2), gridspace)) == [(2, 1), (1, 2), (3, 2), (2, 3)]
    @test collect(nearby_positions((1, 1), gridspace)) == [(2, 1), (1, 2)]
    a = add_agent!((2, 2), gridspace, rand())
    add_agent!((3, 2), gridspace, rand())
    @test collect(nearby_ids((1, 2), gridspace)) == [1]
    @test sort!(collect(nearby_ids((1, 2), gridspace, 2))) == [1, 2]
    @test sort!(collect(nearby_ids((2, 2), gridspace))) == [1, 2]
    @test collect(nearby_ids(a, gridspace)) == [2]

    continuousspace = ABM(Agent6, ContinuousSpace(2; extend = (1, 1)))
    a = add_agent!((0.5, 0.5), continuousspace, (0.2, 0.1), 0.01)
    b = add_agent!((0.6, 0.5), continuousspace, (0.1, -0.1), 0.01)
    @test_throws ErrorException nearby_positions(1, continuousspace)
    @test nearby_ids(a, continuousspace, 0.05) == []
    @test nearby_ids(a, continuousspace, 0.1) == [2]
    @test sort!(nearby_ids((0.55, 0.5), continuousspace, 0.05)) == [1, 2]
    move_agent!(a, continuousspace)
    move_agent!(b, continuousspace)
    @test nearby_ids(a, continuousspace, 0.1) == []
    # Checks for type instability #208
    @test typeof(nearby_ids(a, continuousspace, 0.1)) <: Vector{Int}
    @test typeof(nearby_ids((0.55, 0.5), continuousspace, 0.05)) <: Vector{Int}
end

@testset "Discrete space mutability" begin
    model1 = ABM(Agent1, GridSpace((3,3)))

    agent = add_agent!((1,1), model1)
    @test agent.pos == (1, 1)
    @test agent.id == 1
    pos1 = ids_in_position((1,1), model1)
    @test length(pos1) == 1
    @test pos1[1] == 1

    move_agent!(agent, (2,2), model1)
    @test agent.pos == (2,2)
    pos1 = ids_in_position((1,1), model1)
    @test length(pos1) == 0
    pos2 = ids_in_position((2,2), model1)
    @test pos2[1] == 1

    # %% get/set testing
    model = ABM(Agent1, GridSpace((10,10)); properties=Dict(:number => 1, :nested => BadAgent(1,1)))
    add_agent!(model)
    add_agent!(model)
    @test (model.number += 1) == 2
    @test (model.nested.pos = 5) == 5
    @test_throws ErrorException (model.space = ContinuousSpace(2))
end

end


@with_kw mutable struct Sampler{P, V, Pol <: Policy, T1 <: AbstractSpace, T2 <: AbstractSpace}
    mdp::P
    agent::PolicyParams{Pol, T2} # The agent
    adversary = nothing # The adversary
    s::V = rand(initialstate(mdp)) # Current State
    S::T1 = state_space(initial_observation(mdp, s)) # State space
    svec::AbstractArray = tovec(initial_observation(mdp, s), S) # Current observation
    max_steps::Int = 100
    required_columns::Array{Symbol} = !isnothing(adversary) ? [:x] : []
    γ::Float32 = discount(mdp)
    λ::Float32 = NaN32
    episode_length::Int64 = 0
    episode_checker::Function = (data, start, stop) -> true

    # Parameters for cost constraints
    Vc::Union{ContinuousNetwork, Nothing} = nothing
    λcost::Float32 = NaN32
end

Sampler(mdp, π::T; kwargs...) where {T <: Policy} = Sampler(;mdp=mdp, agent=PolicyParams(π), kwargs...)
Sampler(mdp, agent::T; kwargs...) where {T <: PolicyParams} = Sampler(;mdp=mdp, agent=agent, kwargs...)

# Construct a vector of samplers from a vector of mdps
Sampler(mdps::AbstractVector, π::T; kwargs...) where {T <: Policy} = [Sampler(mdps[i], agent; kwargs...) for i in 1:length(mdps)]
Sampler(mdps::AbstractVector, agent::T; kwargs...) where {T <: PolicyParams} = [Sampler(mdps[i], agent; kwargs...) for i in 1:length(mdps)]

function reset_sampler!(sampler::Sampler)
    sampler.s = rand(initialstate(sampler.mdp))
    sampler.svec = tovec(initial_observation(sampler.mdp, sampler.s), sampler.S)
    sampler.episode_length = 0
end

function initial_observation(mdp, s)
    if mdp isa POMDP
        return convert_o(AbstractArray, rand(initialobs(mdp, s)), mdp)
    else
        return convert_s(AbstractArray, s, mdp)
    end
end

function terminate_episode!(sampler::Sampler, data, j)
    haskey(data, :episode_end) && (data[:episode_end][1,j] = true)
    ep = j - sampler.episode_length + 1 : j
    haskey(data, :advantage) && fill_gae!(data, ep, sampler.agent.π, sampler.λ, sampler.γ)
    haskey(data, :return) && fill_returns!(data, ep, sampler.γ)

    # Dealing with cost constraints
    haskey(data, :cost_advantage) && fill_gae!(data, ep, sampler.Vc, sampler.λ, sampler.γ, source=:cost, target=:cost_advantage)
    haskey(data, :cost_return) && fill_returns!(data, ep, sampler.γ, source=:cost, target=:cost_return)

    reset_sampler!(sampler)
end

function step!(data, j::Int, sampler::Sampler; explore=false, i=0)
    a, logprob = explore ? exploration(sampler.agent.π_explore, sampler.svec, π_on=sampler.agent.π, i=i) : (action(sampler.agent.π, sampler.svec), NaN)
    (a isa AbstractArray || a isa Tuple) && length(a) == 1 && (a = a[1])

    # This implements the ability to get cost information from safety gym
    info = Dict()
    kwargs = haskey(data, :cost) ? (info=info,) : ()

    args = (a,)
    if !isnothing(sampler.adversary)
        x, xlogprob = explore ? exploration(sampler.adversary.π_explore, sampler.svec, π_on=sampler.adversary.π, i=i) : (action(sampler.adversary.π, sampler.svec), NaN)
        (x isa AbstractArray || x isa Tuple) && length(x) == 1 && (x = x[1]) # disturbances always come out as an array
        data[:x][:, j:j] .= tovec(x, sampler.adversary.space)
        haskey(data, :xlogprob) && (data[:xlogprob][:, j] .= xlogprob)
        args = (a, x)
    end

    if sampler.mdp isa POMDP
        sp, o, r = gen(sampler.mdp, sampler.s, args...; kwargs...)
        spvec = convert_o(AbstractArray, o, sampler.mdp)
    else
        sp, r = gen(sampler.mdp, sampler.s, args...; kwargs...)
        spvec = convert_s(AbstractArray, sp, sampler.mdp)
    end
    spvec = tovec(spvec, sampler.S)
    done = isterminal(sampler.mdp, sp)

    # Save the tuple
    bslice(data[:s], j:j) .= sampler.svec
    data[:a][:, j:j] .= tovec(a, sampler.agent.space)
    bslice(data[:sp], j:j) .= spvec
    data[:r][1, j] = r
    data[:done][1, j] = done

    # Handle optional data storage
    haskey(data, :logprob) && (data[:logprob][:, j] .= logprob)
    haskey(data, :t) && (data[:t][1, j] = sampler.episode_length + 1)
    haskey(data, :i) && (data[:i][1, j] = i+1)
    haskey(data, :cost) && (data[:cost][1,j] = info["cost"])
    haskey(data, :fail) && (data[:fail][1,j] = extra_functions["isfailure"](sampler.mdp, sp)) #TODO Changed this to "s" instead of "sp" for the continuum world

    # Cut the episode short if needed
    sampler.episode_length += 1
    if done || sampler.episode_length >= sampler.max_steps
        terminate_episode!(sampler, data, j)
    else
        sampler.s = sp
        sampler.svec = spvec
    end
end

function steps!(sampler::Sampler; Nsteps = 1, explore=false, i=0, reset=false, return_episodes=false, return_at_episode_end=false)
    data = mdp_data(sampler.S, sampler.agent.space, Nsteps, return_episodes ? [sampler.required_columns..., :episode_end] : sampler.required_columns)
    for j=1:Nsteps
        step!(data, j, sampler, explore=explore, i=i + (j-1))
        if return_at_episode_end && sampler.episode_length == 0
            trim!(data, j)
            break
        end
    end
    reset && terminate_episode!(sampler, data, Nsteps)
    return_episodes ? (data, episodes(data)) : data
end

function steps!(samplers::Vector{T}; Nsteps = 1, explore = false, i = 0, reset = false, return_episodes = false) where {T<:Sampler}
    data = mdp_data(samplers[1].S, samplers[1].agent.space, Nsteps*length(samplers), return_episodes ? [samplers[1].required_columns..., :episode_end] : samplers[1].required_columns)
    j = 1
    for s=1:Nsteps
        for sampler in samplers
            step!(data, j, sampler, explore = explore, i = i + (j-1))
            j += 1
        end
    end
    reset && terminate_episode!(sampler, data, Nsteps)
    return_episodes ? (data, episodes(data)) : data
end

function episodes!(sampler::Sampler; Neps=1, explore=false, i=0, return_episodes=false)
    reset_sampler!(sampler)
    data = mdp_data(sampler.S, sampler.agent.space,
     Neps*sampler.max_steps, sampler.required_columns)
    episode_starts = zeros(Int, Neps)
    episode_ends = zeros(Int, Neps)
    j, k = 0, 1
    while k <= Neps
        episode_starts[k] = j+1
        while true
            j = j+1
            step!(data, j, sampler, explore=explore, i=i + (i-1))
            if sampler.episode_length == 0
                episode_ends[k] = j
                sampler.episode_checker(data, episode_starts[k], j) ? (k = k+1) : (j = episode_starts[k])
                break
            end
        end
    end
    trim!(data, j)
    return_episodes ? (data, zip(episode_starts, episode_ends)) : data
end

function fillto!(b::ExperienceBuffer, s::Union{Sampler, Vector{T}}, N::Int; i=1, explore=false) where {T <: Sampler}
    Nfill = max(0, N - length(b))
    Nfill > 0 && push!(b, steps!(s, i=i, Nsteps=Nfill, explore=explore))
    Nfill
end

## metric

# Recover multiple metrics from a single sampler
metrics_by_key(data, start, stop; keys) = [sum(data[key][1,start:stop]) for key in keys]

function metrics_by_key(s::Sampler; keys, Neps=100, kwargs...)
    data = episodes!(s, Neps=Neps; kwargs...)
    [sum(data[key]) / Neps for key in keys]
end

# recover a single metric
metric_by_key(data, start, stop; key) = metrics_by_key(data, start, stop; keys=[key])[1]

metric_by_key(s::Sampler; key, Neps=100, kwargs...) = metrics_by_key(s; keys=[key], Neps=Neps, kwargs...)[1]

# Get the undiscounted return
undiscounted_return(data, start, stop) = metric_by_key(data, start, stop; key=:r)
undiscounted_return(s::Sampler; Neps=100, kwargs...) = metric_by_key(s; Neps=Neps, key=:r, kwargs...)

## Discounted returns
function discounted_return(data, start, stop, γ)
    r = 0f0
    for i in reverse(start:stop)
        r = data[:r][1, i] + γ*r
    end
    r
end

function discounted_return(s::Sampler; Neps=100, kwargs...)
    data, episodes = episodes!(s, Neps=Neps, return_episodes=true; kwargs...)
    mean([discounted_return(data, e..., discount(s.mdp)) for e in episodes])
end

## Failures
failure(data, start, stop; threshold = 0.) = undiscounted_return(data, start, stop) < threshold

function failure(s::Sampler; threshold=0., Neps=100, kwargs...)
    data, episodes = episodes!(s, Neps = Neps, return_episodes = true; kwargs...)
    mean([failure(data, e..., threshold = threshold) for e in episodes])
end

## Generalized Advantage Estimation
function fill_gae!(d::ExperienceBuffer, V, λ::Float32, γ::Float32)
    eps = episodes(d)
    for ep in eps
        fill_gae!(d, ep, V, λ, γ)
    end
end

function fill_gae!(d, episode_range, V, λ::Float32, γ::Float32; source = :r, target = :advantage)
    A, c = 0f0, λ*γ
    nd = ndims(d[:s])
    for i in reverse(episode_range)
        Vsp = value(V, bslice(d[:sp], i:i))
        Vs = value(V, bslice(d[:s], i:i))
        @assert length(Vs) == 1
        A = c*A + d[source][1,i] + (1.f0 - d[:done][1,i])*γ*Vsp[1] - Vs[1]
        @assert !isnan(A)
        d[target][:, i] .= A
    end
end

function fill_returns!(data, episode_range, γ::Float32; source=:r, target=:return)
    r = 0f0
    for i in reverse(episode_range)
        r = data[source][1, i] + γ*r
        data[target][:, i] .= r
    end
end

# Utils
function trim!(data::Dict{Symbol, Array}, N)
    for k in keys(data)
        data[k] = bslice(data[k], 1:N)
    end
    data
end


module ExportNURBS

using BasicBSpline
using Luxor
using IntervalSets
using Colors
using ColorVectorSpace

export save_png, save_svg, save_pov

include("_common.jl")
include("_luxor.jl")
include("_povray.jl")

end # module


using Test, WAVI

@testset "Melting" begin
    @testset "test binary input file melt rate construction and update" begin

        @info "Testing binary input file construction and update"
        grid = Grid()
        bed_elevation = -900 .* ones(grid.nx, grid.ny) #low bed so we can make it float everywhere
        initial_conditions = InitialConditions(initial_thickness = 100.0*ones(grid.nx, grid.ny)) #initialize to 100m thickn

        #write a binary file
        m = ones(grid.nx, grid.ny)
        #m .= hton.(m)
        filename =  "melt_test_file.bin"
        mfileID =  open(filename,"w")
          write(mfileID, m[:,:])
        close(mfileID)

        #create a melt rate model from this input
        binfile_melt_rate = BinfileMeltRate(input_filename =filename)
        @test binfile_melt_rate isa BinfileMeltRate

        #create model
        model = Model(grid = grid,
                      bed_elevation = bed_elevation,
                      initial_conditions = initial_conditions,
                      melt_rate = binfile_melt_rate,
                      solver_params = SolverParams(maxiter_picard = 1))
        @test model.melt_rate isa BinfileMeltRate

        #check that melt rate read OK
        update_state!(model)
        @test all(model.fields.gh.basal_melt .== 1)

        #change the binary file and check we can update
        m = 2 .* ones(grid.nx, grid.ny)
        #m .= hton.(m)
        mfileID =  open(filename,"w")
          write(mfileID, m[:,:])
        close(mfileID)
        update_state!(model)
        @test all(model.fields.gh.basal_melt .== 2)

        #change the binary file to one of the wrong size, and check we return an error when updating
        m = 2 .* ones(grid.nx+1, grid.ny+1)
        m .= hton.(m)
        mfileID =  open(filename,"w")
          write(mfileID, m[:,:])
        close(mfileID)
        @test_throws DimensionMismatch update_state!(model)

        #remove the file we just made
        rm(filename)

    end

end


abstract type Data end

struct IncludeExcludePairs{V}
    mode::Symbol
    point_ids::Vector{V}
    include_pairs::Matrix{V}
end
function IncludeExcludePairs(V)
    IncludeExcludePairs(:undef, V[], Matrix{V}(undef,0,0))
end

struct NetworkData{T,V} <: Data
    coords::Tuple{Vector{V},Vector{V},Vector{T}}
    fp::Vector{V}
    source_map::Matrix{T}
    ground_map::Matrix{T}
end

struct RasterMeta
    ncols::Int
    nrows::Int
    xllcorner::Float64
    yllcorner::Float64
    cellsize::Float64
    nodata::Float64
    file_type::Int
    crs::String
    affine_map::AffineMap
end
function RasterMeta()
    RasterMeta(0,0,0,0,0,0,0,"",AffineMap([0 0; 0 0], [0, 0]))
end

struct RasData{T,V} <: Data
    cellmap::Matrix{T}
    polymap::Matrix{V}
    source_map::Matrix{T}
    ground_map::Matrix{T}
    points_rc::Tuple{Vector{V},Vector{V},Vector{V}}
    strengths::Matrix{T}
    included_pairs::IncludeExcludePairs{V}
    hbmeta::RasterMeta
end

function load_graph(V, gpath::String, ::Type{T}) where {T}
    g = readdlm(gpath, T)
    i = zeros(V, size(g, 1))
    j = zeros(V, size(g, 1))
    v = zeros(T, size(g, 1))
    for iter = 1:size(g, 1)
        i[iter] = g[iter,1] + 1
        j[iter] = g[iter,2] + 1
        v[iter] = g[iter,3]
    end
    i,j,v
end

read_focal_points(V, path::String) = V.(vec(readdlm(path)) .+ 1)

function read_point_strengths(T, path::String, inc = true)
    a = readdlm(path, T)
    if inc
        @. a[:,1] = a[:,1] + 1
    end
    a
end

function read_cellmap(habitat_file::String, is_res::Bool, ::Type{T}) where {T}

    if is_geotiff(habitat_file)
        cell_map, rastermeta = read_geotiff(T, habitat_file)
    else
        cell_map, rastermeta = read_ascii_grid(T, habitat_file)
    end

    gmap = similar(cell_map)
    ind = findall(x -> x == -9999, cell_map)
    if is_res
        if count(x -> x == 0, cell_map) > 0
            throw("Error: zero resistance values are not currently supported for habitat maps. Use a short-circuit region file instead.")
        else
            for i in eachindex(cell_map)
                gmap[i] = 1 ./ cell_map[i]
            end
            gmap[ind] .= 0
        end
    else
        copyto!(gmap, cell_map)
        gmap[ind] .= 0
    end
    gmap, rastermeta
end

function _ascii_grid_reader(T, file)
    f = endswith(file, ".gz") ? GZip.open(file, "r") : open(file, "r")
    rastermeta = _ascii_grid_read_header(file, f)
    c = Matrix{T}(undef,0,0)
    ss = 6
    if rastermeta.nodata == -Inf
        ss = 5
    end
    try
        c = readdlm(f, T; skipstart = ss)
    catch
        seek(f, 0)
        try
            d = readdlm(f; skipstart = ss)
            d = d[:, 1:end-1]
            c = map(T, d)
        catch
            error("Failed to read habitat map. There may be errors in your file.")
        end
    end
    map!(x -> x == rastermeta.nodata ? -9999. : x , c, c)
    close(f)
    c, rastermeta
end

function _ascii_grid_read_header(habitat_file, f)
    file_type = _guess_file_type(habitat_file, f)
    ncols = parse(Int, split(readline(f))[2])
    nrows = parse(Int, split(readline(f))[2])
    xllcorner = parse(Float64, split(readline(f))[2])
    yllcorner = parse(Float64, split(readline(f))[2])
    cellsize = parse(Float64, split(readline(f))[2])
    nodata = -Inf
    s = split(readline(f))
    if occursin("NODATA", s[1]) || occursin("nodata", s[1])
        nodata = parse(Float64, s[2])
    end
    seek(f, 0)
    RasterMeta(ncols, nrows, xllcorner, yllcorner, cellsize, nodata, file_type ,"", AffineMap([0 0; 0 0], [0, 0]))
end

#=function _guess_file_type(filename, f)
    s = readline(f)
    seek(f, 0)

    if startswith(s, "min")
        return PAIRS_AAGRID
    elseif startswith(s, "mode")
        return PAIRS_LIST
    elseif occursin(".asc", filename)
        return AAGRID
    elseif endswith(filename, ".txt")
        return TXTLIST
    else
        throw("Check file format")
    end

end=#

function is_geotiff(filename)
    endswith(filename,"tif") | endswith(filename,"tiff")
end

function is_aagrid(filename)
    endswith(filename,"asc")
end

function read_geotiff(T, f)
    gt=GeoArrays.read(f)
    ncols = size(gt,1)
    nrows = size(gt,2)
    xll = gt.f.translation[1]
    yll= gt.f.translation[2]-gt.f.linear[1,1]-(gt.f.linear[1,1]*(size(gt,2)-1))
    cellsize = gt.f.linear[1,1]
    crs = gt.crs
    affine_map = gt.f
    file_type = FILE_TYPE_GEOTIFF
    nodata = -Inf
    Array{T}(permutedims(gt.A[:,:,1])), RasterMeta(ncols, nrows, xll, yll, cellsize, nodata, file_type, crs, affine_map)
end

function read_ascii_grid(T, f)
    gt=GeoArrays.read(f)
    ncols = size(gt,1)
    nrows = size(gt,2)
    xll = gt.f.translation[1]
    yll= gt.f.translation[2]-gt.f.linear[1,1]-(gt.f.linear[1,1]*(size(gt,2)-1))
    cellsize = gt.f.linear[1,1]
    crs = gt.crs
    affine_map = gt.f
    file_type = FILE_TYPE_AAGRID
    nodata = -Inf
    Array{T}(permutedims(gt.A[:,:,1])), RasterMeta(ncols, nrows, xll, yll, cellsize, nodata, file_type, crs, affine_map)
end

function _guess_file_type(filename, f)

    if is_geotiff(filename)
        filetype = FILE_TYPE_GEOTIFF
    elseif is_aagrid(filename)
        filetype = FILE_TYPE_AAGRID
    else
        hdr = readline(f)
        seek(f, 0)

        if startswith(hdr, FILE_HDR_NPY)
            filetype = FILE_TYPE_NPY
        elseif startswith(lowercase(hdr), FILE_HDR_AAGRID)
            filetype = FILE_TYPE_AAGRID
        elseif startswith(hdr, FILE_HDR_INCL_PAIRS_AAGRID)
            filetype = FILE_TYPE_INCL_PAIRS_AAGRID
        elseif startswith(hdr, FILE_HDR_INCL_PAIRS)
            filetype = FILE_TYPE_INCL_PAIRS
        else
            filetype = FILE_TYPE_TXTLIST
        end
    end

    return filetype
end

function read_polymap(T, file::String, habitatmeta;
                            nodata_as = 0, resample = true)

    if is_geotiff(file)
        polymap, rastermeta = read_geotiff(T, file)
    else
        polymap, rastermeta = read_ascii_grid(T, file)
    end

    ind = findall(x -> x == rastermeta.nodata, polymap)
    if nodata_as != -1
        polymap[ind] .= nodata_as
    end

    if(is_geotiff(file) || is_aagrid(file))
        if rastermeta.yllcorner != habitatmeta.yllcorner
            cswarn("yllcorner is not the same")
        elseif rastermeta.xllcorner != habitatmeta.xllcorner
            cswarn("xllcorner is not the same")
        elseif rastermeta.ncols != habitatmeta.ncols
            cswarn("ncols is not the same")
        elseif rastermeta.nrows != habitatmeta.nrows
            cswarn("nrows is not the same")
        elseif rastermeta.cellsize != habitatmeta.cellsize
            cswarn("cellsize is not the same")
        elseif rastermeta.crs != habitatmeta.crs
            cswarn("reference system is not the same")
        end
    else
        if rastermeta.cellsize != habitatmeta.cellsize
            cswarn("cellsize is not the same")
        elseif rastermeta.ncols != habitatmeta.ncols
            cswarn("ncols is not the same")
        elseif rastermeta.nrows != habitatmeta.nrows
            cswarn("nrows is not the same")
        elseif rastermeta.yllcorner != habitatmeta.yllcorner
            cswarn("yllcorner is not the same")
        elseif rastermeta.xllcorner != habitatmeta.xllcorner
            cswarn("xllcorner is not the same")
        end
    end

    polymap
end

function read_point_map(V, file, habitatmeta)

    # Advanced mode
    if file == "none"
        return (V[], V[], V[])
    end

    f = endswith(file, ".gz") ? GZip.open(file, "r") : open(file, "r")
    filetype = _guess_file_type(file, f)
    _points_rc = filetype == FILE_TYPE_TXTLIST ? readdlm(file) :
                        read_polymap(V, file, habitatmeta)

    i = V[]
    j = V[]
    v = V[]
    if filetype == FILE_TYPE_TXTLIST
        I = _points_rc[:,2]
        J = _points_rc[:,3]
        v = _points_rc[:,1]
        i  = ceil.(V, habitatmeta.nrows .- (J .- habitatmeta.yllcorner) ./ habitatmeta.cellsize)
        j = ceil.(V, (I .- habitatmeta.xllcorner) ./ habitatmeta.cellsize)
    else
        _I = findall(!iszero, _points_rc)
        (i,j,v) =  (getindex.(_I, 1), getindex.(_I, 2), _points_rc[_I])
    end

    ind = findall(x -> x < 0, v)

    # Get rid of negative resistances
    for index in ind
        deleteat!(i, index)
        deleteat!(j, index)
        deleteat!(v, index)
    end

    # Sort them
    idx = sortperm(v)
    i = i[idx]
    j = j[idx]
    v = v[idx]

    if (minimum(i) < 0) || (minimum(j) < 0) ||
            (maximum(i) > (habitatmeta.nrows)) ||
            (maximum(j) > (habitatmeta.ncols))
         throw("At least one focal node location falls outside of habitat map")
    end

    if size(unique(v),1) < 2
        throw("Less than two valid focal nodes found. Please check focal node location file.")
    end

    close(f)
    V.(i), V.(j), V.(v)
end

function read_source_and_ground_maps(T, V, source_file, ground_file, habitatmeta,
                                        is_res)

    ground_map = Matrix{T}(undef,0,0)
    source_map = Matrix{T}(undef,0,0)

    f = endswith(ground_file, "gz") ? Gzip.open(ground_file, "r") : open(ground_file, "r")
    filetype = _guess_file_type(ground_file, f)

    if filetype == FILE_TYPE_AAGRID || filetype == FILE_TYPE_GEOTIFF
        ground_map = read_polymap(T, ground_file, habitatmeta; nodata_as = -1)
        ground_map = map(T, ground_map)
    else
        rc = readdlm(ground_file, V)
        ground_map = -9999 * ones(T, habitatmeta.nrows, habitatmeta.ncols)
        ground_map[rc[:,2], rc[:,3]] = rc[:,1]
    end
    close(f)

    f = endswith(source_file, "gz") ? Gzip.open(source_file, "r") : open(source_file, "r")
    filetype = _guess_file_type(source_file, f)

    if filetype == FILE_TYPE_AAGRID || filetype == FILE_TYPE_GEOTIFF
        source_map = read_polymap(T, source_file, habitatmeta)
        source_map = map(T, source_map)
    else
        rc = readdlm(source_file, V)
        source_map = -9999 * ones(T, habitatmeta.nrows, habitatmeta.ncols)
        source_map[rc[:,2], rc[:,3]] = rc[:,1]
    end
    close(f)

    if is_res
        ind = findall(x -> x == -9999, ground_map)
        ground_map = 1 ./ ground_map
        ground_map[ind] .= 0
    else
        ind = findall(x -> x == -9999, ground_map)
        ground_map[ind] .= 0
    end

    source_map, ground_map
end

function read_included_pairs(V, filename)

    f = endswith(filename, "gz") ? Gzip.open(filename, "r") : open(filename, "r")
    filetype = _guess_file_type(filename, f)
    minval = 0
    maxval = 0
    mode = :undef

    if filetype == FILE_TYPE_INCL_PAIRS_AAGRID
        open(filename, "r") do fV
            minval = parse(Float64, split(readline(f))[2])
            maxval = parse(Float64, split(readline(f))[2])
        end
        included_pairs = readdlm(filename, skipstart=2)
        point_ids = V.(included_pairs[:,1])
        deleteat!(point_ids, 1)
        included_pairs = included_pairs[2:end, 2:end]
        map!(x -> x > maxval ? 0 : x, included_pairs, included_pairs)
        idx = findall(x -> x >= minval, included_pairs)
        mode = :include
        bin = map(x -> x >= minval ? V(1) : V(0), included_pairs)
        return IncludeExcludePairs(mode, point_ids, bin)
    elseif filetype == FILE_TYPE_INCL_PAIRS
        open(filename, "r") do f
            mode = Symbol(split(readline(f))[2])
        end
        included_pairs = readdlm(filename, skipstart = 1)
        if size(included_pairs, 1) == 1
            pl = zeros(V, 1,2)
            pl[1,:] = included_pairs
            included_pairs = pl
        end
        point_ids = V.(sort!(unique(included_pairs)))
        idx = findall(x -> x == 0 , point_ids)
        if length(idx) > 0
            deleteat!(point_ids, idx)
            cswarn("Code to include pairs is activated, some entries did not match with focal node file. Some focal nodes may have been dropped")
        end

        mat = zeros(V, size(point_ids, 1), size(point_ids, 1))

        for i = 1:size(included_pairs, 1)
            idx1 = findfirst(x -> x == included_pairs[i, 1], point_ids)
            idx2 = findfirst(x -> x == included_pairs[i, 2], point_ids)
            if idx1 != nothing && idx2 != nothing
                mat[idx1,idx2] = 1
                mat[idx2,idx1] = 1
            end
        end

        return IncludeExcludePairs(mode, point_ids, mat)
    end

end

function get_network_data(T, V, cfg)::NetworkData{T,V}

    hab_is_res = cfg["habitat_map_is_resistances"] in TRUELIST
    hab_file = cfg["habitat_file"]
    fp_file = cfg["point_file"]
    source_file = cfg["source_file"]
    ground_file = cfg["ground_file"]

    is_pairwise = cfg["scenario"] in PAIRWISE

    i,j,v = load_graph(V, hab_file, T)
    if hab_is_res
        v = 1 ./ v
    end

    if is_pairwise
        fp = read_focal_points(V, fp_file)
    else
        fp = V[]
    end

    if !is_pairwise
        source_map = read_point_strengths(T, source_file)
        ground_map = read_point_strengths(T, ground_file)
    else
        source_map = Matrix{T}(undef,0,0)
        ground_map = Matrix{T}(undef,0,0)
    end

    NetworkData((i,j,v), fp, source_map, ground_map)
end

function load_raster_data(T, V, cfg)::RasData{T,V}

    # Habitat file
    hab_file = cfg["habitat_file"]
    hab_is_res = cfg["habitat_map_is_resistances"] in TRUELIST

    # Polygons
    use_polygons = cfg["use_polygons"] in TRUELIST
    polygon_file = cfg["polygon_file"]

    # Mask file
    use_mask = cfg["use_mask"] in TRUELIST
    mask_file = cfg["mask_file"]

    # Point file
    point_file = cfg["point_file"]

    # Variable source strengths
    use_var_source = cfg["use_variable_source_strengths"] in TRUELIST
    var_source_file = cfg["variable_source_file"]

    # Included Pairs
    use_inc_pairs = cfg["use_included_pairs"] in TRUELIST
    inc_pairs_file = cfg["included_pairs_file"]

    # Advanced mode
    is_pairwise = cfg["scenario"] in PAIRWISE
    is_advanced = cfg["scenario"] in ADVANCED
    source_file = cfg["source_file"]
    ground_file = cfg["ground_file"]
    ground_is_res = cfg["ground_file_is_resistances"] in TRUELIST

    csinfo("Reading maps")

    # Read cell map
    cellmap, hbmeta = read_cellmap(hab_file, hab_is_res, T)
    c = count(x -> x > 0, cellmap)
    csinfo("Resistance/Conductance map has $c nodes")

    # Read polymap
    if use_polygons
        polymap = read_polymap(V, polygon_file, hbmeta)
    else
        polymap = Matrix{V}(undef,0,0)
    end

    # Read and update cellmap with mask file
    if use_mask
        update!(cellmap, mask_file, hbmeta)
        sum(cellmap) == 0 && throw("Mask file deleted everything!")
    end

    # Read point file
    if !is_advanced
        points_rc = read_point_map(V, point_file, hbmeta)
    else
        points_rc = (V[], V[], V[])
    end

    # Advanced mode reading
    if is_advanced
        source_map, ground_map =
        read_source_and_ground_maps(T, V, source_file, ground_file,
                                    hbmeta, ground_is_res)
    else
        source_map, ground_map = Matrix{T}(undef,0,0), Matrix{T}(undef,0,0)
    end

    # Included Pairs
    if use_inc_pairs
        included_pairs = read_included_pairs(V, inc_pairs_file)
    else
        included_pairs = IncludeExcludePairs(V)
    end

    # Variable source strengths
    if use_var_source
        strengths = read_point_strengths(T, var_source_file)
    else
        strengths = Matrix{T}(undef, 0,0)
    end

    RasData(cellmap, polymap, source_map, ground_map, points_rc, strengths,
                    included_pairs, hbmeta)
end

function update!(cellmap::Matrix{T}, m::String, hbmeta) where {T}
    mask = read_polymap(T, m, hbmeta)
    map!(x -> x > 0 ? 1 : 0, mask, mask)
    cellmap .= cellmap .* mask
end


using PSCOPF

using Test
using Dates

@testset verbose=true "test_network" begin

    @testset "test_empty_network" begin
        network = PSCOPF.Networks.Network()
        @test length(network.buses) == 0
        @test length(network.branches) == 0
        @test length(network.generators) == 0
    end

    @testset "test_network_read" begin
        data_path = joinpath(@__DIR__, "..", "..", "data", "2buses_usecase")
        network = PSCOPF.Data.pscopfdata2network(data_path)
        @test length(network.buses) == 2
        @test length(network.branches) == 1
        @test length(network.generators) == 4

        @test network.generators["alta"].type == PSCOPF.Networks.PILOTABLE
        @test network.generators["park_city"].type == PSCOPF.Networks.PILOTABLE
        @test network.generators["wind_1"].type == PSCOPF.Networks.LIMITABLE
        @test network.generators["wind_2"].type == PSCOPF.Networks.LIMITABLE

        @test network.generators["alta"].id == "alta"
        @test network.generators["alta"].bus_id == "poste_1_0"
        @test network.generators["alta"].p_min == 10.
        @test network.generators["alta"].p_max == 200.
        @test network.generators["alta"].start_cost == 45000.
        @test network.generators["alta"].prop_cost == 30.
        @test network.generators["alta"].dmo == Dates.Second(14400)
        @test network.generators["alta"].dp == Dates.Second(900)

        @test network.generators["park_city"].id == "park_city"
        @test network.generators["park_city"].bus_id == "poste_1_0"
        @test network.generators["park_city"].p_min == 10.
        @test network.generators["park_city"].p_max == 100.
        @test network.generators["park_city"].start_cost == 12000.
        @test network.generators["park_city"].prop_cost == 120.
        @test network.generators["park_city"].dmo == Dates.Second(3600)
        @test network.generators["park_city"].dp == Dates.Second(900)

        @test network.generators["wind_1"].id == "wind_1"
        @test network.generators["wind_1"].bus_id == "poste_1_0"
        @test network.generators["wind_1"].p_min == 0.
        @test network.generators["wind_1"].p_max == 50.
        @test network.generators["wind_1"].start_cost == 0.
        @test network.generators["wind_1"].prop_cost == 50.
        @test network.generators["wind_1"].dmo == Dates.Second(0)
        @test network.generators["wind_1"].dp == Dates.Second(0)

        @test network.generators["wind_2"].id == "wind_2"
        @test network.generators["wind_2"].bus_id == "poste_2_0"
        @test network.generators["wind_2"].p_min == 0.
        @test network.generators["wind_2"].p_max == 20.
        @test network.generators["wind_2"].start_cost == 0.
        @test network.generators["wind_2"].prop_cost == 50.
        @test network.generators["wind_2"].dmo == Dates.Second(0)
        @test network.generators["wind_2"].dp == Dates.Second(0)

    end

end


export CircularArrayBuffer, capacity, isfull, consecutive_view, select_frame

mutable struct CircularArrayBuffer{T,N} <: AbstractArray{T,N}
    buffer::Array{T,N}
    first::Int
    length::Int
    step_size::Int

    function CircularArrayBuffer{T}(d::Integer...) where {T}
        N = length(d)
        N > 0 || throw(ArgumentError("dimension must be greater than 0"))
        new{T,N}(Array{T}(undef, d...), 1, 0, N == 1 ? 1 : *(d[1:end-1]...))
    end
end

Base.IndexStyle(::CircularArrayBuffer) = IndexLinear()
Base.size(cb::CircularArrayBuffer{T,N}) where {T,N} = (size(cb.buffer)[1:N-1]..., cb.length)
Base.getindex(cb::CircularArrayBuffer{T, N}, i::Int) where {T, N} = getindex(cb.buffer, _buffer_index(cb, i))
Base.setindex!(cb::CircularArrayBuffer{T, N}, v, i::Int) where {T, N} = setindex!(cb.buffer, v, _buffer_index(cb, i))
capacity(cb::CircularArrayBuffer{T, N}) where {T, N} = size(cb.buffer, N)
isfull(cb::CircularArrayBuffer) = cb.length == capacity(cb)
Base.isempty(cb::CircularArrayBuffer) = cb.length == 0

@inline function _buffer_index(cb::CircularArrayBuffer, i::Int)
    ind = (cb.first - 1) * cb.step_size + i
    if ind > length(cb.buffer)
        ind - length(cb.buffer)
    else
        ind
    end
end

@inline function _buffer_frame(cb::CircularArrayBuffer, i::Int)
    n = capacity(cb)
    idx = cb.first + i - 1
    if idx > n
        idx - n
    else
        idx
    end
end

_buffer_frame(cb::CircularArrayBuffer, I::Vector{Int}) = map(i -> _buffer_frame(cb, i), I)

function Base.empty!(cb::CircularArrayBuffer)
    cb.first = 1
    cb.length = 0
    cb
end

function Base.push!(cb::CircularArrayBuffer{T,N}, data) where {T,N}
    # length(data) == cb.step_size || throw(ArgumentError("length of , $(cb.step_size) != $(length(data))"))
    if cb.length == capacity(cb)
        cb.first = (cb.first == capacity(cb) ? 1 : cb.first + 1)
    else
        cb.length += 1
    end
    selectdim(cb.buffer, N, _buffer_frame(cb, cb.length)) .= data
    cb
end

Base.push!(cb::CircularArrayBuffer{T, N}, data::CircularArrayBuffer{T, N}) where {T, N} = push!(cb, select_frame(data, data.length))

select_frame(cb::CircularArrayBuffer{T, 1}, i) where {T} = getindex(cb.buffer, _buffer_frame(cb, i))
select_frame(cb::CircularArrayBuffer{T, 2}, i) where {T} = view(cb.buffer, :, _buffer_frame(cb, i))
select_frame(cb::CircularArrayBuffer{T, 3}, i) where {T} = view(cb.buffer, :, :, _buffer_frame(cb, i))
select_frame(cb::CircularArrayBuffer{T, 4}, i) where {T} = view(cb.buffer, :, :, :, _buffer_frame(cb, i))

select_frame(cb::Array{T, 1}, i::Int) where {T} = getindex(cb, i)
select_frame(cb::Array{T, N}, i::Int) where {T, N} = selectdim(cb, N, i)

consecutive_view(cb::CircularArrayBuffer, inds::Vector{Int}) = select_frame(cb, inds)

consecutive_view(cb::CircularArrayBuffer{T,1}, inds::Vector{Int}, n::Int) where {T} = reshape(view(cb.buffer, [_buffer_frame(cb, i) for x in inds for i in x:x+n-1]), n, length(inds))
consecutive_view(cb::CircularArrayBuffer{T,2}, inds::Vector{Int}, n::Int) where {T} = reshape(view(cb.buffer, :, [_buffer_frame(cb, i) for x in inds for i in x:x+n-1]), size(cb.buffer, 1), n, length(inds))
consecutive_view(cb::CircularArrayBuffer{T,3}, inds::Vector{Int}, n::Int) where {T} = reshape(view(cb.buffer, :, :, [_buffer_frame(cb, i) for x in inds for i in x:x+n-1]), size(cb.buffer, 1), size(cb.buffer, 2), n, length(inds))
consecutive_view(cb::CircularArrayBuffer{T,4}, inds::Vector{Int}, n::Int) where {T} = reshape(view(cb.buffer, :, :, :, [_buffer_frame(cb, i) for x in inds for i in x:x+n-1]), size(cb.buffer, 1), size(cb.buffer, 2), size(cb.buffer, 3), n, length(inds))

consecutive_view(cb::CircularArrayBuffer{T,1}, inds::Vector{Int}, n::Int, c::Int) where {T} = reshape(view(cb.buffer, [_buffer_frame(cb, c) for x in inds for i in x:x+n-1 for c in i-c+1:i]), c, n, length(inds))
consecutive_view(cb::CircularArrayBuffer{T,2}, inds::Vector{Int}, n::Int, c::Int) where {T} = reshape(view(cb.buffer, :, [_buffer_frame(cb, c) for x in inds for i in x:x+n-1 for c in i-c+1:i]), size(cb.buffer, 1), c, n, length(inds))
consecutive_view(cb::CircularArrayBuffer{T,3}, inds::Vector{Int}, n::Int, c::Int) where {T} = reshape(view(cb.buffer, :, :, [_buffer_frame(cb, c) for x in inds for i in x:x+n-1 for c in i-c+1:i]), size(cb.buffer, 1), size(cb.buffer, 2), c, n, length(inds))
consecutive_view(cb::CircularArrayBuffer{T,4}, inds::Vector{Int}, n::Int, c::Int) where {T} = reshape(view(cb.buffer, :, :, :, [_buffer_frame(cb, c) for x in inds for i in x:x+n-1 for c in i-c+1:i]), size(cb.buffer, 1), size(cb.buffer, 2), size(cb.buffer, 3), c, n, length(inds))

consecutive_view(cb::CircularArrayBuffer, inds::Vector{Int}, n::Int, ::Nothing) = consecutive_view(cb, inds, n)
consecutive_view(cb::CircularArrayBuffer, inds::Vector{Int}, ::Nothing, ::Nothing) = consecutive_view(cb, inds)

Base.getindex(cb::CircularArrayBuffer{T, 2}, i1::Int, i2::Int) where T = cb.buffer[i1, _buffer_frame(cb, i2)]
Base.getindex(cb::CircularArrayBuffer{T, 3}, i1::Int, i2::Int, i3::Int) where T = cb.buffer[i1, i2, _buffer_frame(cb, i3)]
Base.getindex(cb::CircularArrayBuffer{T, 4}, i1::Int, i2::Int, i3::Int, i4::Int) where T = cb.buffer[i1, i2, i3, _buffer_frame(cb, i4)]


using RCall, CSV, DataFrames

########################################################################
#Chronic Infected (Figure 6)
########################################################################

Det_Homo = CSV.read("noStochFigure6Data.csv")
Det_Hetero = CSV.read("VirusFigure6_Det_Hetero.csv")
Stoch_Homo = CSV.read("StochFigure6Data.csv")

@rput Det_Homo
@rput Det_Hetero
@rput Stoch_Homo
R"""
library(ggplot2)
library(ggpubr)

 commonFigureOptions <- list(
   theme_pubr(border=TRUE),
   xlab("Time (hours)"),
   labs(fill=expression(paste("IFN", beta, " (nM)"))),
   theme_pubr(border=TRUE,base_size = 14),
   scale_x_continuous(expand=c(0,0)),
   scale_y_continuous(expand=c(0,0)),
   xlab(bquote("Interferon Degradation" ~ tau[7])),
   ylab(bquote("JAK/STAT Activity" ~ k[cat8])),
   theme(plot.title = element_text(hjust = 0.5),aspect.ratio = 1) )

p1 <- ggplot(Det_Homo, aes(kcat8, tau7, fill=IFN)) +
  geom_raster(aes(fill=IFN)) +
  ggtitle("ISD \n Deterministic \n Homogeneous") +
  scale_fill_distiller(palette = "Spectral",
   guide = guide_colorbar(frame.colour = "black",
   ticks.colour = "black")) +
  commonFigureOptions

p2 <- ggplot(Det_Hetero, aes(kcat8, tau7, fill=IFN)) +
  geom_raster(aes(fill=IFN)) +
  ggtitle("ISD \n Deterministic \n Heterogeneous") +
  scale_fill_distiller(palette = "Spectral",
   guide = guide_colorbar(frame.colour = "black",
   ticks.colour = "black")) +
  commonFigureOptions

Stoch_Homo$IFN[Stoch_Homo$IFN<0] <- 0

p3 <- ggplot(Stoch_Homo, aes(kcat8, tau7, fill=IFN)) +
  geom_raster(aes(fill=IFN)) +
  ggtitle("ISD \n Stochastic \n Homogeneous") +
  scale_fill_distiller(palette = "Spectral",
   guide = guide_colorbar(frame.colour = "black",
   ticks.colour = "black"),
   limits = c(0,max(Stoch_Homo$IFN))) +
  commonFigureOptions

figure <- ggarrange(p1, p2, p3,
                    labels = "AUTO",
                    legend = "right",
                    align = "hv",
                    ncol = 3, nrow = 1)

ggsave("../Figures/Figure6_All.pdf",width = 15,height=5)
"""

########################################################################
#Stochastic IFN (Figure 7)
########################################################################

R"""
library(tidyverse)
library(ggpubr)
library(scales)

df <- read.csv("VirusFigure7DataHomo50cells.csv")

#And that's when I discovered how to use tidyverse

df_tidy_median <- df %>%
  pivot_longer(c(Healthy,Infected,Dead),names_to = "CellState",values_to = "CellPercent") %>%
  filter(Percent==0.2) %>%
  group_by(Time,CellState) %>%
  summarise(medianPercentCells = median(CellPercent),
            q25 = quantile(CellPercent,0.25),
            q75 = quantile(CellPercent,0.75))

df_tidy_median$CellState = factor(df_tidy_median$CellState,ordered=TRUE,levels = c("Healthy","Infected","Dead"))

color_list <- c("#377EB8","#4DAF4A","#E41A1C")

p1 <- ggplot(df_tidy_median, aes(x=Time, y=medianPercentCells)) +
  geom_line(aes(x=Time, y=medianPercentCells, color=CellState)) +
  geom_ribbon(aes(ymin=q25, ymax=q75, fill=CellState),alpha=0.3) +
  scale_fill_manual(values=color_list) +
  scale_color_manual(values=color_list) +
  scale_x_continuous(breaks=seq(0, 48, 12)) +
  scale_y_continuous(labels = scales::percent) +
  xlab("Time (hours)") +
  ylab("Cell Percent") +
  theme_pubr(border=TRUE)

df_tidy_box <- df %>%
  filter(Time==48)

p2 <-ggplot(df_tidy_box, aes(x=as.factor(Percent), y=Dead)) +
  stat_boxplot(geom = "errorbar", width = 0.2,size=0.1) +
  geom_boxplot(fatten=1, fill = "#f6b9ba") +
  #geom_jitter(shape=16, position=position_jitter(0.2)) +
  scale_y_continuous(labels = scales::percent,limits=c(0,1)) +
  scale_x_discrete(labels = percent((0:10)/10)) +
  xlab(expression("IFN"*beta*" Producing Cells")) +
  ylab("Dead Cells") +
  theme_pubr(border=TRUE)

df <- read.csv("VirusFigure7DataHetero50cells.csv")

df_tidy_median <- df %>%
  pivot_longer(c(Healthy,Infected,Dead),names_to = "CellState",values_to = "CellPercent") %>%
  filter(Percent==0.2) %>%
  group_by(Time,CellState) %>%
  summarise(medianPercentCells = median(CellPercent),
            q25 = quantile(CellPercent,0.25),
            q75 = quantile(CellPercent,0.75))

df_tidy_median$CellState = factor(df_tidy_median$CellState,ordered=TRUE,levels = c("Healthy","Infected","Dead"))

color_list <- c("#377EB8","#4DAF4A","#E41A1C")

p3 <- ggplot(df_tidy_median, aes(x=Time, y=medianPercentCells)) +
  geom_line(aes(x=Time, y=medianPercentCells, color=CellState)) +
  geom_ribbon(aes(ymin=q25, ymax=q75, fill=CellState),alpha=0.3) +
  scale_fill_manual(values=color_list) +
  scale_color_manual(values=color_list) +
  scale_x_continuous(breaks=seq(0, 48, 12)) +
  scale_y_continuous(labels = scales::percent) +
  xlab("Time (hours)") +
  ylab("Cell Percent") +
  theme_pubr(border=TRUE)

df_tidy_box <- df %>%
  filter(Time==48)

p4 <-ggplot(df_tidy_box, aes(x=as.factor(Percent), y=Dead)) +
  stat_boxplot(geom = "errorbar", width = 0.2,size=0.1) +
  geom_boxplot(fatten=1, fill = "#f6b9ba") +
  scale_y_continuous(labels = scales::percent,limits=c(0,1)) +
  scale_x_discrete(labels = percent((0:10)/10)) +
  xlab(expression("IFN"*beta*" Producing Cells")) +
  ylab("Dead Cells") +
  theme_pubr(border=TRUE)

figure <- ggarrange(p1, p2, p3, p4,
                    labels = "AUTO",
                    common.legend = TRUE, legend = "right",
                    widths = c(3,6),
                    align = "hv",
                    ncol = 2, nrow = 2)

ggsave("../Figures/Figure7.pdf",width = 9,height=6,units="in")
"""


@testset "Independent Bounds" begin
    m = BabyPOMDP()
    planner = solve(AdaOPSSolver(), m)
    belief = WPFBelief([true], [1.0], 1, 1, AdaOPSTree(planner, initialstate(m)), true)
    b = IndependentBounds(0.0, -1e-5)
    @test bounds(b, m, belief, 90, true) == (0.0, -1e-5)
    b = IndependentBounds(0.0, -1e-5, consistency_fix_thresh=1e-5)
    @test bounds(b, m, belief, 90, true) == (0.0, 0.0)
    b = IndependentBounds(0.0, -1e-4, consistency_fix_thresh=1e-5)
    @test bounds(b, m, belief, 90, true) == (0.0, -1e-4)
end


function _precompile_()
  ccall(:jl_generating_output, Cint, ()) == 1 || return nothing
  precompile(
    Tuple{
      Core.kwftype(typeof(ITensors.eigen)),
      NamedTuple{
        (
          :ishermitian,
          :which_decomp,
          :tags,
          :maxdim,
          :mindim,
          :cutoff,
          :eigen_perturbation,
          :ortho,
        ),
        Tuple{Bool,Nothing,TagSet,Int64,Int64,Float64,Nothing,String},
      },
      typeof(eigen),
      ITensor{2},
      IndexSet{1,Index{Int64},Tuple{Index{Int64}}},
      IndexSet{1,Index{Int64},Tuple{Index{Int64}}},
    },
  )
  precompile(
    Tuple{
      Core.kwftype(typeof(ITensors.factorize)),
      NamedTuple{
        (:which_decomp, :tags, :maxdim, :mindim, :cutoff, :eigen_perturbation, :ortho),
        Tuple{Nothing,TagSet,Int64,Int64,Float64,Nothing,String},
      },
      typeof(factorize),
      ITensor{3},
      IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
    },
  )
  precompile(
    Tuple{
      Core.kwftype(typeof(ITensors.replacebond!)),
      NamedTuple{
        (:maxdim, :mindim, :cutoff, :eigen_perturbation, :ortho, :which_decomp),
        Tuple{Int64,Int64,Float64,Nothing,String,Nothing},
      },
      typeof(replacebond!),
      MPS,
      Int64,
      ITensor{3},
    },
  )
  precompile(
    Tuple{
      Core.kwftype(typeof(KrylovKit.eigsolve)),
      NamedTuple{
        (:ishermitian, :tol, :krylovdim, :maxiter),Tuple{Bool,Float64,Int64,Int64}
      },
      typeof(KrylovKit.eigsolve),
      ProjMPO,
      ITensor{3},
      Int64,
      Symbol,
    },
  )
  precompile(
    Tuple{
      Core.kwftype(typeof(KrylovKit.expand!)),
      NamedTuple{(:verbosity,),Tuple{Int64}},
      typeof(KrylovKit.expand!),
      KrylovKit.LanczosIterator{ProjMPO,ITensor{3},KrylovKit.ModifiedGramSchmidt2},
      KrylovKit.LanczosFactorization{ITensor{3},Float64},
    },
  )
  precompile(
    Tuple{
      Core.kwftype(typeof(NDTensors.eigen)),
      NamedTuple{
        (
          :ishermitian,
          :which_decomp,
          :tags,
          :maxdim,
          :mindim,
          :cutoff,
          :eigen_perturbation,
          :ortho,
        ),
        Tuple{Bool,Nothing,TagSet,Int64,Int64,Float64,Nothing,String},
      },
      typeof(eigen),
      LinearAlgebra.Hermitian{
        Float64,
        NDTensors.Tensor{
          Float64,
          2,
          NDTensors.Dense{Float64,Array{Float64,1}},
          IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
        },
      },
    },
  )
  precompile(
    Tuple{
      Core.kwftype(typeof(NDTensors.svd)),
      NamedTuple{(:alg,),Tuple{String}},
      typeof(svd),
      NDTensors.Tensor{
        Complex{Float64},
        2,
        NDTensors.Dense{Complex{Float64},Array{Complex{Float64},1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
    },
  )
  precompile(
    Tuple{
      Core.kwftype(typeof(NDTensors.svd)),
      NamedTuple{(:alg,),Tuple{String}},
      typeof(svd),
      NDTensors.Tensor{
        Float64,
        2,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
    },
  )
  precompile(Tuple{ProjMPO,ITensor{3}})
  precompile(Tuple{Type{MPO},AutoMPO,Array{Index{Int64},1}})
  precompile(Tuple{Type{TagSet},String})
  precompile(
    Tuple{
      typeof(*),
      Float64,
      NDTensors.Tensor{
        Float64,
        3,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
      },
    },
  )
  precompile(Tuple{typeof(*),ITensor{2},ITensor{3}})
  precompile(Tuple{typeof(*),ITensor{3},ITensor{1}})
  precompile(Tuple{typeof(*),ITensor{3},ITensor{2}})
  precompile(Tuple{typeof(*),ITensor{3},ITensor{3}})
  precompile(Tuple{typeof(*),ITensor{4},ITensor{1}})
  precompile(Tuple{typeof(*),ITensor{4},ITensor{4}})
  precompile(Tuple{typeof(==),Array{Index{Int64},1},Array{Index{Int64},1}})
  precompile(
    Tuple{
      typeof(==),
      IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
      IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
    },
  )
  precompile(Tuple{typeof(==),TagSet,TagSet})
  precompile(
    Tuple{
      typeof(Base.Broadcast.materialize!),
      ITensor{4},
      Base.Broadcast.Broadcasted{
        ITensors.ITensorStyle,Nothing,typeof(+),Tuple{ITensor{4},ITensor{4}}
      },
    },
  )
  precompile(Tuple{typeof(Base.afoldl),typeof(*),ITensor{4},ITensor{3}})
  precompile(
    Tuple{typeof(ITensors.setstore!),ITensor{2},NDTensors.Dense{Float64,Array{Float64,1}}}
  )
  precompile(
    Tuple{
      typeof(KrylovKit.orthogonalize!),ITensor{3},ITensor{3},KrylovKit.ModifiedGramSchmidt2
    },
  )
  precompile(
    Tuple{
      typeof(NDTensors.contract!!),
      NDTensors.Tensor{
        Complex{Float64},
        2,
        NDTensors.Dense{Complex{Float64},Array{Complex{Float64},1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
      NDTensors.Tensor{
        Complex{Float64},
        2,
        NDTensors.Dense{Complex{Float64},Array{Complex{Float64},1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
      NDTensors.Tensor{
        Complex{Float64},
        2,
        NDTensors.Dense{Complex{Float64},Array{Complex{Float64},1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
    },
  )
  precompile(
    Tuple{
      typeof(NDTensors.contract!!),
      NDTensors.Tensor{
        Complex{Float64},
        2,
        NDTensors.Dense{Complex{Float64},Array{Complex{Float64},1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
      NDTensors.Tensor{
        Complex{Float64},
        2,
        NDTensors.Dense{Complex{Float64},Array{Complex{Float64},1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
      NDTensors.Tensor{
        Number,
        2,
        NDTensors.Combiner,
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
    },
  )
  precompile(
    Tuple{
      typeof(NDTensors.contract!!),
      NDTensors.Tensor{
        Complex{Float64},
        2,
        NDTensors.Dense{Complex{Float64},Array{Complex{Float64},1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
      NDTensors.Tensor{
        Float64,
        2,
        NDTensors.Diag{Float64,Array{Float64,1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
      NDTensors.Tensor{
        Complex{Float64},
        2,
        NDTensors.Dense{Complex{Float64},Array{Complex{Float64},1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
    },
  )
  precompile(
    Tuple{
      typeof(NDTensors.contract!!),
      NDTensors.Tensor{
        Float64,0,NDTensors.Dense{Float64,Array{Float64,1}},IndexSet{0,Union{},Tuple{}}
      },
      Tuple{},
      NDTensors.Tensor{
        Float64,
        3,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64,Int64},
      NDTensors.Tensor{
        Float64,
        3,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64,Int64},
    },
  )
  precompile(
    Tuple{
      typeof(NDTensors.contract!!),
      NDTensors.Tensor{
        Float64,
        2,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
      NDTensors.Tensor{
        Float64,
        2,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
      NDTensors.Tensor{
        Float64,
        2,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
    },
  )
  precompile(
    Tuple{
      typeof(NDTensors.contract!!),
      NDTensors.Tensor{
        Float64,
        2,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
      NDTensors.Tensor{
        Float64,
        2,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
      NDTensors.Tensor{
        Number,
        2,
        NDTensors.Combiner,
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
    },
  )
  precompile(
    Tuple{
      typeof(NDTensors.contract!!),
      NDTensors.Tensor{
        Float64,
        2,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
      NDTensors.Tensor{
        Float64,
        2,
        NDTensors.Diag{Float64,Array{Float64,1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
      NDTensors.Tensor{
        Float64,
        2,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
    },
  )
  precompile(
    Tuple{
      typeof(NDTensors.contract!!),
      NDTensors.Tensor{
        Float64,
        2,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
      NDTensors.Tensor{
        Float64,
        3,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64,Int64},
      NDTensors.Tensor{
        Float64,
        1,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{1,Index{Int64},Tuple{Index{Int64}}},
      },
      Tuple{Int64},
    },
  )
  precompile(
    Tuple{
      typeof(NDTensors.contract!!),
      NDTensors.Tensor{
        Float64,
        2,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
      NDTensors.Tensor{
        Float64,
        3,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64,Int64},
      NDTensors.Tensor{
        Float64,
        3,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64,Int64},
    },
  )
  precompile(
    Tuple{
      typeof(NDTensors.contract!!),
      NDTensors.Tensor{
        Float64,
        3,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64,Int64},
      NDTensors.Tensor{
        Float64,
        2,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
      NDTensors.Tensor{
        Float64,
        3,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64,Int64},
    },
  )
  precompile(
    Tuple{
      typeof(NDTensors.contract!!),
      NDTensors.Tensor{
        Float64,
        3,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64,Int64},
      NDTensors.Tensor{
        Float64,
        3,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64,Int64},
      NDTensors.Tensor{
        Float64,
        2,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
    },
  )
  precompile(
    Tuple{
      typeof(NDTensors.contract!!),
      NDTensors.Tensor{
        Float64,
        3,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64,Int64},
      NDTensors.Tensor{
        Float64,
        4,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{4,Index{Int64},NTuple{4,Index{Int64}}},
      },
      NTuple{4,Int64},
      NDTensors.Tensor{
        Float64,
        1,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{1,Index{Int64},Tuple{Index{Int64}}},
      },
      Tuple{Int64},
    },
  )
  precompile(
    Tuple{
      typeof(NDTensors.contract!!),
      NDTensors.Tensor{
        Float64,
        3,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64,Int64},
      NDTensors.Tensor{
        Float64,
        4,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{4,Index{Int64},NTuple{4,Index{Int64}}},
      },
      NTuple{4,Int64},
      NDTensors.Tensor{
        Float64,
        3,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64,Int64},
    },
  )
  precompile(
    Tuple{
      typeof(NDTensors.contract!!),
      NDTensors.Tensor{
        Float64,
        4,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{4,Index{Int64},NTuple{4,Index{Int64}}},
      },
      NTuple{4,Int64},
      NDTensors.Tensor{
        Float64,
        2,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
      NDTensors.Tensor{
        Float64,
        2,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
    },
  )
  precompile(
    Tuple{
      typeof(NDTensors.contract!!),
      NDTensors.Tensor{
        Float64,
        4,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{4,Index{Int64},NTuple{4,Index{Int64}}},
      },
      NTuple{4,Int64},
      NDTensors.Tensor{
        Float64,
        3,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64,Int64},
      NDTensors.Tensor{
        Float64,
        3,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64,Int64},
    },
  )
  precompile(
    Tuple{
      typeof(NDTensors.contract!!),
      NDTensors.Tensor{
        Float64,
        4,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{4,Index{Int64},NTuple{4,Index{Int64}}},
      },
      NTuple{4,Int64},
      NDTensors.Tensor{
        Float64,
        4,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{4,Index{Int64},NTuple{4,Index{Int64}}},
      },
      NTuple{4,Int64},
      NDTensors.Tensor{
        Float64,
        4,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{4,Index{Int64},NTuple{4,Index{Int64}}},
      },
      NTuple{4,Int64},
    },
  )
  precompile(
    Tuple{
      typeof(NDTensors.contraction_output),
      NDTensors.Tensor{
        Float64,
        3,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
      },
      NDTensors.Tensor{
        Float64,
        3,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
      },
      IndexSet{0,Union{},Tuple{}},
    },
  )
  precompile(
    Tuple{
      typeof(NDTensors.contraction_output),
      NDTensors.Tensor{
        Float64,
        4,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{4,Index{Int64},NTuple{4,Index{Int64}}},
      },
      NDTensors.Tensor{
        Float64,
        1,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{1,Index{Int64},Tuple{Index{Int64}}},
      },
      IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
    },
  )
  precompile(
    Tuple{
      typeof(NDTensors.getperm),
      IndexSet{4,Index{Int64},NTuple{4,Index{Int64}}},
      IndexSet{4,Index{Int64},NTuple{4,Index{Int64}}},
    },
  )
  precompile(
    Tuple{
      typeof(NDTensors.permutedims!!),
      NDTensors.Tensor{
        Float64,
        3,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
      },
      NDTensors.Tensor{
        Float64,
        3,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64,Int64},
      Function,
    },
  )
  precompile(
    Tuple{
      typeof(NDTensors.permutedims!!),
      NDTensors.Tensor{
        Float64,
        4,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{4,Index{Int64},NTuple{4,Index{Int64}}},
      },
      NDTensors.Tensor{
        Float64,
        4,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{4,Index{Int64},NTuple{4,Index{Int64}}},
      },
      NTuple{4,Int64},
      Function,
    },
  )
  precompile(
    Tuple{
      typeof(Random.randn!),
      NDTensors.Tensor{
        Float64,
        3,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
      },
    },
  )
  precompile(Tuple{typeof(axpy!),Float64,ITensor{3},ITensor{3}})
  precompile(Tuple{typeof(combiner),Index{Int64}})
  precompile(
    Tuple{
      typeof(contract),
      NDTensors.Tensor{
        Complex{Float64},
        2,
        NDTensors.Dense{Complex{Float64},Array{Complex{Float64},1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
      NDTensors.Tensor{
        Complex{Float64},
        2,
        NDTensors.Dense{Complex{Float64},Array{Complex{Float64},1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
      Tuple{Int64,Int64},
    },
  )
  precompile(
    Tuple{
      typeof(contract),
      NDTensors.Tensor{
        Complex{Float64},
        2,
        NDTensors.Dense{Complex{Float64},Array{Complex{Float64},1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
      NDTensors.Tensor{
        Complex{Float64},
        2,
        NDTensors.Dense{Complex{Float64},Array{Complex{Float64},1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
    },
  )
  precompile(
    Tuple{
      typeof(contract),
      NDTensors.Tensor{
        Complex{Float64},
        2,
        NDTensors.Dense{Complex{Float64},Array{Complex{Float64},1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
      NDTensors.Tensor{
        Number,
        2,
        NDTensors.Combiner,
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
    },
  )
  precompile(
    Tuple{
      typeof(contract),
      NDTensors.Tensor{
        Float64,
        2,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
      NDTensors.Tensor{
        Float64,
        2,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
      NTuple{4,Int64},
    },
  )
  precompile(
    Tuple{
      typeof(contract),
      NDTensors.Tensor{
        Float64,
        2,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
      NDTensors.Tensor{
        Float64,
        2,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
      Tuple{Int64,Int64},
    },
  )
  precompile(
    Tuple{
      typeof(contract),
      NDTensors.Tensor{
        Float64,
        2,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
      NDTensors.Tensor{
        Float64,
        2,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
    },
  )
  precompile(
    Tuple{
      typeof(contract),
      NDTensors.Tensor{
        Float64,
        2,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
      NDTensors.Tensor{
        Float64,
        3,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64,Int64},
      Tuple{Int64,Int64,Int64},
    },
  )
  precompile(
    Tuple{
      typeof(contract),
      NDTensors.Tensor{
        Float64,
        2,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
      NDTensors.Tensor{
        Float64,
        3,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64,Int64},
    },
  )
  precompile(
    Tuple{
      typeof(contract),
      NDTensors.Tensor{
        Float64,
        2,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
      NDTensors.Tensor{
        Number,
        2,
        NDTensors.Combiner,
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
      Tuple{Int64,Int64},
    },
  )
  precompile(
    Tuple{
      typeof(contract),
      NDTensors.Tensor{
        Float64,
        2,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
      NDTensors.Tensor{
        Number,
        2,
        NDTensors.Combiner,
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
    },
  )
  precompile(
    Tuple{
      typeof(contract),
      NDTensors.Tensor{
        Float64,
        2,
        NDTensors.Diag{Float64,Array{Float64,1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
      NDTensors.Tensor{
        Complex{Float64},
        2,
        NDTensors.Dense{Complex{Float64},Array{Complex{Float64},1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
      Tuple{Int64,Int64},
    },
  )
  precompile(
    Tuple{
      typeof(contract),
      NDTensors.Tensor{
        Float64,
        2,
        NDTensors.Diag{Float64,Array{Float64,1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
      NDTensors.Tensor{
        Complex{Float64},
        2,
        NDTensors.Dense{Complex{Float64},Array{Complex{Float64},1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
    },
  )
  precompile(
    Tuple{
      typeof(contract),
      NDTensors.Tensor{
        Float64,
        2,
        NDTensors.Diag{Float64,Array{Float64,1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
      NDTensors.Tensor{
        Float64,
        2,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
      Tuple{Int64,Int64},
    },
  )
  precompile(
    Tuple{
      typeof(contract),
      NDTensors.Tensor{
        Float64,
        2,
        NDTensors.Diag{Float64,Array{Float64,1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
      NDTensors.Tensor{
        Float64,
        2,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
    },
  )
  precompile(
    Tuple{
      typeof(contract),
      NDTensors.Tensor{
        Float64,
        3,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64,Int64},
      NDTensors.Tensor{
        Float64,
        1,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{1,Index{Int64},Tuple{Index{Int64}}},
      },
      Tuple{Int64},
      Tuple{Int64,Int64},
    },
  )
  precompile(
    Tuple{
      typeof(contract),
      NDTensors.Tensor{
        Float64,
        3,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64,Int64},
      NDTensors.Tensor{
        Float64,
        1,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{1,Index{Int64},Tuple{Index{Int64}}},
      },
      Tuple{Int64},
    },
  )
  precompile(
    Tuple{
      typeof(contract),
      NDTensors.Tensor{
        Float64,
        3,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64,Int64},
      NDTensors.Tensor{
        Float64,
        2,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
      Tuple{Int64,Int64,Int64},
    },
  )
  precompile(
    Tuple{
      typeof(contract),
      NDTensors.Tensor{
        Float64,
        3,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64,Int64},
      NDTensors.Tensor{
        Float64,
        2,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64},
    },
  )
  precompile(
    Tuple{
      typeof(contract),
      NDTensors.Tensor{
        Float64,
        3,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64,Int64},
      NDTensors.Tensor{
        Float64,
        3,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64,Int64},
      NTuple{4,Int64},
    },
  )
  precompile(
    Tuple{
      typeof(contract),
      NDTensors.Tensor{
        Float64,
        3,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64,Int64},
      NDTensors.Tensor{
        Float64,
        3,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64,Int64},
      Tuple{Int64,Int64},
    },
  )
  precompile(
    Tuple{
      typeof(contract),
      NDTensors.Tensor{
        Float64,
        3,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64,Int64},
      NDTensors.Tensor{
        Float64,
        3,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64,Int64},
      Tuple{},
    },
  )
  precompile(
    Tuple{
      typeof(contract),
      NDTensors.Tensor{
        Float64,
        3,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64,Int64},
      NDTensors.Tensor{
        Float64,
        3,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64,Int64},
    },
  )
  precompile(
    Tuple{
      typeof(contract),
      NDTensors.Tensor{
        Float64,
        4,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{4,Index{Int64},NTuple{4,Index{Int64}}},
      },
      NTuple{4,Int64},
      NDTensors.Tensor{
        Float64,
        1,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{1,Index{Int64},Tuple{Index{Int64}}},
      },
      Tuple{Int64},
      Tuple{Int64,Int64,Int64},
    },
  )
  precompile(
    Tuple{
      typeof(contract),
      NDTensors.Tensor{
        Float64,
        4,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{4,Index{Int64},NTuple{4,Index{Int64}}},
      },
      NTuple{4,Int64},
      NDTensors.Tensor{
        Float64,
        1,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{1,Index{Int64},Tuple{Index{Int64}}},
      },
      Tuple{Int64},
    },
  )
  precompile(
    Tuple{
      typeof(contract),
      NDTensors.Tensor{
        Float64,
        4,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{4,Index{Int64},NTuple{4,Index{Int64}}},
      },
      NTuple{4,Int64},
      NDTensors.Tensor{
        Float64,
        3,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64,Int64},
      Tuple{Int64,Int64,Int64},
    },
  )
  precompile(
    Tuple{
      typeof(contract),
      NDTensors.Tensor{
        Float64,
        4,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{4,Index{Int64},NTuple{4,Index{Int64}}},
      },
      NTuple{4,Int64},
      NDTensors.Tensor{
        Float64,
        3,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
      },
      Tuple{Int64,Int64,Int64},
    },
  )
  precompile(
    Tuple{
      typeof(contract),
      NDTensors.Tensor{
        Float64,
        4,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{4,Index{Int64},NTuple{4,Index{Int64}}},
      },
      NTuple{4,Int64},
      NDTensors.Tensor{
        Float64,
        4,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{4,Index{Int64},NTuple{4,Index{Int64}}},
      },
      NTuple{4,Int64},
      NTuple{4,Int64},
    },
  )
  precompile(
    Tuple{
      typeof(contract),
      NDTensors.Tensor{
        Float64,
        4,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{4,Index{Int64},NTuple{4,Index{Int64}}},
      },
      NTuple{4,Int64},
      NDTensors.Tensor{
        Float64,
        4,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{4,Index{Int64},NTuple{4,Index{Int64}}},
      },
      NTuple{4,Int64},
    },
  )
  precompile(
    Tuple{typeof(hassameinds),ITensor{3},Tuple{Index{Int64},Index{Int64},Index{Int64}}}
  )
  precompile(Tuple{typeof(hastags),TagSet,String})
  precompile(Tuple{typeof(itensor),Array{Float64,2},Index{Int64},Index{Int64}})
  precompile(
    Tuple{
      typeof(itensor),Array{Float64,3},Index{Int64},Index{Int64},Vararg{Index,N} where N
    },
  )
  precompile(Tuple{typeof(mul!),ITensor{3},ITensor{3},Float64})
  precompile(
    Tuple{
      typeof(mul!),
      ITensor{3},
      KrylovKit.OrthonormalBasis{ITensor{3}},
      SubArray{Float64,1,Array{Float64,2},Tuple{UnitRange{Int64},Int64},true},
    },
  )
  precompile(Tuple{typeof(prime),ITensor{3},ITensors.Not{TagSet}})
  precompile(
    Tuple{
      typeof(prime),
      IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
      Int64,
      ITensors.Not{IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}}},
    },
  )
  precompile(Tuple{typeof(randomITensor),Index{Int64},Index{Int64},Index{Int64}})
  precompile(Tuple{typeof(randomMPS),Array{Index{Int64},1},Int64})
  precompile(Tuple{typeof(replaceind!),ITensor{2},Index{Int64},Index{Int64}})
  precompile(Tuple{typeof(replaceind!),ITensor{3},Index{Int64},Index{Int64}})
  precompile(
    Tuple{
      typeof(replaceind),
      IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      Index{Int64},
      Index{Int64},
    },
  )
  precompile(
    Tuple{
      typeof(replaceinds),
      IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      Tuple{Index{Int64},Index{Int64}},
      Tuple{Index{Int64},Index{Int64}},
    },
  )
  precompile(
    Tuple{
      typeof(replaceinds),
      IndexSet{3,Index{Int64},Tuple{Index{Int64},Index{Int64},Index{Int64}}},
      IndexSet{1,Index{Int64},Tuple{Index{Int64}}},
      IndexSet{1,Index{Int64},Tuple{Index{Int64}}},
    },
  )
  precompile(
    Tuple{
      typeof(setindex!),
      Dict{Array{ITensors.SiteOp,1},Array{Float64,2}},
      Array{Float64,2},
      Array{ITensors.SiteOp,1},
    },
  )
  precompile(
    Tuple{
      typeof(svd),
      NDTensors.Tensor{
        Complex{Float64},
        2,
        NDTensors.Dense{Complex{Float64},Array{Complex{Float64},1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
    },
  )
  precompile(
    Tuple{
      typeof(svd),
      NDTensors.Tensor{
        Float64,
        2,
        NDTensors.Dense{Float64,Array{Float64,1}},
        IndexSet{2,Index{Int64},Tuple{Index{Int64},Index{Int64}}},
      },
    },
  )
  return precompile(
    Tuple{
      typeof(zip),
      IndexSet{1,Index{Int64},Tuple{Index{Int64}}},
      Vararg{IndexSet{1,Index{Int64},Tuple{Index{Int64}}},N} where N,
    },
  )
end


# broadcasting

using Base.Broadcast: BroadcastStyle, Broadcasted

struct CuArrayStyle{N} <: AbstractGPUArrayStyle{N} end
CuArrayStyle(::Val{N}) where N = CuArrayStyle{N}()
CuArrayStyle{M}(::Val{N}) where {N,M} = CuArrayStyle{N}()

BroadcastStyle(::Type{CuArray{T,N}}) where {T,N} = CuArrayStyle{N}()

Base.similar(bc::Broadcasted{CuArrayStyle{N}}, ::Type{T}) where {N,T} =
    similar(CuArray{T}, axes(bc))

Base.similar(bc::Broadcasted{CuArrayStyle{N}}, ::Type{T}, dims) where {N,T} =
    CuArray{T}(undef, dims)

## replace base functions with libdevice alternatives

cufunc(f) = f
cufunc(::Type{T}) where T = (x...) -> T(x...) # broadcasting type ctors isn't GPU compatible

Broadcast.broadcasted(::CuArrayStyle{N}, f, args...) where {N} =
  Broadcasted{CuArrayStyle{N}}(cufunc(f), args, nothing)

const device_intrinsics = :[
  cos, cospi, sin, sinpi, tan, acos, asin, atan,
  cosh, sinh, tanh, acosh, asinh, atanh,
  log, log10, log1p, log2, logb, ilogb,
  exp, exp2, exp10, expm1, ldexp,
  erf, erfinv, erfc, erfcinv, erfcx,
  brev, clz, ffs, byte_perm, popc,
  isfinite, isinf, isnan, nearbyint,
  nextafter, signbit, copysign, abs,
  sqrt, rsqrt, cbrt, rcbrt, pow,
  ceil, floor, saturate,
  lgamma, tgamma,
  j0, j1, jn, y0, y1, yn,
  normcdf, normcdfinv, hypot,
  fma, sad, dim, mul24, mul64hi, hadd, rhadd, scalbn].args

for f in device_intrinsics
  isdefined(Base, f) || continue
  @eval cufunc(::typeof(Base.$f)) = $f
end

# broadcast ^

culiteral_pow(::typeof(^), x::T, ::Val{0}) where {T<:Real} = one(x)
culiteral_pow(::typeof(^), x::T, ::Val{1}) where {T<:Real} = x
culiteral_pow(::typeof(^), x::T, ::Val{2}) where {T<:Real} = x * x
culiteral_pow(::typeof(^), x::T, ::Val{3}) where {T<:Real} = x * x * x
culiteral_pow(::typeof(^), x::T, ::Val{p}) where {T<:Real,p} = pow(x, Int32(p))

cufunc(::typeof(Base.literal_pow)) = culiteral_pow
cufunc(::typeof(Base.:(^))) = pow

using MacroTools

const _cufuncs = [copy(device_intrinsics); :^]
cufuncs() = (global _cufuncs; _cufuncs)

_cuint(x::Int) = Int32(x)
_cuint(x::Expr) = x.head == :call && x.args[1] == :Int32 && x.args[2] isa Int ? Int32(x.args[2]) : x
_cuint(x) = x

function _cupowliteral(x::Expr)
  if x.head == :call && x.args[1] == :(CUDA.cufunc(^)) && x.args[3] isa Int32
    num = x.args[3]
    if 0 <= num <= 3
      sym = gensym(:x)
      new_x = Expr(:block, :($sym = $(x.args[2])))

      if iszero(num)
        push!(new_x.args, :(one($sym)))
      else
        unroll = Expr(:call, :*)
        for x = one(num):num
          push!(unroll.args, sym)
        end
        push!(new_x.args, unroll)
      end

      x = new_x
    end
  end
  x
end
_cupowliteral(x) = x

function replace_device(ex)
  global _cufuncs
  MacroTools.postwalk(ex) do x
    x = x in _cufuncs ? :(CUDA.cufunc($x)) : x
    x = _cuint(x)
    x = _cupowliteral(x)
    x
  end
end

macro cufunc(ex)
  global _cufuncs
  def = MacroTools.splitdef(ex)
  f = def[:name]
  def[:name] = Symbol(:cu, f)
  def[:body] = replace_device(def[:body])
  push!(_cufuncs, f)
  quote
    $(esc(MacroTools.combinedef(def)))
    CUDA.cufunc(::typeof($(esc(f)))) = $(esc(def[:name]))
  end
end