"""
main : the code you run to double-encode... the "main loop" of optimization.
"""

module main

include("utils.jl")
include("types.jl")
include("lookup.jl")
include("std_setup.jl")
include("math.jl")
include("optimize.jl")
include("mrf.jl")

using ArgParse, Logging, JLD, StatsBase, Distributions, Random, Statistics

function convert_to_saveable(cur_pop)
	new_pop = types.SaveChrome[]

	for indiv in cur_pop
		push!(new_pop, types.SaveChrome(indiv.path, indiv.full_sequence,
										indiv.deg_nuc, indiv.deg_map, indiv.deg_trns, indiv.deg_trne, indiv.deg_d, indiv.deg_skip, indiv.deg_insert,
										indiv.mark_nuc, indiv.mark_map, indiv.mark_trns, indiv.mark_trne, indiv.mark_d, indiv.mark_skip, indiv.mark_insert,
										indiv.deg_prob, indiv.deg_base_E, indiv.deg_seq,
										indiv.mark_prob, indiv.mark_base_E, indiv.mark_seq, indiv.first_weight))
	end

	return new_pop
end

"""
set_up_and_optimize is the code that sets up appropriate information about genes we want to target for double-encoding and
orchestrates the initial HMM double-encoding solution and then iteratively improves on that with greedy pseudolikelihood improvements.
Code also logs information through Logging package.
Genes are often referred to as mark/deg. This is historic from times of thinking of "marker" genes and "designated essential gene".
At this point this distinction is not meaningful but is kept to avoid equally arbitrary x_name, y_name conventions.
Arguments:
	mark_name : gene ID for 'mark' gene. Needed as a key for looking up some values associated with genes in text files.
	deg_name : gene ID for 'deg' gene.
	mark/deg_grem : path to MRF parameter files, usually a JLD/CSV file.
	mark/deg_hmm : path to HMM files, usually .hmm files.
	pop_size : size of population, i.e. number of individual HMM solutions to greedily optimize.
	frame : p1/p2.
	max_iter : maximum number of iterations.
	X_range/Y_range: positions along gene to consider for double-encoding... useful if you want to restrict where double-encoding can occur.
						Ranges are defined in terms of the end of the sequence, so if you have a 70 aa sequence you want to start at position 80,
						set the range to be something like 150:151. Generally some buffer around this value (i.e. 150:160) is useful.
						Don't set, or set to false if you don't care where the genes overlap.
"""
function set_up_and_optimize(log_io, rand_barcode, out_path, mark_name, deg_name, mark_grem, deg_grem, mark_hmm, deg_hmm, pop_size, frame, max_iter, X_range=false, Y_range=false; rand_weights = false, actually_mrf = true)

	@debug("Beginning run.")
	@debug(Libc.strftime(time()))
	@debug(Libc.strftime(time()))

	println("The random barcode on this run is: $rand_barcode")
	@debug("The random barcode on this run is: $rand_barcode")

	do_cull = false #culling reduces number of sequences we optimize over time. Just a trick to save time if you want top-performers only.
	full_even_window = 20
	half_window = div(full_even_window, 2) #window statistics to see if scores still decreasing (sometimes used).
	last_few_sf = Float64[]
	last_few_cs = Float64[]

	#Looks up mean/std. dev of family pseudolikelihoods, pre-computed.
	mu_mark, sig_mark = lookup.mu_sig(deg_name, "psls/")
	mu_deg, sig_deg = lookup.mu_sig(mark_name, "psls/")

	@debug("Stat param vars (mark / deg):\t$(mu_mark)\t$(sig_mark)\t$(mu_deg)\t$(sig_deg)")

	if true #replacing try/catch.
		gen_samples = true #this is true if starting from scratch (general case), false if we have some candidates to optimize ahead of time.
		population = types.Chromosome[]
		local mark_grem_prot, deg_grem_prot
		if !gen_samples
			population = load("trpE_population.jld")["population"][1:100] #[1:100]
			mark_gremodel, deg_gremodel = std_setup.short_set_up(mark_name, deg_name, mark_grem, deg_grem, mark_hmm, deg_hmm, length(population), rand_barcode, frame)
			deg_nNodes, mark_nNodes = deg_gremodel.nNodes, mark_gremodel.nNodes
		else
			#Generate population of sampled hmm starting points.
			if X_range == false || Y_range == false #we require both to be there...
				@debug("Doing standard full set up...")
				mark_gremodel, deg_gremodel, population, mark_grem_prot, deg_grem_prot = std_setup.full_set_up(out_path, mark_name, mark_hmm, mark_grem, deg_name, deg_hmm, deg_grem, pop_size, 1200, 1200, rand_barcode, frame)
			else
				@debug("The x range is $X_range")
				@debug("The y range is $Y_range")
				mark_gremodel, deg_gremodel, population, mark_grem_prot, deg_grem_prot = std_setup.full_sample_set_up(out_path, mark_name, mark_hmm, mark_grem, deg_name, deg_hmm, deg_grem, pop_size, 1200, 1200, rand_barcode, X_range, Y_range, frame)
			end
			deg_nNodes, mark_nNodes = deg_gremodel.nNodes, mark_gremodel.nNodes
		end

		#Look through generated candidates, check their fitness values and their correctness.
		success_chrom, fdp, fmp, fitness_values, indie_deg_maps, indie_mark_maps, deg_ull, deg_pv_w1, deg_pv_w2, deg_prot_mat, mark_ull, mark_pv_w1, mark_pv_w2, mark_prot_mat = optimize.assess_founders(population, mark_gremodel.nNodes, mark_gremodel.w1, mark_gremodel.w2, deg_gremodel.nNodes, deg_gremodel.w1, deg_gremodel.w2)

		@debug("Num success... $(length(success_chrom))")

		if actually_mrf
			founding_fitness_values = fitness_values[1:end]

			mark_base_energy = mrf.basic_energy_calc(mark_grem_prot, mark_gremodel.w1, mark_gremodel.w2)
			deg_base_energy = mrf.basic_energy_calc(deg_grem_prot, deg_gremodel.w1, deg_gremodel.w2)

			#Now we load this info into an ExChrome type object.
			cur_pop = types.ExChrome[]
			for i in 1:length(success_chrom)
				old = success_chrom[i]
				if rand_weights
					new_chrom = types.ExChrome(old.path, old.full_sequence, old.deg_nuc, indie_deg_maps[i], old.deg_trns, old.deg_trne,
															old.deg_d, old.deg_skip, old.deg_insert, old.mark_nuc, indie_mark_maps[i], old.mark_trns, old.mark_trne,
															old.mark_d, old.mark_skip, old.mark_insert, fdp[i], deg_base_energy, deg_prot_mat[i:i, 1:end], utils.aa_vec_to_seq(deg_prot_mat[i:i, 1:end]), deg_ull[i], deg_pv_w1[i], deg_pv_w2[i:i, 1:end], fmp[i], mark_base_energy,
															mark_prot_mat[i:i, 1:end], utils.aa_vec_to_seq(mark_prot_mat[i:i, 1:end]), mark_ull[i], mark_pv_w1[i], mark_pv_w2[i:i, 1:end], rand())
				else
					new_chrom = types.ExChrome(old.path, old.full_sequence, old.deg_nuc, indie_deg_maps[i], old.deg_trns, old.deg_trne,
															old.deg_d, old.deg_skip, old.deg_insert, old.mark_nuc, indie_mark_maps[i], old.mark_trns, old.mark_trne,
															old.mark_d, old.mark_skip, old.mark_insert, fdp[i], deg_base_energy, deg_prot_mat[i:i, 1:end], utils.aa_vec_to_seq(deg_prot_mat[i:i, 1:end]), deg_ull[i], deg_pv_w1[i], deg_pv_w2[i:i, 1:end], fmp[i], mark_base_energy,
															mark_prot_mat[i:i, 1:end], utils.aa_vec_to_seq(mark_prot_mat[i:i, 1:end]), mark_ull[i], mark_pv_w1[i], mark_pv_w2[i:i, 1:end], 0.5)
				end
				push!(cur_pop, new_chrom)
			end

			#There are occasionally differences between HMM/MRF models of the proteins, mostly because HMMs are more flexible with insertions.
			#This generates mappings between positions in either data type.
			mark_hmm_to_grem_map = std_setup.get_explicit_mapping(mark_grem_prot, mark_hmm)
			deg_hmm_to_grem_map = std_setup.get_explicit_mapping(deg_grem_prot, deg_hmm)

			@debug("Hey the explicit mapping for mark_hmm_to_grem is $mark_hmm_to_grem_map")
			@debug("Hey the explicit mapping for deg_hmm_to_grem is $deg_hmm_to_grem_map")
			#So now everything is in an extended chromosome object.

			changed_seq = length(cur_pop)
			no_change_iters = 0
			iter = 0;

			#Let's keep track of optimization history.
			history_matrix = zeros(Float32, round(Int64, pop_size * 1.2), max_iter)

			println("Beginning long-range optimization.")
			@debug("About to start optimization.")
			@debug(Libc.strftime(time()))

			stop_early = false #a condition to evaluate every (few) iteration(s) to see if it's worth continuing a run. When true, we clean up and move to next gene pair.
			not_worth_continuing = false

			#Okay let's set up the energy normal distributions...

			deg_energy_mu, deg_energy_sig = lookup.get_energy_params(deg_name, "energies/")
			mark_energy_mu, mark_energy_sig = lookup.get_energy_params(mark_name, "energies/")

			mark_energy_normal = Normal(mark_energy_mu, mark_energy_sig)
			deg_energy_normal = Normal(deg_energy_mu, deg_energy_sig)

			while (iter == 0 || (iter < max_iter)) #&& !(stop_early) # && minimum((fitness_values)) > 400 && no_change_iters < 100)
				sfv = sort(fitness_values)
				#best_50 = sfv[50]

				if iter % 50 == 0
					println("Step $(iter) of $(max_iter)...")
				end
				@debug("Iteration $iter out of $max_iter")
				@debug("Mean fitness is: $(mean((fitness_values)))")
				@debug("Bottom five are $(sfv[1:5])")
				@debug("... and top five are $(sfv[end-5:end]).")
				@debug("Number of sequences that were changed: $changed_seq")
				flush(log_io)

				#iterated conditional modes, multi-site keep-track = icm_multi_kt. keep track = store values of changes to sequence.
				if do_cull
					if iter < div(max_iter, 4)
						#println("$iter : no cutoff")
						cur_pop, changed_seq = optimize.icm_multi_kt(cur_pop, deg_gremodel, mark_gremodel) #, deg_energy_normal, mark_energy_normal)
					elseif iter >= div(max_iter, 4) && iter < div(max_iter, 2)
						#println("$iter : cutoff is $(sfv[div(length(sfv), 2)])")
						cur_pop, changed_seq = optimize.icm_multi_kt(cur_pop, deg_gremodel, mark_gremodel; score_cutoff = sfv[div(length(sfv), 2)])
					elseif iter >= div(max_iter, 2)
						#println("$iter : cutoff is $(sfv[div(length(sfv), 4)])")
						cur_pop, changed_seq = optimize.icm_multi_kt(cur_pop, deg_gremodel, mark_gremodel; score_cutoff = sfv[div(length(sfv), 4)])
					end
				else
					cur_pop, changed_seq = optimize.icm_multi_kt(cur_pop, deg_gremodel, mark_gremodel)
				end

				fitness_values = optimize.assess_pop(cur_pop)
				summed_fitness = sum(fitness_values) #one of the things we should be watching.

				push!(last_few_sf, summed_fitness)
				push!(last_few_cs, changed_seq)

				indiv_count = 1
				for indiv in fitness_values #history matrix stores values of function being optimized over time.
					history_matrix[indiv_count, iter + 1] = indiv
					indiv_count += 1
				end

				#This was originally just code to test the pseudolikelihoods decreasing correctly. It does.
				#Now it's still sorta somewhat useful to log a few values of sequences in both deg/mark as optimization goes.
				if (mod(iter, 50) == 0 && iter != 0)
					@debug("Second opinion, these should be ~0.")
					for i_count in 1:5
						#for the record this is not a good way to generate random numbers.
						#doesn't go from 1 to end.
						i = i_count #round(Int64, rand() * (length(cur_pop) - 1) + 1)

						bel_deg = cur_pop[i].deg_prob
						bel_mrk = cur_pop[i].mark_prob
						#Sort of assuming that end-1 is because of a "*" at stop codon being added.
						if length(strip(cur_pop[i].deg_seq, '*')) == deg_nNodes && length(strip(cur_pop[i].mark_seq, '*')) == mark_nNodes
							cal_deg = mrf.psl(strip(cur_pop[i].deg_seq, '*'), deg_gremodel.w1, deg_gremodel.w2, false)
							cal_mrk = mrf.psl(strip(cur_pop[i].mark_seq, '*'), mark_gremodel.w1, mark_gremodel.w2, false)
							@debug("DEG: $bel_deg vs. $cal_deg")
							@debug("MRK: $bel_mrk vs. $cal_mrk")
							@debug(abs(bel_deg - cal_deg) < 1e-2)
							@debug(abs(bel_mrk - cal_mrk) < 1e-2)
						else
							@debug("Proteins not correct length!")
							@debug(cur_pop[i].mark_seq)
							@debug(cur_pop[i].deg_seq)
						end
					end
				end
				if changed_seq == 0
					no_change_iters += 1
				else
					no_change_iters = 0
				end
				iter += 1
				@debug("\n")
			end

			#So we're done the iterated conditional modes step. Now we report everything...
			@debug("Done while loop...")
			@debug(Libc.strftime(time()))

			@debug("Let us save the history matrix!")
			save("$out_path/$(mark_name)_$(deg_name)_$frame/opt_his_mat_$(rand_barcode).jld", "hist_mat", history_matrix)

			deg_significance = false
			mark_significance = false #this sees if we get any hits worth keeping at all. If yes we save full population.
			all_cal_deg_scores = Float64[]
			all_cal_mark_scores = Float64[]
			#If no we save just a few (top 10) individuals. These are top-6 overall + top-2 in deg and top-2 in mark.
			out_file = open("$out_path/$(mark_name)_$(deg_name)_$frame/all_final_fitness_$(rand_barcode).txt", "w")
			write(out_file, "Ind. #\tMark Score\tDeg Score\tMark Sign\tDeg Sign\n")
			for last_indi in 1:length(cur_pop)
				rep_deg = cur_pop[last_indi].deg_prob
				rep_mrk = cur_pop[last_indi].mark_prob
				if length(strip(cur_pop[last_indi].deg_seq, '*')) == deg_nNodes && length(strip(cur_pop[last_indi].mark_seq, '*')) == mark_nNodes
					cal_deg = abs(mrf.psl(strip(cur_pop[last_indi].deg_seq, '*'), deg_gremodel.w1, deg_gremodel.w2, false))
					cal_mrk = abs(mrf.psl(strip(cur_pop[last_indi].mark_seq, '*'), mark_gremodel.w1, mark_gremodel.w2, false))
					special_deg = "-"
					if cal_deg < mu_deg + sig_deg
						special_deg = "***!"
						deg_significance = true
					elseif cal_deg < mu_deg + 2*sig_deg
						special_deg = "**."
						deg_significance = true
					elseif cal_deg < mu_deg + 3*sig_deg
						special_deg = "*?"
						deg_significance = true
					elseif cal_deg < mu_deg + 4*sig_deg #we don't mind keeping these.
						deg_significance = true
					end
					special_mark = "-"
					if cal_mrk < mu_mark + sig_mark
						special_mark = "***!"
						mark_significance = true
					elseif cal_mrk < mu_mark + 2*sig_mark
						special_mark = "**."
						mark_significance = true
					elseif cal_mrk < mu_mark + 3*sig_mark
						special_mark = "*?"
						mark_significance = true
					elseif cal_mrk < mu_mark + 4*sig_mark #we don't mind keeping these.
						mark_significance = true
					end
					push!(all_cal_mark_scores, cal_mrk)
					push!(all_cal_deg_scores, cal_deg)
					write(out_file, "$(last_indi)\t$(cal_mrk)\t$(cal_deg)\t$special_mark\t$special_deg\n")
				else
					push!(all_cal_mark_scores, Inf)
					push!(all_cal_deg_scores, Inf)
					write(out_file, "$(last_indi)\t-1.0\t-1.0\t/\t/\n")
				end
			end
			close(out_file)

			cal_deg_p = sortperm(all_cal_deg_scores)
			cal_mark_p = sortperm(all_cal_mark_scores)
			cal_both_p = sortperm(fitness_values)

			@debug(Libc.strftime(time()))
			if true #deg_significance && mark_significance
				@debug("Let's save our optimized population...")
				saved_pop = convert_to_saveable(cur_pop)
				save("$out_path/$(mark_name)_$(deg_name)_$frame/saved_pop_$(rand_barcode).jld", "variants", saved_pop)
			else #We want just a subset.
				@debug("We'll save 12 interesting members of the population.")
				selected_pop = ExChrome[]
				for ijk in 1:3
					push!(selected_pop, cur_pop[cal_deg_p[ijk]])
					push!(selected_pop, cur_pop[cal_mark_p[ijk]])
				end
				for ijk in 1:6
					push!(selected_pop, cur_pop[cal_both_p[ijk]])
				end
				save("$out_path/$(mark_name)_$(deg_name)_$frame/top_pop_$(rand_barcode).jld", "variants", selected_pop)
			end

			@debug("And we'll also output our top twelve sequences...")
			out_file = open("$out_path/$(mark_name)_$(deg_name)_$frame/top_twelve_$(rand_barcode).fa", "w")

			for ijk in 1:3
				deg_ind = cal_deg_p[ijk]
				write(out_file, ">TOP_DEG_d (ind $deg_ind ) (fit: $(fitness_values[deg_ind])) (deg: $(all_cal_deg_scores[deg_ind])) (mark: $(all_cal_mark_scores[deg_ind]))\n")
				write(out_file, "$(cur_pop[deg_ind].deg_seq)\n")
				write(out_file, ">TOP_DEG_m (ind $deg_ind ) (fit: $(fitness_values[deg_ind])) (deg: $(all_cal_deg_scores[deg_ind])) (mark: $(all_cal_mark_scores[deg_ind]))\n")
				write(out_file, "$(cur_pop[deg_ind].mark_seq)\n")
			end
			for ijk in 1:3
				deg_ind = cal_mark_p[ijk]
				write(out_file, ">TOP_MARK_d (ind $deg_ind ) (fit: $(fitness_values[deg_ind])) (deg: $(all_cal_deg_scores[deg_ind])) (mark: $(all_cal_mark_scores[deg_ind]))\n")
				write(out_file, "$(cur_pop[deg_ind].deg_seq)\n")
				write(out_file, ">TOP_MARK_m (ind $deg_ind ) (fit: $(fitness_values[deg_ind])) (deg: $(all_cal_deg_scores[deg_ind])) (mark: $(all_cal_mark_scores[deg_ind]))\n")
				write(out_file, "$(cur_pop[deg_ind].mark_seq)\n")
			end
			for ijk in 1:6
				deg_ind = cal_both_p[ijk]
				write(out_file, ">TOP_GEN_d (ind $deg_ind ) (fit: $(fitness_values[deg_ind])) (deg: $(all_cal_deg_scores[deg_ind])) (mark: $(all_cal_mark_scores[deg_ind]))\n")
				write(out_file, "$(cur_pop[deg_ind].deg_seq)\n")
				write(out_file, ">TOP_GEN_m (ind $deg_ind ) (fit: $(fitness_values[deg_ind])) (deg: $(all_cal_deg_scores[deg_ind])) (mark: $(all_cal_mark_scores[deg_ind]))\n")
				write(out_file, "$(cur_pop[deg_ind].mark_seq)\n")
			end
			close(out_file)
		end

		@debug("REALLY DONE!")

		@debug(Libc.strftime(time()))

	end
