import Base

"""Base class for a single optimization objective.

All objectives must have a field `initial_state` and a field `generator`, at
minimum.
"""
abstract type AbstractControlObjective end

"""Standard optimization objective.

```julia
Objective(;
    initial_state=<initial_state>,
    generator=<genenerator>,
    target_state=<target_state>
)
```

describes an optimization objective where the time evoluation of the given
`initial_state` under the given `generator` aims towards `target_state`. The
`generator` here is e.g. a time-dependent Hamiltonian or Liouvillian.

The most common control problems in quantum control, e.g. state-to-state
transitions or quantum gate implementations can be expressed by simultaneously
fulfilling multiple objectives of this type.

Note that the objective can only be instantiated via keyword arguments.
"""
struct Objective{ST, GT} <: AbstractControlObjective
    initial_state :: ST
    generator :: GT
    target_state :: ST
    function Objective(;initial_state::ST, generator::GT, target_state::ST) where {ST, GT}
        new{ST, GT}(initial_state, generator, target_state)
    end
end

"""Standard optimization objective with a weight.

```julia
WeightedObjective(;
    initial_state=<initial_state>,
    generator=<genenerator>,
    target_state=<target_state>,
    weight=<weight>
)
```

initializes a control objective like [`Objective`](@ref), but with an
additional `weight` parameter (a float generally between 0 and 1) that weights
the objective relative to other objectives that are part of the same control
problem.
"""
struct WeightedObjective{ST, GT} <: AbstractControlObjective
    initial_state :: ST
    generator :: GT
    target_state :: ST
    weight :: Float64
    function WeightedObjective(;initial_state::ST, generator::GT, target_state::ST, weight::Float64) where {ST, GT}
        new{ST, GT}(initial_state, generator, target_state, weight)
    end
end

"""A full control problem with multiple objectives.

```julia
ControlProblem(
   objectives=<list of objectives>,
   pulse_options=<dict of controls to pulse options>,
   tlist=<time grid>,
   kwargs...
)
```

Note that the control problem can only be instantiated via keyword arguments.

The `objectives` are a list of [`AbstractControlObjective`](@ref) instances,
each defining an initial state and a dynamical generator for the evolution of
that state. Usually, the objective will also include a target state (see
[`Objective`](@ref)) and possibly a weight (see [`WeightedObjective`](@ref)).

The `pulse_options` are a dictionary (`IdDict`) mapping controls that occur in
the `objectives` to properties specific to the control method.

The `tlist` is the time grid on which the time evolution of the initial states
of each objective should be propagated.

The remaining `kwargs` are keyword arguments that are passed directly to the
optimal control method. These typically include e.g. the optimization
functional.

The control problem is solved by finding a set of controls that simultaneously
fulfill all objectives.
"""
struct ControlProblem{OT<:AbstractControlObjective, OPT<:AbstractDict, KWT}
    objectives :: Vector{OT}
    pulse_options :: OPT # TODO: make this optional?
    tlist :: Vector{Float64}
    kwargs :: KWT
    function ControlProblem(;objectives, pulse_options, tlist, kwargs...)
        kwargs_dict = Dict(kwargs)  # make the kwargs mutable
        new{typeof(objectives[1]), typeof(pulse_options), typeof(kwargs_dict)}(
            objectives, pulse_options, tlist, kwargs_dict
        )
    end
end

function Base.copy(problem::ControlProblem)
    return ControlProblem(
        objectives=problem.objectives,
        pulse_options=problem.pulse_options,
        tlist=problem.tlist;
        problem.kwargs...
    )
end

# adjoint for the nested-tuple dynamical generator (e.g. `(H0, (H1, ϵ))`)
function dynamical_generator_adjoint(G::Tuple)
    result = []
    for part in G
        if isa(part, Tuple)
            push!(result, (Base.adjoint(part[1]), part[2]))
        else
            push!(result, Base.adjoint(part))
        end
    end
    return Tuple(result)
end

# fallback adjoint
dynamical_generator_adjoint(G) = Base.adjoint(G)

"""
```julia
adjoint(objective)
```

Adjoint of a control objective. The adjoint objective contains the adjoint of
the dynamical generator `obj.generator`, and adjoints of the
`obj.initial_state` / `obj.target_state` if these exist and have an adjoint.
"""
function Base.adjoint(obj::Objective)
    initial_state_adj = obj.initial_state
    try initial_state_adj = Base.adjoint(obj.initial_state) catch end
    generator_adj = dynamical_generator_adjoint(obj.generator)
    target_adj = obj.target_state
    try target_adj = Base.adjoint(obj.target_state) catch end
    return Objective(
        initial_state=initial_state_adj,
        generator=generator_adj,
        target_state=target_adj
    )
end

function Base.adjoint(obj::WeightedObjective)
    initial_state_adj = obj.initial_state
    try initial_state_adj = Base.adjoint(obj.initial_state) catch end
    generator_adj = dynamical_generator_adjoint(obj.generator)
    target_adj = obj.target_state
    try target_adj = Base.adjoint(obj.target_state) catch end
    return WeightedObjective(
        initial_state=initial_state_adj,
        generator=generator_adj,
        target_state=target_adj,
        weight=obj.weight
    )
end


# new_LP.jl
# New LP formulation for travel time estimation
# Authored by Arthur J Delarue on 7/2/15

MODEL = "relaxed_smooth"		# relaxed/strict _ smooth/nothing/random
LAST_SMOOTH = false 			# true if last iteration is smooth, false otherwise. Irrelevant if model is smooth already.
ERROR_COMPUTATION = "average" 		# average/single-ride/both

MAX_ROUNDS = 2					# max number of iterations

MIN_RIDES = 1					# min number of rides
RADIUS = 140					# radius (4D Euclidian)
YEAR = 2013
START_TIME = 12
END_TIME = 14
START_MONTH = 1
END_MONTH = 1
WEEKDAYS = true 				# true for weekdays, false for weekends
METHOD = "average"				# average/nearest [neighbor]

RANDOM_CONSTRAINTS = false 		# true for purely random constraints, false otw
DYNAMIC_CONSTRAINTS = true 		# true for dynamic constraints, false otw
GLOBAL_CONSTRAINT_UPDATE = false # true if we remove worst paths globally, false otw
SAMPLE_SIZE = 1000 				# starting number of constraints
NUM_OD_ADDED = 1000 			# number of (O,D) pairs to add
NUM_OD_REMOVED = 1000			# number of (O,D) pairs to remove
UPDATE_EVERY_N_ITERATIONS = 1 	# add number of (O,D) above every $N iterations
ADD_IF_ABOVE_PERCENTILE = 0.7		# only add rides that have an error above this percentile
REMOVE_IF_BELOW_PERCENTILE = 0.3	# only remove rides that have an error below this percentile

TURN_COST = 0.	 				# turning cost initial value
TURN_COST_AS_VARIABLE = false 	# true if LP updates turning cost, false otw

DELTA_BOUND = [0.06, 0.05, 0.04]	# relaxation base (sketchy)

MAX_NUM_PATHS_PER_OD = 3 		# at least 1 please

COMPUTE_FINAL_SHORTEST_PATHS = true 	# Should always be true unless trying to save time somehow

START_SIMPLE = false 			# true if initial simple LP is used, false otw

METROPOLIS = false 				# always set this to false unless you wish to crash everything

WARMSTART = false

LOCALIZED_TESTING = false

