using Flux
using AlphaGo
using AlphaGo: select_leaf, incorporate_results!, child_U, inject_noise!,
                Q, N, child_Q, initialize_game!, tree_search!, MCTSRules, extract_data
using AlphaGo: pass_move!, play_move!, PlayerMove, BLACK, WHITE, PlayerMove,
                to_kgs, score, from_kgs

include("test_utils.jl")

struct DummyNet
  fake_priors
  fake_value

  function DummyNet(env; fake_priors = nothing, fake_value = 0)
    if fake_priors == nothing
      fake_priors = ones(env.action_space) / (env.action_space)
    end
    new(fake_priors, fake_value)
  end
end

(dn::DummyNet)(position::Position) = param(dn.fake_priors),
                                          param(dn.fake_value)

function (dn::DummyNet)(positions = nothing)
  if positions == nothing || length(positions) == 0
    throw(ArgumentError("No positions passed!"))
  end
  len = length(positions)
  param(repeat(dn.fake_priors, outer = [1, len])),
          param(repeat([dn.fake_value], len))
end

env = GoEnv(9)

ALMOST_DONE_BOARD = load_board("""
                              .XO.XO.OO
                              X.XXOOOO.
                              XXXXXOOOO
                              XXXXXOOOO
                              .XXXXOOO.
                              XXXXXOOOO
                              .XXXXOOO.
                              XXXXXOOOO
                              XXXXOOOOO
                              """, env)

SEND_TWO_RETURN_ONE = Position(env,
    board = ALMOST_DONE_BOARD,
    n = 70,
    komi = 2.5,
    caps = (1, 4),
    ko = nothing,
    recent = [PlayerMove(BLACK, (1, 2)),
    PlayerMove(WHITE, (1, 9))],
    to_play = BLACK
    );

function initialize_basic_player(env)
  player = MCTSPlayer(env, DummyNet(env))
  initialize_game!(player)
  first_node = select_leaf(player.root)
  p, v = player.network(player.root.position)
  incorporate_results!(first_node, p.data, v.data, player.root)
  return player
end

function initialize_almost_done_player(env)
  probs = ones(env.action_space) * 0.001
  probs[3:5] .= 0.2  # some legal moves along the top.
  probs[end] = 0.2  # passing is also ok
  net = DummyNet(env; fake_priors = probs)
  player = MCTSPlayer(env, net)
  # root position is white to play with no history == white passed.
  initialize_game!(player, SEND_TWO_RETURN_ONE)
  return player
end

@testset "MCTSPlayer" begin
  # Tromp taylor means black can win if we hit the move limit.
  TT_FTW_BOARD = load_board("""
                            .XXOOOOOO
                            X.XOO...O
                            .XXOO...O
                            X.XOO...O
                            .XXOO..OO
                            X.XOOOOOO
                            .XXOOOOOO
                            X.XXXXXXX
                            XXXXXXXXX
                            """, env)

  @testset "inject_noise" begin
    player = initialize_basic_player(env)
    sum_priors = sum(player.root.child_prior)
    # dummyNet should return normalized priors.
    @test sum_priors ≈ 1
    u = child_U(player.root)
    @test all(u .== u[1])

    inject_noise!(player.root)
    new_sum_priors = sum(player.root.child_prior)
    # priors should still be normalized after injecting noise
    @test sum_priors ≈ new_sum_priors

    # With dirichelet noise, majority of density should be in one node.
    max_p = maximum(player.root.child_prior)
    @test max_p > 3 / (env.action_space)
  end

  @testset "pick_moves" begin
    player = initialize_basic_player(env)
    root = player.root
    root.child_N[to_flat((3, 1), player.root.position)] = 10
    root.child_N[to_flat((2, 1), player.root.position)] = 5
    root.child_N[to_flat((4, 1), player.root.position)] = 1

    root.position.n = env.action_space  # move 81, or 361, or... Endgame.

    # Assert we're picking deterministically
    @test root.position.n > player.τ_threshold
    move = pick_move(player)
    @test move == (3, 1)

    # But if we're in the early part of the game, pick randomly
    root.position.n = 3
    @test player.root.position.n ≤ player.τ_threshold

    #TODO: complete this test
    #with mock.patch('random.random', lambda: .5)
    #move = pick_move(player)
    #@test move == (3, 1)

    #with mock.patch('random.random', lambda: .99):
    #move = pick_move(player)
    #@test move == (4, 1)
  end

 @testset "dont_pass_if_losing" begin
    player = initialize_almost_done_player(env)

    # check -- white is losing.
    @test score(player.root.position) == -0.5

    for i = 1:20
      tree_search!(player)
    end

    # uncomment to debug this test
    #println(describe(player.root))

    # Search should converge on D9 as only winning move.
    flattened = to_flat(from_kgs("D9", env), player.root.position)
    best_move = findmax(player.root.child_N)[2]
    @test best_move == flattened
    # D9 should have a positive value
    @test Q(player.root.children[flattened]) > 0
    @test N(player.root) ≥ 20
    # passing should be ineffective.
    @test child_Q(player.root)[end] < 0
    # no virtual losses should be pending
    @test assertNoPendingVirtualLosses(player.root)
    # uncomment to debug this test
    #println(describe(player.root))
  end

  @testset "parallel_tree_search" begin
    player = initialize_almost_done_player(env)
    # check -- white is losing.
    @assert score(player.root.position) == -0.5
    # initialize the tree so that the root node has populated children.
    tree_search!(player, 1)
    # virtual losses should enable multiple searches to happen simultaneously
    # without throwing an error...
    for i = 1:6
      tree_search!(player, 10)
    end
    # uncomment to debug this test
    # print(player.root.describe())

    # Search should converge on D9 as only winning move.
    flattened = to_flat(from_kgs("D9", env), player.root.position)
    best_moves = findall(x -> x .== maximum(player.root.child_N), player.root.child_N)
    @test flattened ∈ best_moves
    # D9 should have a positive value
    @test Q(player.root.children[flattened]) > 0
    @test N(player.root) ≥ 20
    # passing should be ineffective.
    child_Q(player.root)[end] < 0
    # no virtual losses should be pending
    @test assertNoPendingVirtualLosses(player.root)
  end

  @testset "ridiculously_parallel_tree_search" begin
    player = initialize_almost_done_player(env)
    # Test that an almost complete game
    # will tree search with # parallelism > # legal moves.
    for i = 1:10
      tree_search!(player, 50)
    end
    @test assertNoPendingVirtualLosses(player.root)
  end

  @testset "long_game_tree_search" begin
    mcts_rules = MCTSRules(env)
    player = MCTSPlayer(env, DummyNet(env))
    endgame = Position(env,
        board = TT_FTW_BOARD,
        n = mcts_rules.max_game_length - 2,
        komi = 2.5,
        ko = nothing,
        recent = [PlayerMove(BLACK, (1, 2)),
                PlayerMove(WHITE, (1, 9))],
        to_play = BLACK
    )
    initialize_game!(player, endgame)

    # Test that MCTS can deduce that B wins because of TT-scoring
    # triggered by move limit.
    for i = 1:10
      tree_search!(player, 8)
    end
    @test assertNoPendingVirtualLosses(player.root)
    @test Q(player.root) > 0
  end

  @testset "cold_start_parallel_tree_search" begin
    # Test that parallel tree search doesn't trip on an empty tree
    player = MCTSPlayer(env, DummyNet(env, fake_value = 0.17))
    initialize_game!(player)
    @test N(player.root) == 0
    @test !player.root.is_expanded
    tree_search!(player, 4)
    @test assertNoPendingVirtualLosses(player.root)
    # Even though the root gets selected 4 times by tree search, its
    # final visit count should just be 1.
    N(player.root) == 1
    # 0.085 = average(0, 0.17), since 0 is the prior on the root.
    @test Q(player.root) ≈ 0.085
  end

  @testset "tree_search_failsafe" begin
    # Test that the failsafe works correctly. It can trigger if the MCTS
    # repeatedly visits a finished game state.
    probs = ones(env.N * env.N + 1) * 0.001
    probs[end] = 1  # Make the dummy net always want to pass
    player = MCTSPlayer(env, DummyNet(env, fake_priors = probs))
    pass_position = pass_move!(Position(env))
    initialize_game!(player, pass_position)
    tree_search!(player, 1)
    @test assertNoPendingVirtualLosses(player.root)
  end

  @testset "only_check_game_end_once" begin
    # When presented with a situation where the last move was a pass,
    # and we have to decide whether to pass, it should be the first thing
    # we check, but not more than that.

    white_passed_pos = pass_move!(
                        play_move!(
                          play_move!(
                            play_move!(
                              Position(env), (4,4) # b plays
                                ), (4,5)  # w plays
                              ), (5,4)  # b plays
                            ) # w passes - if B passes too, B would lose by komi.
                          )

    player = MCTSPlayer(env, DummyNet(env))
    initialize_game!(player, white_passed_pos)
    # initialize the root
    for i = 1:15
      tree_search!(player)
    end
    # explore a child - should be a pass move.
    #tree_search!(player, 10)
    pass_move = env.N * env.N + 1
    @test N(player.root.children[pass_move]) == 1
    @test player.root.child_N[pass_move] == 1
    tree_search!(player)
    # check that we didn't visit the pass node any more times.
    @test player.root.child_N[pass_move] == 1
  end

  @testset "extract_data_normal_end" begin
    player = MCTSPlayer(env, DummyNet(env))
    initialize_game!(player)
    tree_search!(player)
    play_move!(player, nothing)
    tree_search!(player)
    play_move!(player, nothing)
    @test is_done(player.root)
    set_result!(player, AlphaGo.result(player.root.position), false)

    positions, πs, results = extract_data(player)
    @test length(positions) == length(πs) == length(results) == 2
    position, pi, result = positions[1], πs[1], results[1]
    # White wins by komi
    @test result == WHITE
    @test player.result_string == "W+$(player.root.position.komi)"
  end

  @testset "extract_data_resign_end" begin
    player = MCTSPlayer(env, DummyNet(env))
    initialize_game!(player)
    tree_search!(player)
    play_move!(player, (1, 1))
    tree_search!(player)
    play_move!(player, nothing)
    tree_search!(player)
    # Black is winning on the board
    @test AlphaGo.result(player.root.position) == BLACK
    # But if Black resigns
    set_result!(player, WHITE, true)

    data = extract_data(player)
    position, pi, result = data[1], data[2], data[3]
    # Result should say White is the winner
    @test result[1] == WHITE
    @test player.result_string == "W+R"
  end
end


Network = NetLib.SimpleNet

netparams = NetLib.SimpleNetHP(
  width=200,
  depth_common=6,
  use_batch_norm=true,
  batch_norm_momentum=1.)

self_play = SelfPlayParams(
  sim=SimParams(
    num_games=1000,
    num_workers=128,
    batch_size=128,
    use_gpu=false,
    reset_every=4,
    flip_probability=0.,
    alternate_colors=false),
  mcts=MctsParams(
    num_iters_per_turn=400,
    cpuct=1.0,
    temperature=ConstSchedule(1.0),
    dirichlet_noise_ϵ=0.2,
    dirichlet_noise_α=1.0))

arena = ArenaParams(
  sim=SimParams(
    num_games=100,
    num_workers=100,
    batch_size=100,
    use_gpu=false,
    reset_every=1,
    flip_probability=0.,
    alternate_colors=true),
  mcts = MctsParams(
    self_play.mcts,
    temperature=ConstSchedule(0.3),
    dirichlet_noise_ϵ=0.1),
  update_threshold=0.00)

learning = LearningParams(
  use_gpu=false,
  samples_weighing_policy=LOG_WEIGHT,
  l2_regularization=1e-4,
  optimiser=CyclicNesterov(
    lr_base=1e-3,
    lr_high=1e-2,
    lr_low=1e-3,
    momentum_high=0.9,
    momentum_low=0.8),
  batch_size=32,
  loss_computation_batch_size=2048,
  nonvalidity_penalty=1.,
  min_checkpoints_per_epoch=0,
  max_batches_per_checkpoint=5_000,
  num_checkpoints=1)

params = Params(
  arena=arena,
  self_play=self_play,
  learning=learning,
  num_iters=2,
  memory_analysis=MemAnalysisParams(
    num_game_stages=4),
  ternary_rewards=true,
  use_symmetries=true,
  mem_buffer_size=PLSchedule(80_000))

benchmark_sim = SimParams(
  arena.sim;
  num_games=400,
  num_workers=100,
  batch_size=100)

benchmark = [
  Benchmark.Duel(
    Benchmark.Full(self_play.mcts),
    Benchmark.MctsRollouts(self_play.mcts),
    benchmark_sim)
  ]

experiment = Experiment(
  "ospiel_ttt", GameSpec(), params, Network, netparams, benchmark)


using ShaderToy, GLAbstraction, FileIO, ColorTypes
shadertoy("submerged.frag", Dict{Symbol, Any}(
	:iChannel0 => Texture(rand(Float32, 64,64), x_repeat=:repeat, minfilter=:linear),
))


using Colors
using FixedPointNumbers
using Test
using Random
using NeuroCore.ColorChannels
using AxisIndices
using MappedArrays
using TimeAxes