end

function parse_commandline()
	s = ArgParseSettings()
	@add_arg_table s begin
		"commands"
			help = "positional argument, path to file with tab-delimited commands to run (default: example.txt)"
			default = "example.txt"
		"--num"
			help = "optional, experimentally used for running multiple jobs at once"
			default = "0"
		"--threads"
			help = "optional, experimentally used for running multiple jobs at once"
			default = "1"
	end
	return parse_args(s)
end

function run_file()
	parsed_args = parse_commandline()

	command_file = parsed_args["commands"]
	num = parsed_args["num"]
	threads = parsed_args["threads"]

	run_i = parse(Int64, num)
	NUM_THREADS = parse(Int64, threads)

	if NUM_THREADS > 1
		sleep(rand() * 2) #randomly delay start to de-synchronize starts of runs.
	end

	if isfile(command_file)
		in_file = open(command_file)
		in_read = readlines(in_file)
		close(in_file)

		#This is an additional file that records errors independent of individual log files.
		problem_file = open("problem_runs_$(run_i).txt", "w")

		line_count = 0
		for line in in_read #
			line_count += 1
			if line[1] != '#' && (line_count % NUM_THREADS == run_i)
				try
					run_args = split(line, "\t")
					out_dir, short, long, short_jld, long_jld, short_hmm, long_hmm, pop_size, frame, num_iter = run_args

					rand_barcode = Random.randstring()

					the_out_path = "$out_dir/$(short)_$(long)_$frame/"
					if !(isdir(the_out_path))
						run(`mkdir -p $the_out_path`)
					end

					log_io = open("$out_dir/$(short)_$(long)_$frame/log_$(rand_barcode).txt", "w+")

					logger = SimpleLogger(log_io, Logging.Debug)
					global_logger(logger)

					with_logger(logger) do
						pop_size = parse(Int64, pop_size)
						num_iter = parse(Int64, num_iter)
						set_up_and_optimize(log_io, rand_barcode, out_dir, short, long, short_jld, long_jld, short_hmm, long_hmm, pop_size, frame, num_iter)
					end

					flush(log_io)
					close(log_io)
				catch y
					write(problem_file, "$line $y")
				end

			end
		end
		close(problem_file)
	else
		println("$command_file does not exist. Exiting.")
	end

end

@time run_file()

end


import ThreeJS

main(window) =  begin
    push!(window.assets,("ThreeJS","threejs"))
    push!(window.assets,"widgets")
    w = Signal(1.0)
    h = Signal(1.0)
    d = Signal(1.0)
    rx = Signal(0.0)
    ry = Signal(0.0)
    rz = Signal(0.0)
    o = Signal(1.0)
    vbox(
        "ThreeJS Example",
        vskip(2em),
        hbox(
            vbox(
                "Sides",
                hbox("width",slider(1.0:5.0) >>> w),
                hbox("height",slider(1.0:5.0) >>> h),
                hbox("depth",slider(1.0:5.0) >>> d),
            ),
            hskip(2em),
            vbox(
                "Rotations",
                hbox("x", slider(0.0:5.0:360.0) >>> rx),
                hbox("y", slider(0.0:5.0:360.0) >>> ry),
                hbox("z", slider(0.0:5.0:360.0) >>> rz),
            ),
            hbox("opacity",slider(0.0:0.1:1.0, value = 1.0) >>> o),
        ),
        vskip(2em),
        map(w, h, d, rx, ry, rz, o) do w, h, d, rx, ry, rz, o
        ThreeJS.outerdiv() <<
            (ThreeJS.initscene() <<
                [
                    ThreeJS.mesh(0.0, 0.0, 0.0, rx = rx, ry = ry, rz = rz) <<
                    [
                        ThreeJS.box(w, h, d), ThreeJS.material(Dict(:kind=>"lambert",:color=>"red", :transparent=>true, :opacity=>o))
                    ],
                    ThreeJS.pointlight(10.0, 10.0, 10.0),
                    ThreeJS.camera(0.0, 0.0, 10.0)
                ]
            )
        end
    ) |> pad(2em)
end


module RocketSomeOperatorTest

using Test
using Rocket

include("../test_helpers.jl")

@testset "operator: some()" begin

    println("Testing: operator some()")

    run_proxyshowcheck("Some", some())

    run_testset([
        (
            source = from(1:42) |> max() |> some(),
            values = @ts([ 42, c ]),
            source_type = Int,
        ),
        (
            source = from([ 2, 3, 4, nothing ]) |> some(),
            values = @ts([ 2, 3, 4, c ]),
            source_type = Int,
        ),
        (
            source = completed() |> skip_next(),
            values = @ts(c)
        ),
        (
            source = completed(Int) |> max() |> some(),
            values = @ts(c),
            source_type = Int
        ),
        (
            source = faulted(1) |> skip_next(),
            values = @ts(e(1))
        ),
        (
            source = never() |> skip_next(),
            values = @ts()
        )
    ])

    @testset begin
        @test_throws Exception of(1) |> some()
    end

end

end


using NESSie.Format
using NESSie.Format: readmsms_nodes, readmsms_elements

testfiles = ((mktemp()..., mktemp()..., (0,0)), # empty file
             (mktemp()..., mktemp()..., (3,2))) # dummy file

write(testfiles[2][2], """
# MSMS solvent excluded surface vertices
#vertex #sphere density probe_r
  x  x  x  x
  0.0 0.0 0.0 x x x x x x
  1.0 2.0 3.0 x x x x x x
  -2.00001 1.337 42.0 x x x x x x
""")

write(testfiles[2][4], """
# MSMS solvent excluded surface faces
#faces  #sphere density probe_r
  x  x  x  x
  1 2 3 x x
  2 3 1 x x
""")

@testset "readmsms" begin
    try
        for (fnamev, fv, fnamef, ff, len) in testfiles, fn in ("single", "bulk"), T in testtypes
            seekstart(fv)
            seekstart(ff)
            nodes = (); elements = ()
            if fn == "single"
                # check single methods
                nodes = readmsms_nodes(fv, T)
                elements = readmsms_elements(ff, nodes)
            elseif fn == "bulk"
                # check bulk method (implicit single methods with atstart=true)
                model = readmsms(fv, ff, T)
                nodes, elements = (model.nodes, model.elements)
            end
            # check return types
            @test typeof(nodes) == Vector{Vector{T}}
            @test typeof(elements) == Vector{Triangle{T}}
            # check lengths
            @test length(nodes) == len[1]
            @test length(elements) == len[2]

            # check content
            if fnamev == testfiles[2][1]
                @test nodes[1] == T[0, 0, 0]
                @test nodes[2] == T[1, 2, 3]
                @test nodes[3] == T[-2.00001, 1.337, 42]
                @test elements[1].v1 === nodes[1]
                @test elements[1].v2 === nodes[2]
                @test elements[1].v3 === nodes[3]
                @test elements[2].v1 === nodes[2]
                @test elements[2].v2 === nodes[3]
                @test elements[2].v3 === nodes[1]
            end
        end
    finally
        [begin close(fv); rm(fnamev); close(ff); rm(fnamef) end for (fnamev, fv, fnamef, ff, _) in testfiles]
    end
end


const IndexTuple = Union{Tuple, NamedTuple}