function fast_LP(
	manhattan::Manhattan,
	travelTimes::Array{Float64,2},
	numRides::Array{Int,2},
	testingData::Union{DataFrame, Array{Float64,2}},
	startTimes::AbstractArray{Float64,2};
	errorComputation::AbstractString=ERROR_COMPUTATION,
	model_type::AbstractString=MODEL,
	last_smooth::Bool=LAST_SMOOTH,
	max_rounds::Int=MAX_ROUNDS,
	min_rides::Int=MIN_RIDES,
	radius::Int=RADIUS,
	year::Int=YEAR,
	startTime::Int=START_TIME,
	endTime::Int=END_TIME,
	startMonth::Int=START_MONTH,
	endMonth::Int=END_MONTH,
	method::AbstractString=METHOD,
	sample_size::Int=SAMPLE_SIZE,
	turnCost::Float64=TURN_COST,
	turnCostAsVariable::Bool=TURN_COST_AS_VARIABLE,
	delta_bound::Union{Array{Float64,1},Array{Any,1}}=DELTA_BOUND,
	maxNumPathsPerOD::Int=MAX_NUM_PATHS_PER_OD,
	computeFinalSP::Bool = COMPUTE_FINAL_SHORTEST_PATHS,
	startWithSimpleLP::Bool=START_SIMPLE,
	randomConstraints::Bool=RANDOM_CONSTRAINTS,
	dynamicConstraints::Bool=DYNAMIC_CONSTRAINTS,
	numPairsToAdd::Int = NUM_OD_ADDED,
	numPairsToRemove::Int = NUM_OD_REMOVED,
	iterationMultiple::Int = UPDATE_EVERY_N_ITERATIONS,
	addOnlyIfAbovePercentile::Float64 = ADD_IF_ABOVE_PERCENTILE,
	removeOnlyIfBelowPercentile::Float64 = REMOVE_IF_BELOW_PERCENTILE,
	globalConstraintUpdate::Bool=GLOBAL_CONSTRAINT_UPDATE,
	warmstart::Bool=WARMSTART,
	metropolis::Bool=METROPOLIS,
	localized_testing::Bool=LOCALIZED_TESTING,
	real_TOD_metropolis::AbstractArray{Float64}=zeros(1,1),
	real_tij_metropolis::AbstractArray{Float64}=zeros(1,1),
	prob::Float64=0.0,
	testingData2::Array{Float64,2}=zeros(1,1))

	graph = manhattan.network
	roadTimes = deepcopy(manhattan.roadTime)
	distances = manhattan.distances
	positions = manhattan.positions
	roads = edges(graph)
	nodes = vertices(graph)
	out = [copy(out_neighbors(graph,i)) for i in nodes]
	inn = [copy(in_neighbors(graph,i)) for i in nodes]

	# Find nonzero travel times
	if DYNAMIC_CONSTRAINTS
		numDataPoints = sample_size
		pairs = [copy(Int[]) for i = 1:length(nodes)]
	else
		pairs = [find(travelTimes[i,:]) for i=nodes]
		numDataPoints = sum([length(pairs[i]) for i=nodes])
	end

	# Create output directory name (sorry this is so complicated)
	if ((!last_smooth) && (!metropolis))
		TESTDIR = "fast_r$(radius)_minr$(min_rides)_i$(max_rounds)_wd_$(startTime)$(endTime)_m$(startMonth)$(endMonth)_y$(year)_$(model_type)_ppc$(maxNumPathsPerOD)_data$(method)_error$errorComputation"
	elseif last_smooth && (!metropolis)
		TESTDIR = "fast_r$(radius)_minr$(min_rides)_i$(max_rounds)_wd_$(startTime)$(endTime)_m$(startMonth)$(endMonth)_y$(year)_$(model_type)_lsmooth_ppc$(maxNumPathsPerOD)_data$(method)_error$errorComputation"
	elseif (!last_smooth) && metropolis
		TESTDIR = "fast_metropolis$(prob)_minr$(min_rides)_i$(max_rounds)_$(model_type)_ppc$(maxNumPathsPerOD)"
	elseif last_smooth && metropolis
		TESTDIR = "fast_metropolis$(prob)_minr$(min_rides)_i$(max_rounds)_$(model_type)_lsmooth_ppc$(maxNumPathsPerOD)"
	end
	if randomConstraints
		TESTDIR=string(TESTDIR, "_rnd_rides$(sample_size)")
	elseif dynamicConstraints
		TESTDIR=string(TESTDIR, "_dynConstr_st$(sample_size)_add$(numPairsToAdd)_thr$(addOnlyIfAbovePercentile)_rem$(numPairsToRemove)_thr$(removeOnlyIfBelowPercentile)_every$(iterationMultiple)")
	end
	if globalConstraintUpdate
		TESTDIR=string(TESTDIR, "_gcu")
	end
	if turnCostAsVariable
		TESTDIR=string(TESTDIR, "_tcvar_start$(turnCost)")
	else
		TESTDIR=string(TESTDIR, "_tc$(turnCost)")
	end
	if startWithSimpleLP
		TESTDIR=string(TESTDIR, "_ss")
	end
	if warmstart
		TESTDIR=string(TESTDIR, "_ws")
	end
	if localized_testing
		TESTDIR=string(TESTDIR, "_loctest")
	end

	# Create directory if necessary:
	if !isdir("Outputs/$TESTDIR")
		mkdir("Outputs/$TESTDIR")
	end
	# Add vital information to file
	writeDataToFile("Outputs/$TESTDIR", model_type, max_rounds, min_rides, radius, false, 0, sample_size, turnCost)
	# Turn cost file
	if turnCostAsVariable
		turnCostFile = open("Outputs/$(TESTDIR)/tc.csv", "w")
	end
	# Metropolis error output
	if metropolis
		# specific error output
	else
		errorFile = open("Outputs/$(TESTDIR)/errorstats.csv", "w")
	end
	costFunctionFile = open("Outputs/$(TESTDIR)/cf.csv", "w")
	# Runtime info output
	timeFile = open("Outputs/$(TESTDIR)/timestats.csv", "w")
	if dynamicConstraints
		numConstraintsFile = open("Outputs/$(TESTDIR)/numODpairs.csv", "w")
	end

	# Tell us where output is going
	println("-- Saving outputs to directory Outputs/$(TESTDIR)/")

	# Create path storing array and initialize. By convention totalPaths[i,1] is the equality constraint
	numPaths = zeros(Int, numDataPoints)
	totalPaths = Array(Any, numDataPoints)
	totalNumExpensiveTurns = Array(Any, numDataPoints)
	for i = 1:numDataPoints
		totalPaths[i] = Array{Int}[]
		totalNumExpensiveTurns[i] = Int[]
		sizehint!(totalPaths[i], maxNumPathsPerOD)
		sizehint!(totalNumExpensiveTurns[i], maxNumPathsPerOD)
	end

	if errorComputation == "single-ride" || errorComputation == "both"
		println("**** Building KDTree of nodes for testing ****")
		nodeTree, nodePairs = buildNodeKDTree(manhattan.positions)
	end

	println("**** Initializing dataset ****")
	if dynamicConstraints
		srcs, dsts, newTravelTimes, newNumRides, pairs = chooseStartingConstraints(travelTimes, numRides, pairs, sample_size = sample_size, num_nodes = length(nodes))
	else
		# Run over all pairs of nodes that have data
		srcs = Int[]
		dsts = Int[]
		sizehint!(srcs, sample_size)
		sizehint!(dsts, sample_size)
		for i in nodes, j in nodes
			if travelTimes[i,j] > 0
				# Load sources and destinations: will be useful when adding constraints
				push!(srcs, i)
				push!(dsts, j)
			end
		end
	end

	status = 0
	newTimes = startTimes
	old_objective = 0

	# Compute shortest paths (with turn cost)
	if startWithSimpleLP
		println("**** Finding initial set of paths ****")
		status, newTimes, turnCost, totalPaths, totalNumExpensiveTurns = simple_LP(manhattan, travelTimes, numRides, testingData, TESTDIR, model_type=split(model_type, "_")[1], max_rounds=max_rounds, turnCost=turnCost, turnCostAsVariable=turnCostAsVariable, delta_bound=delta_bound, maxNumPathsPerOD=maxNumPathsPerOD, computeFinalSP=false)
	else
		println("**** Computing shortest paths ****")
		new_graph, new_edge_dists, new_nodes, new_edge_isExpensive = modifyGraphForDijkstra(graph, newTimes, positions, turn_cost=turnCost)
		old_nodes = getInverseMapping(new_nodes, nv(new_graph))
		@time new_sp = parallelShortestPathsWithTurnsAuto(graph, new_graph, new_edge_dists, new_nodes)
	end

	# l is round number
	l = 1
	while l <= max_rounds
		println("###### ROUND $l ######")
		tic()
		actual_l = l
		# Update bound on relaxation variables
		if split(model_type, "_")[1] == "relaxed"
			delta = (1/2)^(div(actual_l-1,length(delta_bound)))*delta_bound[((actual_l-1)%length(delta_bound)) + 1]
			# Avoid numerical issues
			if delta < 1e-4 || l == max_rounds
				delta = 0
			end
		end

		# Create JuMP model for LP and QP
		println("**** Creating LP instance ****")
		if l == 1 && !startWithSimpleLP && !dynamicConstraints
			tolerance = 1e-4
		else
			tolerance = 1e-6
		end
		m = Model(solver=GurobiSolver(TimeLimit=10000, OutputFlag=1, Method=3, BarConvTol=tolerance))

		# Add one variable for each road, for each model
		if !metropolis
			@defVar(m, t[i=nodes,j=out[i]] >= 0.038 * manhattan.distances[i,j])
		else
			@defVar(m, t[i=nodes,j=out[i]] >= startTimes[i,j])
		end

		# Add decision variable for turn cost if necessary
		if turnCostAsVariable
			@defVar(m, tc >= 0)
		else
			tc = turnCost
		end

		# Epsilon is the variable used for the absolute value difference between T and \hat{T}
		@defVar(m, epsilon[i=nodes,j=pairs[i]] >= 0)

		# Create handles for inequality constraints
		if maxNumPathsPerOD > 1
			@defConstrRef inequalityPath[1:numDataPoints, 1:(max_rounds-1)]
		end

		# Set first LP objective
		@setObjective(m, Min, sum{ sqrt(numRides[i,j]/travelTimes[i,j]) * epsilon[i,j], i=nodes, j=pairs[i]})

		# Add path constraints
		println("**** Adding constraints ****")
		if !(startWithSimpleLP) || l > 1
			println("---Fetching paths")
			paths, numExpensiveTurns = reconstructMultiplePathsWithExpensiveTurns(new_sp.previous, srcs, dsts, old_nodes, new_sp.real_destinations, new_edge_isExpensive)
			println("---Updating path array")
		end
		# Update paths
		if !(startWithSimpleLP) || l > 1
			totalPaths, totalNumExpensiveTurns, numPaths = updatePaths(paths, numExpensiveTurns, totalPaths, totalNumExpensiveTurns, numPaths, maxNumPathsPerOD=maxNumPathsPerOD, times=newTimes, turnCost=turnCost, travelTimes=travelTimes, numRides=numRides, srcs=srcs, dsts=dsts, dynamicConstraints=dynamicConstraints, globalConstraintUpdate=globalConstraintUpdate)
		else
			for i = 1:numDataPoints
				numPaths[i] = length(totalPaths[i])
			end
		end
		for i=1:numDataPoints
			# Add inequality constraints
			if numPaths[i] > 1
				for j = 1:(numPaths[i]-1)
					if split(model_type, "_")[1] == "strict"
						inequalityPath[i,j] = @addConstraint(m, sum{t[totalPaths[i][j+1][a],totalPaths[i][j+1][a+1]], a=1:(length(totalPaths[i][j+1])-1)} - sum{t[totalPaths[i][1][a],totalPaths[i][1][a+1]], a=1:(length(totalPaths[i][1])-1)} >= - tc * totalNumExpensiveTurns[i][j+1] + tc * totalNumExpensiveTurns[i][1])
					elseif split(model_type, "_")[1] == "relaxed"
						inequalityPath[i,j] = @addConstraint(m, sum{t[totalPaths[i][j+1][a],totalPaths[i][j+1][a+1]], a=1:(length(totalPaths[i][j+1])-1)} - sum{t[totalPaths[i][1][a],totalPaths[i][1][a+1]], a=1:(length(totalPaths[i][1])-1)} >= - tc * totalNumExpensiveTurns[i][j+1] + tc * totalNumExpensiveTurns[i][1] - delta * travelTimes[srcs[i],dsts[i]])
					end
				end
			end
		end
		# Equality constraints (shortest path close to travel time data)
		@addConstraint(m, TLowerBound[i=1:numDataPoints], sum{t[totalPaths[i][1][a],totalPaths[i][1][a+1]], a=1:(length(totalPaths[i][1])-1)} + tc * totalNumExpensiveTurns[i][1] - travelTimes[srcs[i],dsts[i]] >= - epsilon[srcs[i],dsts[i]])
		@addConstraint(m, TUpperBound[i=1:numDataPoints], sum{t[totalPaths[i][1][a],totalPaths[i][1][a+1]], a=1:(length(totalPaths[i][1])-1)} + tc * totalNumExpensiveTurns[i][1] - travelTimes[srcs[i],dsts[i]] <= epsilon[srcs[i],dsts[i]])

		# Define objective variables and constraints for second part of m
		if split(model_type, "_")[2] == "smooth" || (last_smooth && l == max_rounds)
			@defVar(m, delta2[i=nodes,j=out[i]] >= 0)
			@addConstraint(m, objConstrLower[i=nodes,j=out[i]], -1 * t[i,j]/distances[i,j] + 1/(length(inn[i]) + length(out[j])) * (sum{1/distances[j,k] * t[j,k], k = out[j]} + sum{1/distances[h,i] * t[h,i], h=inn[i]}) <= delta2[i,j])
			@addConstraint(m, objConstrUpper[i=nodes,j=out[i]], t[i,j]/distances[i,j] - 1/(length(inn[i]) + length(out[j])) * (sum{1/distances[j,k] * t[j,k], k = out[j]} + sum{1/distances[h,i] * t[h,i], h=inn[i]}) <= delta2[i,j])
		elseif split(model_type, "_")[2] == "random"
			coeffVec = 2 * rand(length(roads)) - 1
			coeffs = spzeros(length(nodes), length(nodes))
			k = 1
			for i = nodes, j = out[i]
				coeffs[i,j] = coeffVec[k]
				k += 1
			end
			if turnCostAsVariable
				coeff_tc = 2 * rand() - 1
			end
		end
		# Solve LP
		println("**** Solving LP ****")
		status = solve(m)
		if status == :Infeasible
			buildInternalModel(m)
			print_iis_gurobi(m)
			break
		end
		result = getValue(t)
		objective = getObjectiveValue(m)
		# Update cost function log
		write(costFunctionFile, string(l, ",", objective, "\n"))
		flush(costFunctionFile)
		if turnCostAsVariable
			println("Left turn cost: ", getValue(tc))
		end

		println("**** Setting up second LP ****")
		if split(model_type, "_")[2] == "smooth" || (last_smooth && l == max_rounds)
			if turnCostAsVariable
				@setObjective(m, Min, sum{delta2[i,j], i=nodes, j=out[i]} + tc)
			else
				@setObjective(m, Min, sum{delta2[i,j], i=nodes, j=out[i]})
			end
		elseif split(model_type, "_")[2] == "random"
			# Need to minimize absolute value (otherwise unbounded) so define new variable
			@defVar(m, zAbsVal >= 0)
			# Depending on whether turn cost is a variable, add it (or not) to the objective
			if turnCostAsVariable
				@addConstraint(m, randomObjectiveConstraint1, sum{coeffs[i,j] * t[i,j]/distances[i,j], i=nodes, j=out[i]} + coeff_tc * tc <= zAbsVal)
				@addConstraint(m, randomObjectiveConstraint2, sum{coeffs[i,j] * t[i,j]/distances[i,j], i=nodes, j=out[i]} + coeff_tc * tc >= - zAbsVal)
			else
				@addConstraint(m, randomObjectiveConstraint1, sum{coeffs[i,j] * t[i,j]/distances[i,j], i=nodes, j=out[i]} <= zAbsVal)
				@addConstraint(m, randomObjectiveConstraint2, sum{coeffs[i,j] * t[i,j]/distances[i,j], i=nodes, j=out[i]} >= - zAbsVal)
			end
			@setObjective(m, Min, zAbsVal)
		end

		if split(model_type, "_")[2] == "nothing" && !(last_smooth && l == max_rounds)
			# do nothing
			if turnCostAsVariable
				turnCost = getValue(tc)
				write(turnCostFile, string(l,",", turnCost, "\n"))
			end
		else
			# Solve second LP
			println("**** Solving second LP ****")
			@addConstraint(m, fixObjective, sum{sqrt(numRides[i,j]/travelTimes[i,j]) * epsilon[i,j], i=nodes, j=pairs[i]} <=objective + 1e-1)
			status = solve(m)
			# Get new turning cost
			if turnCostAsVariable
				turnCost = getValue(tc)
				println(turnCost)
				write(turnCostFile, string(l,",", turnCost, "\n"))
			end
		end

		# Debug if infeasible
		if status == :Infeasible
			println("!!!! Computing IIS !!!!")
			buildInternalModel(m)
			print_iis_gurobi(m)
			break
		# Prepare output
		elseif status == :Optimal || status == :Suboptimal || status == :UserLimit
			if status != :UserLimit
				result = getValue(t)
			end
			newTimes = spzeros(length(nodes), length(nodes))
			for key in keys(result)
				newTimes[key[1], key[2]] = result[key[1], key[2]]
			end
			# Save updated Manhattan road times to file
			if metropolis
				saveRoadTimes(newTimes, "$TESTDIR/metropolis-times-$l")
			else
				saveRoadTimes(newTimes, "$TESTDIR/manhattan-times-$l")
			end
			if abs(objective - old_objective)/old_objective < 1e-6
				save("Outputs/$TESTDIR/end.jld", "num_iter", l)
				l = max_rounds
			else
				old_objective = objective
			end
		end
		if l < max_rounds || computeFinalSP
			println("**** Computing shortest paths ****")
			new_graph, new_edge_dists, new_nodes, new_edge_isExpensive = modifyGraphForDijkstra(graph, newTimes, manhattan.positions, turn_cost=turnCost)
			old_nodes = getInverseMapping(new_nodes, nv(new_graph))
			@time new_sp = parallelShortestPathsWithTurnsAuto(graph, new_graph, new_edge_dists, new_nodes)
			if metropolis
				printErrorStatsToFile("Outputs/$TESTDIR/errorstats-$(actual_l).txt", real_TOD_metropolis, travelTimes, new_sp.traveltime, numRides, real_tij_metropolis, newTimes, num_nodes = nv(graph))
			else
				if errorComputation == "single-ride"
					@time avg_sq_err, avg_rel_err, avg_bias = computeTestingError(testingData, nodeTree, nodePairs, new_sp.traveltime, method)
					write(errorFile, string(l, ",", avg_sq_err, ",", avg_rel_err, ",", avg_bias,"\n"))
					flush(errorFile)
				elseif errorComputation == "average"
					@time avg_sq_err, avg_rel_err, avg_bias = computeAverageTestingError(testingData, new_sp.traveltime, num_nodes=nv(graph))
					write(errorFile, string(l, ",", avg_sq_err, ",", avg_rel_err, ",", avg_bias,"\n"))
					flush(errorFile)
				elseif errorComputation == "both"
					@time avg_sq_err, avg_rel_err, avg_bias = computeTestingError(testingData, nodeTree, nodePairs, new_sp.traveltime, method)
					@time avg_sq_err_2, avg_rel_err_2, avg_bias_2 = computeAverageTestingError(testingData2, new_sp.traveltime, num_nodes=nv(graph))
					write(errorFile, string(l, ",", avg_sq_err, ",", avg_rel_err, ",", avg_bias, ",", avg_sq_err_2, ",", avg_rel_err_2, ",", avg_bias_2,"\n"))
					flush(errorFile)
				end
			end
			if dynamicConstraints
				write(numConstraintsFile, string(l,",",length(srcs),",", count_number_edges_used(totalPaths, manhattan),"\n"))
				flush(numConstraintsFile)
				if l < max_rounds && l % iterationMultiple == 0
					println("**** Updating constraint set ****")
					srcs, dsts, totalPaths, totalNumExpensiveTurns, numPaths, pairs = updateConstraints(travelTimes, numRides, new_sp.traveltime, totalPaths, totalNumExpensiveTurns, numPaths, srcs, dsts, pairs, numNodePairsToAdd = numPairsToAdd, numNodePairsToRemove = numPairsToRemove, addOnlyIfAbovePercentile = addOnlyIfAbovePercentile, removeOnlyIfBelowPercentile = removeOnlyIfBelowPercentile)
					numDataPoints = length(srcs)
				end
			end
		end
		write(timeFile, string(l, ",", toc(), "\n"))
		flush(timeFile)
		l += 1
	end
	println("-- Saved outputs to directory Outputs/$(TESTDIR)/")
	# Close log files
	if turnCostAsVariable
		close(turnCostFile)
	end
	if !metropolis
		close(errorFile)
	end
	close(costFunctionFile)
	close(timeFile)
	if dynamicConstraints
		close(numConstraintsFile)
	end
	return status, newTimes, "Outputs/$TESTDIR/"
end

# manhattan = loadCityGraph()
# travel_times, num_rides, trainTestDf, testing_travel_times, testing_num_rides, testDf = loadInputTravelTimes(manhattan.positions, METHOD, year=YEAR, startTime=START_TIME, endTime=END_TIME, startMonth=START_MONTH, endMonth=END_MONTH, radius = RADIUS)
# if RANDOM_CONSTRAINTS
# 	travel_times, num_rides = chooseConstraints(travel_times, num_rides, sample_size=SAMPLE_SIZE);
# end
# if ERROR_COMPUTATION == "single-ride"
# 	@time status, new_times = fast_LP(manhattan, travel_times, num_rides, trainTestDf, manhattan.roadTime)
# elseif ERROR_COMPUTATION == "average"
# 	@time status, new_times = fast_LP(manhattan, travel_times, num_rides, testing_travel_times, manhattan.roadTime)
# elseif ERROR_COMPUTATION == "both"
# 	@time status, new_times = fast_LP(manhattan, travel_times, num_rides, trainTestDf, manhattan.roadTime, testingData2 = testing_travel_times)
# end


