#!/usr/bin/env julia
###############################################################################

using Tables
using DataFrames

###############################################################################
# Config.

on_stock = true
# PCB, PSU...
fixed_cost = 0

###############################################################################

families = read_families()

#Read offers with params from XLS.
if !isfile("tmp/offers.xls")
	include("get_offers.jl")
end

offers = read_table("tmp/offers.xls", "offers")

###############################################################################

# Super table with all kind of params needed for check bellow.
uber_table = offers

# Hack for having some FPGAs in pinout zips.
a = families["Artix-7"]["dev_pack_banks"]
a[("XC7A12T", "CPG236")] = a[("XC7A15T", "CPG236")]
a[("XC7A12T", "CSG325")] = a[("XC7A15T", "CSG325")]
a[("XC7A25T", "CPG236")] = a[("XC7A35T", "CPG236")]
a[("XC7A25T", "CSG325")] = a[("XC7A35T", "CSG325")]

N_rows = size(uber_table)[1]

#TODO Should make metric for PCB price depending on package.
# Then little larger FPGAs will be better.
cost = uber_table[:price] .+ fixed_cost

###############################################################################

uber_table[:pins] = Vector{Float64}(N_rows)
for r in eachrow(uber_table)
	dpp_t = families[r[:family]]["dev_pack_pins"]
	# Make device-package pair.
	dp = (r[:device], r[:package])

	# Search for pins info with such device-package pair.
	i = findfirst(dpp_t[:dev_pack], dp)

	# No such device-package pair.
	if i == 0
		return 0
	end

	# Extract row for device-package pair, to harvest pins data.
	dpp_r = dpp_t[i, :]

	c = 0
	if haskey(dpp_r, :HR)
		c += dpp_r[:HR][1]
	end
	if haskey(dpp_r, :HP)
		c += dpp_r[:HP][1]
	end

	r[:pins] = c
end

uber_table[:cost_per_pin] = cost./uber_table[:pins]

###############################################################################

uber_table[:HR_DDR_MTps] = Vector{Int}(N_rows)
uber_table[:HP_DDR_MTps] = Vector{Int}(N_rows)
for r in eachrow(uber_table)
	ds = families[r[:family]]["ddr_speeds"]
	r[:HR_DDR_MTps] = ds["HR"][r[:speed_grade]]
	if haskey(ds, "HP")
		r[:HP_DDR_MTps] = ds["HP"][r[:speed_grade]]
	else
		r[:HP_DDR_MTps] = 0
	end
end

###############################################################################

pins_for_byte_group = 12 #TODO Could it be lower?

function DDR_bytes_groups_per_cols(pin_type, col)
	counts = Int[]
	for r in eachrow(uber_table)
		dpb = families[r[:family]]["dev_pack_banks"]
		# Make device-package pair.
		dp = (r[:device], r[:package])

		count = 0

		# No such device-package pair.
		if !haskey(dpb, dp)
			warn("No bank data for ", dp)
		else
			# Bank table.
			bt = dpb[dp]

			pin_type = string(pin_type)

			for r in eachrow(bt)
				bank = r[:bank]
				if bank != "NA" && length(bank) == 2 &&
					r[:pin_type] == pin_type &&	parse(Int, bank[1]) == col

					# All 4 byte groups needed.
					if r[:byte_group_0_pin_num] >= pins_for_byte_group &&
						r[:byte_group_1_pin_num] >= pins_for_byte_group &&
						r[:byte_group_2_pin_num] >= pins_for_byte_group &&
						r[:byte_group_3_pin_num] >= pins_for_byte_group
						count += 4
					end
				end
			end
		end

		push!(counts, count)
	end

	return counts
end
uber_table[:HR_DDR_BG_c1] = DDR_bytes_groups_per_cols(:HR, 1)
uber_table[:HR_DDR_BG_c2] = DDR_bytes_groups_per_cols(:HR, 2)
uber_table[:HR_DDR_BG_c3] = DDR_bytes_groups_per_cols(:HR, 3)
uber_table[:HR_DDR_BG_c4] = DDR_bytes_groups_per_cols(:HR, 4)
uber_table[:HP_DDR_BG_c1] = DDR_bytes_groups_per_cols(:HP, 1)
uber_table[:HP_DDR_BG_c2] = DDR_bytes_groups_per_cols(:HP, 2)
uber_table[:HP_DDR_BG_c3] = DDR_bytes_groups_per_cols(:HP, 3)
uber_table[:HP_DDR_BG_c4] = DDR_bytes_groups_per_cols(:HP, 4)
@assert sum(DDR_bytes_groups_per_cols(:HP, 5)) == 0

###############################################################################

function BG_to_B(BG)
	#TODO What if could make 6 bytes?
	if BG >= 11*4
		warn("Could set 4 DIMMs on one column?")
		return 8*4
	elseif BG >= 11*3
		return 8*3
	elseif BG >= 11*2
		return 8*2
	elseif BG >= 11
		return 8
	elseif BG >= 7
		return 4
	elseif BG >= 5
		return 2
	elseif BG >= 4
		return 1
	else
		return 0
	end
end

uber_table[:HR_DDR_B] = Vector{Int}(N_rows)
uber_table[:HP_DDR_B] = Vector{Int}(N_rows)
for r in eachrow(uber_table)
	r[:HR_DDR_B] =
		BG_to_B(r[:HR_DDR_BG_c1]) +
		BG_to_B(r[:HR_DDR_BG_c2]) +
		BG_to_B(r[:HR_DDR_BG_c3]) +
		BG_to_B(r[:HR_DDR_BG_c4])
	r[:HP_DDR_B] =
		BG_to_B(r[:HP_DDR_BG_c1]) +
		BG_to_B(r[:HP_DDR_BG_c2]) +
		BG_to_B(r[:HP_DDR_BG_c3]) +
		BG_to_B(r[:HP_DDR_BG_c4])
end
uber_table[:DDR_B] = uber_table[:HR_DDR_B] + uber_table[:HP_DDR_B]

###############################################################################

function BG_to_DIMM(BG)
	if BG >= 11*4
		warn("Could set 4 DIMMs on one column?")
		return 4
	elseif BG >= 11*3
		return 3
	elseif BG >= 11*2
		return 2
	elseif BG >= 11
		return 1
	else
		return 0
	end
end

N = size(uber_table)[1]
uber_table[:HR_DIMM] = Vector{Int}(N)
uber_table[:HP_DIMM] = Vector{Int}(N)
for r in eachrow(uber_table)
	r[:HR_DIMM] =
		BG_to_DIMM(r[:HR_DDR_BG_c1]) +
		BG_to_DIMM(r[:HR_DDR_BG_c2]) +
		BG_to_DIMM(r[:HR_DDR_BG_c3]) +
		BG_to_DIMM(r[:HR_DDR_BG_c4])
	r[:HP_DIMM] =
		BG_to_DIMM(r[:HP_DDR_BG_c1]) +
		BG_to_DIMM(r[:HP_DDR_BG_c2]) +
		BG_to_DIMM(r[:HP_DDR_BG_c3]) +
		BG_to_DIMM(r[:HP_DDR_BG_c4])
end
uber_table[:DIMM] = uber_table[:HR_DIMM] + uber_table[:HP_DIMM]

###############################################################################

uber_table[:cost_per_DDR_B] = cost./uber_table[:DDR_B]
uber_table[:cost_per_DIMM] = cost./uber_table[:DIMM]

###############################################################################

uber_table[:HR_DDR_MBps] = uber_table[:HR_DDR_B] .* uber_table[:HR_DDR_MTps]
uber_table[:HP_DDR_MBps] = uber_table[:HP_DDR_B] .* uber_table[:HP_DDR_MTps]
uber_table[:HR_DIMM_MBps] = 8*uber_table[:HR_DIMM] .* uber_table[:HR_DDR_MTps]
uber_table[:HP_DIMM_MBps] = 8*uber_table[:HP_DIMM] .* uber_table[:HP_DDR_MTps]

uber_table[:DDR_MBps] = uber_table[:HR_DDR_MBps] + uber_table[:HP_DDR_MBps]
uber_table[:DIMM_MBps] = uber_table[:HR_DIMM_MBps] + uber_table[:HP_DIMM_MBps]

###############################################################################

uber_table[:cost_per_DDR_MBps] = cost./uber_table[:DDR_MBps]
uber_table[:cost_per_DIMM_MBps] = cost./uber_table[:DIMM_MBps]

###############################################################################

macro save(table)
	quote
		n = $(string(table))
		write_table("tmp/" * n * ".xlsx", n, $table)
	end
end

# Save it just for documentation.
@save uber_table

###############################################################################

# Select those who have needed stock.
if on_stock
	uber_table = uber_table[uber_table[:stock_vs_need] .>= 0, :]
end

###############################################################################
# Cost.

price = deepcopy(uber_table)
sort!(price, cols = :price)

###############################################################################
# Cost per pin.

cost_per_pin = deepcopy(uber_table)
sort!(cost_per_pin, cols = :cost_per_pin)

if false
	# Cheapest from family:
	println("Cheapest Artix-7:")
	println(
		sort(
			cost_per_pin[cost_per_pin[:family] .== "Artix-7", :],
			cols = :price
		)[1, :]
	)
	println("Cheapest Kintex-7:")
	println(
		sort(
			cost_per_pin[cost_per_pin[:family] .== "Kintex-7", :],
			cols = :price
		)[1, :]
	)
	println("Cheapest Virtex-7:")
	println(
		sort(
			cost_per_pin[cost_per_pin[:family] .== "Virtex-7", :],
			cols = :price
		)[1, :]
	)

	@save cost_per_pin
	println("Cheapers per pin:")
	best = cost_per_pin[1, :]
	println(best)
end

###############################################################################
# Cost per memory width.

cost_per_DDR_B = deepcopy(uber_table)
sort!(cost_per_DDR_B, cols = :cost_per_DDR_B)

cost_per_DIMM = deepcopy(uber_table[uber_table[:DIMM] .!= 0, :])
sort!(cost_per_DIMM, cols = :cost_per_DIMM)

###############################################################################
# Cost per memory bandwidth.

cost_per_DDR_MBps = deepcopy(uber_table)
sort!(cost_per_DDR_MBps, cols = :cost_per_DDR_MBps)

cost_per_DIMM_MBps = deepcopy(uber_table[uber_table[:DIMM] .!= 0, :])
sort!(cost_per_DIMM_MBps, cols = :cost_per_DIMM_MBps)
@save cost_per_DIMM_MBps

best_for_DDR_is_best_for_DIMM =
	cost_per_DDR_MBps[1, :] == cost_per_DIMM_MBps[1, :]

@show best_for_DDR_is_best_for_DIMM

println("Best (cost per bandwidth):")
println(cost_per_DIMM_MBps[1, :])

###############################################################################

#TODO Check DSP performance and bandwidth with Wang ratio.

#TODO Check:
# + cost per pin
# + cost per pin bandwidth
# + cost per SO-DIMM badwidth (sharing common pins or not)
# - compare price with and without stock
# - cost with and without PCIe
# - without PCIe and user rest of SO-DIMM pins for parallel bus.

###############################################################################

println("End")

###############################################################################


if !@isdefined USE_GPU
    const USE_GPU = length(ARGS) > 0 && ARGS[1] == "gpu"
end
using CUDA
using LinearAlgebra
using DataFrames, CSV
using BenchmarkTools, Statistics
using ParallelStencil
using ParallelStencil.FiniteDifferences2D
if !ParallelStencil.is_initialized()
    @static if USE_GPU
        @init_parallel_stencil(CUDA, Float64, 2)
    else
        @init_parallel_stencil(Threads, Float64, 2)
    end
end
include("../scripts-part2/multigrid.jl")

const FILENAME = joinpath("benchmark-results", "bench_multigrid_$(USE_GPU ? "gpu" : "cpu").csv")
if isfile(FILENAME)
    bench_df = DataFrame(CSV.File(FILENAME))
else
    bench_df = DataFrame()
end
#for execution_policy in [parallel, parallel_shmem], l in 2:3, k in 7:10
for k in 4:13, l in 2:min(k-4, 8), solver in [jacobi, conjugate_gradient], execution_policy in [parallel, parallel_shmem]
    @info "Running with params" USE_GPU, execution_policy, l, k, solver
    n = (2^k)+1
    h = 1 / (n - 1)
    c = 0.0
    tol = 1e-6
    Nmax = 100

    opt = MGOpt()
    opt.execution_policy = execution_policy
    opt.coarse_solve_size = (2^l)+1
    opt.coarse_solver = solver

    for seed in 1:1
        try
            bench_res = @benchmark MGsolve_2DPoisson!(x, b, $h, $c, $tol, $Nmax, false; opt=$opt, verbose=false) setup=((x, b) = (@zeros($n, $n), @rand($n, $n))) evals=1
            bench_row = (
                 execution_policy = string(execution_policy),
                 coarse_solver = string(solver),
                 k = k,  # exponent for computing size of largest grid with nx = (2^k)+1
                 l = l,  # same as k but for the coarsest grid
                 median_time = median(bench_res).time / 1e9,
                 mean_time = mean(bench_res).time / 1e9,
                 std_time = std(bench_res).time / 1e9,
                 seed = seed,
                 use_gpu = USE_GPU,
                 nthreads = Threads.nthreads(),
                )
            push!(bench_df, bench_row)

            open(FILENAME, "w") do io
                CSV.write(io, bench_df)
            end
        catch e
            @warn "Failed with " e
        end
    end
end


#Setup

include("make_plots_distributed.jl")

using .MakePlots

println("Ready to set up workers.")
flush(stdout)
MakePlots.setup() # setup workers based on cpu counts

#Run simulation

d_tasks = parse(Int, ENV["SLURM_ARRAY_TASK_ID"]) # 1 - 3

job_id = "rz_0.001"
samples = 1000*20000 # 20,000 samples take about 1min per cpu (1 plot, 1 dist, 10 angles)
dists = [(1 + 2 * d_tasks):2:(2 + 2 * d_tasks)...] # code distance: 3 or 5 or 7 separately
plots = [4] # list of plot types
plot_id = 4 # type of plots to be made

println("Simulation begins... (", samples, " samples)")
flush(stdout)

results = MakePlots.dist_calc_all(job_id, dists, samples, plots) # start computing and wait until complete
println("Simulation completed.")

collected = MakePlots.fetch_finished(results, job_id, plot_id)
println(collected)