"""
    IndexSpace(interval, boundingrect, nrows)

Metadata about an chunk.

- `interval`: An `Interval` object with the first and the last index tuples.
- `boundingrect`: An `Interval` object with the lowest and the highest indices as tuples.
- `nrows`: A `Nullable{Int}` of number of rows in the NDSparse, if knowable.
"""
struct IndexSpace{T<:IndexTuple}
    interval::Interval{T}
    boundingrect::Interval{T}
    nrows::Nullable{Int}
end

"""
A distributed table
"""
mutable struct DIndexedTable{T,K} <: AbstractIndexedTable
    # primary key columns
    pkey::Vector{Int}
    # extent of values in the pkeys
    domains::Vector{IndexSpace}
    chunks::Vector
end

noweakref(w::WeakRefString) = string(w)
noweakref(x) = x
function Dagger.domain(t::IndexedTable)
    ks = pkeys(t)
    T = eltype(ks)

    if isempty(t)
        return t.pkey => EmptySpace{T}()
    end

    wrap = T<:NamedTuple ? IndexedTables.namedtuple(fieldnames(T)...)∘tuple : tuple

    interval = Interval(map(noweakref, first(ks)), map(noweakref, last(ks)))
    cs = astuple(columns(ks))
    extr = map(extrema, cs[2:end]) # we use first and last value of first column
    boundingrect = Interval(wrap(noweakref(first(cs[1])), noweakref.(map(first, extr))...),
                            wrap(noweakref(last(cs[1])), noweakref.(map(last, extr))...))
    return t.pkey => IndexSpace(interval, boundingrect, Nullable{Int}(length(t)))
end

# if one of the input vectors is a Dagger operation / array
# chose the distributed implementation.
IndexedTables._impl(::Val, ::Dagger.ArrayOp, x...) = Val(:distributed)

function table(::Val{:distributed}, tup::Tup; chunks=nothing, kwargs...)
    if chunks === nothing
        # this means the vectors are distributed.
        # pick the first distributed vector and distribute
        # all others similarly
        idx = findfirst(x->isa(x, ArrayOp), tup)
        if idx == 0
            error("Don't know how to distribute. specify `chunks`")
        end
        darr = compute(tup[idx])
        chunks = domainchunks(darr)
    end

    darrays = map(x->distribute(x, chunks), tup)

    if isempty(darrays)
        error("Table must be constructed with at least one column")
    end

    nchunks = length(darrays[1].chunks)
    cs = Array{Any}(undef, nchunks)
    names = isa(tup, NamedTuple) ? [keys(tup)...] : nothing
    f = delayed((cs...) -> table(cs...; names=names, kwargs...))
    for i = 1:nchunks
        cs[i] = f(map(x->x.chunks[i], darrays)...)
    end
    fromchunks(cs)
end

# Copying constructor
function table(t::Union{IndexedTable, DIndexedTable};
               columns=columns(t),
               pkey=t.pkey,
               presorted=false,
               copy=true, kwargs...)

    table(columns;
          pkey=pkey,
          presorted=presorted,
          copy=copy, kwargs...)
end

Base.eltype(dt::DIndexedTable{T}) where {T} = T
function colnames(t::DIndexedTable{T}) where T
    fieldnames(T)
end

function trylength(t)::Nullable{Int}
    len = 0
    for l in map(x->x.nrows, t.domains)
        if !isnull(l)
            len = len + get(l)
        else
            return nothing
        end
    end
    return len
end

function Base.length(t::DIndexedTable)
    l = trylength(t)
    if isnull(l)
        error("The length of the DNDSparse is not yet known since some of its parts are not yet computed. Call `compute` to compute them, and then call `length` on the result of `compute`.")
    else
        get(l)
    end
end

"""
    fromchunks(cs)

Construct a distributed object from chunks. Calls `fromchunks(T, cs)`
where T is the type of the data in the first chunk. Computes any thunks.
"""
function fromchunks(cs::AbstractArray, args...; output=nothing, fnoffset=0, kwargs...)
    if output !== nothing
        if !isdir(output)
            mkdir(output)
        end
        cs = Any[(
            fn = lpad(idx+fnoffset, 5, "0");
            delayed(Dagger.savechunk, get_result=true)(
                c, output, fn);
        ) for (idx, c) in enumerate(cs)]

        vec_thunk = treereduce(delayed(vcat, meta=true), cs)
        if length(cs) == 1
            cs = [compute(get_context(), vec_thunk)]
        else
            cs = compute(get_context(), vec_thunk)
        end
        return fromchunks(cs)
    elseif any(x->isa(x, Thunk), cs)
        vec_thunk = delayed((refs...) -> [refs...]; meta=true)(cs...)
        cs = compute(get_context(), vec_thunk)
    end
    T = chunktype(first(cs))
    fromchunks(T, cs, args...; kwargs...)
end

function fromchunks(::Type{<:AbstractVector},
                    cs::AbstractArray, args...; kwargs...)
    lengths = length.(domain.(cs))
    dmnchunks = DomainBlocks((1,), (cumsum(lengths),))
    T = promote_eltype_chunktypes(cs)
    DArray(T, ArrayDomain(1:sum(lengths)),
           dmnchunks, cs, dvcat)
end

function fromchunks(::Type{<:IndexedTable}, chunks::AbstractArray;
                    domains::AbstractArray = last.(domain.(chunks)),
                    pkey=first(domain(first(chunks))),
                    K = promote_eltypes(domains),
                    T = promote_eltype_chunktypes(chunks))

    nzidxs = findall(!isempty, domains)
    domains = domains[nzidxs]

    DIndexedTable{T, K}(pkey, domains, chunks[nzidxs])
end

function promote_eltypes(ts::AbstractArray)
    nonempty_domains = Iterators.filter(!has_empty_values, ts)
    mapreduce(eltype, _promote_type, nonempty_domains, init = Union{})
end

function promote_eltype_chunktypes(ts::AbstractArray)
    nonempty_chunks = Iterators.filter(!has_empty_values∘domain, ts)
    mapreduce(eltype∘chunktype, _promote_type, nonempty_chunks, init = Union{})
end

"""
    distribute(t::Table, chunks)

Distribute a table in `chunks` pieces. Equivalent to `table(t, chunks=chunks)`.
"""
distribute(t::IndexedTable, chunks) = table(t, chunks=chunks, copy=false)

compute(t::DIndexedTable; kwargs...) = compute(get_context(), t; kwargs...)

function compute(ctx, t::DIndexedTable; output=nothing)
    if any(Dagger.istask, t.chunks)
        fromchunks(IndexedTable, cs, output=output)
    else
        map(Dagger.unrelease, t.chunks) # don't let this be freed
        foreach(Dagger.persist!, t.chunks)
        t
    end
end

distribute(t::DIndexedTable, cs) = table(t, chunks=cs)

collect(t::DIndexedTable) = collect(get_context(), t)

function collect(ctx::Context, dt::DIndexedTable{T}) where T
    cs = dt.chunks
    if length(cs) > 0
        collect(ctx, treereduce(delayed(_merge), cs))
    else
        table(similar(IndexedTables.arrayof(T), 0), pkey=dt.pkey)
    end
end

# merging two tables

function _merge(f, a::IndexedTable, b::IndexedTable)
    if isempty(a.pkey) && isempty(b.pkey)
        return table(vcat(rows(a), rows(b)))
    end

    @assert a.pkey == b.pkey
    ia = pkeys(a)
    ib = pkeys(b)

    if isempty(a)
        b
    elseif isempty(b)
        a
    elseif last(ia) < first(ib)
        # can vcat
        table(map(vcat, columns(a), columns(b)), pkey=a.pkey, presorted=true, copy=false)
    elseif last(ib) < first(ia)
        _merge(f, b, a)
    else
        f(a,b)
    end
end

_merge(f, x::IndexedTable) = x
function _merge(f, x::IndexedTable, y::IndexedTable, ys::IndexedTable...)
    treereduce((a,b)->_merge(f, a, b), [x,y,ys...])
end

_merge(x::IndexedTable, y::IndexedTable...) = _merge((a,b) -> merge(a, b), x, y...)

function Base.show(io::IO, big::DIndexedTable)
    h, w = displaysize(io)
    showrows = h - 5 # This will trigger an ellipsis when there's
                     # more to see than the screen fits
    t = first(Iterators.partition(big, showrows))
    len = trylength(big)
    vals = isnull(len) ? "" : " with $(get(len)) rows"
    header = "Distributed Table$vals in $(length(big.chunks)) chunks:"
    cstyle = Dict(i=>:bold for i in t.pkey)
    showtable(io, t; header=header, ellipsis=:end, cstyle=cstyle)
end


abstract type Setter <: KaleidoLens end

Setfield.get(::Any, setter::Setter) = error("Setters do not support `get`.")

struct NullSetter <: Setter end
Setfield.set(obj, ::NullSetter, ::Any) = obj

"""
    nullsetter :: Setter

A setter that does nothing; i.e., `set(x, nullsetter, y) === x` for any
`x` and `y`.

# Examples
```jldoctest
julia> using Setfield, Kaleido

julia> set(1, nullsetter, 2)
1
```
"""
const nullsetter = NullSetter()

Base.show(io::IO, ::NullSetter) = print(io, "nullsetter")
# TODO: handle prefix

"""
    ToField(f) :: Setter

Apply `f` when setting.  Use `x -> get(x, f)` if `f` is a `Lens`.

# Examples
```jldoctest
julia> using Setfield, Kaleido

julia> setter = (@lens _.x) ∘ ToField(@lens _.a)
(@lens _.x) ∘ (←(@lens _.a)|❌→)

julia> set((x = 1, y = 2), setter, (a = 10, b = 20))
(x = 10, y = 2)
```
"""
struct ToField{F} <: Setter
    f::F
end

_eval(f, x) = f(x)
_eval(f::Lens, x) = get(x, f)

Setfield.set(::Any, setter::ToField, x) = _eval(setter.f, x)


__precompile__()

module MXNet

# we put everything in the namespace mx, because there are a lot of
# functions with the same names as built-in utilities like "zeros", etc.
export mx
module mx

using Compat
import Compat.String
import Compat.view

using Formatting

# Functions from base that we can safely extend and that are defined by libmxnet.
import Base: round, ceil, floor, cos, sin, abs, sign, exp, sqrt, exp, log, norm,
             transpose

include("base.jl")
include("context.jl")
include("util.jl")

include("ndarray.jl")
include("random.jl")

include("name.jl")
include("symbolic-node.jl")
include("executor.jl")

include("metric.jl")
include("optimizer.jl")
include("initializer.jl")

include("io.jl")
include("kvstore.jl")

include("callback.jl")
include("model.jl")

include("visualize.jl")

include("nn-factory.jl")

end # mx

end # module MXNet


#
# Copyright (c) 2021 Tobias Thummerer, Lars Mikelsons, Josef Kircher
# Licensed under the MIT license. See LICENSE file in the project root for details.
#

###############
# Prepare FMU #
###############

cd(dirname(@__FILE__))
pathToFMU = joinpath(pwd(), "../model/IO.fmu")

myFMU = fmiLoad(pathToFMU)

#########################
# Testing Setter/Getter #
#########################

# test reference conversion
@test fmi2String2ValueReference(myFMU, "p_real") == 16777216
@test fmi2String2ValueReference(myFMU, "p_integer") == 16777217
@test fmi2ValueReference2String(myFMU, Cint(16777217)) == ["p_integer"]

# create arrays for in place getter
vR = zeros(Real, 3)
vI = zeros(Integer, 3)
vB = zeros(Bool, 3)
vS = ["test","test"]

c1 = fmiInstantiate!(myFMU; loggingOn=true)

fmiEnterInitializationMode(myFMU)
@test fmiSetReal(myFMU, "p_real", 5.0) == 0
@test fmiSetReal(myFMU, ["u_real", "p_real"], [7.0, 8.0]) == 0

@test fmiSetInteger(myFMU, "p_integer", 6) == 0
@test fmiSetInteger(myFMU, ["u_integer"], [4]) == 0

@test fmiSetBoolean(myFMU, "p_boolean", false) == 0
@test fmiSetBoolean(myFMU, ["u_boolean"], [true]) == 0

@test fmiSetString(myFMU, ["p_string"], ["OldString"]) == 0
@test fmiSetString(myFMU, "p_string", "New String") == 0

fmiGetReal!(myFMU, ["p_real", "y_real", "u_real"], vR)
@test vR == [8.0, 7.0, 7.0]
@test fmiGetReal(myFMU, "p_real") == 8.0

fmiGetInteger!(myFMU, ["u_integer", "y_integer", "p_integer"], vI)
@test vI == [4, 4, 6]
@test fmiGetInteger(myFMU, "y_integer") == 4

fmiGetBoolean!(myFMU, ["u_boolean", "y_boolean", "p_boolean"], vB)
@test vB == [true, true, false]
@test fmiGetBoolean(myFMU, "p_boolean") == false

fmiGetString!(myFMU, ["p_string", "p_string"], vS)
@test vS == ["New String", "New String"]
@test fmiGetString(myFMU, "p_string") == "New String"
fmiExitInitializationMode(myFMU)

############
# Clean up #
############

fmiReset(myFMU)
fmiTerminate(myFMU)
fmiUnload(myFMU)


struct MadNLPExecutionStats{T} <: AbstractExecutionStats
    status::Status
    solution::StrideOneVector{T}
    objective::T
    constraints::StrideOneVector{T}
    dual_feas::T
    primal_feas::T
    multipliers::StrideOneVector{T}
    multipliers_L::StrideOneVector{T}
    multipliers_U::StrideOneVector{T}
    iter::Int
    counters::NLPModelsCounters
    elapsed_time::Real