# See https://udel.edu/~arthij/MD.pdf for information on forces
# See https://arxiv.org/pdf/1401.1181.pdf for applying forces to atoms
# See OpenMM documentation and Gromacs manual for other aspects of forces

export
    ustrip_vec,
    accelerations,
    force,
    SpecificForce2Atoms,
    SpecificForce3Atoms,
    SpecificForce4Atoms,
    forces

"""
    ustrip_vec(x)

Broadcasted form of `ustrip` from Unitful.jl, allowing e.g. `ustrip_vec.(coords)`.
"""
ustrip_vec(x) = ustrip.(x)

function check_force_units(fdr, force_units)
    if unit(first(fdr)) != force_units
        error("System force units are ", force_units, " but encountered force units ", unit(first(fdr)))
    end
end

"""
    accelerations(system, neighbors=nothing; parallel=true)

Calculate the accelerations of all atoms using the pairwise and specific
interactions and Newton's second law.
If the interactions use neighbor lists, the neighbors should be computed
first and passed to the function.
"""
function accelerations(s, neighbors=nothing; parallel::Bool=true)
    return forces(s, neighbors; parallel=parallel) ./ mass.(s.atoms)
end

"""
    force(inter::PairwiseInteraction, vec_ij, coord_i, coord_j,
          atom_i, atom_j, box_size)
    force(inter::SpecificInteraction, coord_i, coord_j,
          box_size)
    force(inter::SpecificInteraction, coord_i, coord_j,
          coord_k, box_size)
    force(inter::SpecificInteraction, coord_i, coord_j,
          coord_k, coord_l, box_size)

Calculate the force between atoms due to a given interation type.
For `PairwiseInteraction`s returns a single force vector and for
`SpecificInteraction`s returns a type such as `SpecificForce2Atoms`.
Custom pairwise and specific interaction types should implement
this function.
"""
function force end

@inline @inbounds function force!(fs, inter, s::System, i::Integer, j::Integer, force_units, weight_14::Bool=false)
    dr = vector(s.coords[i], s.coords[j], s.box_size)
    if weight_14
        fdr = force(inter, dr, s.coords[i], s.coords[j], s.atoms[i], s.atoms[j], s.box_size, true)
    else
        fdr = force(inter, dr, s.coords[i], s.coords[j], s.atoms[i], s.atoms[j], s.box_size)
    end
    check_force_units(fdr, force_units)
    fdr_ustrip = ustrip.(fdr)
    fs[i] -= fdr_ustrip
    fs[j] += fdr_ustrip
    return nothing
end

@inline @inbounds function force_nounit(inters, coord_i, coord_j, atom_i, atom_j,
                                        box_size, force_units, weight_14::Bool=false)
    dr = vector(coord_i, coord_j, box_size)
    sum(inters) do inter
        if weight_14
            fdr = force(inter, dr, coord_i, coord_j, atom_i, atom_j, box_size, true)
        else
            fdr = force(inter, dr, coord_i, coord_j, atom_i, atom_j, box_size)
        end
        check_force_units(fdr, force_units)
        return ustrip.(fdr)
    end
end

accumulateadd(x) = accumulate(+, x)

# Accumulate values in an array based on the ordered boundaries in bounds
# Used to speed up views with repeated indices on the GPU when you have the bounds
@views @inbounds function accumulate_bounds(arr, bounds)
    zf = zero(arr[1:1])
    accum_pad = vcat(zf, accumulateadd(arr))
    accum_bounds = accum_pad[bounds]
    accum_bounds_offset = vcat(zf, accum_bounds[1:(end - 1)])
    return accum_bounds .- accum_bounds_offset
end

# Uses the Zygote path which gives wrong gradients on the GPU for repeated indices
# Hence only used when we know the indices don't contain repeats
# See https://github.com/FluxML/Zygote.jl/pull/1131
unsafe_getindex(arr, inds) = @view arr[inds]

# Sum forces on neighboring atom pairs to get forces on each atom
# Neighbor forces are accumulated and then atom forces extracted by subtraction
#   at the atom boundaries
# The forces are re-arranged for the other atom index and the process is repeated
# This isn't pretty but it works on the GPU
@views @inbounds function sum_forces(nb_forces, neighbors)
    fs_accum_i = accumulate_bounds(nb_forces, neighbors.atom_bounds_i)
    fs_accum_j = accumulate_bounds(unsafe_getindex(nb_forces, neighbors.sortperm_j), neighbors.atom_bounds_j)
    return fs_accum_j .- fs_accum_i
end

# Functions defined to allow us to write rrules
getindices_i(arr, neighbors) = @view arr[neighbors.nbsi]
getindices_j(arr, neighbors) = @view arr[neighbors.nbsj]

@views function forces_inters(inters, coords, atoms, neighbors, box_size, force_units, weights_14)
    coords_i, atoms_i = getindices_i(coords, neighbors), getindices_i(atoms, neighbors)
    coords_j, atoms_j = getindices_j(coords, neighbors), getindices_j(atoms, neighbors)
    @inbounds nb_forces = force_nounit.((inters,), coords_i, coords_j, atoms_i, atoms_j,
                                        (box_size,), force_units, weights_14)
    return sum_forces(nb_forces, neighbors)
end

"""
    SpecificForce2Atoms(f1, f2)

Forces on two atoms arising from an interaction.
"""
struct SpecificForce2Atoms{D, T}
    f1::SVector{D, T}
    f2::SVector{D, T}
end

"""
    SpecificForce3Atoms(f1, f2, f3)

Forces on three atoms arising from an interaction.
"""
struct SpecificForce3Atoms{D, T}
    f1::SVector{D, T}
    f2::SVector{D, T}
    f3::SVector{D, T}
end

"""
    SpecificForce4Atoms(f1, f2, f3, f4)

Forces on four atoms arising from an interaction.
"""
struct SpecificForce4Atoms{D, T}
    f1::SVector{D, T}
    f2::SVector{D, T}
    f3::SVector{D, T}
    f4::SVector{D, T}
end

function SpecificForce2Atoms(f1::StaticArray{Tuple{D}, T}, f2::StaticArray{Tuple{D}, T}) where {D, T}
    return SpecificForce2Atoms{D, T}(f1, f2)
end

function SpecificForce3Atoms(f1::StaticArray{Tuple{D}, T}, f2::StaticArray{Tuple{D}, T},
                            f3::StaticArray{Tuple{D}, T}) where {D, T}
    return SpecificForce3Atoms{D, T}(f1, f2, f3)
end

function SpecificForce4Atoms(f1::StaticArray{Tuple{D}, T}, f2::StaticArray{Tuple{D}, T},
                            f3::StaticArray{Tuple{D}, T}, f4::StaticArray{Tuple{D}, T}) where {D, T}
    return SpecificForce4Atoms{D, T}(f1, f2, f3, f4)
end

Base.:+(x::SpecificForce2Atoms, y::SpecificForce2Atoms) = SpecificForce2Atoms(x.f1 + y.f1, x.f2 + y.f2)
Base.:+(x::SpecificForce3Atoms, y::SpecificForce3Atoms) = SpecificForce3Atoms(x.f1 + y.f1, x.f2 + y.f2, x.f3 + y.f3)
Base.:+(x::SpecificForce4Atoms, y::SpecificForce4Atoms) = SpecificForce4Atoms(x.f1 + y.f1, x.f2 + y.f2, x.f3 + y.f3, x.f4 + y.f4)

getf1(x) = x.f1
getf2(x) = x.f2
getf3(x) = x.f3
getf4(x) = x.f4

@views function specific_force(inter_list::InteractionList2Atoms, coords, box_size, force_units, n_atoms)
    sparse_fs = Array(force.(inter_list.inters, coords[inter_list.is], coords[inter_list.js], (box_size,)))
    fis, fjs = getf1.(sparse_fs), getf2.(sparse_fs)
    check_force_units(first(first(fis)), force_units)
    return sparsevec(inter_list.is, fis, n_atoms) + sparsevec(inter_list.js, fjs, n_atoms)
end

@views function specific_force(inter_list::InteractionList3Atoms, coords, box_size, force_units, n_atoms)
    sparse_fs = Array(force.(inter_list.inters, coords[inter_list.is], coords[inter_list.js],
                                coords[inter_list.ks], (box_size,)))
    fis, fjs, fks = getf1.(sparse_fs), getf2.(sparse_fs), getf3.(sparse_fs)
    check_force_units(first(first(fis)), force_units)
    return sparsevec(inter_list.is, fis, n_atoms) + sparsevec(inter_list.js, fjs, n_atoms) + sparsevec(inter_list.ks, fks, n_atoms)
end

@views function specific_force(inter_list::InteractionList4Atoms, coords, box_size, force_units, n_atoms)
    sparse_fs = Array(force.(inter_list.inters, coords[inter_list.is], coords[inter_list.js],
                                coords[inter_list.ks], coords[inter_list.ls], (box_size,)))
    fis, fjs, fks, fls = getf1.(sparse_fs), getf2.(sparse_fs), getf3.(sparse_fs), getf4.(sparse_fs)
    check_force_units(first(first(fis)), force_units)
    return sparsevec(inter_list.is, fis, n_atoms) + sparsevec(inter_list.js, fjs, n_atoms) + sparsevec(inter_list.ks, fks, n_atoms) + sparsevec(inter_list.ls, fls, n_atoms)
end

"""
    forces(system, neighbors=nothing; parallel=true)

Calculate the forces on all atoms in the system.
If the interactions use neighbor lists, the neighbors should be computed
first and passed to the function.

    forces(inter, system, neighbors=nothing)

Calculate the forces on all atoms in the system arising from a general
interaction.
If the interaction uses neighbor lists, the neighbors should be computed
first and passed to the function.
Custom general interaction types should implement this function.
"""
function forces(s::System{D, false}, neighbors=nothing; parallel::Bool=true) where D
    n_atoms = length(s)

    if parallel && nthreads() > 1 && n_atoms >= 100
        fs_threads = [ustrip_vec.(zero(s.coords)) for i in 1:nthreads()]

        # Loop over interactions and calculate the acceleration due to each
        for inter in values(s.pairwise_inters)
            if inter.nl_only
                if isnothing(neighbors)
                    error("An interaction uses the neighbor list but neighbors is nothing")
                end
                @threads for ni in 1:neighbors.n
                    i, j, w = neighbors.list[ni]
                    force!(fs_threads[threadid()], inter, s, i, j, s.force_units, w)
                end
            else
                @threads for i in 1:n_atoms
                    for j in 1:(i - 1)
                        force!(fs_threads[threadid()], inter, s, i, j, s.force_units)
                    end
                end
            end
        end

        fs = sum(fs_threads)
    else
        fs = ustrip_vec.(zero(s.coords))

        for inter in values(s.pairwise_inters)
            if inter.nl_only
                if isnothing(neighbors)
                    error("An interaction uses the neighbor list but neighbors is nothing")
                end
                for ni in 1:neighbors.n
                    i, j, w = neighbors.list[ni]
                    force!(fs, inter, s, i, j, s.force_units, w)
                end
            else
                for i in 1:n_atoms
                    for j in 1:(i - 1)
                        force!(fs, inter, s, i, j, s.force_units)
                    end
                end
            end
        end
    end

    for inter_list in values(s.specific_inter_lists)
        sparse_vec = specific_force(inter_list, s.coords, s.box_size, s.force_units, n_atoms)
        fs += ustrip_vec.(Array(sparse_vec))
    end

    for inter in values(s.general_inters)
        # Force type not checked
        fs += ustrip_vec.(forces(inter, s, neighbors))
    end

    return fs * s.force_units
end

function forces(s::System{D, true}, neighbors=nothing; parallel::Bool=true) where D
    fs = ustrip_vec.(zero(s.coords))

    pairwise_inters_nonl = filter(inter -> !inter.nl_only, values(s.pairwise_inters))
    @views if length(pairwise_inters_nonl) > 0
        fs += Zygote.checkpointed(forces_inters, pairwise_inters_nonl, s.coords, s.atoms, neighbors.all,
                                    s.box_size, s.force_units, false)
    end

    pairwise_inters_nl = filter(inter -> inter.nl_only, values(s.pairwise_inters))
    if length(pairwise_inters_nl) > 0 && length(neighbors.close.nbsi) > 0
        fs += Zygote.checkpointed(forces_inters, pairwise_inters_nl, s.coords, s.atoms, neighbors.close,
                                    s.box_size, s.force_units, neighbors.close.weights_14)
    end

    for inter_list in values(s.specific_inter_lists)
        sparse_vec = specific_force(inter_list, s.coords, s.box_size, s.force_units, length(s))
        # Move back to GPU if required
        fs += convert(typeof(fs), ustrip_vec.(Array(sparse_vec)))
    end

    for inter in values(s.general_inters)
        # Force type not checked
        fs += ustrip_vec.(forces(inter, s, neighbors))
    end

    return fs * s.force_units
end


"""

Predict and plot out vaccination impact in a county given initial condition and vaccination parameter inputs
"""