doc"""
Various parameters for the solve! call.
```
  maxiter := maximum total number of iterations
  erroriter := number of iterations between updating the SolverStats object and checking stopping conditions
  breakiter := number of iterations between the `onbreak` callback
  stopepochs := number of epochs (1 epoch == 1 update to stats) without improvement for early stopping
  minerror := when validation error drops below this, stop
  onbreak := callback function, run once every breakiter iterations. function args: (net::NeuralNet, solverParams::SolverParams, stats::SolverStats)
```
"""
type SolverParams
  maxiter::Int   # maximum total number of iterations
  erroriter::Int  # number of iterations be
  breakiter::Int
  plotiter::Int
  plotfields::Vector{Symbol}
  stopepochs::Int
  minerror::Float64
  onbreak::Function
end

function SolverParams(; maxiter = 1000,
                        erroriter = 1000,
                        breakiter = 10000,
                        plotiter = -1,
                        plotfields = Symbol[:x, :xhat, :y, :Σ, :a],
                        stopepochs = 100,
                        minerror = 1e-5,
                        onbreak = donothing)
  SolverParams(maxiter, erroriter, breakiter, plotiter, plotfields, stopepochs, minerror, onbreak)
end

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

@enum SolverStatus RUNNING CONVERGED STOPPEDEARLY MAXITER

type SolverStats
  numiter::Int
  trainError::Float64
  validationError::Float64
  bestValidationError::Float64
  epochSinceImprovement::Int
  status::SolverStatus
  bestModel
  plotter
end
SolverStats() = SolverStats(0, Inf, Inf, Inf, 0, RUNNING, nothing, nothing)

Base.string(stats::SolverStats) = "SolverStats{$(stats.status) n=$(stats.numiter), trainerr=$(stats.trainError), valerr=$(stats.validationError), besterr=$(stats.bestValidationError), epochSinceImprovement=$(stats.epochSinceImprovement)}"
Base.print(io::IO, stats::SolverStats) = print(io, string(stats))
Base.show(io::IO, stats::SolverStats) = print(io, string(stats))

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

doc"""
Batch stochastic gradient descent solver.  Sample from training data to update net, stop if converged
or if validation error is not improving.  Returns SolverStats summary object.
"""
function solve!(net::NetStat, traindata::DataSampler, validationdata::DataSampler, transformY::Bool = false)

  stats = SolverStats()

  # optionally set up the plotter (only if plotiter ≥ 0)
  if net.solverParams.plotiter >= 0
    stats.plotter = NetProgressPlotter(net, stats, net.solverParams.plotfields)
  end

  println("\nsolve: $net\n")

  # loop through maxiter times
  for i in 1:net.solverParams.maxiter

    stats.numiter += 1

    # randomly sample one data item and update the network
    data = sample(traindata)
    fit!(net, data, transformY)

    # update the plot?
    piter = net.solverParams.plotiter
    if (piter > 0 && i % net.solverParams.plotiter == 0) || piter == 0
      fit!(stats.plotter)
    end

    # # check for convergence
    if i % net.solverParams.erroriter == 0
      stats.trainError = totalCost(net, traindata)
      stats.validationError = totalCost(net, validationdata)
      println("Status: $stats")

      # check for improvement in validation error
      if stats.validationError < stats.bestValidationError
        stats.bestValidationError = stats.validationError
        stats.bestModel = copy(net)
        stats.epochSinceImprovement = 0
      else
        stats.epochSinceImprovement += 1

        # early stopping... no improvement
        if stats.epochSinceImprovement >= net.solverParams.stopepochs
          println("Early stopping: $stats")
          stats.status = STOPPEDEARLY
          return stats
        end
      end

      # check if our error is low enough
      if stats.validationError <= net.solverParams.minerror
        println("Converged, breaking: $stats")
        stats.status = CONVERGED
        return stats
      end
    end

    # take a break?
    if i % net.solverParams.breakiter == 0
      net.solverParams.onbreak(net, stats)
    end
  end

  println("Maxiter reached: $stats")
  stats.status = MAXITER
  stats
end


mutable struct SwitchSimulateState{T}
    score::Float64
    noise::Float64
    index::Int
    subtrace::Trace
    retval::T
    SwitchSimulateState{T}(score::Float64, noise::Float64) where T = new{T}(score, noise)
end

function process!(gen_fn::Switch{C, N, K, T},
                  index::Int,
                  args::Tuple,
                  state::SwitchSimulateState{T}) where {C, N, K, T}
    local retval::T
    subtrace = simulate(getindex(gen_fn.branches, index), args)
    state.index = index
    state.noise += project(subtrace, EmptySelection())
    state.subtrace = subtrace
    state.score += get_score(subtrace)
    state.retval = get_retval(subtrace)
end

@inline process!(gen_fn::Switch{C, N, K, T}, index::C, args::Tuple, state::SwitchSimulateState{T}) where {C, N, K, T} = process!(gen_fn, getindex(gen_fn.cases, index), args, state)

function simulate(gen_fn::Switch{C, N, K, T},
                  args::Tuple) where {C, N, K, T}

    index = args[1]
    state = SwitchSimulateState{T}(0.0, 0.0)
    process!(gen_fn, index, args[2 : end], state)
    return SwitchTrace(gen_fn, state.subtrace,
                       state.retval, args,
                       state.score, state.noise)
end


@testset "Count2Reification" begin
@testset "CountEqualTo2Reification: $(fct_type), $(array_size) items, $(T)" for fct_type in ["vector of variables", "vector affine function"], array_size in [2, 3], T in [Int, Float64]
    base_model = if T == Int
        IntReificationEqualToModel{T}()
    elseif T == Float64
        FloatReificationEqualToModel{T}()
    else
        @assert false
    end
    mock = MOIU.MockOptimizer(base_model)
    model = COIB.CountEqualTo2Reification{T}(mock)

    @test MOI.supports_constraint(model, MOI.VariableIndex, MOI.ZeroOne)
    @test MOI.supports_constraint(model, MOI.VariableIndex, MOI.Integer)
    @test MOI.supports_constraint(
        model,
        MOI.VectorAffineFunction{T},
        CP.Reification{MOI.EqualTo{T}},
    )
    @test MOIB.supports_bridging_constraint(
        model,
        MOI.VectorOfVariables,
        CP.Count{MOI.EqualTo{T}},
    )

    x_count, _ = MOI.add_constrained_variable(model, MOI.Integer())

    if T == Int
        x_array, _ = MOI.add_constrained_variables(model, [MOI.Integer() for _ in 1:array_size])
    elseif T == Float64
        x_array = MOI.add_variables(model, array_size)
    end

    fct = if fct_type == "vector of variables"
        MOI.VectorOfVariables([x_count, x_array...])
    elseif fct_type == "vector affine function"
        MOIU.vectorize([x_count, x_array...])
    else
        @assert false
    end
    c = MOI.add_constraint(model, fct, CP.Count(array_size, MOI.EqualTo(zero(T))))

    @test MOI.is_valid(model, x_count)
    for i in 1:array_size
        @test MOI.is_valid(model, x_array[i])
    end
    @test MOI.is_valid(model, c)

    bridge = MOIBC.bridges(model)[MOI.ConstraintIndex{MOI.VectorOfVariables, CP.Count{MOI.EqualTo{T}}}(-1)]

    @testset "Bridge properties" begin
        @test MOIBC.concrete_bridge_type(typeof(bridge), MOI.VectorOfVariables, CP.Count{MOI.EqualTo{T}}) == typeof(bridge)
        @test MOIB.added_constrained_variable_types(typeof(bridge)) == [(MOI.ZeroOne,)]
        @test MOIB.added_constraint_types(typeof(bridge)) == [
            (MOI.VectorAffineFunction{T}, CP.Reification{MOI.EqualTo{T}}),
            (MOI.ScalarAffineFunction{T}, MOI.EqualTo{T}),
        ]

        @test MOI.get(bridge, MOI.NumberOfVariables()) == array_size
        @test MOI.get(bridge, MOI.NumberOfConstraints{MOI.VariableIndex, MOI.ZeroOne}()) == array_size
        @test MOI.get(bridge, MOI.NumberOfConstraints{MOI.VectorAffineFunction{T}, CP.Reification{MOI.EqualTo{T}}}()) == array_size
        @test MOI.get(bridge, MOI.NumberOfConstraints{MOI.ScalarAffineFunction{T}, MOI.EqualTo{T}}()) == 1

        @test MOI.get(bridge, MOI.ListOfVariableIndices()) == bridge.vars
        @test MOI.get(bridge, MOI.ListOfConstraintIndices{MOI.VariableIndex, MOI.ZeroOne}()) == bridge.vars_bin
        @test MOI.get(bridge, MOI.ListOfConstraintIndices{MOI.VectorAffineFunction{T}, CP.Reification{MOI.EqualTo{T}}}()) == bridge.cons
        @test MOI.get(bridge, MOI.ListOfConstraintIndices{MOI.ScalarAffineFunction{T}, MOI.EqualTo{T}}()) == [bridge.con_sum]
    end

    @testset "Set of variables" begin
        @test length(bridge.vars) == array_size
        @test length(bridge.vars_bin) == array_size
        for i in 1:array_size
            @test MOI.is_valid(model, bridge.vars[i])
            @test MOI.is_valid(model, bridge.vars_bin[i])
            @test MOI.get(model, MOI.ConstraintFunction(), bridge.vars_bin[i]) == bridge.vars[i]
            @test MOI.get(model, MOI.ConstraintSet(), bridge.vars_bin[i]) == MOI.ZeroOne()
        end
    end

    @testset "Reification" begin
        @test length(bridge.cons) == array_size
        for i in 1:array_size
            @test MOI.is_valid(model, bridge.cons[i])
            f = MOI.get(model, MOI.ConstraintFunction(), bridge.cons[i])
            @test length(f.terms) == 2
            @test MOI.get(model, MOI.ConstraintSet(), bridge.cons[i]) == CP.Reification(MOI.EqualTo(zero(T)))

            t1 = f.terms[1]
            @test t1.output_index == 1
            @test t1.scalar_term.coefficient === one(T)
            @test t1.scalar_term.variable == bridge.vars[i]

            t2 = f.terms[2]
            @test t2.output_index == 2
            @test t2.scalar_term.coefficient === one(T)
            @test t2.scalar_term.variable == x_array[i]
        end
    end

    @testset "Sum" begin
        @test MOI.is_valid(model, bridge.con_sum)
        f = MOI.get(model, MOI.ConstraintFunction(), bridge.con_sum)
        @test length(f.terms) == array_size + 1
        @test MOI.get(model, MOI.ConstraintSet(), bridge.con_sum) == MOI.EqualTo(zero(T))

        t = f.terms[1]
        @test t.coefficient === -one(T)
        @test t.variable == x_count

        for i in 1:array_size
            t = f.terms[1 + i]
            @test t.coefficient === one(T)
            @test t.variable == bridge.vars[i]
        end
    end
end
end


using aoc2019.computer: load_program, io, run
using Match

input = joinpath(@__DIR__, "input")
p = load_program(input)

function camera!(channel, view)
    i, j = 0, 0
    for msg in channel
        if msg == 10
            i = 0
            j += 1
        else
            view[Pair(i, j)] = msg
            i += 1
        end
    end
end

function calibrate(view)
    scaffolds = collect(keys(filter(kv -> kv.second != 46, view)))
    acc = 0
    for kv in scaffolds
        x, y = kv
        up = Pair(x, y + 1)
        down = Pair(x, y - 1)
        right = Pair(x + 1, y)
        left = Pair(x - 1, y)
        if all(dir -> dir in scaffolds, [up, down, right, left])
            acc += x * y
        end
    end
    return acc
end

function display(view)
    i, j = 0, 0
    mx = maximum(map(kv -> kv.first, collect(keys(view))))
    my = maximum(map(kv -> kv.second, collect(keys(view))))
    for j in 0:my
        for i in 0:mx
            print(convert(Char, view[Pair(i,j)]))
            i += 1
        end
        println()
        i = 0
        j += 1
    end
end

function path(view)
    p = []
    robot = filter(kv -> kv.second != 46 && kv.second != 35, view)
    position, facing = first(robot)
    d = direction(facing)
    while true
        # must turn, try each and see if there's a scaffold in that direction
        rx, ry = turn_right(d)
        r = Pair(position.first + rx, position.second + ry)
        lx, ly = turn_left(d)
        l = Pair(position.first + lx, position.second + ly)
        if haskey(view, r) && view[r] != 46
            d = Pair(rx, ry)
            steps = 0
            while true
                next = Pair(position.first + d.first, position.second + d.second)
                if haskey(view, next) && view[next] != 46
                    steps += 1
                    position = next
                else
                    break
                end
            end
            push!(p, ("R", string(steps)))
        elseif haskey(view, l) && view[l] != 46
            d = Pair(lx, ly)
            steps = 0
            while true
                next = Pair(position.first + d.first, position.second + d.second)
                if haskey(view, next) && view[next] != 46
                    steps += 1
                    position = next
                else
                    break
                end
            end
            push!(p, ("L", string(steps)))
        else
            break
        end
    end
    return p
end

function turn_right(m)
    adj = Dict(
        Pair(0, 1) => Pair(-1, 0),
        Pair(1, 0) => Pair(0, 1),
        Pair(0, -1) => Pair(1, 0),
        Pair(-1, 0) => Pair(0, -1)
    )
    return adj[m]
end

function turn_left(m)
    adj = Dict(
        Pair(0, 1) => Pair(1, 0),
        Pair(1, 0) => Pair(0, -1),
        Pair(0, -1) => Pair(-1, 0),
        Pair(-1, 0) => Pair(0, 1)
    )
    return adj[m]
end

function direction(char)
    if char == convert(Int, '^')
        return Pair(0, -1)
    elseif char == convert(Int, '>')
        return Pair(1, 0)
    elseif char == convert(Int, 'v')
        return Pair(0, 1)
    elseif char == convert(Int, '<')
        return Pair(-1, 0)
    else
        throw("unreachable?")
    end
end

function send_input!(channel, instructions)
    reverse!(instructions)
    while true
        if isempty(instructions)
            put!(channel, 10)
            break
        else
            put!(channel, pop!(instructions))
        end
    end
end

function how_much_dust(channel)
    last = nothing
    while true
        v = take!(stdout)
        ## ignore ascii output
        if v < 128
            nothing
        else
            last = v
            break
        end
    end
    return last
end

stdin, stdout = io()
task = @async run(copy(p), stdin, stdout)

