text:
#  Unit testing of (bounded) univariate discrete distributions
#
#  Here, bounded means the sample values are bounded.
#
#  Distributions covered by this suite:
#
# 	- Bernoulli
#  	- Categorical
#   - DiscreteUniform
#

using Distributions
using Base.Test
import StatsBase: entropy

distlist = [
    Bernoulli(0.1),
    Bernoulli(0.5),
    Bernoulli(0.9),
    Categorical([0.1, 0.9]),
    Categorical([0.5, 0.5]),
    Categorical([0.9, 0.1]),
    Categorical([0.2, 0.5, 0.3]),
    DiscreteUniform(0, 3),
    DiscreteUniform(2.0, 5.0),
    Binomial(1, 0.5),
    Binomial(100, 0.1),
    Binomial(100, 0.9)]

if length(ARGS) > 0
    newdistlist = {}
    for arg in ARGS
        a = eval(parse(arg))
        if isa(a, DataType)
            append!(newdistlist, filter(x -> isa(x,a),distlist))
        elseif isa(a,Distribution)
            push!(newdistlist, a)
        end
    end
    distlist = newdistlist
end

for d in distlist
    length(ARGS) > 0 && println(d)

    xmin = minimum(d)
    xmax = maximum(d)
    @assert isa(xmin, Int)
    @assert isa(xmax, Int)
    @assert xmin <= xmax

    ####
    #
    #  Part 1:  testing the capability of sampling
    #
    ####

    n = 10000

    # check that we can generate a single random draw
    draw = rand(d)
    @test isa(draw, Int)
    @test xmin <= draw <= xmax

    # check that we can generate many random draws at once
    x = rand(d, n)
    @test isa(x, Vector{Int})
    @test xmin <= minimum(x) <= maximum(x) <= xmax

    # check that we can generate many random draws in-place
    rand!(d, x)
    @test xmin <= minimum(x) <= maximum(x) <= xmax

    ####
    #
    #  Part 2: testing insupport
    #
    ####

    x = [xmin:xmax]
    n = length(x)

    @test !insupport(d, xmin-1)
    @test !insupport(d, xmax+1)
    for i = 1:n
        @test insupport(d, x[i])
    end
    @test all(insupport(d, x))

    ####
    #
    #  Part 3: testing evaluation
    #
    ####

    p = Array(Float64, n)
    c = Array(Float64, n)
    cc = Array(Float64, n)

    lp = Array(Float64, n)
    lc = Array(Float64, n)
    lcc = Array(Float64, n)

    ci = 0.

    for i in 1 : n
        p[i] = pdf(d, x[i])
        ci += p[i]

        c[i] = cdf(d, x[i])
        cc[i] = ccdf(d, x[i])

        @test_approx_eq ci c[i]
        @test_approx_eq c[i] + cc[i] 1.0

        lp[i] = logpdf(d, x[i])
        lc[i] = logcdf(d, x[i])
        lcc[i] = logccdf(d, x[i])

        @test_approx_eq_eps lp[i] log(p[i]) 1.0e-12
        @test_approx_eq_eps lc[i] log(c[i]) 1.0e-12
        @test_approx_eq_eps lcc[i] log(cc[i]) 1.0e-12

        if !isa(d, Binomial)
            @test quantile(d, c[i] - 1.0e-8) == x[i]
            @test cquantile(d, cc[i] + 1.0e-8) == x[i]
            @test invlogcdf(d, lc[i] - 1.0e-8) == x[i]

            if 0.0 < c[i] < 1.0
                @test invlogccdf(d, lcc[i] + 1.0e-8) == x[i]
            end
        end
    end

    # check consistency of scalar-based and vectorized evaluation

    @test_approx_eq pdf(d, x) p
    @test_approx_eq cdf(d, x) c
    @test_approx_eq ccdf(d, x) cc

    @test_approx_eq logpdf(d, x) lp
    @test_approx_eq logcdf(d, x) lc
    @test_approx_eq logccdf(d, x) lcc

    ####
    #
    #  Part 4: testing statistics
    #
    ####

    xf = float64(x)
    xmean = dot(p, xf)
    xvar = dot(p, abs2(xf .- xmean))
    xstd = sqrt(xvar)
    xentropy = entropy(p)
    xskew = dot(p, (xf .- xmean).^3) / (xstd.^3)
    xkurt = dot(p, (xf .- xmean).^4) / (xvar.^2) - 3.0

    @test_approx_eq mean(d)     xmean
    @test_approx_eq var(d)      xvar
    @test_approx_eq std(d)      xstd
    @test_approx_eq skewness(d) xskew
    @test_approx_eq kurtosis(d) xkurt
    @test_approx_eq entropy(d)  xentropy

end


text:
"""
    AbstractAttenMaskOp

Trait-like abstract type for holding operation related argument, defined how the mask should be apply to input array
"""
abstract type AbstractAttenMaskOp end

"""
    AbstractAttenMask

Abstract type for mask data, can be viewed as `AbstractArray{Bool}`
"""
abstract type AbstractAttenMask end

apply_mask(::Nothing, s) = s
apply_mask(_, ::Nothing, s) = s
apply_mask(::Nothing, m, s) = apply_mask(m, s)
apply_mask(m, s) = apply_mask(MaskOp(m), m, s)
MaskOp(m) = NaiveAttenMaskOp()
MaskOp(::Nothing) = nothing

struct NaiveAttenMaskOp <: AbstractAttenMaskOp end

"""
    apply_mask(op::NaiveAttenMaskOp, mask::AbstractAttenMask, score)

Directly broadcast multiply mask to attention score, i.e. `score .* mask`.
"""
apply_mask(op::NaiveAttenMaskOp, mask::AbstractAttenMask, score) = score .* mask

struct GenericAttenMaskOp{F, B<:StaticBool, T} <: AbstractAttenMaskOp
    apply::F
    flip::B
    scale::T
end

GenericAttenMaskOp(apply, flip::Bool, scale) = GenericAttenMaskOp(apply, static(flip), scale)

GenericAttenMaskOp(::typeof(+), flip::Bool, scale) = GenericAttenMaskOp(.+, flip, scale)
GenericAttenMaskOp(::typeof(-), flip::Bool, scale) = GenericAttenMaskOp(.+, flip, -scale)
GenericAttenMaskOp(::typeof(.-), flip::Bool, scale) = GenericAttenMaskOp(.+, flip, -scale)

# softmax norm default value
GenericAttenMaskOp() = GenericAttenMaskOp(.+, static(true), -1e9)

getmask(m::AbstractAttenMask, score) = getmask(m, score, one(eltype(score)))
function getmask(m::AbstractAttenMask, score, scale)
    tmp = similar(score)
    @. tmp = m * scale
    return tmp
end

function apply_broadcast_mask(f, mask::AbstractAttenMask, score, scale)
    @. f(score, mask * scale)
end

"""
    apply_mask(op::GenericAttenMaskOp, mask::AbstractAttenMask, score)

Equivalent to `op.apply(score, op.scale .* (op.flip ? .! mask : mask))`.

# Example

```julia
julia> x = randn(10, 10);

julia> m = CausalMask()
CausalMask()

julia> apply_mask(GenericAttenMaskOp(.+, true, -1e9), m, x) ==  @. x + (!m * -1e9)
true

```
"""
function apply_mask(op::GenericAttenMaskOp, mask::AbstractAttenMask, score)
    scale = convert(eltype(score), op.scale)
    apply = op.apply
    m = Bool(op.flip) ? !mask : mask

    if apply isa Base.BroadcastFunction
        masked_score = apply_broadcast_mask(apply.f, m, score, scale)
    else
        tmp = getmask(m, score, scale)
        masked_score = apply(tmp, score)
    end
    return masked_score
end

abstract type AbstractDatalessMask <: AbstractAttenMask end
abstract type AbstractArrayMask <: AbstractAttenMask end
abstract type AbstractWrapperMask <: AbstractAttenMask end

Broadcast.broadcastable(m::AbstractAttenMask) = m
Base.eltype(::AbstractAttenMask) = Bool
Base.@propagate_inbounds Broadcast.newindex(arg::AbstractAttenMask, I::CartesianIndex) = I
Base.@propagate_inbounds Broadcast.newindex(arg::AbstractAttenMask, I::Integer) = I

const MaskIndexer = Indexer{<:AbstractAttenMask}
Base.@propagate_inbounds Broadcast.newindex(arg::MaskIndexer, I::CartesianIndex) = I
Base.@propagate_inbounds Broadcast.newindex(arg::MaskIndexer, I::Integer) = I
Base.eltype(::MaskIndexer) = Bool

GetIndexer(m::AbstractDatalessMask) = m

Base.@propagate_inbounds Base.getindex(m::AbstractAttenMask, i::CartesianIndex) = m[Tuple(i)]
Base.@propagate_inbounds Base.getindex(m::AbstractAttenMask, I::Tuple) = GetIndexer(m)[I...]
Base.@propagate_inbounds Base.getindex(m::M, I::Integer...) where {M <: Union{<:AbstractWrapperMask, <:AbstractArrayMask}} = m[I]
Base.@propagate_inbounds Base.getindex(m::MaskIndexer, i::CartesianIndex) = m[Tuple(i)]
Base.@propagate_inbounds Base.getindex(m::MaskIndexer, I::Tuple) = m[I...]

@init @require CUDA="052768ef-5323-5732-b1bb-66c8b64840ba" begin
    Adapt.adapt(to::CUDA.Adaptor, m::AbstractArrayMask) = Indexer{typeof(m)}(map(Base.Fix1(Adapt.adapt, to), GetIndexer(m).__fields))
end

randomness(::AbstractAttenMask) = static(false)

ext: jl


text:
using Compat.Test
include("../Inference.jl")
include("../dataset/wavefunctions.jl")

