# using Distributions, PyPlot
push!(LOAD_PATH, "../../VICAS/src/")
push!(LOAD_PATH, "../../DeepCorrection")
push!(LOAD_PATH, "../../DeepCorrection/src")
push!(LOAD_PATH, "../TrajectoryViz")

include("../TrajectoryViz/trajViz.jl")

th = linspace(0, 2*pi, 50)

conflictRange = SENSING_RANGE * 1.5

function check_NMAC(i, acArr, finished)
	ac = acArr[i]
	for j in 1:length(acArr)
		if j != i && finished[i] == false && finished[j] == false
			d = sqrt((ac.x - acArr[j].x)^2 + (ac.y - acArr[j].y)^2)
			if d <= NMAC_RANGE
				acArr[i].dmin = d
				acArr[j].dmin = d
				return true, d
			end
		end
	end
	return false, Inf
end

function check_colli(ac, acArr)
	for existing_ac in acArr
		if sqrt((ac.x - existing_ac.x)^2 + (ac.y - existing_ac.y)^2) <= NMAC_RANGE
			return true
		end
	end
end

function create_ac(spawn_range, CASType, policy)
	# Annulus type, centralized encounter
	spawn_range = rand(2000:4000)
	start_angle = rand(th)
	x0 = spawn_range * cos(start_angle)
	y0 = spawn_range * sin(start_angle)
	# conflictPoint_x = conflictRange * (- 2 * rand() + 1)
	# conflictPoint_y  = sqrt(conflictRange^2 - conflictPoint_x^2) * (- 2 * rand() + 1)
	# dest_angle = atan2(conflictPoint_y - y0, conflictPoint_x - x0)
	# dest_dist = (4000 - 2 * spawn_range) * rand() + 2 * spawn_range
	dest_dist = rand(2000:4000)
	# x_dest = x0 + cos(dest_angle) * dest_dist
	# y_dest = y0 + sin(dest_angle) * dest_dist
	x_dest = dest_dist * cos(start_angle + pi)
	y_dest = dest_dist * sin(start_angle + pi)

	# Circular type encounter
	# x0 = spawn_range * (- 2 * rand() + 1)
	# y0max  = sqrt(spawn_range^2 - x0^2)
	# y0     = y0max * (- 2 * rand() + 1)
	# dest_angle = rand(th)
	# dest_dist = 2*spawn_range
	# x_dest = cos(dest_angle) * dest_dist
	# y_dest = sin(dest_angle) * dest_dist

	# dest_angle = rand(th)
	# x_dest = spawn_range * cos(dest_angle)
	# y_dest = spawn_range * sin(dest_angle)
	ac = initOwnship(x0, y0, x_dest, y_dest, CASType, policy)
	return ac
end

function resRatio(;
				CASType::Symbol=:correctedSector,
				policyFile::AbstractString="",
				numAC::Int64=2,
				numSim::Int64=10000,
				timeout::Int64=200,
				correction_weight::Float64=0.09)

	problem_file = "../../DeepCorrection/logs/" * policyFile * "/final_problem.jld"
	weights_file = "../../DeepCorrection/logs/" * policyFile * "/final_weights.jld"
	policy = restore(problem_file=problem_file, weights_file=weights_file)
	policy.env.correction_weight = correction_weight
	spawn_range = SENSING_RANGE * 1.5
	resolvedStats = 0
	unresolved_NMAC_Stats = Int64[]
	unresolved_Timeout_Stats = 0
	resTimeArr = Vector{Float64}()
	avgResTime = 0
	criticality = Float64[]
	v_closing_v_cruise = Float64[]

	for sim in 1 : numSim
		# Initiate
		t = 1
		NMACFlag = false
		fail = false
		finished = falses(numAC)
		acArr = Ownship[]
		for i = 1 : numAC
			ac = create_ac(spawn_range, CASType, policy)
			while check_colli(ac, acArr) == true
				ac = create_ac(spawn_range, CASType, policy)
			end
			push!(acArr, ac)
		end

		push!(unresolved_NMAC_Stats, 0)

		while !all(finished) && NMACFlag == false
			for i in 1 : numAC
				if !finished[i]
					ownship = acArr[i]
					# Get intruder coords for ownship
					intCoords = Vector{Float64}()
					for j in 1 : length(acArr)
						if j != i && !finished[j] # not ownship and not finished
							intCoords = vcat(intCoords, [acArr[j].x, acArr[j].y], acArr[j].heading)
						end
					end
					numInt = Int(length(intCoords) / 3)

					if numInt > 0 && CASType != :NOCAS
						# Get state
						numObservedInt = getState!(ownship, numInt, intCoords)
						if numObservedInt > 0
							# Get action and advisory
							getAdvisory!(ownship)
						else
							ownship.advisory = :COC
						end
					else
						ownship.advisory = :COC
					end
					# print(i, " : ", ownship.state, " : ", ownship.advisory, "\n")
				end # if !finished[i]
			end # for

			for i in 1 : numAC
				if !finished[i]
					dest_angle = norm_angle(atan2(acArr[i].y_dest - acArr[i].y, acArr[i].x_dest - acArr[i].x))
					heading_error = norm_angle(dest_angle - acArr[i].heading)
					push!(v_closing_v_cruise, cos(heading_error))
					updateOwnshipCoords!(acArr[i])
				end

				if sqrt((acArr[i].x - acArr[i].x_dest)^2 + (acArr[i].y - acArr[i].y_dest)^2) < 80
					finished[i] = true
				end
			end # for

			for i in 1 : numAC
				nmac, d = check_NMAC(i, acArr, finished)
				if finished[i] == false && nmac == true # FAIL: NMAC
					push!(criticality, max(0.0, 1 - acArr[i].dmin / NMAC_RANGE))
					fail = true
					NMACFlag = true
					unresolved_NMAC_Stats[sim] += 1
				end
			end

			if all(finished) && NMACFlag == false # SUCCESS
				resolvedStats += 1
			end

			t += 1
			if t > timeout && NMACFlag == false # FAIL: timeout
				fail = true
				unresolved_Timeout_Stats += 1
				break # while
			end

		end #while

		if t <= timeout && NMACFlag == false # only count SUCCESS cases
			avgResTime += t
		end
	end # for sim

	return resolvedStats, unresolved_NMAC_Stats ./ 2, unresolved_Timeout_Stats, avgResTime / resolvedStats, criticality, v_closing_v_cruise
end