end

struct InvalidNumberException <: Exception end
struct NotEnoughDegreesOfFreedomException <: Exception end

MadNLPExecutionStats(ips::InteriorPointSolver) =MadNLPExecutionStats(
    ips.status,view(ips.x,1:get_nvar(ips.nlp)),ips.obj_val,ips.c,
    ips.inf_du, ips.inf_pr,
    ips.l,view(ips.zl,1:get_nvar(ips.nlp)),view(ips.zu,1:get_nvar(ips.nlp)),
    ips.cnt.k, ips.nlp.counters,ips.cnt.total_time)
getStatus(result::MadNLPExecutionStats) = STATUS_OUTPUT_DICT[result.status]

# Utilities
has_constraints(ips) = ips.m != 0

# Print functions -----------------------------------------------------------
function print_init(ips::AbstractInteriorPointSolver)
    @notice(ips.logger,@sprintf("Number of nonzeros in constraint Jacobian............: %8i",get_nnzj(ips.nlp)))
    @notice(ips.logger,@sprintf("Number of nonzeros in Lagrangian Hessian.............: %8i\n",get_nnzh(ips.nlp)))

    num_fixed = length(ips.ind_fixed)
    num_var = get_nvar(ips.nlp) - num_fixed
    num_llb_vars = length(ips.ind_llb)
    num_lu_vars = sum((get_lvar(ips.nlp).!=-Inf).*(get_uvar(ips.nlp).!=Inf)) - num_fixed
    num_uub_vars = length(ips.ind_uub)
    num_eq_cons = sum(get_lcon(ips.nlp).==get_ucon(ips.nlp))
    num_ineq_cons = sum(get_lcon(ips.nlp).!=get_ucon(ips.nlp))
    num_ue_cons = sum((get_lcon(ips.nlp).!=get_ucon(ips.nlp)).*(get_lcon(ips.nlp).==-Inf).*(get_ucon(ips.nlp).!=Inf))
    num_le_cons = sum((get_lcon(ips.nlp).!=get_ucon(ips.nlp)).*(get_lcon(ips.nlp).!=-Inf).*(get_ucon(ips.nlp).==Inf))
    num_lu_cons = sum((get_lcon(ips.nlp).!=get_ucon(ips.nlp)).*(get_lcon(ips.nlp).!=-Inf).*(get_ucon(ips.nlp).!=Inf))
    get_nvar(ips.nlp) < num_eq_cons && throw(NotEnoughDegreesOfFreedomException())

    @notice(ips.logger,@sprintf("Total number of variables............................: %8i",num_var))
    @notice(ips.logger,@sprintf("                     variables with only lower bounds: %8i",num_llb_vars))
    @notice(ips.logger,@sprintf("                variables with lower and upper bounds: %8i",num_lu_vars))
    @notice(ips.logger,@sprintf("                     variables with only upper bounds: %8i",num_uub_vars))
    @notice(ips.logger,@sprintf("Total number of equality constraints.................: %8i",num_eq_cons))
    @notice(ips.logger,@sprintf("Total number of inequality constraints...............: %8i",num_ineq_cons))
    @notice(ips.logger,@sprintf("        inequality constraints with only lower bounds: %8i",num_le_cons))
    @notice(ips.logger,@sprintf("   inequality constraints with lower and upper bounds: %8i",num_lu_cons))
    @notice(ips.logger,@sprintf("        inequality constraints with only upper bounds: %8i\n",num_ue_cons))
    return
end

function print_iter(ips::AbstractInteriorPointSolver;is_resto=false)
    mod(ips.cnt.k,10)==0&& @info(ips.logger,@sprintf(
        "iter    objective    inf_pr   inf_du lg(mu)  ||d||  lg(rg) alpha_du alpha_pr  ls"))
    @info(ips.logger,@sprintf(
        "%4i%s% 10.7e %6.2e %6.2e %5.1f %6.2e %s %6.2e %6.2e%s  %i",
        ips.cnt.k,is_resto ? "r" : " ",ips.obj_val/ips.obj_scale[],
        is_resto ? ips.RR.inf_pr_R : ips.inf_pr,
        is_resto ? ips.RR.inf_du_R : ips.inf_du,
        is_resto ? log(10,ips.RR.mu_R) : log(10,ips.mu),
        ips.cnt.k == 0 ? 0. : norm(primal(ips.d),Inf),
        ips.del_w == 0 ? "   - " : @sprintf("%5.1f",log(10,ips.del_w)),
        ips.alpha_z,ips.alpha,ips.ftype,ips.cnt.l))
    return
end

function print_summary_1(ips::AbstractInteriorPointSolver)
    @notice(ips.logger,"")
    @notice(ips.logger,"Number of Iterations....: $(ips.cnt.k)\n")
    @notice(ips.logger,"                                   (scaled)                 (unscaled)")
    @notice(ips.logger,@sprintf("Objective...............:  % 1.16e   % 1.16e",ips.obj_val,ips.obj_val/ips.obj_scale[]))
    @notice(ips.logger,@sprintf("Dual infeasibility......:   %1.16e    %1.16e",ips.inf_du,ips.inf_du/ips.obj_scale[]))
    @notice(ips.logger,@sprintf("Constraint violation....:   %1.16e    %1.16e",norm(ips.c,Inf),ips.inf_pr))
    @notice(ips.logger,@sprintf("Complementarity.........:   %1.16e    %1.16e",
                                ips.inf_compl*ips.obj_scale[],ips.inf_compl))
    @notice(ips.logger,@sprintf("Overall NLP error.......:   %1.16e    %1.16e\n",
                                max(ips.inf_du*ips.obj_scale[],norm(ips.c,Inf),ips.inf_compl),
                                max(ips.inf_du,ips.inf_pr,ips.inf_compl)))
    return
end

function print_summary_2(ips::AbstractInteriorPointSolver)
    ips.cnt.solver_time = ips.cnt.total_time-ips.cnt.linear_solver_time-ips.cnt.eval_function_time
    @notice(ips.logger,"Number of objective function evaluations             = $(ips.cnt.obj_cnt)")
    @notice(ips.logger,"Number of objective gradient evaluations             = $(ips.cnt.obj_grad_cnt)")
    @notice(ips.logger,"Number of constraint evaluations                     = $(ips.cnt.con_cnt)")
    @notice(ips.logger,"Number of constraint Jacobian evaluations            = $(ips.cnt.con_jac_cnt)")
    @notice(ips.logger,"Number of Lagrangian Hessian evaluations             = $(ips.cnt.lag_hess_cnt)")
    @notice(ips.logger,@sprintf("Total wall-clock secs in solver (w/o fun. eval./lin. alg.)  = %6.3f",
                                ips.cnt.solver_time))
    @notice(ips.logger,@sprintf("Total wall-clock secs in linear solver                      = %6.3f",
                                ips.cnt.linear_solver_time))
    @notice(ips.logger,@sprintf("Total wall-clock secs in NLP function evaluations           = %6.3f",
                                ips.cnt.eval_function_time))
    @notice(ips.logger,@sprintf("Total wall-clock secs                                       = %6.3f\n",
                                ips.cnt.total_time))
end

function print_ignored_options(logger,option_dict)
    @warn(logger,"The following options are ignored: ")
    for (key,val) in option_dict
        @warn(logger," - "*string(key))
    end
end
function string(ips::AbstractInteriorPointSolver)
    """
                Interior point solver

                number of variables......................: $(get_nvar(ips.nlp))
                number of constraints....................: $(get_ncon(ips.nlp))
                number of nonzeros in lagrangian hessian.: $(get_nnzh(ips.nlp))
                number of nonzeros in constraint jacobian: $(get_nnzj(ips.nlp))
                status...................................: $(ips.status)
                """
end
print(io::IO,ips::AbstractInteriorPointSolver) = print(io, string(ips))
show(io::IO,ips::AbstractInteriorPointSolver) = print(io,ips)


module FieldModule

using Printf: @printf
using UserExceptionModule: UserError
using CoordsModule: Coords,LatLonCoords,DirectionIndicator,Generic1DCoords
using GridModule: Grid,LatLonGrid, for_all, wrap_coords, coords_in_grid,for_all_with_line_breaks
using GridModule: UnstructuredGrid
using InteractiveUtils: subtypes
using SpecialDirectionCodesModule
using MergeTypesModule
import Base.maximum
import Base.*
import Base./
import Base.+
import Base.==
import Base.isapprox
import Base.fill!
import Base.show
import Base.round
import Base.convert
import Base.sum
import Base.Broadcast
using InteractiveUtils

abstract type Field{T} end

function get(field::Field{T},coords::Coords) where {T}
  throw(UserError())
end

function set!(field::Field,coords::Coords)
  throw(UserError())
end

function fill!(field::Field)
  throw(UserError())
end

function maximum(field::Field)
  throw(UserError())
end

function +(lfield::Field,rfield::Field)
  throw(UserError())
end

function *(lfield::Field,value::T) where {T<:Number}
  throw(UserError())
end

function /(lfield::Field,value::T) where {T<:Number}
  throw(UserError())
end

function ==(lfield::Field,rfield::Field)
  throw(UserError())
end

function sum(field::Field)
  throw(UserError())
end

# This could be done by overloading broadcast but this is overcomplicated
# Assumes result is fractional and hence returns a field of floating point numbers
function elementwise_divide(lfield::Field{T},rfield::Field{T}) where {T}
  return Field{Float64}(lfield.grid,lfield.data./rfield.data)::Field{Float64}
end

# This could be done by overloading broadcast but this is overcomplicated
function equals(lfield::Field,value::T) where {T}
  return Field{Bool}(lfield.grid,Bool.(lfield.data .== value))::Field{Bool}
end

function isapprox(lfield::Field,rfield::Field;
                  rtol::Real=atol>0 ? 0 : sqrt(eps),
                  atol::Real,nans::Bool)
  throw(UserError())
end

function invert(field::Field)
    throw(UserError())
end

function add_offset(field::Field,offset::T,skip_offset_on_values::Vector{T}) where {T}
    throw(UserError())
end

function round!(type::DataType,field::Field{T}) where {T}
  throw(UserError())
end

function add_offset(field::T2,offset::T,skip_offset_on_values::Vector{T}) where {T,T2<:Field{T}}
  for_all(get_grid(field)) do coords::Coords
    if ! (field(coords) in skip_offset_on_values)
      set!(field,coords,field(coords)+offset)
    end
  end
end

function divide(field::T2,divisor::T) where {T,T2<:Field{T}}
  for_all(get_grid(field)) do coords::Coords
    set!(field,coords,field(coords)/divisor)
  end
end

get_grid(obj::T) where {T <: Field} =
  obj.grid::Grid

get_data(obj::T) where {T <: Field} =
  obj.data::Array

function repeat(field::Field{T}, count::Integer) where {T}
  array::Array{Field{T},1} = Array{Field{T},1}(undef,count)
  for i in 1:count
    array[i] = deepcopy(field)
  end
  return array
end

struct LatLonField{T} <: Field{T}
  data::Array{T,2}
  grid::LatLonGrid
  function LatLonField{T}(grid::LatLonGrid,value::T) where {T}
    data::Array{T,2} = value == zero(T) ? zeros(T,grid.nlat,grid.nlon) :
                                          ones(T,grid.nlat,grid.nlon)*value
    return new(data,grid)
  end
  function LatLonField{MergeTypes}(grid::LatLonGrid,value::MergeTypes)
    data::Array{MergeTypes,2} =  zeros(MergeTypes,grid.nlat,grid.nlon)
    fill!(data,value)
    return new(data,grid)
  end
  function LatLonField{T}(grid::LatLonGrid,values::AbstractArray{T,2}) where {T}
    if size(values,1) != grid.nlat ||
       size(values,2) != grid.nlon
       error("Values provided don't match selected grid")
    end
    return new(values,grid)
  end
  function LatLonField{MergeTypes}(grid::LatLonGrid,
                                   integer_field::LatLonField{T}) where {T<:Signed}
    grid = get_grid(integer_field)
    merge_type_field = LatLonField{MergeTypes}(grid,null_mtype)
    for_all(grid) do coords::Coords
      element::MergeTypes = MergeTypes(get(integer_field,coords))
      set!(merge_type_field,coords,element)
    end
    return new(merge_type_field.data,grid)
  end
end

struct UnstructuredField{T} <: Field{T}
  data::Array{T,1}
  grid::UnstructuredGrid
  function UnstructuredField{T}(grid::UnstructuredGrid,value::T) where {T}
    data::Array{T,1} = value == zero(T) ? zeros(T,grid.ncells) :
                                          ones(T,grid.ncells)*value
    return new(data,grid)
  end
  function UnstructuredField{MergeTypes}(grid::UnstructuredGrid,value::MergeTypes)
    data::Array{MergeTypes,1} =  zeros(MergeTypes,grid.ncells)
    fill!(data,value)
    return new(data,grid)
  end
  function UnstructuredField{T}(grid::UnstructuredGrid,values::AbstractArray{T,1}) where {T}
    if size(values,1) != grid.ncells
       error("Values provided don't match selected grid")
    end
    return new(values,grid)
  end
  function UnstructuredField{MergeTypes}(grid::UnstructuredGrid,
                                   integer_field::UnstructuredField{T}) where {T<:Signed}
    grid = get_grid(integer_field)
    merge_type_field = UnstructuredField{MergeTypes}(grid,null_mtype)
    for_all(grid) do coords::Coords
      element::MergeTypes = MergeTypes(get(integer_field,coords))
      set!(merge_type_field,coords,element)
    end
    return new(merge_type_field.data,grid)
  end