function predict_vacc_county_new(name::String,linelist_data_with_pos_neg,serology_data,
                    M_Kenya_ho,M_Kenya_other,M_Kenya_school,M_Kenya_work,N_kenya;
                    PCR_array,relative_testing_rate,
                    rapid, scenario,
                    prob_sus_w,
                    prob_sus,
                    prob_symp,
                    H_rate_by_age,
                    p_ID,p_IH, all_IFR_fits,ICU_fit,hosp_fit,
                    deaths_data,
                    priors,
                    T_to_full_efficacy,
                    vacc_end_date = Date(2022,6,30), # date vaccination program ends
                    startdate = Date(2020,12,31),
                    enddate = Date(2022,6,30),
                    obs_factor)

    fit_dict = load("modelfits/$(name)_model.jld2")
    model = fit_dict[first(keys(fit_dict))]
    up_county = uppercase(model.areaname)

    #Set times, all times after model_start_time are in reference to the model_start_time
    model_start_time = model.model_start_time # start in 1st Dec 2020
    model_end_time = model.model_end_time   # end of model run for fitting
    marchendpoint = model.marchendpoint # 10th March when roudn 3 serosurvey data ends
    alpha_variant_time = model.alpha_variant_time  # earliest time of alpha variant introduction
    delta_variant_time = model.delta_variant_time # earliest time of delta variant introduction
    janstarttime = model.janstarttime  # time t start including observations in likelihood
    vacc_end_point =  (vacc_end_date - model.startdate).value # End of model run
    pre_vacc_period = (Date(2021,3,6) - model.startdate).value # vaccination starts on 6th March 2021
    model_end_point =  (enddate - startdate).value # End of model run

    # Variants data
    kenya_variants = CSV.File("data/kenya_gisaid_variants.csv") |> DataFrame
    kenya_variants.date = [Date(d,DateFormat("dd/mm/yyyy")) for d in kenya_variants.date]
    kenya_variants.rel_time = [(d - Date(2020,2,20)).value for d in kenya_variants.date]

    idxs_alpha = kenya_variants.variant .== "Alpha"
    idxs_beta = kenya_variants.variant .== "Beta"
    idxs_delta = kenya_variants.variant .== "Delta"

    #Daily prediction of variant prevalence using linear interpolation
    interp_alpha = DataInterpolations.LinearInterpolation(kenya_variants.perc_sequences[idxs_alpha]./100,kenya_variants.rel_time[idxs_alpha])
    alpha_perc = [interp_alpha(t) for t in 1:maximum(kenya_variants.rel_time[idxs_alpha])]
    interp_beta = DataInterpolations.LinearInterpolation(kenya_variants.perc_sequences[idxs_beta]./100,kenya_variants.rel_time[idxs_beta])
    beta_perc = [interp_beta(t) for t in 1:maximum(kenya_variants.rel_time[idxs_beta])]
    interp_delta = DataInterpolations.LinearInterpolation(kenya_variants.perc_sequences[idxs_delta]./100,kenya_variants.rel_time[idxs_delta])
    delta_perc = [interp_delta(t) for t in 1:maximum(kenya_variants.rel_time[idxs_delta])]

    delta_perc = delta_perc[model_start_time:end]  # pick data from 1st Dec 2020
    alpha_perc = alpha_perc[model_start_time:end]
    beta_perc  = beta_perc[model_start_time:end]

    n = model_end_point

    delta_perc = [delta_perc;ones(n-length(delta_perc))]
    alpha_beta_perc = alpha_perc .+ beta_perc
    alpha_beta_perc = [alpha_beta_perc;zeros(n-length(alpha_beta_perc))]
    other_perc = 1.0 .- alpha_beta_perc .- delta_perc

    # fitted parameters for county
    MCMCdraws = get(model.MCMC_results.chain,[:β₀,:β_home,:β_school,:β_other,:β_work,:ϵ,:χ,:p_test,:E₀,:inc_R_αvar,:time_scale_αvar,:mid_point_αvar,:inc_R_δvar,:time_scale_δvar,:mid_point_δvar,:init_scale])
    prior = priors[countynames.==model.areaname][1]

    # Data sets to populate and save after loop
    crit_dis_inc_samples  = zeros(model_end_point,6,2000) # 6 age groups, 2000 samples from HMC
    severe_dis_inc_samples = zeros(model_end_point,6,2000)
    asymp_dis_inc_samples = zeros(model_end_point,6,2000)
    mild_dis_inc_samples = zeros(model_end_point,6,2000)
    death_inc_samples = zeros(model_end_point,6,2000)

    for k in 1:2000
        α= model.α  # 1/mean latent period
        ϵ= MCMCdraws.ϵ[k] #0.5    # Relative infectiousness of asymptomatic cases
        αP= model.αP # rate of movement from Pre symptomatic state
        γA = model.γA #Recovery rate from Asymptomatic state
        γM= model.γM  #Recovery rate from Mild disease

        relative_testing_rate = KenyaCoVaccines.relative_testing_rate
        PCR_array= KenyaCoVaccines.PCR_array #Relative Sensitivity of PCR test on days post-infection
        PCR_sensitivity = model.PCR_sensitivity #Base sensitivity for RT-PCR test
        PCR_specificity = model.PCR_specificity #Base specificity for RT_PCR test
        baseline_sero_array = KenyaCoVaccines.rel_sero_array_26days #Relative sensitivity of serology test on days post-infection
        sero_sensitivity = model.sero_sensitivity#Base sensitivity for serology test
        sero_specificity= model.sero_specificity #Base specificity for serology test
        M_BB= model.M_BB

        ω= model.ω # Waning of natural protection, rate per day
        σω= model.p_sus_w # relative suceptibility after first episode: 0.16 for all ages for now.
        σ = model.p_sus
        δ = model.p_symp
        υ = model.hosp_rate_by_age

        β₀ = MCMCdraws.β₀[k]
        β_home = MCMCdraws.β_home[k]
        β_school = MCMCdraws.β_school[k]
        β_other = MCMCdraws.β_other[k]
        β_work = MCMCdraws.β_work[k]

        #Set variance scalers
        clustering_factor_PCR = 0.5
        M_PCR = 30.
        T = eltype(β₀)

        inc_R_αvar= MCMCdraws.inc_R_αvar[k]
        time_scale_αvar= MCMCdraws.time_scale_αvar[k]
        mid_point_αvar= MCMCdraws.mid_point_αvar[k]
        inc_R_δvar= MCMCdraws.inc_R_δvar[k]
        time_scale_δvar= MCMCdraws.time_scale_δvar[k]
        mid_point_δvar= MCMCdraws.mid_point_δvar[k]
        init_scale = MCMCdraws.init_scale[k]

        p = convert.(T,[β₀,β_home,0.0,β_other,β_work,α,ϵ,αP,γA,γM,ω,inc_R_αvar,time_scale_αvar,alpha_variant_time + mid_point_αvar*30,inc_R_δvar,time_scale_δvar, delta_variant_time + mid_point_δvar*30,
         init_scale])  # start with schools closed
        p = vcat(p,model.VE_acquisition,model.VE_infectiousness,model.VE_severe_disease_risk) # add vaccine efficacies to the paramter vector

        init_sero = model.init_sero
        γV = γM

        N= model.N
        E₀ = MCMCdraws.E₀[k]

        θ = (β₀=β₀,β_home=β_home,β_school=β_school,β_other=β_other,β_work=β_work,E₀=E₀,ϵ=ϵ,init_scale=init_scale)
        u0 = KenyaCoVaccines.create_initial_conditions(θ,model)

        #Set up the callback for lockdown
        schoolsclosure_times = [Float64((Date(2021, 3, 19) - Date(2020, 12, 1)).value),
            Float64((Date(2021, 7, 16) - Date(2020, 12, 1)).value),
            Float64((Date(2021, 10, 1) - Date(2020, 12, 1)).value),
            Float64((Date(2021, 12, 23) - Date(2020, 12, 1)).value)]
        function affect_schoolsclosed!(integrator)
            integrator.p[3] = 0.0 #Schools shut
        end
        schoolsclosed_cb = PresetTimeCallback(schoolsclosure_times, affect_schoolsclosed!,save_positions=(false,false))

        #Set up callback for reopening Schools
        schoolsreopening_times = [Float64((Date(2021, 1, 4) - Date(2020, 12, 1)).value),
            Float64((Date(2021, 5, 10) - Date(2020, 12, 1)).value),
            Float64((Date(2021, 7, 26) - Date(2020, 12, 1)).value),
            Float64((Date(2021, 10, 11) - Date(2020, 12, 1)).value)]
        function affect_schoolsopening!(integrator)
            integrator.p[3] = β_school # Schools open #
        end
        schoolsreopening_cb = PresetTimeCallback(schoolsreopening_times, affect_schoolsopening!,save_positions=(false,false))

        cb = CallbackSet(schoolsclosed_cb, schoolsreopening_cb)

        immune_escape_variant_time = Float64((Date(2022,1,1) - model.startdate).value) #  December 1 2021, time to implement variant immunity escape

        function immune_escape_variant!(integrator)
            # Reduce vaccine efficacies by 50%
            integrator.p[19:27] *= 0.5

            # Start with 'variant_E' exposed number of individuals
            variant_E = 1
            prop_vacc = sum(integrator.u[:,1:8,:],dims=2)[:,1,:] ./ sum(integrator.u[:,1:8,:],dims=[2,3])[:]
            integrator.u[:,2,:] .= variant_E .* prop_vacc # number exposed in each age group and vaccinatin status
            integrator.u[:,1,:] .= integrator.u[:,1,:] .- integrator.u[:,2,:] #substract E from S

            # move 50% from R and W -> S
            loss_susceptible_R = 0.5 .* integrator.u[:,7,:]
            loss_susceptible_W = 0.5 .* integrator.u[:,8,:]
            integrator.u[:,1,:] .+= loss_susceptible_R
            integrator.u[:,7,:] .-= loss_susceptible_R
            integrator.u[:,1,:] .+= loss_susceptible_W
            integrator.u[:,8,:] .-= loss_susceptible_W
        end
        immune_escape_variant_cb = PresetTimeCallback([immune_escape_variant_time],immune_escape_variant!,save_positions=(false,false))

        # Update callbacks based on scenario
        if  scenario==1 || scenario>4  # For Scenario 4-6
         cb = CallbackSet(cb, immune_escape_variant_cb)
        end

        county_model = KenyaCoVaccines.CoVAreaModel(name,KenyaCoVaccines.ll_onegroup_newvariant_infboost,prior;
                                PCR_cases = linelist_data_with_pos_neg,
                                sero_cases = serology_data,
                                average_sero_init = zeros(6),
                                deaths = deaths_data,
                                pop_data= N_kenya,
                                M_county_ho=M_Kenya_ho,
                                M_county_other=M_Kenya_other,
                                M_county_school=M_Kenya_school,
                                M_county_work=M_Kenya_work,
                                rapid = rapid,
                                scenario= scenario,
                                startdate=startdate,
                                enddate=enddate)

        #Solve dynamical system
        sol = solve(county_model.prob, BS3();tspan = (0,model_end_point),callback = cb,u0=u0, p=p,saveat = 1) #isoutofdomain=(u,p,t)->any(x->x<0,u)

        # Number of unscalled (for testing) cases over time by severity
        ι = KenyaCoVaccines.get_incidence_time_array_by_age_dose(sol,10)

        # Number of tests over time by age
        PCR_cases = sum(linelist_data_with_pos_neg.cases[:,linelist_data_with_pos_neg.areas .== up_county,:,:],dims = [2])[:,1,:,:]
        p_vec = [[1,2,3,4],[5,6,7,8,9,10],[11,12],[13,14],[15,16],[17]]
        PCR_cases_collaped =  zeros(Int64,size(PCR_cases,1),6,2)
        for a in 1:6
            PCR_cases_collaped[:,a:a,:] .= sum(PCR_cases[:,p_vec[a][1]:p_vec[a][end],:], dims=2)
        end
        PCR_cases = PCR_cases_collaped
        num_tests =  PCR_cases[model_start_time:end,:,1] .+ PCR_cases[model_start_time:end,:,2]
        num_tests_age = zeros(size(ι,1),6)
        for a in 1:6
            num_tests_age[:,a] = vcat(num_tests[:,a],fill(mean(num_tests[:,a]),(length(ι)-length(num_tests[:,a]))))  # use mean number of teste for forcast periods, could be changed
        end

        # scalling parameters
        p_test = MCMCdraws.p_test[k]
        χ = MCMCdraws.χ[k]
        p_test_scale = 5e-4

        # Apply vaccine efficacy against disease
        v = (1 .- model.VE_severe_disease_risk') .* υ  # efficacy

        l_inc = [inc .* v for inc in ι]
        cases_wt = similar(ι)
        cases_αβ = similar(ι)
        cases_δ  = similar(ι)
        crit_incidence = similar(ι)
        sev_incidence = similar(ι)

        for i in 1:size(ι,1)
            cases_wt[i] = sum(l_inc[i] .* other_perc[i],dims=2) # sum over dose
            cases_αβ[i] = sum(l_inc[i] .* alpha_beta_perc[i], dims=2)
            cases_δ[i]  = sum(l_inc[i] .* delta_perc[i], dims=2)

            PCR_num_wt = cases_wt[i] .* p_test_scale .* p_test
            PCR_denom_wt = PCR_num_wt .+ p_test_scale*((p_test/χ)*(N .- cases_wt[i]))
            p_PCR_pred_wt = PCR_num_wt./PCR_denom_wt
            cases_wt[i]  =  p_PCR_pred_wt .* num_tests_age[i,:]

            PCR_num_αβ = cases_αβ[i] .* p_test_scale .* p_test
            PCR_denom_αβ = PCR_num_αβ .+ p_test_scale*((p_test/χ)*(N .- cases_αβ[i]))
            p_PCR_pred_αβ = PCR_num_αβ./PCR_denom_αβ
            cases_αβ[i]  =  p_PCR_pred_αβ .* num_tests_age[i,:]

            PCR_num_δ = cases_δ[i] .* p_test_scale .* p_test
            PCR_denom_δ = PCR_num_wt .+ p_test_scale*((p_test/χ)*(N .- cases_δ[i]))
            p_PCR_pred_δ = PCR_num_δ./PCR_denom_δ
            cases_δ[i]  =  p_PCR_pred_δ .* num_tests_age[i,:]

            # need to edit structure on ICU and hosp fit to include kth sample fit rathert than point est (single value)
            crit_incidence[i] = ICU_fit[1]*(cases_wt[i] .+ ICU_fit[2].*cases_αβ[i] .+ ICU_fit[3].*cases_δ[i])
            sev_incidence[i] = hosp_fit[1]*(cases_wt[i] .+ hosp_fit[2].*cases_αβ[i] .+ hosp_fit[3].*cases_δ[i])
        end

        crit_incidence =  Matrix(VectorOfArray([inc[:] for inc in crit_incidence])')  # convert to a matrix
        sev_incidence =  Matrix(VectorOfArray([inc[:] for inc in sev_incidence])')

        for a in 1:6 # convolve with PCR array to get right timing of PCR detectable incidence
            crit_incidence[:,a] = KenyaCoVaccines.simple_conv(crit_incidence[:,a],KenyaCoVaccines.PCR_array)
            sev_incidence[:,a] = KenyaCoVaccines.simple_conv(sev_incidence[:,a],KenyaCoVaccines.PCR_array)
        end

        # Saving severe and critical disease incidence
        crit_dis_inc_samples[:,:,k]   .= crit_incidence
        severe_dis_inc_samples[:,:,k] .= sev_incidence

        # Asymptomatic cases
        asymp_dis_inc = [inc .* (ones(3)'.* (1 .- δ )) for inc in ι]
        asymp_dis_inc = Matrix(VectorOfArray([sum(inc,dims=2)[:] for inc in asymp_dis_inc])')  #sum over dose
        asymp_dis_inc_samples[:,:,k] .= asymp_dis_inc

        # Mild disease
        symp_dis_inc  =  [inc .* (ones(3)'.* δ) for inc in ι]
        symp_dis_inc =  Matrix(VectorOfArray([sum(inc,dims=2)[:] for inc in symp_dis_inc])')  #sum over dose
        for a in 1:6
            symp_dis_inc[:,a] =  KenyaCoVaccines.simple_conv(symp_dis_inc[:,a],p_IH)  # convolve with time to hsop function to get right timing of Hospitalisation
        end
        mild_dis_inc = symp_dis_inc .- sev_incidence .- crit_incidence
        mild_dis_inc_samples[:,:,k] .= mild_dis_inc

        # Number of unscalled deaths by age group over time
        IFR = all_IFR_fits[[nm[1] .== name for nm in all_IFR_fits]][1].IFR # EDIT to include HMC samples, so there shoudl be a 'k' index on the IFR
        d = (1 .- VE_severe_disease_risk') .* IFR # efficacy multiplied by risk of death by age, assumes only those with severe disease die

        death_inc =  [inc .* d for inc in ι]
        death_inc =  Matrix(VectorOfArray([sum(inc,dims=2)[:] for inc in death_inc])')  #sum over dose
        for a in 1:6
            death_inc[:,a] =  KenyaCoVaccines.simple_conv(death_inc[:,a],p_ID)  # convolve with time to death function to get right timing of deaths
        end
        death_inc = death_inc ./ obs_factor

    end # end loop over HMC chains

    # pick data from Jan for CEA analysis
    asymp_dis_inc_samples = asymp_dis_inc_samples[janstarttime:end,:,:]
    crit_dis_inc_samples = crit_dis_inc_samples[janstarttime:end,:,:]
    severe_dis_inc_samples = severe_dis_inc_samples[janstarttime:end,:,:]
    mild_dis_inc_samples = mild_dis_inc_samples[janstarttime:end,:,:]
    death_inc_samples = death_inc_samples[janstarttime:end,:,:]

    if rapid  # becasue we now have rapid scenario we save them in different directory: I have created this filing structure in the repository
        rapid_text = "rapid"
    else
        rapid_text = "non_rapid"
    end

    @save("scenario_$(scenario)/$(rapid_text)/mild_cases/severe_cases_$(name).jld2",mild_dis_inc_samples)
    @save("scenario_$(scenario)/$(rapid_text)/severe_cases/severe_cases_$(name).jld2",severe_dis_inc_samples)
    @save("scenario_$(scenario)/$(rapid_text)/critical_cases/critical_cases_$(name).jld2",crit_dis_inc_samples)
    @save("scenario_$(scenario)/$(rapid_text)/asymp_cases/asymp_cases_$(name).jld2",asymp_dis_inc_samples)
    @save("scenario_$(scenario)/$(rapid_text)/deaths/deaths_$(name).jld2",death_inc_samples)

    # Number of vaccines administered over time, this is the same for all 2000 chains
    doses_Kenya_baseline = 0.
    if scenario==1
        doses_Kenya = 0.
    else
        doses_Kenya = (1.25 + 9.76 +4.9) .* 1000000  #Number of doses planned for Kenya over the ynext one year under a 30% cover plan
    end

    age_dist_vacc, doses_daily = KenyaCoVaccines.age_dist([name],doses_Kenya,N_kenya,vacc_end_point,pre_vacc_period,rapid)
    vacc_rate_1,vacc_rate_2 = KenyaCoVaccines.vacc_rate_data(doses_daily,vacc_end_point,pre_vacc_period,age_dist_vacc,N_kenya,name,scenario,rapid)
    num_doses = sum(vacc_rate_1, dims=2)[:] .+ sum(vacc_rate_2, dims=2)[:]
    num_doses = num_doses[janstarttime:end]
    complete_doses_age = vacc_rate_2[janstarttime:end,:]  # number with complete doses

    age_dist_vacc, doses_daily = KenyaCoVaccines.age_dist([name],doses_Kenya_baseline,N_kenya,vacc_end_point,pre_vacc_period,rapid)
    vacc_rate_1,vacc_rate_2 = KenyaCoVaccines.vacc_rate_data(doses_daily,vacc_end_point,pre_vacc_period,age_dist_vacc,N_kenya,name,baseline_scenario,rapid)
    num_doses_baseline = sum(vacc_rate_1, dims=2)[:] .+ sum(vacc_rate_2, dims=2)[:]
    num_doses_baseline = num_doses_baseline[janstarttime:end]
    complete_doses_age_baseline = vacc_rate_2[janstarttime:end,:]

    @save("scenario_$(scenario)/$(rapid_text)/doses/doses_$(name).jld2",num_doses)
    @save("scenario_$(scenario)/$(rapid_text)/doses_age/doses_$(name).jld2",complete_doses_age)

    return nothing
end


# To-do: derive high-fidelity gradient
using DwsimOpt
using JuMP
using Ipopt
using FiniteDifferences
using LinearAlgebra
# using GAMS
# using MadNLP

## SIMULATION OPTIMIZATION PROBLEM DEFINITION
path2sim = "c:/Users/lfsfr/Desktop/DwsimOpt.jl/examples/PRICO/PRICO_composite.dwxmz"
path2dwsim = "C:/Users/lfsfr/AppData/Local/DWSIM7/"

include("c:/Users/lfsfr/Desktop/DwsimOpt.jl/src/SimOpt.jl")
sim_jl = py"sim"

# Declare the simulation optimization problem
py"""
# Import dwsim-python data exchange interface (has to be after Automation2)
from dwsimopt.py2dwsim import create_pddx, assign_pddx

# Assign DoF:
create_pddx( ["MR-1", "CompoundMassFlow", "Nitrogen", "kg/s"],    sim, element="dof" )
create_pddx( ["MR-1", "CompoundMassFlow", "Methane", "kg/s"],     sim, element="dof" )
create_pddx( ["MR-1", "CompoundMassFlow", "Ethane", "kg/s"],      sim, element="dof" )
create_pddx( ["MR-1", "CompoundMassFlow", "Propane", "kg/s"],     sim, element="dof" )
create_pddx( ["VALV-01", "OutletPressure", "Mixture", "Pa"],      sim, element="dof" )
create_pddx( ["COMP-4", "OutletPressure", "Mixture", "Pa"],       sim, element="dof" )

# Assign F
create_pddx( ["Sum_W", "EnergyFlow", "Mixture", "kW"], sim, element="fobj" )

# adding constraints (g_i <= 0):
g1 = create_pddx( ["MITA1-Calc", "OutputVariable", "mita", "°C"], sim, element="constraint", assign=False )
r = 1.0/3
assign_pddx( lambda: r*(3-g1[0]()) , ["MITA1-Calc", "OutputVariable", "mita", "°C"], sim, element="constraint" )
create_pddx( ["MSTR-27", "MassFraction", "Liquid", "x"], sim, element="constraint" )
create_pddx( ["MR-1", "MassFraction", "Liquid", "x"], sim, element="constraint" )
create_pddx( ["MSTR-03", "MassFraction", "Liquid", "x"], sim, element="constraint" )
create_pddx( ["MSTR-05", "MassFraction", "Liquid", "x"], sim, element="constraint" )

# decision variables bounds
# x0 = np.array( [1.07249112e-04, 1.93057244e-04, 2.16881626e-04, 7.44968925e-04, 3.43386624e+05, 6.39810411e+06] )
# x0 = np.array( [1.12083333e-04, 1.30420113e-04, 2.57333445e-04, 8.00266697e-04, 3.45000000e+05, 1.53000000e+02] )
x0 = np.array( [1.12083333e-04, 2.20416667e-04, 1.42216680e-04, 7.54165016e-04, 3.45000000e+05, 7.20000000e+06] )
x0 = np.array( [1.3*1.12083333e-04, 1.3*2.20416667e-04, 1.3*1.42216680e-04, 1.3*7.54165016e-04, 3.45000000e+05, 7.20000000e+06] )
bounds_raw = np.array( [0.5*np.asarray(x0), 1.5*np.asarray(x0)] )   # 50 % around base case

# regularizer calculation
regularizer = np.zeros(x0.size)
regularizer_b = np.zeros(x0.size)
for i in range(len(regularizer)):
    regularizer[i] = 10**(-1*math.floor(math.log(x0[i],10))) # regularizer for magnitude order of 1e0

# bounds regularized
bounds_reg = regularizer*bounds_raw + regularizer_b

# objective and constraints lambda definitions
f = lambda x: sim.calculate_optProblem(np.asarray(x)/regularizer)[0:sim.n_f]
g = lambda x: sim.calculate_optProblem(np.asarray(x)/regularizer)[sim.n_f:(sim.n_f+sim.n_g)]
"""

f = py"f"
g = py"g"
x0 = py"x0*regularizer"
searchSpace = py"bounds_reg"
dim = sim_jl.n_dof
op1 = optProblem(f, g, x0, searchSpace, dim, sim_jl, py"sim.n_dof", py"sim.n_f", py"sim.n_g")
op1.sim_jl.verbose = false;
save_sim() = py"""sim.interface.SaveFlowsheet(sim.flowsheet,$pwd()+"/examples/PRICO/PRICO_composite2.dwxmz",True)"""

# external functions and jacobian calculations
x0 = 1.0 * (op1.searchSpace[1, :] + op1.searchSpace[2, :]) ./ 2
# x0 = py"x0*regularizer"
iter = 0
function f_block(x)
    global iter
    iter += 1
    return [op1.f(x)[1]; op1.g(x)[:]]
end
first_time = true
h = 0.01
flag_diff_method = "2nd_central"
flag_FiniteDifference = false
function jac_block(x)
    # println("hi")
    global x0, ∇f_bkp, ∇g1_bkp, ∇g2_bkp, ∇g3_bkp, ∇g4_bkp, ∇g5_bkp, first_time
    if norm(x0 - x) > 1e-6 || first_time == true
        f0 = f_block(x)
        println("Calculating jacobian of black-box functions at x=", x, " and f^{sim}=", f0)
        # flag_diff_method = "1st_forward"
        global h, flag_FiniteDifference, flag_diff_method
        tmp = zeros(6, length(x))
        if flag_FiniteDifference
            tmp = jacobian(central_fdm(2, 1), f_block, x)[1]
        else
            for i = 1:length(x)
                if flag_diff_method == "1st_forward"
                    tmp[:, i] = (f_block([x[1:i-1]; x[i] + h; x[i+1:end]]) - f0) / h
                elseif flag_diff_method == "2nd_central"
                    tmp[:, i] = (f_block([x[1:i-1]; x[i] + h; x[i+1:end]]) - f_block([x[1:i-1]; x[i] - h; x[i+1:end]])) / (2 * h)
                end
            end
        end
        println("∇f_block=", tmp)
        # tmp = jacobian(forward_fdm(2, 1), f_block, x)[1]
        ∇f_bkp = tmp[1, :]
        ∇g1_bkp = tmp[2, :]
        ∇g2_bkp = tmp[3, :]
        ∇g3_bkp = tmp[4, :]
        ∇g4_bkp = tmp[5, :]
        ∇g5_bkp = tmp[6, :]
        x0 = x
        first_time = false
        return tmp
    end
end

function setup_model()
    ## JuMP MODEL STUFF
    model = Model(optimizer_with_attributes(Ipopt.Optimizer))
    # model = Model(Ipopt.Optimizer)
    # model = Model(GAMS.Optimizer)
    # set_optimizer_attribute(model, GAMS.ModelType(), "NLP")
    # set_optimizer_attribute(model, "NLP", "CONOPT")
    # set_optimizer_attribute(model, "resLIM", 1200.0)
    # model = Model(() -> MadNLP.Optimizer(print_level = MadNLP.INFO, max_iter = 100))

    # VARIABLES
    i = 1
    @variable(model, op1.searchSpace[1, i]' <= x1 <= op1.searchSpace[2, i]', start = (op1.searchSpace[1, i] + op1.searchSpace[2, i]) ./ 2)
    i = 2
    @variable(model, op1.searchSpace[1, i]' <= x2 <= op1.searchSpace[2, i]', start = (op1.searchSpace[1, i] + op1.searchSpace[2, i]) ./ 2)
    i = 3
    @variable(model, op1.searchSpace[1, i]' <= x3 <= op1.searchSpace[2, i]', start = (op1.searchSpace[1, i] + op1.searchSpace[2, i]) ./ 2)
    i = 4
    @variable(model, op1.searchSpace[1, i]' <= x4 <= op1.searchSpace[2, i]', start = (op1.searchSpace[1, i] + op1.searchSpace[2, i]) ./ 2)
    i = 5
    @variable(model, op1.searchSpace[1, i]' <= x5 <= op1.searchSpace[2, i]', start = (op1.searchSpace[1, i] + op1.searchSpace[2, i]) ./ 2)
    i = 6
    @variable(model, op1.searchSpace[1, i]' <= x6 <= op1.searchSpace[2, i]', start = (op1.searchSpace[1, i] + op1.searchSpace[2, i]) ./ 2)

    # OBJECTIVE FUNCTION
    fobj_bb(x1, x2, x3, x4, x5, x6) = op1.f([x1, x2, x3, x4, x5, x6])[1]
    function ∇fobj_bb(g::AbstractVector{T}, x1::T, x2::T, x3::T, x4::T, x5::T, x6::T) where {T}
        # println("I am in ∇fobj_bb with x=", x1)
        global x0, ∇f_bkp, ∇g1_bkp, ∇g2_bkp, ∇g3_bkp, ∇g4_bkp, ∇g5_bkp
        jac_block([x1, x2, x3, x4, x5, x6])
        ∇f_tmp = ∇f_bkp
        # ∇f_tmp = grad(forward_fdm(2, 1), op1.f, [x1, x2, x3, x4, x5, x6, x7, x8])[1]
        for i = 1:length([x1, x2, x3, x4, x5, x6])
            g[i] = ∇f_tmp[i]
        end
    end
    register(model, :fobj_bb, op1.dim, fobj_bb, ∇fobj_bb; autodiff = false)
    @NLobjective(model, Min, fobj_bb(x1, x2, x3, x4, x5, x6))

    # CONSTRAINTS
    g1(x) = op1.g(x)[1]
    g2(x) = op1.g(x)[2]
    g3(x) = op1.g(x)[3]
    g4(x) = op1.g(x)[4]
    g5(x) = op1.g(x)[5]
    g1_bb(x1, x2, x3, x4, x5, x6) = op1.g([x1, x2, x3, x4, x5, x6])[1]
    g2_bb(x1, x2, x3, x4, x5, x6) = op1.g([x1, x2, x3, x4, x5, x6])[2]
    g3_bb(x1, x2, x3, x4, x5, x6) = op1.g([x1, x2, x3, x4, x5, x6])[3]
    g4_bb(x1, x2, x3, x4, x5, x6) = op1.g([x1, x2, x3, x4, x5, x6])[4]
    g5_bb(x1, x2, x3, x4, x5, x6) = op1.g([x1, x2, x3, x4, x5, x6])[5]
    function ∇g1_bb(g::AbstractVector{T}, x1::T, x2::T, x3::T, x4::T, x5::T, x6::T) where {T}
        # println("I am in ∇g1_bb with x=", x1)
        global x0, ∇f_bkp, ∇g1_bkp, ∇g2_bkp, ∇g3_bkp, ∇g4_bkp, ∇g5_bkp
        jac_block([x1, x2, x3, x4, x5, x6])
        ∇g1_tmp = ∇g1_bkp
        for i = 1:length([x1, x2, x3, x4, x5, x6])
            g[i] = ∇g1_tmp[i]
        end
    end
    function ∇g2_bb(g::AbstractVector{T}, x1::T, x2::T, x3::T, x4::T, x5::T, x6::T) where {T}
        # println("I am in ∇g2_bb with x=", x1)
        global x0, ∇f_bkp, ∇g1_bkp, ∇g2_bkp, ∇g3_bkp, ∇g4_bkp, ∇g5_bkp
        jac_block([x1, x2, x3, x4, x5, x6])
        ∇g2_tmp = ∇g2_bkp
        for i = 1:length([x1, x2, x3, x4, x5, x6])
            g[i] = ∇g2_tmp[i]
        end
    end
    function ∇g3_bb(g::AbstractVector{T}, x1::T, x2::T, x3::T, x4::T, x5::T, x6::T) where {T}
        # println("I am in ∇g2_bb with x=", x1)
        global x0, ∇f_bkp, ∇g1_bkp, ∇g2_bkp, ∇g3_bkp, ∇g4_bkp, ∇g5_bkp
        jac_block([x1, x2, x3, x4, x5, x6])
        ∇g3_tmp = ∇g3_bkp
        for i = 1:length([x1, x2, x3, x4, x5, x6])
            g[i] = ∇g3_tmp[i]
        end
    end
    function ∇g4_bb(g::AbstractVector{T}, x1::T, x2::T, x3::T, x4::T, x5::T, x6::T) where {T}
        # println("I am in ∇g2_bb with x=", x1)
        global x0, ∇f_bkp, ∇g1_bkp, ∇g2_bkp, ∇g3_bkp, ∇g4_bkp, ∇g5_bkp
        jac_block([x1, x2, x3, x4, x5, x6])
        ∇g4_tmp = ∇g4_bkp
        for i = 1:length([x1, x2, x3, x4, x5, x6])
            g[i] = ∇g4_tmp[i]
        end
    end
    function ∇g5_bb(g::AbstractVector{T}, x1::T, x2::T, x3::T, x4::T, x5::T, x6::T) where {T}
        # println("I am in ∇g2_bb with x=", x1)
        global x0, ∇f_bkp, ∇g1_bkp, ∇g2_bkp, ∇g3_bkp, ∇g4_bkp, ∇g5_bkp
        jac_block([x1, x2, x3, x4, x5, x6])
        ∇g5_tmp = ∇g5_bkp
        for i = 1:length([x1, x2, x3, x4, x5, x6])
            g[i] = ∇g5_tmp[i]
        end
    end
    register(model, :g1_bb, op1.dim, g1_bb, ∇g1_bb; autodiff = false)
    @NLconstraint(model, g1_bb(x1, x2, x3, x4, x5, x6) <= 0)
    register(model, :g2_bb, op1.dim, g2_bb, ∇g2_bb; autodiff = false)
    @NLconstraint(model, g2_bb(x1, x2, x3, x4, x5, x6) == 0)
    register(model, :g3_bb, op1.dim, g2_bb, ∇g2_bb; autodiff = false)
    @NLconstraint(model, g3_bb(x1, x2, x3, x4, x5, x6) == 0)
    register(model, :g4_bb, op1.dim, g2_bb, ∇g2_bb; autodiff = false)
    @NLconstraint(model, g4_bb(x1, x2, x3, x4, x5, x6) == 0)
    register(model, :g5_bb, op1.dim, g2_bb, ∇g2_bb; autodiff = false)
    @NLconstraint(model, g5_bb(x1, x2, x3, x4, x5, x6) == 0)

    return model
end

model = setup_model()
function set_opt_params(model)
    set_optimizer_attribute(model, "tol", 1e-2)
    set_optimizer_attribute(model, "dual_inf_tol", 1e-2)
    set_optimizer_attribute(model, "constr_viol_tol", 1e-3)
    set_optimizer_attribute(model, "compl_inf_tol", 1e-2)
    set_optimizer_attribute(model, "acceptable_tol", 1e-2)
    set_optimizer_attribute(model, "acceptable_iter", 3)
    # set_optimizer_attribute(model, "mu_target", 1e-6)

    set_optimizer_attribute(model, "hessian_approximation", "limited-memory")

    set_optimizer_attribute(model, "max_iter", 50)
    set_optimizer_attribute(model, "max_cpu_time", 3600.0)
    set_optimizer_attribute(model, "print_level", 5)
    return model
end
model = set_opt_params(model)
JuMP.optimize!(model)

iter
println("Model used h=", h, " and tol=", get_optimizer_attribute(model, "tol"))
println("Initial value x0=", (op1.searchSpace[1, :] + op1.searchSpace[2, :]) ./ 2)
op1.searchSpace

@show JuMP.termination_status(model)
@show JuMP.primal_feasibility_report(model)
@show JuMP.dual_status(model)

# retrieve the objective value, corresponding x values and the status
println([value(model[:x1]), value(model[:x2]), value(model[:x3]), value(model[:x4]), value(model[:x5]), value(model[:x6])])
println(JuMP.objective_value(model))
println(JuMP.termination_status(model))

f_block(JuMP.value.([value(model[:x1]), value(model[:x2]), value(model[:x3]), value(model[:x4]), value(model[:x5]), value(model[:x6])]))


# This file is a part of Julia. License is MIT: https://julialang.org/license

module Random

using Base.dSFMT
using Base.GMP: Limb, MPZ
import Base: copymutable, copy, copy!, ==, hash

export srand,
       rand, rand!,
       randn, randn!,
       randexp, randexp!,
       bitrand,
       randstring,
       randsubseq, randsubseq!,
       shuffle, shuffle!,
       randperm, randperm!,
       randcycle, randcycle!,
       AbstractRNG, MersenneTwister, RandomDevice,
       GLOBAL_RNG, randjump

abstract type AbstractRNG end

abstract type FloatInterval{T<:AbstractFloat} end

struct CloseOpen{  T<:AbstractFloat} <: FloatInterval{T} end # interval [0,1)
struct Close1Open2{T<:AbstractFloat} <: FloatInterval{T} end # interval [1,2)

const FloatInterval_64 = FloatInterval{Float64}
const CloseOpen_64     = CloseOpen{Float64}
const Close1Open2_64   = Close1Open2{Float64}

CloseOpen(  ::Type{T}=Float64) where {T<:AbstractFloat} = CloseOpen{T}()
Close1Open2(::Type{T}=Float64) where {T<:AbstractFloat} = Close1Open2{T}()

const BitFloatType = Union{Type{Float16},Type{Float32},Type{Float64}}

function __init__()
    try
        srand()
    catch ex
        Base.showerror_nostdio(ex,
            "WARNING: Error during initialization of module Random")
    end
end

include("RNGs.jl")
include("generation.jl")
include("normal.jl")
include("misc.jl")

## rand & rand! & srand docstrings

"""
    rand([rng=GLOBAL_RNG], [S], [dims...])

Pick a random element or array of random elements from the set of values specified by `S`;
`S` can be

* an indexable collection (for example `1:n` or `['x','y','z']`),
* an `Associative` or `AbstractSet` object,
* a string (considered as a collection of characters), or
* a type: the set of values to pick from is then equivalent to `typemin(S):typemax(S)` for
  integers (this is not applicable to [`BigInt`](@ref)), and to ``[0, 1)`` for floating
  point numbers;

`S` defaults to [`Float64`](@ref).

# Examples
```julia-repl
julia> rand(Int, 2)
2-element Array{Int64,1}:
 1339893410598768192
 1575814717733606317

julia> rand(MersenneTwister(0), Dict(1=>2, 3=>4))
1=>2
```

!!! note
    The complexity of `rand(rng, s::Union{Associative,AbstractSet})`
    is linear in the length of `s`, unless an optimized method with
    constant complexity is available, which is the case for `Dict`,
    `Set` and `IntSet`. For more than a few calls, use `rand(rng,
    collect(s))` instead, or either `rand(rng, Dict(s))` or `rand(rng,
    Set(s))` as appropriate.
"""
rand

"""
    rand!([rng=GLOBAL_RNG], A, [S=eltype(A)])

Populate the array `A` with random values. If `S` is specified
(`S` can be a type or a collection, cf. [`rand`](@ref) for details),
the values are picked randomly from `S`.
This is equivalent to `copy!(A, rand(rng, S, size(A)))`
but without allocating a new array.

# Examples
```jldoctest
julia> rng = MersenneTwister(1234);

julia> rand!(rng, zeros(5))
5-element Array{Float64,1}:
 0.590845
 0.766797
 0.566237
 0.460085
 0.794026
```
"""
rand!

"""
    srand([rng=GLOBAL_RNG], seed) -> rng
    srand([rng=GLOBAL_RNG]) -> rng

Reseed the random number generator. If a `seed` is provided, the RNG will give a
reproducible sequence of numbers, otherwise Julia will get entropy from the system. For
`MersenneTwister`, the `seed` may be a non-negative integer or a vector of [`UInt32`](@ref)
integers. `RandomDevice` does not support seeding.

# Examples
```jldoctest
julia> srand(1234);

julia> x1 = rand(2)
2-element Array{Float64,1}:
 0.590845
 0.766797

julia> srand(1234);

julia> x2 = rand(2)
2-element Array{Float64,1}:
 0.590845
 0.766797

julia> x1 == x2
true
```
"""
srand

end # module


# with xr and xs
function fdtdKernel!{T}(n::Int,
		     p_out::Array{T},
		     p0::Array{T,3},
		     p1::Array{T,3},
		     p2::Array{T,3},
		     xr::Vector{Tuple{Int,Int,Int}},
		     xs::Vector{Tuple{Int,Int,Int}},
		     s::Array{T},
		     d1::T,
		     d2::T,
		     d3::T,
		     d4::T,
		     q::Vector{T},
		     G::Array{UInt8,3} )
	for i = 1:size(p0,3), m = 1:size(p0,2), l = 1:size(p0,1)
		if G[l,m,i] == 0x01
			air!(l,m,i,d1,d2,d3,d4,p0,p1,p2)
			#walls
		elseif G[l,m,i] == 0x02
			l1!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		elseif G[l,m,i] == 0x03
			le!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		elseif G[l,m,i] == 0x04
			m1!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		elseif G[l,m,i] == 0x05
			me!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		elseif G[l,m,i] == 0x06
			i1!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		elseif G[l,m,i] == 0x07
			ie!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)

			#edges
		elseif G[l,m,i] == 0x08
			l1m1!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		elseif G[l,m,i] == 0x09
			l1me!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		elseif G[l,m,i] == 0x0a
			l1i1!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		elseif G[l,m,i] == 0x0b
			l1ie!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)

		elseif G[l,m,i] == 0x0c
			lem1!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		elseif G[l,m,i] == 0x0d
			leme!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		elseif G[l,m,i] == 0x0e
			lei1!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		elseif G[l,m,i] == 0x0f
			leie!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)

		elseif G[l,m,i] == 0x10
			m1i1!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		elseif G[l,m,i] == 0x11
			m1ie!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		elseif G[l,m,i] == 0x12
			mei1!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		elseif G[l,m,i] == 0x13
			meie!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)

			#corners
		elseif G[l,m,i] == 0x14
			l1m1i1!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		elseif G[l,m,i] == 0x15
			l1m1ie!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		elseif G[l,m,i] == 0x16
			l1mei1!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		elseif G[l,m,i] == 0x17
			l1meie!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)

		elseif G[l,m,i] == 0x18
			lem1i1!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		elseif G[l,m,i] == 0x19
			lem1ie!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		elseif G[l,m,i] == 0x1a
			lemei1!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		elseif G[l,m,i] == 0x1b
			lemeie!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		end
		## add source
		for k in eachindex(xs)
			if xs[k][1] == l && xs[k][2] == m && xs[k][3] == i
				p2[l,m,i] += s[n,k]
			end
		end
		## get pressure
		for k in eachindex(xr)
			if xr[k][1] == l && xr[k][2] == m && xr[k][3] == i
				p_out[n,k] = p2[l,m,i]
			end
		end
	end
end

# full and xs
function fdtdKernel!{T}(n::Int,
		     p_out::Array{T,4}, #full output
		     p0::Array{T,3},
		     p1::Array{T,3},
		     p2::Array{T,3},
		     xs::Vector{Tuple{Int,Int,Int}},
		     s::Array{T},
		     d1::T,
		     d2::T,
		     d3::T,
		     d4::T,
		     q::Vector{T},
		     G::Array{UInt8,3} )
	for i = 1:size(p0,3), m = 1:size(p0,2), l = 1:size(p0,1)
		if G[l,m,i] == 0x01
			air!(l,m,i,d1,d2,d3,d4,p0,p1,p2)

			#walls
		elseif G[l,m,i] == 0x02
			l1!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		elseif G[l,m,i] == 0x03
			le!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		elseif G[l,m,i] == 0x04
			m1!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		elseif G[l,m,i] == 0x05
			me!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		elseif G[l,m,i] == 0x06
			i1!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		elseif G[l,m,i] == 0x07
			ie!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)

			#edges
		elseif G[l,m,i] == 0x08
			l1m1!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		elseif G[l,m,i] == 0x09
			l1me!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		elseif G[l,m,i] == 0x0a
			l1i1!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		elseif G[l,m,i] == 0x0b
			l1ie!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)

		elseif G[l,m,i] == 0x0c
			lem1!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		elseif G[l,m,i] == 0x0d
			leme!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		elseif G[l,m,i] == 0x0e
			lei1!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		elseif G[l,m,i] == 0x0f
			leie!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)

		elseif G[l,m,i] == 0x10
			m1i1!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		elseif G[l,m,i] == 0x11
			m1ie!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		elseif G[l,m,i] == 0x12
			mei1!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		elseif G[l,m,i] == 0x13
			meie!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)

			#corners
		elseif G[l,m,i] == 0x14
			l1m1i1!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		elseif G[l,m,i] == 0x15
			l1m1ie!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		elseif G[l,m,i] == 0x16
			l1mei1!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		elseif G[l,m,i] == 0x17
			l1meie!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)

		elseif G[l,m,i] == 0x18
			lem1i1!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		elseif G[l,m,i] == 0x19
			lem1ie!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		elseif G[l,m,i] == 0x1a
			lemei1!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		elseif G[l,m,i] == 0x1b
			lemeie!(l,m,i,d1,d2,d3,d4,p0,p1,p2,q)
		end
		## add source
		for k in eachindex(xs)
			if xs[k][1] == l && xs[k][2] == m && xs[k][3] == i
				p2[l,m,i] += s[n,k]
			end
		end
		p_out[l,m,i,n] = p2[l,m,i]
	end
end

# xr and full s
function fdtdKernel!{T}(n::Int,
		     p_out::Array{T},
		     p0::Array{T,3},
		     p1::Array{T,3},
		     p2::Array{T,3},
		     xr::Vector{Tuple{Int,Int,Int}},
		     s::Array{T,4},
		     d1::T,
		     d2::T,
		     d3::T,
		     d4::T,
		     q::Vector{T},
		     G::Array{UInt8,3} )
	for i = 1:size(p0,3), m = 1:size(p0,2), l = 1:size(p0,1)
		if G[l,m,i] == 0x01
			air!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,s)
			#walls
		elseif G[l,m,i] == 0x02
			l1!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		elseif G[l,m,i] == 0x03
			le!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		elseif G[l,m,i] == 0x04
			m1!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		elseif G[l,m,i] == 0x05
			me!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		elseif G[l,m,i] == 0x06
			i1!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		elseif G[l,m,i] == 0x07
			ie!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)

			#edges
		elseif G[l,m,i] == 0x08
			l1m1!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		elseif G[l,m,i] == 0x09
			l1me!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		elseif G[l,m,i] == 0x0a
			l1i1!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		elseif G[l,m,i] == 0x0b
			l1ie!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)

		elseif G[l,m,i] == 0x0c
			lem1!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		elseif G[l,m,i] == 0x0d
			leme!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		elseif G[l,m,i] == 0x0e
			lei1!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		elseif G[l,m,i] == 0x0f
			leie!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)

		elseif G[l,m,i] == 0x10
			m1i1!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		elseif G[l,m,i] == 0x11
			m1ie!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		elseif G[l,m,i] == 0x12
			mei1!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		elseif G[l,m,i] == 0x13
			meie!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)

			#corners
		elseif G[l,m,i] == 0x14
			l1m1i1!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		elseif G[l,m,i] == 0x15
			l1m1ie!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		elseif G[l,m,i] == 0x16
			l1mei1!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		elseif G[l,m,i] == 0x17
			l1meie!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)

		elseif G[l,m,i] == 0x18
			lem1i1!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		elseif G[l,m,i] == 0x19
			lem1ie!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		elseif G[l,m,i] == 0x1a
			lemei1!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		elseif G[l,m,i] == 0x1b
			lemeie!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		end
		## get pressure
		for k in eachindex(xr)
			if xr[k][1] == l && xr[k][2] == m && xr[k][3] == i
				p_out[n,k] = p2[l,m,i]
			end
		end
	end