# function add_ref_dcgrid!(ref::Dict{Symbol,<:Any}, data::Dict{String,<:Any})
#     for (n, nw_ref) in ref[:nw]
#         if haskey(nw_ref, :convdc)
#             #Filter converters & DC branches with status 0 as well as wrong bus numbers
#             nw_ref[:convdc] = Dict([x for x in nw_ref[:convdc] if (x.second["status"] == 1 && x.second["busdc_i"] in keys(nw_ref[:busdc]) && x.second["busac_i"] in keys(nw_ref[:bus]))])
#             nw_ref[:branchdc] = Dict([x for x in nw_ref[:branchdc] if (x.second["status"] == 1 && x.second["fbusdc"] in keys(nw_ref[:busdc]) && x.second["tbusdc"] in keys(nw_ref[:busdc]))])
#
#             # DC grid arcs for DC grid branches
#             nw_ref[:arcs_dcgrid_from] = [(i,branch["fbusdc"],branch["tbusdc"]) for (i,branch) in nw_ref[:branchdc]]
#             nw_ref[:arcs_dcgrid_to]   = [(i,branch["tbusdc"],branch["fbusdc"]) for (i,branch) in nw_ref[:branchdc]]
#             nw_ref[:arcs_dcgrid] = [nw_ref[:arcs_dcgrid_from]; nw_ref[:arcs_dcgrid_to]]
#             nw_ref[:arcs_conv_acdc] = [(i,conv["busac_i"],conv["busdc_i"]) for (i,conv) in nw_ref[:convdc]]
#             #bus arcs of the DC grid
#             bus_arcs_dcgrid = Dict([(bus["busdc_i"], []) for (i,bus) in nw_ref[:busdc]])
#             for (l,i,j) in nw_ref[:arcs_dcgrid]
#                 push!(bus_arcs_dcgrid[i], (l,i,j))
#             end
#             nw_ref[:bus_arcs_dcgrid] = bus_arcs_dcgrid
#
#                 # bus_convs for AC side power injection of DC converters
#             bus_convs_ac = Dict([(i, []) for (i,bus) in nw_ref[:bus]])
#             for (i,conv) in nw_ref[:convdc]
#                 push!(bus_convs_ac[conv["busac_i"]], i)
#             end
#             nw_ref[:bus_convs_ac] = bus_convs_ac
#
#             # bus_convs for AC side power injection of DC converters
#             bus_convs_dc = Dict([(bus["busdc_i"], []) for (i,bus) in nw_ref[:busdc]])
#             for (i,conv) in nw_ref[:convdc]
#                 push!(bus_convs_dc[conv["busdc_i"]], i)
#             end
#             nw_ref[:bus_convs_dc] = bus_convs_dc
#             # Add DC reference buses
#             ref_buses_dc = Dict{String, Any}()
#             for (k,v) in nw_ref[:convdc]
#                 if v["type_dc"] == 2
#                     ref_buses_dc["$k"] = v
#                 end
#             end
#
#             if length(ref_buses_dc) == 0
#                 for (k,v) in nw_ref[:convdc]
#                     if v["type_ac"] == 2
#                         ref_buses_dc["$k"] = v
#                     end
#                 end
#                 Memento.warn(_PM._LOGGER, "no reference DC bus found, setting reference bus based on AC bus type")
#             end
#
#             for (k,conv) in nw_ref[:convdc]
#                 conv_id = conv["index"]
#                 if conv["type_ac"] == 2 && conv["type_dc"] == 1
#                     Memento.warn(_PM._LOGGER, "For converter $conv_id is chosen P is fixed on AC and DC side. This can lead to infeasibility in the PF problem.")
#                 elseif conv["type_ac"] == 1 && conv["type_dc"] == 1
#                     Memento.warn(_PM._LOGGER, "For converter $conv_id is chosen P is fixed on AC and DC side. This can lead to infeasibility in the PF problem.")
#                 end
#                 convbus_ac = conv["busac_i"]
#                 if conv["Vmmax"] < nw_ref[:bus][convbus_ac]["vmin"]
#                     Memento.warn(_PM._LOGGER, "The maximum AC side voltage of converter $conv_id is smaller than the minimum AC bus voltage")
#                 end
#                 if conv["Vmmin"] > nw_ref[:bus][convbus_ac]["vmax"]
#                     Memento.warn(_PM._LOGGER, "The miximum AC side voltage of converter $conv_id is larger than the maximum AC bus voltage")
#                 end
#             end
#
#             if length(ref_buses_dc) > 1
#                 ref_buses_warn = ""
#                 for (rb) in keys(ref_buses_dc)
#                     ref_buses_warn = ref_buses_warn*rb*", "
#                 end
#                 Memento.warn(_PM._LOGGER, "multiple reference buses found, i.e. "*ref_buses_warn*"this can cause infeasibility if they are in the same connected component")
#             end
#
#
#             # if haskey(pm.setting, "find_all_ac_grids") && pm.tetting["find_all_ac_grids"] == true
#             #     ACgrids = find_all_ac_grids(nw_ref[:branch], nw_ref[:bus])
#             #     for (i, grid) in ACgrids
#             #         a = 0
#             #         for (j, bus) in nw_ref[:ref_buses]
#             #             if (bus["bus_i"] in grid["Buses"])
#             #                 a = 1
#             #             end
#             #         end
#             #         if a == 0
#             #             Memento.warn(_PM._LOGGER, "Grid $i does not have any voltage reference bus, this might cause infeasibility")
#             #         end
#             #     end
#             # end
#             nw_ref[:ref_buses_dc] = ref_buses_dc
#             nw_ref[:buspairsdc] = buspair_parameters_dc(nw_ref[:arcs_dcgrid_from], nw_ref[:branchdc], nw_ref[:busdc])
#         else
#             nw_ref[:convdc] = Dict{String, Any}()
#             nw_ref[:busdc] = Dict{String, Any}()
#             nw_ref[:branchdc] = Dict{String, Any}()
#             # DC grid arcs for DC grid branches
#             nw_ref[:arcs_dcgrid] = Dict{String, Any}()
#             nw_ref[:arcs_dcgrid_from] = Dict{String, Any}()
#             nw_ref[:arcs_dcgrid_to] = Dict{String, Any}()
#             nw_ref[:arcs_conv_acdc] = Dict{String, Any}()
#             nw_ref[:bus_arcs_dcgrid] = Dict{String, Any}()
#             nw_ref[:arcs_reserves_syn] = Dict{String, Any}()
#             bus_convs_ac = Dict([(i, []) for (i,bus) in nw_ref[:bus]])
#             for (i,conv) in nw_ref[:convdc]
#                 push!(bus_convs_ac[conv["busac_i"]], i)
#             end
#             nw_ref[:bus_convs_ac] = bus_convs_ac
#             nw_ref[:bus_convs_dc] = Dict{String, Any}()
#             nw_ref[:ref_buses_dc] = Dict{String, Any}()
#             nw_ref[:buspairsdc] = Dict{String, Any}()
#         end
#     end
# end
function add_ref_dcgrid!(ref::Dict{Symbol,<:Any}, data::Dict{String,<:Any})
    for (n, nw_ref) in ref[:nw]
        if haskey(nw_ref, :convdc)
            #Filter converters & DC branches with status 0 as well as wrong bus numbers
            nw_ref[:convdc] = Dict([x for x in nw_ref[:convdc] if (x.second["status"] == 1 && x.second["busdc_i"] in keys(nw_ref[:busdc]) && x.second["busac_i"] in keys(nw_ref[:bus]))])
            nw_ref[:branchdc] = Dict([x for x in nw_ref[:branchdc] if (x.second["status"] == 1 && x.second["fbusdc"] in keys(nw_ref[:busdc]) && x.second["tbusdc"] in keys(nw_ref[:busdc]))])
            # DC grid arcs for DC grid branches
            nw_ref[:arcs_dcgrid_from] = [(i,branch["fbusdc"],branch["tbusdc"]) for (i,branch) in nw_ref[:branchdc]]
            nw_ref[:arcs_dcgrid_to]   = [(i,branch["tbusdc"],branch["fbusdc"]) for (i,branch) in nw_ref[:branchdc]]
            nw_ref[:arcs_dcgrid] = [nw_ref[:arcs_dcgrid_from]; nw_ref[:arcs_dcgrid_to]]
            nw_ref[:arcs_conv_acdc] = [(i,conv["busac_i"],conv["busdc_i"]) for (i,conv) in nw_ref[:convdc]]
            #bus arcs of the DC grid
            bus_arcs_dcgrid = Dict([(bus["busdc_i"], []) for (i,bus) in nw_ref[:busdc]])
            for (l,i,j) in nw_ref[:arcs_dcgrid]
                push!(bus_arcs_dcgrid[i], (l,i,j))
            end

            nw_ref[:bus_arcs_dcgrid] = bus_arcs_dcgrid
        # for OPF OPF OPF OPF##########

        bus_convs_dirct =  Dict([(reserve["syncarea"], []) for (i,reserve) in nw_ref[:reserves]])
        #to determine which converters are connected to particular dc line contingency
        for (b,br) in nw_ref[:branchdc]
            if br["rateA"] == 0
                for (c, cv) in nw_ref[:convdc]
                    if cv["busdc_i"] == br["fbusdc"] ||  cv["busdc_i"] == br["tbusdc"]
                        push!(bus_convs_dirct[cv["syncarea"]], c)
                    end
                end
            end
        end

        nw_ref[:bus_convs_dirct] =   bus_convs_dirct
        display("bus_convs_dirct:$n")
        display(bus_convs_dirct)

            # display("nw)ref")
            # display(nw_ref)
            if haskey(nw_ref, :reserves)
            # display("Multinetwork:reserves")

            nw_ref[:arcs_sync_conv] = [(i,conv["busac_i"],conv["syncarea"]) for (i,conv) in nw_ref[:convdc]]
            # display("arcs_sync_conv")
            # display(nw_ref[:load])
            nw_ref[:reserves] = Dict([x for x in nw_ref[:reserves]])

            bus_arcs_load = Dict([(reserve["syncarea"], []) for (i,reserve) in nw_ref[:reserves]])

            for (la,load) in nw_ref[:load]
                    ac_bus = load["load_bus"]
                    # display("la:$la")
                    # display("ac_bus:$ac_bus")
                    for (c,conv) in nw_ref[:convdc]
                        # display("c:$c")
                        if conv["busac_i"] == ac_bus
                            # display("busac_i")
                            # display(conv["busac_i"])
                            # display("syncarea")
                            # display(conv["syncarea"])
                            # display("la")
                            # display(la)
                        push!(bus_arcs_load[conv["syncarea"]], la)
                        end
                    end
                end

            nw_ref[:bus_arcs_load] = bus_arcs_load
            # display("bus_arcs_load")
            #
            # display(bus_arcs_load)

            bus_arcs_conv = Dict([(reserve["syncarea"], []) for (i,reserve) in nw_ref[:reserves]])
            for (l,i,j) in nw_ref[:arcs_sync_conv]
                push!(bus_arcs_conv[j], l)
            end
            nw_ref[:bus_arcs_conv] = bus_arcs_conv
            # display("bus_arcs_conv")
            # display(bus_arcs_conv)

            bus_arcs_sync = Dict([(reserve["syncarea"], []) for (i,reserve) in nw_ref[:reserves]])
            for (l,i,j) in nw_ref[:arcs_sync_conv]
                push!(bus_arcs_sync[j], i)
            end
            nw_ref[:bus_arcs_sync] = bus_arcs_sync
            # display(bus_arcs_sync)

            arcs_reserves_syn = Dict([(i, []) for (i,reserve) in nw_ref[:reserves]])
            for (i,reserve) in nw_ref[:reserves]
                push!(arcs_reserves_syn[reserve["syncarea"]], i)
            end
            nw_ref[:arcs_reserves_syn] = arcs_reserves_syn
            # display("arcs_reserves_syn")
            # display(arcs_reserves_syn)

            bus_syc_arc = Dict([(i, []) for (i,bus) in nw_ref[:bus]])
            for (i,conv) in nw_ref[:convdc]
                 # display(conv)
                 push!(bus_syc_arc[conv["busac_i"]],conv["syncarea"])
            end
            # display(bus_syc_arc)
        end
    # end of  of OPF OPF OPF OPF

             # bus_convs for AC side power injection of DC converters
            bus_convs_ac = Dict([(i, []) for (i,bus) in nw_ref[:bus]])
            for (i,conv) in nw_ref[:convdc]
                push!(bus_convs_ac[conv["busac_i"]], i)
            end
            nw_ref[:bus_convs_ac] = bus_convs_ac

            # bus_convs for AC side power injection of DC converters
            bus_convs_dc = Dict([(bus["busdc_i"], []) for (i,bus) in nw_ref[:busdc]])
            for (i,conv) in nw_ref[:convdc]
                push!(bus_convs_dc[conv["busdc_i"]], i)
            end
            nw_ref[:bus_convs_dc] = bus_convs_dc
            # Add DC reference buses
            ref_buses_dc = Dict{String, Any}()
            for (k,v) in nw_ref[:convdc]
                if v["type_dc"] == 2
                    ref_buses_dc["$k"] = v
                end
            end

            if length(ref_buses_dc) == 0
                for (k,v) in nw_ref[:convdc]
                    if v["type_ac"] == 2
                        ref_buses_dc["$k"] = v
                    end
                end
                Memento.warn(_PM._LOGGER, "no reference DC bus found, setting reference bus based on AC bus type")
            end

            for (k,conv) in nw_ref[:convdc]
                conv_id = conv["index"]
                if conv["type_ac"] == 2 && conv["type_dc"] == 1
                    Memento.warn(_PM._LOGGER, "For converter $conv_id is chosen P is fixed on AC and DC side. This can lead to infeasibility in the PF problem.")
                elseif conv["type_ac"] == 1 && conv["type_dc"] == 1
                    Memento.warn(_PM._LOGGER, "For converter $conv_id is chosen P is fixed on AC and DC side. This can lead to infeasibility in the PF problem.")
                end
                convbus_ac = conv["busac_i"]
                if conv["Vmmax"] < nw_ref[:bus][convbus_ac]["vmin"]
                    Memento.warn(_PM._LOGGER, "The maximum AC side voltage of converter $conv_id is smaller than the minimum AC bus voltage")
                end
                if conv["Vmmin"] > nw_ref[:bus][convbus_ac]["vmax"]
                    Memento.warn(_PM._LOGGER, "The miximum AC side voltage of converter $conv_id is larger than the maximum AC bus voltage")
                end
            end

            if length(ref_buses_dc) > 1
                ref_buses_warn = ""
                for (rb) in keys(ref_buses_dc)
                    ref_buses_warn = ref_buses_warn*rb*", "
                end
                Memento.warn(_PM._LOGGER, "multiple reference buses found, i.e. "*ref_buses_warn*"this can cause infeasibility if they are in the same connected component")
            end

            # if haskey(pm.setting, "find_all_ac_grids") && pm.tetting["find_all_ac_grids"] == true
            #     ACgrids = find_all_ac_grids(nw_ref[:branch], nw_ref[:bus])
            #     for (i, grid) in ACgrids
            #         a = 0
            #         for (j, bus) in nw_ref[:ref_buses]
            #             if (bus["bus_i"] in grid["Buses"])
            #                 a = 1
            #             end
            #         end
            #         if a == 0
            #             Memento.warn(_PM._LOGGER, "Grid $i does not have any voltage reference bus, this might cause infeasibility")
            #         end
            #     end
            # end
            nw_ref[:ref_buses_dc] = ref_buses_dc
            nw_ref[:buspairsdc] = buspair_parameters_dc(nw_ref[:arcs_dcgrid_from], nw_ref[:branchdc], nw_ref[:busdc])
        else
            nw_ref[:convdc] = Dict{String, Any}()
            nw_ref[:busdc] = Dict{String, Any}()
            nw_ref[:branchdc] = Dict{String, Any}()
            nw_ref[:arcs_reserves_syn] = Dict{String, Any}()
            # DC grid arcs for DC grid branches
            nw_ref[:arcs_dcgrid] = Dict{String, Any}()
            nw_ref[:arcs_dcgrid_from] = Dict{String, Any}()
            nw_ref[:arcs_dcgrid_to] = Dict{String, Any}()
            nw_ref[:arcs_conv_acdc] = Dict{String, Any}()
            nw_ref[:bus_arcs_dcgrid] = Dict{String, Any}()
            bus_convs_ac = Dict([(i, []) for (i,bus) in nw_ref[:bus]])
            for (i,conv) in nw_ref[:convdc]
                push!(bus_convs_ac[conv["busac_i"]], i)
            end
            nw_ref[:bus_convs_ac] = bus_convs_ac
            nw_ref[:bus_convs_dc] = Dict{String, Any}()
            nw_ref[:ref_buses_dc] = Dict{String, Any}()
            nw_ref[:buspairsdc] = Dict{String, Any}()
        end
    end
end

"compute bus pair level structures"
function buspair_parameters_dc(arcs_dcgrid_from, branches, buses)
    buspair_indexes = collect(Set([(i,j) for (l,i,j) in arcs_dcgrid_from]))

    bp_branch = Dict([(bp, Inf) for bp in buspair_indexes])

    for (l,branch) in branches
        i = branch["fbusdc"]
        j = branch["tbusdc"]

        bp_branch[(i,j)] = min(bp_branch[(i,j)], l)
    end

    buspairs = Dict([((i,j), Dict(
    "branch"=>bp_branch[(i,j)],
    "vm_fr_min"=>buses[i]["Vdcmin"],
    "vm_fr_max"=>buses[i]["Vdcmax"],
    "vm_to_min"=>buses[j]["Vdcmin"],
    "vm_to_max"=>buses[j]["Vdcmax"]
    )) for (i,j) in buspair_indexes])

    return buspairs
end

function find_all_ac_grids(branches_ac, buses_ac)
    ACgrids = Dict{String, Any}()

    if isempty(branches_ac)
        for (i, bus) in buses_ac
            ACgrids["$i"] = Dict{String, Any}()
            ACgrids["$i"]["Buses"] = bus["index"]
        end
    else
        ACgrids["1"] = Dict{String, Any}()
        ACgrids["1"]["Buses"] = [branches_ac[1]["f_bus"] branches_ac[1]["t_bus"]]
        closed_buses = [branches_ac[1]["f_bus"] branches_ac[1]["t_bus"]]
        closed_branches = [1]
        connections = []
        buses = []
        for (i, bus) in buses_ac
            if VERSION < v"0.7.0-"
                buses = cat(1,buses,bus["index"])
            else
                buses = cat(buses, bus["index"], dims = 1)
            end
        end
        grid_id = 1
        iter_id = 1
        branch_iter = 1
        while length(closed_buses) != length(buses) && iter_id < 10
            while branch_iter <= length(branches_ac)
                for (i, branch) in branches_ac
                    for (index, grid) in ACgrids
                        if (branch["t_bus"] in grid["Buses"]) && (branch["f_bus"] in grid["Buses"])
                            if !(branch["index"] in closed_branches)
                                closed_branches = [closed_branches branch["index"]]
                            end
                        elseif (branch["f_bus"] in grid["Buses"])
                            if !(branch["t_bus"] in grid["Buses"])
                                ACgrids["$index"]["Buses"] = [grid["Buses"] branch["t_bus"]]
                                closed_buses = [closed_buses branch["t_bus"]]
                                closed_branches = [closed_branches branch["index"]]
                            end
                        elseif (branch["t_bus"] in grid["Buses"])
                            if !(branch["f_bus"] in grid["Buses"])
                                ACgrids["$index"]["Buses"] = [grid["Buses"] branch["f_bus"]]
                                closed_buses = [closed_buses branch["f_bus"]]
                                closed_branches = [closed_branches branch["index"]]
                            end
                        end
                    end
                end
                branch_iter = branch_iter + 1
            end
            if length(closed_branches) < length(branches_ac)
                grid_id = grid_id + 1
                branch_iter = 1
                ACgrids["$grid_id"] = Dict{String, Any}()
                for (i, branch) in branches_ac
                    if !(branch["index"] in closed_branches) && isempty(ACgrids["$grid_id"])
                        ACgrids["$grid_id"]["Buses"] = [branch["f_bus"] branch["t_bus"]]
                        closed_branches = [closed_branches branch["index"]]
                    end
                end
            end
            iter_id = iter_id + 1 # to avoid infinite loop -> if not all subgrids detected
        end
    end
    return ACgrids
end