view = Dict{Pair{Int, Int}, Int}()

# let program run
while !istaskdone(task); yield() end
# close channel so it can be iterated over
close(stdout)
# populate view
camera!(stdout, view)
# display(view)

p1 = calibrate(view)
@assert p1 == 7280

# generate the path from start to end as turn-step pairs
# route = path(view)

## i compressed this by hand. no good ideas on how to do it progmatically
routine = map(c -> convert(Int, c), ['A', ',', 'B', ',', 'A', ',', 'C', ',', 'B', ',', 'C', ',', 'A', ',', 'B', ',', 'A', ',', 'C'])
a = map(c -> convert(Int, c), ['R', ',', '1', '0', ',', 'L', ',', '8', ',', 'R', ',', '1', '0', ',', 'R', ',', '4'])
b = map(c -> convert(Int, c), ['L', ',', '6', ',', 'L', ',', '6', ',', 'R', ',', '1', '0'])
c = map(c -> convert(Int, c), ['L', ',', '6', ',', 'R', ',', '1', '2', ',', 'R', ',', '1', '2', ',', 'R', ',', '1', '0'])

p[1] = 2
stdin, stdout = io()
task = @async run(copy(p), stdin, stdout)

send_input!(stdin, routine)
send_input!(stdin, a)
send_input!(stdin, b)
send_input!(stdin, c)
# no live output
send_input!(stdin, [convert(Int, 'n')])

# wait for program to run
while !istaskdone(task); yield() end
# close channel so we don't block
close(stdout)

p2 = how_much_dust(stdout)
@assert p2 == 1045393

print("-----------------------------------------------------------------------\n")
print("set and forget -- part one\n    calibration : $p1\n")
print("set and forget -- part two\n    dust        : $p2\n")
print("-----------------------------------------------------------------------\n")


symbolmatch(s::Symbol, st::AS) where {AS <: AbstractString} = ismatch(Regex("$(st)"), string(s))
symbolmatch(s::Symbol, st::Char) = symbolmatch(s, string(st))

landmarksymbol(i::Int) = Symbol("l$i")
landmarknumber(sy::Symbol) = parse(Int, string(sy)[2:end])

function addoccurance!(lmo::Dict{Symbol, Int}, lm::Symbol)
  lmo[lm] = haskey(lmo, lm) ? lmo[lm]+1 : 1
  nothing
end

"""
    $(SIGNATURES)

For multiple measurements::Dict{Int,T} with unique Int signature, add "LANDMARK" `;softtype=Point2` to the fg<:AbstractDFG.  Type T not used in this function.
"""
function findAddNewLandmarks!(fgl::G, lmoccur::Dict{Symbol, Int}, measurements::Dict{Int, T}; min_occurance::Int=2, softtype=Point2) where {G <: AbstractDFG, T}
  # Count occurances of each landmark
  lmsyms = landmarksymbol.(collect(keys(measurements)))
  addoccurance!.(lmoccur, lmsyms)

  # add new landmarks not already in factor graph
  lmsinfgl = ls(fgl)[2]
  for lm in lmsyms
    if lmoccur[lm] >= min_occurance
      if !(lm in lmsinfgl)
        # add a new landmark
        println("adding lm=$(lm)")
        addVariable!(fgl, lm, softtype, tags=[:LANDMARK;])
      end
      # @show measurements[landmarknumber(lm)]
    end
  end
  nothing
end

function findAddLandmFactorsByPose!(fgl::G, psym::Symbol, meas::Dict{Int, T};
                    range_sigma::Float64=1.0, bearing_sigma::Float64=0.05, N::Int=100) where {G <: AbstractDFG, T}
  lmsyms = landmarksymbol.(collect(keys(meas)))
  posels = ls(fgl, psym)
  lmsinfg = ls(fgl)[2]
  lms = intersect(lmsyms, lmsinfg)
  for lm in lms
    btwfcts = intersect(ls(fgl, lm), posels)
    if length(btwfcts) == 0
      z_rb = meas[landmarknumber(lm)]
      println("addFactor! $(psym) -- $(lm), z_rb=$(z_rb)")
      pp = Pose2Point2BearingRange( Normal(z_rb[2],bearing_sigma),Normal(z_rb[1],range_sigma) )
      addFactor!(fgl, [psym; lm], pp)
    end
  end
  nothing
end

function drive!(fgl::G, idx::Int, odos::Dict, meas::Dict, lmoccurl::Dict, Podo;
                N::Int=100, min_occurance::Int=2 ) where G <: AbstractDFG
  #
  prev, X, nextn = getLastPose2D(fgl)
  vp, fp = addOdoFG!(fgl, nextn, odos[idx][1:3], Podo, N=N)

  findAddNewLandmarks!(fgl, lmoccurl, meas[idx], min_occurance=min_occurance)

  # walk backwards
  donels = Symbol[]
  for i in 0:(min_occurance-1)
    @show psym = Symbol("x$(idx-i-1)") # off by one
    if idx-i > 0
      currmeas = meas[idx-i]
      findAddLandmFactorsByPose!(fgl, psym, currmeas, N=N)
    end
  end
  nothing
end

function donextframe!(fgl::G, idx::Int, odos::Dict, meas::Dict, lmoccurl::Dict, Podo; N::Int=100) where G <: AbstractDFG
  drive!(fgl, idx, odos, meas, lmoccurl, Podo, N=N)

  tree = wipeBuildNewTree!(fgl, drawpdf=true)
  inferOverTree!(fgl, tree,N=N)

  drawPosesLandms(fgl)
end


export OfflinePolicy, RLTransition

struct RLTransition
    state
    action
    reward
    terminal
    next_state
end

Base.@kwdef struct OfflinePolicy{L,T} <: AbstractPolicy
    learner::L
    dataset::T
    continuous::Bool
    batch_size::Int
end

(π::OfflinePolicy)(env) = π(env, ActionStyle(env), action_space(env))

function (π::OfflinePolicy)(env, ::MinimalActionSet, ::Base.OneTo)
    if π.continuous
        π.learner(env)
    else
        findmax(π.learner(env))[2]
    end
end
(π::OfflinePolicy)(env, ::FullActionSet, ::Base.OneTo) = findmax(π.learner(env), legal_action_space_mask(env))[2]

function (π::OfflinePolicy)(env, ::MinimalActionSet, A)
    if π.continuous
        π.learner(env)
    else
        A[findmax(π.learner(env))[2]]
    end
end
(π::OfflinePolicy)(env, ::FullActionSet, A) = A[findmax(π.learner(env), legal_action_space_mask(env))[2]]

function RLBase.update!(
    p::OfflinePolicy,
    traj::AbstractTrajectory,
    ::AbstractEnv,
    ::PreExperimentStage,
)
    l = p.learner
    if in(:pretrain_step, fieldnames(typeof(l)))
        println("Pretrain...")
        for _ in 1:l.pretrain_step
            inds, batch = sample(l.rng, p.dataset, p.batch_size)
            update!(l, batch)
        end
    end
end

function RLBase.update!(
    p::OfflinePolicy,
    traj::AbstractTrajectory,
    ::AbstractEnv,
    ::PreActStage,
)
    l = p.learner
    l.update_step += 1

    if in(:target_update_freq, fieldnames(typeof(l))) && l.update_step % l.target_update_freq == 0
        copyto!(l.target_approximator, l.approximator)
    end

    l.update_step % l.update_freq == 0 || return

    inds, batch = sample(l.rng, p.dataset, p.batch_size)

    update!(l, batch)
end

function StatsBase.sample(rng::AbstractRNG, dataset::Vector{T}, batch_size::Int) where {T}
    valid_range = 1:length(dataset)
    inds = rand(rng, valid_range, batch_size)
    batch_data = dataset[inds]
    s_length = size(batch_data[1].state)[1]

    s = Array{Float32}(undef, s_length, batch_size)
    s′ = Array{Float32}(undef, s_length, batch_size)
    a = []
    r = []
    t = []
    for (i, data) in enumerate(batch_data)
        s[:, i] = data.state
        push!(a, data.action)
        s′[:, i] = data.next_state
        push!(r, data.reward)
        push!(t, data.terminal)
    end
    batch = NamedTuple{SARTS}((s, a, r, t, s′))
    inds, batch
end

"""
    calculate_CQL_loss(q_value, action; method)
See paper: [Conservative Q-Learning for Offline Reinforcement Learning](https://arxiv.org/abs/2006.04779)
"""
function calculate_CQL_loss(q_value::Matrix{T}, action::Vector{R}; method = "CQL(H)") where {T, R}
    if method == "CQL(H)"
        cql_loss = mean(log.(sum(exp.(q_value), dims=1)) .- q_value[action])
    else
        @error Wrong method parameter
    end
    return cql_loss
end


@testset "Test MSTM v4" begin
    using STMMRunner
    using STMMRunner.MSTM.V4

    sphere1 = Sphere(r = 1.0, x = 4.2, y = 4.3, z = 4.0, m = 1.5)
    sphere2 = Sphere(r = 1.0, x = 0.0, y = 0.3, z = 1.0, m = 1.5)

    spheres = [sphere1, sphere2]
    layers = [Layer(), Layer(d = 10.0, m = 2.0), Layer(m = 1.2)]

    base = STMMConfig(
        number_processors = Sys.CPU_THREADS ÷ 2,
        scattering_map_model = 0,
        α = 5.0,
        β = 12.0,
        spheres = spheres,
        frame = IncidentFrame,
    )

    @testset "Fixed orientation" begin
        fixed_base = STMMConfig(
            base;
            orientation = FixedOrientation,
            calculate_near_field = true,
            near_field_x_min = -5.0,
            near_field_x_max = 5.0,
            near_field_y_min = -5.0,
            near_field_y_max = 5.0,
            near_field_z_min = -5.0,
            near_field_z_max = 5.0,
            near_field_step_size = 1.0
        )

        @testset "No layer boundaries, scattering_map_model = 0, frame = incident" begin
            param = STMMConfig(fixed_base)
            @test !isnothing(run_mstm(param; keep = true))
        end

        @testset "No layer boundaries, scattering_map_model = 0, frame = incident, use_fft = true" begin
            param = STMMConfig(fixed_base; use_fft_translation = true)
            @test !isnothing(run_mstm(param; keep = true))
        end

        @testset "No layer boundaries, scattering_map_model = 0, frame = target" begin
            param = STMMConfig(fixed_base; frame = TargetFrame)
            @test !isnothing(run_mstm(param; keep = true))
        end

        @testset "No layer boundaries, periodic" begin
            param = STMMConfig(fixed_base; periodic = true)
            @test !isnothing(run_mstm(param; keep = true))
        end

        @testset "No layer boundaries, scattering_map_model = 1, frame = incident" begin
            param = STMMConfig(
                fixed_base;
                scattering_map_model = 1,
                scattering_map_dimension = 20
            )
            @test !isnothing(run_mstm(param; keep = true))
        end

        @testset "With layer boundaries, scattering_map_model = 0, frame = incident" begin
            param = STMMConfig(fixed_base; layers = layers)
            @test !isnothing(run_mstm(param; keep = true))
        end

        @testset "With layer boundaries, periodic" begin
            param = STMMConfig(fixed_base; layers = layers, periodic = true)
            @test !isnothing(run_mstm(param; keep = true))
        end

        @testset "With layer boundaries, scattering_map_model = 1, frame = incident" begin
            param = STMMConfig(
                fixed_base;
                layers = layers,
                scattering_map_model = 1,
                scattering_map_dimension = 20
            )
            @test !isnothing(run_mstm(param; keep = true))
        end
    end

    @testset "Random orientation" begin
        random_base = STMMConfig(base; orientation = RandomOrientation)

        @testset "frame = incident" begin
            param = STMMConfig(random_base)
            @test !isnothing(run_mstm(param; keep = true))
        end

        @testset "frame = target" begin
            param = STMMConfig(random_base; frame = TargetFrame)
            @test !isnothing(run_mstm(param; keep = true))
        end

        @testset "use Monte Carlo" begin
            param = STMMConfig(
                random_base;
                use_monte_carlo_integration = true,
                number_incident_directions = 10
            )
            @test !isnothing(run_mstm(param; keep = true))
        end
    end

    try
        for path in readdir()
            if isdir(path) && occursin("mstm_tmp", path)
                rm(path; force = true, recursive = true)
            end
        end
    catch _ # The unlink operation might fail on Windows so we need this try-catch
    end
end


## Esta función da la distancia de seguridad que debe de haber entre dos vehículos
function distancia_seguridad(i, Carretera1, N, delta, alfa::Float32 = 0.75f0)

    pvfo = Vehiculo_Adelante(i, Carretera1, N)
    d_ifo = (pvfo < N ? (pvfo-i-1) : delta*5 ) # 5 = vmax[2]
    dfo = d_ifo + int8( floor( (1.-alfa) * Carretera1[pvfo].velocidad + 0.5) )
    alfa = 0; pvfo = 0; d_ifo = 0
    return dfo
end

## Esta función hace que un vehículo cambie de carril
function Cambiar_Carril!(v0::Vehiculo, Cuenta0, C1::Carretera1D)

    i = v0.posicion
    if C1.carretera[i].tipo == -1

        pvbo = Vehiculo_Atras(i, C1.carretera, C1.N)
        ############################ cambio ###################################
        d_bo = (pvbo < (C1.N) ? (i-pvbo-1) : 5 ) # 5 = vmax[2] (i-pvbo-longitud)
        vb = C1.carretera[pvbo].velocidad
        if vb <= d_bo

            Cambiar_Vehiculo!( C1.carretera[i], v0.velocidad, v0.posicion, v0.tipo, 1, v0.num)
            Celda_Vacia!( v0 )
            Cuenta0 -= 1
            C1.cuenta += 1
        end
    end
end

## Esta función hace los cambios de carril desde un carril izquierdo hacia uno derecho
function CambioIzq_Der!(Ck, Ckmenos1, k)

    for v in Ck.carretera[end:-1:9]
        if (v.tipo == 1 || v.tipo == 2) && v.cambio == 0

            delta = v.tipo == 2 ? int8(3) : int8(1)
            l = v.velocidad
            dfo = distancia_seguridad(v.posicion, Ckmenos1.carretera, Ckmenos1.N, delta)
            dfs = distancia_seguridad(v.posicion, Ck.carretera, Ck.N, delta)

            if v.tipo == 2 && dfo >= l*delta && (dfs  >= delta*l || l > dfs)
                Cambiar_Carril!(v, Ck.cuenta, Ckmenos1)
            end

            if v.tipo == 1 && dfo >= l*delta #(dfs  >= delta*l || l > dfs)
                Cambiar_Carril!(v, Ck.cuenta, Ckmenos1)
            end
            l = dfo = dfs = delta = 0
        end
    end