end

# full p and full s
function fdtdKernel!{T}(n::Int,
		     p_out::Array{T,4}, #full output
		     p0::Array{T,3},
		     p1::Array{T,3},
		     p2::Array{T,3},
		     s::Array{T,4},
		     d1::T,
		     d2::T,
		     d3::T,
		     d4::T,
		     q::Vector{T},
		     G::Array{UInt8,3} )
	for i = 1:size(p0,3), m = 1:size(p0,2), l = 1:size(p0,1)
		if G[l,m,i] == 0x01
			air!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,s)
			#walls
		elseif G[l,m,i] == 0x02
			l1!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		elseif G[l,m,i] == 0x03
			le!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		elseif G[l,m,i] == 0x04
			m1!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		elseif G[l,m,i] == 0x05
			me!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		elseif G[l,m,i] == 0x06
			i1!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		elseif G[l,m,i] == 0x07
			ie!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)

			#edges
		elseif G[l,m,i] == 0x08
			l1m1!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		elseif G[l,m,i] == 0x09
			l1me!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		elseif G[l,m,i] == 0x0a
			l1i1!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		elseif G[l,m,i] == 0x0b
			l1ie!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)

		elseif G[l,m,i] == 0x0c
			lem1!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		elseif G[l,m,i] == 0x0d
			leme!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		elseif G[l,m,i] == 0x0e
			lei1!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		elseif G[l,m,i] == 0x0f
			leie!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)

		elseif G[l,m,i] == 0x10
			m1i1!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		elseif G[l,m,i] == 0x11
			m1ie!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		elseif G[l,m,i] == 0x12
			mei1!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		elseif G[l,m,i] == 0x13
			meie!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)

			#corners
		elseif G[l,m,i] == 0x14
			l1m1i1!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		elseif G[l,m,i] == 0x15
			l1m1ie!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		elseif G[l,m,i] == 0x16
			l1mei1!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		elseif G[l,m,i] == 0x17
			l1meie!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)

		elseif G[l,m,i] == 0x18
			lem1i1!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		elseif G[l,m,i] == 0x19
			lem1ie!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		elseif G[l,m,i] == 0x1a
			lemei1!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		elseif G[l,m,i] == 0x1b
			lemeie!(n,l,m,i,d1,d2,d3,d4,p0,p1,p2,q,s)
		end
		p_out[l,m,i,n] = p2[l,m,i]
	end