@testset "Grover Search" begin
    # alway use sorted CSC format.
    oracle = inference_oracle([2,-1,3])(3)
    v = ones(1<<3)
    v[Int(0b110)+1] *= -1
    @test mat(oracle) ≈ Diagonal(v)

    ####### Construct Grover Search Using Reflection Block
    num_bit = 12
    oracle = inference_oracle(push!(collect(Int, 1:num_bit-1), num_bit))(num_bit)

    psi = GroverSearch(oracle, 12)
    target_state = zeros(1<<num_bit); target_state[end] = 1
    @test isapprox(abs(statevec(psi)'*target_state), 1, atol=1e-3)
end

max_stars_repo_name: Roger-luo/QCBM

max_issues_repo_name: Roger-luo/QCBM

max_forks_repo_name: Roger-luo/QCBM


text:
"""
    Exo2Real{T <: Real} = Parabolic{Exo1Real{T}}

A 2-exo-real number is a parabolic Cayley-Dickson construct with 1-exo-real numbers.
"""
const Exo2Real{T <: Real} = Parabolic{Exo1Real{T}}

function Exo2Real(z::Exo1Real{T}) where T <: Real
    Exo2Real{T}(z, zero(Exo1Real{T}))
end

function Exo2Real(a::T, b::T, c::T, d::T) where T <: Real
    Exo2Real{T}(Exo1Real{T}(a, b), Exo1Real{T}(c, d))
end

function Exo2Real(a::T, b::T, c::T) where T <: Real
    Exo2Real{T}(Exo1Real{T}(a, b), Exo1Real{T}(c, zero(T)))
end

function Exo2Real(a::T, b::T) where T <: Real
    Exo2Real{T}(Exo1Real{T}(a, b), zero(Exo1Real{T}))
end

function Exo2Real(a::T) where T <: Real
    Exo2Real{T}(Exo1Real{T}(a, zero(T)), zero(Exo1Real{T}))
end

function Exo2Real(a::Real, b::Real)
    Exo2Real(promote(a, b)...)
end

function Exo2Real(a::Real, b::Real, c::Real)
    Exo2Real(promote(a, b, c)...)
end

function Exo2Real(a::Real, b::Real, c::Real, d::Real)
    Exo2Real(promote(a, b, c, d)...)
end

function selfstar(z::Exo2Real)
    (z + star(z) + dagger(z) + star(dagger(z))) / 4
end

function anti_selfstar(z::Exo2Real)
    (z - star(z) + dagger(z) - star(dagger(z))) / 4
end

"""
    unreal(z::Exo2Real)

The unreal part of a 2-exo-real number is a 3-vector.
"""
function unreal(z::Exo2Real)
    vcat(unreal(z.l), array(z.r))
end

function show(io::IO, z::Exo2Real)
    print(io, "[1: ")
    print(io, z.l.l)
    print(io, ", a: ")
    print(io, z.l.r)
    print(io, ", b: ")
    print(io, z.r.l)
    print(io, ", ab: ")
    print(io, z.r.r)
    print(io, "]")
end

function random(::Type{Exo2Real{T}}) where T <: Real
    Exo2Real{T}(random(Exo1Real{T}), random(Exo1Real{T}))
end


text:
include("init-namedarrays.jl")

print("hcat/vcat, ")

letters = [string(Char(96+i)) for i=1:26]

## vectors
## same names
m = @inferred NamedArray(rand(10), (letters[1:10],))
m2 = @inferred NamedArray(rand(10), (letters[1:10],))
mm = hcat(m, m2)
@test mm.array == hcat(m.array, m2.array)
@test namesanddim(mm,1) == namesanddim(m,1)

mm = @inferred vcat(m, m2)
@test mm.array == vcat(m.array, m2.array)
@test names(mm, 1) == [string(i) for i in 1:length(mm)]

## different names
m2 = @inferred NamedArray(rand(10), (letters[11:20],))
mm = @inferred hcat(m, m2)
@test mm.array == hcat(m.array, m2.array)
@test names(mm, 1) != names(m2, 1)

mm = @inferred vcat(m, m2)
@test mm.array == vcat(m.array, m2.array)
@test names(mm, 1) == vcat(names(m, 1), names(m2, 1))

## matrix
## same names
m = @inferred NamedArray(rand(10, 10), (letters[1:10], letters[11:20]))
m2 = @inferred NamedArray(rand(10, 10), (letters[1:10], letters[11:20]))
mm = @inferred hcat(m, m2)
@test mm.array == hcat(m.array, m2.array)
@test namesanddim(mm,1) == namesanddim(m,1)

mm = @inferred vcat(m, m2)
@test mm.array == vcat(m.array, m2.array)
@test names(mm, 1) == [string(i) for i in 1:size(mm, 1)]
@test names(mm, 2) == names(m, 2)

## different names
m2 = @inferred NamedArray(rand(10, 10), (letters[11:20], letters[1:10]))
mm = @inferred hcat(m, m2)
@test mm.array == hcat(m.array, m2.array)
@test names(mm, 1) != names(m2, 1)

mm = @inferred vcat(m, m2)
@test mm.array == vcat(m.array, m2.array)
## @test names(mm, 1) == vcat(names(m, 1), names(m2, 1))
@test names(mm, 2) == [string(i) for i in 1:size(mm, 2)]


text:
using LinearAlgebra

@testset "Circle" begin
    @testset "CirclePointRadius" begin
        img = zeros(Gray{N0f8}, 10, 10)

        expected = Gray{N0f8}[  0 0 1 1 1 1 1 0 0 0
                                0 1 1 1 1 1 1 1 0 0
                                1 1 1 1 1 1 1 1 1 0
                                1 1 1 1 1 1 1 1 1 0
                                1 1 1 1 1 1 1 1 1 0
                                1 1 1 1 1 1 1 1 1 0
                                1 1 1 1 1 1 1 1 1 0
                                0 1 1 1 1 1 1 1 0 0
                                0 0 1 1 1 1 1 0 0 0
                                0 0 0 0 0 0 0 0 0 0 ]

        res = @inferred draw(img, CirclePointRadius(Point(5,5), 5))
        @test all(expected .== res) == true

        res = @inferred draw(img, CirclePointRadius(5, 5, 5))
        @test all(expected .== res) == true

        res = @inferred draw(img, CirclePointRadius(CartesianIndex(5,5), 5))
        @test all(expected .== res) == true

        res = @inferred draw(img, CirclePointRadius(Point(5,5), 5), Gray{N0f8}(0.5))
        expected = map(i->(i==1 ? Gray{N0f8}(0.5) : Gray{N0f8}(0)), expected)
        @test all(expected .== res) == true

    end

    @testset "CircleThreePoints" begin
        img = zeros(Gray{N0f8}, 10, 10)

        expected = Gray{N0f8}[ 0 0 0 0 0 0 0 0 0 0
                               0 0 0 1 1 1 1 1 0 0
                               0 0 1 1 1 1 1 1 1 0
                               0 1 1 1 1 1 1 1 1 1
                               0 1 1 1 1 1 1 1 1 1
                               0 1 1 1 1 1 1 1 1 1
                               0 1 1 1 1 1 1 1 1 1
                               0 1 1 1 1 1 1 1 1 1
                               0 0 1 1 1 1 1 1 1 0
                               0 0 0 1 1 1 1 1 0 0 ]

        res = @inferred draw(img, CircleThreePoints(Point(1,5), Point(5,1), Point(10,5)))
        @test all(expected .== res)

        res = @inferred draw(img, CircleThreePoints(1,5, 5,1, 10,5))
        @test all(expected .== res)

        res = @inferred draw(img, CircleThreePoints(CartesianIndex(5,1), CartesianIndex(1,5), CartesianIndex(5,10)))
        @test all(expected .== res)

        res = @inferred draw(img, CircleThreePoints(Point(1,5), Point(5,1), Point(10,5)), Gray{N0f8}(0.5))
        expected = map(i->(i==1 ? Gray{N0f8}(0.5) : Gray{N0f8}(0)), expected)
        @test all(expected .== res) == true

        invalid_circle = CircleThreePoints(Point(1,1), Point(2,2), Point(3,3))
        @test_throws ErrorException draw(img, invalid_circle)

        invalid_circle = CircleThreePoints(Point(1,1), Point(5,5), Point(10,1))
        @test_throws ErrorException draw(img, invalid_circle)

    end
end


text:
export
    PlanningCache,
    isps_queue_cost,
    initialize_planning_cache,
    reset_cache!

@with_kw_noshow struct PlanningCache
    closed_set::Set{Int}                    = Set{Int}()    # nodes that are completed
    active_set::Set{Int}                    = Set{Int}()    # active nodes
    node_queue::PriorityQueue{Int,Tuple{Int,Float64}} = PriorityQueue{Int,Tuple{Int,Float64}}() # active nodes prioritized by slack
end
function sprint_cache(io::IO, cache::PlanningCache;label_pad=14,pad=5)
    lpad(str) = sprint_padded(str;pad=label_pad,leftaligned=true)
    rpad(str) = sprint_padded(str;pad=label_pad,leftaligned=false)
    spad(str;kwargs...) = sprint_padded_list(str;pad=pad,leftaligned=false,kwargs...)
    print(io,"PlanningCache:","\n")
    print(io,"\t",lpad("closed_set:  "),cache.closed_set,"\n")
    print(io,"\t",lpad("active_set:  "),cache.active_set,"\n")
    print(io,"\t",lpad("node_queue:  "),cache.node_queue,"\n")
end

function Base.show(io::IO, cache::PlanningCache)
    sprint_cache(io, cache)
end

function isps_queue_cost(sched::OperatingSchedule,v::Int)
    path_spec = get_path_spec(sched,v)
    return (Int(path_spec.plan_path), minimum(get_slack(sched,v)))
end

function initialize_planning_cache(sched::OperatingSchedule)
    cache = PlanningCache()
    for v in get_all_root_nodes(sched)
        push!(cache.active_set,v)
        enqueue!(cache.node_queue,v=>isps_queue_cost(sched,v)) # need to store slack
    end
    cache
end

"""
    `reset_cache!(cache,sched)`

    Resets the cache so that a solution can be repaired (otherwise calling
    low_level_search!() will return immediately because the cache says it's
    complete)
"""
function reset_cache!(cache::PlanningCache,sched::OperatingSchedule)
    process_schedule!(sched)
    empty!(cache.closed_set)
    empty!(cache.active_set)
    empty!(cache.node_queue)
    for v in vertices(get_graph(sched))
        if is_root_node(get_graph(sched),v)
            push!(cache.active_set,v)
            enqueue!(cache.node_queue,v=>isps_queue_cost(sched,v))
        end
    end
    cache
end

export
    EnvGraphDict,
    construct_env_graph_dict

const EnvGraphDict = Dict{Symbol,GridFactoryEnvironment}
function construct_env_graph_dict(sched::OperatingSchedule,env::GridFactoryEnvironment)
    dict = EnvGraphDict()
    dict[:Default] = env
    for v in vertices(sched)
        node = get_node_from_vtx(sched,v)
        k = graph_key(node)
        if !haskey(dict,k)
            dict[k] = GridFactoryEnvironment(
                env,
                graph=deepcopy(env.graph),
                dist_function=deepcopy(get_dist_matrix(env))
            )
        end
    end
    return dict
end

function initialize_route_plan end

const State = CRCBS.GraphState
const Action = CRCBS.GraphAction

export
    construct_cost_model,
    construct_heuristic_model

function construct_cost_model end
function construct_heuristic_model end

export
    EnvironmentLayer,
    SearchEnv,
    construct_search_env,
    update_env!

@with_kw_noshow struct EnvironmentLayer
    graph::GridFactoryEnvironment   = GridFactoryEnvironment()
    problem_spec::ProblemSpec       = ProblemSpec()
end
GraphUtils.get_graph(layer::EnvironmentLayer)            = layer.graph
get_problem_spec(layer::EnvironmentLayer)           = layer.problem_spec
CRCBS.get_cost_model(layer::EnvironmentLayer)       = layer.cost_model
CRCBS.get_heuristic_model(layer::EnvironmentLayer)  = layer.heuristic_model

function construct_environment_layer(env::GridFactoryEnvironment,problem_spec::ProblemSpec)
    env_graph = GridFactoryEnvironment(env, graph=deepcopy(env.graph),
        dist_function=deepcopy(get_dist_matrix(env))
        )
    # prob_spec = ProblemSpec(problem_spec, D=get_dist_matrix(env_graph))
    prob_spec = ProblemSpec(problem_spec, D=env_graph)
    EnvironmentLayer(env_graph,prob_spec)
end

function populate_environment_layer_dict!(layers,
        sched::OperatingSchedule,
        env::GridFactoryEnvironment,
        prob_spec::ProblemSpec,
    )
    gkeys = unique(map(graph_key, map(n->n.node, sched.nodes)))
    for k in gkeys
        layers[k] = construct_environment_layer(env,prob_spec)
    end
    return layers
end
function construct_environment_layer_dict(args...)
    layers = Dict{Symbol,EnvironmentLayer}()
    populate_environment_layer_dict!(layers,args...)
end

"""
    SearchEnv{C,H,S} <: GraphEnv{State,Action,C}

Contains all of the information needed to fully define `PC_TAPF` and related
problems.
"""
@with_kw_noshow struct SearchEnv{C,H,S} <: GraphEnv{State,Action,C}
    schedule::OperatingSchedule     = OperatingSchedule()
    cache::PlanningCache            = PlanningCache()
    # Each agent type needs its own set of these #
    env_layers::Dict{Symbol,EnvironmentLayer} = Dict{Symbol,EnvironmentLayer}()
    # ###################################### #
    cost_model::C                   = C()
    heuristic_model::H              = H()
    num_agents::Int                 = length(get_robot_ICs(schedule))
    route_plan::S                   = initialize_route_plan(schedule,cost_model)
end

export
    get_schedule,
    get_cache,
    get_problem_spec,
    get_route_plan

GraphUtils.get_graph(env::SearchEnv,k=graph_key())     = get_graph(env.env_layers[k])
get_problem_spec(env::SearchEnv,k=graph_key())         = get_problem_spec(env.env_layers[k])
CRCBS.get_cost_model(env::SearchEnv)        = env.cost_model
CRCBS.get_heuristic_model(env::SearchEnv)   = env.heuristic_model
get_schedule(env::SearchEnv) = env.schedule
get_schedule(sched::OperatingSchedule) = sched
get_cache(env::SearchEnv) = env.cache
get_route_plan(env::SearchEnv) = env.route_plan

for op in path_spec_accessor_interface
    @eval $op(env::SearchEnv,v) = $op(get_schedule(env),v)
end
for op in path_spec_mutator_interface
    @eval $op(env::SearchEnv,v,val) = $op(get_schedule(env),v,val)
end

function CRCBS.get_start(env::SearchEnv,v::Int)
    node = get_node(get_schedule(env),v)
    start_vtx = get_id(get_default_initial_location_id(node))
    # start_vtx   = get_path_spec(get_schedule(env),v).start_vtx
    start_time  = get_t0(env,v)
    state_type(env)(start_vtx,start_time)
end
CRCBS.num_agents(env::SearchEnv) = env.num_agents
for op in [
    :cost_type,:state_type,:action_type,:path_type,:get_cost,:get_paths,
    :get_path_costs,:set_cost!,:set_solution_path!,:set_path_cost!,
    :convert_to_vertex_lists,:detect_conflicts,:detect_conflicts!
    ]
    @eval CRCBS.$op(env::SearchEnv,args...) = $op(get_route_plan(env),args...)
end
function sprint_search_env(io::IO,env::SearchEnv)
    # print(io,"schedule: ",get_schedule(env),"\n")
    print(io,"SearchEnv: \n")
    print(io,"cache: ",sprint(sprint_cache,get_cache(env)))
    print(io,"active task nodes:","\n")
    for v in get_cache(env).active_set
        print(io,"\t","v = ",
            sprint_padded(v)," => ",
            string(get_node_from_vtx(get_schedule(env),v)),"\n")
    end
    print(io,"route_plan: ",get_route_plan(env),"\n")
end
function Base.show(io::IO,env::SearchEnv)
    sprint_search_env(io,env)
end

export
    get_env_snapshot,
    trim_route_plan

"""
    get_env_snapshot(route_plan::S,t)
"""
function get_env_snapshot(env::SearchEnv,t)
    sched = get_schedule(env)
    route_plan = get_route_plan(env)
    Dict(id=>BOT_AT(id,
        LocationID(get_sp(get_path_node(get_paths(route_plan)[get_id(id)],t)).vtx)
        ) for (id,r) in get_robot_ICs(sched))
end

export
    EnvState,
    get_env_state

"""
    EnvState

Reflects the state of the SearchEnv environment at a given time step.
"""
@with_kw_noshow struct EnvState
    robot_positions::Dict{BotID,BOT_AT} = Dict{BotID,BOT_AT}()
    object_positions::Dict{ObjectID,OBJECT_AT} = Dict{ObjectID,OBJECT_AT}()
    objects_active::Dict{ObjectID,Bool} = Dict{ObjectID,Bool}()
end

export
    robot_positions,
    object_positions,
    objects_active,
    object_position,
    object_active,
    robot_position

robot_positions(s::EnvState) = s.robot_positions
object_positions(s::EnvState) = s.object_positions
objects_active(s::EnvState) = s.objects_active
object_position(s::EnvState,k) = object_positions(s)[k]
robot_position(s::EnvState,k) = robot_positions(s)[k]
object_active(s::EnvState,k) = objects_active(s)[k]

function Base.show(io::IO,s::EnvState)
    print(io,"EnvState")
    for (n,dict) in [
        ("robot_positions",robot_positions(s)),
        ("object_positions",object_positions(s)),
        ("objects_active",objects_active(s)),
        ]
        print(io,"\n\t",n," : ")
        for k in sort(collect(keys(dict)))
            val = dict[k]
            print(io,string(val),", ")
        end
    end
end
function get_env_state(env,t)
    sched = get_schedule(env)
    # cache = get_cache(env)
    robot_positions = get_env_snapshot(env,t)
    object_positions = Dict{ObjectID,OBJECT_AT}()
    objects_active = Dict{ObjectID,Bool}()
    for (o,o_node) in get_object_ICs(sched)
        v = get_vtx(sched,o)
        if t <= get_t0(env,v)
            object_positions[o] = o_node
            objects_active[o] = t >= get_t0(env,v)
        else
            object_positions[o] = o_node
            vtxs = capture_connected_nodes(
                sched,v,v->check_object_id(get_node_from_vtx(sched,v),o)
                )
            node_ids = map(v->get_vtx_id(sched,v),collect(vtxs))
            # @show node_ids
            v_deposit = filter(v->isa(get_node_from_vtx(sched,v),
            Union{BOT_DEPOSIT,TEAM_DEPOSIT}),
                collect(vtxs))[1]
            v_collect = filter(v->isa(get_node_from_vtx(sched,v),
            Union{BOT_COLLECT,TEAM_COLLECT}),
                collect(vtxs))[1]

            node = get_node_from_vtx(sched,v_deposit)
            if t >= get_t0(env,v_deposit)
                object_positions[o] = OBJECT_AT(o,
                    map(n->get_location_id(n),sub_nodes(node))
                    )
            elseif t >= get_t0(env,v_collect)
                r_ids = get_valid_robot_ids(node)
                object_positions[o] = OBJECT_AT(o,
                    map(r->get_location_id(robot_positions[r]),r_ids)
                    )
            # else
            #     r_ids = get_valid_robot_ids(node)
            #     @show object_positions[o] = OBJECT_AT(o,
            #         map(r->get_location_id(robot_positions[r]),r_ids)
            #         )
            end
            objects_active[o] = (t <= get_tF(env,v_deposit))
        end
    end
    EnvState(
        robot_positions = robot_positions,
        object_positions = object_positions,
        objects_active = objects_active
    )
end

export get_node_start_and_end_times
"""
    get_start_and_end_maps(sched,cache,default=0)

Return dictionaries mapping each node id in schedule to its start and end time
"""
function get_node_start_and_end_times(sched::OperatingSchedule)
    t0 = Dict{AbstractID,Int}(get_vtx_id(sched,v)=>get_t0(sched, v) for v in vertices(sched))
    tF = Dict{AbstractID,Int}(get_vtx_id(sched,v)=>get_tF(sched, v) for v in vertices(sched))
    return t0,tF
end
get_node_start_and_end_times(env::SearchEnv) = get_node_start_and_end_times(get_schedule(env))

export
    AbstractPC_MAPF,
    PC_MAPF,
    C_PC_MAPF,
    AbstractPC_TAPF,
    PC_TAPF,
    PC_TA,
    C_PC_TAPF

export
    construct_routing_problem

"""
    AbstractPC_MAPF

An abstract type of which all Precedence-Constrained Multi-Agent Path-Finding
problems are concrete subtypes.
"""
abstract type AbstractPC_MAPF <: AbstractMAPF end

"""
    AbstractPC_TAPF <: AbstractPC_MAPF

An abstract type of which all Precedence-Constrained Multi-Agent Task Assignment
and Path-Finding problems are concrete subtypes.
"""
abstract type AbstractPC_TAPF <: AbstractPC_MAPF end

"""
    PC_TAPF

Precedence-Constrained Multi-Agent Task Assignment and Path-Finding problem.
"""
struct PC_TAPF{E<:SearchEnv} <: AbstractPC_TAPF
    env::E
end

"""
    PC_TA

Precedence-Constrained Multi-Agent Task Assignment problem (no route planning).
"""
struct PC_TA{E<:SearchEnv} <: AbstractPC_TAPF
    env::E
end

"""
    PC_MAPF

A precedence-constrained multi-agent path-finding problem (no task assignment).
"""
struct PC_MAPF{E<:SearchEnv} <: AbstractPC_MAPF
    env::E
end
construct_routing_problem(prob::PC_TAPF,env) = PC_MAPF(env)

"""
    C_PC_MAPF

A collaborative precedence-constrained multi-agent path-finding problem. All
agents have assigned tasks, there are precedence constraints between tasks, and
some tasks must be done in teams.
"""
struct C_PC_MAPF{E<:SearchEnv} <: AbstractPC_MAPF
    env::E
end

"""
    C_PC_TAPF

Defines an instance of a Collaborative Precedence-Constrained Multi-Agent Task
    Assignment and Path-Finding problem, where agents must sometimes transport
    objects in teams.
"""
struct C_PC_TAPF{E<:SearchEnv} <: AbstractPC_TAPF
    env::E
end
C_PC_TAPF(p::PC_TAPF) = C_PC_TAPF(p.env)
construct_routing_problem(prob::C_PC_TAPF,env) = C_PC_MAPF(env)

for T in [:PC_TAPF,:PC_MAPF,:PC_TA]
    @eval $T(prob::$T,env::SearchEnv) = $T(env)
end

initialize_planning_cache(env::SearchEnv) = initialize_planning_cache(get_schedule(env))

"""
    get_next_vtx_matching_agent_id(schedule,cache,agent_id)

Return the node_id of the active node assigned to an agent.
"""
function get_next_vtx_matching_agent_id(env::SearchEnv,agent_id)
    @assert isa(agent_id,BotID)
    for v in get_cache(env).active_set
        n = get_node(get_schedule(env),v)
        has_robot_id(n) ? nothing : continue
        if agent_id == get_robot_id(n)
        # if agent_id == get_path_spec(get_schedule(env), v).agent_id
        # if agent_id == get_robot_id(get_node_from_vtx(get_schedule(env), v))
            return v
        end
    end
    return -1
end

"""
    get_next_node_matching_agent_id(schedule,cache,agent_id)

Return the node_id of the active node assigned to an agent.
"""
function get_next_node_matching_agent_id(env::SearchEnv,agent_id)
    @assert isa(agent_id,BotID)
    v = get_next_vtx_matching_agent_id(env,agent_id)
    if has_vertex(get_schedule(env),v)
        return get_vtx_id(get_schedule(env),v)
    end
    return agent_id
end

export
    update_planning_cache!

"""
    update_planning_cache!(solver,env,v,path)
"""
update_planning_cache!(solver,env::SearchEnv,v::Int,path::Path) = update_planning_cache!(solver,env,v,get_t(get_final_state(path)))
function update_planning_cache!(solver,env::SearchEnv,v::Int,t::Int=-1)
    cache = get_cache(env)
    sched = get_schedule(env)
    active_set = cache.active_set
    closed_set = cache.closed_set
    node_queue = cache.node_queue
    Δt = t - get_tF(env,v)
    if Δt > 0
        set_tF!(sched,v,t)
        process_schedule!(sched)
    end
    # update closed_set
    activated_vtxs = Int[]
    push!(closed_set,v)
    # update active_set
    setdiff!(active_set, v)
    for v2 in outneighbors(sched,v)
        active = true
        for v1 in inneighbors(sched,v2)
            if !(v1 in closed_set)
                active = false
                break
            end
        end
        if active
            push!(activated_vtxs,v2)
            push!(active_set, v2)               # add to active set
        end
    end
    # update priority queue
    for v2 in active_set
        node_queue[v2] = isps_queue_cost(sched,v2)
    end
    @log_info(2,verbosity(solver),"moved ",v," to closed set, moved ",activated_vtxs," to active set")
    @log_info(3,verbosity(solver),string("get_tF(sched,v) = ",get_tF(sched,v)))
    return cache
end
"""
    update_planning_cache!(solver,env)

Update cache continually. After a call to this function, the start and end times
of all schedule nodes will be updated to reflect the progress of active schedule
nodes (i.e., if a robot had not yet completed a GO task, the predicted final
time for that task will be updated based on the robot's current state and
distance to the goal).
All active nodes that don't require planning will be automatically marked as
complete.
"""
function update_planning_cache!(solver,env)
    cache = get_cache(env)
    sched = get_schedule(env)
    node = initialize_root_node(env)
    # Skip over nodes that are already planned or just don't need planning
    while true
        done = true
        for v in collect(cache.active_set)
            if get_path_spec(sched,v).plan_path == false
                update_planning_cache!(solver,env,v)
                done = false
            end
        end
        if done
            break
        end
    end
    for v in collect(cache.active_set)
        n = get_node(sched,v)
        agent_id = get_id(get_default_robot_id(n))
        # path_spec = get_path_spec(sched,v)
        if 1 <= agent_id <= num_agents(env)
            path = get_paths(env)[agent_id]
            s = get_final_state(path)
            t = get_t(s)
            goal_vtx = get_id(get_destination_location_id(n))
            goal_state = state_type(env)(goal_vtx,-1)
            d = get_distance(env,s,goal_state)
            set_tF!(env,v,max(get_tF(env,v),t+d))
        end
    end
    process_schedule!(sched)
    cache
end

for op in [:SumOfMakeSpans,:MakeSpan]
    @eval $op(sched::OperatingSchedule,cache::PlanningCache) = $op(
        get_tF(sched),
        sched.terminal_vtxs,
        map(k->sched.weights[k], sched.terminal_vtxs),
        map(v->get_tF(sched,v),sched.terminal_vtxs),
    )
end

function initialize_route_plan(sched::OperatingSchedule,cost_model)
    starts = State[]
    robot_ics = get_robot_ICs(sched)
    for k in sort(collect(keys(robot_ics)))
        push!(starts, State(vtx = get_id(get_location_id(robot_ics[k])), t = 0))
    end
    paths = map(s->Path{State,Action,cost_type(cost_model)}(s0=s, cost=get_initial_cost(cost_model)), starts)
    costs = map(p->get_cost(p), paths)
    cost = aggregate_costs(cost_model, costs)
    LowLevelSolution(paths=paths, cost_model=cost_model,costs=costs, cost=cost)
end
function initialize_route_plan(env::SearchEnv)
    cost_model=get_cost_model(env)
    paths=deepcopy(get_paths(get_route_plan(env)))
    costs = map(p->get_cost(p), paths)
    LowLevelSolution(
        paths=paths,
        cost_model=cost_model,
        costs=costs,
        cost=aggregate_costs(cost_model, costs)
        )
end
function initialize_route_plan(env::SearchEnv,cost_model)
    paths = [Path(
        path_nodes=deepcopy(p.path_nodes),
        s0=p.s0,
        cost = compute_path_cost(cost_model,get_graph(env),p,i)
    ) for (i,p) in enumerate(get_paths(get_route_plan(env)))]
    costs=map(p->get_cost(p),paths)
    LowLevelSolution(
        paths=paths,
        cost_model=cost_model,
        costs=costs,
        cost=aggregate_costs(cost_model, costs)
        )
end
function construct_search_env(
        solver,
        sched::OperatingSchedule,
        problem_spec::ProblemSpec,
        env_graph,
        cache::PlanningCache=initialize_planning_cache(sched),
        primary_objective=problem_spec.cost_function,
        ;
        extra_T=400,
        kwargs...
    )
    cost_model, heuristic_model = construct_cost_model(
        solver,
        sched,
        cache,
        problem_spec,
        env_graph,
        primary_objective;
        extra_T=extra_T,
        )
    layers = construct_environment_layer_dict(sched,env_graph,problem_spec)
    search_env = SearchEnv(
        schedule=sched,
        cache=cache,
        env_layers=layers,
        cost_model=cost_model,
        heuristic_model=heuristic_model,
        )
    return search_env
end
"""
    construct_search_env(solver,schedule,env,...)

Constructs a new search env by combining the new `schedule` with the pre-
existing `get_route_plan(env)`. This involves constructing a new cost function that
reflects the new schedule structure.
TODO: Carry over information about `get_cache(search_env)`
"""
function construct_search_env(
        solver,
        sched::OperatingSchedule,
        env::SearchEnv,
        cache::PlanningCache=initialize_planning_cache(sched),
        env_graph=get_graph(env),
        problem_spec = get_problem_spec(env),
        primary_objective = problem_spec.cost_function,
        args...
        ;
        extra_T=400,
        kwargs...
    )
    cost_model, heuristic_model = construct_cost_model(
        solver,
        sched,
        cache,
        problem_spec,
        env_graph,
        primary_objective;
        extra_T=extra_T,
        )
    layers = construct_environment_layer_dict(sched,env_graph,problem_spec)
    route_plan = initialize_route_plan(env,cost_model)
    N = length(get_paths(route_plan))
    search_env = SearchEnv(
        schedule=sched,
        cache=cache,
        env_layers=layers,
        cost_model=cost_model,
        heuristic_model=heuristic_model,
        num_agents=N,
        route_plan=route_plan)
end

"""
    function construct_search_env(solver, env::SearchEnv, ... )

Construct a new SearchEnv, with cost_model and heuristic_model defined by the
solver type.
"""
function construct_search_env(solver,
    env::SearchEnv,
    sched=get_schedule(env),
    cache=get_cache(env),
    env_graph=get_graph(env),
    prob_spec=get_problem_spec(env),
    primary_objective = problem_spec.cost_function,
    ;
    kwargs...
    )
    cost_model, heuristic_model = construct_cost_model(
        solver, sched, cache, prob_spec, env_graph, primary_objective;kwargs...
    )
    SearchEnv(
        schedule=sched,
        cache=cache,
        env_layers=env.env_layers,
        cost_model=cost_model,
        heuristic_model=heuristic_model,
    )

end

update_cost_model!(model::C,env::S) where {C,S<:SearchEnv} = nothing
function update_cost_model!(model::C,env::S) where {C<:MultiDeadlineCost,S<:SearchEnv}
    model.tF .= get_tF(get_schedule(env))
end
function update_cost_model!(model::C,env::S) where {C<:CompositeCostModel,S<:SearchEnv}
    for m in model.cost_models
        update_cost_model!(m,env)
    end
end
update_cost_model!(env::S) where {S<:SearchEnv} = update_cost_model!(get_cost_model(env),env)

"""
    `update_env!`

    `v` is the vertex id
"""
function update_env!(solver,env::SearchEnv,v::Int,path::P,
        agent_id = get_id(get_default_robot_id(get_node(get_schedule(env),v)))
        ) where {P<:Path}
    route_plan = get_route_plan(env)
    # UPDATE CACHE
    update_planning_cache!(solver,env,v,get_t(get_final_state(path)))
    update_cost_model!(env)
    # ADD UPDATED PATH TO HEURISTIC MODELS
    # n = get_node(get_schedule(env),v)
    # agent_id = has_robot_id(n) ? get_id(get_robot_id(n)) : -1
    # agent_id = get_id(get_default_robot_id(n))
    if agent_id != -1
        partially_set_path!(get_heuristic_model(env),agent_id,convert_to_vertex_lists(get_paths(route_plan)[agent_id]))
        partially_set_path!(get_cost_model(env),agent_id,convert_to_vertex_lists(get_paths(route_plan)[agent_id]))
    end

    env
end

################################################################################
############################### Low Level Search ###############################
################################################################################

include("pccbs.jl")

export evaluate_path_gap

"""
    evaluate_path_gap(search_env::SearchEnv,path,v)

Returns the gap between a path's length and it's expected length (based on times
stored in `get_cache(env).t0`)
"""
function evaluate_path_gap(search_env::SearchEnv,path,v)
    t0 = get_t0(search_env,v)
    gap = t0 - get_end_index(path)
    return gap
end

function replan_path!(solver, pc_mapf::AbstractPC_MAPF, env::SearchEnv, constraint_node::ConstraintTreeNode, vtx::VtxID,t0)
    v = get_id(vtx)
    sched = get_schedule(env)
    node_id = get_vtx_id(sched,v)
    schedule_node = get_node_from_id(sched,node_id)
    ### trim path
    cbs_env = build_env(solver,pc_mapf,env,constraint_node,VtxID(v))
    path = get_base_path(solver,env,cbs_env)
    trim_path!(cbs_env,path,get_t0(env,v))
    ### plan path with new goal time
    status = plan_path!(low_level(solver),pc_mapf,env,constraint_node,schedule_node,v)
    return status
end

export tighten_gaps!

"""
    tighten_gaps!(solver, pc_mapf, env::SearchEnv, constraint_node::ConstraintTreeNode)

If any path ends before it should (based on times stored in `get_cache(env)`),
recomputes the path segment for the final node in that line.
"""
function tighten_gaps!(solver, pc_mapf::AbstractPC_MAPF, env::SearchEnv, constraint_node::ConstraintTreeNode)
    solver.tighten_paths ? nothing : return env
    sched = get_schedule(env)
    active_nodes = robot_tip_map(sched,get_cache(env).active_set)
    for (robot_id, node_id) in active_nodes
        path = get_paths(env)[get_id(robot_id)]
        v = get_vtx(sched,node_id)
        gap = evaluate_path_gap(env,path,v)
        if gap > 0
            t0 = get_t0(env,v)
            @log_info(2, verbosity(solver), "tighten_gaps!: base path for ",
                string(get_node_from_vtx(sched,v)),
                ", v = ",v," ends at t=",get_end_index(path),
                " but should end at t=",t0," (gap = ", gap,").")
            vtxs = backtrack_node(sched,v)
            for vp in vtxs
                if get_tF(env,vp) < t0
                    @log_info(2,verbosity(solver)," Re-launching planner on ",string(get_node_from_vtx(sched,vp))," (v = ",v,")"," with extended horizon ",t0," ...")
                    if replan_path!(solver, pc_mapf, env, constraint_node, VtxID(vp),t0)
                        @log_info(2,verbosity(solver),"tightening succeeded on ",string(get_node_from_vtx(sched,vp))," (v = ",v,")"," with extended horizon ",t0)
                    else
                        @log_info(2,verbosity(solver),"tightening failed on ",string(get_node_from_vtx(sched,vp))," (v = ",v,")"," with extended horizon ",t0)
                        return false
                    end
                end
            end
        end
    end
    return true
end

function get_base_path(solver,search_env::SearchEnv,env::PCCBSEnv)
    base_path = get_paths(search_env)[get_agent_id(env)]
    return base_path
end
function CRCBS.build_env(
        solver,
        pc_mapf::AbstractPC_MAPF,
        env::SearchEnv,
        constraint_node::ConstraintTreeNode,
        node::T,
        v::Int,
        path_spec=get_path_spec(get_schedule(env), v),
        ;
        heuristic = get_heuristic_model(env),
        cost_model = get_cost_model(env)
    ) where {T}
    agent_id= get_id(get_default_robot_id(node))
    # goal_vtx = path_spec.final_vtx
    goal_vtx = has_robot_id(node) ? get_id(get_destination_location_id(node)) : -1
    goal_time = get_tF(env,v) # time after which goal can be satisfied
    deadline = get_tF(env,v) .+ get_slack(env,v)         # deadline for DeadlineCost
    # Adjust deadlines if necessary:
    if path_spec.tight == true
        goal_time += minimum(get_local_slack(env,v))
    end
    for v_next in outneighbors(get_graph(get_schedule(env)),v)
        if get_path_spec(get_schedule(env), v_next).static == true
            duration_next = get_path_spec(get_schedule(env),v_next).min_duration
            for c in sorted_state_constraints(env,get_constraints(constraint_node, agent_id)) #.sorted_state_constraints
                if get_sp(get_path_node(c)).vtx == goal_vtx
                    if 0 < get_time_of(c) - goal_time < duration_next
                        @log_info(1,verbosity(solver),"extending goal_time for node ",v,
                            " from ",goal_time," to ",get_time_of(c),
                            " to avoid constraints")
                        goal_time = max(goal_time, get_time_of(c))
                    end
                end
            end
        end
    end
    if (path_spec.free == true) && is_terminal_node(get_graph(get_schedule(env)),v)
        goal_time = makespan(get_schedule(env))
        goal_vtx = -1
        # deadline = Inf # already taken care of, perhaps?
        @log_info(3,verbosity(solver),string("BUILD ENV: ",string(node),
            " - setting goal_vtx = ",goal_vtx,", t = maximum(cache.tF) = ",goal_time))
    end
    @assert goal_time != Inf "goal time set to $goal_time for node $(string(node))"
    cbs_env = PCCBSEnv(
        search_env = env,
        schedule_node = node,
        node_id     = get_vtx_id(get_schedule(env),v),
        agent_idx   = agent_id, # this is only used for the HardConflictTable, which can be updated via the combined search node
        constraints = get_constraints(constraint_node, agent_id), # agent_id represents the whole path
        goal        = state_type(env)(goal_vtx,goal_time),
        )
    # update deadline in DeadlineCost
    set_deadline!(get_cost_model(cbs_env),deadline)
    @log_info(3,verbosity(solver),"build_env: ",string(node),
        ", goal_time=",goal_time,", deadline=",deadline)
    return cbs_env
end
function CRCBS.build_env(
    solver,
    pc_mapf::AbstractPC_MAPF,
    env::SearchEnv,
    constraint_node::ConstraintTreeNode,
    v::VtxID,
    node=get_node_from_vtx(get_schedule(env),get_id(v)),
    args...
    ;
    kwargs...
    )
    id = get_id(v)
    env = build_env(solver,pc_mapf,env,constraint_node,node,id,args...;kwargs...)
end
function CRCBS.build_env(
        solver,
        pc_mapf::AbstractPC_MAPF,
        env::SearchEnv,
        constraint_node::ConstraintTreeNode,
        agent_id::AgentID
        )
    node_id = get_next_node_matching_agent_id(env,RobotID(get_id(agent_id)))
    build_env(solver,pc_mapf,env,constraint_node,VtxID(get_vtx(get_schedule(env),node_id)))
end
CRCBS.build_env(env::SearchEnv) = PCCBSEnv(search_env=env)

function get_base_path(solver,search_env::SearchEnv,meta_env::MetaAgentCBS.AbstractMetaEnv)
    base_paths = map(env->get_base_path(solver,search_env,env),MetaAgentCBS.get_envs(meta_env))
    cost = aggregate_costs(get_cost_model(meta_env), map(get_cost,base_paths))
    MetaAgentCBS.construct_meta_path(base_paths,cost)
end

"""
    For COLLABORATIVE transport problems
"""
function CRCBS.build_env(
    solver,
    pc_mapf::C_PC_MAPF,
    env::E,
    constraint_node::N,
    node::T,
    v::Int,
    ;
    kwargs...) where {E<:SearchEnv,N<:ConstraintTreeNode,T}
    envs = []
    agent_idxs = Int[]
    for (i, sub_node) in enumerate(sub_nodes(node))
        ph = PerfectHeuristic(get_team_config_dist_function(get_graph(env),team_configuration(node),i))
        heuristic = construct_heuristic_model(solver,get_graph(env),ph)
        cbs_env = build_env(solver,PC_MAPF(get_env(pc_mapf)),env,constraint_node,VtxID(v),sub_node,generate_path_spec(get_schedule(env),get_problem_spec(env),sub_node);
            heuristic=heuristic,
            kwargs...
        )
        push!(envs, cbs_env)
        push!(agent_idxs, get_id(get_robot_id(sub_node)))
    end
    meta_env = MetaAgentCBS.construct_team_env(
        [envs...],
        agent_idxs,
        get_cost_model(env)
        )
    return meta_env
end

export
    reset_route_plan!

"""
    Helper to reset the solution in a constraint node between re-runs of ISPS
"""
function reset_route_plan!(node::N,base_route_plan) where {N<:ConstraintTreeNode}
    for (agent_id,path) in enumerate(get_paths(base_route_plan))
        set_solution_path!(node.solution, deepcopy(path), agent_id)
        set_path_cost!(node.solution, get_cost(path), agent_id)
    end
    set_cost!(node, get_cost(base_route_plan))
    node
end

################################################################################
############################## CBS Wrapper Stuff ###############################
################################################################################

CRCBS.build_env(solver, pc_mapf::AbstractPC_MAPF, args...) = build_env(solver,pc_mapf,get_env(pc_mapf),args...)
CRCBS.build_env(prob::Union{PC_TAPF,PC_MAPF}) = build_env(get_env(prob))
CRCBS.get_initial_solution(pc_mapf::PC_MAPF) = get_env(pc_mapf)
function CRCBS.initialize_root_node(env::SearchEnv)
    solution = SearchEnv(
        env,
        cache=initialize_planning_cache(env),
        route_plan=initialize_route_plan(env)
        )
    ConstraintTreeNode(
        solution    = solution,
        constraints = Dict(
            i=>discrete_constraint_table(env,i) for i in 1:num_agents(env)
            ),
        id = 1)
end
function CRCBS.discrete_constraint_table(env::SearchEnv,agent_id=-1,tf=2*makespan(get_schedule(env))+100*num_agents(env))
    discrete_constraint_table(num_states(env),num_actions(env),agent_id,tf)
end
CRCBS.initialize_root_node(solver,pc_mapf::AbstractPC_MAPF) = initialize_root_node(get_env(pc_mapf))
function Base.copy(env::SearchEnv)
    SearchEnv(
        env,
        schedule=deepcopy(get_schedule(env)),
        cache=deepcopy(get_cache(env)),
        route_plan=deepcopy(get_route_plan(env))
        )
end
function CRCBS.default_solution(env::SearchEnv)
    solution = deepcopy(env)
    set_cost!(get_route_plan(solution),get_infeasible_cost(get_route_plan(solution)))
    solution, get_cost(solution)
end
CRCBS.default_solution(pc_mapf::M) where {M<:AbstractPC_MAPF} = default_solution(get_env(pc_mapf))
function CRCBS.cbs_update_conflict_table!(solver,mapf::AbstractPC_MAPF,node,constraint)
    search_env = node.solution
    idxs = collect(1:num_agents(search_env))
    # t0 = max(minimum(get_cache(search_env).t0), 1) # This is particularly relevant for replanning, where we don't care to look for conflicts way back in the past.
    t0 = max(minimum(get_t0(get_schedule(search_env))),1) # This is particularly relevant for replanning, where we don't care to look for conflicts way back in the past.
    detect_conflicts!(node.conflict_table,get_route_plan(search_env),idxs,t0)
end
CRCBS.detect_conflicts!(table,env::SearchEnv,args...) = detect_conflicts!(table,get_route_plan(env),args...)

# CRCBS.serialize(pc_mapf::PC_MAPF,args...) = serialize(get_env(pc_mapf),args...)
# CRCBS.deserialize(pc_mapf::PC_MAPF,args...) = serialize(get_env(pc_mapf),args...)
CRCBS.get_env(pc_mapf::AbstractPC_MAPF)             = pc_mapf.env
for op in [
    :cost_type,:state_type,:action_type,:path_type,:num_states,:num_actions,
    :num_agents,:serialize,:deserialize,
    :discrete_constraint_table,
    ]
    @eval CRCBS.$op(prob::AbstractPC_MAPF,args...) = $op(get_env(prob),args...)
end
CRCBS.base_env_type(pc_mapf::AbstractPC_MAPF)               = PCCBSEnv
CRCBS.base_env_type(pc_mapf::Union{C_PC_MAPF,C_PC_TAPF})    = MetaAgentCBS.TeamMetaEnv


text:
left(rect::Rect{2}) = minimum(rect)[1]
right(rect::Rect{2}) = maximum(rect)[1]
bottom(rect::Rect{2}) = minimum(rect)[2]
top(rect::Rect{2}) = maximum(rect)[2]

width(rect::Rect{2}) = right(rect) - left(rect)
height(rect::Rect{2}) = top(rect) - bottom(rect)

function BBox(left::Number, right::Number, bottom::Number, top::Number)
    mini = (left, bottom)
    maxi = (right, top)
    return FRect2D(mini, maxi .- mini)
end

function RowCols(ncols::Int, nrows::Int)
    return RowCols(
        zeros(ncols),
        zeros(ncols),
        zeros(nrows),
        zeros(nrows)
    )
end

Base.getindex(rowcols::RowCols, ::Left) = rowcols.lefts
Base.getindex(rowcols::RowCols, ::Right) = rowcols.rights
Base.getindex(rowcols::RowCols, ::Top) = rowcols.tops
Base.getindex(rowcols::RowCols, ::Bottom) = rowcols.bottoms

"""
    eachside(f)
Calls f over all sides (Left, Right, Top, Bottom), and creates a BBox from the result of f(side)
"""
function eachside(f)
    return BBox(f(Left()), f(Right()), f(Bottom()), f(Top()))
end

"""
mapsides(
       f, first::Union{Rect{2}, RowCols}, rest::Union{Rect{2}, RowCols}...
   )::FRect2D
Maps f over all sides of the rectangle like arguments.
e.g.
```
mapsides(BBox(left, right, bottom, top)) do side::Side, side_val::Number
    return ...
end::FRect2D
```
"""
function mapsides(
        f, first::Union{Rect{2}, RowCols}, rest::Union{Rect{2}, RowCols}...
    )
    return eachside() do side
        f(side, getindex.((first, rest...), (side,))...)
    end
end


text:
let
    @test degrees_to_deg_min_sec(30.0) == (30,0,0)

    lla = LatLonAlt(deg2rad(deg_min_sec_to_degrees(73, 0, 0)),
                    deg2rad(deg_min_sec_to_degrees(45, 0, 0)),
                    0.0)
    utm = convert(UTM, lla, INTERNATIONAL)
    # @printf("%3d %1d %2.3f  %3d %1d %2.3f → ", degrees_to_deg_min_sec(rad2deg(lla.lat))..., degrees_to_deg_min_sec(rad2deg(lla.lon))...)
    # @printf("%2d %7.2f  %6.2f\n", utm.zone, utm.n, utm.e)

    lla = LatLonAlt(deg2rad(deg_min_sec_to_degrees( 30, 0, 0)),
                    deg2rad(deg_min_sec_to_degrees(102, 0, 0)),
                    0.0)
    utm = convert(UTM, lla, INTERNATIONAL)
    # @printf("%3d %1d %2.3f  %3d %1d %2.3f → ", degrees_to_deg_min_sec(rad2deg(lla.lat))..., degrees_to_deg_min_sec(rad2deg(lla.lon))...)
    # @printf("%2d %7.2f  %6.2f\n", utm.zone, utm.n, utm.e)

    utm = UTM(210577.93, 3322824.35, 0.0, 48)
    lla_target = LatLonAlt(deg2rad(deg_min_sec_to_degrees(30, 0, 6.489)),
                           deg2rad(deg_min_sec_to_degrees(101, 59, 59.805)), # E
                           0.0)
    lla = convert(LatLonAlt, utm, INTERNATIONAL)
    # @printf("%3d %1d %2.3f  %3d %1d %2.3f ← ", degrees_to_deg_min_sec(rad2deg(lla.lat))..., degrees_to_deg_min_sec(rad2deg(lla.lon))...)
    # @printf("%2d %7.2f  %6.2f\n", utm.zone, utm.n, utm.e)
    # @printf("%3d %1d %2.3f  %3d %1d %2.3f (should be this)\n", degrees_to_deg_min_sec(rad2deg(lla_target.lat))..., degrees_to_deg_min_sec(rad2deg(lla_target.lon))...)
    @test isapprox(lla.lat, lla_target.lat, atol=1e-8)
    @test isapprox(lla.lon, lla_target.lon, atol=1e-8)

    utm = UTM(789411.59, 3322824.08, 0.0, 47)
    lla_target = LatLonAlt(deg2rad(deg_min_sec_to_degrees(30, 0, 6.489)),
                           deg2rad(deg_min_sec_to_degrees(101, 59, 59.805)), # E
                           0.0)
    lla = convert(LatLonAlt, utm, INTERNATIONAL)
    # @printf("%3d %1d %2.3f  %3d %1d %2.3f ← ", degrees_to_deg_min_sec(rad2deg(lla.lat))..., degrees_to_deg_min_sec(rad2deg(lla.lon))...)
    # @printf("%2d %7.2f  %6.2f\n", utm.zone, utm.n, utm.e)
    # @printf("%3d %1d %2.3f  %3d %1d %2.3f (should be this)\n", degrees_to_deg_min_sec(rad2deg(lla_target.lat))..., degrees_to_deg_min_sec(rad2deg(lla_target.lon))...)
    @test isapprox(lla.lat, lla_target.lat, atol=1e-8)
    @test isapprox(lla.lon, lla_target.lon, atol=1e-8)

    utm = UTM(200000.00, 1000000.00, 0.0, 31)
    lla_target = LatLonAlt(deg2rad(deg_min_sec_to_degrees(9, 2, 10.706)),
                           deg2rad(deg_min_sec_to_degrees(0, 16, 17.099)), # E
                           0.0)
    lla = convert(LatLonAlt, utm, INTERNATIONAL)
    # @printf("%3d %1d %2.3f  %3d %1d %2.3f ← ", degrees_to_deg_min_sec(rad2deg(lla.lat))..., degrees_to_deg_min_sec(rad2deg(lla.lon))...)
    # @printf("%2d %7.2f  %6.2f\n", utm.zone, utm.n, utm.e)
    # @printf("%3d %1d %2.3f  %3d %1d %2.3f (should be this)\n", degrees_to_deg_min_sec(rad2deg(lla_target.lat))..., degrees_to_deg_min_sec(rad2deg(lla_target.lon))...)
    @test isapprox(lla.lat, lla_target.lat, atol=1e-8)
    @test isapprox(lla.lon, lla_target.lon, atol=1e-8)

    buf = IOBuffer()
    print(buf,  lla)
    close(buf)

    # should NOT create a warning
    check_is_in_radians(LatLonAlt(0.1,0.1,0.1))

    lla = LatLonAlt(0.1,0.2-2π,0.3)
    @test isapprox(ensure_lon_between_pies(lla).lon, 0.2, atol=1e-10)
    lla = LatLonAlt(0.1,0.2+2π,0.3)
    @test isapprox(ensure_lon_between_pies(lla).lon, 0.2, atol=1e-10)

    @test isapprox(get_earth_radius(0.0), 6378137.0, atol=0.1)
    @test isapprox(get_earth_radius(π/2), 6356752.3, atol=0.1)

    v3 = convert(VecE3, ECEF(0.1,0.2,0.3))
    @test v3.x == 0.1
    @test v3.y == 0.2
    @test v3.z == 0.3

    ecef = convert(ECEF, VecE3(0.1,0.2,0.3))
    @test ecef.x == 0.1
    @test ecef.y == 0.2
    @test ecef.z == 0.3

    v3 = convert(VecE3, UTM(0.1,0.2,0.3,1))
    @test v3.x == 0.1
    @test v3.y == 0.2
    @test v3.z == 0.3

    utm = UTM(0.1,0.2,0.3,1) + UTM(0.2,0.3,0.4,1)
    @test utm.e ≈ 0.3
    @test utm.n ≈ 0.5
    @test utm.u ≈ 0.7
    @test utm.zone == 1

    utm = UTM(0.1,0.2,0.3,1) - UTM(0.2,0.3,0.4,1)
    @test utm.e ≈ -0.1
    @test utm.n ≈ -0.1
    @test utm.u ≈ -0.1
    @test utm.zone == 1

    v3 = convert(VecE3, ENU(0.1,0.2,0.3))
    @test v3.x == 0.1
    @test v3.y == 0.2
    @test v3.z == 0.3

    enu = convert(ENU, VecE3(0.1,0.2,0.3))
    @test enu.e == 0.1
    @test enu.n == 0.2
    @test enu.u == 0.3

    enu = ENU(0.1,0.2,0.3) + ENU(0.2,0.3,0.4)
    @test enu.e ≈ 0.3
    @test enu.n ≈ 0.5
    @test enu.u ≈ 0.7

    enu = ENU(0.1,0.2,0.3) - ENU(0.2,0.3,0.4)
    @test enu.e ≈ -0.1
    @test enu.n ≈ -0.1
    @test enu.u ≈ -0.1

    for (lla, ecef) in [(LatLonAlt(deg2rad( 0),deg2rad( 0),  0.0), ECEF(6378137.0,0.0,0.0)),
                        (LatLonAlt(deg2rad(90),deg2rad( 0),  0.0), ECEF(0.0,0.0,6356752.31)),
                        (LatLonAlt(deg2rad(20),deg2rad(40),110.0), ECEF(4593156.33, 3854115.78, 2167734.41))]

        lla2 = convert(LatLonAlt, ecef)
        @test isapprox(lla2.lat, lla.lat, atol=1e-2)
        @test isapprox(lla2.lon, lla.lon, atol=1e-2)
        @test isapprox(lla2.alt, lla.alt, atol=1e-2)

        ecef2 = convert(ECEF, lla)
        @test isapprox(ecef2.x, ecef.x, atol=1.0)
        @test isapprox(ecef2.y, ecef.y, atol=1.0)
        @test isapprox(ecef2.z, ecef.z, atol=1.0)
    end
end


text:
# --------------------------------------------------------------------------
# ACE.jl and SHIPs.jl: Julia implementation of the Atomic Cluster Expansion
# Copyright (c) 2019 Christoph Ortner <christophortner0@gmail.com>
# All rights reserved.
# --------------------------------------------------------------------------

module Random

# TODO: rename rand_radial -> rand???

import LinearAlgebra: norm
import ACE: ScalarBasis, ZList, rand_radial, scaling

using Random: shuffle
using JuLIP: JVecF, AbstractCalculator, rnn, chemical_symbol, bulk, rattle!,
             fltype, rfltype
using JuLIP.MLIPs: combine
using JuLIP.Potentials: zlist, ZList, SZList
using StaticArrays: @SMatrix

export rand_nhd, rand_config, rand_sym, randcoeffs, randcombine

# -------------------------------------------
#   random neighbourhoods and  configurations

# TODO: JVecF is hard-coded
function rand_sphere()
   R = randn(JVecF)
   return R / norm(R)
end

# TODO: this could be rand(PSH1BasisFcn) ...
#       rand_sphere = rand(SphericalHarmonics)
rand_vec(J::ScalarBasis) where T = rand_radial(J) *  rand_sphere()
rand_vec(J::ScalarBasis, N::Integer) = [ rand_vec(J) for _ = 1:N ]

# -> rand_config?
function rand_nhd(Nat, J::ScalarBasis, species = :X)
   zlist = ZList(species)
   Rs = [ rand_vec(J) for _ = 1:Nat ]
   Zs = [ rand(zlist.list) for _ = 1:Nat ]
   z0 = rand(zlist.list)
   return Rs, Zs, z0
end

rand_config(species; kwargs...) =
      rand_config(ZList(species); kwargs...)

rand_config(V::AbstractCalculator; kwargs...) =
      rand_config(zlist(V); kwargs...)

function rand_config(zlist::Union{ZList, SZList};
                     absrattle = 0.0, relrattle = 0.2, repeat = 3,
                     kwargs...)
   # start with the longest rnn
   rnns = rnn.(zlist.list)
   sym = chemical_symbol( zlist.list[findmax(rnns)[2]] )
   at = bulk(sym; kwargs...) * repeat
   for n = 1:length(at)
      at.Z[n] = rand(zlist.list)
   end
   rattle!(at, maximum(rnns) * relrattle + absrattle)
   return at
end

# --------------------------------------------------------------
# random operations on neighbourhoods, mostly for testing

function rand_perm(Rs, Zs)
   @assert length(Rs) == length(Zs)
   p = shuffle(1:length(Rs))
   return Rs[p], Zs[p]
end

function rand_rot(Rs, Zs)
   @assert length(Rs) == length(Zs)
   K = (@SMatrix rand(3,3)) .- 0.5
   K = K - K'
   Q = exp(K)
   return [ Q * R for R in Rs ], Zs
end

rand_refl(Rs, Zs) = (-1) .* Rs, Zs

rand_sym(Rs, Zs) = rand_refl(rand_rot(rand_perm(Rs, Zs)...)...)

# -------------------------------------------
#     random potentials

# TODO: we have an issue with eltypes that needs to be fixed!!!

function randcoeffs(basis; diff = 2)
   ww = scaling(basis, diff)
   c = 2 * (rand(rfltype(basis), length(basis)) .- 0.5) ./ ww
   return c / norm(c)
end

randcombine(basis; diff = 2) =
   combine(basis, randcoeffs(basis; diff = diff))

# # move to utility???
# function rand(::Type{ACE.RPI.RPIBasis}; kwargs...)
#
# end

end


text:
"""
	function fit_psgld_step!(θ,∇L,precond,t, a, b, γ)

Perform one inplace draw of parameters using preconditioned stochastic gradient
Langevin dynamics.
"""
function fit_psgld_step!(θ,∇L,precond,step, a, b, γ)
    θ = θ₀
	##########DEFINE########################
	beta = 0.9;
	λ =1e-8;

	################GRADIENT CALCULATIONS
	 ϵ = a*(b + step)^-γ

	###############PRECONDITIONING#####################
	 if step == 1
	   precond[:] .= ∇L.*∇L
     else
	   precond .*= beta
	   precond .+= (1-beta).*(∇L .*∇L)
     end
	 #m = λ .+ sqrt.(precond/((1-(beta)^t)))
	 G_inv_diag = λ .+ sqrt.(precond)

	 ###############ASCENT###############################
	 for i in 1:length(∇L)
		 noise = sqrt(ϵ/G_inv_diag[i])*randn()
		 θ[i] += (0.5*ϵ*∇L[i]/G_inv_diag[i] +  noise)
	 end
	 return nothing
end


text:
"""

```
map_naive!(f,output,x,box)
```

Function that uses the naive pairwise mapping algorithm, for testing.

"""
function map_naive!(f,output,x,box::Box)
  @unpack unit_cell, cutoff_sq = box
  for i in 1:length(x)-1
    xᵢ = x[i]
    for j in i+1:length(x)
      xⱼ = wrap_relative_to(x[j],xᵢ,box)
      d2 = norm_sqr(xᵢ - xⱼ)
      if d2 <= cutoff_sq
        output = f(xᵢ,xⱼ,i,j,d2,output)
      end
    end
  end
  return output
end

function map_naive!(f,output,x,y,box::Box)
  @unpack unit_cell, cutoff_sq = box
  for i in 1:length(x)
    xᵢ = x[i]
    for j in 1:length(y)
      yⱼ = wrap_relative_to(y[j],xᵢ,box)
      d2 = norm_sqr(xᵢ - yⱼ)
      if d2 <= cutoff_sq
        output = f(xᵢ,yⱼ,i,j,d2,output)
      end
    end
  end
  return output
end

"""

```
view_celllist_particles(cl::CellList)
```

Auxiliary function to view the particles of a computing box, including images created
for computing purposes.

### Example

```julia
julia> box = Box([ 100 50; 50 100 ],10);

julia> x = [ box.unit_cell_max .* rand(SVector{2,Float64}) for i in 1:1000 ];

julia> cl = CellList(x,box);

julia> p = CellListMap.view_celllist_particles(cl);

julia> using Plots

julia> scatter(Tuple.(p),label=nothing,xlims=(-10,180),ylims=(-10,180))

```

"""
function view_celllist_particles(cl::CellList{N,T}) where {N,T}
  @unpack ncp, np = cl
  x = Vector{SVector{N,T}}(undef,ncp[1])
  ip = 0
  for p in cl.fp
    while p.index > 0
      ip += 1
      x[ip] = p.coordinates
      p = np[p.index]
    end
  end
  return [SVector{N,T}(ntuple(j -> x[i][j],N)) for i in 1:ncp[1]]
end

test_map(box,cl) = map_pairwise!((x,y,i,j,d2,s) -> s += d2, 0., box, cl, parallel=false)
test_naive(box,x) = CellListMap.map_naive!((x,y,i,j,d2,s) -> s += d2, 0., x, box)

function check_random_cells(N,M=2)
  local x, box
  ntrial = 0
  p = Progress(10000,0.5)
  while ntrial < 10000
    unit_cell_matrix = 10*rand(SMatrix{N,N,Float64})
    cutoff = max(0.01,rand())
    if !check_unit_cell(unit_cell_matrix,cutoff,printerr=false)
      continue
    end
    box = Box(unit_cell_matrix,cutoff)
    if prod(box.nc) > 100000
      continue
    end
    ntrial += 1
    next!(p)
    x = 100 .* rand(SVector{N,Float64},M)
    for i in eachindex(x)
      x[i] = x[i] .- 50
    end
    cl = CellList(x,box)
    if !(test_map(box,cl) ≈ test_naive(box,x))
      println("FOUND PROBLEMATIC SETUP.")
      return x, box
    end
  end
  println(" ALL PASSED! ")
  return x, box
end

function drawbox(box::Box{UnitCellType,2}) where UnitCellType
  S = SVector{2,Float64}
  m = box.unit_cell.matrix
  x = [S(0.,0.)]
  push!(x,S(m[:,1]))
  push!(x,S(m[:,1] + m[:,2]))
  push!(x,S(m[:,2]))
  push!(x,S(0.,0.))
  return x
end

function drawbox(box::Box{UnitCellType,3}) where UnitCellType
  S = SVector{3,Float64}
  m = box.unit_cell.matrix
  x = [S(0.,0.,0.)]
  push!(x,S(m[:,1]))
  push!(x,S(m[:,1] + m[:,2]))
  push!(x,S(m[:,2]))
  push!(x,S(0.,0.,0.))
  push!(x,S(m[:,1]))
  push!(x,S(m[:,1] + m[:,3]))
  push!(x,S(m[:,3]))
  push!(x,S(0.,0.,0.))
  push!(x,S(m[:,2]))
  push!(x,S(m[:,2] + m[:,3]))
  push!(x,S(m[:,2]))
  push!(x,S(0.,0.,0.))
  push!(x,S(m[:,1]))
  push!(x,S(m[:,1] + m[:,2]))
  push!(x,S(m[:,1] + m[:,2]) + m[:,3])
  push!(x,S(m[:,1] + m[:,3]))
  push!(x,S(m[:,3]))
  push!(x,S(m[:,3]) + m[:,2])
  push!(x,S(m[:,1] + m[:,2]) + m[:,3])
  return x
end

"""

```
draw_computing_cell(x,box::Box{UnitCellType,2}) where UnitCellType
```

This function creates a plot of the computing cell, in two dimensions.

"""
function draw_computing_cell(x,box::Box{UnitCellType,2}) where UnitCellType
  cl = CellList(x,box)
  box_points = drawbox(box)
  p = view_celllist_particles(cl)
  plt = Main.plot()
  Main.plot!(plt,Tuple.(box_points),label=:none)
  Main.scatter!(plt,Tuple.(p),label=:none,markeralpha=0.3)
  Main.scatter!(plt,Tuple.(wrap_to_first.(x,Ref(box))),label=:none)
  xmin = minimum(el[1] for el in p) - 3*box.cell_size[1]
  xmin = minimum(el[1] for el in p) - 3*box.cell_size[1]
  xmax = maximum(el[1] for el in p) + 3*box.cell_size[1]
  ymin = minimum(el[2] for el in p) - 3*box.cell_size[2]
  ymax = maximum(el[2] for el in p) + 3*box.cell_size[2]
  Main.plot!(plt,
    aspect_ratio=1,framestyle=:box,xrotation=60,
    xlims=(xmin,xmax),
    ylims=(ymin,ymax),
    xticks=(round.(digits=3,xmin:box.cell_size[1]:xmax)),
    yticks=(round.(digits=3,ymin:box.cell_size[2]:ymax)),
  )
  return plt
end

"""

```
draw_computing_cell(x,box::Box{UnitCellType,3}) where UnitCellType
```

This function creates a plot of the computing cell, in three dimensions.

"""
function draw_computing_cell(x,box::Box{UnitCellType,3}) where UnitCellType
  cl = CellList(x,box)
  box_points = drawbox(box)
  p = view_celllist_particles(cl)
  plt = Main.plot()
  Main.plot!(plt,Tuple.(box_points),label=:none)
  Main.scatter!(plt,Tuple.(p),label=:none,markeralpha=0.3)
  Main.scatter!(plt,Tuple.(wrap_to_first.(x,Ref(box))),label=:none,markeralpha=0.3)
  lims = Vector{Float64}[]
  for i in 1:3
    push!(lims, [ -2*box.cell_size[i], box.nc[i] + 2*box.cell_size[i] ])
  end
  Main.plot!(plt,
    aspect_ratio=1,framestyle=:box,xrotation=60,yrotation=-70,zrotation=0,
    xlims=lims[1],
    ylims=lims[2],
    zlims=lims[3],
    xticks=(round.(digits=3,lims[1][1]:box.cell_size[1]:lims[1][2])),
    yticks=(round.(digits=3,lims[2][1]:box.cell_size[2]:lims[2][2])),
    zticks=(round.(digits=3,lims[3][1]:box.cell_size[3]:lims[3][2])),
  )
  return plt
end


text:
# generate a spectra with multiple peaks for calibration testing
dist_energies = [4000,4700,5500,6000,6800,7500,8200,9600,10000]
energy_value_distribution = MixtureModel([Normal(mu,15) for mu in dist_energies])

be = 0:1:20000

events = rand(energy_value_distribution, 10000)
be, c = hist(events, be)
k=UnivariateKDE(be, float(c))

findpeaks(h::Histogram, fwhm) = findpeaks(bin_edges(h), counts(h), fwhm, typemax(Int))
findpeaks(h::Histogram, fwhm,n) = findpeaks(bin_edges(h), counts(h), fhwm,n)

findpeak(bin_edges::Range, x::Vector, fwhm) = findpeaks(bin_edges, x, fwhm, typemax(Int))
function findpeaks(bin_edges::Range, x::Vector, fwhm,n)
	k = UnivariateKDE(bin_edges, convert(Array{Float64,1},x))
	findpeaks(k,fwhm,n)
end

function findpeaks(x::UnivariateKDE, fwhm,n)
	σ = fwhm/sqrt(8*log(2))
	d = Normal(0,σ)
	smoothed = conv(x,d)
	peakinds = findpeaks_(smoothed.density)
	if length(peakinds)>n
		peakinds = peakinds[end-n+1:end]
	end
	peakinds
end

function findpeaks_{T}(x::Vector{T})
	# find all local maxima, return a vector of indicies sorted by value of x at that index
	peakinds = Array(Int,0)
	for i = 2:length(x)-1
		if x[i-1]<x[i]>x[i+1]
			push!(peakinds, i)
		end
	end
	p = sortperm(x[peakinds])
	peakinds[p]
end

function plotpeaks(bin_edges::Range, x::Vector, fwhm, n)
	k = UnivariateKDE(bin_edges, convert(Array{Float64,1},x))
	peakinds = findpeaks(bin_edges, x, fwhm, n)
	peakx = midpoints(bin_edges)[peakinds]
	figure()
	plot(midpoints(bin_edges), x, label="x")
	plot(midpoints(k.x), conv(k,Normal(0,fwhm/2.355)).density, label="smoothed")
	plot(peakx, x[peakinds],"o",label="raw peaks")
	plot(peakx, conv(k,Normal(0,fwhm/2.355)).density[peakinds],"o",label="smoothed peaks")
	legend()
end

peakinds = findpeaks(k,50,length(dist_energies))
peakx = sort(midpoints(be)[peakinds])

plotpeaks(be, c, 10,10)

figure()
plot(midpoints(be),c)
plot(midpoints(k.x), conv(k,Normal(0,10/2.355)).density)
plot(midpoints(be)[peakinds], c[peakinds],"o")
plot(midpoints(k.x)[peakinds], conv(k,Normal(0,10/2.355)).density[peakinds],"o")

@test all(abs(peakx-dist_energies).<5);


text:
using LinearAlgebra

function proj_dual_iso2d!(p, temp, weight)
    p1, p2 = view(p,:,:,1), view(p,:,:,2)
    temp .= 1 ./ max.(1.0, sqrt.( p1 .^2 + p2 .^ 2 ) ./ (weight+1e-8) )
    p1 .*= temp
    p2 .*= temp
end

"Project to dual l1 norm"
function proj_dual_l1!(x, weight)
    x ./= max.(1.0, abs.(x) / weight)
end

function div2d!(divp, p)
    p1, p2 = view(p,:,:,1), view(p,:,:,2)
    p1_x, p2_y = view(divp,:,:,1), view(divp,:,:,2)

    # p2_y is temp variable to save memory
    p1_x         .=  p1 - circshift!(p2_y, p1, [0, 1])
    p1_x[:, end] .= -p1[:, end-1]
    p1_x[:,   1] .=  p1[:, 1]

    circshift!(p2_y, p2, [1, 0])
    p2_y         .= p2 - p2_y
    p2_y[end, :] .= -p2[end-1, :]
    p2_y[1,   :] .=  p2[1, :]

    p1_x .+= p2_y
    return p1_x
end

function grad2d!(du, u) where {T<:AbstractFloat}
    du1, du2 = view(du, :, :, 1), view(du, :, :, 2)
    du1 .= circshift!(du1, u, [0, -1]) .- u
    du2 .= circshift!(du2, u, [-1, 0]) .- u
    du1[:, end] .= 0.0
    du2[end, :] .= 0.0
end

"Compute opeartor norm based on the Power method"
function compute_opnorm_block(As,shape, niter=3)
    λ = 0
    size,count = shape
    x = rand(size*count)

    for i=1:niter
        x_next = zeros(size*count)
        for j=1:count
            A = As[j]
            Ax = A*x[((j-1)*size)+1:((j-1)*size)+size]
            x_next_j = A'*Ax
            x_next[((j-1)*size)+1:((j-1)*size)+size] = x_next_j
        end
        λ_prev = λ
        λ = sqrt(sum(x_next.^2))

        if (abs(λ_prev-λ) < 1e-9)
            println("break at iteration: ", i)
            break
        end
        x = x_next / λ
    end
    return sqrt(λ)
end

function _recon2d_tv_primaldual_flow(As,A_norm,W_list,W_norm,u0s,bs,w_tv,w_flow,c,niter,p1,p2,p3)
    #A variational reconstruction method for undersampled dynamic x-ray tomography based on physical motion models (Burger)#
    height, width, frames = size(u0s)

    u = u0s
    u_prev = similar(u)
    ubar = deepcopy(u)

    p_adjoint  = similar(p3)
    p3_adjoint = similar(p3)
    p_adjoint_prev = similar(p3)
    opsnorm = [A_norm, sqrt(8), W_norm]
    tau = c /sum(opsnorm)

    data1 = similar(bs)
    u_grad = similar(p2)
    u_temp = similar(u)
    divp2 = similar(p2)

    Wus = similar(u)

    sigmas = map(n -> 1.0/(n*c), opsnorm)
    #@info sigmas

    for it=1:niter
        u_prev .= u
        p_adjoint_prev .= p_adjoint
        Threads.@threads for t=1:frames
            A = As[t]
            #sigmas = map(n-> 1.0/(n*c), opsnorm)

            @views data1[:,t] .= mul!(data1[:,t], A, vec(ubar[:,:,t])) .- bs[:,t]
            @views p1[:,t] .= (p1[:,t] .+ sigmas[1] * data1[:,t]) / (sigmas[1] + 1.0)
            @views mul!( view(p_adjoint[:,:,t], :), A', p1[:,t])

            @views grad2d!(u_grad[:,:,:,t], ubar[:,:,t])
            @views p2[:,:,:,t] .+= sigmas[2] .* u_grad[:,:,:,t]
            @views proj_dual_iso2d!(p2[:,:,:,t], u_temp[:,:,t], w_tv)

            @views p_adjoint[:,:,t] .-= div2d!(divp2[:,:,:,t], p2[:,:,:,t])

            if t < frames
                _p3 = view(p3, :, :, t)
                Wu = view(Wus, :, :, t)
                Wuv = vec(Wu)
                @views Wuv .= mul!(Wuv, W_list[t], vec(ubar[:,:,t+1])) .- vec(ubar[:,:,t])

                @views _p3 .+= sigmas[3] .* Wu
                proj_dual_l1!(_p3, w_flow)
                p3_adj = view(p3_adjoint[:,:,t], :)

                mul!(p3_adj, W_list[t]', vec(_p3))
                # @views p_adjoint[:,:,t+1] .+= p3_adjoint[:,:,t] # race condition
                @views p_adjoint[:,:,t] .-= _p3
            end
        end

        # add p3_adjoint to prevent race condition
        @views p_adjoint[:,:,2:end] .+= p3_adjoint[:,:,1:end-1]

        # primal update
        u .-= tau * p_adjoint
        u .= max.(u, 0.0) # positivity constraint

        # acceleration
        ubar .= 2*u .- u_prev

        if it % 50 == 0
            #du = u_prev - u
            primal_gap = sum(abs.(-p_adjoint+p_adjoint_prev + (u_prev-u)/tau)) / length(p_adjoint)
            #@info "primal gap:" primal_gap
            @info it primal_gap
        end
    end
    return u
end

"""
    recon2d_tv_primaldual_flow(As, bs, u0s, niter, w_tv, w_flow, c=1.0)

Reconstruct a 2d image by total variation model using Primal Dual optimization method and optical flow estimation

# Args
As : List of forward opeartors
bs : List of projection data
u0s: Initial guess of images
niter: number of iterations
w_tv: weight for TV term
W_flow: weight for flow term
c : See 61 page in 2016_Chambolle,Pock_An_introduction_to_continuous_optimization_for_imagingActa_Numerica

"""
function recon2d_tv_primaldual_flow(A_list, bs::Array{R, 2}, u0s::Array{R, 3}, niter1::Int, niter2::Int, w_tv::R, w_flow::R, c=10.0) where {R <: AbstractFloat}

    height,width,frames = size(u0s)
    v = zeros(height, width, 2, frames)
    u = u0s
    shape = (height*width, frames)
    A_norm = compute_opnorm_block(A_list, shape)

    p1 = zeros(size(bs))
    p2 = zeros(height, width, 2,frames)
    p3 = zeros(height, width, frames)
    for i=1:niter1
        @info i
        u_prev = u
        v_prev = v
        W_list = mapslices(f -> compute_warping_operator(f), v,dims=[1,2,3])
        # check if W_list is correct
        println(sum(abs.(W_list[1]*vec(u[:,:,2]) - vec(u[:,:,1]))))
        @views @info "check if W_list is correct. data term after and before warping" sum(abs.(W_list[1]*vec(u[:,:,2]) - vec(u[:,:,1])))/(height*width) sum(abs.(u[:,:,2] - u[:,:,1]))/(height*width)

        W_norm = compute_opnorm_block(W_list, shape)
        u = _recon2d_tv_primaldual_flow(A_list,A_norm,W_list,W_norm,u,bs,w_tv,w_flow,c,niter2,p1,p2,p3)
        v = get_flows(u)
    end
    return u
end

max_stars_repo_name: QPImaging/xfromProjections

max_issues_repo_name: QPImaging/xfromProjections

max_forks_repo_name: QPImaging/xfromProjections


text:
struct PermutingPartitioner{Prm, Prt}
    prm::Prm
    prt::Prt
end

function partition_plaid(A::SparseMatrixCSC, K, method::PermutingPartitioner; adj_A = nothing, kwargs...)
    @inbounds begin
        if adj_A === nothing
            adj_A = adjointpattern(A)
        end
        σ, τ = permute_plaid(A, method.prm; adj_A = adj_A)
        B = A[perm(σ), perm(τ)]
        Π, Φ = partition_plaid(B, K, method.prt; adj_A = adj_A)
        return (σ[Π], τ[Φ])
    end
end

function partition_stripe(A::SparseMatrixCSC, K, method::PermutingPartitioner; kwargs...)
    @inbounds begin
        τ = permute_stripe(A, method.prm)
        B = A[:, perm(τ)]
        Φ = partition_stripe(B, K, method.prt)
        return τ[Φ]
    end
end

function partition_stripe(A::SparseMatrixCSC, K, Π, method::PermutingPartitioner; kwargs...)
    @inbounds begin
        τ = permute_stripe(A, method.prm)
        B = A[:, perm(τ)]
        Φ = partition_stripe(B, K, Π, method.prt)
        return τ[Φ]
    end
end

function pack_plaid(A::SparseMatrixCSC, method::PermutingPartitioner; adj_A = nothing, kwargs...)
    @inbounds begin
        if adj_A === nothing
            adj_A = adjointpattern(A)
        end
        σ, τ = permute_plaid(A, method.prm; adj_A = adj_A)
        B = A[perm(σ), perm(τ)]
        Π, Φ = pack_plaid(B, method.prt; adj_A = adj_A)
        return (σ[Π], τ[Φ])
    end
end

function pack_stripe(A::SparseMatrixCSC, method::PermutingPartitioner; kwargs...)
    @inbounds begin
        τ = permute_stripe(A, method.prm)
        B = A[:, perm(τ)]
        Φ = pack_stripe(B, method.prt)
        return τ[Φ]
    end
end

function pack_stripe(A::SparseMatrixCSC, Π, method::PermutingPartitioner; kwargs...)
    @inbounds begin
        τ = permute_stripe(A, method.prm)
        B = A[:, perm(τ)]
        Φ = pack_stripe(B, Π, method.prt)
        return τ[Φ]
    end
end


text:
###############################################################
######################SURVIVAL MODEL###########################
###############################################################

abstract type SurvivalModel end

function show(io::IO, obj::SurvivalModel)
    print(io,"\nModel: ", obj.model, obj.formula,"\n\n")
    print(io,obj.coefmat)
end

coef(SM::SurvivalModel) = SM.coefmat.cols[1]

immutable CoxModel <: SurvivalModel
    model::AbstractString
    formula::Formula
    coefmat::CoefTable
    M::ModelFrame
    loglik::Float64
    score::Array{Float64,1}
    fischer_info::Array{Float64,2}
end

immutable AftModel{D<:Distribution} <: SurvivalModel
    model::AbstractString
    formula::Formula
    coefmat::CoefTable
    M::ModelFrame
    loglik::Float64
    score::Array{Float64,1}
    fischer_info::Array{Float64,2}
    dist::D
end

coef_reg(AM::AftModel) = coef(AM)[length(AM.dist.params)+1:end]
coef_dist(AM::AftModel) = AM.dist.params

function predict(AM::AftModel)
    model_matrix = DataFrames.ModelMatrix(AM.M)
    X = model_matrix.m
    coefs = coef_reg(AM)
    exp.(X*coefs)
end

function rand(AM::AftModel)
    preds = predict(AM)
    [rand(AM.dist)*pred for pred in preds]
end
###############################################################
######################EVENT####################################
###############################################################

immutable Event{T<:Real}
    time::T
    censored::Bool
end

Event(x) = Event(x,false)

function Base.show(io::IO, obj::Event)
    print(io, obj.time, obj.censored ? "+":"")
end

function Base.parse(T::Type{Event},str::String)
    if str[end] == '+'
        return Event(parse(str[1:(end-1)]), true)
    else
        return Event(parse(str), false)
    end
end

Base.isless(a::Event, b::Event) = isless((a.time, a.censored), (b.time,b.censored))

###############################################################
######################EVENT WINDOW#############################
###############################################################

immutable EventWindow
    t₀::Float64
    t₁::Float64
end

EventWindow{S<:Real}(t::S) = EventWindow(t, t)
EventWindow(ev::Event) = ev.censored ? EventWindow(ev.time,Inf) : EventWindow(ev.time,ev.time)

function Base.show(io::IO, obj::EventWindow)
    if obj.t₀ == obj.t₁
        print(io, obj.t₀)
    elseif obj.t₁ == Inf
        print(io, obj.t₀, "+")
    else
        print(io, "[$(obj.t₀),$(obj.t₁))")
    end
end

function Base.parse(T::Type{EventWindow},str::String)
    if str[end] == '+'
        t₀ = parse(str[1:(end-1)])
        return EventWindow(t₀,Inf)
    elseif str[end] == ')'
        substr = str[2:end-1]
        splitstr = split(substr,',')
        return EventWindow(parse.(splitstr)...)
    else
        t₀ = parse(str)
        return EventWindow(t₀, t₀)
    end
end

###############################################################
######################PRERPROCESSING###########################
###############################################################

immutable SortedEvents{R<:Real}
    events::Vector{Event{R}}
    fs::Vector{Int64}
    ls::Vector{Int64}
end

function SortedEvents(events)
    sorted = issorted(events) ? copy(events) : sort(events)
    return SortedEvents(sorted, find(firsts(sorted)), find(lasts(sorted)))
end

###############################################################
######################AFTRESP##################################
###############################################################
immutable AftResp{R<:Number, D<:Distribution}
    W::Array{R,2}
    X::Array{R,2}
    Xβ::Array{R,1}
    Θ::Array{R,1}
    dist::D
    M::Int64
    N::Int64
    I1::Array{Int64,1}
    I2::Array{Int64,1}
    τs::Array{R,2}
    cdfs::Array{R,2}
    Δcdf::Array{R,1}
    loglik::Array{R,1}
end

function AftResp(S, X, dist)
    M = length(dist.params)
    N = size(X,2)
    Xβ = zeros(size(X,1))
    Θ = ones(size(X,1))
    L = length(S)
    W = [(k == 1 ? s.t₀ : s.t₁) for s in S, k in 1:2]
    instants = [s.t₁ == s.t₀ for s in S]
    I1 = find(instants)
    I2 = find(.!(instants))
   return AftResp(W, X, Xβ, Θ, dist, M, N, I1, I2,
                   zeros(L,2), ones(L,2),
                   zeros(L), zeros(L))
end
###############################################################
######################INTCOEFS#################################
###############################################################

immutable IntCoefs{R<:Number, N}
    aux::Array{R,1}
    g::Array{SVector{N,R},1}
    hggt::Array{SVector{N,R},2}
end

function IntCoefs{N}(pdist::Distribution, degreetype::Val{N} = Val{50}())
    aux = auxvec(pdist)
    int_coefs = IntCoefs(aux,[SVector{N, Float64}(zeros(N)) for i in eachindex(pdist.params)],
    [SVector{N, Float64}(zeros(N)) for i in eachindex(pdist.params), j in eachindex(pdist.params)])
    for i in eachindex(pdist.params)
        int_coefs.g[i] = clenshaw_coefs(pdist, t -> ds(pdist,t, int_coefs, i), degreetype)
    end
    for i in eachindex(pdist.params), j in eachindex(pdist.params)
        int_coefs.hggt[i, j] = clenshaw_coefs(pdist, t -> ds(pdist,t, int_coefs, i)*ds(pdist,t,int_coefs, j) +
        d²s(pdist,t, int_coefs,i, j), degreetype)
    end
    return int_coefs
end


text:
using DataTableaux
using Base.Test

# construct a `DataTableau` object directly:
employees = ["Bob", "Ann", "Bob"]
hours = [10.0, 12.0, 15.0]
dt = DataTableau(columns = Any[employees, hours], names = [:employee, :hours])

# or, with automatic naming of columns:
dt = DataTableau(columns = Any[employees, hours])

# load a reduced Ames House Price data set as a `DataFrame`:
df = DataTableaux.load_reduced_ames()

# make integer types float (else considered categorical):
for ftr in names(df)
    if eltype(df[ftr]) <: Integer
        df[ftr] = convert(Array{Float64}, df[ftr])
    end
end

# split into train and test sets:
ntrain = round(Int, 0.8*size(df,1))
train = 1:ntrain
test = (ntrain + 1):size(df,1)
df_test = df[test,:]
df = df[train,:]

# construct a `DataTableau` object:
dt = DataTableau(df)
@test dt.names == names(df)

# recover the scheme encoding the transformation from `df` to `dt`:
s = FrameToTableauScheme(dt)

# transform the test data according to the same scheme:
dt_test = DataTableaux.transform(s, df_test)

# recover the original test `DataFrame` object from `dt`:
@test df == DataTableaux.inverse_transform(s, dt)

# construct an unfitted transformation scheme
s = FrameToTableauScheme()

# fit to the training `DataFrame` (return value:
DataTableaux.fit!(s, df)

# obtain the corresponding `DataTableaux` object:
dt2 = DataTableaux.transform(s, df)

# which amounts to the same construction given above:
@test dt2.raw == dt.raw

# alternatively, perform the preceding construction in one less step
# with `fit_transform!`:
s = FrameToTableauScheme()
dt3 = DataTableaux.fit_transform!(s, df)
@test dt3.raw == dt.raw

# test `getindex`
dt[10:12,:]


text:
"""
	FeedbackPolicy

Struct representing a feedback control policy.

# Fields
- `control`: function u = control(x) describing the control law.
"""
struct FeedbackPolicy <: Policy
    control
end

"""
	(k::FeedbackPolicy)(x)

Evaluate FeedbackPolicy at state x.
"""
function (k::FeedbackPolicy)(x)
    return k.control(x)
end


text:
using Plots
p = [ 10*rand(2) for i in 1:100 ]
x = [ p[i][1] for i in 1:100 ]
y = [ p[i][2] for i in 1:100 ]
scatter(x,y,
        xlabel="x",ylabel="y",label="",
        xlim=[-15,25],ylim=[-15,25],
        size=(400,400),framestyle=:box)
savefig("pbc1.pdf")


text:
"""
    appear(s::Symbol)

Appear can be used inside a [`SubAction`](@ref)

# Example
```
Action(101:200, (args...)->house_of_nicholas(); subactions = [
    SubAction(1:20, appear(:fade)),
    SubAction(81:100, disappear(:fade))
])
```
In this case the `house_of_nicholas` will fade in during the first 20 frames
of the [`Action`](@ref) so `101-120`.

# Arguments
- `s::Symbol`: the symbol defines the animation of appearance
    The only symbols that are currently supported are:
    - `:fade_line_width` which increases the line width up to the default value
       or the value specified by [`setline`](@ref)
    - `:fade` which increases the opcacity up to the default value
       or the value specified by [`setopacity`](@ref)
"""
function appear(s::Symbol)
    (video, action, subaction, rel_frame) ->
        _appear(video, action, subaction, rel_frame, Val(s))
end

function _appear(video, action, subaction, rel_frame, symbol::Val{:fade_line_width})
    # t is between 0 and 1
    t = (rel_frame - first(get_frames(subaction))) / (length(get_frames(subaction)) - 1)
    action.current_setting.mul_line_width = t
end

function _appear(video, action, subaction, rel_frame, symbol::Val{:fade})
    # t is between 0 and 1
    t = (rel_frame - first(get_frames(subaction))) / (length(get_frames(subaction)) - 1)
    action.current_setting.mul_opacity = t
end

function _appear(video, action, subaction, rel_frame, symbol::Val{:scale})
    # t is between 0 and 1
    t = (rel_frame - first(get_frames(subaction))) / (length(get_frames(subaction)) - 1)
    action.current_setting.mul_scale = t
end

"""
    disappear(s::Symbol)

Disappear can be used inside a [`SubAction`](@ref)

# Example
```
Action(101:200, (args...)->house_of_nicholas(); subactions = [
    SubAction(1:20, appear(:fade)),
    SubAction(81:100, disappear(:fade))
])
```
In this case the `house_of_nicholas` will fade out during the last 20 frames
of the [`Action`](@ref) so `181-200`.

# Arguments
- `s::Symbol`: the symbol defines the animation of disappearance
    The only symbols that are currently supported are:
    - `:fade_line_width` which descreases the line width up to the default value
        or the value specified by [`setline`](@ref)
    - `:fade` which decreases the opcacity up to the default value
        or the value specified by [`setopacity`](@ref)
"""
function disappear(s::Symbol)
    (video, action, subaction, rel_frame) ->
        _disappear(video, action, subaction, rel_frame, Val(s))
end

function _disappear(video, action, subaction, rel_frame, symbol::Val{:fade_line_width})
    # t is between 0 and 1
    t = (rel_frame - first(get_frames(subaction))) / (length(get_frames(subaction)) - 1)
    action.current_setting.mul_line_width = 1 - t
end

function _disappear(video, action, subaction, rel_frame, symbol::Val{:fade})
    # t is between 0 and 1
    t = (rel_frame - first(get_frames(subaction))) / (length(get_frames(subaction)) - 1)
    action.current_setting.mul_opacity = 1 - t
end

function _disappear(video, action, subaction, rel_frame, symbol::Val{:scale})
    # t is between 0 and 1
    t = (rel_frame - first(get_frames(subaction))) / (length(get_frames(subaction)) - 1)
    action.current_setting.mul_scale = 1 - t
end


text:
include("magpoly.jl")

using .magpoly
const mp = magpoly
using DelimitedFiles, Random, Distributions, LinearAlgebra

function generate_saw(n_samples::Int, samples_lag::Int)
    io = open("features.txt", "r")
    features = readdlm(io,Float64)
    close(io)
    n_mono = size(features, 1)

    burn_in = ceil(Int, 20*n_mono^(1.11))
    #burn_in = 3
    pol = mp.Magnetic_polymer(n_mono, 1.0, 0.0, 1.0)
    mp.initialize_poly!(pol)

    for i_mono in 1:n_mono
        for j in 1:3
            pol.trial_coord[j,i_mono] = pol.coord[j,i_mono]
        end
    end

    hash_base = 2*n_mono +1

    # equilibration of the chain
    for i in 1:burn_in
        pivot = rand(2:n_mono-1)
        p_move = rand(1:47)
        mp.try_pivot!(pivot, p_move, pol.coord, pol.trial_coord, pol.n_mono)
        still_saw = mp.checksaw!(pivot, pol, hash_base)
        #println(still_saw)
        if still_saw
            for i_mono in 1:n_mono
                for j in 1:3
                    pol.coord[j,i_mono] = pol.trial_coord[j,i_mono]
                end
            end
        end
        empty!(pol.hash_saw)
    end
    println("End of burn in!")
    n_acc = 0

    n_steps = n_samples*samples_lag

    saw_configs = zeros(Int, n_samples, 3, n_mono)

    for i in 1:n_steps
        pivot = rand(2:n_mono-1)
        p_move = rand(1:47)
        if i%10000 == 0
            println("Pivot: ", pivot)
            println("Attmpted move: ", p_move)
            println("step number: ", i, "\n")
        end
        mp.try_pivot!(pivot, p_move, pol.coord, pol.trial_coord, pol.n_mono)
        still_saw = mp.checksaw!(pivot, pol, hash_base)
        if still_saw
            for i_mono in 1:n_mono
                for j in 1:3
                    pol.coord[j,i_mono] = pol.trial_coord[j,i_mono]
                end
            end
            n_acc += 1
        end

        if i%samples_lag == 0
            i_sample = div(i,samples_lag)
            for i_mono in 1:n_mono
                for j in 1:3
                    saw_configs[i_sample,j,i_mono] = pol.coord[j,i_mono]
                end
            end
        end
        empty!(pol.hash_saw)
    end
    println("Fraction of accepted pivots: ", n_acc/n_steps)

    !isdir("SAW_data") && mkdir("SAW_data")
    io = open("SAW_data/saw_configs.txt", "w")
    for i_sample in 1:n_samples
        writedlm(io, saw_configs[i_sample,:,:])
    end
    close(io)
end

# This function generates samples from the posterior distribution of the weights
# Assuming that the data were generated with zero spins coupling i.e. a logistic regression problem
# The sampling is done with a simple metroplis hastings sampler.
function vanilla_sampler(n_samples::Int, data_file::String, features_file::String)
    io = open(data_file,"r")
    data_spins = readdlm(io, Int64; header=true)[1]
    close(io)
    n_mono = size(data_spins, 2)
    n_data = size(data_spins, 1)
    println(n_mono)
    println(n_data)
    println(typeof(data_spins))

    io = open(features_file, "r")
    features = readdlm(io,Float64)
    close(io)
    n_mono = size(features, 1)
    n_feats = size(features, 2)

    weights_series = zeros(n_samples,2)
    energies = zeros(n_samples)

    w_std = 3.0
    weights = zeros(n_feats)
    trial_weights = zeros(n_feats)
    weights = randn(2) .* w_std  # vector of gaussian numbers with standard deviation = w_std
    weights_series[1,:] .= weights
    delta_w = 0.1

    fields = zeros(n_mono)
    for i in 1:n_feats
        fields .+= features[:,i] .* weights[i]
    end

    #energy = -0.5*(weights[1]^2 + weights[2]^2)/(w_std^2)  # prior contribution
    energy = 0.0
    for i_data in 1:n_data
        for i_mono in 1:n_mono
            energy += data_spins[i_data, i_mono]*fields[i_mono]*0.5 - log(1+exp(0.5*fields[i_mono]))
        end
    end
    energies[1] = energy

    w_acceptance = 0.0
    n_acc = 0

    for i_sample in 2:n_samples
        trial_weights .= weights .+ randn(2).* delta_w
        fields = zeros(n_mono)
        for i in 1:n_feats
            fields .+= features[:,i] .* trial_weights[i]
        end

        #trial_energy = -0.5*(trial_weights[1]^2 + trial_weights[2]^2)/(w_std^2)
        trial_energy = 0.0
        for i_data in 1:n_data
            for i_mono in 1:n_mono
                trial_energy += data_spins[i_data, i_mono]*fields[i_mono]*0.5 - log(1+exp(0.5*fields[i_mono]))
            end
        end

        delta_energy = trial_energy - energy
        delta_energy >=0 ? w_acceptance=1 : w_acceptance=exp(delta_energy)

        if i_sample%1000 == 0
            println(i_sample)
            println("w1: ",weights[1]," ---> ",trial_weights[1])
            println("w2: ",weights[2]," ---> ",trial_weights[2])
            println("delta_syn_like: ", delta_energy)
            println("acceptance: ", w_acceptance)
        end

        if rand() <= w_acceptance
            energy = trial_energy
            weights .= trial_weights
            n_acc += 1
        end

        energies[i_sample] = energy
        weights_series[i_sample,:] .= weights
    end

    println("Acceptance fraction: ", n_acc/(n_samples-1))

    !isdir("IS_data") && mkdir("IS_data")
    open("IS_data/weights.txt", "w") do io
        writedlm(io, weights_series)
    end
    open("IS_data/neg_energies.txt", "w") do io
        writedlm(io, energies)
    end
end


text:
# plot.jl: functions for plotting
using PyCall: PyObject, PyDict
using PyPlot: matplotlib, Figure, plt, withfig

# plotting parameters
rcParams = PyDict(matplotlib."rcParams")
rcParams["font.sans-serif"] = ["Arial"];
rcParams["font.size"] = 12;
colpalette = ["red", "orange", "lightblue", "blue"];

# functions for plotting
function plot_strat_dist(
    pop::Population,                        # need pop.sets.conditional
    data::Array{Float64, 2},                # data to plot
    ax::PyObject,                           # axis on which to plot
    ylabel::String="Frequency",             # y-axis label # change for cumulative distribution
    gen_start::Int64=1,                     # starting time step
    gen_skip::Int64=1,                      # time step interval
    generations::Int64=size(data,1),        # final time step / total sim length
    xlabel::String="Time step",             # x-axis label
    colpalette::Array{String, 1}=colpalette # color palette
    )
    # function to plot strategy distribution
    strat_ids = ["DD","DC","CD","CC"] # names of strategies

    [ax.plot(gen_start:gen_skip:generations, data[gen_start:gen_skip:generations,x],
        label=strat_ids[x], linewidth=0.75,
        color=colpalette[x]) for x in 1:length(pop.all_strategies)]
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_xlim([gen_start,generations])
    ax.set_ylim([-0.05,1.05])
    ax.legend(bbox_to_anchor=(1.15, 0.5), loc=5, borderaxespad=0., fontsize=10)

end

# another method, with effective coopeartion
function plot_strat_dist(
    pop::Population,                        # need pop.sets.conditional
    data::Array{Float64, 2},                # data to plot
    data2::Array{Float64, 1},               # data to plot
    ax::PyObject,                           # axis on which to plot
    ylabel::String="Frequency",             # y-axis label # change for cumulative distribution
    gen_start::Int64=1,                     # starting time step
    gen_skip::Int64=1,                      # time step interval
    generations::Int64=size(data,1),        # final time step / total sim length
    xlabel::String="Time step",             # x-axis label
    colpalette::Array{String, 1}=colpalette # color palette
    )
    # function to plot strategy distribution
    strat_ids = ["DD","DC","CD","CC","Cooperation"] # names of strategies

    [ax.plot(gen_start:gen_skip:generations, data[gen_start:gen_skip:generations,x],
        label=strat_ids[x], linewidth=0.75,
        color=colpalette[x]) for x in 1:length(pop.all_strategies)]
    ax.plot(gen_start:gen_skip:generations, data2[gen_start:gen_skip:generations],
        label=strat_ids[5], linewidth=0.5, color="Black", linestyle="dotted")
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_xlim([gen_start,generations])
    ax.set_ylim([-0.05,1.05])
    ax.legend(bbox_to_anchor=(1.2, 0.5), loc=5, borderaxespad=0., fontsize=8)

end

function plot_set_dist(
    pop::Population,                     # need pop.sets.conditional
    data::Array{Int64, 2},               # data to plot
    ax::PyObject,                        # axis on which to plot
    ylabel::String="No. of individuals", # y-axis label # change for cumulative distribution
    gen_start::Int64=1,                  # starting time step
    gen_skip::Int64=1,                   # time step interval
    generations::Int64=size(data,1),     # final time step / total sim length
    xlabel::String="Time step"           # x-axis label
    )
    # function to plot strategy distribution
    set_ids = ["Set $(x)" for x in 1:pop.sets.M] # define names of sets for legend

    [ax.plot(gen_start:gen_skip:generations, data[gen_start:gen_skip:generations,x],
        label=set_ids[x], linewidth=0.75) for x in 1:(pop.sets.M)]
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_xlim([gen_start,generations])
    ax.set_ylim([-1,(pop.sets.N)+1])
    ax.legend(bbox_to_anchor=(1.15, 0.5), loc=5, borderaxespad=0., fontsize=8)

end

function plot_opn_dist(
    pop::Population,                          # need pop.sets.conditional
    data::Array{Float64, 2},                  # data to plot
    ax::PyObject,                             # axis on which to plot
    ylabel::String="Frequency",               # y-axis label # change for cumulative distribution
    gen_start::Int64=1,                       # starting time step
    gen_skip::Int64=1,                        # time step interval
    generations::Int64=size(data,1),          # final time step / total sim length
    xlabel::String="Time step"                # x-axis label
    )
    # function to plot average and variance of total opinions
    opinion_ids::Array{String, 1} = ["$(base3todec(x, pop))" for x in pop.all_opinions]

    # ax.plot(gen_start:gen_skip:generations, data[gen_start:gen_skip:generations], linewidth=0.25)
    [ax.plot(gen_start:gen_skip:generations,
        data[gen_start:gen_skip:generations,x],
        label=opinion_ids[x], linewidth=0.75) for x in 1:length(pop.all_opinions)]
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_xlim([gen_start,generations])
    ax.set_ylim([-0.05,1.05])
    ax.legend(bbox_to_anchor=(1.15, 0.5), loc=5, borderaxespad=0., fontsize=7) # note fontsize
end

function plot_total_opn(
    pop::Population,                          # need pop.sets.conditional
    data::Array{Float64, 1},                  # data to plot
    ax::PyObject,                             # axis on which to plot
    ylabel::String="Average\ntotal opinion",  # y-axis label # change for cumulative distribution
    gen_start::Int64=1,                       # starting time step
    gen_skip::Int64=1,                        # time step interval
    generations::Int64=size(data,1),          # final time step / total sim length
    xlabel::String="Time step"                # x-axis label
    )
    # function to plot average and variance of total opinions
    ax.plot(gen_start:gen_skip:generations,
            data[gen_start:gen_skip:generations],
            label="value", linewidth=0.75, color="purple")
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_xlim([gen_start,generations])
    ax.set_ylim([min(0,minimum(data))-0.1, max(pop.sets.K,maximum(data))+0.1])
    ax.legend(bbox_to_anchor=(1.15, 0.5), loc=5, borderaxespad=0., fontsize=8)
end

function base3todec( base3::Int64, pop::Population )
    # convert opinion in int to vector of 1, 0, -1
    reverse( digits( base3 + vectobase3( repeat([1],pop.sets.M) ), base = 3, pad = pop.sets.M )) .- 1
end

function plotter( tracker::Tracker, gen_start::Int64, gen_skip::Int64, gen_end::Int64 )
    # function to make a quick plot
    fig, axs = plt.subplots(6, 2, figsize = (16,12));

    plot_strat_dist( tracker.pop, tracker.ens_strat_dist, tracker.ens_total_cooperation, axs[1], "Frequency", gen_start, gen_skip, gen_end );
    plot_opn_dist( tracker.pop, tracker.ens_opn_dist, axs[3], "Frequency", gen_start, gen_skip, gen_end); # opinion distribution
    plot_set_dist( tracker.pop, tracker.ens_pop_dist, axs[9], "No. of individuals", gen_start, gen_skip, gen_end );
    plot_strat_dist( tracker.pop, tracker.ens_running_strat_dist, axs[2], "Running\navg frequency", gen_start, gen_skip, gen_end );
    plot_total_opn( tracker.pop, tracker.ens_total_means, axs[7], "Average\ntotal opinion", gen_start, gen_skip, gen_end );
    plot_total_opn( tracker.pop, tracker.ens_total_vars, axs[8], "Variance\ntotal opinion", gen_start, gen_skip, gen_end );
    # plot_total_opn( tracker.pop, tracker.ens_total_cooperation, axs[7], "Effective cooperation", gen_start, gen_skip, gen_end );
    plot_total_opn( tracker.pop, tracker.ens_opn_simpson, axs[4], "Opinion diversity\n(1 = low)", gen_start, gen_skip, gen_end );
    plot_total_opn( tracker.pop, tracker.ens_sets_simpson, axs[10], "Set diversity\n (1 = low)", gen_start, gen_skip, gen_end );
    plot_total_opn( tracker.pop, tracker.ens_hamming_ind, axs[5], "Pairwise\nHamming distance", gen_start, gen_skip, gen_end );
    plot_total_opn( tracker.pop, tracker.ens_hamming_party, axs[6], "Within-party\nHamming distance",
    gen_start, gen_skip, gen_end );
    plot_total_opn( tracker.pop, tracker.ens_cityblock_ind, axs[11], "Pairwise\nCityblock distance", gen_start, gen_skip, gen_end );
    plot_total_opn( tracker.pop, tracker.ens_cityblock_party, axs[12], "Within-party\nCityblock distance", gen_start, gen_skip, gen_end );

    return fig
end