end

## Esta función hace los cambios de carril desde un carril derecho hacia uno izquierdo
function CambioDer_Izq!(Ck, Ckmas1, k)

    for v in Ck.carretera[end:-1:9]
        #if (v.tipo == 2 || (v.tipo == 1 && k < 2)) && v.cambio == 0
        if v.tipo == 2 && v.cambio == 0
            l = v.velocidad
            dfs = distancia_seguridad(v.posicion, Ck.carretera, Ck.N, int8(1))

            if dfs < l
                if Ckmas1.carretera[v.posicion].tipo == -1
                    dfo = distancia_seguridad(v.posicion, Ckmas1.carretera, Ckmas1.N, int8(1))
                    if dfo >= l
                        Cambiar_Carril!(Ck.carretera[v.posicion], Ck.cuenta, Ckmas1)
                    end
                end
            end
            dfs = l = dfo = 0
        end
    end
end

## Esta función simula una rampa de entrada (tipo == 1) o de salida (tipo == 0)
function Rampa!( tipo, x0, lramp, pin, p, Ck, num, vmax::Array{Int8, 1} = Int8[3, 5])
    #vmax = [3, 5]

    # Rampa de salida, tipo == 0
    if tipo == 0
    ############################ cambio ###################################
        x = Pos_der(Ck.carretera[x0:x0+lramp])  # if longitud > 1 Ck.carretera[xi] && ... && Ck.carretera[xi+londitud] esten vacia
        if x <= x0+lramp && rand() < pin
            if Ck.cuenta > 0
                Ck.cuenta -= 1
            end
            Celda_Vacia!(Ck.carretera[x])
        end

    # Rampa de entrada, tipo == 1
    elseif tipo == 1
    ############################ cambio ###################################
        x = Pos_siguiente(Ck.carretera[x0:x0+lramp])
        if x <= x0+lramp && rand() < pin
          Ck.cuenta += 1
          v_nueva = (rand() <= p ? vmax[1] : vmax[2])
          tipo_nuevo = (v_nueva == vmax[1] ? 1 : 2)
          Cambiar_Vehiculo!( Ck.carretera[x], v_nueva, x, tipo_nuevo, 1, num)
        end
        v_nueva = 0; tipo_nuevo = 0
    end
    x = 0
end


# cannot go over the speed limit
# cannot go backward
function speed_limit_actions(mdp::GenerativeMergingMDP, s::AugScene)
    valid_acts = Int64[]
    sizehint!(valid_acts, 7)
    ego = get_by_id(s.scene, EGO_ID)
    acc = s.ego_info.acc
    for a in actions(mdp)
        next_v = ego.state.v + action_map(mdp, acc, a).a*mdp.dt
        if 0 <= next_v <= mdp.env.main_lane_vmax
            push!(valid_acts, a)
        end
    end
    return valid_acts
end

# jerk constraints
function acceleration_limit_actions(mdp::GenerativeMergingMDP, s::AugScene)
    acc = s.ego_info.acc
    valid_acts = Int64[]
    # if we were hard braking we must release the brake
    # if isapprox(acc, mdp.max_deceleration)
    #     for a in actions(mdp)
    #         if action_map(mdp, 0.0, a).a >= 0.0
    #             push!(valid_acts, a)
    #         end
    #     end
    #     return valid_acts
    # else
    for a in actions(mdp)
        next_acc = action_map(mdp, acc, a).a
        if abs(next_acc) <= mdp.comfortable_acceleration
            push!(valid_acts, a)
        end
    end
    return valid_acts
    # end
end

function collision_lookahead(mdp::GenerativeMergingMDP, s::AugScene)
    valid_acts = Int64[]
    action_list = LatLonAccel[LatLonAccel(0.,0.0) for veh in s.scene if veh.id != EGO_ID]
    for a in actions(mdp)
        sp = deepcopy(s.scene)
        acc = action_map(mdp, s.ego_info.acc, a)
        tick!(sp, mdp.env.roadway, [acc, action_list...], mdp.dt)
        if !icollision_checker(sp, EGO_ID)
            push!(valid_acts, a)
        end
    end
    return valid_acts
end

function idm_limit_actions(mdp::GenerativeMergingMDP, s::AugScene)
    scene = s.scene
    ego_ind = findfirst(EGO_ID, scene)
    ego = scene[ego_ind]
    acc = s.ego_info.acc
    fore_res = get_front_neighbor(mdp.env, s.scene, EGO_ID)
    if fore_res.ind == nothing
        return actions(mdp)
    end
    v_ego = scene[ego_ind].state.v
    v_oth = scene[fore_res.ind].state.v
    headway = fore_res.Δs
    track_longitudinal!(mdp.ego_idm, v_ego, v_oth, headway)
    des_acc = mdp.ego_idm.a
    valid_acts = Int64[]
    for a in actions(mdp)
        next_acc = action_map(mdp, acc, a)
        if next_acc.a <= des_acc
            push!(valid_acts, a)
        end
    end
    return valid_acts
end

function safe_actions(mdp::GenerativeMergingMDP, s::AugScene)
    speed = speed_limit_actions(mdp, s)
    # jerk = acceleration_limit_actions(mdp, s)
    collision_free = collision_lookahead(mdp, s)
    available_acts = intersect(speed, collision_free)
    # available_acts = intersect(speed, jerk, collision_free)
    # available_acts = intersect(jerk, speed, collision_free)
    # ego = get_by_id(s.scene, EGO_ID)
    # if get_lane(mdp.env.roadway, ego) == main_lane(mdp.env)
    #     idm = idm_limit_actions(mdp, s)
    #     available_acts = intersect(speed, jerk, idm, collision_free)
    # end
    # push!(available_acts, HARD_BRAKE)
    # if isempty(available_acts)
        # return Int64[HARD_BRAKE]
    # end
    sort!(available_acts)
    return available_acts
end

function safe_actions(mdp::MDP{S,A}, o::AbstractArray) where {S,A}
    s = POMDPs.convert_s(S, o, mdp)
    return safe_actions(mdp, s)
end

function best_action(acts::Vector{A}, val::AbstractArray{T}, problem::M) where {A, T <: Real, M <: Union{POMDP, MDP}}
    best_a = acts[1]
    best_ai = actionindex(problem, best_a)
    best_val = val[best_ai]
    for a in acts
        ai = actionindex(problem, a)
        if val[ai] > best_val
            best_val = val[ai]
            best_ai = ai
            best_a = a
        end
    end
    return best_a::A
end

function masked_linear_epsilon_greedy(max_steps::Int64, eps_fraction::Float64, eps_end::Float64)
    # define function that will be called to select an action in DQN
    # only supports MDP environments
    function action_masked_epsilon_greedy(policy::AbstractNNPolicy, env::MDPEnvironment, obs, global_step::Int64, rng::AbstractRNG)
        eps = DeepQLearning.update_epsilon(global_step, eps_fraction, eps_end, max_steps)
        acts = safe_actions(policy.problem, obs) #XXX using pomdp global variable replace by safe_actions(mask, obs)
        val = actionvalues(policy, obs) #change this
        if rand(rng) < eps
            return (rand(rng, acts), eps)
        else
            return (best_action(acts, val, env.problem), eps)
        end
    end
    return action_masked_epsilon_greedy
end

function masked_evaluation()
    function masked_evaluation_policy(policy::AbstractNNPolicy, env::MDPEnvironment, n_eval::Int64, max_episode_length::Int64, verbose::Bool)
        avg_r = 0
        violations = 0
        avg_steps = 0
        for i=1:n_eval
            done = false
            r_tot = 0.0
            step = 0
            obs = reset!(env)
            DeepQLearning.resetstate!(policy)
            while !done && step <= max_episode_length
                acts = safe_actions(env.problem, obs)
                val = actionvalues(policy, obs)
                act = best_action(acts, val, env.problem)
                obs, rew, done, info = step!(env, act)
                r_tot += rew
                step += 1
            end
            # println("Episode reward: $r_tot")
            r_tot < 0. ? violations += 1 : nothing
            avg_r += r_tot
            avg_steps += step
        end
        # log_value(LOGGER, "eval_reward", avg_r/n_eval)
        # log_value(LOGGER, "collisions", violations/n_eval*100)
        # log_value(LOGGER, "avg_steps", avg_steps/n_eval)
        if verbose
            @printf("Evaluation ... Avg Reward %1.2f | Violations (%%) %2.2f | Avg Steps %3.2f", avg_r/n_eval, violations/n_eval*100, avg_steps/n_eval)
        end
        avg_r /= n_eval
        return  avg_r
    end
    return masked_evaluation_policy
end

struct MaskedPolicy{P<:AbstractNNPolicy} <: AbstractNNPolicy
    policy::P
end

function POMDPs.action(p::MaskedPolicy, s)
    acts = safe_actions(p.policy.problem, s)
    val = actionvalues(p.policy, s)
    act = best_action(acts, val, p.policy.problem)
    @assert in(act, acts)
    return act
end

struct MaskedRandomPolicy <: Policy
    problem::GenerativeMergingMDP
    rng::MersenneTwister
end

function POMDPs.action(p::MaskedRandomPolicy, s)
    acts = safe_actions(p.problem, s)
    return rand(p.rng, acts)
end


"""
This function creates the expected market data for each investor for each scenario using CEM.
"""
function create_expected_marketdata(investor_dir::String,
                                    market_names::Vector{Symbol},
                                    expected_portfolio::Vector{<: Project{<: BuildPhase}},
                                    zones::Vector{String},
                                    lines::Vector{ZonalLine},
                                    peak_load::Float64,
                                    rep_hour_weight::Vector{Float64},
                                    average_capital_cost_multiplier::Float64,
                                    scenario::Scenario,
                                    iteration_year::Int64,
                                    yearly_horizon::Int64,
                                    solver::JuMP.MOI.OptimizerWithAttributes,
                                    sys_results_dir::String,
                                    investor_name::String)

    system = create_cem_mkt_clr_problem(investor_dir,
                                        market_names,
                                        expected_portfolio,
                                        zones,
                                        lines,
                                        peak_load,
                                        rep_hour_weight,
                                        average_capital_cost_multiplier,
                                        scenario,
                                        iteration_year,
                                        yearly_horizon)

    capacity_price,
    energy_price,
    reserve_up_price,
    reserve_down_price,
    rec_price,
    capacity_factors,
    capacity_accepted_perc,
    new_options = cem(
                        system,
                        solver,
                        "C:/Users/manwar2/Documents/GitRepos/emt-tests/data/simulation_data/results/cem_results.txt"
                     )

    output_file = joinpath(investor_dir, "expected_market_data", "$(get_name(scenario))_year_$(iteration_year).jld2")

    FileIO.save(output_file,
                     "capacity_price", capacity_price,
                     "energy_price", energy_price,
                     "reserve_up_price", reserve_up_price,
                     "reserve_down_price", reserve_down_price,
                     "rec_price", rec_price,
                     "capacity_factors", capacity_factors,
                     "capacity_accepted_perc", capacity_accepted_perc,
                     "new_options", new_options
        )

    sys_results_file = joinpath(sys_results_dir, investor_name, "expected_market_data", "$(get_name(scenario))_year_$(iteration_year).jld2")

    FileIO.save(sys_results_file,
                     "capacity_price", capacity_price,
                     "energy_price", energy_price,
                     "reserve_up_price", reserve_up_price,
                     "reserve_down_price", reserve_down_price,
                     "rec_price", rec_price,
                     "capacity_factors", capacity_factors,
                     "capacity_accepted_perc", capacity_accepted_perc,
                     "new_options", new_options
        )

    return
end

"""
This function updates the maximum new options which can be built by the investor based on CEM predictions.
"""
function update_max_new_options!(max_new_options::Dict{String, Int64},
                                scenario_new_options::Dict{String, Int64},
                                option_projects::Vector{Project})
    for project in option_projects
        type = get_type(get_tech(project))
        zone = get_zone(get_tech(project))
        if scenario_new_options["option_$(type)_$(zone)"] > max_new_options[get_name(project)]
            max_new_options[get_name(project)] = scenario_new_options["option_$(type)_$(zone)"]
        end
    end
    return max_new_options
end


baremodule DZBase

export +, ⊕, -, ⊖, *, ⊗, /, ⊘, %, ==, !=, <, >, <=, >=,
    Array, Vector, Matrix, length

##################################################### UNARY ARITHMETIC OPERATORS

+(x::Int8   ) = x
+(x::Int16  ) = x
+(x::Int32  ) = x
+(x::Int64  ) = x
+(x::Int128 ) = x
+(x::UInt8  ) = x
+(x::UInt16 ) = x
+(x::UInt32 ) = x
+(x::UInt64 ) = x
+(x::UInt128) = x
+(x::Float16) = x
+(x::Float32) = x
+(x::Float64) = x

-(x::Int8   ) = Core.Intrinsics.neg_int(x)
-(x::Int16  ) = Core.Intrinsics.neg_int(x)
-(x::Int32  ) = Core.Intrinsics.neg_int(x)
-(x::Int64  ) = Core.Intrinsics.neg_int(x)
-(x::Int128 ) = Core.Intrinsics.neg_int(x)
-(x::UInt8  ) = Core.Intrinsics.neg_int(x)
-(x::UInt16 ) = Core.Intrinsics.neg_int(x)
-(x::UInt32 ) = Core.Intrinsics.neg_int(x)
-(x::UInt64 ) = Core.Intrinsics.neg_int(x)
-(x::UInt128) = Core.Intrinsics.neg_int(x)