function add_candidate_dcgrid!(ref::Dict{Symbol,<:Any}, data::Dict{String,<:Any})
    for (n, nw_ref) in ref[:nw]
        if haskey(nw_ref, :convdc_ne)
            nw_ref[:arcs_dcgrid_from_ne] = [(i,branch["fbusdc"],branch["tbusdc"]) for (i,branch) in nw_ref[:branchdc_ne]]
            nw_ref[:arcs_dcgrid_to_ne]   = [(i,branch["tbusdc"],branch["fbusdc"]) for (i,branch) in nw_ref[:branchdc_ne]]
            nw_ref[:arcs_dcgrid_ne] = [nw_ref[:arcs_dcgrid_from_ne]; nw_ref[:arcs_dcgrid_to_ne]]
            nw_ref[:arcs_conv_acdc_ne] = [(i,conv["busac_i"],conv["busdc_i"]) for (i,conv) in nw_ref[:convdc_ne]]
            nw_ref[:arcs_conv_acdc_acbus_ne] = [(i,conv["busac_i"]) for (i,conv) in nw_ref[:convdc_ne]]
            nw_ref[:arcs_incident_ne] = [(i, incident["convdc"]) for (i,incident) in nw_ref[:incident]]

            #bus arcs of the DC grid
            bus_arcs_dcgrid_ne = Dict([(bus["busdc_i"], []) for (i,bus) in nw_ref[:busdc_ne]])
            for (l,i,j) in nw_ref[:arcs_dcgrid_ne]
                if haskey(bus_arcs_dcgrid_ne, i)
                    push!(bus_arcs_dcgrid_ne[i], (l,i,j))
                elseif  haskey(nw_ref[:bus_arcs_dcgrid], i)
                    bus_arcs_dcgrid_ne[i] = []
                    push!(bus_arcs_dcgrid_ne[i], (l,i,j))
                else
                    bus_arcs_dcgrid_ne[i] = []
                end
            end
            nw_ref[:bus_arcs_dcgrid_ne] = bus_arcs_dcgrid_ne
            # bus_convs for AC side power injection of DC converters
            bus_convs_ac = Dict([(i, []) for (i,bus) in nw_ref[:bus]])
            for (i,conv) in nw_ref[:convdc_ne]
                push!(bus_convs_ac[conv["busac_i"]], i)
            end
            nw_ref[:bus_convs_ac_ne] = bus_convs_ac

            # add new converters to existting DC buses
            bus_convs_dc_ne = Dict([(bus["busdc_i"], []) for (i,bus) in nw_ref[:busdc]])
            for (i,conv) in nw_ref[:convdc_ne]
                if haskey(bus_convs_dc_ne, conv["busdc_i"])
                    push!(bus_convs_dc_ne, i)
                end
            end
            nw_ref[:bus_convs_dc_ne] = bus_convs_dc_ne
            # Bus converters for candidate DC buses
            bus_ne_convs_dc_ne = Dict([(bus["busdc_i"], []) for (i,bus) in nw_ref[:busdc_ne]])
            for (i,conv) in nw_ref[:convdc_ne]
                if haskey(bus_ne_convs_dc_ne, conv["busdc_i"])
                    push!(bus_ne_convs_dc_ne[conv["busdc_i"]], i)
                end
            end
            nw_ref[:bus_ne_convs_dc_ne] = bus_ne_convs_dc_ne
            nw_ref[:ref_buses_dc_ne] = Dict{String, Any}()
            nw_ref[:buspairsdc_ne] = buspair_parameters_dc_ne(nw_ref[:arcs_dcgrid_from_ne], nw_ref[:branchdc_ne], nw_ref[:busdc_ne], nw_ref[:busdc])
        else
            nw_ref[:convdc_ne] = Dict{String, Any}()
            nw_ref[:busdc_ne] = Dict{String, Any}()
            nw_ref[:branchdc_ne] = Dict{String, Any}()
            # DC grid arcs for DC grid branches
            nw_ref[:arcs_dcgrid_ne] = Dict{String, Any}()
            nw_ref[:arcs_conv_acdc_ne] = Dict{String, Any}()
            nw_ref[:bus_arcs_dcgrid_ne] = Dict{String, Any}()
            bus_convs_ac = Dict([(i, []) for (i,bus) in nw_ref[:bus]])
            for (i,conv) in nw_ref[:convdc_ne]
                push!(bus_convs_ac[conv["busac_i"]], i)
            end
            nw_ref[:bus_convs_ac_ne] = bus_convs_ac
            nw_ref[:bus_convs_dc_ne] = Dict{String, Any}()
            nw_ref[:ref_buses_dc_ne] = Dict{String, Any}()
            nw_ref[:buspairsdc_ne] = Dict{String, Any}()
        end
    end
end

function add_ref_frequency!(ref::Dict{Symbol,<:Any}, data::Dict{String,<:Any})
    for (n, nw_ref) in ref[:nw]
        nw_ref[:reserves] = Dict([x for x in nw_ref[:reserves]])

        bus_convs_dirct =  Dict([(reserve["syncarea"], []) for (i,reserve) in nw_ref[:reserves]])
            #to determine which converters are connected to particular dc line contingency - FS strategies
            for (b,br) in nw_ref[:branchdc]
                if br["rateA"] == 0
                    for (c, cv) in nw_ref[:convdc]
                        if cv["busdc_i"] == br["fbusdc"] ||  cv["busdc_i"] == br["tbusdc"]
                            push!(bus_convs_dirct[cv["syncarea"]], c)
                        end
                    end
                end
            end
         nw_ref[:bus_convs_dirct] =   bus_convs_dirct

         if isempty(bus_convs_dirct) == true
             Memento.error(_LOGGER, "no contigency defined for frequency constraint. set rateA to zero for the line that has contigency")
         end

        nw_ref[:arcs_sync_conv] = [(i,conv["busac_i"],conv["syncarea"]) for (i,conv) in nw_ref[:convdc]]

        bus_arcs_load = Dict([(reserve["syncarea"], []) for (i,reserve) in nw_ref[:reserves]])
         for (la,load) in nw_ref[:load]
            ac_bus = load["load_bus"]
              for (c,conv) in nw_ref[:convdc]
                 if conv["busac_i"] == ac_bus
                 push!(bus_arcs_load[conv["syncarea"]], la)
                 end
             end
         end
         nw_ref[:bus_arcs_load] = bus_arcs_load

         bus_arcs_conv = Dict([(reserve["syncarea"], []) for (i,reserve) in nw_ref[:reserves]])
         for (l,i,j) in nw_ref[:arcs_sync_conv]
             push!(bus_arcs_conv[j], l)
         end
         nw_ref[:bus_arcs_conv] = bus_arcs_conv

         #which reserves belongs to which synchronous area
         arcs_reserves_syn = Dict([(i, []) for (i,reserve) in nw_ref[:reserves]])
         for (i,reserve) in nw_ref[:reserves]
           push!(arcs_reserves_syn[reserve["syncarea"]], i)
         end
         nw_ref[:arcs_reserves_syn] = arcs_reserves_syn
     end
end

function add_ref_frequency_candidates!(ref::Dict{Symbol,<:Any}, data::Dict{String,<:Any})
    for (n, nw_ref) in ref[:nw]
    bus_convs_dirct_ne =  Dict([(reserve["syncarea"], []) for (i,reserve) in nw_ref[:reserves]])
    #to determine which converters are connected to particular dc line contingency
    for (b,br) in nw_ref[:branchdc_ne]
        if br["rateA"] == 0
            for (c, cv) in nw_ref[:convdc_ne]
                if cv["busdc_i"] == br["fbusdc"] ||  cv["busdc_i"] == br["tbusdc"]
                    push!(bus_convs_dirct_ne[cv["syncarea"]], c)
                end
            end
        end
    end
        nw_ref[:bus_convs_dirct_ne] =   bus_convs_dirct_ne

                nw_ref[:arcs_sync_conv_ne] = [(i,conv["busac_i"],conv["syncarea"]) for (i,conv) in nw_ref[:convdc_ne]]
                nw_ref[:reserves] = Dict([x for x in nw_ref[:reserves]])
                bus_arcs_load_ne = Dict([(reserve["syncarea"], []) for (i,reserve) in nw_ref[:reserves]])
                for (la,load) in nw_ref[:load]
                        ac_bus = load["load_bus"]
                        for (c,conv) in nw_ref[:convdc_ne]
                            if conv["busac_i"] == ac_bus
                             push!(bus_arcs_load_ne[conv["syncarea"]], la)
                             end
                        end
                end
                nw_ref[:bus_arcs_load_ne] = bus_arcs_load_ne

                #old
                bus_arcs_conv_ne = Dict([(reserve["syncarea"], []) for (i,reserve) in nw_ref[:reserves]])
                for (l,i,j) in nw_ref[:arcs_sync_conv_ne]
                    push!(bus_arcs_conv_ne[j], l)
                end
                nw_ref[:bus_arcs_conv_ne] = bus_arcs_conv_ne

                #new
                bus_arcs_incident_ne = Dict([(i, []) for (i,branch) in nw_ref[:branchdc_ne]])
                for (i,incident) in nw_ref[:incident]
                    for b = 1: length(incident)-3
                        if incident["br$(b)"] == 1
                        push!(bus_arcs_incident_ne[b], incident["convdc"])
                        end
                    end
                end
                nw_ref[:bus_arcs_incident_ne] = bus_arcs_incident_ne

                bus_arcs_sync_ne = Dict([(reserve["syncarea"], []) for (i,reserve) in nw_ref[:reserves]])
                for (l,i,j) in nw_ref[:arcs_sync_conv_ne]
                    push!(bus_arcs_sync_ne[j], i)
                end
                nw_ref[:bus_arcs_sync_ne] = bus_arcs_sync_ne

                arcs_reserves_syn_ne = Dict([(i, []) for (i,reserve) in nw_ref[:reserves]])
                for (i,reserve) in nw_ref[:reserves]
                    push!(arcs_reserves_syn_ne[reserve["syncarea"]], i)
                end
                nw_ref[:arcs_reserves_syn_ne] = arcs_reserves_syn_ne

                bus_syc_arc_ne = Dict([(i, []) for (i,bus) in nw_ref[:bus]])
                for (i,conv) in nw_ref[:convdc_ne]
                     push!(bus_syc_arc_ne[conv["busac_i"]],conv["syncarea"])
                end
            end
end

"compute bus pair level structures"
function buspair_parameters_dc_ne(arcs_dcgrid_from, branches, busesdc_ne, busesdc)
    buspair_indexes = collect(Set([(i,j) for (l,i,j) in arcs_dcgrid_from]))

    bp_branch = Dict([(bp, Inf) for bp in buspair_indexes])
    for (l,branch) in branches
        i = branch["fbusdc"]
        j = branch["tbusdc"]
        bp_branch[(i,j)] = min(bp_branch[(i,j)], l)
    end
    buspairs = Dict([((i,j), Dict{String,Any}(
    "branch" => nothing,
    "vm_fr_min" => nothing,
    "vm_fr_max" => nothing,
    "vm_to_min" => nothing,
    "vm_to_max" => nothing
    )) for (i,j) in buspair_indexes])

    for (i,j) in buspair_indexes
        buspairs[(i,j)]["branch"] = bp_branch[(i,j)]
        for (k, bus) in busesdc_ne
            if i == bus["busdc_i"]
                buspairs[(i,j)]["vm_fr_min"] = bus["Vdcmin"]
                buspairs[(i,j)]["vm_fr_max"] = bus["Vdcmax"]
            end
            if j == bus["busdc_i"]
                buspairs[(i,j)]["vm_to_min"] = bus["Vdcmin"]
                buspairs[(i,j)]["vm_to_max"] = bus["Vdcmax"]
            end
        end
        for (k, bus) in busesdc
            if i == bus["busdc_i"]
                buspairs[(i,j)]["vm_fr_min"] = bus["Vdcmin"]
                buspairs[(i,j)]["vm_fr_max"] = bus["Vdcmax"]
            end
            if j == bus["busdc_i"]
                buspairs[(i,j)]["vm_to_min"] = bus["Vdcmin"]
                buspairs[(i,j)]["vm_to_max"] = bus["Vdcmax"]
            end
        end
    end

    return buspairs
end


module BlockMatching

using OffsetArrays
using Distances
using Requires

export FullSearch, best_match, multi_match

abstract type AbstractBlockMatchingStrategy end

"""
    best_match(S::AbstractBlockMatchingStrategy, ref, [frame=ref]; offset=false)
    best_match(S::AbstractBlockMatchingStrategy, ref, [frame=ref], p; offset=false)

For given pixel `p` in `frame`, find the best match pixel `q` in `ref` with block matching strategy
`S`. If `p` is not provided, it operates on all pixels `p ∈ CartesianIndices(ref)` and returns an array.

!!! note
    For clarification, "pixel" here means the position `CartesianIndex` instead of the value `Colorant`.

# Arguments

* `S::AbstractBlockMatchingStrategy`: required
  The concrete block matching strategy and relevant configs, e.g., patch size, search window size.
  See `subtypes(BlockMatching.AbstractBlockMatchingStrategy)` for possible strategies.
* `ref`::AbstractArray: required
  Reference image where the matched pixel `q` belongs to.
* `frame::AbstractArray`: optional
  The current image/frame where the input pixel `p` belongs to. If `frame` is not provided, it will
  be `ref`.
* `p::CartesianIndex`: optional
  The given pixel `p` that block matching operates on. If not provided, block matching is operating
  on the whole image and thus returns an array of pixels `q`s.

# Parameters

* `offset::Bool`: optional
  if `true`, it returns the motion vector `q-p` instead of the pixel `q`. The default value is
  `false`.

# Output

If `p` is given, it returns `CartesianIndex{N}`. Otherwise the block matching operates on
the whole image and returns `Array{CartesianIndex{N}, N}`, i.e., the motion field.

# Examples

If pixel `p` is provided, block matching only operates on the given pixel, the output is the matched
pixels `q`s for `p`.

```jldoctest best_match
using Images, TestImages
using BlockMatching

ref = imresize(testimage("cameraman"), (64, 64))
img = imrotate(ref, 0.2, CartesianIndices(ref).indices)

S = FullSearch(SqEuclidean(), ndims(img), patch_radius=5, search_radius=11)
p = CartesianIndex(17, 17)
best_match(S, img, ref, p)

# output

CartesianIndex(17, 14)
```

If `p` is not provided, the block matching operates on the whole image. Each item of the output array
is the matched pixels `q`s for `p`, i.e., `best_match(S, ref, frame)[p] = best_match(S, ref, frame, p)`.

```jldoctest best_match
julia> matches = best_match(S, img, ref);

julia> summary(matches)
"64×64 $(Matrix{CartesianIndex{2}})"

julia> matches[p] == best_match(S, img, ref, p)
true
```

!!! tip

    `best_match(S, ref, frame)` is usually more performant than `map(p->best_match(S, ref, frame,
    p), CartesianIndices(frame))` because intermediate results can be cached/pre-allocated so as
    to reduce unnecessary computation and memory allocation.

# See also

* [`multi_match`](@ref) can be used when multiple candidates are wanted.
"""
best_match(S::AbstractBlockMatchingStrategy, ref; kwargs...) = best_match(S, ref, ref; kwargs...)
best_match(S::AbstractBlockMatchingStrategy, ref, p::CartesianIndex; kwargs...) = best_match(S, ref, ref, p; kwargs...)