end

LatLonField{T}(grid::LatLonGrid) where {T} = LatLonField{T}(grid,zeros(T,grid.nlat,grid.nlon))

UnstructuredField{T}(grid::UnstructuredGrid) where {T} = UnstructuredField{T}(grid,zeros(T,grid.ncells))

Field{T}(grid::LatLonGrid,value::T) where {T} = LatLonField{T}(grid::LatLonGrid,value::T)
Field{T}(grid::LatLonGrid) where {T} = LatLonField{T}(grid::LatLonGrid)

Field{T}(grid::UnstructuredGrid,value::T) where {T} = UnstructuredField{T}(grid::UnstructuredGrid,
                                                                           value::T)
Field{T}(grid::UnstructuredGrid) where {T} = UnstructuredField{T}(grid::UnstructuredGrid)

Field{T}(grid::LatLonGrid,values::AbstractArray{T,2}) where {T} =
  LatLonField{T}(grid::LatLonGrid,values::AbstractArray{T,2})
Field{T}(grid::UnstructuredGrid,values::AbstractArray{T,1}) where {T} =
  UnstructuredField{T}(grid::UnstructuredGrid,
                       values::AbstractArray{T,1})

Field{MergeTypes}(grid::LatLonGrid,integer_field::LatLonField{T}) where {T<:Signed} =
  LatLonField{MergeTypes}(grid::LatLonGrid,integer_field::LatLonField{T})
Field{MergeTypes}(grid::UnstructuredGrid,integer_field::UnstructuredField{T}) where {T<:Signed} =
  UnstructuredField{MergeTypes}(grid::UnstructuredGrid,integer_field::UnstructuredField{T})

convert(::Type{Field{MergeTypes}},x::T2) where {T<:Signed,T2<:Field{T}} = Field{MergeTypes}(get_grid(x),x)

function (latlon_field::LatLonField{T})(coords::LatLonCoords) where {T}
    return latlon_field.data[coords.lat,coords.lon]::T
end

function get(latlon_field::LatLonField{T},coords::LatLonCoords) where {T}
    return latlon_field(coords)::T
end

function set!(latlon_field::LatLonField{T},coords::LatLonCoords,value::T) where {T}
    latlon_field.data[coords.lat,coords.lon] = value
end

function +(lfield::LatLonField{T},rfield::LatLonField{T}) where {T}
  return LatLonField{T}(lfield.grid,lfield.data + rfield.data)
end

function *(lfield::Field,value::T) where {T<:Number}
  return LatLonField{T}(lfield.grid,lfield.data * value)
end

function /(lfield::Field,value::T) where {T<:Number}
  return LatLonField{T}(lfield.grid,lfield.data / value)
end

function ==(lfield::LatLonField{T},rfield::LatLonField{T}) where {T}
  return (lfield.data == rfield.data)::Bool
end

LatLonFieldOrUnstructuredField = Union{LatLonField{T},UnstructuredField{T}} where {T}

function sum(field::T) where {T<:LatLonFieldOrUnstructuredField}
  return sum(field.data)
end

function isapprox(lfield::LatLonField{T},rfield::LatLonField{T};
                  rtol::Real=atol>0 ? 0 : sqrt(eps),
                  atol::Real=0.0,nans::Bool=false) where {T}
  return isapprox(lfield.data,rfield.data,
                 rtol=rtol,atol=atol,nans=nans)::Bool
end

function invert(field::LatLonField{T}) where {T}
  return LatLonField{T}(field.grid,.!field.data)
end

function round(type::DataType,field::LatLonField{T}) where {T}
  return LatLonField{type}(field.grid,round.(type,field.data))
end

function maximum(field::LatLonField)
  return maximum(field.data)
end

function show(io::IO,field::LatLonField{T}) where {T <: Number}
  for_all_with_line_breaks(field.grid) do coords::Coords
    if T <: AbstractFloat
      @printf(io,"%6.2f",field(coords))
    elseif T == Bool
      print(io,field(coords) ? " X " : " - ")
    else
      @printf(io,"%4.0f",field(coords))
    end
    print(io," ")
  end
end

function fill!(field::LatLonField{T},value::T) where {T}
  return fill!(field.data,value)
end

function (unstructured_field::UnstructuredField{T})(coords::Generic1DCoords) where {T}
    return unstructured_field.data[coords.index]::T
end

function get(unstructured_field::UnstructuredField{T},coords::Generic1DCoords) where {T}
    return unstructured_field(coords)::T
end

function set!(unstructured_field::UnstructuredField{T},coords::Generic1DCoords,value::T) where {T}
    unstructured_field.data[coords.index] = value
end

function +(lfield::UnstructuredField{T},rfield::UnstructuredField{T}) where {T}
  return UnstructuredField{T}(lfield.grid,lfield.data + rfield.data)
end

function *(lfield::UnstructuredField,value::T) where {T<:Number}
  return UnstructuredField{T}(lfield.grid,lfield.data * value)
end

function /(lfield::UnstructuredField,value::T) where {T<:Number}
  return UnstructuredField{T}(lfield.grid,lfield.data / value)
end

function ==(lfield::UnstructuredField{T},rfield::UnstructuredField{T}) where {T}
  return (lfield.data == rfield.data)::Bool
end

function isapprox(lfield::UnstructuredField{T},rfield::UnstructuredField{T};
                  rtol::Real=atol>0 ? 0 : sqrt(eps),
                  atol::Real=0.0,nans::Bool=false) where {T}
  return isapprox(lfield.data,rfield.data,
                 rtol=rtol,atol=atol,nans=nans)::Bool
end

function invert(field::UnstructuredField{T}) where {T}
  return UnstructuredField{T}(field.grid,.!field.data)
end

function round(type::DataType,field::UnstructuredField{T}) where {T}
  return UnstructuredField{type}(field.grid,round.(type,field.data))
end

function maximum(field::UnstructuredField)
  return maximum(field.data)
end

function fill!(field::UnstructuredField{T},value::T) where {T}
  return fill!(field.data,value)
end

function show(io::IO,field::UnstructuredField{T}) where {T <: Number}
  for_all_with_line_breaks(field.grid) do coords::Coords
    if T <: AbstractFloat
      @printf(io,"%6.2f",field(coords))
    elseif T == Bool
      print(io,field(coords) ? " X " : " - ")
    else
      @printf(io,"%4.0f",field(coords))
    end
    print(io," ")
  end
end

abstract type DirectionIndicators end

function get(field::DirectionIndicators,coords::Coords)
  throw(UserError())
end

function set!(field::DirectionIndicators,coords::Coords)
  throw(UserError())
end

struct LatLonDirectionIndicators <: DirectionIndicators
  next_cell_coords::Array{LatLonField{Int64},1}
  function LatLonDirectionIndicators(dir_based_rdirs::LatLonField{Int64})
    grid = dir_based_rdirs.grid
    nlat = dir_based_rdirs.grid.nlat
    nlon = dir_based_rdirs.grid.nlon
    lat_indices::LatLonField{Int64} =
      LatLonField{Int64}(grid,zeros(Int64,nlat,nlon))
    lon_indices::LatLonField{Int64} =
      LatLonField{Int64}(grid,zeros(Int64,nlat,nlon))
    for_all(grid) do coords::LatLonCoords
      dir_based_rdir::Int64 = dir_based_rdirs(coords)
      if dir_based_rdir == dir_based_indicator_ocean
        set!(lat_indices,coords,coord_base_indicator_ocean)
        set!(lon_indices,coords,coord_base_indicator_ocean)
      elseif dir_based_rdir == dir_based_indicator_outflow
        set!(lat_indices,coords,coord_base_indicator_outflow)
        set!(lon_indices,coords,coord_base_indicator_outflow)
      elseif dir_based_rdir == dir_based_indicator_truesink
        set!(lat_indices,coords,coord_base_indicator_truesink)
        set!(lon_indices,coords,coord_base_indicator_truesink)
      elseif dir_based_rdir == dir_based_indicator_lake
        set!(lat_indices,coords,coord_base_indicator_lake)
        set!(lon_indices,coords,coord_base_indicator_lake)
      elseif 1 <= dir_based_rdir <= 9
        lat_offset::Int64 = dir_based_rdir in [7,8,9] ? -1 :
                           (dir_based_rdir in [1,2,3] ?  1 : 0)
        lon_offset::Int64 = dir_based_rdir in [7,4,1] ? -1 :
                           (dir_based_rdir in [9,6,3] ?  1 : 0)
        destination_coords::LatLonCoords = LatLonCoords(coords.lat + lat_offset,
                                                        coords.lon + lon_offset)
        destination_coords = wrap_coords(grid,destination_coords)
        if ! coords_in_grid(grid,destination_coords)
          error("Direction based direction indicator points out of grid")
        end
        set!(lat_indices,coords,destination_coords.lat)
        set!(lon_indices,coords,destination_coords.lon)
      else
        error("Invalid direction based direction indicator")
      end
    end
    new(LatLonField{Int64}[lat_indices,lon_indices])
  end
end

struct UnstructuredFieldDirectionIndicators <: DirectionIndicators
  next_cell_coords::UnstructuredField{Int64}
end

function (latlon_direction_indicators::LatLonDirectionIndicators)(coords::LatLonCoords)
  return DirectionIndicator(LatLonCoords(latlon_direction_indicators.next_cell_coords[1](coords),
                                         latlon_direction_indicators.next_cell_coords[2](coords)))
end

function get(latlon_direction_indicators::LatLonDirectionIndicators,coords::LatLonCoords)
  return latlon_direction_indicators(coords)
end

function set!(latlon_direction_indicators::DirectionIndicators,
               coords::LatLonCoords,direction::DirectionIndicator)
  next_cell_coords =  get_next_cell_coords(direction,coords)
  set!(latlon_direction_indicators.next_cell_coords[1],coords,next_cell_coords.lat)
  set!(latlon_direction_indicators.next_cell_coords[2],coords,next_cell_coords.lon)
end

function fill!(field::LatLonField{T},value::T) where {T}
  return fill!(field.data,value)
end

struct UnstructuredDirectionIndicators <: DirectionIndicators
  next_cell_coords::UnstructuredField{Int64}
end

function (unstructured_direction_indicators::UnstructuredDirectionIndicators)(coords::Generic1DCoords)
  return DirectionIndicator(Generic1DCoords(unstructured_direction_indicators.next_cell_coords(coords)))
end

function get(unstructured_direction_indicators::UnstructuredDirectionIndicators,
             coords::Generic1DCoords)
  return unstructured_direction_indicators(coords)
end

function set!(unstructured_direction_indicators::UnstructuredDirectionIndicators,
              coords::Generic1DCoords,direction::DirectionIndicator)
  next_cell_coords =  get_next_cell_coords(direction,coords)
  set!(unstructured_direction_indicators.next_cell_coords,coords,next_cell_coords.index)
end

#Cannot define functors on abstract types in Julia yet. Use this
#workaround - needs to be at end of file as it only acts on subtypes of
#Field/DirectionIndicators that are already defined

for T in subtypes(Field)
 @eval function (field::$T)(coords::Coords) throw(UserError()) end
end

for T in subtypes(DirectionIndicators)
  @eval function (direction_indicators::$T)(coords::Coords) throw(UserError()) end
end

end


export Integer, Node, InputNode, InteriorNode, OutputNode, print_node_tuple, gate_tuple

abstract type Node end
mutable struct Func
    func::Function
    arity::Integer
    name::AbstractString
end
Func(f::Function, a::Integer) = Func(f, a, string(f))

mutable struct InputNode <: Node
    index::Integer
    active::Bool
    cache::MyInt
end

function InputNode(index::Integer)
    return InputNode(index, false, 0 )
end

mutable struct InteriorNode <: Node
    func::Func
    inputs::Vector{Integer}
    active::Bool
    cache::MyInt
end

function InteriorNode(func::Func, inputs::Vector{Int64})
    return InteriorNode(func, inputs, false, MyInt(0) )
end

function InteriorNode(func::Func, inputs::Vector{Integer})
    return InteriorNode(func, inputs, false, MyInt(0) )
end

mutable struct OutputNode <: Node
    input::Integer
end

#
function print_node_tuple( f::IO, node_vect::Vector{InputNode} )
  print(f,"(")
  len = length(node_vect)
  for i in 1:(len-1)
    print(f,node_vect[i].index,",")
  end
  if len > 1
    print(f,node_vect[len].index,")")
  else
    print(f,node_vect[len].index,",)") # add comma so 1-element tuple will be read as a tuple
  end
end

function print_node_tuple( node_vect::Vector{InputNode} )
  print_node_tuple( Base.stdout, node_vect )
end

function print_node_tuple( f::IO, node_vect::Vector{InteriorNode} )
  print(f," (")
  len = length(node_vect)
  for i in 1:(len-1)
    print(f,"(",node_vect[i].func.name,",",node_vect[i].inputs,"),")
  end
  if len > 1
    print(f,"(",node_vect[len].func.name,",",node_vect[len].inputs,"))")
  else
    # add comma so 1-element tuple will be read as a tuple
    print(f,"(",node_vect[len].func.name,",",node_vect[len].inputs,"),)")
  end
end

function print_node_tuple( node_vect::Vector{InteriorNode} )
  print_node_tuple( Base.stdout, node_vect )
end

function print_node_tuple( f::IO, node_vect::Vector{OutputNode} )
  print(f," (")
  len = length(node_vect)
  for i in 1:(len-1)
    print(f,node_vect[i].input,",")
  end
  if len > 1
    print(f,node_vect[len].input,")")
  else
    print(f,node_vect[len].input,",)")
  end