end


#!usr/bin/env julia

include("call_cc.jl")

@show @call_cc cc begin
    5*6+cc(7)
end # ⇒ 9

block(content::Function, except::Function) =
    (@call_cc success begin
        except(@call_cc fail begin
            success(content(fail))
        end)
    end)

block(
    raise -> begin
        function safe_div(a, b)
            if b == 0
                raise("divide by zero")
            else
                div(a, b)
            end
        end
        @show safe_div(3/4)
        @show safe_div(9/6)
        @show safe_div(2/0)
    end,
    errmsg -> begin
        print_with_color(:red, errmsg)
    end
)


#using Distributions
#include("constants.jl")
#include("orbit.jl")
#include("planet.jl")

if !isdefined(:PlanetarySystemAbstract)
  @compat abstract type PlanetarySystemAbstract end

  immutable PlanetarySystem{StarT<:StarAbstract} <: PlanetarySystemAbstract
    star::StarT
    planet::Vector{Planet}         # TODO OPT: If want to optimize, consider using  something like ImmmutableArrays?
    orbit::Vector{Orbit}           # TODO OPT: If want to optimize, consider using  something like ImmmutableArrays?
    # TODO DETAIL: Setup inner constructor to enforce equal number of planets & orbits
    #PlanetarySystem{StarT}(s::StarAbstract, p::Vector{Planet}, o::Vector{Orbit})
    #  @assert(length(p)==length(o)) # else error(string("Number of planets must match number of orbits: Np= ",length(p)," No= ",length(o)))
    #  new(s,p,o)
    #end
  end
  PlanetarySystemSingleStar = PlanetarySystem{SingleStar}