"""
    multi_match(S::AbstractBlockMatchingStrategy, ref, [frame=ref]; num_patches, offset=true)
    multi_match(S::AbstractBlockMatchingStrategy, ref, [frame=ref], p::CartesianIndex; num_patches, offset=true)

For given pixel `p` in `frame`, find matched pixels `q`s in `ref` with block matching strategy `S`.

# Arguments

* `S::AbstractBlockMatchingStrategy`: required
  The concrete block matching strategy and relevant configs, e.g., similarity mesaure, patch size,
  search window size. See `subtypes(BlockMatching.AbstractBlockMatchingStrategy)` for possible
  strategies.
* `ref`::AbstractArray: required
  Reference image where the matched pixel `q` belongs to.
* `frame::AbstractArray`: optional
  The current image/frame where the input pixel `p` belongs to. If `frame` is not provided, it will
  be `ref`.
* `p::CartesianIndex`: optional
  The given pixel `p` that block matching operates on. If not provided, block matching is operating
  on the whole image and thus returns an array of pixels `q`s.

# Parameters

* `num_patches`: required
  The number of matched pixels. If the candidates number is smaller than `num_patches`, the behavior
  is undefined; it may or may not throw errors.
* `offset::Bool`: optional
  if `true`, it returns the motion vector `q-p` instead of the pixel `q`. The default value is
  `false`.

# Output

If `p` is given, it returns `Vector{CartesianIndex{N}}`. Otherwise the block matching operates on
the whole image and returns `Array{CartesianIndex{N}, N+1}`, where the extra N+1 dimension stores
the matched vector.

# Examples

If pixel `p` is provided, block matching only operates on the given pixel, the output is the matched
pixels `q`s for `p`.

```jldoctest multi_match
using Images, TestImages
using BlockMatching

ref = imresize(testimage("cameraman"), (64, 64))
img = imrotate(ref, 0.2, CartesianIndices(ref).indices)

S = FullSearch(SqEuclidean(), ndims(img), patch_radius=5, search_radius=11)
p = CartesianIndex(17, 17)
multi_match(S, img, ref, p; num_patches=2)

# output

2-element $(Vector{CartesianIndex{2}}):
 CartesianIndex(17, 14)
 CartesianIndex(18, 13)
```

Otherwise, the block matching operates on the whole image. Each item of the output array is the
matched pixels `q`s for `p`, i.e., `multi_match(S, ref, frame)[p] = multi_match(S, ref, frame, p)`.

```jldoctest multi_match
julia> matches = multi_match(S, img, ref; num_patches=2);

julia> summary(matches)
"64×64 $(Matrix{Vector{CartesianIndex{2}}})"

julia> matches[p, :] == multi_match(S, img, ref, p; num_patches=2)
true
```

!!! tip

    `multi_match(S, ref, frame)` is usually more performant than `map(p->multi_match(S, ref, frame,
    p), CartesianIndices(frame))` because intermediate results can be cached so as to reduce
    computation and memory allocation.

# See also

[`best_match`](@ref) is more performant a choice when `num_patches==1`. It also supports more
sophisticated strategies, e.g., ThreeStepSearch.
"""
multi_match(S::AbstractBlockMatchingStrategy, ref; kwargs...) = multi_match(S, ref, ref; kwargs...)
multi_match(S::AbstractBlockMatchingStrategy, ref, p::CartesianIndex; kwargs...) = multi_match(S, ref, ref, p; kwargs...)

include("utils.jl")
include("strategies/full_search.jl")

function __init__()
    @require CUDA="052768ef-5323-5732-b1bb-66c8b64840ba" begin
        include("strategies/full_search_cuda.jl")
    end
end

end


#= List of Variables

tr      trace dataset
wv      waveform dataset
sta     station dataset
idₛ     station id
idₜ     trace id
stas    list of station names
traces  list of trace names
trₛ     names of traces in sta

idⱼ     S[j] :id
iⱼ      S[j] first index in :x to write
jⱼ      S[j] last index in :x to write
lxⱼ     S[j] value of :t[end,1]
ntⱼ     S[j] number of times in :t[:,1]
nⱼ      S[j] length(:x)
tⱼ      S[j] :t

eᵣ      request window end time [ns from epoch]
nᵣ      request window length in samples
sᵣ      request window start time [ns from epoch]

eₜ      trace end time [ns from epoch]
fs      trace sampling frequency [Hz]
iₜ      trace first index to read
jₜ      trace last index to read
nₜ      trace total number of samples
sₜ      trace start time [ns from epoch]

nₒ      overlap window number of samples
sₒ      overlap window start time [ns from epoch]

j         index to channel in S
v         verbosity
Δ         sampling interval [ns]

because I managed to confuse even myself with the bookkeeping, and that's not
happened since coding SeisIO began
=#

function read_asdf!(  S::GphysData,
                      hdf::String,
                      id::String,
                      s::TimeSpec,
                      t::TimeSpec,
                      msr::Bool,
                      v::Integer  )

  SX = SeisData() # for XML
  idr = isa(id, String) ? id_to_regex(id) : id
  (v > 2) && println("Reading IDs that match ", idr)

  if typeof(s) == String && typeof(t) == String
    d0 = s
    d1 = t
    sᵣ = DateTime(s).instant.periods.value*1000 - dtconst
    eᵣ = DateTime(t).instant.periods.value*1000 - dtconst
  else
    (d0, d1) = parsetimewin(s, t)
    sᵣ = DateTime(d0).instant.periods.value*1000 - dtconst
    eᵣ = DateTime(d1).instant.periods.value*1000 - dtconst
  end
  sᵣ *= 1000
  eᵣ *= 1000
  Δ = 0
  fs = 0.0

  # this nesting is a mess
  netsta = netsta_to_regex(id)
  idr = id_to_regex(id)
  f = h5open(hdf, "r")
  wv = f["Waveforms"]
  stas = names(wv)
  sort!(stas)
  (v > 2) && println("Net.sta found: ", stas)

  for idₛ in stas
    if occursin(netsta, idₛ)
      sta = wv[idₛ]
      traces = names(sta)
      sort!(traces)
      (v > 2) && println("Traces found: ", traces)
      for idₜ in traces
        if idₜ == "StationXML"
          sxml = String(UInt8.(read(sta[idₜ])))
          read_station_xml!(SX, sxml, d0, d1, msr, v)
        elseif occursin(idr, idₜ)
          tr = sta[idₜ]
          nₜ = length(tr)
          sₜ = read(tr["starttime"])
          fs = read(tr["sampling_rate"])

          # convert fs to sampling interval in ns
          Δ = round(Int64, 1.0e9/fs)
          eₜ = sₜ + (nₜ-1)*Δ
          (v > 2) && println("sₜ = ", sₜ,"; eₜ = ", eₜ, "; sᵣ = ", sᵣ, "; eᵣ = ", eᵣ)

          if (sᵣ ≤ eₜ) && (eᵣ ≥ sₜ)
            lxⱼ = 0
            iₜ, jₜ, sₒ = get_trace_bounds(sᵣ, eᵣ, sₜ, eₜ, Δ, nₜ)
            nₒ = jₜ-iₜ+1
            idⱼ = String(split(idₜ, "_", limit=2, keepempty=true)[1])
            j = findid(idⱼ, S.id)
            (v > 2) && println("idⱼ = ", idⱼ, " (found at index in S = ", j, ")")
            nᵣ = div(eᵣ-sᵣ, Δ)+1
            if j == 0
              T = eltype(tr)
              push!(S, SeisChannel(id = idⱼ,
                                   fs = fs,
                                   x = Array{T,1}(undef, nᵣ)))
              j = S.n
              if has(tr, "event_id")
                S.misc[j]["event_id"] = read(tr["event_id"])
              end
            else
              tⱼ = getindex(getfield(S, :t), j)
              ntⱼ = div(lastindex(tⱼ), 2)
              nⱼ = lastindex(S.x[j])
              if ntⱼ > 0
                lxⱼ = getindex(tⱼ, ntⱼ)
                check_for_gap!(S, j, div(sₒ, 1000), nₒ, v)
              end
              if lxⱼ + nₒ > nⱼ
                resize!(S.x[j], lxⱼ + max(nₒ, nᵣ))
              end
              if S.fs[j] == 0.0
                S.fs[j] = fs
              end
            end

            if lxⱼ == 0
              tⱼ = mk_t(nₒ, div(sₒ, 1000))
              setindex!(getfield(S, :t), tⱼ, j)
            end

            iⱼ = lxⱼ+1
            jⱼ = iⱼ+nₒ-1
            load_data!(S.x[j], tr, iₜ:jₜ, iⱼ:jⱼ)
          end
          HDF5.h5d_close(tr)
        end
      end
      HDF5.h5g_close(sta)
    end
  end

  # merge in the XML that we read
  sxml_mergehdr!(S, SX, false, true, v)
  trunc_x!(S)

  # Done
  HDF5.h5g_close(wv)
  close(f)
  return S
end

function read_asdf(hdf::String, id::Union{String,Regex}, s::TimeSpec, t::TimeSpec, msr::Bool, v::Integer)
  S = SeisData()
  read_asdf!(S, hdf, id, s, t, msr, v)
  return S
end


using MAT

## Colors
C = [1 0 0;
0 .75 0;
0 0 1;
0 .5 .5;
.5 0 .5;
1 .5 0;
0 0 0;
.75 .75 0;
0.1 0.5 0.6;
0 .75 .75;
.4 .5 .3]

using Plots
# Must be kept in this order, unless you are careful what you do
labels = [1; 2; 3; 12; 18; 17; 25; 15; 24]
names = ["Amazon Fashion", "All Beauty", "Appliances",
"Gift Cards", "Magazine Subsciptions", "Luxury Beauty", "Software",
"Industrial and Scientific","Prime Pantry" ]
sizes = [31; 85; 48; 148; 157; 1581; 802; 5334; 4970]
lnum = length(labels)

# See outer parameters
seednum = 10   # 5
grownum = 300  #[200; 300; 500]
epsis = [10.0 1.0 0.1 0.01 0.001]

plot()
for lab = [1;3;4;5]                 # For each cluster
label = labels[lab]
outputmat = "Output_VaryEps/Label_$label"*"_$seednum"*"_$grownum.mat"

mat = matread(outputmat)
hl_time = mean(mat["hl_time"],dims = 2)
hl = round.(mean(mat["hl_f1"],dims = 2),digits = 2)
b1 = round.(mat["b1_f1"],digits = 2)
b2 = round.(mat["b2_f1"],digits = 2)
r = round.(mat["r_f1"],digits = 2)
re = round.(mat["r_re"],digits = 2)
pr = round.(mat["r_pr"],digits = 2)
nS = mat["newS"]

lw = 3
plot!(epsis',hl,linewidth = lw, grid = false, markershape =:circle,
    legend = :bottomleft, color = RGBA(C[lab,1],C[lab,2],C[lab,3],1),
    xaxis=:log10, label = "label $(label)")

end
savefig("Plots/Smallest_VaryEps_seednum_$seednum"*"_grownum_$grownum.pdf")


using Random
using CUDAdrv
using CUDAnative
using CuArrays
using StaticArrays
using GPUifyLoops # for @unroll macro

function volumerhs!(::Val{N}, rhs, Q, nelem) where {N}
    Nq = N + 1

    s_F = @cuStaticSharedMem eltype(Q) (Nq, Nq)

    r_rhsρ = MArray{Tuple{Nq},eltype(rhs)}(undef)

    e = blockIdx().x
    j = threadIdx().y
    i = threadIdx().x

    @inbounds @unroll for k in 1:Nq
        r_rhsρ[k] = zero(eltype(rhs))
    end

    @inbounds @unroll for k in 1:Nq
        s_F[j, i] = Q[i, j, k, e]

        sync_threads()

        # loop of ξ-grid lines
        @unroll for n = 1:Nq
            r_rhsρ[k] += s_F[j, n]
            r_rhsρ[k] += s_F[i, n]
            r_rhsρ[k] += s_F[n, j]
            r_rhsρ[k] += s_F[n, i]
        end
        sync_threads()
    end

    @inbounds @unroll for k in 1:Nq
        rhs[i, j, k, e] += r_rhsρ[k]
    end

    nothing
end

function main()
    N = 4
    nelem = 4000
    ntrials = 1
    DFloat = Float32

    rnd = MersenneTwister(0)

    Nq = N + 1
    Q = 1 .+ CuArray(rand(rnd, DFloat, Nq, Nq, Nq, nelem))

    rhs = CuArray(zeros(DFloat, Nq, Nq, Nq, nelem))

    @cuda(threads=(N+1, N+1), blocks=nelem, maxregs=255,
          volumerhs!(Val(N), rhs, Q, nelem))

    CUDAdrv.@profile  begin
        @cuda(threads=(N+1, N+1), blocks=nelem, maxregs=255,
              volumerhs!(Val(N), rhs, Q, nelem))
    end

    CUDAnative.@device_code dir="jlir" begin
        @cuda(threads=(N+1, N+1), blocks=nelem, maxregs=255,
              volumerhs!(Val(N), rhs, Q, nelem))
    end

    nothing