end

function print_node_tuple( node_vect::Vector{OutputNode} )
  print_node_tuple( Base.stdout, node_vect )
end

# Outputs an vector of InteriorNodes to IO stream f in the format needed by circuit() and print_circuit() in Chromosome.jl
function gate_tuple( f::IO, node_vect::Vector{InteriorNode}, numinputs::Int64 )
  print(f," (")
  len = length(node_vect)
  for i in 1:(len-1)
    print(f,"(",numinputs+i,",",node_vect[i].func.name,",",node_vect[i].inputs[1],",",node_vect[i].inputs[2],"), ")
  end
  if len > 1
    print(f,"(",numinputs+len,",",node_vect[len].func.name,",",node_vect[len].inputs[1],",",node_vect[len].inputs[2],"))")
  else
    # add comma so 1-element tuple will be read as a tuple
    print(f,"(",numinputs+len,",",node_vect[len].func.name,",",node_vect[len].inputs[1],",",node_vect[len].inputs[2],"),)")
  end
end

#= moved to Chromosome.jl
function print_build_chromosome( f::IO, c::Chromosome )
  println(f, "build_chromosome(")
  print_node_tuple(f, c.inputs )
  print(f,",")
  print_node_tuple(f, c.interiors )
  print(f,",")
  print_node_tuple(f, c.outputs )
  println(f, ")")
end
=#


include("../../configs.jl")
include("../core/metrics.jl")

function encode_gt_and_get_indices(gt, priors, losses, pos_thold, neg_thold)
    variances = size(priors, 2) == 102300 ? [1, 1] : [0.2, 0.1]
    iou_vals = iou(gt[1:4,:], _to_min_max_form(priors))
    # selecting positive prior boxes
    pos_pairs = findall(iou_vals .>= pos_thold)
    pos_gt = zeros(15, length(pos_pairs))
    gt_idx = getindex.(pos_pairs, [1 2])[:, 2]
    prior_idx = getindex.(pos_pairs, [1 2])[:, 1]

    # enlarging gt matrix to match selected anchors
    pos_gt .= gt[:,gt_idx]
    num_poses = length(prior_idx)
    if num_poses == 0
        # if no positive anchor boxes are found, then no loss will be calculated
        return nothing, nothing, nothing
    end

    #selecting negative prior boxes
    max_prior_vals, max_prior_idx = findmax(iou_vals; dims=2)
    neg_indices = getindex.(findall(max_prior_vals .<= neg_thold), [1 2])[:,1]
    neg_indices = neg_indices[sortperm(losses[neg_indices])]

    neg_cnt = ohem_ratio * num_poses
    neg_cnt = neg_cnt < 1 ? 1 : neg_cnt
    neg_indices = neg_indices[1:neg_cnt]

    # gt bbox conversion
    selected_priors = priors[:,prior_idx]

    pos_gt[1:4,:] = _to_center_length_form(pos_gt[1:4,:])
    pos_gt[3:4,:] = log.(pos_gt[3:4,:] ./ selected_priors[3:4,:]) ./ variances[1]

    pos_gt[1:2,:] = (pos_gt[1:2,:] .- selected_priors[1:2,:]) ./ (variances[2] .* selected_priors[3:4,:])
    pos_gt[5:6,:] = (pos_gt[5:6,:] .- selected_priors[1:2,:]) ./ (variances[2] .* selected_priors[3:4,:])
    pos_gt[7:8,:] = (pos_gt[7:8,:] .- selected_priors[1:2,:]) ./ (variances[2] .* selected_priors[3:4,:])
    pos_gt[9:10,:] = (pos_gt[9:10,:] .- selected_priors[1:2,:]) ./ (variances[2] .* selected_priors[3:4,:])
    pos_gt[11:12,:] = (pos_gt[11:12,:] .- selected_priors[1:2,:]) ./ (variances[2] .* selected_priors[3:4,:])
    pos_gt[13:14,:] = (pos_gt[13:14,:] .- selected_priors[1:2,:]) ./ (variances[2] .* selected_priors[3:4,:])

    # print("Pos_indices: ", prior_idx, "&& Neg_indices: ", neg_indices, "\n")
    return pos_gt, prior_idx, neg_indices
end

function nms(scores, points; thold=0.4)
    x1 = points[1,:]; y1 = points[2,:]; x2 = points[3,:]; y2 = points[4,:];

    areas = (x2 - x1) .* (y2 - y1)
    order = sortperm(scores, rev=true)
    keep = []

    for i in order
        if keep == []
            push!(keep, i)
        else
            iou_vals = iou(points[:,i:i], points[:,keep])
            val, idx = findmax(iou_vals, dims=1)
            if val[1] <= thold
                push!(keep, i)
            end
        end
    end
    return keep
end

"""
Conversion from:
(center_x, center_y, width, height) --> (min_x, min_y, max_x, max_y)
"""
function _to_min_max_form(boxes)
    init_shape = size(boxes)
    boxes = reshape(boxes, (size(boxes)[1], prod(size(boxes)[2:end]))) # converting to 2D
    half_lens = boxes[3:4,:] ./ 2
    return reshape(cat(boxes[1:2,:] .- half_lens, boxes[1:2,:] .+ half_lens, dims=1), init_shape)
end

"""
Conversion from:
(min_x, min_y, max_x, max_y) --> (center_x, center_y, width, height)
"""
function _to_center_length_form(boxes)
    init_shape = size(boxes)
    boxes = reshape(boxes, (size(boxes)[1], prod(size(boxes)[2:end]))) # converting to 2D
    lengths = boxes[3:4,:] .- boxes[1:2,:]
    return reshape(cat(boxes[1:2,:] .+ (lengths ./ 2), lengths, dims=1), init_shape)
end

"""
Returns the anchor boxes with their center_x, center_y, width, height information.
"""
function _get_priorboxes(num_anchors, anchor_info, img_size)
    feature_maps = [Int(ceil(img_size / scale["stride"])) for scale in anchor_info]
    num_proposals = num_anchors * sum([i*i for i in feature_maps])
    anchors = zeros(4, num_proposals)

    counter = 1
    for (idx, f) in enumerate(feature_maps)
        scaler = anchor_info[idx]["stride"]
        for h in 1:f
            cy = (h - 0.5) * scaler
            for w in 1:f
                cx = (w - 0.5) * scaler
                for s in anchor_info[idx]["anchors"]
                    anchors[:,counter] = [cx cy s s]
                    counter += 1
                end
            end
        end
    end
    return anchors
end

"""
Decoder functions:
The main motivation is that the network does not directly predict the proposals but
the combination of prediction and prior anchor box style and rescaling gives the actual
bounding box and landmark coordinations.
"""
function decode_points(bboxes, landmarks, priors)
    if length(size(priors)) == 2
        priors = reshape(priors, (size(priors)..., 1))
    end
    decoded_bboxes = _decode_bboxes(bboxes, priors)
    decoded_landmarks = _decode_landmarks(landmarks, priors)
    return decoded_bboxes, decoded_landmarks
end

function _decode_bboxes(bbox, priors)
    variances = size(priors, 2) > 100000 ? [1, 1] : [0.2, 0.1]
    if length(size(priors)) == 2
        priors = reshape(priors, (size(priors)..., 1))
    end
    centers = priors[1:2,:,:] .+ bbox[1:2,:,:] .* variances[2] .* priors[3:4,:,:]
    lengths = exp.(bbox[3:4,:,:] .* variances[1]) .* priors[3:4,:,:]
    return cat(centers, lengths, dims=1)
end

function _decode_landmarks(landmarks, priors)
    variances = size(priors, 2) > 100000 ? [1, 1] : [0.2, 0.1]
    if length(size(priors)) == 2
        priors = reshape(priors, (size(priors)..., 1))
    end
    lm1 = priors[1:2,:,:] .+ landmarks[1:2,:,:]  .* variances[2] .* priors[3:4,:,:]
    lm2 = priors[1:2,:,:] .+ landmarks[3:4,:,:]  .* variances[2] .* priors[3:4,:,:]
    lm3 = priors[1:2,:,:] .+ landmarks[5:6,:,:]  .* variances[2] .* priors[3:4,:,:]
    lm4 = priors[1:2,:,:] .+ landmarks[7:8,:,:]  .* variances[2] .* priors[3:4,:,:]
    lm5 = priors[1:2,:,:] .+ landmarks[9:10,:,:] .* variances[2] .* priors[3:4,:,:]
    return cat(lm1, lm2, lm3, lm4, lm5, dims=1)
end


using PATH

const PTH_OPTIMIZER = PATH.Optimizer()#output = "no")
# MOI.set(PTH_OPTIMIZER, MOI.Silent(), true)
const PTH_CACHE = MOIU.UniversalFallback(MOIU.Model{Float64}())
const PTH_CACHED = MOIU.CachingOptimizer(PTH_CACHE, PTH_OPTIMIZER)
const PTH_BRIDGED = MOI.Bridges.full_bridge_optimizer(PTH_CACHED, Float64)

push!(solvers_nlp, (opt = PTH_BRIDGED, mode = BilevelJuMP.ComplementMode{Float64}(1e-9)))
# push!(solvers_nlp, (opt = PTH_BRIDGED, mode = BilevelJuMP.ProductWithSlackMode{Float64}(1e-9)))

push!(solvers, (opt = PTH_BRIDGED, mode = BilevelJuMP.ComplementMode{Float64}(1e-9)))
push!(solvers_quad, (opt = PTH_BRIDGED, mode = BilevelJuMP.ComplementMode{Float64}(1e-9)))


using ..JuliaRL.FeatureCreators
import ..build_algorithm_dict
using Flux

"""
    Agent for learning online in the four rooms environment.
"""
mutable struct NNFourRoomsContAgent{O, P<:Resampling.AbstractPolicy, F} <: JuliaRL.AbstractAgent
    μ::P
    ER::ExperienceReplay
    WER::WeightedExperienceReplay
    opt::O
    normalize::F
    gvf::GVF
    training_gap::Int64
    buffer_size::Int64
    batch_size::Int64
    warm_up::Int64
    algo_dict::Dict{String, Resampling.LearningUpdate}
    sample_dict::Dict{String, String}
    value_dict::Dict{String, Flux.Chain}
    counter::Int64
    state_t::Tuple{Array{Float64, 1}, Bool}
    action_t::Int64
    ones::Array{Array{Float32, 1}, 1}
    to_update::Dict{String, Bool}
    max_is::Float32
    function NNFourRoomsContAgent(μ::P, gvf::GVF, opt::O,
                                  training_gap, buffer_size, batch_size,
                                  warm_up, parsed, env_size;
                                  max_is_ratio=1.0,
                                  rng=Random.GLOBAL_RNG,
                                  normalize::F=identity) where {O, P<:Resampling.AbstractPolicy, F}

        # build tile coder....
        algo_dict, sample_dict, value_type_dict = build_algorithm_dict(parsed; max_is=max_is_ratio)
        value_dict = Dict{String, Flux.Chain}()
        to_update = Dict{String, Bool}()
        init_func=(dims...)->Resampling.glorot_uniform(rng, dims...)
        base_network = Chain(
            Dense(2, 10, relu; initW=init_func),
            Dense(10, 10, relu; initW=init_func),
            Dense(10, 1; initW=init_func)
        )
        for key in keys(algo_dict)
            if value_type_dict[key] == "State"
                value_dict[key] = deepcopy(base_network)
            else
                throw("Not Implemented")
            end
            to_update[key] = true
        end
        er_names = [:s_t, :s_tp1, :a_t, :r, :γ_tp1, :terminal, :mu_t, :pi_t, :ρ]
        state_type = Array{Float64, 1}
        er_types = [state_type, state_type, Int64, Array{Float32,1}, Array{Float32,1}, Bool, Float64, Float64, Array{Float32,1}]
        ER = ExperienceReplay(buffer_size, er_types; column_names=er_names)
        WER = WeightedExperienceReplay(buffer_size, er_types; column_names=er_names)
        new{O, P, F}(μ, ER, WER, opt, normalize, gvf, training_gap, buffer_size, batch_size, warm_up,
                     algo_dict, sample_dict, value_dict, warm_up, ([0.0,0.0], false), 0,
                     [ones(Float32, 1) for i in 1:batch_size], to_update,
                     max_is_ratio)
    end
end

JuliaRL.get_action(agent::NNFourRoomsContAgent, s; rng=Random.GLOBAL_RNG) =
    StatsBase.sample(rng, agent.μ, s[1])

function JuliaRL.start!(agent::NNFourRoomsContAgent, env_s_tp1; rng=Random.GLOBAL_RNG, kwargs...)
    agent.counter = agent.warm_up
    agent.state_t = env_s_tp1
    agent.action_t = JuliaRL.get_action(agent, env_s_tp1; rng=rng)
    return agent.action_t
end

function JuliaRL.step!(agent::NNFourRoomsContAgent, env_s_tp1, r, terminal; rng=Random.GLOBAL_RNG, kwargs...)

    # After taking action at timestep t
    μ_prob = get(agent.μ, agent.state_t[1], agent.action_t, nothing, nothing, nothing)
    c, γ, π_prob = get(agent.gvf, agent.state_t, agent.action_t, env_s_tp1, nothing, nothing)

    experience = (agent.normalize(agent.state_t[1]),
                  agent.normalize(env_s_tp1[1]),
                  agent.action_t, [c], [γ], terminal,
                  μ_prob, π_prob, [π_prob/μ_prob])

    add!(agent.ER, experience)
    add!(agent.WER, experience, π_prob/μ_prob)

    if agent.counter == 0
        agent.counter = agent.training_gap
        train_value_functions(agent; rng=rng)
    end
    agent.counter -= 1

    agent.state_t = env_s_tp1
    agent.action_t = JuliaRL.get_action(agent, agent.state_t; rng=rng)
    return agent.action_t