-(x::Float16) = Core.Intrinsics.llvmcall("
    %res = fneg fast half %0
    ret half %res
", Float16, Tuple{Float16}, x)

-(x::Float32) = Core.Intrinsics.llvmcall("
    %res = fneg fast float %0
    ret float %res
", Float32, Tuple{Float32}, x)

-(x::Float64) = Core.Intrinsics.llvmcall("
    %res = fneg fast double %0
    ret double %res
", Float64, Tuple{Float64}, x)

⊖(x::Float16) = Core.Intrinsics.llvmcall("
    %res = fneg nnan ninf nsz half %0
    ret half %res
", Float16, Tuple{Float16}, x)

⊖(x::Float32) = Core.Intrinsics.llvmcall("
    %res = fneg nnan ninf nsz float %0
    ret float %res
", Float32, Tuple{Float32}, x)

⊖(x::Float64) = Core.Intrinsics.llvmcall("
    %res = fneg nnan ninf nsz double %0
    ret double %res
", Float64, Tuple{Float64}, x)

*(x::Int8   ) = x
*(x::Int16  ) = x
*(x::Int32  ) = x
*(x::Int64  ) = x
*(x::Int128 ) = x
*(x::UInt8  ) = x
*(x::UInt16 ) = x
*(x::UInt32 ) = x
*(x::UInt64 ) = x
*(x::UInt128) = x
*(x::Float16) = x
*(x::Float32) = x
*(x::Float64) = x

#################################################### BINARY ARITHMETIC OPERATORS

+(x::Int8   , y::Int8   ) = Core.Intrinsics.add_int(x, y)
+(x::Int16  , y::Int16  ) = Core.Intrinsics.add_int(x, y)
+(x::Int32  , y::Int32  ) = Core.Intrinsics.add_int(x, y)
+(x::Int64  , y::Int64  ) = Core.Intrinsics.add_int(x, y)
+(x::Int128 , y::Int128 ) = Core.Intrinsics.add_int(x, y)
+(x::UInt8  , y::UInt8  ) = Core.Intrinsics.add_int(x, y)
+(x::UInt16 , y::UInt16 ) = Core.Intrinsics.add_int(x, y)
+(x::UInt32 , y::UInt32 ) = Core.Intrinsics.add_int(x, y)
+(x::UInt64 , y::UInt64 ) = Core.Intrinsics.add_int(x, y)
+(x::UInt128, y::UInt128) = Core.Intrinsics.add_int(x, y)

+(x::Float16, y::Float16) = Core.Intrinsics.llvmcall("
    %res = fadd fast half %0, %1
    ret half %res
", Float16, Tuple{Float16, Float16}, x, y)

+(x::Float32, y::Float32) = Core.Intrinsics.llvmcall("
    %res = fadd fast float %0, %1
    ret float %res
", Float32, Tuple{Float32, Float32}, x, y)

+(x::Float64, y::Float64) = Core.Intrinsics.llvmcall("
    %res = fadd fast double %0, %1
    ret double %res
", Float64, Tuple{Float64, Float64}, x, y)

⊕(x::Float16, y::Float16) = Core.Intrinsics.llvmcall("
    %res = fadd nnan ninf nsz half %0, %1
    ret half %res
", Float16, Tuple{Float16, Float16}, x, y)

⊕(x::Float32, y::Float32) = Core.Intrinsics.llvmcall("
    %res = fadd nnan ninf nsz float %0, %1
    ret float %res
", Float32, Tuple{Float32, Float32}, x, y)

⊕(x::Float64, y::Float64) = Core.Intrinsics.llvmcall("
    %res = fadd nnan ninf nsz double %0, %1
    ret double %res
", Float64, Tuple{Float64, Float64}, x, y)

-(x::Int8   , y::Int8   ) = Core.Intrinsics.sub_int(x, y)
-(x::Int16  , y::Int16  ) = Core.Intrinsics.sub_int(x, y)
-(x::Int32  , y::Int32  ) = Core.Intrinsics.sub_int(x, y)
-(x::Int64  , y::Int64  ) = Core.Intrinsics.sub_int(x, y)
-(x::Int128 , y::Int128 ) = Core.Intrinsics.sub_int(x, y)
-(x::UInt8  , y::UInt8  ) = Core.Intrinsics.sub_int(x, y)
-(x::UInt16 , y::UInt16 ) = Core.Intrinsics.sub_int(x, y)
-(x::UInt32 , y::UInt32 ) = Core.Intrinsics.sub_int(x, y)
-(x::UInt64 , y::UInt64 ) = Core.Intrinsics.sub_int(x, y)
-(x::UInt128, y::UInt128) = Core.Intrinsics.sub_int(x, y)

-(x::Float16, y::Float16) = Core.Intrinsics.llvmcall("
    %res = fsub fast half %0, %1
    ret half %res
", Float16, Tuple{Float16, Float16}, x, y)

-(x::Float32, y::Float32) = Core.Intrinsics.llvmcall("
    %res = fsub fast float %0, %1
    ret float %res
", Float32, Tuple{Float32, Float32}, x, y)

-(x::Float64, y::Float64) = Core.Intrinsics.llvmcall("
    %res = fsub fast double %0, %1
    ret double %res
", Float64, Tuple{Float64, Float64}, x, y)

⊖(x::Float16, y::Float16) = Core.Intrinsics.llvmcall("
    %res = fsub nnan ninf nsz half %0, %1
    ret half %res
", Float16, Tuple{Float16, Float16}, x, y)

⊖(x::Float32, y::Float32) = Core.Intrinsics.llvmcall("
    %res = fsub nnan ninf nsz float %0, %1
    ret float %res
", Float32, Tuple{Float32, Float32}, x, y)

⊖(x::Float64, y::Float64) = Core.Intrinsics.llvmcall("
    %res = fsub nnan ninf nsz double %0, %1
    ret double %res
", Float64, Tuple{Float64, Float64}, x, y)

*(x::Int8   , y::Int8   ) = Core.Intrinsics.mul_int(x, y)
*(x::Int16  , y::Int16  ) = Core.Intrinsics.mul_int(x, y)
*(x::Int32  , y::Int32  ) = Core.Intrinsics.mul_int(x, y)
*(x::Int64  , y::Int64  ) = Core.Intrinsics.mul_int(x, y)
*(x::Int128 , y::Int128 ) = Core.Intrinsics.mul_int(x, y)
*(x::UInt8  , y::UInt8  ) = Core.Intrinsics.mul_int(x, y)
*(x::UInt16 , y::UInt16 ) = Core.Intrinsics.mul_int(x, y)
*(x::UInt32 , y::UInt32 ) = Core.Intrinsics.mul_int(x, y)
*(x::UInt64 , y::UInt64 ) = Core.Intrinsics.mul_int(x, y)
*(x::UInt128, y::UInt128) = Core.Intrinsics.mul_int(x, y)

*(x::Float16, y::Float16) = Core.Intrinsics.llvmcall("
    %res = fmul fast half %0, %1
    ret half %res
", Float16, Tuple{Float16, Float16}, x, y)

*(x::Float32, y::Float32) = Core.Intrinsics.llvmcall("
    %res = fmul fast float %0, %1
    ret float %res
", Float32, Tuple{Float32, Float32}, x, y)

*(x::Float64, y::Float64) = Core.Intrinsics.llvmcall("
    %res = fmul fast double %0, %1
    ret double %res
", Float64, Tuple{Float64, Float64}, x, y)

⊗(x::Float16, y::Float16) = Core.Intrinsics.llvmcall("
    %res = fmul nnan ninf nsz half %0, %1
    ret half %res
", Float16, Tuple{Float16, Float16}, x, y)

⊗(x::Float32, y::Float32) = Core.Intrinsics.llvmcall("
    %res = fmul nnan ninf nsz float %0, %1
    ret float %res
", Float32, Tuple{Float32, Float32}, x, y)

⊗(x::Float64, y::Float64) = Core.Intrinsics.llvmcall("
    %res = fmul nnan ninf nsz double %0, %1
    ret double %res
", Float64, Tuple{Float64, Float64}, x, y)

/(x::Int8   , y::Int8   ) = Core.Intrinsics.sdiv_int(x, y)
/(x::Int16  , y::Int16  ) = Core.Intrinsics.sdiv_int(x, y)
/(x::Int32  , y::Int32  ) = Core.Intrinsics.sdiv_int(x, y)
/(x::Int64  , y::Int64  ) = Core.Intrinsics.sdiv_int(x, y)
/(x::Int128 , y::Int128 ) = Core.Intrinsics.sdiv_int(x, y)
/(x::UInt8  , y::UInt8  ) = Core.Intrinsics.udiv_int(x, y)
/(x::UInt16 , y::UInt16 ) = Core.Intrinsics.udiv_int(x, y)
/(x::UInt32 , y::UInt32 ) = Core.Intrinsics.udiv_int(x, y)
/(x::UInt64 , y::UInt64 ) = Core.Intrinsics.udiv_int(x, y)
/(x::UInt128, y::UInt128) = Core.Intrinsics.udiv_int(x, y)

/(x::Float16, y::Float16) = Core.Intrinsics.llvmcall("
    %res = fdiv fast half %0, %1
    ret half %res
", Float16, Tuple{Float16, Float16}, x, y)

/(x::Float32, y::Float32) = Core.Intrinsics.llvmcall("
    %res = fdiv fast float %0, %1
    ret float %res
", Float32, Tuple{Float32, Float32}, x, y)

/(x::Float64, y::Float64) = Core.Intrinsics.llvmcall("
    %res = fdiv fast double %0, %1
    ret double %res
", Float64, Tuple{Float64, Float64}, x, y)

⊘(x::Float16, y::Float16) = Core.Intrinsics.llvmcall("
    %res = fdiv nnan ninf nsz half %0, %1
    ret half %res
", Float16, Tuple{Float16, Float16}, x, y)

⊘(x::Float32, y::Float32) = Core.Intrinsics.llvmcall("
    %res = fdiv nnan ninf nsz float %0, %1
    ret float %res
", Float32, Tuple{Float32, Float32}, x, y)

⊘(x::Float64, y::Float64) = Core.Intrinsics.llvmcall("
    %res = fdiv nnan ninf nsz double %0, %1
    ret double %res
", Float64, Tuple{Float64, Float64}, x, y)

%(x::Int8   , y::Int8   ) = Core.Intrinsics.srem_int(x, y)
%(x::Int16  , y::Int16  ) = Core.Intrinsics.srem_int(x, y)
%(x::Int32  , y::Int32  ) = Core.Intrinsics.srem_int(x, y)
%(x::Int64  , y::Int64  ) = Core.Intrinsics.srem_int(x, y)
%(x::Int128 , y::Int128 ) = Core.Intrinsics.srem_int(x, y)
%(x::UInt8  , y::UInt8  ) = Core.Intrinsics.urem_int(x, y)
%(x::UInt16 , y::UInt16 ) = Core.Intrinsics.urem_int(x, y)
%(x::UInt32 , y::UInt32 ) = Core.Intrinsics.urem_int(x, y)
%(x::UInt64 , y::UInt64 ) = Core.Intrinsics.urem_int(x, y)
%(x::UInt128, y::UInt128) = Core.Intrinsics.urem_int(x, y)

########################################################### COMPARISON OPERATORS

==(x::Int8   , y::Int8   ) = Core.Intrinsics.eq_int(x, y)
==(x::Int16  , y::Int16  ) = Core.Intrinsics.eq_int(x, y)
==(x::Int32  , y::Int32  ) = Core.Intrinsics.eq_int(x, y)
==(x::Int64  , y::Int64  ) = Core.Intrinsics.eq_int(x, y)
==(x::Int128 , y::Int128 ) = Core.Intrinsics.eq_int(x, y)
==(x::UInt8  , y::UInt8  ) = Core.Intrinsics.eq_int(x, y)
==(x::UInt16 , y::UInt16 ) = Core.Intrinsics.eq_int(x, y)
==(x::UInt32 , y::UInt32 ) = Core.Intrinsics.eq_int(x, y)
==(x::UInt64 , y::UInt64 ) = Core.Intrinsics.eq_int(x, y)
==(x::UInt128, y::UInt128) = Core.Intrinsics.eq_int(x, y)

!=(x::Int8   , y::Int8   ) = Core.Intrinsics.ne_int(x, y)
!=(x::Int16  , y::Int16  ) = Core.Intrinsics.ne_int(x, y)
!=(x::Int32  , y::Int32  ) = Core.Intrinsics.ne_int(x, y)
!=(x::Int64  , y::Int64  ) = Core.Intrinsics.ne_int(x, y)
!=(x::Int128 , y::Int128 ) = Core.Intrinsics.ne_int(x, y)
!=(x::UInt8  , y::UInt8  ) = Core.Intrinsics.ne_int(x, y)
!=(x::UInt16 , y::UInt16 ) = Core.Intrinsics.ne_int(x, y)
!=(x::UInt32 , y::UInt32 ) = Core.Intrinsics.ne_int(x, y)
!=(x::UInt64 , y::UInt64 ) = Core.Intrinsics.ne_int(x, y)
!=(x::UInt128, y::UInt128) = Core.Intrinsics.ne_int(x, y)

<(x::Int8   , y::Int8   ) = Core.Intrinsics.slt_int(x, y)
<(x::Int16  , y::Int16  ) = Core.Intrinsics.slt_int(x, y)
<(x::Int32  , y::Int32  ) = Core.Intrinsics.slt_int(x, y)
<(x::Int64  , y::Int64  ) = Core.Intrinsics.slt_int(x, y)
<(x::Int128 , y::Int128 ) = Core.Intrinsics.slt_int(x, y)
<(x::UInt8  , y::UInt8  ) = Core.Intrinsics.ult_int(x, y)
<(x::UInt16 , y::UInt16 ) = Core.Intrinsics.ult_int(x, y)
<(x::UInt32 , y::UInt32 ) = Core.Intrinsics.ult_int(x, y)
<(x::UInt64 , y::UInt64 ) = Core.Intrinsics.ult_int(x, y)
<(x::UInt128, y::UInt128) = Core.Intrinsics.ult_int(x, y)

>(x::Int8   , y::Int8   ) = Core.Intrinsics.slt_int(y, x)
>(x::Int16  , y::Int16  ) = Core.Intrinsics.slt_int(y, x)
>(x::Int32  , y::Int32  ) = Core.Intrinsics.slt_int(y, x)
>(x::Int64  , y::Int64  ) = Core.Intrinsics.slt_int(y, x)
>(x::Int128 , y::Int128 ) = Core.Intrinsics.slt_int(y, x)
>(x::UInt8  , y::UInt8  ) = Core.Intrinsics.ult_int(y, x)
>(x::UInt16 , y::UInt16 ) = Core.Intrinsics.ult_int(y, x)
>(x::UInt32 , y::UInt32 ) = Core.Intrinsics.ult_int(y, x)
>(x::UInt64 , y::UInt64 ) = Core.Intrinsics.ult_int(y, x)
>(x::UInt128, y::UInt128) = Core.Intrinsics.ult_int(y, x)

<=(x::Int8   , y::Int8   ) = Core.Intrinsics.sle_int(x, y)
<=(x::Int16  , y::Int16  ) = Core.Intrinsics.sle_int(x, y)
<=(x::Int32  , y::Int32  ) = Core.Intrinsics.sle_int(x, y)
<=(x::Int64  , y::Int64  ) = Core.Intrinsics.sle_int(x, y)
<=(x::Int128 , y::Int128 ) = Core.Intrinsics.sle_int(x, y)
<=(x::UInt8  , y::UInt8  ) = Core.Intrinsics.ule_int(x, y)
<=(x::UInt16 , y::UInt16 ) = Core.Intrinsics.ule_int(x, y)
<=(x::UInt32 , y::UInt32 ) = Core.Intrinsics.ule_int(x, y)
<=(x::UInt64 , y::UInt64 ) = Core.Intrinsics.ule_int(x, y)
<=(x::UInt128, y::UInt128) = Core.Intrinsics.ule_int(x, y)

>=(x::Int8   , y::Int8   ) = Core.Intrinsics.sle_int(y, x)
>=(x::Int16  , y::Int16  ) = Core.Intrinsics.sle_int(y, x)
>=(x::Int32  , y::Int32  ) = Core.Intrinsics.sle_int(y, x)
>=(x::Int64  , y::Int64  ) = Core.Intrinsics.sle_int(y, x)
>=(x::Int128 , y::Int128 ) = Core.Intrinsics.sle_int(y, x)
>=(x::UInt8  , y::UInt8  ) = Core.Intrinsics.ule_int(y, x)
>=(x::UInt16 , y::UInt16 ) = Core.Intrinsics.ule_int(y, x)
>=(x::UInt32 , y::UInt32 ) = Core.Intrinsics.ule_int(y, x)
>=(x::UInt64 , y::UInt64 ) = Core.Intrinsics.ule_int(y, x)
>=(x::UInt128, y::UInt128) = Core.Intrinsics.ule_int(y, x)

################################################################################

const Array = Core.Array
const Vector{T} = Core.Array{T, 1}
const Matrix{T} = Core.Array{T, 2}

length(a::Array{T, N}) where {T, N} = Core.Intrinsics.arraylen(a)

################################################################################

function two_sum(a::T, b::T) where {T}
    sum = a ⊕ b
    b_prime = sum ⊖ a
    a_prime = sum ⊖ b_prime
    a_err = a ⊖ a_prime
    b_err = b ⊖ b_prime
    return (sum, a_err ⊕ b_err)
end

function two_diff(a::T, b::T) where {T}
    diff = a ⊖ b
    b_prime = a ⊖ diff
    a_prime = diff ⊕ b_prime
    a_err = a ⊖ a_prime
    b_err = b ⊖ b_prime
    return (sum, a_err ⊖ b_err)
end

function two_prod(a::T, b::T) where {T}
    p = a * b
    p, Core.Intrinsics.fma_float(a, b, -p)
end

end # baremodule DZBase


using Test
using EmergencyBrakingSystem
using AutomotiveDrivingModels
using AutoViz
using AutomotiveSensors
using AutomotivePOMDPs
using Parameters
using StaticArrays
using Reel
using Random

@testset "EmergencyBrakingSystem.jl" begin

Random.seed!(1)

ped_v = 5.0 / 3.6
ped_theta = π/2
ped_right_side = true
hitpoint = 0
ped_t_collision = 1.0
ego_y = 0.0
ped_x = 100.0
ped_y = 0.0

obstacle_right_offset = -VehicleDef().width/2 - 1.0
obstacle_1 = ConvexPolygon([VecE2(44, obstacle_right_offset), VecE2(44, obstacle_right_offset-2), VecE2(48.5, obstacle_right_offset-2), VecE2(48.5, obstacle_right_offset)],4)
obstacles = [obstacle_1]

# unavoidable situation
ego_v = 50.0 / 3.6
ped_v = 0.0
ego_x = ped_x - ego_v * ped_t_collision - VehicleDef().length/2;
(scenes, timestep, env, ego_vehicle, sensor, sensor_observations, risk, ttc, collision_rate, brake_request, prediction, collision) = EmergencyBrakingSystem.evaluate_scenario(ego_x, ego_y, ego_v, ped_x, ped_y, ped_v, ped_theta, obstacles);
@test collision[end] == true

# avoidable situation
ego_v = 30.0 / 3.6
ego_x = ped_x - ego_v * ped_t_collision - VehicleDef().length/2;
(scenes, timestep, env, ego_vehicle, sensor, sensor_observations, risk, ttc, collision_rate, brake_request, prediction, collision) = EmergencyBrakingSystem.evaluate_scenario(ego_x, ego_y, ego_v, ped_x, ped_y, ped_v, ped_theta, obstacles);
@test collision[end] == false

end


# utility functions needed for sandshark

using DelimitedFiles
using Interpolations
using ProgressMeter

include(joinpath(@__DIR__, "CommonUtils.jl"))

datadir = joinpath(ENV["HOME"],"data","sandshark","full_wombat_2018_07_09","extracted")
matcheddir = joinpath(datadir, "matchedfilter", "particles")
beamdir = joinpath(datadir, "beamformer", "particles")

function loaddircsvs(datadir)
  # https://docs.julialang.org/en/v0.6.1/stdlib/file/#Base.Filesystem.walkdir
  datadict = Dict{Int, Array{Float64}}()
  for (root, dirs, files) in walkdir(datadir)
    # println("Files in $root")
    for file in files
      # println(joinpath(root, file)) # path to files
      data = readdlm(joinpath(root, file),',')
      datadict[parse(Int,split(file,'.')[1])/1] = data
    end
  end
  return datadict
end

rangedata = loaddircsvs(matcheddir)
azidata = loaddircsvs(beamdir)
timestamps = intersect(sort(collect(keys(rangedata))), sort(collect(keys(azidata))))

# NAV data
navdata = Dict{Int, Vector{Float64}}()
navfile = readdlm(joinpath(datadir, "nav_data.csv"))
for row in navfile
    s = split(row, ",")
    id = round(Int, 1e9*parse(Float64, s[1]))
    # round(Int, 1000 * parse(s[1])) = 1531153292381
    navdata[id] = parse.(Float64,s)
end
navkeys = sort(collect(keys(navdata)))
# NAV colums are X,Y = 7,8
# lat,long = 9,10
# time,pitch,roll,heading,speed,[Something], internal_x,internal_y,internal_lat,internal_long, yaw_rad

# LBL data - note the timestamps need to be exported as float in future.
lbldata = Dict{Int,  Vector{Float64}}()
lblfile = readdlm(joinpath(datadir, "lbl.csv"))
for row in lblfile
    s = split(row, ",")
    id = round(Int, 1e9*parse(Float64, s[1]))
    if s[2] != "NaN"
        lbldata[id] = parse.(Float64, s)
    end
end
lblkeys = sort(collect(keys(lbldata)))

# function heading2yaw(heading)
# heading = map( x->deg2rad(navdata[x][4]), navkeys )
# wrapheading = TU.wrapRad.(heading)
# end

# GET Y = north,  X = East,  Heading along +Y clockwise [0,360)]
# east = Float64[]
# north = Float64[]
# heading = Float64[]

# WANT X = North,  Y = West,  Yaw is right and rule from +X (0) towards +Y pi/2, [-pi,pi)
# so the drawPoses picture will look flipped from Nicks picture
# remember theta = atan2(y,x)    # this is right hand rule
X = Float64[]
Y = Float64[]
yaw = Float64[]
for id in navkeys
  # push!(east, getindex(navdata[id],7)) # x-column csv
  # push!(north, getindex(navdata[id],8)) # y-column csv
  # push!(heading, getindex(navdata[id],4))
  # push!(yaw, TU.wrapRad(-deg2rad(getindex(navdata[id],4))))

  push!(X, 0.7*getindex(navdata[id],7) ) # 8
  push!(Y, 0.7*getindex(navdata[id],8) ) # 7
  push!(yaw, TU.wrapRad(pi/2-deg2rad(getindex(navdata[id],4))) )  # rotation about +Z
end

lblX = Float64[]
lblY = Float64[]
for id in lblkeys
    push!(lblX, getindex(lbldata[id],2) )
    push!(lblY, getindex(lbldata[id],3) )
end

# Build interpolators for x, y, yaw
interp_x = LinearInterpolation(navkeys, X)
interp_y = LinearInterpolation(navkeys, Y)
interp_yaw = LinearInterpolation(navkeys, yaw)

function poorMansDeconv_BF(XX = collect(range(-pi,pi,length=1000)),
                           YY = ppbrDict[epochs[1]].bearing(XX)      )
  #
  YY1 = YY.^4
  YY1 ./= maximum(YY1)
  YY2 = YY1.^4
  YY2 ./= maximum(YY2)
  # Gadfly.plot(x=XX, y=YY2, Geom.line)
  bss = AliasingScalarSampler(XX, YY2)
  pts = reshape(rand(bss, 200), 1, :)
  pc = manikde!(pts, Sphere1)
  # plotKDECircular(pc)

  return pc
end

# Step: Selecting a subset for processing and build up a cache of the factors.

function doEpochs(timestamps, rangedata, azidata, interp_x, interp_y, interp_yaw, odonoise; TSTART=356, TEND=1200, SNRfloor::Float64=0.6, STRIDE::Int=4)
  #
  ## Caching factors
  ppbrDict = Dict{Int, Pose2Point2BearingRange}()
  ppbDict = Dict{Int, Pose2Point2Bearing}()
  pprDict = Dict{Int, Pose2Point2Range}()
  odoDict = Dict{Int, Pose2Pose2}()
  NAV = Dict{Int, Vector{Float64}}()

  XX = collect(range(-pi,pi,length=1000))

  epochs = timestamps[TSTART:STRIDE:TEND]
  lastepoch = 0
  @showprogress "preparing data" for ep in epochs
    # @show ep
    if lastepoch != 0
      # @show interp_yaw(ep)
      deltaAng = interp_yaw(ep) - interp_yaw(lastepoch)

      wXi = TU.SE2([interp_x(lastepoch);interp_y(lastepoch);interp_yaw(lastepoch)])
      wXj = TU.SE2([interp_x(ep);interp_y(ep);interp_yaw(ep)])
      iDXj = se2vee(wXi\wXj)
      NAV[ep] = iDXj
      # NAV[ep][1:2] .*= 0.7
      # println("$(iDXj[1]), $(iDXj[2]), $(iDXj[3])")

      odoDict[ep] = Pose2Pose2(MvNormal(NAV[ep], odonoise) )
    end
    rangepts = rangedata[ep][:]
    rangeprob = kde!(rangepts)

    # azipts = azidata[ep][:,1]
    azipts = collect(azidata[ep][:,1:1]')
    aziptsw = TU.wrapRad.(azipts)

    # direct
    aziprobl = kde!(azipts)
    # npts = rand(aziprobl, 200)
    # aziprob = manikde!(npts, Sphere1)

    # with deconv
    aziprob = poorMansDeconv_BF(XX, aziprobl(XX))

    # alternative range probability
    rawmf = readdlm("/home/dehann/data/sandshark/full_wombat_2018_07_09/extracted/matchedfilter/raw/$(ep).csv",',')
    dvmf = exp.(rawmf[:,2])
    dvmf .= dvmf.^4
    dvmf ./= cumsum(dvmf)[end]
    dvmf .= dvmf.^2
    dvmf ./= cumsum(dvmf)[end]
    range_bss = AliasingScalarSampler(rawmf[:,1], dvmf, SNRfloor=SNRfloor) # exp.(rawmf[:,2])

    # prep the factor functions
    ppbrDict[ep] = Pose2Point2BearingRange(aziprob, range_bss) # rangeprob
    ppbDict[ep] = Pose2Point2Bearing(aziprob) # rangeprob
    pprDict[ep] = Pose2Point2Range(range_bss) # rangeprob
    lastepoch = ep
  end
  return epochs, odoDict, ppbrDict, ppbDict, pprDict, NAV
end

function initializeAUV_noprior(dfg::AbstractDFG,
                               dashboard::Dict;
                               stride_range::Int=4,
                               magStdDeg::Float64=5.0,
                               stride_solve::Int=30 )
  #
  addVariable!(dfg, :x0, Pose2)

  # Pinger location is [17; 1.8]
  addVariable!(dfg, :l1, Point2, solvable=0)
  beaconprior = PriorPoint2( MvNormal([17; 1.8], Matrix(Diagonal([0.1; 0.1].^2)) ) )
  addFactor!(dfg, [:l1], beaconprior, autoinit=true, solvable=0)

  addVariable!(dfg, :drt_0, Pose2, solvable=0)
  drec = MutablePose2Pose2Gaussian(MvNormal(zeros(3), Matrix{Float64}(LinearAlgebra.I, 3,3)))
  addFactor!(dfg, [:x0; :drt_0], drec, solvable=0, autoinit=false)
  # NOTE have not added to dashboard[:drtMpp] yet

  # reference odo solution
  addVariable!(dfg, :drt_ref, Pose2, solvable=0)
  drec = MutablePose2Pose2Gaussian(MvNormal(zeros(3), Matrix{Float64}(LinearAlgebra.I, 3,3)))
  addFactor!(dfg, [:x0; :drt_ref], drec, solvable=0, autoinit=false)
  dashboard[:drtOdoRef] = (:x0, :drt_ref, drec)

  # store current real time tether factor
  dashboard[:drtMpp] = Dict{Symbol,MutablePose2Pose2Gaussian}(:x0 => drec)
  # dashboard[:drtCurrent] = (:x0, :drt_0)
  # standard odo process noise levels
  # TODO, debug and refine cont2disc
  dashboard[:Qc_odo] = Diagonal([0.01;0.01;0.001].^2) |> Matrix

  dashboard[:odoTime] = unix2datetime(0)
  dashboard[:poseRate] = Second(1)
  dashboard[:lastPose] = :x0

  # dashboard[:solvables] = Channel{Vector{Symbol}}(100)

  # dashboard[:loopSolver] = true

  # dashboard[:SOLVESTRIDE] = stride_solve
  # add a range measurement every xth pose
  dashboard[:poseStride] = 0
  # dashboard[:canTakePoses] = HSMReady
  # dashboard[:solveInProgress] = SSMReady

  # dashboard[:poseSolveToken] = Channel{Symbol}(3)

  dashboard[:RANGESTRIDE] = stride_range # add a range measurement every xth pose
  dashboard[:rangesBuffer] = CircularBuffer{Tuple{DateTime, Array{Float64,2}, Vector{Bool}}}(dashboard[:RANGESTRIDE]+4)
  dashboard[:rangeCount] = 0

  dashboard[:SNRfloor] = 0.6

  dashboard[:realTimeSlack] = Millisecond(0)

  dashboard[:magBuffer] = CircularBuffer{Tuple{DateTime, Float64, Vector{Bool}}}(20)
  dashboard[:magNoise] = deg2rad(magStdDeg)

  dashboard[:lblBuffer] = CircularBuffer{Tuple{DateTime, Array{Float64,1}, Vector{Bool}}}(dashboard[:RANGESTRIDE]+4)

  dashboard[:odoCov] = Matrix{Float64}(Diagonal([0.5;0.3;0.2].^2))

  dashboard[:savePlotting] = false

  dashboard[:solveSettings] = ManageSolveSettings(
    stride_solve,
    true,
    Channel{Vector{Symbol}}(100),
    SSMReady,
    Channel{Symbol}(3),
    HSMReady,
    (:x0, :drt_0)
  )

  nothing
end

## NOTE, moved to RoME/src/Slam.jl
# # Requires
# #  dashboard: SOLVESTRIDE, loopSolver, solvables, loopSolver, solveInProgress, poseSolveToken, canTakePoses, drtCurrent
# #
# function manageSolveTree!(dfg::AbstractDFG,
#                           dashboard::Dict;
#                           dbg::Bool=false,
#                           timinglog=Base.stdout,
#                           limitfixeddown::Bool=true  )
#
#   @info "logpath=$(getLogPath(dfg))"
#   getSolverParams(dfg).drawtree = true
#   getSolverParams(dfg).qfl = dashboard[:SOLVESTRIDE]
#   getSolverParams(dfg).isfixedlag = true
#   getSolverParams(dfg).limitfixeddown = limitfixeddown
#   getSolverParams(dfg).maxincidence = 1000
#
#   # allow async process
#   # getSolverParams(dfg).async = true
#
#   # prep with empty tree
#   tree = emptyBayesTree()
#
#   # needs to run asynchronously
#   ST = @async begin
#     while @show length(ls(dfg, :x0, solvable=1)) == 0
#       "waiting for prior on x0" |> println
#       sleep(1)
#     end
#     solvecycle = 0
#     # keep solving
#     while dashboard[:loopSolver]
#       t0 = time_ns()
#       solvecycle += 1
#       # add any newly solvables (atomic)
#       while !isready(dashboard[:solvables]) && dashboard[:loopSolver]
#         sleep(0.2)
#       end
#       dt_wait = (time_ns()-t0)/1e9
#
#       #add any new solvables
#       while isready(dashboard[:solvables]) && dashboard[:loopSolver]
#         dashboard[:solveInProgress] = SSMConsumingSolvables
#         @show tosolv = take!(dashboard[:solvables])
#         for sy in tosolv
#           # setSolvable!(dfg, sy, 1) # see DFG #221
#           # TODO temporary workaround
#           getfnc = occursin(r"f", string(sy)) ? getFactor : getVariable
#           getfnc(dfg, sy).solvable = 1
#         end
#       end
#       dt_solvable = (time_ns()-t0)/1e9
#
#       @info "Ensure all new variables initialized"
#       ensureAllInitialized!(dfg)
#       dt_init = (time_ns()-t0)/1e9
#       dt_disengage = 0.0
#       dt_save1 = 0.0
#       dt_solve = 0.0
#
#       dashboard[:solveInProgress] = SSMReady
#
#       # solve only every 10th pose
#       if 0 < length(dashboard[:poseSolveToken].data)
#       # if 10 <= dashboard[:poseStride]
#         @info "reduce problem size by disengaging older parts of factor graph"
#         setSolvableOldPoses!(dfg, youngest=getSolverParams(dfg).qfl+round(Int,dashboard[:SOLVESTRIDE]/2), oldest=100, solvable=0)
#         dt_disengage = (time_ns()-t0)/1e9
#
#         # set up state machine flags to allow overlapping or block
#         dashboard[:solveInProgress] = SSMSolving
#         # dashboard[:poseStride] = 0
#
#         # do the actual solve (with debug saving)
#         lasp = getLastPoses(dfg, filterLabel=r"x\d", number=1)[1]
#         !dbg ? nothing : saveDFG(dfg, joinpath(getLogPath(dfg), "fg_before_$(lasp)"))
#         dt_save1 = (time_ns()-t0)/1e9
#         # constrain solve with the latest pose at the top
#         # @show latestPose = intersect(getLastPoses(dfg, filterLabel=r"x\d", number=12), ls(dfg, r"x\d", solvable=1))[end]
#         tree, smt, hist = solveTree!(dfg, tree) # , variableConstraints=[latestPose;]
#         dt_solve = (time_ns()-t0)/1e9
#         !dbg ? nothing : saveDFG(dfg, joinpath(getLogPath(dfg), "fg_after_$(lasp)"))
#
#         # unblock LCMLog reader for next STRIDE segment
#         dashboard[:solveInProgress] = SSMReady
#         # de-escalate handler state machine
#         dashboard[:canTakePoses] = HSMHandling
#
#         # adjust latest RTT after solve, latest solved -- hard coded pose stride 10
#         lastList = sortDFG(ls(dfg, r"x\d+9\b|x9\b", solvable=1))
#         if 0 < length(lastList)
#           lastSolved = lastList[end]
#           dashboard[:drtCurrent] = (lastSolved, Symbol("drt_"*string(lastSolved)[2:end]))
#         end
#
#         # remove a token to allow progress to continue
#         gotToken = take!(dashboard[:poseSolveToken])
#         "end of solve cycle, token=$gotToken" |> println
#       else
#         "sleep a solve cycle" |> println
#         sleep(0.2)
#       end
#       dt_finish = (time_ns() - t0)/1e9
#
#       # store timing results
#       println(timinglog, "$solvecycle, $t0, $dt_wait, $dt_solvable, $dt_init, $dt_disengage, $dt_save1, $dt_solve, $(tree.buildTime), $dt_finish")
#     end
#   end
#   return ST
# end

#


################################################################################
#
#                       arf.jl: wrapper for arf
#
# This file is part of hecke.
#
# Copyright (c) 2015: Claus Fieker, Tommy Hofmann
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
# * Redistributions of source code must retain the above copyright notice, this
#   list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
#
#  Copyright (C) 2015 Tommy Hofmann
#
################################################################################
#
#  TODO:
#  - Sensible way to deal with rounding modes (at least documentation)
#
################################################################################

export arf, ArfField

export pos_inf, neg_inf, nan, isposinf, isneginf, isnan, isnormal, max, min

################################################################################
#
#  Types and memory management
#
################################################################################

const ArfFieldID = ObjectIdDict()

mutable struct ArfField <: Field
  prec::Int
  rndmode::Cint

  function ArfField(p::Int = 256, r::Cint = Cint(4))
    try
      return ArfFieldID[p,r]::ArfField
    catch
      z = new(p,r)
      ArfFieldID[p,r] = z
      return z
    end
  end
end

mutable struct arf
  exp::Int # fmpz
  size::UInt64 # mp_size_t
  d1::Int64 # mantissa_struct
  d2::Int64
  parent::ArfField

  function arf()
    z = new()
    ccall((:arf_init, :libarb), Void, (Ptr{arf}, ), &z)
    finalizer(z, _arf_clear_fn)
    return z
  end

  function arf(i::Int)
    z = new()
    ccall((:arf_init, :libarb), Void, (Ptr{arf}, ), &z)
    ccall((:arf_set_si, :libarb), Void, (Ptr{arf}, Int), &z, i)
    finalizer(z, _arf_clear_fn)
    return z
  end

  function arf(i::UInt)
    z = new()
    ccall((:arf_init, :libarb), Void, (Ptr{arf}, ), &z)
    ccall((:arf_set_ui, :libarb), Void, (Ptr{arf}, UInt), &z, i)
    finalizer(z, _arf_clear_fn)
    return z
  end

  function arf(i::fmpz)
    z = new()
    ccall((:arf_init, :libarb), Void, (Ptr{arf}, ), &z)
    ccall((:arf_set_fmpz, :libarb), Void, (Ptr{arf}, Ptr{fmpz}), &z, &i)
    finalizer(z, _arf_clear_fn)
    return z
  end

  function arf(x::BigFloat)
    z = new()
    ccall((:arf_init, :libarb), Void, (Ptr{arf}, ), &z)
    ccall((:arf_set_mpfr, :libarb), Ptr{arf}, (Ptr{arf}, Ptr{BigFloat}), &z, &x)
    finalizer(z, _arf_clear_fn)
    return z
  end

  function arf(x::Cdouble)
    z = new()
    ccall((:arf_init, :libarb), Void, (Ptr{arf}, ), &z)
    ccall((:arf_set_d, :libarb), Void, (Ptr{arf}, Cdouble), &z, x)
    finalizer(z, _arf_clear_fn)
    return z
  end

  function arf(x::mag)
    z = new()
    ccall((:arf_init, :libarb), Void, (Ptr{arf}, ), &z)
    ccall((:arf_set_mag, :libarb), Void, (Ptr{arf}, Ptr{mag}), &z, &x)
    finalizer(z, _arf_clear_fn)
    return z
  end

  function arf(x::arf, p::Int, r::Cint)
    z = new()
    ccall((:arf_init, :libarb), Void, (Ptr{arf}, ), &z)
    ccall((:arf_set_round, :libarb), Void,
                (Ptr{arf}, Ptr{arf}, Int, Cint), &z, &x, p, r)
    finalizer(z, _arf_clear_fn)
    return z
  end

  function arf(x::Int, p::Int, r::Cint)
    z = new()
    ccall((:arf_init, :libarb), Void, (Ptr{arf}, ), &z)
    ccall((:arf_set_round_si, :libarb), Void,
                  (Ptr{arf}, Int, Int, Cint), &z, x, p, r)
    finalizer(z, _arf_clear_fn)
    return z
  end

 function arf(x::UInt, p::Int, r::Cint)
    z = new()
    ccall((:arf_init, :libarb), Void, (Ptr{arf}, ), &z)
    ccall((:arf_set_round_ui, :libarb), Void,
                  (Ptr{arf}, UInt, Int, Cint), &z, x, p, r)
    finalizer(z, _arf_clear_fn)
    return z
  end

 function arf(x::fmpz, p::Int, r::Cint)
    z = new()
    ccall((:arf_init, :libarb), Void, (Ptr{arf}, ), &z)
    ccall((:arf_set_round_fmpz, :libarb), Void,
                  (Ptr{arf}, Ptr{fmpz}, Int, Cint), &z, &x, p, r)
    finalizer(z, _arf_clear_fn)
    return z
  end
end

function _arf_clear_fn(x::arf)
  ccall((:arf_clear, :libarb), Void, (Ptr{arf}, ), &x)
end

parent(x::arf) = x.parent

################################################################################
#
#  Parent object overloads
#
################################################################################

function (r::ArfField)()
  z = arf()
  z.parent = r
  return z
end

function (r::ArfField)(x::Int)
  z = arf(x, r.prec, r.rndmode)
  z.parent = r
  return z
end

function (r::ArfField)(x::UInt)
  z = arf(x, r.prec, r.rndmode)
  z.parent = r
  return z
end

function (r::ArfField)(x::fmpz)
  z = arf(x, r.prec, r.rndmode)
  z.parent = r
  return z
end

function (r::ArfField)(x::Float64)
  z = arf(x)
  z = arf(z, r.prec, r.rndmode)
  z.parent = r
  return z
end

function (r::ArfField)(x::BigFloat)
  z = arf(x)
  z = arf(z, r.prec, r.rndmode)
  z.parent = r
  return z
end

function (r::ArfField)(x::mag)
  z = arf(x)
  z = arf(z, r.prec, r.rndmode)
  z.parent = r
  return z
end

################################################################################
#
#  Special values
#
################################################################################

for (s,f) in (("zero", "arf_zero"), ("one", "arf_one"),
              ("pos_inf", "arf_pos_inf"), ("neg_inf", "arf_neg_inf"),
              ("nan", "arf_nan"))
  @eval begin
    function($(Symbol(s)))(r::ArfField)
      z = r()
      ccall(($f, :libarb), Void, (Ptr{arf}, ), &z)
      return z
    end
  end
end

################################################################################
#
#  String I/O
#
################################################################################

# this function is crap
function show(io::IO, a::arf)
  #cstr = ccall((:arf_get_str, :libarb), Ptr{UInt8}, (Ptr{arf}, ), a.data)
  #print(io, bytestring(cstr))
  #ccall((:flint_free, :libflint), Void, (Ptr{UInt8},), cstr)
  return show(io, BigFloat(a))
end

function show(io::IO, a::ArfField)
  print(io, "Field of arf's with precision $(a.prec)")
  print(io, " and rounding mode $(a.rndmode)")
end

################################################################################
#
#  Predicates
#
################################################################################

for (s,f) in (("iszero", "arf_iszero"), ("isone", "arf_is_one"),
              ("isposinf", "arf_is_pos_inf"),
              ("isneginf", "arf_is_neg_inf"),
              ("isinf", "arf_is_inf"),
              ("isnan", "arf_is_nan"),
              ("isnormal", "arf_is_normal"),
              ("isfinite", "arf_isfinite"),
              ("isspecial", "arf_is_special"))
  @eval begin
    function($(Symbol(s)))(x::arf)
      return Bool(ccall(($f, :libarb), Cint, (Ptr{arf},), x.data))
    end
  end
end

################################################################################
#
#  Comparison
#
################################################################################

function ==(x::arf, y::arf)
  r = ccall((:arf_equal, :libarb), Cint, (Ptr{arf}, Ptr{arf}), &x, &y)
  return Bool(r)
end

function compare(x::arf, y::arf)
  r = ccall((:arf_cmp, :libarb), Cint, (Ptr{arf}, Ptr{arf}), &x, &y)
  return r
end

function <(x::arf, y::arf)
  r = compare(x,y)
  r < 0 ? true : false
end

function >(x::arf, y::arf)
  r = compare(x,y)
  r > 0 ? true : false
end

function max(x::arf, y::arf)
  z = parent(x)()
  ccall((:arf_max, :libarb), Void, (Ptr{arf}, Ptr{arf}, Ptr{arf}), &z, &x, &y)
  return z
end

function min(x::arf, y::arf)
  z = parent(x)()
  ccall((:arf_min, :libarb), Void, (Ptr{arf}, Ptr{arf}, Ptr{arf}), &z, &x, &y)
  return z
end

################################################################################
#
#  Unary operations
#
################################################################################

function -(x::arf)
  z = parent(x)()
  ccall((:arf_neg, :libarb), Void, (Ptr{arf}, Ptr{arf}), &z, &x)
  return z
end

################################################################################
#
#  Binary operations
#
################################################################################

function +(x::arf, y::arf)
  check_parent(x,y)
  z = parent(x)()
  ccall(("arf_add", :libarb), Void,
              (Ptr{arf}, Ptr{arf}, Ptr{arf}, Int, Cint),
              &z, &x, &y, parent(x).prec, parent(x).rndmode)
  return z
end

function +(x::arf, y::UInt)
  z = parent(x)()
  ccall(("arf_add_ui", :libarb), Void,
              (Ptr{arf}, Ptr{arf}, UInt, Int, Cint),
              &z, &x, y, parent(x).prec, parent(x).rndmode)
  return z
end

+(x::UInt, y::arf) = +(y, x)

function +(x::arf, y::Int)
  z = parent(x)()
  ccall(("arf_add_si", :libarb), Void,
              (Ptr{arf}, Ptr{arf}, Int, Int, Cint),
              &z, &x, y, parent(x).prec, parent(x).rndmode)
  return z
end

+(x::Int, y::arf) = +(y, x)

function +(x::arf, y::fmpz)
  z = parent(x)()
  ccall(("arf_add_fmpz", :libarb), Void,
              (Ptr{arf}, Ptr{arf}, Ptr{fmpz}, Int, Cint),
              &z, &x, &y, parent(x).prec, parent(x).rndmode)
  return z
end

+(x::fmpz, y::arf) = +(y, x)

function *(x::arf, y::arf)
  z = parent(x)()
  ccall(("_arf_mul", :libarb), Void,
              (Ptr{arf}, Ptr{arf}, Ptr{arf}, Int, Cint),
              &z, &x, &y, parent(x).prec, parent(x).rndmode)
  return z
end

function *(x::arf, y::UInt)
  z = parent(x)()
  ccall(("arf_mul_ui", :libarb), Void,
              (Ptr{arf}, Ptr{arf}, UInt, Int, Cint),
              &z, &x, y, parent(x).prec, parent(x).rndmode)
  return z
end

*(x::UInt, y::arf) = *(y, x)

function *(x::arf, y::Int)
  z = parent(x)()
  ccall(("arf_mul_si", :libarb), Void,
              (Ptr{arf}, Ptr{arf}, Int, Int, Cint),
              &z, &x, y, parent(x).prec, parent(x).rndmode)
  return z
end

*(x::Int, y::arf) = *(y, x)

function *(x::arf, y::fmpz)
  z = parent(x)()
  ccall(("arf_mul_fmpz", :libarb), Void,
              (Ptr{arf}, Ptr{arf}, Ptr{fmpz}, Int, Cint),
              &z, &x, &y, parent(x).prec, parent(x).rndmode)
  return z
end

*(x::fmpz, y::arf) = *(y, x)

function -(x::arf, y::arf)
  check_parent(x,y)
  z = parent(x)()
  ccall(("arf_sub", :libarb), Void,
                (Ptr{arf}, Ptr{arf}, Ptr{arf}, Int, Cint),
                &z, &x, &y, parent(x).prec, parent(x).rndmode)
  return z
end

function -(x::arf, y::UInt)
  z = parent(x)()
  ccall(("arf_sub_ui", :libarb), Void,
              (Ptr{arf}, Ptr{arf}, UInt, Int, Cint),
              &z, &x, y, parent(x).prec, parent(x).rndmode)
  return z
end

-(x::UInt, y::arf) = -(y - x)

function -(x::arf, y::Int)
  z = parent(x)()
  ccall(("arf_sub_si", :libarb), Void,
              (Ptr{arf}, Ptr{arf}, Int, Int, Cint),
              &z, &x, y, parent(x).prec, parent(x).rndmode)
  return z
end

-(x::Int, y::arf) = -(y - x)

function -(x::arf, y::fmpz)
  z = parent(x)()
  ccall(("arf_sub_fmpz", :libarb), Void,
              (Ptr{arf}, Ptr{arf}, Ptr{fmpz}, Int, Cint),
              &z, &x, &y, parent(x).prec, parent(x).rndmode)
  return z
end

-(x::fmpz, y::arf) = -(y - x)

function /(x::arf, y::arf)
  check_parent(x,y)
  z = parent(x)()
  ccall((:arf_div, :libarb), Void,
              (Ptr{arf}, Ptr{arf}, Ptr{arf}, Int, Cint),
              &z, &x, &y, parent(x).prec, parent(x).rndmode)
  return z
end

################################################################################
#
#  Unsafe arithmetic
#
################################################################################

function add!(x::arf, y::arf)
  ccall(("arf_add", :libarb), Void,
              (Ptr{arf}, Ptr{arf}, Ptr{arf}, Int, Cint),
              &z, &x, &y, parent(x).prec, parent(x).rndmode)
end

function div!(z::arf, x::arf, y::arf)
  ccall((:arf_div, :libarb), Void,
              (Ptr{arf}, Ptr{arf}, Ptr{arf}, Int, Cint),
              &z, &x, &y, parent(x).prec, parent(x).rndmode)
end

function sub!(x::arf, y::arf)
  ccall((:arf_sub, :libarb), Void,
              (Ptr{arf}, Ptr{arf}, Ptr{arf}, Int, Cint),
              &z, &x, &y, parent(x).prec, parent(x).rndmode)
end

function mul!(x::arf, y::arf)
  ccall((:_arf_mul, :libarb), Void,
              (Ptr{arf}, Ptr{arf}, Ptr{arf}, Int, Cint),
              &z, &x, &y, parent(x).prec, parent(x).rndmode)
end

################################################################################
#
#  Some functions
#
################################################################################

function sign(x::arf)
  r = ccall((:arf_sgn, :libarb), Cint, (Ptr{arf}, ), &x)
  return r
end

function abs(x::arf)
  z = parent(x)()
  ccall((:arf_abs, :libarb), Void, (Ptr{arf}, Ptr{arf}), &z, &x)
  return z
end

function Base.sqrt(x::arf)
  z = parent(x)()
  ccall((:arf_sqrt, :libarb), Void,
              (Ptr{arf}, Ptr{arf}, Int, Cint),
              &z, &x, parent(x).prec, parent(x).rndmode)
  return z
end

################################################################################
#
#  Conversion
#
################################################################################

function Cdouble(a::arf, rnd::Cint = 4)
  z = ccall((:arf_get_d, :libarb), Cdouble, (Ptr{arf}, Cint), a.data, rnd)
  return z
end

function BigFloat(x::arf)
  old_prec = get_bigfloat_precision()
  set_bigfloat_precision(parent(x).prec)
  z = BigFloat(0)
  r = ccall((:arf_get_mpfr, :libarb), Cint,
                (Ptr{BigFloat}, Ptr{arf}, Cint), &z, &x, Cint(0))
  set_bigfloat_precision(old_prec)
  return z
end


using Plots, DelimitedFiles

const CLASSES = filter(e-> e in ["cutest","power"], ARGS)
const SOLVERS = filter(e-> e in ["current","master","ipopt","knitro"], ARGS)

const LABELS = Dict(
    "current" => "MadNLP (dev)",
    "master" => "MadNLP (master)",
    "ipopt" => "Ipopt",
    "knitro" => "Knitro",
)

for class in CLASSES

    time = Dict()
    status = Dict()
    mem = Dict()
    iter = Dict()

    name = readdlm("name-$class.csv")[:]
    for solver in SOLVERS
        status[solver] = readdlm("status-$class-$solver.csv")[:]
        time[solver] = readdlm("time-$class-$solver.csv")[:]
        mem[solver] = readdlm("mem-$class-$solver.csv")[:]
        iter[solver] = readdlm("iter-$class-$solver.csv")[:]
    end

    for (str,metric) in [("time",time),("iterations",iter),("memory",mem)]
        relmetric = deepcopy(metric)

        for i=1:length(name)
            top = []
            topstatus = 3
            for solver in SOLVERS
                if status[solver][i] < topstatus
                    empty!(top)
                    push!(top,solver)
                    topstatus = status[solver][i]
                elseif status[solver][i] == topstatus
                    push!(top,solver)
                end
            end

            fastest = Inf
            for solver in top
                fastest = min(fastest,metric[solver][i])
            end

            for solver in SOLVERS
                if status[solver][i] == 3
                    relmetric[solver][i] = NaN
                else
                    relmetric[solver][i] = log2(relmetric[solver][i]/fastest)
                end
            end
        end

        for solver in SOLVERS
            filter!(a->!isnan(a),relmetric[solver])
        end

        p = plot(;
                 ylim=(0,1),
                 xlim=(
                     minimum(minimum(relmetric[solver]) for solver in SOLVERS),
                     maximum(maximum(relmetric[solver]) for solver in SOLVERS)
                 ),
                 xlabel="Not More Than 2ˣ-Times Worse Than Best Solver ($str)",
                 ylabel="Fraction of Problems Solved",
                 framestyle=:box,
                 legend=:bottomright)

        for solver in SOLVERS
            y = [0:length(relmetric[solver]);length(relmetric[solver])]/length(name)
            push!(relmetric[solver],maximum(maximum(relmetric[s]) for s in SOLVERS))
            push!(relmetric[solver],0)
            plot!(p,sort(relmetric[solver]),y;
                  qqline=:none,
                  linetype=:steppost,
                  label=LABELS[solver])
        end
        savefig(p,"$str-$class.pdf")
    end
end


using CairoMakie
using Images, FileIO
using ImageTransformations
using WindMixing
using JLD2
using WindMixing: animate_profiles_fluxes_final

# convective_adjustment = parse(Bool, ARGS[1])
convective_adjustment = false
num = parse(Int, ARGS[1])

if convective_adjustment
    CA_str = "CA"
else
    CA_str = "noCA"
end

u_img = rotr90(load(assetpath(joinpath(pwd(), "images_plot", "u.png"))))
v_img = rotr90(load(assetpath(joinpath(pwd(), "images_plot", "v.png"))))
T_img = rotr90(load(assetpath(joinpath(pwd(), "images_plot", "T.png"))))
uw_img = rotr90(load(assetpath(joinpath(pwd(), "images_plot", "uw.png"))))
vw_img = rotr90(load(assetpath(joinpath(pwd(), "images_plot", "vw.png"))))
wT_img = rotr90(load(assetpath(joinpath(pwd(), "images_plot", "wT.png"))))
Ri_img = rotr90(load(assetpath(joinpath(pwd(), "images_plot", "Ri.png"))))
z_img = load(assetpath(joinpath(pwd(), "images_plot", "z.png")))

axis_images = (
    u = u_img,
    v = v_img,
    T = T_img,
    z = z_img,
    uw = uw_img,
    vw = vw_img,
    wT = wT_img,
    Ri = Ri_img,
)

EXTRACTED_FILE_DIR = joinpath(pwd(), "extracted_training_output")

EXTRACTED_FILE_PATH = joinpath(EXTRACTED_FILE_DIR, "NDE_3sim_diurnal_18simBFGST0.8nograd_divide1f5_gradient_smallNN_leakyrelu_rate_2e-4_T0.8_extracted.jld2")
FILE_DIR = joinpath(pwd(), "final_data_2", "NDE_3sim_diurnal_18simBFGST0.8nograd_divide1f5_gradient_smallNN_leakyrelu_rate_2e-4_T0.8_$(CA_str)")

# if !ispath(FILE_DIR)
#     mkdir(FILE_DIR)
# end

file = jldopen(EXTRACTED_FILE_PATH)
train_files = file["training_info/train_files"]
loss_scalings = file["training_info/loss_scalings"]
train_parameters = file["training_info/parameters"]
close(file)

test_files_all = [
    "wind_-5e-4_diurnal_5e-8",
    "wind_-3.5e-4_diurnal_5e-8",
    "wind_-2e-4_diurnal_5e-8",

    "wind_-4.5e-4_diurnal_4e-8",
    "wind_-4.5e-4_diurnal_2e-8",
    "wind_-3e-4_diurnal_4e-8",
    "wind_-3e-4_diurnal_2e-8",

    "wind_-5.5e-4_diurnal_5.5e-8",
    "wind_-1.5e-4_diurnal_5.5e-8",

    "wind_-5.5e-4_new",

    "wind_-5.5e-4_heating_-3.5e-8",
    "wind_-1.5e-4_heating_-3.5e-8",
    "wind_-5.5e-4_cooling_3.5e-8",
    "wind_-1.5e-4_cooling_3.5e-8",
]

test_files = [test_files_all[num]]

ν₀ = train_parameters["ν₀"]
ν₋ = train_parameters["ν₋"]
ΔRi = train_parameters["ΔRi"]
Riᶜ = train_parameters["Riᶜ"]
Pr = train_parameters["Pr"]

# solve_oceananigans_modified_pacanowski_philander_nn(test_files, EXTRACTED_FILE_PATH, FILE_DIR,
#                                                         timestep=1, convective_adjustment=convective_adjustment)

diurnal = occursin("diurnal", test_files[1])

if num <= 3
    DIR_NAME = "train_$(test_files[1])"
    animation_type = "Training"
elseif num >= 4 && num <= 7
    DIR_NAME = "test_$(test_files[1])"
    animation_type = "Interpolating"
else
    DIR_NAME = "test_$(test_files[1])"
    animation_type = "Extrapolating"
end

n_trainings = length(train_files)

plot_data = NDE_profile_oceananigans(joinpath(FILE_DIR, DIR_NAME), train_files, test_files,
                                  ν₀=ν₀, ν₋=ν₋, ΔRi=ΔRi, Riᶜ=Riᶜ, Pr=Pr,
                                  loss_scalings=loss_scalings,
                                  OUTPUT_PATH=joinpath(FILE_DIR, DIR_NAME, "profiles_fluxes_oceananigans.jld2"))

animate_profiles_fluxes_final(plot_data, axis_images,
                    joinpath(FILE_DIR, DIR_NAME, "$(test_files[1])_60fps"),
                    animation_type=animation_type, n_trainings=n_trainings, training_types="Wind + Diurnal Heating & Cooling", fps=60)