end

main()


function ftest()
   unpack_expr = Meta.parse(string(join(string.(p), ","), ", = a"))
   G = :(a1*t + a2)
   J = :(a2*t + a1)
   @eval return(TT(
      ((t, a) -> ($unpack_expr; $G)) |> FunctionWrapper{Float64, Tuple{Float64, Vector{Float64}}},
      ((t, a) -> ($unpack_expr; $J)) |> FunctionWrapper{Float64, Tuple{Float64, Vector{Float64}}},

      ((ta, a) -> ($unpack_expr; [$G for t in ta])) |> FunctionWrapper{Vector{Float64}, Tuple{Vector{Float64}, Vector{Float64}}},
      ((ta, a) -> ($unpack_expr; [$J for t in ta])) |> FunctionWrapper{Vector{Float64}, Tuple{Vector{Float64}, Vector{Float64}}},
      ))
end

macro mtest()
   # test base case
   p=(:a1, :a2)
   unpack_expr = Meta.parse(string(join(string.(p), ","), ", = a"))
   G = :(a1*t + a2)
   J = :(a2*t + a1)
   return quote TT(
      ((t, a) -> ($unpack_expr; $G)) |> FunctionWrapper{Float64, Tuple{Float64, Vector{Float64}}},
      ((t, a) -> ($unpack_expr; $J)) |> FunctionWrapper{Float64, Tuple{Float64, Vector{Float64}}},

      ((ta, a) -> ($unpack_expr; [$G for t in ta])) |> FunctionWrapper{Vector{Float64}, Tuple{Vector{Float64}, Vector{Float64}}},
      ((ta, a) -> ($unpack_expr; [$J for t in ta])) |> FunctionWrapper{Vector{Float64}, Tuple{Vector{Float64}, Vector{Float64}}},
     ) end
end

ftest()
@mtest


include("utils.jl")

struct TouchstoneSnP
    frequency_unit::String
    parameter_type::String
    parameter_format::String
    reference_resistance::Int # [Ohm]
    number_of_ports::Int
    number_of_frequencies::Int
    number_of_noise_frequencies::Int
    reference::Array{Float64}
    two_port_data_order::String
    matrix_format::String
    network_data::Dict#{Array{Float64}, 1}
    noise_data::Dict#{Array{Float64}, 1}

    function TouchstoneSnP(snp::AbstractString)

        sentences = map(uppercase, readlines(snp))
        # Default parameters
        frequency_unit = "GHZ"
        parameter_type = "S"
        parameter_format = "MA"
        reference_resistance = 50 # [Ohm]
        number_of_ports = 2 # Required
        number_of_frequencies = 2 # Required
        number_of_noise_frequencies = 0 # Required if noise data is present
        reference = []
        two_port_data_order = ""
        matrix_format = "FULL"
        network_data = Dict()
        noise_data = Dict()
        t_network = []
        t_noise = []

        i = 1
        i = skip_comments(sentences, i)
        m = match(r"^\s{0,}\[VERSION\]\s(2\.0)", sentences[i])[1]
        if m != "2.0"
            throw("Only version 2.0 is currently supported")
        end
        i = skip_comments(sentences, i+1)

        # Parse option line
        freq_unit = match(r"((?:GHZ)|(?:KHZ)|(?:MHZ)|(?:HZ))", sentences[i])
        if !(freq_unit === nothing)
            frequency_unit = freq_unit[1][1:end-1]*"z"
        end

        param_type = match(r"\s(S|Y|Z|H|G)\s", sentences[i])
        if !(param_type === nothing)
            parameter_type = param_type[1]
        end

        param_format = match(r"((?:DB)|(?:MA)|(?:RI))", sentences[i])
        if !(param_format === nothing)
            parameter_format = param_format[1]
        end

        ref_resistance = match(r"R (\d+)", sentences[i])

        if !(ref_resistance === nothing)
            try
                reference_resistance = parse(Int, ref_resistance[1])
            catch
                throw("Incorrect number of reference resistance")
            end
        end

        i = skip_comments(sentences, i+1)

        # Parse keywords that don't appear in a specific order
        while true
            keyword = match(r"^\s*\[(.+)\]", sentences[i])[1]
            if keyword == "NUMBER OF PORTS"
                number_of_ports = parse(Int, match(r"^\[NUMBER OF PORTS\]\s+(\d+)", sentences[i])[1])
            elseif keyword == "REFERENCE"
                reference_resistance = parse(Int, match(r"(\d+)", sentences[i])[1])
            elseif keyword == "TWO-PORT ORDER"
                two_port_data_order = match(r"^\[TWO-PORT DATA ORDER\]\s+((?:12_21)|(?:21_12))", sentences[i])[1]
            elseif keyword == "NUMBER OF FREQUENCIES"
                number_of_frequencies = parse(Int, match(r"^\[NUMBER OF FREQUENCIES\]\s+(\d+)", sentences[i])[1])
            elseif keyword == "NUMBER OF NOISE FREQUENCIES"
                number_of_noise_frequencies = parse(Int, match(r"^\[NUMBER OF NOISE FREQUENCIES\]\s+(\d+)", sentences[i])[1])
            elseif keyword == "REFERENCE"
                # TODO: Think about a more elegant way to do this
                reference = []
                while true
                   for k in split(split(sentences[i], "!")[1])
                        if k != "[REFERENCE]"
                            append!(reference, parse(Float64, k))
                        end
                    end
                    if length(reference) >= number_of_ports
                        break
                    end
                    i = skip_comments(sentences, i+1)
                end
            elseif keyword == "MATRIX FORMAT"
                matrix_format = match(r"((?:FULL)|(?:LOWER)|(?:UPPER))", sentences[i])[1]
            elseif keyword == "NETWORK DATA"
                # This should always come at the end along with network noise data
                network_data_len = number_of_frequencies * ((matrix_format=="FULL") ? 2*number_of_ports^2 + 1 : number_of_ports^2 + number_of_ports + 1)
                i = skip_comments(sentences, i+1)
                while true
                    t_network = vcat(t_network, parse.(Float64, split(split(sentences[i], "!")[1])))
                    (length(t_network) < network_data_len) || break
                    i = skip_comments(sentences, i+1)
                end
            elseif keyword == "NOISE DATA"
                # This should always come at the end along with network data
                noise_data_len = number_of_noise_frequencies * ((matrix_format=="FULL") ? 2*number_of_ports^2 + 1 : number_of_ports^2 + number_of_ports + 1)

                i = skip_comments(sentences, i+1)
                while true
                    t_noise = vcat(t_noise, parse.(Float64, split(split(sentences[i], "!")[1])))
                    (length(t_noise) < noise_data_len) || break
                    i = skip_comments(sentences, i+1)
                end
            end
        try
            i = skip_comments(sentences, i+1)
        catch
            break
        end

        end

        network_data = transform_vector_to_matrix_mode(t_network, String(matrix_format), number_of_frequencies)
        if number_of_noise_frequencies > 0
            noise_data = transform_vector_to_matrix_mode(t_noise, String(matrix_format), number_of_noise_frequencies)
        end

        new(frequency_unit, parameter_type, parameter_format, reference_resistance, number_of_ports,
        number_of_frequencies, number_of_noise_frequencies, reference, two_port_data_order, matrix_format,
        network_data, noise_data)
    end
end

"""
Converts Network/Noise data to complex matrix
...
# Arguments
- `snp::TouchstonerSnP`: provided Touchstone file
- `frequency::Int`: target frequency (must be specified in snp)
- `noise_data::Bool` if set to `true` returns results for Noise Data. Else,
    Network Data are assumed
...
"""
function get_ri_form(snp::TouchstoneSnP, frequency::Float64, noise_data::Bool=false)
data = noise_data ? snp.noise_data : snp.network_data
if !haskey(data, frequency)
    throw("Frequency not specified in data")
end

parameters = data[frequency]
n, m = size(parameters)
if snp.parameter_format == "RI"
    return [to_complex([i, 2j - 1:2j]) for i = 1:n, j = 1:n]
elseif snp.parameter_format == "MA"
   return [to_complex(ma_to_ri(parameters[i, 2j - 1:2j])) for i = 1:n, j = 1:n]
elseif snp.parameter_format == "DB"
    return [to_complex(ma_to_ri(db_to_ma(parameters[i, 2j - 1:2j]))) for i = 1:n, j = 1:n]
end
end

"""
Converts Network/Noise parameters to specified parameters
...
# Arguments
- `snp::TouchstonerSnP`: provided Touchstone file
- `frequency::Int`: target frequency (must be specified in snp)
- `parameter_type`::Str : parameter typ to convert to (source type being the  one specified in SnP)
- `noise_data::Bool` if set to `true` returns results for Noise Data. Else,
    Network Data are assumed
...
"""
function convert_param_type(snp::TouchstoneSnP, frequency::Float64, parameters_type::String, noise_data::Bool=false)
    data = noise_data ? snp.noise_data : snp.network_data
    if !haskey(data, frequency)
        throw("Frequency not specified in data")
    end

    # Convert to R-I format for easier conversion
    parameters = get_ri_form(snp, frequency, noise_data)
    # todo: Mixed-port parameters
    # Also I don't like this at all :P
    if parameters_type == "S"
        if snp.parameter_type == "S"
            return parameters
        elseif snp.parameter_type == "Y"
            return admittance_params_to_scatter(parameters, snp.reference_resistance)
        elseif snp.parameter_type == "Z"
            return impedance_params_to_scatter(parameters, snp.reference_resistance)
        end
    elseif parameters_type == "Y"
        if snp.parameter_type == "S"
            return scatter_params_to_admittance(parameters, snp.reference_resistance)
        elseif snp.parameter_type == "Y"
            return parameters
        elseif snp.parameter_type == "Z"
            return inv(parameters)
        end
    elseif parameters_type == "Z"
        if snp.parameter_type == "S"
            return scatter_params_to_impedance(parameters, snp.reference_resistance)
        elseif snp.parameter_type == "Y"
            return inv(parameters)
        elseif snp.parameter_type == "Z"
            return parameters
    end
end

end


module MazeUndirectedModule

import ..GridWorlds as GW
import Random
import ReinforcementLearningBase as RLBase

#####
##### game logic
#####

const NUM_OBJECTS = 3
const AGENT = 1
const WALL = 2
const GOAL = 3
const NUM_ACTIONS = 4

mutable struct MazeUndirected{R, RNG} <: GW.AbstractGridWorld
    tile_map::BitArray{3}
    agent_position::CartesianIndex{2}
    reward::R
    rng::RNG
    done::Bool
    terminal_reward::R
    goal_position::CartesianIndex{2}
end

function MazeUndirected(; R = Float32, height = 9, width = 9, rng = Random.GLOBAL_RNG)
    @assert isodd(height) && isodd(width) "height and width must be odd numbers"

    tile_map = falses(NUM_OBJECTS, height, width)

    tile_map[WALL, :, :] .= true
    for j in 2 : 2 : width - 1
        for i in 2 : 2 : height - 1
            tile_map[WALL, i, j] = false
        end
    end

    agent_position = CartesianIndex(2, 2)
    tile_map[AGENT, agent_position] = true

    goal_position = CartesianIndex(height - 1, width - 1)
    tile_map[GOAL, goal_position] = true

    reward = zero(R)
    done = false
    terminal_reward = one(R)

    env = MazeUndirected(tile_map, agent_position, reward, rng, done, terminal_reward, goal_position)

    GW.reset!(env)

    return env
end

function GW.reset!(env::MazeUndirected)
    tile_map = env.tile_map
    rng = env.rng
    _, height, width = size(tile_map)

    tile_map[AGENT, env.agent_position] = false
    tile_map[GOAL, env.goal_position] = false

    tile_map[WALL, :, :] .= true
    for j in 2 : 2 : width - 1
        for i in 2 : 2 : height - 1
            tile_map[WALL, i, j] = false
        end
    end

    generate_maze!(env)

    new_agent_position = GW.sample_empty_position(rng, tile_map)
    env.agent_position = new_agent_position
    tile_map[AGENT, new_agent_position] = true

    new_goal_position = GW.sample_empty_position(rng, tile_map)
    env.goal_position = new_goal_position
    tile_map[GOAL, new_goal_position] = true

    env.reward = zero(env.reward)
    env.done = false

    return nothing
end

function GW.act!(env::MazeUndirected, action)
    @assert action in Base.OneTo(NUM_ACTIONS) "Invalid action $(action). Action must be in Base.OneTo($(NUM_ACTIONS))"

    tile_map = env.tile_map
    agent_position = env.agent_position

    if action == 1
        new_agent_position = GW.move_up(agent_position)
    elseif action == 2
        new_agent_position = GW.move_down(agent_position)
    elseif action == 3
        new_agent_position = GW.move_left(agent_position)
    else
        new_agent_position = GW.move_right(agent_position)
    end

    if !tile_map[WALL, new_agent_position]
        tile_map[AGENT, agent_position] = false
        env.agent_position = new_agent_position
        tile_map[AGENT, new_agent_position] = true
    end

    if tile_map[GOAL, env.agent_position]
        env.reward = env.terminal_reward
        env.done = true
    else
        env.reward = zero(env.reward)
        env.done = false
    end

    return nothing
end