using AxisIndices.PaddedViews: filltype, paddedviews, PaddedView

using NeuroCore.ColorChannels: GenericImage, GenericGrayImage
using NeuroCore.ColorChannels: NumberLike, Pixel

include("channels.jl")
include("views.jl")
include("convert_reinterpret.jl")
include("traits.jl")

using FFTW

@testset "functions" begin
    ag = rand(Gray{Float32}, 4, 5)
    ac = rand(RGB{Float32}, 4, 5)
    for (f, args) in ((fft, (ag,)), (fft, (ag, 1:2)), (plan_fft, (ag,)),
                      (rfft, (ag,)), (rfft, (ag, 1:2)), (plan_rfft, (ag,)),
                      (fft, (ac,)), (fft, (ac, 1:2)), (plan_fft, (ac,)),
                      (rfft, (ac,)), (rfft, (ac, 1:2)), (plan_rfft, (ac,)))
        ret = @test_throws ErrorException f(args...)
        @test occursin("channel_view", ret.value.msg)
        @test occursin(eltype(args[1])<:Gray ? "1:2" : "2:3", ret.value.msg)
    end
    for (a, dims) in ((ag, 1:2), (ac, 2:3))
        @test ifft(fft(channel_view(a), dims), dims) ≈ channel_view(a)
        ret = @test_throws ErrorException rfft(a)
        @test occursin("channel_view", ret.value.msg)
        @test occursin("$dims", ret.value.msg)
        @test irfft(rfft(channel_view(a), dims), 4, dims) ≈ channel_view(a)
    end

    a = [RGB(1,0,0), RGB(0,1,0), RGB(0,0,1)]
    @test a' == [RGB(1,0,0) RGB(0,1,0) RGB(0,0,1)]

    a = [RGB(1,0,0) RGB(0,1,0); RGB(0,0,1) RGB(0, 0, 0)]
    @test a' == [RGB(1,0,0) RGB(0,0,1); RGB(0,1,0) RGB(0, 0, 0)]
end


############################################################################################
function init(
    ::Type{HMC},
    config::ConfigHMC,
    objective::Objective,
    difftune::AbstractDifferentiableTune,
)
    return HMC(
        ModelWrappers.log_density_and_gradient(objective, difftune),
        difftune,
        config.energy,
        StepNumberTune(
            config.stepnumber.stepnumberadaption, config.stepnumber.steps, config.stepnumber.maxsteps, config.stepnumber.∫dt
        ),
    )
end

function infer(
    _rng::Random.AbstractRNG,
    diagnostics::Type{AbstractDiagnostics},
    kernel::HMC,
    model::ModelWrapper,
    data::D,
) where {D}
    return DiagnosticsHMC{model.info.reconstruct.default.output}
end
############################################################################################
#export
export init, infer


using ACEds
using ACE
using ACE.RPI.Rotations3D
#using ACEds.Rot3D2
using StaticArrays

ll = @SVector [1,1]
L=1
#a= ACEds.Rot3D2._mrange(ll, L)

#for k in ACEds.Rot3D2._mrange(ll, 1)
#    print(k,"\n")
#end

a = @SVector [1,1.,1]

len=4
A = zeros(SArray{Tuple{3},Float64,1,3}, len, len)
for i=1:3
    A[1,2] += @SVector [1,1.,1]
end

A[1,2]
#ACEds.Rot3D2.get0(1)

#ACEds.Rot3D2.Rot3DCoeffs(0)


######################################################################
# plotting.jl
# Handles all the calls to PyPlot to generate plots
#
# hold is deprecated as of Matplotlib 2.0
#  Figures now default to hold (true)
#  I think we should adhere to this for future compliance
######################################################################

"""
`plot(m::SearchDomain, f::AbstractFilter, x::Vehicle)`

Plots the belief, jammer, and vehicles.
"""
function plot(m::SearchDomain, f::AbstractFilter, x::Vehicle; show_mean::Bool=false, show_cov::Bool=false, alpha=0.2, color="k")
	plot(m, f, (x.x, x.y, x.heading); show_mean=show_mean, show_cov=show_cov, alpha=alpha, color=color)
end
# allow plotting of multiple vehicles with a single shared filter
function plot(m::SearchDomain, f::AbstractFilter, vx::Vector{Vehicle})
	plot(m, f, vx[1])
	num_vehicles = length(vx)
	for i = 2:num_vehicles
		plot_vehicle(m, vx[i])
	end
end
function plot(m::SearchDomain, f::AbstractFilter, p::Pose; show_mean::Bool=false, show_cov::Bool=false, alpha=1.0, color="b")
	plot_theta(m)
	#hold(true) # deprecated
	plot_vehicle(m, p; color=color)
	if show_mean
		plot_mean(centroid(f), color=color)
	end
	if show_cov
		plot(m, centroid(f), covariance(f), color=color)
	end
	plot(m, f, alpha=alpha)
	return # so it doesn't spit out result of axis
end

function plot(m::SearchDomain, b::Matrix{Float64}, x::Vehicle; show_mean::Bool=false, show_cov::Bool=false, alpha=1.0, color="b", obs=nothing)
	plot(m, b, (x.x, x.y, x.heading); show_mean=show_mean, show_cov=show_cov, alpha=alpha, color=color, obs=obs)
end

function plot(m::SearchDomain, b::Matrix{Float64}, p::Pose; show_mean::Bool=false, show_cov::Bool=false, alpha=1.0, color="b", obs=nothing)
	plot_theta(m)
	#hold(true)  # deprecated
	plot_vehicle(m, p; color=color)
	#title_string = "plot: "
	title_string = ""
	if obs != nothing
		# TODO: just temporary for rss submission
		#title_string = "$(title_string)o = $(obs), "
		title_string = "t = $(obs) s"
	end
	if show_mean
		plot_mean(centroid(b, m.length), color=color)
	end
	if show_cov
		cov = covariance(b,m.length)
		evals = eigvals(cov)
		e1 = round(evals[1],1)
		e2 = round(evals[2],1)
		plot(m, centroid(b,m.length), cov, color=color)
		title_string = "$(title_string)e1 = $e1, e2 = $e2"
		#title("e1 = $e1, e2 = $e2")
	end
	title(title_string, fontname="Times New Roman", fontsize=12)
	plot(m, b, alpha=alpha)
	return # so it doesn't spit out result of axis
end

function plot(m::SearchDomain, f::AbstractFilter; alpha=1.0)
	error(typeof(f), " has not implemented `plot`(::SearchDomain, ::AbstractFilter)")
end

function plot(m::SearchDomain, f::DF; alpha=1.0, cmap="Greys")
	plot(m, f.b, alpha=alpha, cmap=cmap)