end
PlanetarySystem{StarT<:StarAbstract}(s::StarT) = PlanetarySystem(s,Array{Planet}(0),Array{Orbit}(0))  # Constructor for a Planetary System with no planets

PlanetarySystem{StarT<:StarAbstract}(s::StarT, p::Planet, o::Orbit) = PlanetarySystem(s,[p],[o])  # Constructor for a single Planet System

function PlanetarySystem{StarT}(ps::PlanetarySystem{StarT}, keep::Vector{Int64}) # Why doesn't this work?
   PlanetarySystem{StarT}(ps.star,ps.planet[keep],ps.orbit[keep])
end

#function PlanetarySystemSingleStar(ps::PlanetarySystemSingleStar, keep::Vector{Int64})
#function PlanetarySystem{StarT<:StarAbstract}(ps::PlanetarySystem{StarT}, keep::Vector{Int64})
#   PlanetarySystem(ps.star,ps.planet[keep],ps.orbit[keep])
#end

flux(ps::PlanetarySystem{StarAbstract}) = flux(ps.star)
flux(ps::PlanetarySystem{Star}) = flux(ps.star)
flux(ps::PlanetarySystem{BinaryStar}) = flux(ps.star)
flux(ps::PlanetarySystem{MultipleStar}) = flux(ps.star)

function num_planets(s::PlanetarySystem{Star})
  @assert( length(s.planet) == length(s.orbit) )    # TODO OPT: Deactivate inner assert's like this for speed once tested
  return length(s.planet)
end

function generate_planet_mass_from_radius_powerlaw(r::Float64, sim_param::SimParam)
  const mr_power_index::Float64 = get_real(sim_param,"mr_power_index")
  const mr_const::Float64 = get_real(sim_param,"mr_const")
  const mr_max_mass::Float64 = get_real(sim_param,"mr_max_mass")
  m = mr_const*earth_mass*(r/earth_radius)^mr_power_index
  if m > mr_max_mass
     m = mr_max_mass
  end
  return m
end

function generate_planet_mass_from_radius_powerlaw(r::Float64, s::Star, o::Orbit, sim_param::SimParam)  # TODO SCI: IMPORTANT once have stability criteria, replace w/ better M-R relationship
  generate_planet_mass_from_radius_powerlaw(r,sim_param)
end

function generate_num_planets_poisson(lambda::Real, max_planets::Integer; min_planets::Integer = 0)
  if lambda < min_planets*1e-3
      return min_planets
  end
  bug_fixed = false # TODO OPT: true case should work, but Danley found bug in Distributions package.  Revert once fixed for speed.
  local n
  if bug_fixed
     d = Distributions.Truncated(Distributions.Poisson(lambda),min_planets,max_planets)
     n = rand(d)
  else
     if min_planets == 0
        min_planets = -1
     end
     d = Distributions.Truncated(Distributions.Poisson(lambda),min_planets,max_planets)
     n = rand(d)
     #=
     n = -1
     while !(min_planets<=n<=max_planets)
        n = rand(Distributions.Poisson(lambda))
     end
     =#
  end
  return n
end

function generate_num_planets_poisson(s::Star, sim_param::SimParam)
  const lambda::Float64 = exp(get_real(sim_param,"log_eta_pl"))
  const max_tranets_in_sys::Int64 = get_int(sim_param,"max_tranets_in_sys")
  generate_num_planets_poisson(lambda,max_tranets_in_sys)
end

function generate_period_and_sizes_log_normal(s::Star, sim_param::SimParam; num_pl::Integer = 1)  # TODO USER SCI:  User should make sure planetary properties are being drawn appropriately for their scientific purposes
    const mu_log_r::Float64 = get_real(sim_param,"mean_log_planet_radius")
    const sigma_log_r::Float64 = get_real(sim_param,"sigma_log_planet_radius")
    const mu_log_P::Float64 = get_real(sim_param,"mean_log_planet_period")
    const sigma_log_P::Float64 = get_real(sim_param,"sigma_log_planet_period")
    const min_period::Float64 = get_real(sim_param,"min_period")
    const max_period::Float64 = get_real(sim_param,"max_period")
    const min_radius::Float64 = get_real(sim_param,"min_radius")
    const max_radius::Float64 = get_real(sim_param,"max_radius")
    const max_draws::Int64 = 100

    if   sigma_log_r <= 0. || sigma_log_P<=0.
     println("# mu_log_r= ", mu_log_r, " sigma_log_r= ", sigma_log_r, " mu_log_P= ", mu_log_P, " sigma_log_P= ", sigma_log_P)
    end
    const rdist = LogNormal(mu_log_r,sigma_log_r)
    const Pdist = LogNormal(mu_log_P,sigma_log_P)
    #Rlist = rand(rdist,num_pl)
    #Plist = rand(Pdist,num_pl)
    #idx_keep = find(i->(min_radius<=Rlist[i]<=max_radius) && (min_period<=Plist[i]<=max_period), 1:num_pl )
    #return Plist[idx_keep], Rlist[idx_keep]  # replaced because want to return exactly num_pl.  Could use Truncated to restore above code.
    Rlist = zeros(num_pl)
    Plist = zeros(num_pl)
    for i in 1:num_pl
      j = 0
      while ! (min_radius<Rlist[i]<=max_radius) && j<max_draws
            Rlist[i] = rand(rdist)
            j+=1
      end
      if j>=max_draws
         println("# Struggled to draw size for: ",mu_log_r, " ", sigma_log_r)
      end
      j = 0
      while ! (min_period<Plist[i]<=max_period) && j<max_draws
            Plist[i] = rand(Pdist)
            j+=1
      end
      if j>=max_draws
         println("# Struggled to draw period for: ",mu_log_P, " ", sigma_log_P)
      end
    end
    return Plist, Rlist
end

function draw_power_law(n::Real, x0::Real, x1::Real, num_pl::Integer)
     ((x1^(n+1) - x0^(n+1)).*rand(num_pl) .+ x0^(n+1)).^(1/(n+1))
end

function draw_broken_power_law(n1::Real, n2::Real, x0::Real, x1::Real, xb::Real, num_pl::Integer)
    #x0 and x1 are the lower and upper truncation limits, and xb is the break point, i.e. x0 <= xb <= x1 (all must be positive)
    #n1 and n2 are the power law indices between x0 and xb, and xb and x1, respectively (can be positive or negative)
    @assert(x0 <= xb <= x1)
    @assert(num_pl >= 1)

    u_draws = rand(num_pl) #'num_pl' draws from the uniform distribution between 0 and 1
    x_draws = zeros(num_pl)

    if (n1 != -1) & (n2 != -1)
        C1 = 1./(((xb^(n1+1) - x0^(n1+1))/(n1+1)) + ((xb^(n1-n2)*(x1^(n2+1) - xb^(n2+1)))/(n2+1))) #normalization constant
        ub = (C1*(xb^(n1+1) - x0^(n1+1)))/(n1+1) #break point in u, between 0 and 1
        for (i,u) in enumerate(u_draws)
            if u <= ub
                x_draws[i] = (((n1+1)*u)/C1 + x0^(n1+1))^(1/(n1+1))
            else #if u > ub
                x_draws[i] = (((n2+1)/(C1*xb^(n1-n2)))*(u - (C1*(xb^(n1+1) - x0^(n1+1)))/(n1+1)) + xb^(n2+1))^(1/(n2+1))
            end
        end
    elseif (n1 == -1) & (n2 != -1)
        C1 = 1./(log(xb/x0) + ((xb^(-1-n2))*(x1^(n2+1)) - 1)/(n2+1)) #normalization constant
        ub = C1*log(xb/x0) #break point in u, between 0 and 1
        for (i,u) in enumerate(u_draws)
            if u <= ub
                x_draws[i] = x0*exp(u/C1)
            else #if u > ub
                x_draws[i] = (((n2+1)/(C1*xb^(-1-n2)))*(u - C1*log(xb/x0)) + xb^(n2+1))^(1/(n2+1))
            end
        end
    elseif (n1 != -1) & (n2 == -1)
        C1 = 1./(((xb^(n1+1) - x0^(n1+1))/(n1+1)) + (xb^(n1+1))*log(x1/xb)) #normalization constant
        ub = (C1*(xb^(n1+1) - x0^(n1+1)))/(n1+1) #break point in u, between 0 and 1
        for (i,u) in enumerate(u_draws)
            if u <= ub
                x_draws[i] = (((n1+1)*u)/C1 + x0^(n1+1))^(1/(n1+1))
            else #if u > ub
                x_draws[i] = xb*exp((1/(C1*xb^(n1+1)))*(u - (C1*(xb^(n1+1) - x0^(n1+1)))/(n1+1)))
            end
        end
    else #if n1 == -1 and n2 == -1 (i.e. it is a single power-law with index of -1)
        for (i,u) in enumerate(u_draws)
            x_draws[i] = x0*exp(u*log(x1/x0))
        end
    end

    return x_draws