function generate_maze!(env::MazeUndirected)
    tile_map = env.tile_map
    rng = env.rng
    _, height, width = size(tile_map)

    vertical_range = 2:2:height-1
    horizontal_range = 2:2:width-1

    visited = falses(height, width)
    stack = CartesianIndex{2}[]

    first_cell = CartesianIndex(rand(rng, vertical_range), rand(rng, horizontal_range))
    visited[first_cell] = true
    push!(stack, first_cell)

    valid_neighbors = CartesianIndex{2}[]

    while !isempty(stack)
        current_cell = pop!(stack)

        i, j = current_cell.I
        all_neighbors = (CartesianIndex(i - 2, j), CartesianIndex(i, j - 2), CartesianIndex(i, j + 2), CartesianIndex(i + 2, j))

        empty!(valid_neighbors)
        for neighbor in all_neighbors
            if (neighbor.I[1] in vertical_range) && (neighbor.I[2] in horizontal_range) && !visited[neighbor]
                push!(valid_neighbors, neighbor)
            end
        end

        if length(valid_neighbors) > 0
            push!(stack, current_cell)

            next_cell = rand(rng, valid_neighbors)

            mid_pos = CartesianIndex((current_cell.I .+ next_cell.I) .÷ 2)
            tile_map[WALL, mid_pos] = false

            visited[next_cell] = true
            push!(stack, next_cell)
        end
    end

    return nothing
end

#####
##### miscellaneous
#####

GW.get_height(env::MazeUndirected) = size(env.tile_map, 2)
GW.get_width(env::MazeUndirected) = size(env.tile_map, 3)

GW.get_action_names(env::MazeUndirected) = (:MOVE_UP, :MOVE_DOWN, :MOVE_LEFT, :MOVE_RIGHT)
GW.get_object_names(env::MazeUndirected) = (:AGENT, :WALL, :GOAL)

function GW.get_pretty_tile_map(env::MazeUndirected, position::CartesianIndex{2})
    characters = ('☻', '█', '♥', '⋅')

    object = findfirst(@view env.tile_map[:, position])
    if isnothing(object)
        return characters[end]
    else
        return characters[object]
    end
end

function GW.get_pretty_sub_tile_map(env::MazeUndirected, window_size, position::CartesianIndex{2})
    tile_map = env.tile_map
    agent_position = env.agent_position

    characters = ('☻', '█', '♥', '⋅')

    sub_tile_map = GW.get_sub_tile_map(tile_map, agent_position, window_size)

    object = findfirst(@view sub_tile_map[:, position])
    if isnothing(object)
        return characters[end]
    else
        return characters[object]
    end
end

function Base.show(io::IO, ::MIME"text/plain", env::MazeUndirected)
    str = "tile_map:\n"
    str = str * GW.get_pretty_tile_map(env)
    str = str * "\nsub_tile_map:\n"
    str = str * GW.get_pretty_sub_tile_map(env, GW.get_window_size(env))
    str = str * "\nreward: $(env.reward)"
    str = str * "\ndone: $(env.done)"
    str = str * "\naction_names: $(GW.get_action_names(env))"
    str = str * "\nobject_names: $(GW.get_object_names(env))"
    print(io, str)
    return nothing
end

GW.get_action_keys(env::MazeUndirected) = ('w', 's', 'a', 'd')

#####
##### RLBase API
#####

RLBase.StateStyle(env::GW.RLBaseEnv{E}) where {E <: MazeUndirected} = RLBase.InternalState{Any}()
RLBase.state_space(env::GW.RLBaseEnv{E}, ::RLBase.InternalState) where {E <: MazeUndirected} = nothing
RLBase.state(env::GW.RLBaseEnv{E}, ::RLBase.InternalState) where {E <: MazeUndirected} = env.env.tile_map

RLBase.reset!(env::GW.RLBaseEnv{E}) where {E <: MazeUndirected} = GW.reset!(env.env)

RLBase.action_space(env::GW.RLBaseEnv{E}) where {E <: MazeUndirected} = Base.OneTo(NUM_ACTIONS)
(env::GW.RLBaseEnv{E})(action) where {E <: MazeUndirected} = GW.act!(env.env, action)

RLBase.reward(env::GW.RLBaseEnv{E}) where {E <: MazeUndirected} = env.env.reward
RLBase.is_terminated(env::GW.RLBaseEnv{E}) where {E <: MazeUndirected} = env.env.done

end # module


module Move_FourBar_noMacros

import Modia3D
using  Modia3D.ModiaMath
using  Modia3D.Unitful

# Solids
groundMaterial = Modia3D.Material(color="DarkGreen", transparency=0.5)
massProperties = "Aluminium"
vmat1 = Modia3D.Material(color="LightBlue" , transparency=0.5)
vmat2 = Modia3D.Material(color="Red")

# Base dimensions
length1  = 1.0
width    = 0.2
height   = 0.2
diameter = 0.1
groundWidth  = 1.0
groundHeight = 0.01
cyl  = Modia3D.Cylinder(width/2, 1.2*width; material=vmat2)
Tcyl = ModiaMath.rot1(90u"°")
TGround = ModiaMath.rot1(-90u"°")

# -----------------------------------------------------------------------
# @assembly Bar(;Lx = 0.1, Ly=Lx/5, Lz=Ly) begin
#   frame0 = Modia3D.Object3D(Modia3D.Solid(Modia3D.SolidBeam(Lx,Ly,Lz), nothing, vmat1))
#   frame1 = Modia3D.Object3D(frame0; r=[-Lx/2, 0.0, 0.0])
#   frame2 = Modia3D.Object3D(frame0; r=[ Lx/2, 0.0, 0.0])
#   cyl1   = Modia3D.Object3D(frame1, cyl)
#   cyl2   = Modia3D.Object3D(frame2, cyl)
# end
mutable struct Bar <: Modia3D.AbstractAssembly
   _internal::Modia3D.AssemblyInternal
   Lx
   Ly
   Lz
   frame0
   frame1
   frame2
   cyl1
   cyl2

   function Bar(;sceneOptions=nothing,Lx = 0.1, Ly=Lx/5, Lz=Ly)
      this = new( Modia3D.AssemblyInternal(:Bar,sceneOptions) )
      this.Lx = Lx
      this.Ly = Ly
      this.Lz = Lz

      frame0 = Modia3D.Object3D(Modia3D.Solid(Modia3D.SolidBeam(Lx,Ly,Lz), nothing, vmat1))
      Modia3D.initAssemblyComponent!(this, frame0, "frame0")

      frame1 = Modia3D.Object3D(frame0; r=[-Lx/2, 0.0, 0.0])
      Modia3D.initAssemblyComponent!(this, frame1, "frame1")

      frame2 = Modia3D.Object3D(frame0; r=[ Lx/2, 0.0, 0.0])
      Modia3D.initAssemblyComponent!(this, frame2, "frame2")

      cyl1   = Modia3D.Object3D(frame1, cyl)
      Modia3D.initAssemblyComponent!(this, cyl1, "cyl1")

      cyl2   = Modia3D.Object3D(frame2, cyl)
      Modia3D.initAssemblyComponent!(this, cyl2, "cyl2")

      return this
   end
end
# -----------------------------------------------------------------------

# -----------------------------------------------------------------------
# @assembly Fourbar begin
#    world  = Modia3D.Object3D(Modia3D.CoordinateSystem(0.6))
#    frame1 = Modia3D.Object3D(world ; r=[0.5*length1, 0.0, groundWidth/2])
#    frame2 = Modia3D.Object3D(frame1; r=[length1    , 0.0, 0.0])
#    ground = Modia3D.Object3D(world , Modia3D.Box(3*length1, groundWidth, groundHeight; material=groundMaterial);
#                                               r = [1.5*length1, -groundHeight/2, groundWidth/2],
#                                               R = Modia3D.rot(;n1 = [-1,0,0], n2 = [0,0,1]))
#    bar1   = Bar(Lx=length1)
#    bar2   = Bar(Lx=length1)
#    bar3   = Bar(Lx=length1)
#
#    rev1   = Modia3D.Revolute(frame1     , bar1.frame1; phi =  pi/2)
#    rev2   = Modia3D.Revolute(bar1.frame2, bar2.frame1; phi = -pi/2)
#    rev3   = Modia3D.Revolute(frame2     , bar3.frame1; phi =  pi/2)
#    rev4   = Modia3D.Revolute(bar3.frame2, bar2.frame2; phi = -pi/2)
# end
mutable struct Fourbar <: Modia3D.AbstractAssembly
   _internal::Modia3D.AssemblyInternal
   world
   frame1
   frame2
   ground
   bar1
   bar2
   bar3
   rev1
   rev2
   rev3
   rev4

   function Fourbar(;sceneOptions=nothing)
      this = new( Modia3D.AssemblyInternal(:Fourbar, sceneOptions) )

      world  = Modia3D.Object3D(Modia3D.CoordinateSystem(0.6))
      Modia3D.initAssemblyComponent!(this, world , "world")

      frame1 = Modia3D.Object3D(world ; r=[0.5*length1, 0.0, groundWidth/2])
      Modia3D.initAssemblyComponent!(this, frame1, "frame1")

      frame2 = Modia3D.Object3D(frame1; r=[length1, 0.0, 0.0])
      Modia3D.initAssemblyComponent!(this, frame2, "frame2")

      ground = Modia3D.Object3D(world , Modia3D.Box(3*length1, groundWidth, groundHeight; material=groundMaterial);
                                                 r = [1.5*length1, -groundHeight/2, groundWidth/2],
                                                 R = ModiaMath.rot_nxy([-1,0,0], [0,0,1]))
      Modia3D.initAssemblyComponent!(this, ground, "ground")

      bar1 = Bar(Lx=length1)
      Modia3D.initAssemblyComponent!(this, bar1, "bar1")

      bar2 = Bar(Lx=length1)
      Modia3D.initAssemblyComponent!(this, bar2, "bar2")

      bar3 = Bar(Lx=length1)
      Modia3D.initAssemblyComponent!(this, bar3, "bar3")

      rev1 = Modia3D.Revolute(frame1, bar1.frame1; phi_start =  pi/2)
      Modia3D.initAssemblyComponent!(this, rev1, "rev1")

      rev2 = Modia3D.Revolute(bar1.frame2, bar2.frame1; phi_start = -pi/2)
      Modia3D.initAssemblyComponent!(this, rev2, "rev2")

      rev3 = Modia3D.Revolute(frame2, bar3.frame1; phi_start =  pi/2)
      Modia3D.initAssemblyComponent!(this, rev3, "rev3")

      rev4 = Modia3D.Revolute(bar3.frame2, bar2.frame2; phi_start = -pi/2)
      Modia3D.initAssemblyComponent!(this, rev4, "rev4")

      return this
   end
end
# -----------------------------------------------------------------------

# Kinematic simulation
fourbar = Fourbar(sceneOptions=Modia3D.SceneOptions(visualizeFrames=true,
                                                    defaultFrameLength=0.3,
                                                    enableContactDetection=false))
Modia3D.initAnalysis!(fourbar)
tStart=0.0
tEnd  =1.0
Modia3D.driveJoint!(fourbar.rev1)  # Define that rev1 is driven (setAngle! can be called on it)
Modia3D.driveJoint!(fourbar.rev2)  # Define that rev2 is driven (setAngle! can be called on it)

for time = range(tStart, stop=tEnd, length=101)
  # update positional degrees of freedom
  delta_phi = Modia3D.linearMovement(pi/3, tStart, tEnd, time)
  Modia3D.setAngle!(fourbar.rev1,  pi/2 - delta_phi)
  Modia3D.setAngle!(fourbar.rev2, -pi/2 + delta_phi)

  Modia3D.updatePosition!(fourbar)

  # Visualize
  Modia3D.visualize!(fourbar,time)
end

Modia3D.closeAnalysis!(fourbar)

println("... success of Move_FourBar_noMacros.jl!")
end


"""
    CSysModule

Module for management of coordinate systems.
"""
module CSysModule

using FinEtools.FTypesModule: FInt, FFlt, FCplxFlt, FFltVec, FIntVec, FFltMat, FIntMat, FMat, FVec, FDataDict
import LinearAlgebra: norm, cross

"""
Type for coordinate system transformations.
"""
struct CSys{F<:Function}
    isconstant::Bool
    isidentity::Bool
    updatebuffer!::F # function to update the coordinate system matrix
    csmat::Array{FFlt, 2} # the coordinate system matrix (buffer); see
end

"""
    CSys(sdim::FInt, mdim::FInt, computecsmat::F) where {F<:Function}

Construct ccoordinate system when the function to compute the
rotation matrix is given.

The function signature:
`update!(csmatout::FFltMat, XYZ::FFltMat, tangents::FFltMat, fe_label::FInt)`
where
- `csmatout`= output matrix buffer,
- `XYZ`= location  in physical coordinates,
- `tangents`= tangent vector matrix, tangents to the parametric coordinate
  curves  in the element,
- `fe_label`= finite element label.
"""
function CSys(sdim::FInt, mdim::FInt, computecsmat::F) where {F<:Function}
    csmat = fill(zero(FFlt), sdim, mdim); # Allocate buffer, in preparation for the first call
    return CSys(false, false, computecsmat, csmat);
end

"""
    CSys(csmat::FFltMat)

Construct coordinate system when the rotation matrix is given.
"""
function CSys(csmat::FFltMat)
    function updatebuffer!(csmatout::FFltMat, XYZ::FFltMat, tangents::FFltMat, fe_label::FInt)
        return csmatout # nothing to be done here, the matrix is already in the buffer
    end
    return CSys(true, false, updatebuffer!, deepcopy(csmat));# fill the buffer with the given matrix
end

"""
    CSys(dim::FInt)

Construct coordinate system when the rotation matrix is the identity.

`dim` = is the space dimension.
"""
function CSys(dim::FInt)
    return CSys([i==j ? one(FFlt) : zero(FFlt) for i=1:dim, j=1:dim]);
end

"""
    CSys(sdim::FInt, mdim::FInt)

Construct coordinate system for isotropic-material used with isoparametric
finite elements.

- `sdim` = number of space dimensions,
- `mdim` = number of manifold dimensions of the finite element in which the
  coordinate system  is being evaluated.

# Warning
If  the coordinate system matrix  should be identity, better use the constructor
for this specific situation, `CSys(dim::FInt)`. That will be much more efficient.
# See also
`gen_iso_csmat`
"""
function CSys(sdim::FInt, mdim::FInt)
    csmat = fill(zero(FFlt), sdim, mdim); # Allocate buffer, prepare for the first call
    function updatebuffer!(csmatout::FFltMat, XYZ::FFltMat, tangents::FFltMat, fe_label::FInt)
      gen_iso_csmat!(csmatout, XYZ, tangents, fe_label)
      return  csmatout
    end
    return CSys(false, false, updatebuffer!, csmat);