end

function train_value_functions(agent::NNFourRoomsContAgent; rng=Random.GLOBAL_RNG)

    samp_er = sample(agent.ER, agent.batch_size; rng=rng)
    samp_wer = sample(agent.WER, agent.batch_size; rng=rng)

    action_tp1_er = [StatsBase.sample(rng, agent.μ, s) for s in samp_er[:s_tp1]]
    action_tp1_wer = [StatsBase.sample(rng, agent.μ, s) for s in samp_wer[:s_tp1]]
    # action_tp1 = StatsBase.sample(rng, [1, 2], StatsBase.weights(π), batch_size)

    arg_er = (samp_er[:ρ], samp_er[:s_t], samp_er[:s_tp1],
              samp_er[:r], samp_er[:γ_tp1], samp_er[:terminal],
              samp_er[:a_t], action_tp1_er, agent.gvf.policy)

    # println("ER:", samp_er[:s_t])

    arg_wer = (agent.ones, samp_wer[:s_t], samp_wer[:s_tp1],
               samp_wer[:r], samp_wer[:γ_tp1], samp_wer[:terminal],
               samp_er[:a_t], action_tp1_wer, agent.gvf.policy)
    # println("WER:", samp_wer[:s_t])

    # opt = Descent(α)
    for key in filter((k)->agent.to_update[k], keys(agent.algo_dict))
        try
            if agent.sample_dict[key] == "IR"
                corr_term = 1.0
                if key == "BCIR"
                    corr_term = Resampling.total(agent.WER)/size(agent.WER)[1]
                end
                update!(agent.value_dict[key], agent.opt, agent.algo_dict[key], arg_wer...; corr_term=corr_term)
            elseif agent.sample_dict[key] == "ER"
                if key == "WISBuffer"
                    corr_term = Resampling.total(agent.WER)/size(agent.WER)[1]
                    update!(agent.value_dict[key], agent.opt, agent.algo_dict[key], arg_er...; corr_term=1.0/corr_term)
                elseif key == "NormIS"
                    update!(agent.value_dict[key], agent.opt, agent.algo_dict[key], arg_er[1]./agent.max_is, arg_er[2:end]...;)
                else
                    update!(agent.value_dict[key], agent.opt, agent.algo_dict[key], arg_er...;)
                end
            elseif agent.sample_dict[key] == "Optimal"
                samp_opt = agent.ER.buffer
                arg_opt = (samp_opt[:ρ], samp_opt[:s_t], samp_opt[:s_tp1],
                           samp_opt[:r], samp_opt[:γ_tp1], samp_opt[:terminal])
                update!(agent.value_dict[key], opt, agent.algo_dict[key], arg_opt...;)
            end
        catch ex
            println(ex)
            try
                if ex.msg == "Loss is infinite" || ex.msg == "Loss is NaN"
                    agent.to_update[key] = false
                else
                    throw(ex)
                end
            catch
                throw(ex)
            end
        end
    end
end

function predict(agent::NNFourRoomsContAgent, key::AbstractString, eval_states::Array{Array{Float64, 1}, 1})
    # states = [create_features(agent.tilecoder, eval_state./11.0) for eval_state in eval_states]
    return agent.value_dict[key].(states)
end

function predict!(agent::NNFourRoomsContAgent, eval_states::Array{Array{Float64, 1}, 1}, predict_dict::Dict{String, Array{Float64, 1}})
    states = [agent.normalize(eval_state) for eval_state in eval_states]
    if length(keys(predict_dict)) == 0
        for key in keys(agent.value_dict)
            predict_dict[key] = zeros(length(eval_states))
        end
    end
    for key in keys(agent.value_dict)
        predict_dict[key] .= getindex.(Flux.data.(agent.value_dict[key].(states)), 1)
    end
end


function add_vars!(model::Model, obj::Vector, l_in::Bounds, u_in::Bounds)
    nvars = length(obj)
    obj = fvec(obj)
    l = fvecx(copy(l_in), nvars)
    u = fvecx(copy(u_in), nvars)
    for i = 1:nvars
        if l[i] == -Inf
            l[i] = -CPX_INFBOUND
        end
        if u[i] == Inf
            u[i] = CPX_INFBOUND
        end
    end
    if nvars > 0
        stat = @cpx_ccall(newcols, Cint, (
                          Ptr{Void},
                          Ptr{Void},
                          Cint,
                          Ptr{Cdouble},
                          Ptr{Cdouble},
                          Ptr{Cdouble},
                          Ptr{Cchar},
                          Ptr{Ptr{Cchar}}
                          ),
                          model.env.ptr, model.lp, nvars, float(obj), float(l), float(u), C_NULL, C_NULL)
        if stat != 0
            throw(CplexError(model.env, stat))
        end
    end
end

add_var!(model::Model, obj::Vector, l::Vector, u::Vector) = add_vars!(model, obj, l, u)

function add_var!(model::Model, constridx::IVec, constrcoef::FVec, l::FVec, u::FVec, objcoef::FVec)
    nvars = length(objcoef)
    (nvars == length(l) == length(u)) || error("Inconsistent dimensions when adding variables.")
    for i = 1:nvars
        if l[i] == -Inf
            l[i] = -CPX_INFBOUND
        end
        if u[i] == Inf
            u[i] = CPX_INFBOUND
        end
    end
    if nvars > 0
        stat = @cpx_ccall(addcols, Cint, (
                          Ptr{Void},
                          Ptr{Void},
                          Cint,
                          Cint,
                          Ptr{Cdouble},
                          Ptr{Cint},
                          Ptr{Cint},
                          Ptr{Cdouble},
                          Ptr{Cdouble},
                          Ptr{Cdouble},
                          Ptr{Ptr{Cchar}}
                          ),
                          model.env.ptr, model.lp, nvars, length(constridx), objcoef, Cint[0], constridx.-1, constrcoef, l, u, C_NULL)
        if stat != 0
            throw(CplexError(model.env, stat))
        end
    end
end

add_var!(model::Model, obj, l, u) = add_vars!(model, Cdouble[obj], Cdouble[l], Cdouble[u])

function add_var!(model::Model, constridx, constrcoef, l, u, objcoef)
    return add_var!(model,
                    convert(Vector{Cint},   vec(collect(constridx))),
                    convert(Vector{Cdouble},vec(collect(constrcoef))),
                    convert(Vector{Cdouble},vec(collect(l))),
                    convert(Vector{Cdouble},vec(collect(u))),
                    convert(Vector{Cdouble},vec(collect(objcoef))))
end

function add_var!(model::Model, constridx::Vector, constrcoef::Vector, l::Vector, u::Vector, objcoef::Vector)
    return add_var!(model, ivec(constridx), fvec(constrcoef), fvec(l), fvec(u), fvec(objcoef))
end

function get_varLB(model::Model)
    nvars = num_var(model)
    lb = Array(Cdouble, nvars)
    stat = @cpx_ccall(getlb, Cint, (
                      Ptr{Void},
                      Ptr{Void},
                      Ptr{Cdouble},
                      Cint,
                      Cint
                      ),
                      model.env.ptr, model.lp, lb, 0, nvars-1)
    if stat != 0
        throw(CplexError(model.env, stat))
    end
    return lb
end

function set_varLB!(model::Model, l::FVec)
    nvars = num_var(model)
    for i = 1:nvars
        if l[i] == -Inf
            l[i] = -CPX_INFBOUND
        end
    end
    stat = @cpx_ccall(chgbds, Cint, (
                      Ptr{Void},
                      Ptr{Void},
                      Cint,
                      Ptr{Cint},
                      Ptr{Cchar},
                      Ptr{Cdouble}
                      ),
                      model.env.ptr, model.lp, nvars, Cint[0:nvars-1;], fill(convert(Cchar, 'L'), nvars), l)
    if stat != 0
        throw(CplexError(model.env, stat))
    end
end

function get_varUB(model::Model)
    nvars = num_var(model)
    ub = Array(Cdouble, nvars)
    stat = @cpx_ccall(getub, Cint, (
                      Ptr{Void},
                      Ptr{Void},
                      Ptr{Cdouble},
                      Cint,
                      Cint
                      ),
                      model.env.ptr, model.lp, ub, 0, nvars-1)
    if stat != 0
        throw(CplexError(model.env, stat))
    end
    return ub
end

function set_varUB!(model::Model, u::FVec)
    nvars = num_var(model)
        for i = 1:nvars
        if u[i] == Inf
            u[i] = CPX_INFBOUND
        end
    end
    stat = @cpx_ccall(chgbds, Cint, (
                      Ptr{Void},
                      Ptr{Void},
                      Cint,
                      Ptr{Cint},
                      Ptr{Cchar},
                      Ptr{Cdouble}
                      ),
                      model.env.ptr, model.lp, nvars, Cint[0:nvars-1;], fill(convert(Cchar, 'U'), nvars), u)
    if stat != 0
        throw(CplexError(model.env, stat))
    end
end

function set_vartype!(model::Model, vtype::Vector{Char})
    nvars = num_var(model)
    stat = @cpx_ccall(chgctype, Cint, (
                      Ptr{Void},
                      Ptr{Void},
                      Cint,
                      Ptr{Cint},
                      Ptr{Cchar}
                      ),
                      model.env.ptr, model.lp, length(vtype), Cint[0:length(vtype)-1;], convert(Vector{Cchar},vtype))
    if stat != 0
        throw(CplexError(model.env, stat))
    end
    if !isempty(find(!(vtype.=='C')))
        model.has_int = true
    end
end

function get_vartype(model::Model)
    nvars = num_var(model)
    vartypes = Array(Cchar, nvars)
    stat = @cpx_ccall(getctype, Cint, (
                      Ptr{Void},
                      Ptr{Void},
                      Ptr{Cchar},
                      Cint,
                      Cint
                      ),
                      model.env.ptr, model.lp, vartypes, 0, nvars-1)
    if stat != 0
        throw(CplexError(model.env, stat))
    end
    return convert(Vector{Char},vartypes)
end

function num_var(model::Model)
    nvar = @cpx_ccall(getnumcols, Cint, (
                      Ptr{Void},
                      Ptr{Void}
                      ),
                      model.env.ptr, model.lp)
    return(nvar)
end

function set_varname!(model::Model, idx::Integer, name::AbstractString)
    s = bytestring(name)

    stat = @cpx_ccall(chgcolname, Cint, (
                      Ptr{Void},
                      Ptr{Void},
                      Cint,
                      Ptr{Cint},
                      Ptr{Ptr{UInt8}}
                      ),
                      model.env.ptr, model.lp, 1, Cint[idx-1], [pointer(s)])
    if stat != 0
        throw(CplexError(model.env, stat))
    end
end


module IncomeTaxCalculations

import Dates
import Dates: Date, now, TimeType, Year

import Model_Household: Person
import STBParameters: IncomeTaxSys
import TBComponents: TaxResult, calctaxdue, RateBands, delete_thresholds_up_to, *
import Parameters: @with_kw

using Definitions

export calc_income_tax, old_enough_for_mca, apply_allowance, ITResult
export calculate_company_car_charge

@with_kw mutable struct ITResult
    total_tax :: Real = 0.0
    taxable_income :: Real = 0.0
    adjusted_net_income :: Real = 0.0
    total_income :: Real = 0.0
    non_savings :: Real = 0.0
    allowance   :: Real = 0.0
    non_savings_band :: Integer = 0
    savings :: Real = 0.0
    savings_band :: Integer = 0
    dividends :: Real = 0.0
    dividend_band :: Integer = 0
    unused_allowance :: Real = 0.0
    mca :: Real = 0.0
    transferred_allowance :: Real = 0.0
end

## FIXME all these constants should ultimately be parameters
const MCA_DATE = Date(1935,4,6) # fixme make this a parameter

const SAVINGS_INCOME = Incomes_Dict(
    bank_interest => 1.0,
    bonds_and_gilts => 1.0,
    other_investment_income => 1.0
)

const DIVIDEND_INCOME = Incomes_Dict(
    stocks_shares => 1.0
)
const Exempt_Income = Incomes_Dict(
    individual_savings_account=>1.0,
    local_taxes=>1.0,
    free_school_meals => 1.0,
    dlaself_care => 1.0,
    dlamobility => 1.0,
    child_benefit => 1.0,
    pension_credit => 1.0,
    bereavement_allowance_or_widowed_parents_allowance_or_bereavement=> 1.0,
    armed_forces_compensation_scheme => 1.0, # FIXME not in my list check this
    war_widows_or_widowers_pension => 1.0,
    severe_disability_allowance => 1.0,
    attendence_allowance => 1.0,
    industrial_injury_disablement_benefit => 1.0,
    employment_and_support_allowance => 1.0,
    incapacity_benefit => 1.0,## taxable after 29 weeks,
    income_support => 1.0,
    maternity_allowance => 1.0,
    maternity_grant_from_social_fund => 1.0,
    funeral_grant_from_social_fund => 1.0,
    guardians_allowance => 1.0,
    winter_fuel_payments => 1.0,
    dwp_third_party_payments_is_or_pc => 1.0,
    dwp_third_party_payments_jsa_or_esa => 1.0,
    extended_hb => 1.0,
    working_tax_credit => 1.0,
    child_tax_credit => 1.0,
    working_tax_credit_lump_sum => 1.0,
    child_tax_credit_lump_sum => 1.0,
    housing_benefit => 1.0,
    universal_credit => 1.0,
    personal_independence_payment_daily_living => 1.0,
    personal_independence_payment_mobility => 1.0 )