end
function plot(m::SearchDomain, b::Matrix{Float64}; alpha=1.0, cmap="Greys")

	# alpha changed by LD for rss submission 2017
	#alpha = 0.5
	#alpha = 0.6
    # commented out for cdc 2018

	a = [0,m.length,0,m.length]
	imshow(b', interpolation="none",cmap=cmap,origin="lower",extent=a,vmin=0, alpha=alpha)
	#imshow(b', cmap=cmap,origin="lower",extent=a,vmin=0, alpha=alpha)
	labels()
	axis(a)
	tick_params(direction="in")
end

plot(m::SearchDomain, f::EKF; alpha=0.1) = plot(m, f.mu, f.Sigma)
plot(m::SearchDomain, f::UKF; alpha=0.1) = plot(m, f.mu, f.Sigma)

# Plots mean and 95% confidence ellipse
function plot(m::SearchDomain, mu::LocTuple, Sigma::Matrix{Float64}; color="b")
	plot(m, [mu[1],mu[2]], Sigma, color=color)
end
function plot(m::SearchDomain, mu::Vector{Float64}, Sigma::Matrix{Float64}; color="b")
	a = [0,m.length,0,m.length]

	Sigma_half = sqrtm(Sigma)
	Sh11 = Sigma_half[1,1]
	Sh12 = Sigma_half[1,2]
	Sh21 = Sigma_half[2,1]
	Sh22 = Sigma_half[2,2]
	m1 = mu[1]
	m2 = mu[2]

	thetas = 0:.02:2pi
	num_theta = length(thetas)
	xvals = zeros(num_theta)
	yvals = zeros(num_theta)
	c = sqrt(-2.0*log(0.05))
	for i = 1:num_theta
		theta = thetas[i]
		s_theta = sin(theta)
		c_theta = cos(theta)
		xvals[i] = c*(Sh11*s_theta + Sh12*c_theta) + m1
		yvals[i] = c*(Sh21*s_theta + Sh22*c_theta) + m2
	end
	plot(xvals, yvals, color)
	plot(m1, m2, "$(color)x", ms=10)

	labels()
	#axis("square")
	axis(a)
	tick_params(direction="in")
end

#function plot_ellipse()
#end

function plot(m::SearchDomain, f::PF; alpha=0.2, cmap="Greys", color="k")
	mark_size = 12
	a = [0,m.length,0,m.length]
	x = zeros(f.n)
	y = zeros(f.n)
	for i = 1:f.n
		x[i] = particle(f,i)[1]
		y[i] = particle(f,i)[2]
	end
	#scatter(x,y,c=f.W,cmap="Greys",vmin=0)
	#scatter(x,y,c=f.W,cmap="Blues",vmin=0, alpha=0.2)
	scatter(x, y, s=1, alpha=alpha, c=color)
	#xmean, ymean = centroid(f)
	#plot(xmean, ymean, "gx", ms=10, mew=2)
	labels()
	axis("scaled")
	axis(a)
	tick_params(direction="in")
end
function plot(m::SearchDomain, b::Vector{Tuple{Float64,Float64}}; alpha=0.2, cmap="Greys", color="k")
	plot_theta(m)
	mark_size = 12
	a = [0,m.length,0,m.length]
    x = zeros(length(b))
    y = zeros(length(b))
    for i = 1:length(b)
        x[i] = b[i][1]
        y[i] = b[i][2]
	end
	#scatter(x,y,c=f.W,cmap="Greys",vmin=0)
	#scatter(x,y,c=f.W,cmap="Blues",vmin=0, alpha=0.2)
	scatter(x, y, s=1, alpha=alpha, c=color)
	#xmean, ymean = centroid(f)
	#plot(xmean, ymean, "gx", ms=10, mew=2)
	labels()
	axis("scaled")
	axis(a)
	tick_params(direction="in")
end
function plot(m::SearchDomain, b::Vector{Tuple{Float64,Float64}}, p; alpha=0.2, cmap="Greys", color="k")
	plot_theta(m)
	plot_vehicle(m, p; color=color)
	mark_size = 12
	a = [0,m.length,0,m.length]
    x = zeros(length(b))
    y = zeros(length(b))
    for i = 1:length(b)
        x[i] = b[i][1]
        y[i] = b[i][2]
	end
	#scatter(x,y,c=f.W,cmap="Greys",vmin=0)
	#scatter(x,y,c=f.W,cmap="Blues",vmin=0, alpha=0.2)
	scatter(x, y, s=1, alpha=alpha, c=color)
	#xmean, ymean = centroid(f)
	#plot(xmean, ymean, "gx", ms=10, mew=2)
	labels()
	axis("scaled")
	axis(a)
	tick_params(direction="in")
end
function plot(m::SearchDomain, b::ParticleCollection; alpha=1.0)
	mark_size = 12
	a = [0,m.length,0,m.length]
    n = length(b.particles)
	x = zeros(n)
	y = zeros(n)
	for i = 1:n
		x[i] = particle(b,i)[1]
		y[i] = particle(b,i)[2]
	end
	#scatter(x,y,c=f.W,cmap="Greys",vmin=0)
	#scatter(x,y,c=f.W,cmap="Blues",vmin=0, alpha=0.2)
	scatter(x, y, s=1, alpha=0.2, c="k")
	#xmean, ymean = centroid(f)
	#plot(xmean, ymean, "gx", ms=10, mew=2)
	labels()
	axis("scaled")
	axis(a)
	tick_params(direction="in")
end

function plot_mean(mu::LocTuple; color="b")
	plot_mean(mu[1], mu[2], color=color)
end
function plot_mean(xmean::Float64, ymean::Float64; color="b")
	plot(xmean, ymean, "$(color)X", ms=10, mew=1.5, mec="k")
end

# This is needed for plot_sim
#function plot_b(m::SearchDomain, b::Belief, x::Vehicle)
#	plot_theta(m)
#	hold(true)
#	plot_vehicles(X)
#	imshow(b', interpolation="none", cmap="Greys", origin="lower")
#	labels()
#end

#"""
#`plot_eid(m::SearchDomain, X::VehicleSet)`
#
#Plots the eid using the Fisher information matrix.
#"""
#function plot_eid(m::SearchDomain, X::VehicleSet; contours=false)
#	eid = EID(m,m.b)
#	if contours
#		plot_contour(m, eid)
#	end
#	plot_theta(m)
#	plot_vehicles(X)
#	imshow(eid', interpolation="none", cmap="Greys", origin="lower")
#	labels()
#end

#"""
#`plot_mi(m::SearchDomain, X::VehicleSet)`
#
#Plots the mutual information.
#"""
#function plot_mi(m::SearchDomain, X::VehicleSet; contours=false)
#	mut_info = mutual_information(m)
#	plot_vehicles(X)
#	imshow(mut_info', interpolation="none", cmap="Greys", origin="lower")
#	labels()
#	if contours
#		plot_contour(m, mut_info)
#	end
#end

#"""
#`plot_sim(m::SearchDomain, s::Simulation)`
#
#Steps through a simulation.
#"""
#function plot_sim(m::SearchDomain, s::Simulation)
#	for t = 0:s.T
#		hold(false)
#		pause(1)
#		plot_b(m, s.belief_list[t+1], s.state_list[t+1])
#	end
#end

######################################################################
# Helper functions
# mec = markeredgecolor
# mfc = markerfacecolor
# lw = linewidth
######################################################################
# Plots locations of the vehicles
plot_vehicle(m::SearchDomain,x::Vehicle; color="b") = plot_vehicle(m,x.x,x.y,x.heading, color=color)
plot_vehicle(m::SearchDomain, p::Pose; color="b") = plot_vehicle(m, p[1], p[2], p[3], color=color)
function plot_vehicle(m::SearchDomain, x::Float64, y::Float64, h::Float64; color="k")
	mark_size = 10

	# Changed by LD for 01/30/2017 for rss submission
	#c = 0.01 * m.length  *sqrt(2)
	#c = 0.015 * m.length  *sqrt(2)
	#c = 0.012 * m.length  *sqrt(2)
	c = 0.013 * m.length  *sqrt(2)
	dx = c*sind(45 + h)
	dy = c*cosd(45 + h)

	# Plot rotors
	# LD change 01/30/2017 for rss submission
	#rotor_size = 5
	#rotor_size = 7.5
	rotor_size = 5.5
	rotor_size = 6.5
	theta1 = 45.0 + h
	theta2 = 135.0 + h
	theta3 = 225.0 +  h
	theta4 = 315.0 + h
	#cs = "$(color)o"
	cs = "wo"
	#color = color
	plot(x+c*sind(theta1), y+c*cosd(theta1), cs, ms=rotor_size, mew=2.5, mfc="none", mec=color)
	plot(x+c*sind(theta2), y+c*cosd(theta2), cs, ms=rotor_size, mew=2.5, mfc="none", mec=color)
	plot(x+c*sind(theta3), y+c*cosd(theta3), cs, ms=rotor_size, mew=2.5, mfc="none", mec=color)
	plot(x+c*sind(theta4), y+c*cosd(theta4), cs, ms=rotor_size, mew=2.5, mfc="none", mec=color)

	# plot heading direction
	xline = [x, x+2*c*sind(h)]
	yline = [y, y+2*c*cosd(h)]
	# Changed by LD 01/30/2017 for rss submission
	#plot(xline, yline, color, mew=2)
	#plot(xline, yline, color, lw=2)
	#plot(xline, yline, color, lw=2.5)
	plot(xline, yline, color, lw=2.)

	# Plot frame
	plot(x, y, marker=(2,0,-theta1), ms=mark_size, mew=1, mec=color)
	plot(x, y, marker=(2,0,-theta2), ms=mark_size, mew=1, mec=color)

end

# Plots jammer location
function plot_theta(m::SearchDomain)
	# changed by LD 01/30/2017 for RSS submission
	#mark_size = 11
	#mark_size = 13
	mark_size = 11
	#mark_size = 13
	plot(m.theta[1], m.theta[2], "k^", markersize=mark_size, mew=2.5,markerfacecolor="none")
	#mark_size = 11
	#plot(m.theta[1], m.theta[2], "r^", markersize=mark_size, markerfacecolor="none", markeredgecolor="r", mew=2)
end

# Plots contours of some distribution `d` (a matrix).
function plot_contour(m::SearchDomain, d::Matrix{Float64})
	X,Y = meshgrid(linspace(0.,m.length,100), linspace(0.,m.length,100))
	contour(X, Y, d')
end

# Sets the appropriate plot labels
function labels()
	#xlabel("East (m)", fontsize=12)
	#ylabel("North (m)", fontsize=12)
	xlabel("East (m)")
	ylabel("North (m)")
end

######################################################################
# Old, deprecated, or unused
######################################################################
function plot_eid2(m, b, x, theta)
	mark_size = 12
	eid = EID(m,theta[1],theta[2])
	plot(x[1], x[2], "b*", markersize=mark_size)
	plot(theta[1], theta[2], "r^", markersize=mark_size)
	imshow(eid', interpolation="none", cmap="Greys", origin="lower")
	xlabel("x")
	ylabel("y")
end

function plot_eid3(m, b, x, theta)
	mark_size = 12
	eid = EID(m,theta[1],theta[2])
	#plot(x[1], m.num_cells - 1 - x[2], "b*", markersize=mark_size)
	#plot(theta[1], m.num_cells - 1 - theta[2], "r^", markersize=mark_size)
	#eid_plot = (eid')[end:-1:1, :]

	#imshow(eid_plot, interpolation="none", cmap="Greys")
	plot_contour(m, eid)
	xlabel("x")
	ylabel("y")
end

"""
`meshgrid(x, y)`

Returns X and Y.
"""
function meshgrid(x, y)
	nx = length(x)
	ny = length(y)

	X = zeros(ny, nx)
	Y = zeros(ny, nx)
	for xi = 1:nx
		for xj = 1:ny
			X[xj, xi] = x[xi]
			Y[xj, xi] = y[xj]
		end
	end
	return X, Y
end

######################################################################
# plotting multiple plots...
######################################################################
# TODO: make this just call the other version
function plot{TP<:AbstractFilter}(m::SearchDomain, farr::Vector{TP}, xarr::Vector{Vehicle}; show_mean::Bool=false, show_cov::Bool=false)
	for i = 1:length(farr)
		f = farr[i]
		x = xarr[i]
		plot(m, f, (x.x, x.y, x.heading); show_mean=show_mean, show_cov=show_cov)
	end
end

function plot{TP<:AbstractFilter}(m::SearchDomain, farr::Vector{TP}, parr::Vector{Pose}; show_mean::Bool=false, show_cov::Bool=false, alpha=1.0)
	num_vehicles = length(farr)
	cmaps = ["Blues", "Reds"]
	colors = ["b", "r"]
	if num_vehicles == 1
		cmaps = ["Greys"]
		colors = ["b"]
	elseif num_vehicles == 2
		cmaps = ["Blues", "Reds"]
		colors = ["b", "r"]
	elseif num_vehicles == 3
		cmaps = ["Blues", "Reds", "Greens"]
		colors = ["b", "r", "g"]
	elseif num_vehicles == 4
		cmaps = ["Blues", "Reds", "Greens", "Greys"]
		colors = ["b", "r", "g", "k"]
	end
	plot_theta(m)
	#hold(true)  # deprecated
	#cmaps = ["Blues", "Reds"]
	#colors = ["b", "r"]
	for (fi,f) in enumerate(farr)
		p = parr[fi]
		plot_vehicle(m,p, color=colors[fi])
		if show_mean
			plot_mean(centroid(f), color=colors[fi])
		end
		if show_cov
			plot(m, centroid(f), covariance(f), color=colors[fi])
		end
		plot(m, f, alpha=alpha, cmap=cmaps[fi])
	end
	return # so it doesn't spit out result of axis
end

function plot(m::SearchDomain, vsu::Vector{SimUnit};
              show_mean::Bool=false,
              show_cov::Bool=false,
              alpha=1.0
             )
	num_vehicles = length(vsu)
	cmaps = ["Blues", "Reds"]
	colors = ["b", "r"]
	if num_vehicles == 1
		cmaps = ["Greys"]
		colors = ["b"]
	elseif num_vehicles == 2
		cmaps = ["Blues", "Reds"]
		colors = ["b", "r"]
	elseif num_vehicles == 3
		cmaps = ["Blues", "Reds", "Greens"]
		colors = ["b", "r", "g"]
	elseif num_vehicles == 4
		cmaps = ["Blues", "Reds", "Greens", "Greys"]
		colors = ["b", "r", "g", "k"]
	end
	plot_theta(m)
	#hold(true)  # deprecated
	#cmaps = ["Blues", "Reds"]
	#colors = ["b", "r"]
	for (fi,uav) in enumerate(vsu)
		#p = parr[fi]
		plot_vehicle(m, uav.x, color=colors[fi])
		if show_mean
			plot_mean(centroid(uav.f), color=colors[fi])
		end
		if show_cov
			plot(m, centroid(uav.f), covariance(uav.f), color=colors[fi])
		end
        plot(m, uav.f, alpha=alpha, cmap=cmaps[fi], color=colors[fi])
	end
	return # so it doesn't spit out result of axis
end


using CUDA
using Adapt

mutable struct Test_Cfg
    A::CuArray
    B::CuArray
end
function Adapt.adapt_structure(to, test::Test_Cfg)
    A = Adapt.adapt_structure(to, test.A)
    B = Adapt.adapt_structure(to, test.B)
    Test_Cfg(A, B)
end
struct Kernel_test_Cfg
    A::Float32
    B::Int32
end
function Adapt.adapt_structure(to, test::Kernel_test_Cfg)
    A = Adapt.adapt_structure(to, test.A)
    B = Adapt.adapt_structure(to, test.B)
    Kernel_test_Cfg(A, B)
end

function kernel_test(test)

    @cuprintln(test.A)
    @cuprintln(test.B)
    return
end
#test = Kernel_test_Cfg(2.0f0,3)
# @cuda kernel_test(test)

function maze_kernel(maze,offset,maze_columns,maze_rows)
    total_amount_of_cells = maze_columns * maze_rows

    x_coordinate_stack = @cuDynamicSharedMem(Int32,total_amount_of_cells,offset)
    y_coordinate_stack = @cuDynamicSharedMem(Int32,total_amount_of_cells,offset+sizeof(x_coordinate_stack))
    neighbours = @cuDynamicSharedMem(Int32,4,offset+sizeof(x_coordinate_stack)+sizeof(y_coordinate_stack))
    @inbounds fill!(maze,1)
    cell_x_coordinate = 1
    cell_y_coordinate = 1
    amount_of_cells_visited = 1
    cell_stack_index = 1

    while amount_of_cells_visited < total_amount_of_cells

        @inbounds fill!(neighbours,0)

        if cell_y_coordinate-1 >= 1
            for i in 1:4
                if maze[cell_y_coordinate-1,cell_x_coordinate,i,blockIdx().x] == 0
                    break
                end
                if i == 4
                    neighbours[1] = 1
                end
            end
        end
        if cell_x_coordinate+1 <= maze_columns
            for i in 1:4
                if maze[cell_y_coordinate,cell_x_coordinate+1,i,blockIdx().x] == 0
                    break
                end
                if i == 4
                    neighbours[2] = 1
                end
            end
        end
        if cell_y_coordinate+1 <= maze_rows
            for i in 1:4
                if maze[cell_y_coordinate+1,cell_x_coordinate,i,blockIdx().x] == 0
                    break
                end
                if i == 4
                    neighbours[3] = 1
                end
            end
        end
        if cell_x_coordinate-1 >= 1
            for i in 1:4
                if maze[cell_y_coordinate,cell_x_coordinate-1,i,blockIdx().x] == 0
                    break
                end
                if i == 4
                    neighbours[4] = 1
                end
            end
        end
        if neighbours[1] == 0 && neighbours[2] == 0 && neighbours[3] == 0 && neighbours[4] == 0
            cell_stack_index = cell_stack_index - 1
            cell_x_coordinate = x_coordinate_stack[cell_stack_index]
            cell_y_coordinate = y_coordinate_stack[cell_stack_index]
            continue
        end

        move_x_coordinate = 0
        move_y_coordinate = 0
        rand_index = (abs(rand(Int32)) % 4)
        for i in 1:4
            index = ((rand_index+i) % 4) + 1
            if neighbours[index] == 1
                if index == 1
                    move_y_coordinate = -1
                    break
                end
                if index == 2
                    move_x_coordinate = 1
                    break
                end
                if index == 3
                    move_y_coordinate = 1
                    break
                end
                if index == 4
                    move_x_coordinate = -1
                    break
                end
            end
        end
        if move_y_coordinate == -1
            maze[cell_y_coordinate,cell_x_coordinate,1,blockIdx().x] = 0
            maze[cell_y_coordinate+move_y_coordinate,cell_x_coordinate,3,blockIdx().x] = 0
        end
        if move_x_coordinate == 1
            maze[cell_y_coordinate,cell_x_coordinate,2,blockIdx().x] = 0
            maze[cell_y_coordinate,cell_x_coordinate+move_x_coordinate,4,blockIdx().x] = 0
        end
        if move_y_coordinate == 1
            maze[cell_y_coordinate,cell_x_coordinate,3,blockIdx().x] = 0
            maze[cell_y_coordinate+move_y_coordinate,cell_x_coordinate,1,blockIdx().x] = 0
        end
        if move_x_coordinate == -1
            maze[cell_y_coordinate,cell_x_coordinate,4,blockIdx().x] = 0
            maze[cell_y_coordinate,cell_x_coordinate+move_x_coordinate,2,blockIdx().x] = 0
        end

        #step5: add origin cell to stack
        x_coordinate_stack[cell_stack_index] = cell_x_coordinate
        y_coordinate_stack[cell_stack_index] = cell_y_coordinate
        cell_stack_index = cell_stack_index +1
        #step6: set coordinates to new cell
        cell_x_coordinate = cell_x_coordinate + move_x_coordinate
        cell_y_coordinate = cell_y_coordinate + move_y_coordinate
        amount_of_cells_visited = amount_of_cells_visited +1
    end
    return
end

function kernel_env_step(maze,action,input,environment_config_array)
    maze_columns = 5
    maze_rows = 5
    maze_cell_size = 80
    agent_radius = 12
    agent_movement_range = 10.0f0
    point_radius = 8
    reward_per_collected_positive_point = 500.00f0
    reward_per_collected_negative_point = -700.00f0

    screen_height = maze_rows * maze_cell_size
    screen_width = maze_columns * maze_cell_size

    agent_x_coordinate = environment_config_array[1] + clamp(floor(action[1] * env_cfg.agent_movement_range),-env_cfg.agent_movement_range,env_cfg.agent_movement_range)
    agent_x_coordinate = min(agent_x_coordinate,screen_width - env_cfg.agent_radius) # Eastern border
    agent_x_coordinate = max(agent_x_coordinate,env_cfg.agent_radius) # Western Border

    agent_y_coordinate = environment_config_array[2] + clamp(floor(action[2] * env_cfg.agent_movement_range),-env_cfg.agent_movement_range,env_cfg.agent_movement_range)
    agent_y_coordinate = min(agent_y_coordinate,screen_height - env_cfg.agent_radius) # Eastern border
    agent_y_coordinate = max(agent_y_coordinate,env_cfg.agent_radius) # Western Border

    # Get cell indizes of agents current position
    cell_x = convert(Int32,ceil(agent_x_coordinate / env_cfg.maze_cell_size))
    overlay_y = convert(Int32,ceil(agent_y_coordinate / env_cfg.maze_cell_size))
    cell_y = maze_rows - (overlay_y-1)

    # Get edge coordinates of current cell
    x_left = env_cfg.maze_cell_size * (cell_x - 1)
    x_right = env_cfg.maze_cell_size * cell_x
    y_bottom = env_cfg.maze_cell_size * (overlay_y - 1)
    y_top = env_cfg.maze_cell_size * overlay_y

    if maze[cell_y,cell_x,1] == 1 #check for Northern Wall
        agent_y_coordinate = min(agent_y_coordinate,y_top - env_cfg.agent_radius)
    end
    if maze[cell_y,cell_x,3] == 1 #check for Southern Wall
        agent_y_coordinate = max(agent_y_coordinate,y_bottom + env_cfg.agent_radius)
    end
    if maze[cell_y,cell_x,2] == 1 #check for Eastern Wall
        agent_x_coordinate = min(agent_x_coordinate,x_right - env_cfg.agent_radius)
    end
    if maze[cell_y,cell_x,4] == 1 #check for Western Wall
        agent_x_coordinate = max(agent_x_coordinate,x_left + env_cfg.agent_radius)
    end
    # Check agent collision with top-left edge (prevents sneaking through the edge)
    if (agent_x_coordinate - x_left < env_cfg.agent_radius) && (y_top - agent_y_coordinate < env_cfg.agent_radius)
        agent_x_coordinate = x_left + env_cfg.agent_radius
        agent_y_coordinate = y_top - env_cfg.agent_radius
    end
    # Check agent collision with top-right edge (prevents sneaking through the edge)
    if (x_right - agent_x_coordinate < env_cfg.agent_radius) && (y_top - agent_y_coordinate < env_cfg.agent_radius)
        agent_x_coordinate = x_right - env_cfg.agent_radius
        agent_y_coordinate = y_top - env_cfg.agent_radius
    end
    # Check agent collision with bottom-right edge (prevents sneaking through the edge)
    if (x_right - agent_x_coordinate < env_cfg.agent_radius) && (agent_y_coordinate - y_bottom < env_cfg.agent_radius)
        agent_x_coordinate = x_right - env_cfg.agent_radius
        agent_y_coordinate = y_bottom + env_cfg.agent_radius
    end
    # Check agent collision with bottom-left edge (prevents sneaking through the edge)
    if (agent_x_coordinate - x_left < env_cfg.agent_radius) && (agent_y_coordinate - y_bottom < env_cfg.agent_radius)
        agent_x_coordinate = x_left + env_cfg.agent_radius
        agent_y_coordinate = y_bottom + env_cfg.agent_radius
    end
end

maze_columns = 5
maze_rows = 5
number_individuals = 10
maze = fill(0,(maze_rows,maze_columns,4,number_individuals))

maze_gpu = CuArray(maze)
@cuda blocks=number_individuals shmem=2*maze_columns*maze_rows+4 maze_kernel(maze_gpu,0,maze_columns,maze_rows)
CUDA.synchronize()
display(Array(maze_gpu))


using Base.Test, SDDiP, JuMP, GLPKMathProgInterface, Ipopt

@testset "Lagrangian" begin
    include("Lagrangian.jl")
end

@testset "SDDiP" begin
    include("SDDiP.jl")
end

@testset "@binarystate" begin
    include("binarystate.jl")
end


"""
    Base.convert(::Type{Matrix{Char}}, strings::Vector{String})

Converts a `Vector{String}` to a `Matrix{Char}`, where each column contains
the characters from a string in `strings`, truncated to the length of the
shortest string.
"""
function Base.convert(::Type{Matrix{Char}}, strings::Vector{String})
    sorted = sort!(strings, by = length)
    max_length = length(sorted[1])
    return reduce(hcat, collect(s)[1:max_length] for s in strings)
end

"Helper function to identify iterables where all the values are the same"
all_equal(v) = all(x -> x == v[1], v)

"""
    longest_common_prefix(strings::Vector{String})

Given a list of strings, returns a string containing the longest common prefix
among all strings. Returns an empty string if no matching prefix is found.
"""
function longest_common_prefix(strings::Vector{String})
    char_matrix    = convert(Matrix{Char}, strings)
    matching_chars = Iterators.takewhile(all_equal, eachrow(char_matrix))
    return map(x -> first(x), matching_chars) |> join
end

# Tests! -----------------------------------------------------------------------

using Test

@testset "Should return any single letter as-is" begin
    @test longest_common_prefix(["a"]) == "a"
    @test longest_common_prefix(["b"]) == "b"
    @test longest_common_prefix(["z"]) == "z"
end

@testset "Should identify single-letter prefixes" begin
    input = ["apple", "aardvark"]
    @test longest_common_prefix(input) == "a"

    input = ["garbage", "graendal", "gift"]
    @test longest_common_prefix(input) == "g"

    input = ["eagle", "easter", "easy", "egg"]
    @test longest_common_prefix(input) == "e"
end

@testset "Should identify multiple-letter prefixes" begin
    input = ["apple", "apricot"]
    @test longest_common_prefix(input) == "ap"

    input = ["garbage", "garden", "gargantuan"]
    @test longest_common_prefix(input) == "gar"

    input = ["castle", "casting", "castanet", "cast"]
    @test longest_common_prefix(input) == "cast"
end

@testset "Should identify when there is no common prefix" begin
    input = ["apple", "banana"]
    @test longest_common_prefix(input) == ""

    input = ["garbage", "garden", "helicopter"]
    @test longest_common_prefix(input) == ""

    input = ["castle", "antimony", "castanet", "antiquity"]
    @test longest_common_prefix(input) == ""
end

@testset "Should identify repeated instances of a single word" begin
    input = ["apple", "apple"]
    @test longest_common_prefix(input) == "apple"

    input = ["apple", "apple", "apple"]
    @test longest_common_prefix(input) == "apple"
end

@testset "Should work on large inputs" begin
    input = fill("variadic", 10000)
    @test longest_common_prefix(input) == "variadic"

    append!(input, fill("variable", 10000))
    append!(input, fill("variety", 10000))
    @test longest_common_prefix(input) == "vari"

    push!(input, "a")
    @test longest_common_prefix(input) == ""
end


using Test
using NFFT
using Random
using LinearAlgebra
using FFTW
using NFFTTools

Random.seed!(123)

include("accuracy.jl")
include("constructors.jl")
include("performance.jl")
include("testToeplitz.jl")
include("samplingDensity.jl")
include("cuda.jl")
# Need to run after the other tests since the overload plan_*
include("wrappers.jl")


module SputnikUtilities

using GroupedErrors
using JuliaDB, IndexedTables
using Parameters

import IndexedTables: AbstractIndexedTable

export Data2Select, SelectedData
export SelectPredicate, SelectValues
export Analysis, process

include("selectdata.jl")
include("process.jl")
include(joinpath("techniques", "statplotsrecipe.jl"))
include(joinpath("techniques", "groupederrors.jl"))
include(joinpath("techniques", "table.jl"))

end # module


export CuSubArray, CuSubVector, CuSubMatrix, CuSubVecOrMat

immutable CuSubArray{T,N,I}
    parent::CuArray{T}
    indexes::I
    offset1::Int
end

#Base.Array(x::CuSubArray) = x.parent[x.indexes...]
#Base.SubArray(x::CuSubArray) = view(Array(x.parent), x.indexes)

function Base.size{T,N,I}(x::CuSubArray{T,N,I}, dim::Int)
    typeof(x.indexes[dim]) == Colon ? size(x.parent,dim) : length(x.indexes[dim])
end
Base.size{T}(x::CuSubArray{T,1}) = (size(x,1),)
Base.size{T}(x::CuSubArray{T,2}) = (size(x,1),size(x,2))
Base.size{T}(x::CuSubArray{T,3}) = (size(x,1),size(x,2),size(x,3))
Base.size{T}(x::CuSubArray{T,4}) = (size(x,1),size(x,2),size(x,3),size(x,4))

Base.strides(x::CuSubArray) = strides(x.parent)
Base.stride(x::CuSubArray, dim::Int) = stride(x.parent, dim)

Base.length(x::CuSubArray) = prod(size(x))
Base.similar{T,N}(x::CuSubArray{T}, dims::NTuple{N,Int}) = CuArray{T}(dims)

function Base.pointer(x::CuSubArray, index::Int=1)
    index == 1 || throw("Not implemented yet.")
    pointer(x.parent, x.offset1+index)
end

function Base.view{T,N,I<:Tuple}(parent::CuArray{T,N}, indexes::I)
    @assert N == length(indexes)
    n = N
    offset1 = 0
    for i = 1:length(indexes)
        index = indexes[i]
        typeof(index) == Int && (n -= 1)
        f = typeof(index) == Colon ? 0 : first(index)-1
        offset1 += f * stride(parent,i)
    end
    CuSubArray{T,n,I}(parent, indexes, offset1)
end
Base.view(parent::CuArray, indexes::Union{UnitRange{Int},Colon,Int}...) = view(parent,indexes)


sampletype(::Dists.DiscreteUnivariateDistribution) = Int
sampletype(::Dists.DiscreteMultivariateDistribution) = Vector{Int}
sampletype(::Dists.ContinuousUnivariateDistribution) = Float64
sampletype(::Dists.ContinuousMultivariateDistribution) = Vector{Float64}


# Copyright (C) 2021 Ilya Baran.  This program is distributed under the terms of the MIT license.
module ModelBlocks

using OrdinaryDiffEq: Threads
using Base: Number
export Variable, Variables, variablesUnion, variablesSubtract, renameVariables, variablesToMatlab
include("variables.jl")

export AbstractReaction, SimpleReaction, GeneralRateReaction, GeneralReaction
include("reactions.jl")

export AbstractBlock, Block, getVariables, getParameters, runBlock, BlockWithBindings, BlockOutputDefinition, BlockExtraData, computeOutputs
export solutionToMatlab, solutionToMatrix, solutionToVariables, setParameter!, setParameters!, getExtraData, getTimeRange, setTimeRange!
export getOutputDefinition, setOutputDefinition!, getDiscontinuities, setDiscontinuities!, BlockCombo, renameOutputs!
include("block.jl")

export localSensitivity
include("sensitivity.jl")

export fitParameters, fitParameters!
include("fitting.jl")

export generatePPop, subsamplePPop, nBallVolume, normalPDF, weightedSample, expandPPop, generatePPopFarthest
include("vpop.jl")

end


"""
    AbstractFilteringProblem{S1, S2, T1<:TimeType, T2<:TimeType}

Abstract type for a filtering problem for observations of type `S2` in `T2` and hidden states of type `S1` in `T1`.
"""
abstract type AbstractFilteringProblem{S1, S2, T1<:TimeType, T2<:TimeType} end

"""
    state_model(problem::AbstractFilteringProblem)

Return the hidden state model underlying `problem`.
"""
function state_model(problem::AbstractFilteringProblem) end

"""
    obs_model(problem::AbstractFilteringProblem)

Return the observation model underlying `problem`.
"""
function obs_model(problem::AbstractFilteringProblem) end

"""
    state_dim(problem::AbstractFilteringProblem)

Return the dimensionality of the hidden state in `problem`.
"""
state_dim(problem::AbstractFilteringProblem) = state_dim(state_model(problem))

"""
    obs_dim(problem::AbstractFilteringProblem)

Return the dimensionality of the observed state in `problem`.
"""
obs_dim(problem::AbstractFilteringProblem) = obs_dim(obs_model(problem))

state_type(problem::AbstractFilteringProblem{S1, S2, T1, T2}) where {S1, S2, T1, T2}       = S1
obs_type(problem::AbstractFilteringProblem{S1, S2, T1, T2}) where {S1, S2, T1, T2}         = S2
hidden_time_type(problem::AbstractFilteringProblem{S1, S2, T1, T2}) where {S1, S2, T1, T2} = T1
obs_time_type(problem::AbstractFilteringProblem{S1, S2, T1, T2}) where {S1, S2, T1, T2}    = T2

# How to add new filtering problem:
# * Add a struct which is a subtype of AbstractFilteringProblem{S1, S2, T1, T2}, where S1 and S2 are the types of hidden states and observations respectively, and T1 and T2 are the time types of the hidden and observed processes, respectively.
# * Implement methods for
#   - state_model
#   - obs_model


@testset "User defined models" begin
    y = [1 2; 1 2.0]
    Z = [1 2.0; 2.0 1.0]
    T = [1 2.0]
    R = [1 2.0]
    @test_throws ErrorException StateSpaceModel(y, Z, T, R)

    Z = Array{Float64, 3}(undef, 2, 2, 2)
    Z[:, :, 1] = [1 2.0; 2.0 1.0]
    Z[:, :, 2] = [1 2.0; 2.0 1.0]
    T = [1 2.0; 2.0 1.0]
    R = [1 2.0; 2.0 1.0]
    ssm = StateSpaceModel(y, Z, T, R)
    @test ssm.mode == "time-variant"

    Z = Array{Float64, 3}(undef, 2, 2, 2)
    Z[:, :, 1] = [1 2.0; 2.0 1.0]
    Z[:, :, 2] = [1 2.0; 2.0 1.0]
    T = [1 2.0]
    R = [1 2.0; 2.0 1.0]
    @test_throws ErrorException StateSpaceModel(y, Z, T, R)

    y = ones(5, 1)
    Z = ones(1, 3)
    T = ones(3, 3)
    R = ones(3, 3)
    d = zeros(5, 1) # correct d dimension
    c = zeros(5, 3) # correct c dimension
    H = SSM.build_H(1, Float64)
    Q = SSM.build_Q(3, 1, Float64)
    ssm = StateSpaceModel(y, Z, T, R, d, c, H, Q)
    @test ssm.dim == StateSpaceDimensions(5, 1, 3, 3)

    d = zeros(4, 1) # incorrect d dimension
    @test_throws AssertionError StateSpaceModel(y, Z, T, R, d, c, H, Q)

    d = zeros(5, 1) # correct d dimension
    c = zeros(5, 2) # incorrect c dimension
    @test_throws AssertionError StateSpaceModel(y, Z, T, R, d, c, H, Q)
end


"""
    compute_independent_paths(env::CoordinatedMAPFEnv, task_list::Vector{AgentTask})

For Direct baseline. Computes shortest start-goal paths for each agent without any coordination.
"""
function compute_independent_paths(env::CoordinatedMAPFEnv, task_list::Vector{AgentTask})

    independent_paths = Vector{AgentPathInfo}(undef, length(task_list))

    for (agent_id, task) in enumerate(task_list)

        bidir_sp = compute_bidir_astar_euclidean(env.road_graph, task.origin, task.dest, env.road_graph_wts, env.location_list)

        path_states = AgentPathState[AgentPathState(road_vtx_id=bidir_sp.path[1], timeval=0.0)]

        cumul_time = 0.0
        for (idx, idx_next) in zip(bidir_sp.path[1:end-1], bidir_sp.path[2:end])
            cumul_time += env.road_graph_wts[idx, idx_next]/AvgSpeed
            push!(path_states, AgentPathState(road_vtx_id=idx_next, timeval=cumul_time))
        end

        independent_paths[agent_id] = AgentPathInfo(path_states=path_states, total_dist=bidir_sp.dist, total_time=cumul_time)
    end

    return independent_paths
end

struct EdgeCopyQueueType
    hop::Int64
    src::Int64
end

"""
    augment_road_graph_with_aerial_paths!(env::CoordinatedMAPFEnv, direct_drone_paths::Vector{AgentPathInfo}, decay_fn::Function)

Use aerial paths to augment road graph.
"""
function augment_road_graph_with_aerial_paths!(env::CoordinatedMAPFEnv, direct_drone_paths::Vector{AgentPathInfo}, decay_fn::Function = x -> logistic(x))

    # Iterate over drone paths
    for (aerial_id, ddp) in enumerate(direct_drone_paths)

        ddp_states = ddp.path_states

        # Temp data structure for closest distance to actual edge
        edge_copy_min_hop = Dict{Tuple{Int64,Int64},Int64}()
        edge_copy_src_to_hop = Dict{Int64,Int64}()

        # Iterate over path states and run outward search from
        for i = 1:length(ddp_states) - 1

            # First add edge on path to copy dict; must be unique
            edge_copy_min_hop[(ddp_states[i].road_vtx_id, ddp_states[i+1].road_vtx_id)] = 0
            edge_copy_src_to_hop[ddp_states[i].road_vtx_id] = 0

            # Now work outward from ddp_states[i]
            edge_copy_queue = Queue{EdgeCopyQueueType}()
            for onbr in outneighbors(env.road_graph, ddp_states[i].road_vtx_id)
                if onbr != ddp_states[i+1].road_vtx_id
                    enqueue!(edge_copy_queue, EdgeCopyQueueType(1, onbr))
                end
            end

            while ~(isempty(edge_copy_queue))

                qtop = dequeue!(edge_copy_queue)
                curr_id = qtop.src

                # Compute the min of current hop of qtop.src and MaxHop
                min_hop_for_src = get(edge_copy_src_to_hop, curr_id, MaxHop)

                # Only jump into it if it was not already in edge_copy_src_to_hop
                if qtop.hop < min_hop_for_src

                    # First update for the src itself
                    edge_copy_src_to_hop[curr_id] = qtop.hop

                    curr_hop = qtop.hop + 1

                    # Iterate over outneighbors and add to edge_copy_min_dist if curr_hop < min_hop
                    # In any case, add to queue with curr_hop
                    for onbr in outneighbors(env.road_graph, curr_id)

                        # TODO: Check later that curr_hop is necessarily better than whatever edge_copy_min_hop has
                        edge_copy_min_hop[(curr_id, onbr)] = curr_hop

                        # Enqueue onbr with curr_hop
                        enqueue!(edge_copy_queue, EdgeCopyQueueType(curr_hop, onbr))

                    end # onbr in out_neighbors(ddp_states[i])
                end # qtop.hop < min_hop_for_src
            end # ~(isempty(edge_copy_queue))
        end # i = 1:length(ddp_states) - 1

        # now iterate over edge_copy_min_hop and add copies
        for (src_dst, hop) in edge_copy_min_hop

            e = LightGraphs.Edge(src_dst[1], src_dst[2])

            # Initialize dictionary for edge if it doesn't exist
            if ~haskey(env.edge_to_copy_info, e)
                env.edge_to_copy_info[e] = EdgeCopyInfo()
            end

            reduced_wt = env.road_graph_wts[e.src, e.dst] * decay_fn(hop)
            env.edge_to_copy_info[e].aerial_id_to_weight[aerial_id] = reduced_wt

        end # (src_dst, hop) in edge_copy_min_hop

    end # ddp in direct_drone_paths

end # function augment_road_graph_with_aerial_paths!

## Define MAPF functions for ground paths
function MultiAgentPathFinding.add_constraint!(cbase::GroundMAPFConstraint, cadd::GroundMAPFConstraint)

    # iterate over edges and IDs of cadd and add new key to cbase or augment existing set
    for (edge, aerial_ids) in cadd.avoid_edge_copy_set
        if haskey(cbase.avoid_edge_copy_set, edge)
            union!(cbase.avoid_edge_copy_set[edge], aerial_ids)
        else
            cbase.avoid_edge_copy_set[edge] = deepcopy(aerial_ids)
        end
    end # (edge, aerial_ids) in cadd.avoid_edge_copy_set
end

function MultiAgentPathFinding.overlap_between_constraints(cbase::GroundMAPFConstraint, cadd::GroundMAPFConstraint)
    for common_key in intersect(keys(cbase.avoid_edge_copy_set, cadd.avoid_edge_copy_set))
        if ~(isempty(intersect(cbase.avoid_edge_copy_set[common_key], cadd.avoid_edge_copy_set[common_key])))
            return true
        end
    end
    return false
end

get_mapf_ground_action(env::CoordinatedMAPFEnv, u::GroundMAPFState, v::GroundMAPFState) = GroundMAPFAction(LightGraphs.Edge(u.road_vtx_id, v.road_vtx_id), env.unique_gmapf_states[v].aerial_ids)

MultiAgentPathFinding.get_empty_constraint(::Type{GroundMAPFConstraint}) = GroundMAPFConstraint()

function MultiAgentPathFinding.create_constraints_from_conflict(env::CoordinatedMAPFEnv, conflict::GroundMAPFConflict)

    # Iterate over conflicting edge Ids and create one element sets of edges to avoid of a drone ID
    edge_set_to_avoid = Dict{LightGraphs.Edge,Set{Int64}}()
    for (edge, aerial_ids) in conflict.conflicting_edge_aerial_ids
        edge_set_to_avoid[edge] = aerial_ids
    end

    # Now just repeat the constraint for each member of pair
    gid_pair = conflict.ground_id_pair
    constraint = Dict{Int64,GroundMAPFConstraint}(gid_pair[1] => GroundMAPFConstraint(edge_set_to_avoid), gid_pair[2] => GroundMAPFConstraint(edge_set_to_avoid))

    return [constraint]
end

function MultiAgentPathFinding.get_first_conflict(env::CoordinatedMAPFEnv, solution::Vector{PlanResult{GroundMAPFState,GroundMAPFAction,Float64}})

    # For the first two trucks that share one or more edges of same drones, return
    for (i, sol_i) in enumerate(solution[1:end-1])
        for (j, sol_j) in enumerate(solution[i+1:end])

            other_ground_id = i+j
            overlapping_aerial_edges = Vector{Tuple{LightGraphs.Edge,Set{Int64}}}(undef, 0)

            for (i_action, _) in sol_i.actions
                for (j_action, _) in sol_j.actions
                    if i_action.edge_id == j_action.edge_id && ~isempty(i_action.edge_aerial_ids) && i_action.edge_aerial_ids == j_action.edge_aerial_ids
                        push!(overlapping_aerial_edges, (i_action.edge_id, i_action.edge_aerial_ids))
                    end # i_action.edge_aerial_id != 0 && i_action == j_action
                end # (sj, (state_j, _)) in enumerate(sol_j.states[2:end])
            end # (si, (state_i, _)) in enumerate(sol_i.states[2:end])

            if ~isempty(overlapping_aerial_edges)
                env.num_global_conflicts += 1
                println("Ground conflict no. $(env.num_global_conflicts)")
                if env.num_global_conflicts > env.threshold_global_conflicts
                    throw(DomainError("Too many conflicts in Stage 1!"))
                end

                return GroundMAPFConflict((i, other_ground_id), overlapping_aerial_edges)
            end
        end # (j, sol_j) in enumerate(solution[i+1:end])
    end # (i, sol_i) in enumerate(solution[1:end-1])

    return nothing
end

function ground_mapf_heuristic(env::CoordinatedMAPFEnv, road_vtx_goal::Int64, u::GroundMAPFState)
    return distance_lat_lon_euclidean(env.location_list[u.road_vtx_id], env.location_list[road_vtx_goal])
end

function generate_ground_mapf_edge_weight(env::CoordinatedMAPFEnv, u::GroundMAPFState, v::GroundMAPFState, aerial_ids::Set{Int64})
    if isempty(aerial_ids)
        return env.road_graph_wts[u.road_vtx_id, v.road_vtx_id]
    else
        # Subtract cost from aerial IDs carried
        base_wt = env.road_graph_wts[u.road_vtx_id, v.road_vtx_id]
        uv_edge = LightGraphs.Edge(u.road_vtx_id, v.road_vtx_id)
        cumul_wt = base_wt

        # Subtract cost diff from base for each drone edge
        for aid in aerial_ids
            cumul_wt -= (base_wt - env.edge_to_copy_info[uv_edge].aerial_id_to_weight[aid])
        end

        return cumul_wt
    end
end

function get_ground_mapf_edge_weight(env::CoordinatedMAPFEnv, u::GroundMAPFState, v::GroundMAPFState)
    return (env.unique_gmapf_states[v].gval - env.unique_gmapf_states[u].gval)
end

function MultiAgentPathFinding.set_low_level_context!(::CoordinatedMAPFEnv, ::Int64, ::GroundMAPFConstraint)
end

function MultiAgentPathFinding.focal_heuristic(env::CoordinatedMAPFEnv, solution::Vector{PlanResult{GroundMAPFState,GroundMAPFAction,Float64}})

    num_potential_conflicts = 0

    for (i, sol_i) in enumerate(solution[1:end-1])
        for (j, sol_j) in enumerate(solution[i+1:end])

            num_overlapping_aerial_edges = 0

            for (i_action, _) in sol_i.actions
                for (j_action, _) in sol_j.actions
                    if i_action.edge_id == j_action.edge_id && ~isempty(i_action.edge_aerial_ids) && i_action.edge_aerial_ids == j_action.edge_aerial_ids
                        num_overlapping_aerial_edges += 1
                    end # i_action.edge_aerial_id != 0 && i_action == j_action
                end # (sj, (state_j, _)) in enumerate(sol_j.states[2:end])
            end # (si, (state_i, _)) in enumerate(sol_i.states[2:end])

            if num_overlapping_aerial_edges > 0
                num_potential_conflicts += 1
            end
        end # (j, sol_j) in enumerate(solution[i+1:end])
    end # (i, sol_i) in enumerate(solution[1:end-1])

    return num_potential_conflicts
end

"""
Underlying per-agent search over the Stage 1 MAPF graph.
"""
function MultiAgentPathFinding.low_level_search!(solver::ECBSSolver, ground_agent_id::Int64, s::GroundMAPFState, constraint::GroundMAPFConstraint, solution::Vector{PlanResult{GroundMAPFState,GroundMAPFAction,Float64}})

    env = solver.env

    # Create temp graph as thin wrapper on road graph for search
    env.ground_mapf_graph = SimpleVListGraph{GroundMAPFState}()
    empty!(env.unique_gmapf_states)

    Graphs.add_vertex!(env.ground_mapf_graph, s)
    env.unique_gmapf_states[s] = GMAPFStateInfo(1, 0.0, Set{Int64}())

    road_vtx_goal = env.ground_task_list[ground_agent_id].dest
    admissible_heuristic(u) = ground_mapf_heuristic(env, road_vtx_goal, u)
    edge_wt_fn(u, v) = get_ground_mapf_edge_weight(env, u, v)

    vis = GroundMAPFGoalVisitor(env, road_vtx_goal, constraint.avoid_edge_copy_set)

    start_idx = 1
    env.next_gmapfg_goal_idx = 0

    a_star_states = a_star_implicit_shortest_path!(env.ground_mapf_graph, edge_wt_fn, start_idx, vis, admissible_heuristic)

    @assert haskey(a_star_states.parent_indices, env.next_gmapfg_goal_idx)

    states = Tuple{GroundMAPFState,Float64}[]
    actions = Tuple{GroundMAPFAction,Float64}[]

    sp_cost = a_star_states.dists[env.next_gmapfg_goal_idx]

    sp_idxs = shortest_path_indices(a_star_states.parent_indices, env.ground_mapf_graph, start_idx, env.next_gmapfg_goal_idx)

    push!(states, (s, 0.0))

    for i = 2:length(sp_idxs)

        u_idx, v_idx =  (sp_idxs[i-1], sp_idxs[i])

        new_state = deepcopy(env.ground_mapf_graph.vertices[v_idx])
        push!(states, (new_state, a_star_states.dists[v_idx]))

        push!(actions, (get_mapf_ground_action(env, env.ground_mapf_graph.vertices[u_idx], env.ground_mapf_graph.vertices[v_idx]), a_star_states.dists[v_idx] - a_star_states.dists[u_idx]))
    end

    return PlanResult{GroundMAPFState,GroundMAPFAction,Float64}(states, actions, sp_cost, sp_cost)
end

struct GroundMAPFGoalVisitor <: Graphs.AbstractDijkstraVisitor
    env::CoordinatedMAPFEnv
    road_vtx_goal::Int64
    avoid_edge_copy_set::Dict{LightGraphs.Edge,Set{Int64}}
end

function Graphs.include_vertex!(vis::GroundMAPFGoalVisitor, u::GroundMAPFState, v::GroundMAPFState, d::Float64, nbrs::Vector{Int64})

    env = vis.env
    road_vtx_goal = vis.road_vtx_goal
    avoid_edge_copy_set = vis.avoid_edge_copy_set

    # When we expand actual goal, exit
    if v.road_vtx_id == road_vtx_goal
        env.next_gmapfg_goal_idx = env.unique_gmapf_states[v].idx
        return false
    end

    # Look up outneighbors of v.road_vtx_id in the true road graph
    # For each outneighbor, if no edge copies, just add onbr with aerial ID 0
    # Otherwise, add the outnbr with lowest cost among edge copies if not in avoid_edge_set
    idx_ctr = Graphs.num_vertices(env.ground_mapf_graph)

    for onbr in outneighbors(env.road_graph, v.road_vtx_id)

        vnbr_edge = LightGraphs.Edge(v.road_vtx_id, onbr)
        aids_to_add_with = Set{Int64}()
        is_aerial = false
        v_aids = env.unique_gmapf_states[v].aerial_ids

        if haskey(env.edge_to_copy_info, vnbr_edge)

            edge_copies_id_to_wt = env.edge_to_copy_info[vnbr_edge].aerial_id_to_weight

            aerial_ids_to_avoid = get(avoid_edge_copy_set, vnbr_edge, Set{Int64}())
            aerial_ids_to_consider = setdiff(Set{Int64}(keys(edge_copies_id_to_wt)), aerial_ids_to_avoid)

            # Find lowest weight edges up to min(capacity, aerial_ids)
            if ~(isempty(aerial_ids_to_consider))

                # NOTE: First prefer what it is already carrying
                union!(aids_to_add_with, intersect(v_aids, aerial_ids_to_consider))

                # Then remove from aerial_ids_to_consider
                setdiff!(aerial_ids_to_consider, v_aids)

                # Now choose best remaining
                if ~(isempty(aerial_ids_to_consider))
                    sorted_wt_edges = [(aid, edge_copies_id_to_wt[aid]) for aid in aerial_ids_to_consider]
                    sort!(sorted_wt_edges, by=x->x[2])

                    n_aids = min(length(sorted_wt_edges), env.car_capacity-length(aids_to_add_with))
                    for (aid, wt) in sorted_wt_edges[1:n_aids]
                        push!(aids_to_add_with, aid)
                    end
                end

                if ~isempty(aids_to_add_with)
                    is_aerial = true
                end
            end
        end

        # Check if id_to_add_with is already in ground_mapf_graph
        new_vert = GroundMAPFState(onbr, is_aerial)
        gval = d + generate_ground_mapf_edge_weight(env, v, new_vert, aids_to_add_with)
        if ~(haskey(env.unique_gmapf_states, new_vert))
            idx_ctr += 1
            Graphs.add_vertex!(env.ground_mapf_graph, new_vert)
            push!(nbrs, idx_ctr)
            env.unique_gmapf_states[new_vert] = GMAPFStateInfo(idx_ctr, gval, aids_to_add_with)
        else
            # Only add updated if gval improves
            # NOTE: Should only happen for aerial-aerial or drone-drone
            if gval < env.unique_gmapf_states[new_vert].gval
                push!(nbrs, env.unique_gmapf_states[new_vert].idx)
                env.unique_gmapf_states[new_vert].gval = gval
                env.unique_gmapf_states[new_vert].aerial_ids = aids_to_add_with
            end
        end
    end # onbr in out_neighbors(v.road_vtx_id, env.road_graph)

    return true
end

"""
    update_ground_paths_with_ground_mapf_result!(env::CoordinatedMAPFEnv, solution::Vector{PlanResult})

Using the Stage 1 MAPF solution, update the set of solution paths in the env datastructure
"""
function update_ground_paths_with_ground_mapf_result!(env::CoordinatedMAPFEnv, solution::Vector{PlanResult{GroundMAPFState,GroundMAPFAction,Float64}})

    env.ground_paths = Vector{AgentPathInfo}(undef, length(env.ground_task_list))

    for (ground_id, ground_soln) in enumerate(solution)

        ground_soln_states = ground_soln.states

        path_states = AgentPathState[AgentPathState(road_vtx_id=ground_soln_states[1][1].road_vtx_id, timeval=0.0)]

        cumul_time = 0.0
        cumul_dist = 0.0
        for (state, state_next) in zip(ground_soln_states[1:end-1], ground_soln_states[2:end])
            dist = env.road_graph_wts[state[1].road_vtx_id, state_next[1].road_vtx_id]
            cumul_dist += dist
            cumul_time += dist/AvgSpeed
            push!(path_states, AgentPathState(road_vtx_id=state_next[1].road_vtx_id, timeval=cumul_time))
        end

        env.ground_paths[ground_id] = AgentPathInfo(path_states=path_states, total_dist=cumul_dist, total_time=cumul_time)
    end
end


"""
    initialize_solver!(p,o)

Initialize the solver to the optimization [`Problem`](@ref) by initializing
the encapsulated `options` from within the [`RecordOptions`](@ref)` o`.
"""
function initialize_solver!(p::Problem, o::RecordOptions)
    initialize_solver!(p, o.options)
    get(o.recordDictionary, :Start, RecordGroup())(p, get_options(o), 0)
    return get(o.recordDictionary, :All, RecordGroup())(p, get_options(o), 0)
end
"""
    step_solver!(p,o,iter)

Do one iteration step (the `iter`th) for [`Problem`](@ref)` p` by modifying
the values in the [`Options`](@ref)` o.options` and record the result(s).
"""
function step_solver!(p::Problem, o::RecordOptions, i)
    step_solver!(p, o.options, i)
    get(o.recordDictionary, :Step, RecordGroup())(p, get_options(o), i)
    return get(o.recordDictionary, :All, RecordGroup())(p, get_options(o), i)
end
"""
    get_solver_result(o)

Return the final result after all iterations that is stored within the
(modified during the iterations) [`Options`](@ref)` o`.
"""
function get_solver_result(o::RecordOptions)
    return get_solver_result(o.options)
end
"""
    stop_solver!(p,o,i)

determine whether the solver for [`Problem`](@ref) `p` and the
[`RecordOptions`](@ref) `o` should stop at iteration `i`.
If so, do a (final) record to `:All` and `:Stop`.
"""
function stop_solver!(p::Problem, o::RecordOptions, i::Int)
    s = stop_solver!(p, o.options, i)
    if s
        get(o.recordDictionary, :Stop, RecordGroup())(p, get_options(o), typemin(Int))
        get(o.recordDictionary, :All, RecordGroup())(p, get_options(o), typemin(Int))
    end
    return s
end


@testset "interactions" begin
	fh = FloryInteraction()
	set_interaction!(fh, 1, 2, 0.5)
	set_interaction!(fh, 2, 1, 0.5)
	@test length(fh.chis) == 1
	@test_throws AssertionError set_interaction!(fh, 1, 1, 0.5)
end


import DomainSets: Interval, Ball, infimum, supremum

@deprecate IntervalDomain(a, b) Interval(a, b)
@deprecate CircleDomain() Ball()

# type piracy on Interval for downstream compatibility to be reverted once upgrade is complete
function Base.getproperty(domain::Interval, sym::Symbol)
    if sym === :lower
        @warn "domain.lower is deprecated, use infimum(domain) instead"
        return infimum(domain)
    elseif sym === :upper
        @warn "domain.upper is deprecated, use supremum(domain) instead"
        return supremum(domain)
    else
        return getfield(domain, sym)
    end
end


using LegendGeSim

##
# import Pkg
# Pkg.add(url="https://github.com/legend-exp/LegendTestData.jl.git")
# Pkg.build("LegendTestData")
using LegendTestData
testdata_path = joinpath(LegendTestData.legend_test_data_path(), "data", "ldsim")

##

# PET (position-energy-time) information of the simulated energy depositions
pet_file = joinpath(testdata_path, "single-invcoax-th228-geant4.csv")
# detector
det_metadata = joinpath(testdata_path, "invcoax-metadata.json")
# simulation settings
sim_config_name = "configs/SSD_NoiseSim.json"
#sim_config_name = "configs/siggen_NoiseSim.json"

##
raw_table = LegendGeSim.simulate_raw(pet_file, det_metadata, sim_config_name)

##
using HDF5
using LegendHDF5IO

##
mkdir("output")
HDF5.h5open("output/my_name_raw.h5", "w") do f
    LegendHDF5IO.writedata(f, "raw", raw_table)
end

##
using Plots

##
plot_wf = plot(raw_table.waveform[1:5])
# png(plot_wf, "raw.png")

##
stp_table = LegendGeSim.pet_to_stp(pet_file, det_metadata)

##
mkdir("cache")
HDF5.h5open("cache/my_name_stp.h5", "w") do f
    LegendHDF5IO.writedata(f, "stp", stp_table)
end

##
pss_table, pss_truth = LegendGeSim.stp_to_pss("cache/my_name_stp.h5", det_metadata, sim_config_name)

##
pss_table, pss_truth = LegendGeSim.stp_to_pss(stp_table, det_metadata, sim_config_name)

##
using Plots

##
plot_pss = plot(pss_table.waveform[1:5])
# png(plot_pss, "plots_md/pss.png")

##
h5open("cache/my_name_pss.h5", "w") do f
    LegendHDF5IO.writedata(f, "pss/pss", pss_table)
    LegendHDF5IO.writedata(f, "pss/truth", pss_truth)
end

##
raw_table = LegendGeSim.pss_to_raw("cache/my_name_pss.h5", det_metadata, sim_config_name)

##
raw_table = LegendGeSim.pss_to_raw(pss_table, pss_truth, det_metadata, sim_config_name)

##
pss_table, pss_truth = LegendGeSim.pet_to_pss(pet_filename, det_metadata, sim_config_name)


@testset "cmp_search_results" begin
    idxs = shuffle!(randsubseq(1:100, 0.5))
    ds = rand(Float64, length(idxs))

    p = randperm(length(idxs))
    @test cmp_search_results((idxs, ds), (idxs[p], ds[p]))

    @test !cmp_search_results((idxs, ds), (idxs[2:end], ds[2:end]))
    @test !cmp_search_results((idxs, ds), (idxs[p], ds))
end


struct BinfileMeltRate <: AbstractMeltRate
    input_filename::String    #specify melt filename
end

function BinfileMeltRate(;
                        input_filename = nothing)

    #input file exceptions
    ~(input_filename === nothing) || throw(ArgumentError("You must pass an input filename"))
    isfile(input_filename) || throw(ArgumentError("Did not find the specified binary file"))

    return BinfileMeltRate(input_filename)
end

function update_melt_rate!(melt_model::BinfileMeltRate, fields, grid)
    @unpack input_filename = melt_model
    @unpack basal_melt = fields.gh

    file_size = stat(input_filename).size #bytes in input file must match matrix returned
    melt_rate = zeros(grid.nx,grid.ny)
    (file_size == sizeof(melt_rate)) || throw(DimensionMismatch("Size of input file incompatible with specified nx, ny"))
    try
        read!(input_filename, melt_rate)
    catch
        Error("Input file read error")
    end

   # melt_rate .= ntoh.(melt_rate)
    basal_melt[:] = melt_rate[:]
    return nothing
end


function setindex_multiprocesses_worker(block::Array{T,N}, ba::BigArray{D,T},
                                        chunkGlobalRange::CartesianIndices{N}) where {D,T,N}
    C = get_encoding(ba)
    ba.kvStore[ cartesian_range2string(chunkGlobalRange) ] = encode( block, C)
end

"""
    put array in RAM to a BigArray
this version uses channel to control the number of asynchronized request
"""
function setindex_multiprocesses!( ba::BigArray{D,T}, buf::Array{T,N},
                       idxes::Union{UnitRange, Int, Colon} ... ) where {D,T,N}
    idxes = colon2unit_range(buf, idxes)
    # check alignment
    @assert all(map((x,y,z)->mod(first(x) - 1 - y, z), idxes, ba.offset.I, ba.chunkSize).==0) "the start of index should align with BigArray chunk size"
    t1 = time()
    baIter = ChunkIterator(idxes, ba.chunkSize; offset=ba.offset)
    @sync begin
        for (blockID, chunkGlobalRange, globalRange, rangeInChunk, rangeInBuffer) in baIter
            chunkGlobalRange, globalRange, rangeInChunk, rangeInBuffer =
                adjust_volume_boundary(ba, chunkGlobalRange, globalRange,
                                       rangeInChunk, rangeInBuffer)
            block = buf[rangeInBuffer]
            @async remotecall_fetch(setindex_multiprocesses_worker, WORKER_POOL,
                                       block, ba, chunkGlobalRange)
        end
    end
    elapsed = time() - t1 # sec
    println("saving speed: $(sizeof(buf)/1024/1024/elapsed) MB/s")
end

function getindex_multiprocesses_worker(ba::BigArray{D,T}, jobs::RemoteChannel,
                                results::RemoteChannel) where {D,T,N}
    baRange = CartesianIndices(ba)
    blockId, chunkGlobalRange, globalRange, rangeInChunk, rangeInBuffer = take!(jobs)
    if any(map((x,y)->x>y, first(globalRange).I, last(baRange).I)) || any(map((x,y)->x<y, last(globalRange).I, first(baRange).I))
        @warn("out of volume range, keep it as zeros")
        return
    end
    chunkGlobalRange, globalRange, rangeInChunk, rangeInBuffer = adjust_volume_boundary(ba, chunkGlobalRange, globalRange, rangeInChunk, rangeInBuffer)
    # finalize to avoid memory leak, see
    # https://discourse.julialang.org/t/understanding-distributed-memory-garbage-collection/8726/2
    #finalize(jobs)
    chunkSize = (last(chunkGlobalRange) - first(chunkGlobalRange) + one(CartesianIndex{N})).I
    #println("processing block in global range: $(cartesian_range2string(globalRange))")
    try
        data = ba.kvStore[ cartesian_range2string(chunkGlobalRange) ]
        chk = Codings.decode(data, get_encoding(ba))
        chk = reshape(reinterpret(T, chk), chunkSize)
        chk = chk[rangeInChunk]
        arr = OffsetArray(chk, cartesian_range2unit_range(globalRange)...)
        put!(results, arr)
    catch err
        if isa(err, KeyError) && ba.fillMissing
            println("no such key: $(err), will fill with zeros.")
        else
            println("catch an error while get index in remote worker: $err")
            @show typeof(err)
            @show stacktrace()
            rethrow()
        end
    end
end

function getindex_multiprocesses( ba::BigArray{D, T, N}, idxes::Union{UnitRange, Int}...) where {D,T,N}
    t1 = time()
    sz = map(length, idxes)
    ret = OffsetArray(zeros(T, sz), idxes...)

    channelSize = cld( nworkers(), 2 )
    jobs    = RemoteChannel(()->Channel{Tuple}( channelSize ));
    results = RemoteChannel(()->Channel{OffsetArray}( channelSize ));

    baIter = ChunkIterator(idxes, ba.chunkSize; offset=get_offset(ba))

    @sync begin
        @async begin
            for iter in baIter
                put!(jobs, iter)
            end
            close(jobs)
        end
        # control the number of concurrent requests here
        for iter in baIter
            @async remote_do(getindex_multiprocesses_worker, WORKER_POOL, ba, jobs, results)
        end

        @async begin
            for iter in baIter
                block = take!(results)
                ret[axes(block)...] = parent(block)
            end
            close(results)
        end
    end
    elapsed = time() - t1 # seconds
    println("cutout speed: $(sizeof(ret)/1024/1024/elapsed) MB/s")
    # handle single element indexing, return the single value
    ret
end


mutable struct PeriodicValues{T}
    values          ::Vector{Tuple{DateTime, T}}
    period          ::Period
    last_dt         ::DateTime
end

function periodic_collector(
    type        ::Type{T},
    period      ::Period
)::Tuple{Function, PeriodicValues{T}} where T
    values = Vector{Tuple{DateTime, T}}()
    pv = PeriodicValues{T}(values, period, DateTime(0))

    @inline function collector(dt::DateTime, value::T)::Nothing
        if (dt - pv.last_dt) >= period
            push!(values, (dt, value))
            pv.last_dt = dt
        end
        nothing
    end

    return collector, pv
end

# ----------------------------------------------------------

mutable struct PredicateValues{T}
    values          ::Vector{Tuple{DateTime, T}}
    last_dt         ::DateTime
    last_value      ::T
    data            ::Any               # user-defined data for use in predicate
end

function predicate_collector(
    type            ::Type{T},
    predicate       ::Function,
    init_value      ::T
)::Tuple{Function, PredicateValues{T}} where T
    values = Vector{Tuple{DateTime, T}}()
    pv = PredicateValues{T}(values, DateTime(0), init_value, nothing)

    @inline function collector(dt::DateTime, value::T)::Nothing
        if predicate(pv, dt, value)
            push!(values, (dt, value))
            pv.last_dt = dt
            pv.last_value = value
        end
        nothing
    end

    return collector, pv
end

# ----------------------------------------------------------

mutable struct MinValue{T}
    dt         ::DateTime
    min_value  ::T
end

function min_value_collector(type::Type{T})::Tuple{Function, MinValue{T}} where T
    mv = MinValue{T}(DateTime(0), 1e50)

    @inline function collector(dt::DateTime, value::T)::Nothing
        if value < mv.min_value
            mv.min_value = value
            mv.dt = dt
        end
        nothing
    end

    return collector, mv
end

# ----------------------------------------------------------

mutable struct MaxValue{T}
    dt         ::DateTime
    max_value  ::T
end

function max_value_collector(type::Type{T})::Tuple{Function, MaxValue{T}} where T
    mv = MaxValue{T}(DateTime(0), typemin(T))

    @inline function collector(dt::DateTime, value::T)::Nothing
        if value > mv.max_value
            mv.max_value = value
            mv.dt = dt
        end
        nothing
    end

    return collector, mv
end

# ----------------------------------------------------------

@enum DrawdownMode Percentage PnL

mutable struct DrawdownValues
    values          ::Vector{Tuple{DateTime, Price}}
    mode            ::DrawdownMode
    max_equity      ::Price
    last_dt         ::DateTime
    data            ::Any               # user-defined data for use in predicate
end

function drawdown_collector(mode::DrawdownMode, predicate::Function)::Tuple{Function, DrawdownValues}
    dv = DrawdownValues(
        Vector{Tuple{DateTime, Price}}(),
        mode,
        -1e50,
        DateTime(0),
        nothing)

    @inline function collector(dt::DateTime, equity::Price)::Nothing
        # keep track of max equity value
        dv.max_equity = max(dv.max_equity, equity)

        # should collect new drawdown value?
        if predicate(dv, dt, equity)
            drawdown = min(0.0, equity - dv.max_equity)
            if mode == Percentage::DrawdownMode
                drawdown /= dv.max_equity
            end
            push!(dv.values, (dt, drawdown))
            dv.last_dt = dt
        end

        nothing
    end

    return collector, dv
end

# ----------------------------------------------------------


abstract type Score end
# expectation fields, must be mutable!
# - prior::Prior

# make scores iterable
Base.iterate(score::Score, state = 0) = state > 0 ? nothing : (score, state + 1)
Base.length(score::Score) = 1

# reduce Base.show() to Base.string()
Base.show(io::IO, score::Score) = print(io, string(score))
Base.show(io::IO, ::MIME"application/prs.juno.inline", score::Score) = print(io, string(score))
Base.string(score::Score) = @sprintf "ToDo: implement Base.string for %s" typeof(score)

# make scores callable
(score::Score)(args...; kwargs...) = evaluate(score, args...; kwargs...)

# update the prior of a score
update!(score::TS, x::TI, n::TI) where {TS<:Score,TI<:Integer} = score.prior = update(score.prior, x, n)

function update(score::Score, prior::Prior)
    error("not implemented")
end

function evaluate(score::TS, x1::TI, n1::TI, x2::TI, n2::TI, c2::TR, p::TR)::TR where {TI<:Integer,TR<:Real,TS<:Score}

    error(@sprintf "ToDo: implement bad.evaluate for %s" typeof(score))
end
function evaluate(score::TS, x1::TI, n1::TI, x2::TI, n2::TI, c2::TR) where {TI<:Integer,TR<:Real,TS<:Score}
    expectation(
        p -> evaluate(score, x1, n1, x2, n2, c2, p),
        update(score.prior, x1 + x2, n1 + n2)
    )
end

function evaluate(score::TS, design::TD, x1::TI, x2::TI, p::TR) where {TI<:Integer,TR<:Real,TD<:AbstractDesign,TS<:Score}

    evaluate(score, x1, n1(design), x2, n2(design, x1), c2(design, x1), p)
end
function evaluate(score::TS, design::TD, x1::TI, x2::TI) where {TI<:Integer,TD<:AbstractDesign,TS<:Score}
    expectation(
        p -> evaluate(score, design, x1, x2, p),
        update(score.prior, x1 + x2, n1(design) + n2(design, x1))
    )
end

# conditional expectation given x1
function evaluate(score::TS, x1::TI, n1::TI, n2::TI, c2::TR, p::TR; partial_stage_two::Tuple{TI,TI} = (0, 0)) where {TI<:Integer,TR<:Real,TS<:Score}

    @assert 0 <= partial_stage_two[1] <= partial_stage_two[2] <= n2
    X2 = collect(partial_stage_two[1]:n2)
    return sum(
        evaluate.(score, x1, n1, X2, n2, c2, p) .*
        pmf.(X2, n2, p; partial = partial_stage_two)
    )
end
function evaluate(score::TS, x1::TI, n1::TI, n2::TI, c2::TR; partial_stage_two::Tuple{TI,TI} = (0, 0)) where {TI<:Integer,TR<:Real,TS<:Score}
    expectation(
        p -> evaluate(score, x1, n1, n2, c2, p; partial_stage_two = partial_stage_two),
        update(score.prior, x1 + partial_stage_two[1], n1 + partial_stage_two[2])
    )
end

function evaluate(score::TS, design::TD, x1::TI, p::TR; partial_stage_two::Tuple{TI,TI} = (0, 0)) where {TI<:Integer,TR<:Real,TD<:AbstractDesign,TS<:Score}

    return evaluate(score, x1, n1(design), n2(design, x1), c2(design, x1), p; partial_stage_two = partial_stage_two)
end
function evaluate(score::TS, design::TD, x1::TI; partial_stage_two::Tuple{TI,TI} = (0, 0)) where {TI<:Integer,TD<:AbstractDesign,TS<:Score}
    expectation(
        p -> evaluate(score, design, x1, p; partial_stage_two = partial_stage_two),
        update(score.prior, x1 + partial_stage_two[1], n1(design) + partial_stage_two[2])
    )
end

# marginalise X1, X2
function evaluate(score::TS, design::TD, p::TR; partial_stage_one::Tuple{TI,TI} = (0, 0)) where {TI<:Integer,TR<:Real,TD<:AbstractDesign,TS<:Score}

    @assert 0 <= partial_stage_one[1] <= partial_stage_one[2] <= n1(design)
    XX = sample_space(design)
    XX = XX[XX[:,1] .>= partial_stage_one[1], :]
    return sum(
        evaluate.(score, design, XX[:,1], XX[:,2], p) .*
        pmf.(XX[:,2], n2.(design, XX[:,1]), p) .*
        pmf.(XX[:,1], n1(design), p; partial = partial_stage_one)
    )
end
function evaluate(score::TS, design::TD; partial_stage_one::Tuple{TI,TI} = (0, 0)) where {TI<:Integer,TD<:AbstractDesign,TS<:Score}

    expectation(
        p -> evaluate(score, design, p; partial_stage_one = partial_stage_one),
        update(score.prior, partial_stage_one[1], partial_stage_one[1])
    )
end

# compute Pr[X1=x1]*E_prior[s(x1, n1, X2, n2(x1), c2(x1))] = score_integrand(x1, n1, n2, c2)
# this is the only thing that actually used during optmisation
# feel free to implement specific, more efficient versions for each score!
function integrand_x1(score::TS, x1::TI, n1::TI, n2::TI, c2::TR; partial_stage_one::Tuple{TI,TI} = (0, 0)) where {TS<:Score,TI<:Integer,TR<:Real}
    score(x1, n1, n2, c2) * pmf(x1, n1, score.prior; partial = partial_stage_one)
end

mutable struct SampleSize{TP<:Prior} <: Score
    prior::TP
end
Base.string(score::SampleSize) = @sprintf "SampleSize<%s>" string(score.prior)

function evaluate(score::SampleSize, x1::TI, n1::TI, x2::TI, n2::TI, c2::TR, p::TR)::TR where {TI<:Integer,TR<:Real,TD<:AbstractDesign}
    n1 + n2
end
# no custom score_integrand necessary - fallback to default

update(score::SampleSize, prior::Prior) = SampleSize(prior)

mutable struct Power{TP<:Prior,TR<:Real} <: Score
    prior::TP
    bounds::Tuple{TR,TR}
    Power{TP,TR}(prior::TP, bounds::Tuple{TR,TR}) where{TP<:Prior,TR<:Real} = new(condition(prior; low = bounds[1], high = bounds[2]), bounds)
end
Power(prior::Prior; bounds::Tuple{Real,Real} = Tuple(bounds(prior))) = Power{typeof(prior),eltype(bounds)}(prior, bounds)
Base.string(score::Power) = @sprintf "Power<%s>" string(score.prior)

update(score::Power, prior::Prior) = Power(prior; bounds = score.bounds)

function evaluate(score::Power, x1::TI, n1::TI, x2::TI, n2::TI, c2::TR, p::TR)::TR where {TI<:Integer,TR<:Real}
    (score.bounds[1] <= p <= score.bounds[2]) ? (x2 > c2) : 0.0
end

# since we integrate power, it is easier to do that in one wash, also
# closed form cdf is available, note that the score prior is already properly conditioned
function integrand_x1(score::Power, x1::TI, n1::TI, n2::TI, c2::TR; partial_stage_one::Tuple{TI,TI} = (0, 0) ) where {TS<:Score,TI<:Integer,TR<:Real}

    expectation( # conditional expected power
        p -> 1 - cdf(c2, n2, p),
        update(score.prior, x1, n1)
    ) * pmf( # conditional pmf given partial observations
        x1, n1, score.prior; partial = partial_stage_one
    )
end

mutable struct CompositeScore <: Score
    ω
    components
end
Base.string(score::CompositeScore) = "CompositeScore"

Base.convert(CompositeScore, score::Score) = CompositeScore([1.0], [score])
Base.convert(::Type{T}, score::T)  where {T<:Score} = score

*(ω::Real, score::Score) = CompositeScore([ω], [score])
*(ω::Real, score::CompositeScore) = CompositeScore(ω .* score.ω, score.components)
function +(score1::T1, score2::T2) where {T1,T2<:Score}

    score1, score2 = convert.(CompositeScore, (score1, score2))
    return CompositeScore(vcat(score1.ω, score2.ω), vcat(score1.components, score2.components))
end
function -(score1::T1, score2::T2) where {T1,T2<:Score}

    score1, score2 = convert.(CompositeScore, (score1, score2))
    return CompositeScore(vcat(score1.ω, -score2.ω), vcat(score1.components, score2.components))
end
function update!(score::CompositeScore, x::TI, n::TI) where {TS<:Score,TI<:Integer}
    map(component -> update!(component, x, n), score.components)
end

function update(score::CompositeScore, prior::Prior)
    score = deepcopy(score)
    score.components = map(component -> update(component, prior), score.components)
    return score
end

evaluate(score::CompositeScore, x1::TI, n1::TI, x2::TI, n2::TI, c2::TR, p::TR) where {TI<:Integer,TR<:Real} = sum( score.ω .* evaluate.(score.components, x1, n1, x2, n2, c2, p) )
evaluate(score::CompositeScore, x1::TI, n1::TI, x2::TI, n2::TI, c2::TR) where {TI<:Integer,TR<:Real} = sum( score.ω .* evaluate.(score.components, x1, n1, x2, n2, c2) )
evaluate(score::CompositeScore, design::TD, x1::TI, x2::TI, p::TR) where {TI<:Integer,TR<:Real,TD<:AbstractDesign} = sum( score.ω .* evaluate.(score.components, design, x1, x2, p) )
evaluate(score::CompositeScore, design::TD, x1::TI, x2::TI) where {TI<:Integer,TD<:AbstractDesign} = sum( score.ω .* evaluate.(score.components, design, x1, x2, c2) )
evaluate(score::CompositeScore, x1::TI, n1::TI, n2::TI, c2::TR, p::TR) where {TI<:Integer,TR<:Real} = sum( score.ω .* evaluate.(score.components, x1, n1, n2, c2, p) )
evaluate(score::CompositeScore, x1::TI, n1::TI, n2::TI, c2::TR) where {TI<:Integer,TR<:Real} = sum( score.ω .* evaluate.(score.components, x1, n1, n2, c2) )
evaluate(score::CompositeScore, design::TD, x1::TI, p::TR) where {TI<:Integer,TR<:Real,TD<:AbstractDesign} = sum( score.ω .* evaluate.(score.components, design, x1, p) )
evaluate(score::CompositeScore, design::TD, x1::TI) where {TI<:Integer,TD<:AbstractDesign} = sum( score.ω .* evaluate.(score.components, design, x1) )
evaluate(score::CompositeScore, design::TD, p::TR; partial_stage_one::Tuple{TI,TI} = (0, 0)) where {TR<:Real,TI<:Integer,TD<:AbstractDesign} = sum( score.ω .* evaluate.(score.components, design, p; partial_stage_one = partial_stage_one) )
evaluate(score::CompositeScore, design::TD; partial_stage_one::Tuple{TI,TI} = (0, 0)) where {TI<:Integer,TD<:AbstractDesign} = sum( score.ω .* evaluate.(score.components, design; partial_stage_one = partial_stage_one) )

function integrand_x1(score::CompositeScore, x1::TI, n1::TI, n2::TI, c2::TR; partial_stage_one::Tuple{TI,TI} = (0, 0) ) where {TS<:Score,TI<:Integer,TR<:Real}
    return sum( score.ω .* integrand_x1.(score.components, x1, n1, n2, c2; partial_stage_one = partial_stage_one) )
end


import Images

const writedir = joinpath(tempdir(), "Images")

if !isdir(writedir)
    mkdir(writedir)
end

a = rand(2,2)
fn = joinpath(writedir, "2by2.png")
Images.imwrite(a, fn)
sleep(0.2)
b = Images.imread(fn)
@assert convert(Array, b) == scale(scaleinfo(a), a)


using BlockArrays, Test

@testset "broadcast" begin
    @testset "BlockArray" begin
        A = BlockArray(randn(6), 1:3)

        @test BlockArrays.BroadcastStyle(typeof(A)) == BlockArrays.BlockStyle{1}()

        @test exp.(A) == exp.(Vector(A))
        @test blocksizes(A) == blocksizes(exp.(A))

        @test A+A isa BlockArray
        @test blocksizes(A + A) == blocksizes(A .+ A) == blocksizes(A)
        @test blocksizes(A .+ 1) == blocksizes(A)

        A = BlockArray(randn(6,6), 1:3,1:3)

        @test BlockArrays.BroadcastStyle(typeof(A)) == BlockArrays.BlockStyle{2}()

        @test exp.(A) == exp.(Matrix(A))
        @test blocksizes(A) == blocksizes(exp.(A))

        @test blocksizes(A + A) == blocksizes(A .+ A) == blocksizes(A)
        @test blocksizes(A .+ 1) == blocksizes(A)
    end

    @testset "PseudoBlockArray" begin
        A = PseudoBlockArray(randn(6), 1:3)

        @test BlockArrays.BroadcastStyle(typeof(A)) == BlockArrays.PseudoBlockStyle{1}()

        @test exp.(A) == exp.(Vector(A))
        @test blocksizes(A) == blocksizes(exp.(A))

        @test A+A isa PseudoBlockArray
        @test blocksizes(A + A) == blocksizes(A .+ A) == blocksizes(A)
        @test blocksizes(A .+ 1) == blocksizes(A)

        B = PseudoBlockArray(randn(6,6), 1:3,1:3)

        @test BlockArrays.BroadcastStyle(typeof(B)) == BlockArrays.PseudoBlockStyle{2}()

        @test exp.(B) == exp.(Matrix(B))
        @test blocksizes(B) == blocksizes(exp.(B))

        @test blocksizes(B + B) == blocksizes(B .+ B) == blocksizes(B)
        @test blocksizes(B .+ 1) == blocksizes(B)
        @test blocksizes(A .+ 1 .+ B) == blocksizes(B)
        @test A .+ 1 .+ B == Vector(A) .+ 1 .+ B == Vector(A) .+ 1 .+ Matrix(B)
    end

    @testset "Mixed" begin
        A = BlockArray(randn(6), 1:3)
        B = PseudoBlockArray(randn(6), 1:3)

        @test A + B isa BlockArray
        @test B + A isa BlockArray

        @test blocksizes(A + B) == blocksizes(A)

        C = randn(6)

        @test A + C isa BlockVector{Float64}
        @test C + A isa BlockVector{Float64}
        @test B + C isa PseudoBlockVector{Float64}
        @test C + B isa PseudoBlockVector{Float64}

        @test blocksizes(A+C) == blocksizes(C+A) == blocksizes(A)
        @test blocksizes(B+C) == blocksizes(C+B) == blocksizes(B)

        A = BlockArray(randn(6,6), 1:3, 1:3)
        D = Diagonal(ones(6))
        @test blocksizes(A + D) == blocksizes(A)
        @test blocksizes(B .+ D) == BlockArrays.BlockSizes([1,2,3],[6])
    end

    @testset "Mixed block sizes" begin
        A = BlockArray(randn(6), 1:3)
        B = BlockArray(randn(6), fill(2,3))
        @test blocksizes(A+B) == BlockArrays.BlockSizes([1,1,1,1,2])

        A = BlockArray(randn(6,6), 1:3, 1:3)
        B = BlockArray(randn(6,6), fill(2,3), fill(3,2))

        @test blocksizes(A+B) == BlockArrays.BlockSizes([1,1,1,1,2], 1:3)
    end

    @testset "UnitRange" begin
        n = 3
        x = mortar([1:4n, 1:n])
        @test eltype(x.blocks) <: UnitRange
        y = 1:length(x)
        z = randn(size(x))
        x2 = vcat(x.blocks...)
        y2 = copy(y)
        z2 = copy(z)
        @test (@. z = x + y + z; z) == (@. z2 = x2 + y2 + z2; z2)
    end

    @testset "Special broadcast" begin
        v = mortar([1:3,4:7])
        @test broadcast(+, v) isa BlockVector{Int,Vector{UnitRange{Int}}}
        @test broadcast(+, v) == v
        @test broadcast(-, v) isa BlockVector{Int,Vector{StepRange{Int,Int}}}
        @test broadcast(-, v) == -v == -Vector(v)
        @test broadcast(+, v, 1) isa BlockVector{Int,Vector{UnitRange{Int}}}
        @test broadcast(+, v, 1) == Vector(v).+1
        @test broadcast(*, 2, v) isa BlockVector{Int,Vector{StepRange{Int,Int}}}
        @test broadcast(*, 2, v) == 2Vector(v)
    end
end


using Circuitscape
using Test
import Circuitscape: compute_single, compute_cholmod, compute_parallel,
                     runtests

# Unit tests for internals
@testset "Unit tests" begin
    include("internal.jl")
end

for f in (compute, compute_single, compute_cholmod, compute_parallel)
    runtests(f)
end