end

"""
    function updatecsmat!(self::CSys, XYZ::FFltMat, tangents::FFltMat, fe_label::FInt)

Update the coordinate system orientation matrix.

The  coordinate system matrix is updated based upon the location `XYZ` of the
evaluation point, and possibly on the Jacobian matrix `tangents` within the
element in which the coordinate system matrix is evaluated,  or perhaps on the
label `fe_label` of the finite element.

After this function returns, the coordinate system matrix can be retrieved
from the buffer as `self.csmat`.
"""
function updatecsmat!(self::CSys, XYZ::FFltMat, tangents::FFltMat, fe_label::FInt)
    self.updatebuffer!(self.csmat, XYZ, tangents, fe_label)
    return self.csmat
end

"""
    gen_iso_csmat(XYZ::FFltMat, tangents::FFltMat, fe_label::FInt)

Compute the coordinate system  for an isotropic material using information
available  by looking at the coordinate curves of isoparametric finite
elements.

- `XYZ`= location  in physical coordinates,
- `tangents`= tangent vector matrix, tangents to the parametric coordinate
  curves  in the element,
- `fe_label`= finite element label.

The basic assumption here is that the material is isotropic, and
therefore the choice of the material directions does not really matter as
long as they correspond to the dimensionality of the element. For
instance a one-dimensional element (L2 as an example) may be embedded
in a three-dimensional space.

This function assumes that it is being called for an mdim-dimensional manifold
element, which is embedded in a sdim-dimensional Euclidean space. If `mdim ==
sdim`, the coordinate system matrix is the identity; otherwise the local
coordinate directions are aligned with the linear subspace defined by the
tangent vectors.

# Warning

This *cannot* be reliably used to produce consistent stresses because each
quadrature point gets a local coordinate system which depends on the
orientation of the element.
"""
function gen_iso_csmat!(csmatout::FFltMat,
    XYZ::FFltMat, tangents::FFltMat, fe_label::FInt)
    sdim, mdim = size(tangents);
    if sdim == mdim # finite element embedded in space of the same dimension
        copyto!(csmatout, [i==j ? one(FFlt) : zero(FFlt) for i=1:sdim, j=1:sdim]);
    else # lower-dimensional finite element embedded in space of higher dimension
        @assert 0 < mdim < 3
        e1 = tangents[:,1]/norm(tangents[:,1]);
        if mdim == 1 # curve-like finite element
            copyto!(csmatout, e1);
        elseif mdim == 2 # surface-like finite element
            n = cross(e1, vec(tangents[:,2]/norm(tangents[:,2])));
            e2 = cross(n, e1);
            e2 = e2/norm(e2);
            csmatout[:,1] = e1
            csmatout[:,2] = e2
        end
    end
    return csmatout
end

end


using ADI
using BenchmarkTools
using CSV
using DataFrames
using HCIDatasets: BetaPictoris, HR8799
using InvertedIndices
using PyCall

results = []

# julia benchmarks

julia_algs = (Median(), PCA(20), NMF(20))

for dataset in (BetaPictoris, HR8799)
    @info "Benchmarking - $dataset ADI Reduction"
    cube = dataset[:cube]
    angles = dataset[:pa]
    for (alg, name) in zip(julia_algs, ("median", "pca_20", "nmf_20"))
        (alg isa NMF && dataset === HR8799) && continue
        @info "Julia - $alg"

        time_elapsed = @belapsed $alg($cube, $angles)

        @info "Julia - $alg" time=time_elapsed
        push!(results, (framework="ADI.jl", alg=name, N=length(cube), time=time_elapsed))
    end
end

# python benchmarks

vip = pyimport("vip_hci")
py_algs = (
    (args...; kwargs...) -> vip.medsub.median_sub(args..., verbose=false, kwargs...),
    (args...; kwargs...) -> vip.pca.pca(args...; ncomp=20, verbose=false, kwargs...),
    (args...; kwargs...) -> vip.nmf.nmf(args...; ncomp=20, verbose=false, kwargs...),
)

for dataset in (BetaPictoris, HR8799)
    @info "Benchmarking - $dataset ADI Reduction"
    cube = dataset[:cube]
    angles = dataset[:pa]
    for (alg, name) in zip(py_algs, ("median", "pca_20", "nmf_20"))
        (name == "nmf_20" && dataset === HR8799) && continue
        @info "Python - $alg"

        time_elapsed = @belapsed $alg($cube, $angles)

        @info "Python - $alg" time=time_elapsed
        push!(results, (framework="VIP", alg=name, N=length(cube), time=time_elapsed))
    end
end

# load results file and update
path = joinpath(@__DIR__, "adi_benchmarks.csv")
df = isfile(path) ? CSV.File(path) |> DataFrame : DataFrame()
out = append!(DataFrame(results), df)
unique!(out, Not(:time)) |> CSV.write(path)


"""
Run gibbs sampler.
"""
function gibbs_sample!(
        model::SeqModel,
        spikes::Vector{Spike},
        initial_assignments::Vector{Int64},
        num_samples::Int64,
        extra_split_merge_moves::Int64,
        split_merge_window::Float64,
        save_every::Int64;
        verbose::Bool=false
    )

    # Initialize spike assignments.
    assignments = initial_assignments
    recompute!(model, spikes, assignments)

    # ======== THINGS TO SAVE ======== #

    # number of samples to save.
    n_saved_samples = Int(round(num_samples / save_every))

    # log likelihoods.
    log_p_hist = zeros(n_saved_samples)

    # parent assignments.
    assignment_hist = zeros(
        Int64, length(spikes), n_saved_samples
    )

    # latent event assignment ids, times, types, amplitudes.
    latent_event_hist = Vector{EventSummaryInfo}[]

    # global variables (offsets, response amplitudes, etc.)
    globals_hist = SeqGlobals[]

    # Order to iterate over spikes.
    spike_order = collect(1:length(spikes))

    # ======== MAIN LOOP ======== #

    for s = 1:num_samples

        # Update spike assignments in random order.
        Random.shuffle!(spike_order)
        for i = spike_order
            remove_datapoint!(model, spikes[i], assignments[i])
            assignments[i] = gibbs_add_datapoint!(model, spikes[i])
        end

        # Add extra split merge moves.
        split_merge_sample!(
            model,
            spikes,
            extra_split_merge_moves,
            assignments,
            split_merge_window,
        )

        # Update latent events.
        gibbs_update_latents!(model)

        # Update globals
        gibbs_update_globals!(model, spikes, assignments)

        # Store results
        if (s % save_every) == 0

            # Index into hist vectors.
            j = Int(s / save_every)

            # Note that the likelihood should be computed after
            # gibbs_update_latents! and before spike parent assignments
            # are recomputed. Otherwise, newly added latent events are
            # not initialized.
            log_p_hist[j] = log_like(model, spikes)

            # Save assignments.
            assignment_hist[:, j] .= assignments

            # Save latent event information.
            push!(
                latent_event_hist,
                event_list_summary(model)
            )

            # Save global variables.
            push!(globals_hist, deepcopy(model.globals))

            # Display progress.
            verbose && print(s, "-")
        end
    end

    # Finish progress bar.
    verbose && (n_saved_samples > 0) && println("Done")

    return (
        assignments,
        assignment_hist,
        log_p_hist,
        latent_event_hist,
        globals_hist
    )
end

"""
Adds spikes `s` to an existing sequence event, to a new sequence event,
or to the background process.

For each sequence event k = 1 ... K, we compute

    prob[k] = p(x_i | z_i = k, x_{neq i}) * (N_k + alpha)

The probability of forming a new cluster is

    prob[K + 1] propto p(x_i | z_i = K + 1) * alpha * (V(K + 1) / V(K))

where p(x_i | z_i = K + 1) is the marginal probability of a singleton cluster.
See section 6 of Miller & Harrison (2018).

The probability of the background is

    prob[K + 2] propto p(x_i | bkgd) * lambda0 * (1 + beta)

where lambda0 = m.bkgd_rate and (alpha, beta) are the shape and
rate parameters of the gamma prior on sequence event amplitudes.
"""
function gibbs_add_datapoint!(model::SeqModel, x::Spike)

    # Create log-probability vector to sample assignments.
    #
    #  - We need to sample K + 2 possibilities. There are K existing clusters
    #    we could assign to. We could also form a new cluster (index K + 1),
    #    or assign the spike to the background (index K + 2).

    # Shape and rate parameters of gamma prior on latent event amplitude.
    α = model.priors.seq_event_amplitude.α
    β = model.priors.seq_event_amplitude.β

    K = num_sequence_events(model)

    # Grab vector without allocating new memory.
    log_probs = resize!(model._K_buffer, K + 2)

    # Iterate over model events, indexed by k = {1, 2, ..., K}.
    for (k, event) in enumerate(model.sequence_events)

        # Check if event is too far away to be considered. If sampled_type < 0,
        # then the event timestamp hasn't been sampled yet, so we can't give
        # up yet. Be aware that this previously led to a subtle bug, which was
        # very painful to fix.
        too_far = abs(x.timestamp - event.sampled_timestamp) > model.max_sequence_length
        if too_far && (event.sampled_type > 0)
            log_probs[k] = -Inf

        # Compute probability of adding spike to cluster k.
        else
            log_probs[k] = (
                log(event.spike_count + α)
                + log_posterior_predictive(event, x, model)
            )
        end
    end

    # New cluster probability.
    log_probs[K + 1] = (
        model.new_cluster_log_prob
        + log_posterior_predictive(x, model)
    )

    # Background probability
    log_probs[K + 2] = (
        model.bkgd_log_prob
        + model.globals.bkgd_log_proportions[x.neuron]
        - log(model.max_time)
    )

    # Sample new assignment for spike x.
    z = sample_logprobs!(log_probs)

    # New sample corresponds to background, do nothing.
    if z == (K + 2)
        return -1

    # New sample corresponds to new sequence event / cluster.
    elseif z == (K + 1)
        return add_event!(model, x)  # returns new assignment

    # Otherwise, add datapoint to existing sequence event. Note
    # that z is an integer in [1:K], while assignment indices
    # can be larger and non-contiguous.
    else
        k = model.sequence_events.indices[z]  # look up assignment index.
        return add_datapoint!(model, x, k)
    end

end

function gibbs_update_globals!(
        model::SeqModel,
        spikes::Vector{Spike},
        assignments::AbstractVector{Int64}
    )

    K = num_sequence_events(model)
    N = num_neurons(model)
    R = num_sequence_types(model)

    priors = model.priors
    globals = model.globals

    # === RESAMPLE BACKGROUND SPIKE PARAMETERS === #

    num_bkgd_spikes = 0
    bkgd_counts = zeros(N)

    for (i, x) in enumerate(spikes)
        if assignments[i] < 0
            num_bkgd_spikes += 1
            bkgd_counts[x.neuron] += 1
        end
    end

    # Sample proportions of background spikes across neurons
    # as a probability vector then map to log space for
    # future computations.
    rand!(
        posterior(
            bkgd_counts, priors.bkgd_proportions
        ),
        globals.bkgd_log_proportions
    ) # multinomial - symmetric dirichlet conjugate pair.
    map!(
        log,
        globals.bkgd_log_proportions,
        globals.bkgd_log_proportions
    )

    # Sample the rate of background spikes. Here, we need to adjust
    # for the length of the time interval, max_time. The scaling
    # property of the gamma distribution implies that:
    #
    #   bkgd_amp * max_time ~ RateGamma(α_bkgd, β_bkgd / max_time)
    #
    # If we observe num_bkgd, then the posterior (by Poisson-gamma
    # conjugacy) is:
    #
    #   bkgd_amp * T | num_bkgd_spikes ~ RateGamma(num_bkgd + α_bkgd, 1 + β_bkgd / T)
    #
    # Now apply the gamma scaling property again, dividing by T this
    # time, so we get:
    #
    #   bkgd_amp | num_bkgd_spikes ~ RateGamma(num_bkgd + α_bkgd, T + β_bkgd)
    #
    globals.bkgd_amplitude = rand(
        RateGamma(
            num_bkgd_spikes + priors.bkgd_amplitude.α,
            priors.bkgd_amplitude.β + model.max_time
        )
    )

    # === RESAMPLE SEQUENCE TYPE PROPORTIONS === #

    seq_type_counts = zeros(R)
    for event in model.sequence_events
        seq_type_counts[event.sampled_type] += 1
    end

    rand!(
        posterior(
            seq_type_counts, priors.seq_type_proportions
        ),
        model.globals.seq_type_log_proportions
    )  # multinomial - symmetric dirichlet conjugate pair.
    map!(
        log,
        globals.seq_type_log_proportions,
        globals.seq_type_log_proportions
    )

    # === RESAMPLE NEURON RESPONSE PROFILES === #

    # TODO -- preallocate arrays for this? Or is it not worth it?
    spk_count = zeros(Int64, N, R)
    spk_fm = zeros(N, R)
    spk_sm = zeros(N, R)

    for (i, x) = enumerate(spikes)

        # spike assignment to latent event.
        k = assignments[i]

        # skip if spike is assigned to background.
        (k < 0) && continue

        event = model.sequence_events[k]
        n = x.neuron
        r = event.sampled_type
        w = event.sampled_warp
        offset = (x.timestamp - event.sampled_timestamp) / model.priors.warp_values[w]

        # compute sufficient statistics
        spk_count[n, r] += 1
        spk_fm[n, r] += offset
        spk_sm[n, r] += offset * offset

    end

    for r = 1:R

        rand!(
            posterior(
                view(spk_count, :, r),
                priors.neuron_response_proportions
            ),
            view(globals.neuron_response_log_proportions, :, r)
        ) # multinomial - dirichlet conjugate pair.

        for n = 1:N

            (
                globals.neuron_response_offsets[n, r],
                globals.neuron_response_widths[n, r]
            ) =
            rand(
                posterior(
                    spk_count[n, r],
                    spk_fm[n, r],
                    spk_sm[n, r],
                    priors.neuron_response_profile
                )
            ) # normal - norm-inv-chi-squared conjugate pair.

        end

    end

    map!(
        log,
        globals.neuron_response_log_proportions,
        globals.neuron_response_log_proportions
    )

    # === RECOMPUTE NEW CLUSTER PROBABILITIES === #

    α = priors.seq_event_amplitude.α
    β = priors.seq_event_amplitude.β

    # TODO : make this a field of globals, not model.
    model.new_cluster_log_prob = (
        log(α)
        + log(model.priors.seq_event_rate)  # TODO: resample this as well.
        + log(model.max_time)
        + α * (log(β) - log(1 + β))
    )

    # TODO : make this a field of globals, not model.
    model.bkgd_log_prob = (
        log(globals.bkgd_amplitude)
        + log(model.max_time)
        + log(1 + β)
    )

    # === RECOMPUTE SUFFICIENT STATISTICS === #

    # Because the sequence types have been resampled, the sufficient statistics
    # computed for each sequence event must be re-computed (they depend on the
    # neuron offsets, relative amplitudes, etc.!)

    recompute!(model, spikes, assignments)