end

function generate_periods_power_law(s::Star, sim_param::SimParam; num_pl::Integer = 1)
    const power_law_P::Float64 = get_real(sim_param,"power_law_P")
    const min_period::Float64 = get_real(sim_param,"min_period")
    const max_period::Float64 = get_real(sim_param,"max_period")
    Plist = power_law_P!=-1.0 ? draw_power_law(power_law_P,min_period,max_period, num_pl) : exp(log(min_period).+rand()*log(max_period/min_period))
    return Plist
end

function generate_sizes_power_law(s::Star, sim_param::SimParam; num_pl::Integer = 1)
    const power_law_r::Float64 = get_real(sim_param,"power_law_r")
    const min_radius::Float64 = get_real(sim_param,"min_radius")
    const max_radius::Float64 = get_real(sim_param,"max_radius")
    Rlist = power_law_r!=-1.0 ? draw_power_law(power_law_r,min_radius,max_radius, num_pl) : exp(log(min_radius).+rand()*log(max_radius/min_radius))
    return Rlist
end

function generate_sizes_broken_power_law(s::Star, sim_param::SimParam; num_pl::Integer = 1)
    const power_law_r1::Float64 = get_real(sim_param,"power_law_r1")
    const power_law_r2::Float64 = get_real(sim_param,"power_law_r2")
    const min_radius::Float64 = get_real(sim_param,"min_radius")
    const max_radius::Float64 = get_real(sim_param,"max_radius")
    const break_radius::Float64 = get_real(sim_param,"break_radius")
    Rlist = draw_broken_power_law(power_law_r1,power_law_r2,min_radius,max_radius,break_radius, num_pl)
    return Rlist
end

function generate_period_and_sizes_power_law(s::Star, sim_param::SimParam; num_pl::Integer = 1)
    return (generate_periods_power_law(s, sim_param, num_pl=num_pl), generate_sizes_power_law(s, sim_param, num_pl=num_pl))
end

function TruncatedUpper(d::Distributions.UnivariateDistribution, u::Float64)
    zero(u) < u || error("lower bound should be less than upper bound.")
    lcdf = zero(u)
    ucdf = isinf(u) ? one(u) : cdf(d, u)
    tp = ucdf - lcdf
    Distributions.Truncated{typeof(d),Distributions.value_support(typeof(d))}(d, zero(u), u, lcdf, ucdf, tp, log(tp))
end

function generate_e_omega_rayleigh_direct(sigma_hk::Float64; max_e::Float64 = 1.0)
  @assert(0<max_e<=1.0)
  ecc::Float64 = rand( TruncatedUpper(Rayleigh(sigma_hk),max_e) )
  w::Float64 = 2pi*rand()
  return ecc, w
end

function generate_e_omega_rayleigh_two_gaussians(sigma_hk::Float64; max_e::Float64 = 1.0)
  @assert(0<max_e<=1.0)
  h = k = 1.0
  while h*h+k*k >= max_e*max_e
    h = sigma_hk*randn()
    k = sigma_hk*randn()
  end
  ecc::Float64 = sqrt(h*h+k*k)
  w::Float64 = atan2(k,h)
  return ecc, w
end

function generate_e_omega_rayleigh(sigma_hk::Float64; max_e::Float64 = 1.0)
   if max_e > sigma_hk
      return generate_e_omega_rayleigh_two_gaussians(sigma_hk,max_e=max_e)
   else
      return generate_e_omega_rayleigh_direct(sigma_hk,max_e=max_e)
   end
end

function generate_e_omega_rayleigh(sim_param::SimParam; max_e::Float64 = 1.0)
  sigma_hk::Float64 = get_real(sim_param,"sigma_hk")
  generate_e_omega_rayleigh(sigma_hk, max_e=max_e)
end

function generate_planetary_system_hardcoded_example(star::StarAbstract, sim_param::SimParam; verbose::Bool = false)
  # in this version we specify fixed functions that are known at compile time, allowing for additional optimizations (~0.6 second faster per Kepler catalog out of ~3.6 sec on my laptop w/ 1 core)
  const  generate_planet_mass_from_radius = generate_planet_mass_from_radius_powerlaw
  const  generate_num_planets = generate_num_planets_poisson
  const  generate_period_and_sizes = generate_period_and_sizes_log_normal

  const  generate_e_omega =  generate_e_omega_rayleigh

  # const  generate_star = get_function(sim_param,"generate_star")
  # const star::StarAbstract = generate_star(sim_param)
  const num_pl::Int64 = generate_num_planets(star, sim_param)

  if( num_pl==0 )
    return PlanetarySystem(star)
  else
    (Plist::Vector{Float64}, Rlist::Vector{Float64}) = generate_period_and_sizes(star, sim_param, num_pl=num_pl)
    idx = sortperm(Plist)                   # TODO OPT: Check to see if sorting is significant time sink.  If so, it might could be deferred

    min_a_in_rstar = 2.0
    min_P_orbit = day_in_year*sqrt((min_a_in_rstar*star.radius*rsol_in_au)^3 / star.mass) # minimum semi-major axis of two stellar radii
    idx = idx[Plist[idx] .> min_P_orbit]
    if( length(idx)==0 )
        return PlanetarySystem(star)
    end

    pl = Array{Planet}(length(idx))
    orbit = Array{Orbit}(length(idx))
    a = map(i->semimajor_axis(Plist[i],star.mass),idx)
    max_e = ones(length(idx))
    max_e_factor = 0.999 # A factor just less than 1 to prevent numerical issues with near-crossing orbits
    if length(a)>=2
       for i in 1:length(a)
         if i==1
           max_e[i] = max_e_factor*(1-a[i]/a[i+1])/(1+a[i]/a[i+1])
         elseif i==length(a)
           max_e[i] = max_e_factor*(1-a[i-1]/a[i])/(1+a[i-1]/a[i])
         else
           max_e[i] = max_e_factor*min( (1-a[i]/a[i+1])/(1+a[i]/a[i+1]), (1-a[i-1]/a[i])/(1+a[i-1]/a[i]) )
         end
       end
    end
    for i in 1:length(idx)
      # if verbose   println("i=",i," idx=",idx," Plist=",Plist[idx] );     end
      P = Plist[idx[i]]
      Rpl = Rlist[idx[i]]
      (ecc::Float64,  omega::Float64) = generate_e_omega(sim_param, max_e=max_e[i])
      incl::Float64 = acos(rand())
      orbit[i] = Orbit(P,ecc,incl,omega,2pi*rand(),2pi*rand())
      mass::Float64 = generate_planet_mass_from_radius(Rpl, star, orbit[i], sim_param)
      pl[i] = Planet( Rpl,  mass )
    end
  return PlanetarySystem(star,pl,orbit)
  end
end

function generate_planetary_system_uncorrelated_incl(star::StarAbstract, sim_param::SimParam; verbose::Bool = false)
  # load functions to use for drawing parameters
  const  generate_planet_mass_from_radius = get_function(sim_param,"generate_planet_mass_from_radius")
  const  generate_num_planets = get_function(sim_param,"generate_num_planets")
  const  generate_period_and_sizes = get_function(sim_param,"generate_period_and_sizes")
  const  generate_e_omega = get_function(sim_param,"generate_e_omega")

  # const  generate_star = get_function(sim_param,"generate_star")
  # const star::StarAbstract = generate_star(sim_param)
  const num_pl::Int64 = generate_num_planets(star, sim_param)::Int64
  const sigma_ecc::Float64 = haskey(sim_param,"sigma_hk") ? get_real(sim_param,"sigma_hk") : 0.0

  if( num_pl==0 )
    return PlanetarySystem(star)::PlanetarySystem
  else
    (Plist::Vector{Float64}, Rlist::Vector{Float64}) = generate_period_and_sizes(star, sim_param, num_pl=num_pl)
    idx = sortperm(Plist)                   # TODO OPT: Check to see if sorting is significant time sink.  If so, it might could be deferred

    min_a_in_rstar = 2.0
    min_P_orbit = day_in_year*sqrt((min_a_in_rstar*star.radius*rsol_in_au)^3 / star.mass) # minimum semi-major axis of two stellar radii
    idx = idx[Plist[idx] .> min_P_orbit]
    if( length(idx)==0 )
        return PlanetarySystem(star)
    end

    pl = Array{Planet}(length(idx))
    orbit = Array{Orbit}(length(idx))
    a = map(i->semimajor_axis(Plist[i],star.mass),idx)
    max_e = ones(length(idx))
    max_e_factor = 0.999 # A factor just less than 1 to prevent numerical issues with near-crossing orbits
    if length(a)>=2
       for i in 1:length(a)
          if i==1
            max_e[i] = max_e_factor*(1-a[i]/a[i+1])/(1+a[i]/a[i+1])
          elseif i==length(a)
            max_e[i] = max_e_factor*(1-a[i-1]/a[i])/(1+a[i-1]/a[i])
          else
            max_e[i] = max_e_factor*min( (1-a[i]/a[i+1])/(1+a[i]/a[i+1]), (1-a[i-1]/a[i])/(1+a[i-1]/a[i]) )
          end
        end
    end
    for i in 1:length(idx)
      # if verbose   println("i=",i," idx=",idx," Plist=",Plist[idx] );     end
      P = Plist[idx[i]]
      Rpl = Rlist[idx[i]]
      if haskey(sim_param,"sigma_hk_one") && haskey(sim_param,"sigma_hk_multi")
         sigma_ecc = num_pl == 1 ? get_real(sim_param,"sigma_hk_one") : get_real(sim_param,"sigma_hk_multi")
      end
      (ecc::Float64,  omega::Float64) = generate_e_omega(sim_param, max_e=max_e[i])
      incl::Float64 = acos(rand())
      orbit[i] = Orbit(P,ecc,incl,omega,2pi*rand(),2pi*rand())
      # set!(orbit[idx[i]],P,ecc,incl,omega,2pi*rand(),2pi*rand())
      mass::Float64 = generate_planet_mass_from_radius(Rpl, star, orbit[i], sim_param)
      pl[i] = Planet( Rpl,  mass )
    end
  return PlanetarySystem(star,pl,orbit)
  end
end

# This version generates more systems roughly near a common plane, but until incorporate CORBITS data, ABC can't match input param
function generate_planetary_system_simple(star::StarAbstract, sim_param::SimParam; verbose::Bool = false)
  # load functions to use for drawing parameters
  const  generate_planet_mass_from_radius = get_function(sim_param,"generate_planet_mass_from_radius")
  const  generate_num_planets = get_function(sim_param,"generate_num_planets")
  #const  generate_num_planets = generate_num_planets_christiansen
  #const  generate_period_and_sizes = generate_period_and_sizes_christiansen
  const  generate_period_and_sizes = get_function(sim_param,"generate_period_and_sizes")
  const  generate_e_omega = get_function(sim_param,"generate_e_omega")
  const  sigma_incl = deg2rad(get_real(sim_param,"sigma_incl"))

  # const  generate_star = get_function(sim_param,"generate_star")
  # const star::StarAbstract = generate_star(sim_param)
  const num_pl = generate_num_planets(star, sim_param)::Int64
  const sigma_ecc::Float64 = haskey(sim_param,"sigma_hk") ? get_real(sim_param,"sigma_hk") : 0.0

  if( num_pl==0 )
    return PlanetarySystem(star)
  else
    (Plist::Vector{Float64}, Rlist::Vector{Float64}) = generate_period_and_sizes(star, sim_param, num_pl=num_pl)
    idx = sortperm(Plist)                   # TODO OPT: Check to see if sorting is significant time sink.  If so, it might could be deferred
    incl_sys = acos(rand())

    min_a_in_rstar = 2.0
    min_P_orbit = day_in_year*sqrt((min_a_in_rstar*star.radius*rsol_in_au)^3 / star.mass) # minimum semi-major axis of two stellar radii
    idx = idx[Plist[idx] .> min_P_orbit]
    if( length(idx)==0 )
        return PlanetarySystem(star)
    end

    pl = Array{Planet}(length(idx))
    orbit = Array{Orbit}(length(idx))
    for i in 1:length(idx)
      # if verbose   println("i=",i," idx=",idx," Plist=",Plist[idx] );     end
      P = Plist[idx[i]]
      Rpl = Rlist[idx[i]]
      if haskey(sim_param,"sigma_hk_one") && haskey(sim_param,"sigma_hk_multi")
         sigma_ecc = num_pl == 1 ? get_real(sim_param,"sigma_hk_one") : get_real(sim_param,"sigma_hk_multi")
      end
      (ecc,  omega) = generate_e_omega(sim_param)::Tuple{Float64,Float64}
      incl_mut = sigma_incl*sqrt(randn()^2+randn()^2) # rand(Distributions.Rayleigh(sigma_incl)) # sigma_incl*randn()
      asc_node = 2pi*rand()
      mean_anom = 2pi*rand()
      #incl = incl_sys + sigma_incl*randn()
      incl =  incl_mut!=zero(incl_mut) ? acos( cos(incl_sys)*cos(incl_mut) + sin(incl_sys)*sin(incl_mut)*cos(asc_node) ) : incl_sys
      orbit[i] = Orbit(P,ecc,incl,omega,asc_node,mean_anom)
      mass = generate_planet_mass_from_radius(Rpl, star, orbit[i], sim_param)::Float64
      pl[i] = Planet( Rpl,  mass )
    end
  return PlanetarySystem(star,pl,orbit)
  end
end

function test_planetary_system_constructors(sim_param::SimParam)
  generate_star = get_function(sim_param,"generate_star")
  star = generate_star(sim_param)
  empty_sys = PlanetarySystem(star)
  earth = Planet(earth_radius,earth_mass)
  earth_orbit = Orbit(365.2425,0.0167,0.5*pi,0.0,0.0,0.0)
  solar_sys = PlanetarySystem(star, earth,earth_orbit)
  m = generate_planet_mass_from_radius_powerlaw(0.02,star,earth_orbit,sim_param)/earth_mass
  generate_planetary_system_simple(star,sim_param,verbose=true)
end