function make_all_taxable()::Incomes_Dict
    eis = union(Set( keys( Exempt_Income )), Definitions.Expenses )
    all_t = Incomes_Dict()
    for i in instances(Incomes_Type)
        if ! (i ∈ eis )
            all_t[i]=1.0
        end
    end
    all_t
end

function calculate_company_car_charge(
    pers :: Person,
    calculator :: Function = guess_car_percentage_2020_21 ) :: Real
    value = max(0.0, pers.company_car_value-
        pers.company_car_contribution )
    if pers.fuel_supplied > 0.0
        value += sys.fuel_imputation
    end
    pct = calculator( pers.company_car_fuel_type )
    value * pct
end

function make_non_savings()::Incomes_Dict
    excl = union(Set(keys(DIVIDEND_INCOME)), Set( keys(SAVINGS_INCOME)))
    nsi = make_all_taxable()
    for i in excl
        delete!( nsi, i )
    end
    nsi
end

const NON_SAVINGS_INCOME = make_non_savings()
const ALL_TAXABLE = make_all_taxable()

"""
Very rough approximation to MCA age - ignores all months since we don't have that in a typical dataset
TODO maybe overload this with age as a Date?
"""
function old_enough_for_mca(
    age            :: Integer,
    model_run_date :: TimeType = now() ) :: Bool
    (model_run_date - Year(age)) < MCA_DATE
end

function calculate_allowance( pers::Person, sys :: IncomeTaxSys ) :: Real
    allowance = sys.personal_allowance
    if pers.registered_blind
        allowance += sys.blind_persons_allowance
    end
    allowance
end

function apply_allowance( allowance::Real, income::Real )::Tuple
    r = max( 0.0, income - allowance )
    allowance = max(0.0, allowance-income)
    allowance,r
end

"""

Complete(??) income tax calculation, based on the Scottish/UK 2019 system.
Mostly taken from Melville (2019) chs 2-4.

FIXME this is too long and needs broken up.

problems:

1. we do this in strict non-savings, savings, dividends order; see 2(11) for examples where it's now advantageous to use a different order
2.

returns a single total tax liabilty, plus multiple intermediate numbers
in the `intermediate` dict

"""
function calc_income_tax(
    pers   :: Person,
    sys    :: IncomeTaxSys,
    intermediate :: Dict,
    spouse_transfer :: Real = 0.0 ) :: ITResult
    itres :: ITResult = ITResult()
    total_income = ALL_TAXABLE*pers.income;
    non_savings = NON_SAVINGS_INCOME*pers.income;
    savings = SAVINGS_INCOME*pers.income;
    dividends = DIVIDEND_INCOME*pers.income;
    allowance = calculate_allowance( pers, sys )
    # allowance reductions goes here

    non_dividends = non_savings + savings

    adjusted_net_income = total_income
    adjusted_net_income += calculate_company_car_charge(pers)
    # ...

    non_savings_tax = TaxResult(0.0, 0)
    savings_tax = TaxResult(0.0, 0)
    dividend_tax = TaxResult(0.0, 0)

    if adjusted_net_income > sys.personal_allowance_income_limit
        allowance =
            max(0.0,
                allowance -
                    sys.personal_allowance_withdrawal_rate*(
                        adjusted_net_income - sys.personal_allowance_income_limit ))
    end
    taxable_income = adjusted_net_income-allowance
    intermediate["allowance"]=allowance
    intermediate["total_income"]=total_income
    intermediate["adjusted_net_income"]=adjusted_net_income
    intermediate["taxable_income"]=taxable_income
    intermediate["savings"]=savings
    intermediate["non_savings"]=non_savings
    intermediate["dividends"]=dividends

    savings_thresholds = deepcopy( sys.savings_thresholds )
    savings_rates = deepcopy( sys.savings_rates )
    # FIXME model all this with parameters
    toprate = size( savings_thresholds )[1]
    if taxable_income > 0
        allowance,non_savings_taxable = apply_allowance( allowance, non_savings )
        non_savings_tax = calctaxdue(
            taxable=non_savings_taxable,
            rates=sys.non_savings_rates,
            thresholds=sys.non_savings_thresholds )

        # horrific savings calculation see Melville Ch2 "Savings Income" & examples 2-3
        # FIXME Move to separate function
        # delete the starting bands up to non_savings taxabke icome
        savings_rates, savings_thresholds = delete_thresholds_up_to(
            rates=savings_rates,
            thresholds=savings_thresholds,
            upto=non_savings_taxable );
        if sys.personal_savings_allowance > 0
            psa = sys.personal_savings_allowance
            # println( "taxable income $taxable_income sys.savings_thresholds[2] $(sys.savings_thresholds[2])")
            if taxable_income > sys.savings_thresholds[toprate]
                psa = 0.0
            elseif taxable_income > sys.savings_thresholds[2] # above the basic rate
                psa *= 0.5 # FIXME parameterise this
            end
            if psa > 0.0 ## if we haven't deleted the zero band already, just widen it
                if savings_rates[1] == 0.0
                    savings_thresholds[1] += psa;
                else ## otherwise, insert a  new one.
                    savings_thresholds = vcat([psa], savings_thresholds )
                    savings_rates = vcat([0.0], savings_rates )
                end
            end
            intermediate["personal_savings_allowance"] = psa
        end # we have a personal_savings_allowance
        intermediate["savings_rates"] = savings_rates
        intermediate["savings_thresholds"] = savings_thresholds
        allowance,savings_taxable = apply_allowance( allowance, savings )
        savings_tax = calctaxdue(
            taxable=savings_taxable,
            rates=savings_rates,
            thresholds=savings_thresholds )

        # Dividends
        # see around example 8-9 ch2
        allowance,dividends_taxable =
            apply_allowance( allowance, dividends )
        dividend_rates=deepcopy(sys.dividend_rates)
        dividend_thresholds=deepcopy(sys.dividend_thresholds )
        # always preserve any bottom zero rate
        add_back_zero_band = false
        zero_band = 0.0
        used_thresholds = non_savings_taxable+savings_taxable
        copy_start = 1
        # handle the zero rate
        if dividend_rates[1] == 0.0
            add_back_zero_band = true
            zero_band = dividend_thresholds[1]
            used_thresholds += min( zero_band, dividends_taxable )
            copy_start = 2
        end
        dividend_rates, dividend_thresholds =
            delete_thresholds_up_to(
                rates=dividend_rates[copy_start:end],
                thresholds=dividend_thresholds[copy_start:end],
                upto=used_thresholds );
        if add_back_zero_band
            dividend_rates = vcat( [0.0], dividend_rates )
            dividend_thresholds .+= zero_band # push all up
            dividend_thresholds = vcat( zero_band, dividend_thresholds )
        end
        intermediate["dividend_rates"]=dividend_rates
        intermediate["dividend_thresholds"]=dividend_thresholds
        intermediate["add_back_zero_band"]=add_back_zero_band
        intermediate["dividends_taxable"]=dividends_taxable

        dividend_tax = calctaxdue(
            taxable=dividends_taxable,
            rates=dividend_rates,
            thresholds=dividend_thresholds )
    else # some allowance left
        allowance = -taxable_income # e.g. allowance - taxable_income
    end
    intermediate["non_savings_tax"]=non_savings_tax.due
    intermediate["savings_tax"]=savings_tax.due
    intermediate["dividend_tax"]=dividend_tax.due

    #
    # tax reducers
    #
    total_tax = non_savings_tax.due+savings_tax.due+dividend_tax.due
    if spouse_transfer > 0
        sp_reduction =
            sys.non_savings_rates[sys.non_savings_basic_rate]*spouse_transfer
        total_tax = max( 0.0, total_tax - sp_reduction )
    end
    itres.total_tax = total_tax
    itres.taxable_income = taxable_income
    itres.allowance = allowance
    itres.total_income = total_income
    itres.adjusted_net_income = adjusted_net_income
    itres.non_savings = non_savings_tax.due
    itres.non_savings_band = non_savings_tax.end_band
    itres.savings = savings_tax.due
    itres.savings_band = savings_tax.end_band
    itres.dividends = dividend_tax.due
    itres.dividend_band = dividend_tax.end_band
    itres.unused_allowance = allowance
    itres
end

function allowed_to_transfer_allowance(
    sys  :: IncomeTaxSys;
    from :: ITResult,
    to   :: ITResult ) :: Bool

   can_transfer :: Bool = true
   if ! (from.unused_allowance > 0.0 &&
         to.unused_allowance <= 0.0)
         # nothing to transfer - this is actually wrong since
         # you can opt to transfer some allowance even if you
         # can technically use it.
       can_transfer = false
   elseif to.savings_band > sys.savings_basic_rate ||
      to.non_savings_band > sys.non_savings_basic_rate ||
      to.dividend_band > sys.dividend_basic_rate
      can_transfer = false
   end
   ## TODO disallow if mca claimed
   can_transfer
end # can_transfer

function calculate_mca( pers :: Person, tax :: ITResult, sys :: IncomeTaxSys)::Real
    ## FIXME parameterise this
    mca = sys.married_couples_allowance
    if tax.adjusted_net_income > sys.mca_income_maximum
        mca = max( sys.mca_minimum, mca -
           (tax.adjusted_net_income-sys.mca_income_maximum)*sys.mca_withdrawal_rate)
    end
    mca * sys.mca_credit_rate
end

function calc_income_tax(
    head   :: Person,
    spouse :: Union{Nothing,Person},
    sys    :: IncomeTaxSys,
    intermediate :: Dict ) :: NamedTuple
    head_intermed = Dict()
    headtax = calc_income_tax( head, sys, head_intermed )
    intermediate["head_tax"] = head_intermed
    spousetax = nothing
    # FIXME the transferable stuff here
    # is not right as you can elect to transfer more than
    # the surplus allowance in some cases.
    # also - add in restrictions on transferring to
    # higher rate payers.
    if spouse != nothing
        spouse_intermed = Dict()
        intermediate["spouse_tax"] = spouse_intermed
        spousetax = calc_income_tax( spouse, sys, spouse_intermed )
        # This is not quite right - you can't claim the
        # MCA AND transfer an allowance. We're assuming
        # always MCA first (I think it's always more valuable?)
        if old_enough_for_mca( head.age ) || old_enough_for_mca( spouse.age )
            # shoud usually just go to the head but.. some stuff about partner
            # with greater income if married after 2005 and you can elect to do this if
            # married before, so:
            if headtax.adjusted_net_income > spousetax.adjusted_net_income
                headtax.mca = calculate_mca( head, headtax, sys )
                headtax.total_tax = max( 0.0, headtax.total_tax - headtax.mca )
            else
                spousetax.mca = calculate_mca( spouse, spousetax, sys )
                spousetax.total_tax = max( 0.0, spousetax.total_tax - spousetax.mca )
            end
        end
        if spousetax.mca == 0.0 == headtax.mca
            if allowed_to_transfer_allowance( sys, from=spousetax, to=headtax )
                transferable_allow = min( spousetax.unused_allowance, sys.marriage_allowance )
                headtax = calc_income_tax( head, sys, head_intermed, transferable_allow )
                intermediate["transfer_spouse_to_head"] = transferable_allow
            elseif allowed_to_transfer_allowance( sys, from=headtax, to=spousetax )
                transferable_allow = min( headtax.unused_allowance, sys.marriage_allowance )
                spousetax = calc_income_tax( spouse, sys, spouse_intermed, transferable_allow )
                intermediate["transfer_head_to_spouse"] = transferable_allow
            end
        end
    end
    ( head=headtax, spouse=spousetax )
end # calc_income_tax

end # module


# Important that oldest is removed first !
# Example showing why this is important:
# For instance, if the lattice is:
#            :Unbounded - ...
#           /
# :Unbounded
#           \
#            :Optimal - ...
# The lower node will give optimality cuts and fill the cutpruner of the root node
# Meanwhile, the upper node will not give any cut so the root not will stay unbounded
# Since the root node is :Unbounded it will not have dual variables for the stats, so
# it may happen (for instance with AvgCutPruner) that all the cuts are at tie !
# When the upper node stops being :Unbounded and gives its cut that is necessary for
# the root node to stop being :Unbounded. However, if the lower node replaces it with
# another optimality cut it will be ignored !
for algo in [AvgCutPruningAlgo(3), DecayCutPruningAlgo(3)]
    @testset "Oldest removed first for $(typeof(algo))" begin
        pruner = CutPruner{2, Int}(algo, :Max)
        addcuts!(pruner, [1 0], [0], [true])
        # add redundants cuts: this cut will not be added to pruner
        addcuts!(pruner, [1 0], [0], [true])
        @test ncuts(pruner) == 2

        # add redundants cuts: the last cut will not be added to pruner
        addcuts!(pruner, [1 1; 1 1], [0, 0], [true, true])
        @test ncuts(pruner) == 3

        addcuts!(pruner, [2 0], [0], [true])
        addcuts!(pruner, [3 0], [0], [true])
        addcuts!(pruner, [4 0], [0], [true])
        @test sort(pruner.A[:,1]) == [2, 3, 4]
        addcuts!(pruner, [5 0], [0], [true])
        @test sort(pruner.A[:,1]) == [3, 4, 5]
        # However, order should only be used to break ties
        addcuts!(pruner, [6 0], [0], [false])
        @test sort(pruner.A[:,1]) == [3, 4, 5]
        addcuts!(pruner, [7 0; 8 0], [0, 0], [true, true])
        @test sort(pruner.A[:,1]) == [5, 7, 8]
    end
end