end

"""
Resamples sequence type, timestamp, and amplitude. For all
latent events.
"""
function gibbs_update_latents!(model::SeqModel)

    # Grab length-R vector (already pre-allocated).
    log_probs = model._RW_buffer

    for event in model.sequence_events

        # We should only be resampling non-empty events.
        @assert (event.spike_count > 0)

        # Sample sequence type.
        log_probs .= event.seq_type_posterior
        ind = sample_logprobs!(vec(log_probs))
        r, w = Tuple(CartesianIndices(size(log_probs))[ind])
        event.sampled_type = r
        event.sampled_warp = w

        # Sample event time, t ~ N(μ, σ2), given sequence type.
        σ2 = 1 / event.summed_precisions[r, w]
        μ = event.summed_potentials[r, w] * σ2
        event.sampled_timestamp = μ + randn() * sqrt(σ2)

        # Sample event amplitude.
        event.sampled_amplitude = rand(
            posterior(
                event.spike_count,
                model.priors.seq_event_amplitude
            )
        ) # Poisson - gamma conjugate pair.

    end

end


"""
    linearized model-predictive control policy
"""

@with_kw mutable struct LinearizedMPCOptions{T}
	altitude_update::Bool = false
	altitude_impact_threshold::T = 1.0
	altitude_verbose::Bool = false
    ip_max_time::T = 1e5     # maximum time allowed for an InteriorPoint solve
    live_plotting::Bool=false # Use the live plotting tool to debug
end

mutable struct LinearizedMPC <: Policy
	traj
	ref_traj
	im_traj
	H
	stride
	altitude
	κ
	newton
	newton_mode
	s
	q0
	N_sample
	cnt
	opts
end

function linearized_mpc_policy(traj, s, obj;
	H_mpc = traj.H,
	N_sample = 1,
	κ_mpc = traj.κ[1],
	ip_type = :interior_point,
	mode = :configurationforce,
	newton_mode = :direct,
	n_opts = NewtonOptions(
		r_tol = 3e-4,
		max_iter = 5,
		verbose = false,
		live_plotting = false),
	mpc_opts = LinearizedMPCOptions(),
	ip_opts = eval(interior_point_options(ip_type))(
				γ_reg = 0.1,
				undercut = 5.0,
				κ_tol = κ_mpc,
				r_tol = 1.0e-8,
				diff_sol = true,
				solver = :empty_solver,
				max_time = mpc_opts.ip_max_time,)
	)

	traj = deepcopy(traj)
	ref_traj = deepcopy(traj)

	im_traj = ImplicitTraj(traj, s,
		ip_type = ip_type,
		κ = κ_mpc,
		max_time = mpc_opts.ip_max_time,
		opts=ip_opts,
		mode = mode)

	stride = get_stride(s.model, traj)
	altitude = zeros(s.model.dim.c)
	if newton_mode == :direct
		newton = Newton(s, H_mpc, traj.h, traj, im_traj, obj = obj, opts = n_opts)
	elseif newton_mode == :structure
		newton = NewtonStructure(s, H_mpc, traj, obj, κ_mpc, opts = n_opts)
	else
		@error "invalid Newton solver specified"
	end

	LinearizedMPC(traj, ref_traj, im_traj, H_mpc, stride, altitude, κ_mpc, newton, newton_mode, s, copy(ref_traj.q[1]),
		N_sample, N_sample, mpc_opts)
end

function policy(p::LinearizedMPC, x, traj, t)
	# reset
	if t == 1
		p.cnt = p.N_sample
		p.q0 = copy(p.ref_traj.q[1])
	end

    if p.cnt == p.N_sample
		(p.opts.altitude_update && t > 1) && (update_altitude!(p.altitude, p.s,
									traj, t, p.N_sample,
									threshold = p.opts.altitude_impact_threshold,
									verbose = p.opts.altitude_verbose))
		# update!(p.im_traj, p.traj, p.s, p.altitude, κ = p.κ) #@@@ keep the altitude update here
		set_altitude!(p.im_traj, p.altitude) #@@@ keep the altitude update here
		newton_solve!(p.newton, p.s, p.im_traj, p.traj,
			warm_start = t > 1, q0 = copy(p.q0), q1 = copy(x))
		update!(p.im_traj, p.traj, p.s, p.altitude, κ = p.κ) #@@@ only keep the rotation stuff not the altitude update.
		p.opts.live_plotting && live_plotting(p.s.model, p.traj, traj, p.newton, p.q0, copy(x), t)

		rot_n_stride!(p.traj, p.stride)
		p.q0 .= copy(x)
		p.cnt = 0
    end

    p.cnt += 1

	if p.newton_mode == :direct
    	return p.newton.traj.u[1] / p.N_sample # rescale output
	elseif p.newton_mode == :structure
		return p.newton.u[1] / p.N_sample
	else
		@error "newton mode specified not available"
	end
end


# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/FEMBase.jl/blob/master/LICENSE

mutable struct LinearSystem{Tv, Ti<:Integer}
    M :: SparseMatrixCSC{Tv, Ti}
    K :: SparseMatrixCSC{Tv, Ti}
    Kg :: SparseMatrixCSC{Tv, Ti}
    C1 :: SparseMatrixCSC{Tv, Ti}
    C2 :: SparseMatrixCSC{Tv, Ti}
    D :: SparseMatrixCSC{Tv, Ti}
    f :: SparseVector{Tv, Ti}
    fg :: SparseVector{Tv, Ti}
    g :: SparseVector{Tv, Ti}
    u :: SparseVector{Tv, Ti}
    la :: SparseVector{Tv, Ti}
    dim :: Int
end

function LinearSystem(dim::Int)
    return LinearSystem(spzeros(dim, dim), spzeros(dim, dim),
                        spzeros(dim, dim), spzeros(dim, dim),
                        spzeros(dim, dim), spzeros(dim, dim),
                        spzeros(dim), spzeros(dim), spzeros(dim),
                        spzeros(dim), spzeros(dim), dim)
end

abstract type AbstractLinearSystemSolver end

function solve!(::LinearSystem, ::Solver) where Solver<:AbstractLinearSystemSolver
    info("This is a placeholder function for solving linear systems")
    info("To solve linear systems, you must define a function")
    info("      solve!(system::LinearSystem, solver::$Solver)")
end

function can_solve(::LinearSystem, ::Solver) where Solver<:AbstractLinearSystemSolver
    return (true, "OK")
end

function solve!(ls::LinearSystem, solvers::Vector{S}) where S<:AbstractLinearSystemSolver
    for solver in solvers
        Solver = typeof(solver)
        cansolve, msg = can_solve(ls, solver)
        if !cansolve
            info("Solver $Solver cannot solve linear system: $msg")
            continue
        end
        timeit("solve linear system using solver $Solver") do
            solve!(ls, solver)
        end
        return
    end
    error("Failed to solve linear system.")
end


@system Plant(
    Mass,
    Area,
    Count,
    Ratio,
    Carbon,
    #Nitrogen,
    Water,
    Weight,
    Density,
    Photosynthesis
) begin
    calendar(context) ~ ::Calendar
    weather(context, calendar) ~ ::Weather
    sun(context, calendar, weather) ~ ::Sun
    soil(context) ~ ::Soil
    pheno(context, calendar, weather, sun, soil): phenology ~ ::Phenology

    primordia => 5 ~ preserve::int(parameter)

    #bulb => begin end ~ produce::Bulb

    #scape => begin end ~ produce::Scape

    root(pheno, emerging=pheno.emerging) => begin
        if emerging
            #TODO import_carbohydrate(soil.total_root_weight)
            produce(Root, phenology=pheno)
        end
    end ~ produce::Root

    #TODO pass PRIMORDIA as initial_leaves
    NU(NU, pheno, primordia, germinated=pheno.germinated, dead=pheno.dead, l=pheno.leaves_initiated): nodal_units => begin
        if germinated && !dead
            n = length(NU)
            if n == 0
                [produce(NodalUnit, phenology=pheno, rank=i) for i in 1:primordia]
            elseif n < l
                [produce(NodalUnit, phenology=pheno, rank=i) for i in (n+1):l]
            end
        end
    end ~ produce::NodalUnit[]

    # dry_yield(bulb_mass, PD) => bulb_mass * PD ~ track(u"g/m^2")
    # fresh_yield(dry_yield, BMC) => begin
    #     dry_yield * 1 / (1 - BMC)
    # end ~ track(u"g/m^2")
    # BMC: bulb_moisture_content => 0.85 ~ preserve

    SLA(leaf_area, leaf_mass): specific_lear_area => begin
        leaf_area / leaf_mass
    end ~ track(u"cm^2/g")
    LMA(SLA): leaf_mass_per_area => 1 / SLA ~ track(u"g/cm^2")

    DAP(pheno.DAP): day_after_planting ~ track::int(u"d")
    time(calendar.time) ~ track::datetime
    date(calendar.date) ~ track::date
end

@system Model(Plant, Controller)


using Test
using DFTK: load_psp, pymatgen_lattice, pymatgen_structure, load_lattice, load_atoms
using DFTK: ElementCoulomb
using Unitful
using UnitfulAtomic
using PyCall

py"""
from pymatgen.core.lattice import Lattice
from pymatgen.core.structure import Structure
"""

@testset "Lattice DFTK -> pymatgen -> DFTK" begin
    reference = randn(3, 3)
    output = load_lattice(pymatgen_lattice(reference))
    @test output ≈ reference atol=1e-14
end

@testset "Lattice pymatgen -> DFTK -> pymatgen" begin
    data = randn(9)
    reference = py"Lattice($data)"
    output = pymatgen_lattice(load_lattice(reference))
    output = py"$output.matrix.ravel()" .+ 0
    @test data ≈ output atol=1e-14
end

@testset "Lattice DFTK -> pymatgen" begin
    a = randn(3)
    b = randn(3)
    c = randn(3)
    reference = [a b c]

    # Convert the lattice to python, make it flat and convert
    # to Julia array
    output = pymatgen_lattice(reference)
    outlatt = py"$output.matrix.ravel()" .+ 0
    @test a ≈ austrip.(outlatt[1:3] * u"Å") atol=1e-14
    @test b ≈ austrip.(outlatt[4:6] * u"Å") atol=1e-14
    @test c ≈ austrip.(outlatt[7:9] * u"Å") atol=1e-14
end

@testset "Structure pymatgen -> DFTK -> pymatgen" begin
    data = randn(9)
    reflattice = py"Lattice($data)"
    species = [1, 1, 6, 6, 6, 8]
    positions = [randn(3), randn(3), randn(3), randn(3), randn(3), randn(3)]

    reference = py"Structure($reflattice, $species, $positions)"
    atoms = load_atoms(reference)
    @test length(atoms) == 3
    @test all(at isa ElementCoulomb for (at, positions) in atoms)
    @test atoms[1][1].symbol == :H
    @test atoms[2][1].symbol == :C
    @test atoms[3][1].symbol == :O

    output = pymatgen_structure(load_lattice(reference), atoms)
    @test output.lattice == reflattice
    for i in 1:6
        @test output.species[i].number == species[i]
        @test output.sites[i].frac_coords == positions[i]
    end
end

@testset "Structure DFTK -> pymatgen" begin
    a = randn(3)
    b = randn(3)
    c = randn(3)
    lattice = [a b c]

    atoms = [
        ElementCoulomb(1) => [randn(3), randn(3)],
        ElementCoulomb(6) => [randn(3), randn(3), randn(3)],
        ElementCoulomb(8) => [randn(3), randn(3)],
    ]

    # Convert the lattice to python, make it flat and convert
    # to Julia array
    output = pymatgen_structure(lattice, atoms)
    outlatt = py"$output.lattice.matrix.ravel()" .+ 0
    @test a ≈ austrip.(outlatt[1:3] * u"Å") atol=1e-14
    @test b ≈ austrip.(outlatt[4:6] * u"Å") atol=1e-14
    @test c ≈ austrip.(outlatt[7:9] * u"Å") atol=1e-14

    specmap = [1, 1, 2, 2, 2, 3, 3]
    offset = [0, 0, 2, 2, 2, 5, 5]
    for i in 1:6
        specpair = atoms[specmap[i]]
        ired = i - offset[i]
        @test output.species[i].number == specpair.first.Z
        @test output.sites[i].frac_coords == specpair.second[ired]
    end
end