"""
GenX: An Configurable Capacity Expansion Model
Copyright (C) 2021,  Massachusetts Institute of Technology
This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation; either version 2 of the License, or
(at your option) any later version.
This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.
A complete copy of the GNU General Public License v2 (GPLv2) is available
in LICENSE.txt.  Users uncompressing this from an archive may not have
received this license file.  If not, see <http://www.gnu.org/licenses/>.
"""

@doc raw"""
	write_net_revenue(path::AbstractString, inputs::Dict, setup::Dict, EP::Model, dfCap::DataFrame, dfESRRev::DataFrame, dfResRevenue::DataFrame, dfChargingcost::DataFrame, dfPower::DataFrame, dfEnergyRevenue::DataFrame, dfSubRevenue::DataFrame, dfRegSubRevenue::DataFrame)

Function for writing net revenue of different generation technologies.
"""
function write_net_revenue(path::AbstractString, inputs::Dict, setup::Dict, EP::Model, dfCap::DataFrame, dfESRRev::DataFrame, dfResRevenue::DataFrame, dfChargingcost::DataFrame, dfPower::DataFrame, dfEnergyRevenue::DataFrame, dfSubRevenue::DataFrame, dfRegSubRevenue::DataFrame)
	dfGen = inputs["dfGen"]
	T = inputs["T"]     			# Number of time steps (hours)
	Z = inputs["Z"]     			# Number of zones
	G = inputs["G"]     			# Number of generators
	COMMIT = inputs["COMMIT"]		# Thermal units for unit commitment

	# Create a NetRevenue dataframe
 	dfNetRevenue = DataFrame(region = dfGen[!,:region], Resource = inputs["RESOURCES"], zone = dfGen[!,:Zone], Cluster = dfGen[!,:cluster], R_ID = dfGen[!,:R_ID])

	# Add investment cost to the dataframe
	dfNetRevenue.Inv_cost_MW = dfGen[!,:Inv_Cost_per_MWyr] .* dfCap[1:end-1,:NewCap]
	dfNetRevenue.Inv_cost_MWh = dfGen[!,:Inv_Cost_per_MWhyr] .* dfCap[1:end-1,:NewEnergyCap]
	if setup["ParameterScale"] == 1
		dfNetRevenue.Inv_cost_MWh = dfNetRevenue.Inv_cost_MWh * (ModelScalingFactor) # converting Million US$ to US$
		dfNetRevenue.Inv_cost_MW = dfNetRevenue.Inv_cost_MW * (ModelScalingFactor) # converting Million US$ to US$
	end

	# Add operations and maintenance cost to the dataframe
	dfNetRevenue.Fixed_OM_cost_MW = dfGen[!,:Fixed_OM_Cost_per_MWyr] .* dfCap[1:end-1,:EndCap]
 	dfNetRevenue.Fixed_OM_cost_MWh = dfGen[!,:Fixed_OM_Cost_per_MWhyr] .* dfCap[1:end-1,:EndEnergyCap]
 	dfNetRevenue.Var_OM_cost_out = (dfGen[!,:Var_OM_Cost_per_MWh]) .* dfPower[1:end-1,:AnnualSum]
	if setup["ParameterScale"] == 1
		dfNetRevenue.Fixed_OM_cost_MW = dfNetRevenue.Fixed_OM_cost_MW * (ModelScalingFactor) # converting Million US$ to US$
		dfNetRevenue.Fixed_OM_cost_MWh = dfNetRevenue.Fixed_OM_cost_MWh * (ModelScalingFactor) # converting Million US$ to US$
		dfNetRevenue.Var_OM_cost_out = dfNetRevenue.Var_OM_cost_out * (ModelScalingFactor) # converting Million US$ to US$
	end

	# Add fuel cost to the dataframe
	dfNetRevenue.Fuel_cost = zeros(size(dfNetRevenue, 1))
	for i in 1:G
		dfNetRevenue.Fuel_cost[i] = sum(inputs["C_Fuel_per_MWh"][i,:] .* inputs["omega"] .* value.(EP[:vP])[i,:])
	end
	if setup["ParameterScale"] == 1
		dfNetRevenue.Fuel_cost = dfNetRevenue.Fuel_cost * (ModelScalingFactor^2) # converting Million US$ to US$
	end

	# Add storage cost to the dataframe
	dfNetRevenue.Var_OM_cost_in = zeros(size(dfNetRevenue, 1))
 	for y in inputs["STOR_ALL"]
 		dfNetRevenue.Var_OM_cost_in[y] = dfGen[y,:Var_OM_Cost_per_MWh_In] * sum(inputs["omega"] .* value.(EP[:vCHARGE])[y,:])
 	end
	if setup["ParameterScale"] == 1
		dfNetRevenue.Var_OM_cost_in = dfNetRevenue.Var_OM_cost_in * (ModelScalingFactor^2) # converting Million US$ to US$
	end
	# Add start-up cost to the dataframe
	dfNetRevenue.StartCost = zeros(size(dfNetRevenue, 1))
 	if (setup["UCommit"]>=1)
 		for y in COMMIT #dfGen[!,:R_ID]
 			dfNetRevenue.StartCost[y] = sum(value.(EP[:eCStart])[y,:])
 		end
 	end
	if setup["ParameterScale"] == 1
		dfNetRevenue.StartCost = dfNetRevenue.StartCost * (ModelScalingFactor^2) # converting Million US$ to US$
	end
	# Add charge cost to the dataframe
	dfNetRevenue.Charge_cost = zeros(size(dfNetRevenue, 1))
	if has_duals(EP) == 1
		dfNetRevenue.Charge_cost = dfChargingcost[!,:AnnualSum] # Unit is confirmed to be US$
	end

	# Add energy and subsidy revenue to the dataframe
	dfNetRevenue.EnergyRevenue = zeros(size(dfNetRevenue, 1))
	dfNetRevenue.SubsidyRevenue = zeros(size(dfNetRevenue, 1))
	if has_duals(EP) == 1
		dfNetRevenue.EnergyRevenue = dfEnergyRevenue[!,:AnnualSum] # Unit is confirmed to be US$
	 	dfNetRevenue.SubsidyRevenue = dfSubRevenue[!,:SubsidyRevenue] # Unit is confirmed to be US$
	end

	# Add capacity revenue to the dataframe
	dfNetRevenue.ReserveMarginRevenue = zeros(size(dfNetRevenue, 1))
 	if setup["CapacityReserveMargin"] > 0 && has_duals(EP) == 1 # The unit is confirmed to be $
 		dfNetRevenue.ReserveMarginRevenue = dfResRevenue[!,:AnnualSum]
 	end

	# Add RPS/CES revenue to the dataframe
	dfNetRevenue.ESRRevenue = zeros(size(dfNetRevenue, 1))
 	if setup["EnergyShareRequirement"] > 0 && has_duals(EP) == 1 # The unit is confirmed to be $
 		dfNetRevenue.ESRRevenue = dfESRRev[!,:AnnualSum]
 	end

	# Calculate emissions cost
	dfNetRevenue.EmissionsCost = zeros(size(dfNetRevenue, 1))
 	if setup["CO2Cap"] >=1 && has_duals(EP) == 1
 		for y in 1:G
			dfNetRevenue.EmissionsCost[y] = 0.0
			for cap in 1:inputs["NCO2Cap"]
 				if dfGen[y,:Zone] in findall(x->x==1, inputs["dfCO2CapZones"][:,cap])
					if setup["CO2Cap"]==1 # Mass-based
						# Cost = sum(sum(emissions of gen y * dual(CO2 constraint[cap]) for z in Z) for cap in setup["NCO2"])
						dfNetRevenue.EmissionsCost[y] += sum(value(EP[:eEmissionsByPlant][y,t]) * inputs["omega"][t] * (-1) * dual(EP[:cCO2Emissions_systemwide][cap])
							for t in 1:T)

					elseif setup["CO2Cap"]==2 # Demand + Rate-based
						# Cost = sum(sum(emissions for zone z * dual(CO2 constraint[cap]) for z in Z) for cap in setup["NCO2"])
						dfNetRevenue.EmissionsCost[y] += sum(value(EP[:eEmissionsByPlant][y,t]) * inputs["omega"][t] * (-1) * dual(EP[:cCO2Emissions_systemwide][cap])
							for t in 1:T)
					elseif setup["CO2Cap"]==3 # Generation + Rate-based
						if y in union(inputs["THERM_ALL"],inputs["VRE"], inputs["VRE"],inputs["MUST_RUN"],inputs["HYDRO_RES"])
							# Cost = sum( sum(emissions - generatio for zone z * MaxCO2Rate for zone z for z in Z) * dual(CO2 constraint[cap] for cap in setup["NCO2"])
							dfNetRevenue.EmissionsCost[y] += sum( (value(EP[:eEmissionsByPlant][y,t]) - (value(EP[:vP][y,t]) * inputs["dfMaxCO2Rate"][z,cap])) * inputs["omega"][t] * (-1) * dual(EP[:cCO2Emissions_systemwide][cap])
								for z=dfGen[y,:Zone], t in 1:T)
						end
					end
				end
			end
 		end
		if setup["ParameterScale"] == 1
			dfNetRevenue[!,:EmissionsCost] = dfNetRevenue[!,:EmissionsCost] * (ModelScalingFactor^2) # converting Million US$ to US$
		end
 	end

	# Add regional technology subsidy revenue to the dataframe
	dfNetRevenue.RegSubsidyRevenue = zeros(size(dfNetRevenue, 1))
	if setup["MinCapReq"] >= 1 && has_duals(EP) == 1 # The unit is confirmed to be US$
		dfNetRevenue.RegSubsidyRevenue = dfRegSubRevenue[!,:SubsidyRevenue]
	end

	dfNetRevenue.Revenue =	dfNetRevenue.EnergyRevenue + dfNetRevenue.SubsidyRevenue + dfNetRevenue.ReserveMarginRevenue + dfNetRevenue.ESRRevenue + dfNetRevenue.RegSubsidyRevenue
	dfNetRevenue.Cost = dfNetRevenue.Inv_cost_MW + dfNetRevenue.Inv_cost_MWh + dfNetRevenue.Fixed_OM_cost_MW + dfNetRevenue.Fixed_OM_cost_MWh + dfNetRevenue.Var_OM_cost_out + dfNetRevenue.Var_OM_cost_in + dfNetRevenue.Fuel_cost + dfNetRevenue.Charge_cost + dfNetRevenue.EmissionsCost + dfNetRevenue.StartCost
	#dfNetRevenue.Cost = dfNetRevenue.Inv_cost_MW + dfNetRevenue.Inv_cost_MWh + dfNetRevenue.Fixed_OM_cost_MW + dfNetRevenue.Fixed_OM_cost_MWh + dfNetRevenue.Var_OM_cost_out + dfNetRevenue.Var_OM_cost_in + dfNetRevenue.Fuel_cost + dfNetRevenue.Charge_cost + dfNetRevenue.EmissionsCost + dfNetRevenue.StartCost
	dfNetRevenue.Profit = 	dfNetRevenue.Revenue - dfNetRevenue.Cost

	CSV.write(joinpath(path, "NetRevenue.csv"), dfNetRevenue)
end


using KLIEPInference
using JLD

file = jldopen("../exper1/graphs/chain1_25.jld", "r")
γy = pack(read(file, "Θy"))
close(file)

# setting 1. (none)
idx0 = KLIEPInference.trimap(5, 6)
for δ in -0.75:0.05:0.75
    γx = deepcopy(γy)
    γx[idx0] += δ
    @save "graphs/params_exp2_set1_$(δ).jld" γx γy
end

# setting 2. (strong)
idx1 = KLIEPInference.trimap(4, 5)
idx2 = KLIEPInference.trimap(6, 7)
for δ in -0.75:0.05:0.75
    γx = deepcopy(γy)
    γx[idx0] += δ
    γx[idx1] += 0.4
    γx[idx2] -= 0.4
    @save "graphs/params_exp2_set2_$(δ).jld" γx γy
end

# setting 3. (weak)
idx3 = KLIEPInference.trimap(4, 6)
idx4 = KLIEPInference.trimap(5, 7)
for δ in -0.75:0.05:0.75
    γx = deepcopy(γy)
    γx[idx0] += δ
    γx[idx3] += 0.2
    γx[idx4] += 0.2
    @save "graphs/params_exp2_set3_$(δ).jld" γx γy
end

# setting 4. (mixed)
for δ in -0.75:0.05:0.75
    γx = deepcopy(γy)
    γx[idx0] += δ
    γx[idx1] += 0.4
    γx[idx2] -= 0.4
    γx[idx3] += 0.2
    γx[idx4] += 0.2
    @save "graphs/params_exp2_set4_$(δ).jld" γx γy
end


# read a PBF format file
# format documentation: https://wiki.openstreetmap.org/wiki/PBF_Format, https://github.com/mapbox/osmpbf-tutorial

import .OSMPBF_pb
using ProtoBuf
using Mmap
using CodecZlib

# read a blob header (they occur many times in the file). off points to the offset into the PBF of the
# header length (which is four raw bytes before the protobuf header)
# all of these read_* functions return their data, and the offset in the file for the next read
function read_blob_header(pbf::Vector{UInt8}, off::Int64)::Tuple{OSMPBF_pb.BlobHeader,Int64}
    # read the length of the header
    hdr_len::Int32 = convert(Int32, ntoh(reinterpret(Int32, @view pbf[off:off + 3])[1]))
    off += 4

    buf = IOBuffer(@view pbf[off:off + hdr_len - 1])
    out = readproto(buf, OSMPBF_pb.BlobHeader())
    off += hdr_len
    return out, off
end

# This reads the header block which occurs once per file
function read_file_header!(pbf::Vector{UInt8}, off::Int64, blhdr::OSMPBF_pb.BlobHeader, out::OSMPBF_pb.HeaderBlock)::Int64
    # read the blob
    off = read_blob!(pbf, off, blhdr, out)
    return off + b
end

function read_blob(pbf::Vector{UInt8}, off::Int64, hdr::OSMPBF_pb.BlobHeader, T::Type)::Tuple{T, Int64}
    buf = IOBuffer(@view pbf[off : off + hdr.datasize - 1])
    blob = readproto(buf, OSMPBF_pb.Blob())

    if hasproperty(blob, :raw)
        buf = IOBuffer(blob.raw)
        out = readproto(buf, T())
    elseif hasproperty(blob, :zlib_data)
        buf = ZlibDecompressorStream(IOBuffer(blob.zlib_data))
        out = readproto(buf, T())
    else
        error("Blob not in raw or zlib format")
    end

    return out, off + hdr.datasize
end

# Parsing raw nodes (i.e. not packed into dense nodes)
# This function is untested because PBF files with regular rather than dense nodes are few and far between
function parse_nodes(nodelist::Vector{OSMPBF_pb.Node}, block::OSMPBF_pb.PrimitiveBlock, strtab::Vector{String}, cb::Function)
    @info "got $(length(nodelist)) nodes"
    for raw_node in nodelist
        nodeid = raw_node.id::Int64
        lat = 1e-9 * (block.lat_offset + (block.granularity * raw_node.lat))::Float64
        lon = 1e-9 * (block.lon_offset + (block.granularity * raw_node.lon))::Float64
        tags = map(Iterators.zip(raw_node.keys, raw_node.vals)) do t
            return strtab[t[1]] => strtab[t[2]]
        end |> Dict
        node = Node(
            nodeid,
            lat,
            lon,
            tags
        )
        cb(node)
    end

    return nothing
end

function parse_dense_nodes(dense::OSMPBF_pb.DenseNodes, block::OSMPBF_pb.PrimitiveBlock, strtab::Vector{String}, cb::Function)
    # first, parse and de-delta-code
    ids = Vector{Int64}()
    lats = Vector{Float64}()
    lons = Vector{Float64}()
    n = length(dense.id)
    sizehint!(ids, n)
    sizehint!(lats, n)
    sizehint!(lons, n)

    # read and de-delta-code ids
    current_id = 0
    for id in dense.id
        current_id += id
        push!(ids, current_id)
    end

    current_lat = 0
    for lat in dense.lat
        current_lat += lat
        push!(lats, 1e-9 * (block.lat_offset + (block.granularity * current_lat)))
    end

    current_lon = 0
    for lon in dense.lon
        current_lon += lon
        push!(lons, 1e-9 * (block.lon_offset + (block.granularity * current_lon)))
    end

    if length(dense.keys_vals) > 0
        # parse tags and yield simultaneously to save memory
        off = 1
        current_node = 1
        current_tags = Dict{String, String}()
        while off <= length(dense.keys_vals)
            if dense.keys_vals[off] == 0
                # yield this node
                Node(ids[current_node], lats[current_node], lons[current_node], current_tags) |> cb

                # prepare for next node
                current_tags = Dict{String, String}()  # can't just call empty! as caller may have a reference to tags from current node
                current_node += 1
                off += 1
                continue
            end

            key = strtab[dense.keys_vals[off]]
            val = strtab[dense.keys_vals[off + 1]]
            current_tags[key] = val
            off += 2
        end

        # yield final node if we have not already
        # the spec doesn't say whether the last element of keys_vals needs to be a zero. osmium does
        # zero-terminate these arrays, but I don't want to count on this behavior
        if current_node == n
            Node(ids[current_node], lats[current_node], lons[current_node], current_tags) |> cb
            current_node += 1
        end

        # make sure we have yielded everything
        @assert current_node == n + 1
    else
        # special case of no tags
        for node in 1:n
            Node(ids[current_node], lats[current_node], lons[current_node], Dict{String, String}()) |> cb
        end
    end

    return nothing
end

function parse_ways(ways::Vector{OSMPBF_pb.Way}, strtab::Vector{String}, cb::Function)
    for raw_way in ways
        wayid = raw_way.id::Int64

        nodes = Vector{Int64}()
        sizehint!(nodes, length(raw_way.refs))

        current_nodeid = 0
        for nodeid in raw_way.refs
            current_nodeid += nodeid
            push!(nodes, current_nodeid)
        end

        tags = map(Iterators.zip(raw_way.keys, raw_way.vals)) do t
            return strtab[t[1]] => strtab[t[2]]
        end |> Dict

        Way(wayid, nodes, tags) |> cb
    end

    return nothing
end

function parse_relations(relations::Vector{OSMPBF_pb.Relation}, strtab::Vector{String}, cb::Function)
    for raw_rel in relations
        relid = raw_rel.id::Int64
        tags = map(Iterators.zip(raw_rel.keys, raw_rel.vals)) do t
            return strtab[t[1]] => strtab[t[2]]
        end |> Dict
        members = Vector{RelationMember}()
        sizehint!(members, length(raw_rel.memids))

        memid = 0
        for i in 1:length(raw_rel.memids)
            memid += raw_rel.memids[i]  # de-delta-code
            role = strtab[raw_rel.roles_sid[i]]
            raw_type = raw_rel.types[i]
            type = raw_type == OSMPBF_pb.Relation_MemberType.NODE ? node :
                (raw_type == OSMPBF_pb.Relation_MemberType.WAY ? way : relation)
            push!(members, RelationMember(memid, type, role))
        end

        Relation(relid, members, tags) |> cb
    end
end

# hoping that the compiler is smart enough that the Any return types won't cause type instability, since the return values are not used...
function scan_pbf(pbffile; nodes::Union{Function, Missing}=missing, ways::Union{Function, Missing}=missing, relations::Union{Function, Missing}=missing)
    # read the file
    open(pbffile, "r") do stream
        pbf::Vector{UInt8} = Mmap.mmap(stream)

        off::Int64 = 1

        blhdr, off = read_blob_header(pbf, off)

        @assert blhdr._type == "OSMHeader"  "Malformed OSM PBF file, expected OSMData block but found $(blhdr._type)"

         file_header, off = read_blob(pbf, off, blhdr, OSMPBF_pb.HeaderBlock)

        @info "Reading file written by $(file_header.writingprogram)"

        # track these to avoid spammy log messages, only log once
        has_relations = false
        has_changesets = false

        # <= b/c julia starts array indices at 1, so pbf[length(pbf)] is the last byte of the file
        while (off <= length(pbf))
            # TODO should be okay to reuse blob headers right
            blhdr, off = read_blob_header(pbf, off)
            @assert blhdr._type == "OSMData"  "Malformed OSM PBF file, expected OSMData block but found $(blhdr._type)"

            block, off = read_blob(pbf, off, blhdr, OSMPBF_pb.PrimitiveBlock)

            # the primitiveblock contains several primitivegroups which may contain nodes, ways, etc.
            # but they reference string tables
            # drop index 1 (in Julia land)/index 0 (everywhere) which is always empty string and used as a delimiter
            # this also means that if a node etc say string x, strtab[x] is the correct string - no off-by-one errors
            # due to Julia starting arrays at 1.
            strtab::Vector{String} = map(String, block.stringtable.s[2:length(block.stringtable.s)])

            for grp in block.primitivegroup
                if (length(grp.nodes) > 0)
                    if !ismissing(nodes)
                        parse_nodes(grp.nodes, block, strtab, nodes)
                    end
                elseif hasproperty(grp, :dense) && length(grp.dense.id) > 0
                    if !ismissing(nodes)
                        parse_dense_nodes(grp.dense, block, strtab, nodes)
                    end
                elseif (length(grp.ways) > 0)
                    if !ismissing(ways)
                        parse_ways(grp.ways, strtab, ways)
                    end
                elseif (length(grp.relations) > 0)
                    if !ismissing(relations)
                        parse_relations(grp.relations, strtab, relations)
                    end
                elseif (length(grp.changesets) > 0)
                    if !has_changesets
                        @warn "changesets are present in $pbffile, but are not yet supported"
                        has_changesets = true  # don't warn again
                    end
                else
                    error("Malformed PBF file, PrimitiveGroup has no ways, nodes, dense nodes, relations, or changesets")
                end
            end
        end
    end
end


include("src/init.jl")

"""
This file contains the script to execute learning containing
the Value and Critic function in one Network!
"""

θ = init_θ_combined(params.structure_net_pl)
res = zeros(params.n_iter)
precision = zeros(params.n_iter)
all_traj  = init_all_traj(params)

function actor_critic_combined(params::Params, θ::Any, res::Array{Float64,1}, precision::Array{Float64,1}, all_traj::Traj_Data)
    #Init Part:
    #------------------------
    #Initialize the Network
    ĝ = copy(θ)
    opt_para = optimizers(θ,   Adam; lr = params.lr_pl)

    #Init Progress bar
    progress = Progress(params.n_iter,1)

    # Preallocate trajectory
    rms_state = init_rms(params.s_init)

    a_t_old = all_traj.traj[1].a_t
    weights_old = zeros(length(net_to_vec(θ)),2)
    #FIM = Symmetric(eye(length(net_to_vec(θ))))*0 #preallocate FisherMatrix

    for i_iter = 1:params.n_iter
        ĝ *= 0

        #Make a Video
        if floor(i_iter/400) == ceil(i_iter/400)
            figure(1)
            #t_horizon = (all_traj.traj[1].len)*(params.dt)
            t_plot = 0:params.dt:params.t_horizon;
            x1 = all_traj.traj[1].s_t
            u1 = all_traj.traj[1].a_t'
            filename = string("Results/Videos/Actor_Critic_",Dates.today(),"_iter_",i_iter,".mp4")
            visualise(:SegwayRobot, x1', t_plot, filename; options = Dict(:U => u1, :size => (1920, 1080)))
        end

        #-----------------------------------------------------------------------
        #Solve it in Parallel! (How to do asynchron updates..?)
        ĝ, res[i_iter], precision[i_iter] = parallel_episodes_combined(params,θ,all_traj,rms_state)

        #Reorganice Data in minibatches
        mini_batch = build_minibatches(all_traj,params)

        #Do the Policy / and Value updates
        θ = parallel_PPO_combined(params, θ, mini_batch, opt_para)

        #Visualization_stuff
        #-----------------------------------------------------------------------
        ProgressMeter.update!(progress,i_iter) #update progress bar
            a_t_old ,weights_old = plot_stuff(params, all_traj, a_t_old,
            weights_old, i_iter,θ, θ, res, precision/params.n_traj, ĝ, ĝ)
        #-----------------------------------------------------------------------
    end
end

actor_critic_combined(params, θ, res, precision, all_traj)
#save_data(res, precision, θa, θv, all_traj)


"A trait for the type of bins a histogram may have."
abstract type BinType end
"Each bin has the same predetermined width."
struct FixedWidth <: BinType end
"Bins have possibly different predetermined widths."
struct VariableWidth <: BinType end
"Bin widths are not known before computing the histogram (i.e. text data). Only 1D histograms are supported."
struct UnboundedWidth <: BinType end

"A trait for the ways the bin search and bin update steps can be parallelized."
abstract type HistogramParallelization end

"No threading nor vectorization."
struct NoParallelization <: HistogramParallelization end

"""
Threads that have private bin data structures that are reduced after their private updates.

Requires these functions to be defined for real-valued histograms:
- `get_subweights(hist)::AbstractArray{Int,N+1}` Returns the weights (i.e. counts) array for an N-dimensional histogram.

Requires these functions to be defined for text histograms:
- `get_subtable(hist)::AbstractVector{AbstractDict{String,Int}}` Returns a vector of independent tables.
"""
struct PrivateThreads <: HistogramParallelization end

"""
SIMD vectorization.

Requires these functions to be defined for real-valued histograms:
- `get_subweights(hist)::AbstractArray{Int,N+1}` Returns the weights (i.e. counts) array for an N-dimensional histogram.

Requires these functions to be defined for text histograms:
- `get_subtable(hist)::AbstractVector{AbstractDict{String,Int}}` Returns a vector of independent tables.
"""
struct SIMD <: HistogramParallelization end


using CUDA
using LULESH
using MPI
using Printf

if length(ARGS) < 2
    printUsage()
    error("Wrong input arguments")
end

if !occursin("-u", ARGS[1]) && !occursin("-s", ARGS[1])
    printUsage()
    error("Wrong input arguments")
end

num_iters = 10 # AD_LENGTH of simulation usually set to "-1" but for debugging set to 10!
if length(ARGS) == 4
    num_iters = parse(Int, ARGS[4])
end

#   bool structured = ( strcmp(argv[1],"-s") == 0 );
structured = occursin("-s", ARGS[1])
structured = true
@show structured
# assume cube subdomain geometry for now (nx)
nx = parse(IndexT, ARGS[2])
# nx = 45
# TODO: change default nr to 11
nr = 1
balance = 1
cost = 1
devicetype = Vector
floattype = Float64
# To disable MPI, pass 'nothing' as the last argument
prob = LuleshProblem(num_iters, structured, nx, nr, balance, cost, Vector, Float64, nothing)

#   Int_t numRanks ;
#   Int_t myRank ;

# #if USE_MPI
#   Domain_member fieldData ;

#   MPI_Init(&argc, &argv) ;
#   MPI_Comm_size(MPI_COMM_WORLD, &numRanks) ;
#   MPI_Comm_rank(MPI_COMM_WORLD, &myRank) ;
# #else
#   numRanks = 1;
#   myRank = 0;
# #endif

#   cuda_init(myRank);

#   Domain *locDom ;

# Set up the mesh and decompose. Assumes regular cubes for now

# TODO: modify this constructor to account for new fields
# TODO: setup communication buffers

domain = Domain(prob)

# #if USE_MPI
#    // copy to the host for mpi transfer
#    locDom->h_nodalMass = locDom->nodalMass;

#    fieldData = &Domain::get_nodalMass;

#    // Initial domain boundary communication
#    CommRecv(*locDom, MSG_COMM_SBN, 1,
#             locDom->sizeX + 1, locDom->sizeY + 1, locDom->sizeZ + 1,
#             true, false) ;
#    CommSend(*locDom, MSG_COMM_SBN, 1, &fieldData,
#             locDom->sizeX + 1, locDom->sizeY + 1, locDom->sizeZ + 1,
#             true, false) ;
#    CommSBN(*locDom, 1, &fieldData) ;

#    // copy back to the device
#    locDom->nodalMass = locDom->h_nodalMass;

#    // End initialization
#    MPI_Barrier(MPI_COMM_WORLD);
# #endif

#   cudaDeviceSetCacheConfig(cudaFuncCachePreferL1);
#   cudaDeviceSetLimit(cudaLimitMallocHeapSize,1024*1024*1024);

@show domain.time_h
@show domain.stoptime
if getMyRank(prob.comm) == 0
    if (structured)
        @printf("Running until t=%f, Problem size=%dx%dx%d\n", domain.stoptime, nx, nx, nx)
    else
        @printf("Running until t=%f, Problem size=%d \n", domain.stoptime, domain.numElem)
        @warn "Unstructured setup not supported"
    end
end
# timestep to solution
its = 0
#   cudaProfilerStart();
start = getWtime(prob.comm)
while domain.time_h < domain.stoptime
    # this has been moved after computation of volume forces to hide launch latencies
    # TimeIncrement(locDom)
    lagrangeLeapFrog(domain)
    # checkErrors(domain, its, myRank)
    if getMyRank(prob.comm) == 0
        @printf("cycle = %d, time = %e, dt=%e\n", its+1, domain.time_h, domain.deltatime_h)
    end
    global its += 1
    # if its == num_iters
        break
    # end
end

#   // make sure GPU finished its work
#   cudaDeviceSynchronize();

# Use reduced max elapsed time
elapsed_time = getWtime(prob.comm) - start
elapsed_timeG = comm_max(elapsed_time, prob.comm)

#   cudaProfilerStop();

#   if (myRank == 0)
#     VerifyAndWriteFinalOutput(elapsed_timeG, *locDom, its, nx, numRanks, structured);

# #ifdef SAMI
#   DumpDomain(locDom) ;
# #endif
#   cudaDeviceReset();

MPI.Finalize()


#_global_mt_setting = true

# Evaluate code either multi-threaded or single-threaded dependent on the global multithreading setting.
# TODO: Remove, find cleaner solution.
macro mt(expr)
    quote
        if _global_mt_setting
            @onthreads mt_nthreads() begin
                $(esc(expr))
            end
        else
            $(esc(expr))
        end
    end
end

function mt_nthreads()
    _global_mt_setting ? allthreads() : 1
end
const _minpoints_per_dimension = 50

"""
    hm_integrate!(result, settings = HMIPrecisionSettings())

*AHMI-internal, not part of stable public API.*

This function starts the adaptive harmonic mean integration. See arXiv:1808.08051 for more details.
It needs a HMIData struct as input, which holds the samples, in form of a dataset, the integration volumes and other
properties, required for the integration, and the final result.
"""
function hm_integrate!(
    result::HMIData{T, I, V},
    integrationvol::Symbol = :HyperRectangle;
    settings::HMISettings = HMIPrecisionSettings())::HMIData{T, I, V} where {T<:AbstractFloat, I<:Integer, V<:SpatialVolume}

    hm_init(result, settings)

    hm_whiteningtransformation!(result, settings)

    hm_createpartitioningtree!(result)

    notsinglemode = hm_findseeds!(result, settings)

	hm_determinetolerance!(result, settings) # tolerance is never Inf.

    # if !notsinglemode
    #     result.dataset1.tolerance = Inf
    #     result.dataset2.tolerance = Inf
    #     @warn "Tolerance set to Inf for single mode distributions"
    # else
    #     hm_determinetolerance!(result, settings)
    # end

    hm_create_integrationvolumes!(result, settings)
    hm_integrate_integrationvolumes!(result, settings)

    for pair in settings.uncertainty_estimators
        @info "Estimating Uncertainty ($(pair[1]))"
        result.integralestimates[pair[1]] = pair[2](result)
    end

    #Possible Uncertainty Estimators:
    #hm_combineresults_legacy!      see my master thesis for details
    #hm_combineresults_covweighted!  see arXiv:1808.08051 for details
    #hm_combineresults_analyticestimation! based on analytic uncertainty of both Z (see thesis) and r (binomial error)

    result
end

ahmi_integrate(samples::DensitySampleVector) = hm_integrate!(HMIData(samples))

"""
    hm_init!(result, settings)

*AHMI-internal, not part of stable public API.*

Sets the global multithreading setting and ensures that a minimum number of samples, dependent on the number of dimensions, are provided.
"""
function hm_init(
    result::HMIData{T, I, V},
    settings::HMISettings) where {T<:AbstractFloat, I<:Integer, V<:SpatialVolume}

    if result.dataset1.N < result.dataset1.P * _minpoints_per_dimension || result.dataset2.N < result.dataset2.P * _minpoints_per_dimension
        @error "Not enough samples for integration"
    end
    @assert result.dataset1.P == result.dataset2.P

    global _global_mt_setting = settings.useMultiThreading

    @info "Harmonic Mean Integration started. Samples in dataset 1 / 2: \t$(result.dataset1.N) / $(result.dataset2.N)\tParameters:\t$(result.dataset1.P)"
end

"""
    hm_whiteningtransformation!(result, settings)

*AHMI-internal, not part of stable public API.*

Applies a whitening transformation to the samples. A custom whitening method can be used by overriding settings.whitening_function!
"""
function hm_whiteningtransformation!(
    result::HMIData{T, I, V},
    settings::HMISettings) where {T<:AbstractFloat, I<:Integer, V<:SpatialVolume}

    if !isinitialized(result.whiteningresult)
        @info "Data Whitening."
        result.whiteningresult = settings.whitening_function!(result.dataset1)
    end

    if !result.dataset1.iswhitened
        @info "Apply Whitening Transformation to Data Set 1"
        transform_data!(result.dataset1, result.whiteningresult.whiteningmatrix, result.whiteningresult.meanvalue, false)
    end
    if !result.dataset2.iswhitened
        @info "Apply Whitening Transformation to Data Set 2"
        transform_data!(result.dataset2, result.whiteningresult.whiteningmatrix, result.whiteningresult.meanvalue, true)
    end
end

function hm_createpartitioningtree!(
    result::HMIData{T, I, V}) where {T<:AbstractFloat, I<:Integer, V<:SpatialVolume}

    maxleafsize = 200
    progress_steps = ((!isinitialized(result.dataset1.partitioningtree) ? result.dataset1.N / maxleafsize : 0.0)
                    + (!isinitialized(result.dataset2.partitioningtree) ? result.dataset2.N / maxleafsize : 0.0))
    progressbar = Progress(round(Int, progress_steps))
    progress_steps > 0 && @info "Create Space Partitioning Tree"
    !isinitialized(result.dataset1.partitioningtree) && create_search_tree(result.dataset1, progressbar, maxleafsize = maxleafsize)
    !isinitialized(result.dataset2.partitioningtree) && create_search_tree(result.dataset2, progressbar, maxleafsize = maxleafsize)
    finish!(progressbar)
end

function hm_findseeds!(
    result::HMIData{T, I, V},
    settings::HMISettings) where {T<:AbstractFloat, I<:Integer, V<:SpatialVolume}

    #if hyper-rectangles are already created and only need to be populated using different data sets
    if !isempty(result.volumelist1) && !isempty(result.volumelist2)
        return true
    end

    notsinglemode = true
    @info "Determine Hyperrectangle Starting Samples"
    if isempty(result.dataset1.startingIDs)
        notsinglemode &= find_hypercube_centers(result.dataset1, result.whiteningresult, settings)
    end
    if isempty(result.dataset2.startingIDs)
        notsinglemode &= find_hypercube_centers(result.dataset2, result.whiteningresult, settings)
    end

    notsinglemode
end

function hm_determinetolerance!(
    result::HMIData{T, I, V},
    settings::HMISettings) where {T<:AbstractFloat, I<:Integer, V<:SpatialVolume}

    #if hyper-rectangles are already created and only need to be populated using different data sets
    if !isempty(result.volumelist1) && !isempty(result.volumelist2)
        return
    end

    (iszero(result.dataset1.tolerance) || iszero(result.dataset2.tolerance)) && @info "Determine Tolerances for Hyperrectangle Creation"

    if iszero(result.dataset1.tolerance)
        suggTolPts = max(result.dataset1.P * 4, ceil(I, sqrt(result.dataset1.N)))
        findtolerance(result.dataset1, min(10, settings.warning_minstartingids), suggTolPts)
        @debug "Tolerance Data Set 1: $(result.dataset1.tolerance)"
    end
    if iszero(result.dataset2.tolerance)
        suggTolPts = max(result.dataset2.P * 4, ceil(I, sqrt(result.dataset2.N)))
        findtolerance(result.dataset2, min(10, settings.warning_minstartingids), suggTolPts)
        @debug "Tolerance Data Set 2: $(result.dataset2.tolerance)"
    end
end

function hm_combineresults_legacy!(result::HMIData{T, I, V}) where {T<:AbstractFloat, I<:Integer, V<:SpatialVolume}
    result_legacy = HMIResult(T)
    result_legacy.result1, result_legacy.dat2 = hm_combineresults_legacy_dataset!(result.dataset1, result.integrals2, result.volumelist2)
    result_legacy.result2, result_legacy.dat2 = hm_combineresults_legacy_dataset!(result.dataset2, result.integrals1, result.volumelist1)
    result_legacy.final = HMIEstimate(result_legacy.result1, result_legacy.result2)

    result_legacy
end

function hm_combineresults_legacy_dataset!(
    dataset::DataSet{T, I},
    integralestimates::IntermediateResults{T},
    volumes::Array{IntegrationVolume{T, I, V}, 1}) where {T<:AbstractFloat, I<:Integer, V<:SpatialVolume}

    dat = Dict{String, Any}()

    pweights = create_pointweights(dataset, volumes, integralestimates.volumeID)

    weights_overlap::Array{T, 1} = zeros(T, length(integralestimates))

    for i = 1:length(integralestimates)
        vol_id = integralestimates.volumeID[i]
        trw = sum(dataset.weights[volumes[vol_id].pointcloud.pointIDs])
        for id in eachindex(volumes[vol_id].pointcloud.pointIDs)
            weights_overlap[i] += 1.0 / trw / pweights[volumes[vol_id].pointcloud.pointIDs[id]]
        end
    end

    M = length(integralestimates)

    i_std = mean(integralestimates.integrals, ProbabilityWeights(weights_overlap))
    e_std = sqrt(var(integralestimates.integrals) / sum(weights_overlap))

    HMIEstimate(i_std, e_std, weights_overlap), dat
end

function hm_combineresults_covweighted!(result::HMIData{T, I, V}) where {T<:AbstractFloat, I<:Integer, V<:SpatialVolume}
    result_covweighted = HMIResult(T)
    result_covweighted.result1, result_covweighted.dat1 = hm_combineresults_covweighted_dataset!(result.dataset1,
        result.integrals2, result.volumelist2)
    result_covweighted.result2, result_covweighted.dat2 = hm_combineresults_covweighted_dataset!(result.dataset2,
        result.integrals1, result.volumelist1)
    result_covweighted.final = HMIEstimate(result_covweighted.result1, result_covweighted.result2)

    result_covweighted
end

function hm_combineresults_covweighted_dataset!(
    dataset::DataSet{T, I},
    integralestimates::IntermediateResults{T},
    volumes::Array{IntegrationVolume{T, I, V}, 1}) where {T<:AbstractFloat, I<:Integer, V<:SpatialVolume}

    dat = Dict{String, Any}()
    M = length(integralestimates)

    dat["Σ"] = Σ = cov(integralestimates.Y) ./ dataset.nsubsets

    weights_cov = 1 ./ diag(Σ)
    weights_cov /= sum(weights_cov)

    i_cov = mean(integralestimates.integrals, AnalyticWeights(weights_cov))

    e_cov = 0.0
    for i=1:M
        for j=1:M
            e_cov += weights_cov[i] * weights_cov[j] * Σ[i, j]
        end
    end

    HMIEstimate(i_cov, convert(T, sqrt(e_cov)), weights_cov), dat
end

function hm_combineresults_analyticestimation!(result::HMIData{T, I, V}) where {T<:AbstractFloat, I<:Integer, V<:SpatialVolume}
    result_analytic = HMIResult(T)
    result_analytic.result1, result_analytic.dat1 = hm_combineresults_analyticestimation_dataset!(result.dataset1,
        result.whiteningresult.determinant, result.integrals2, result.volumelist2)
    result_analytic.result2, result_analytic.dat2 = hm_combineresults_analyticestimation_dataset!(result.dataset2,
        result.whiteningresult.determinant, result.integrals1, result.volumelist1)
    result_analytic.final = HMIEstimate(result_analytic.result1, result_analytic.result2)

    result_analytic
end

function hm_combineresults_analyticestimation_dataset!(
    dataset::DataSet{T, I},
    determinant::T,
    integralestimates::IntermediateResults{T},
    volumes::Array{IntegrationVolume{T, I, V}, 1}) where {T<:AbstractFloat, I<:Integer, V<:SpatialVolume}

    dat = Dict{String, Any}()

    uncertainty_r = Array{T, 1}(undef, length(integralestimates.integrals))
    uncertainty_Y = Array{T, 1}(undef, length(integralestimates.integrals))
    uncertainty_tot = Array{T, 1}(undef, length(integralestimates.integrals))
    f_min = Array{T, 1}(undef, length(integralestimates.integrals))
    f_max = Array{T, 1}(undef, length(integralestimates.integrals))
    μ_Z = Array{T, 1}(undef, length(integralestimates.integrals))
    σ_μZ_sq = Array{T, 1}(undef, length(integralestimates.integrals))
    integral1 = Array{T, 1}(undef, length(integralestimates.integrals))
    integral2 = Array{T, 1}(undef, length(integralestimates.integrals))
    ess = Array{T, 1}(undef, length(integralestimates.integrals))

    @mt for i in workpart(eachindex(integralestimates.integrals), mt_nthreads(), threadid())
        uncertainty_r[i], uncertainty_Y[i], uncertainty_tot[i], f_min[i], f_max[i], μ_Z[i], σ_μZ_sq[i], integral1[i], integral2[i], ess[i] =
            calculateuncertainty(dataset, volumes[integralestimates.volumeID[i]], determinant, integralestimates.integrals[i])
    end

    dat["uncertainty_r"] = uncertainty_r
    dat["uncertainty_Y"] = uncertainty_Y
    dat["uncertainty_tot"] = uncertainty_tot
    dat["f_min"] = f_min
    dat["f_max"] = f_max
    dat["μ_Z"] = μ_Z
    dat["σ_μZ_sq"] = σ_μZ_sq
    dat["integral1_sq_cuba"] = integral1
    dat["integral2_cuba"] = integral2
    dat["ess"] = ess

    weights_cov = 1 ./ uncertainty_tot
    weights_cov /= sum(weights_cov)

    i_cov = mean(integralestimates.integrals, AnalyticWeights(weights_cov))

    dat["overlap"] = overlap = calculate_overlap(dataset, volumes, integralestimates)
    M = length(integralestimates)

    Σ = Array{T, 2}(undef, M, M)
    for i=1:M
        for j=1:M
            Σ[i,j] =  uncertainty_tot[i] * uncertainty_tot[j] * overlap[i, j]
        end
    end
    dat["Σ"] = Σ

    e_cov = 0.0
    for i=1:M
        for j=1:M
            e_cov += weights_cov[i] * weights_cov[j] * Σ[i, j]
        end
    end

	HMIEstimate(i_cov, convert(T, sqrt(e_cov)), weights_cov), dat
end

function findtolerance(
    dataset::DataSet{T, I},
    ntestcubes::I,
    pts::I) where {T<:AbstractFloat, I<:Integer}

    ntestpts = [2, 4, 8] * pts
    # @debug "Tolerance Test Cube Points: $([pts, ntestpts...])"

    vInc = zeros(T, ntestcubes * (length(ntestpts) + 1))
    pInc = zeros(T, ntestcubes * (length(ntestpts) + 1))

    startingIDs = dataset.startingIDs

    cntr = 1
    for id = 1:ntestcubes
        c = find_density_test_cube(dataset.data[:, startingIDs[id]], dataset, pts)
        prevv = c[1]^dataset.P
        prevp = c[2].pointcloud.points
        for i in ntestpts
            c = find_density_test_cube(dataset.data[:, startingIDs[id]], dataset, i)
            v = c[1]^dataset.P
            p = c[2].pointcloud.points

            vInc[cntr] = v/prevv
            pInc[cntr] = p/prevp
            cntr += 1
            prevv = v
            prevp = p
        end
    end
    tols = vInc ./ pInc

    i = length(tols)
    while i > 0
        if isnan(tols[i]) || isinf(tols[i]) || tols[i] <= 1.0
            deleteat!(tols, i)
        end
        i -= 1
    end

    # @debug "Tolerance List: $tols"

    suggTol::T = if length(tols) < 4
        @warn "Tolerance calculation failed. Tolerance is set to default to 1.5"
        3
    else
        (mean(trim(tols)) - 1) * 4 +1
    end

    dataset.tolerance = suggTol
end

function hm_integrate_integrationvolumes!(
    result::HMIData{T, I, V},
    settings::HMISettings) where {T<:AbstractFloat, I<:Integer, V<:SpatialVolume}

    nRes = length(result.volumelist1) + length(result.volumelist2)
    @info "Integrating $nRes Hyperrectangles"

    progressbar = Progress(nRes)

    result.integrals1, result.rejectedrects1 = hm_integrate_integrationvolumes!_dataset(result.volumelist1, result.dataset2, result.whiteningresult.determinant, progressbar, settings)
    result.integrals2, result.rejectedrects2 = hm_integrate_integrationvolumes!_dataset(result.volumelist2, result.dataset1, result.whiteningresult.determinant, progressbar, settings)

    finish!(progressbar)
end

function hm_integrate_integrationvolumes!_dataset(
    volumes::Array{IntegrationVolume{T, I, V}},
    dataset::DataSet{T, I},
    determinant::T,
    progressbar::Progress,
    settings::HMISettings)::Tuple{IntermediateResults{T}, Vector{I}} where {T<:AbstractFloat, I<:Integer, V<:SpatialVolume}

    if length(volumes) < 1
        @error "No hyper-rectangles could be created. Try integration with more points or different settings."
    end

    integralestimates = IntermediateResults(T, length(volumes))
    integralestimates.Y = zeros(T, dataset.nsubsets, length(volumes))

    @mt for i in workpart(eachindex(volumes), mt_nthreads(), threadid())
        integralestimates.Y[:, i], integralestimates.integrals[i] = integrate_hyperrectangle_cov(dataset, volumes[i], determinant)

        @critical next!(progressbar)
    end

    rejectedids = trim(integralestimates, settings.dotrimming)

    return integralestimates, rejectedids
end

"""
    reduced_volume_hm!(log_prob, sample_weights, volume_size, n_total, weight_total, bias_correction==true)

Estimate reduced volume harmonic mean for given arguments. Includes bias corection by default.
Note: Weights are expected to be frequency weights (i.e. counts).
"""
function reduced_volume_hm(log_prob::Array{T}, sample_weights::Array{T}, volume_size::AbstractFloat, n_total::Integer, weight_total::AbstractFloat; bias_correction::Bool=true) where {T<:AbstractFloat}

	pedestal_llh = maximum(log_prob) # maximum log likelihood to be subtracted for numerical stability

	tmp = log_prob .- pedestal_llh
	tmp = T(1) ./ exp.(tmp)
	mean_x = mean(tmp, weights(sample_weights))

	r = weight_total/sum(sample_weights)
	reduced_volume_hm = r*volume_size/mean_x*exp(pedestal_llh)

	if !bias_correction
		return reduced_volume_hm
	else
 		var_x = var(tmp, weights(sample_weights), corrected=false) / sum(sample_weights)
 		mean_r_bias = length(sample_weights) / n_total
 		var_r_bias = mean_r_bias*(1-mean_r_bias)/n_total
 		corr_factor = 1 - var_x / mean_x^2 - var_r_bias / mean_r_bias^2
		return reduced_volume_hm*corr_factor
	end
end

"""
Estimates reduced_volume_hm quantity for one hyperrectangle including estimates from batches (for covariance calculations).
"""
function integrate_hyperrectangle_cov(
    dataset::DataSet{T, I},
    integrationvol::IntegrationVolume{T, I, V},
    determinant::T)::Tuple{Array{T, 1}, T} where {T<:AbstractFloat, I<:Integer, V<:SpatialVolume}

	indices = integrationvol.pointcloud.pointIDs
	integral = reduced_volume_hm(dataset.logprob[indices], dataset.weights[indices], integrationvol.volume, length(dataset.weights), sum(dataset.weights))

	integrals_batches = zeros(T, dataset.nsubsets)

    for i in 1:dataset.nsubsets
        subset_indices = indices[dataset.ids[indices] .== i]
        tot_subset_mask = dataset.ids .== i

		if length(subset_indices) != 0
			integrals_batches[i] = reduced_volume_hm(dataset.logprob[subset_indices], dataset.weights[subset_indices], integrationvol.volume, length(dataset.weights[tot_subset_mask]), sum(dataset.weights[tot_subset_mask]))
		else
			@warn "The batch contains zero samples. It's integral will be replaced with the average batch integral."
			integrals_batches[i] = NaN
		end
    end

    for i in eachindex(integrals_batches)
        if isnan(integrals_batches[i]) || isinf(integrals_batches[i])
            integrals_batches[i] = integral
        end
    end

    return integrals_batches./determinant, integral/determinant
end


exact_energies = Dict(pc"He" => [-2.8616800, [o"1s" => -0.91795555]],
                      pc"Be" => [-14.573023, [o"1s" => -4.7326698, o"2s" => -0.3092695]],

                      pc"Ne" => [-128.54710, [o"1s" => -32.7724455, o"2s" => -1.93039095, o"2p" => -0.85040965]],
                      PseudoPotentials.NeonHF => [-34.709465, [o"2s" => -1.93039095, o"2p" => -0.85040965]],
                      PseudoPotentials.NeonWB => [-34.708785, [o"2s" => -1.93039095, o"2p" => -0.85040965]],

                      pc"Mg" => [-199.61463, [o"1s" => -49.0317255, o"2s" => -3.767718, o"2p" => -2.2822236, o"3s" => -0.25305275]],

                      pc"Ar" => [-526.81751, [o"1s" => -237.22070/2, o"2s" => -24.644306/2, o"2p" => -19.142932/2,
                                             o"3s" => -2.5547063/2, o"3p" => -1.1820348/2]],
                      PseudoPotentials.ArgonHF => [-20.849867, [o"3s" => -2.5547063/2, o"3p" => -1.1820348/2]],
                      PseudoPotentials.ArgonWB => [-20.884584, [o"3s" => -2.5547063/2, o"3p" => -1.1820348/2]],

                      pc"Ca" => [-676.75818, [o"1s" => -298.72744/2, o"2s" => -33.645481/2, o"2p" => -27.258531/2,
                                          o"3s" => -4.4907488/2, o"3p" => -2.6814114/2, o"4s" => -0.3910594/2]],
                      pc"Zn" => [-1777.8481, [o"1s" => -706.60909/2, o"2s" => -88.723452/2, o"2p" => -77.849691/2,
                                             o"3s" => -11.275642/2, o"3p" => -7.6787581/2, o"3d" => -1.5650843/2,
                                             o"4s" => -0.5850141/2]],
                      pc"Kr" => [-2752.0550, [o"1s" => -1040.3309/2, o"2s" => -139.80617/2, o"2p" => -126.01957/2,
                                             o"3s" => -21.698934/2, o"3p" => -16.663004/2, o"3d" => -7.6504697/2,
                                             o"4s" => -2.3058703/2, o"4p" => -1.0483734/2]],
                      pc"Xe" => [-7232.1384, [o"1s" => -2448.7956/2, o"2s" => -378.68024/2, o"2p" => -355.56490/2,
                                             o"3s" => -80.351324/2, o"3p" => -70.443321/2, o"3d" => -52.237736/2,
                                             o"4s" => -15.712602/2, o"4p" => -12.016674, o"4d" => -5.5557597,
                                             o"5s" => -1.8888279/2, o"5p" => -0.9145793/2]],

                      PseudoPotentials.XenonHF => [-14.989100, [o"5s" => -1.8888279/2, o"5p" => -0.9145793/2]],
                      PseudoPotentials.XenonWB => [-15.277055, [o"5s" => -1.8888279/2, o"5p" => -0.9145793/2]],
                      PseudoPotentials.XenonDF2c => [-328.74543, [o"5s" => -1.0097, ro"5p-" => -0.4915, ro"5p" => -0.4398]])

function shift_unit(u::U, d) where {U<:Unitful.Unit}
    d = Int(3floor(Int, d/3))
    iszero(d) && return u
    for tens = Unitful.tens(u) .+ (d:(-3*sign(d)):0)
        haskey(Unitful.prefixdict, tens) && return U(tens, u.power)
    end
    u
end

function shift_unit(u::Unitful.FreeUnits, d)
    tu = typeof(u)
    us,ds,a = tu.parameters

    uu = shift_unit(us[1], d)

    Unitful.FreeUnits{(uu,us[2:end]...), ds, a}()
end

function si_round(q::Quantity; fspec="{1:+9.4f} {2:s}")
    v,u = ustrip(q), unit(q)
    if !iszero(v)
        u = shift_unit(u, log10(abs(v)))
        q = u(q)
    end
    format(fspec, ustrip(q), unit(q))
end

function energy_errors(fock, exact_energies, Δ, δ)
    atom = fock.quantum_system

    Eexact,orbExact = exact_energies
    exact_energies = if first(atom.orbitals) isa SpinOrbital
        vcat(Eexact, [repeat([E], degeneracy(o)) for (o,E) in orbExact]...)
    else
        vcat(Eexact, [E for (o,E) in orbExact]...)
    end

    orbital_refs = unique(nonrelorbital.(first.(orbExact)))

    H = zeros(1,1)
    Etot = SCF.energy_matrix!(H, fock)[1,1]
    energies = vcat(Etot, collect(SCF.energy(fock.equations.equations[i])/degeneracy(o)
                                  for (i,o) in enumerate(atom.orbitals)
                                  if nonrelorbital(Atoms.getspatialorb(o)) ∈ orbital_refs))
    errors = energies - exact_energies

    labels = vcat("Total", string.(collect(o for o in atom.orbitals
                                           if nonrelorbital(Atoms.getspatialorb(o)) ∈ orbital_refs)))

    Ha_formatter = (v,i) -> si_round(v*u"hartree")
    eV_formatter = (v,i) -> si_round(v*u"eV")

    pretty_table([labels exact_energies energies errors 27.211energies 27.211errors errors./abs.(exact_energies)],
                 ["", "HF limit", "Energy", "Δ", "Energy", "Δ", "Δrel"],
                 formatter=Dict(
                     2 => Ha_formatter,
                     3 => Ha_formatter,
                     4 => Ha_formatter,
                     5 => eV_formatter,
                     6 => eV_formatter,
                     7 => (v,i) -> si_round(100v*u"percent")
                 ),
                 highlighters=(Highlighter((v,i,j) -> abs(v[i,7])>0.2, foreground=:red, bold=true),
                               Highlighter((v,i,j) -> abs(v[i,7])>0.01, foreground=:yellow, bold=true),))

    @test abs(errors[1]) < Δ
    @test all(abs.(errors[2:end]) .< δ)
end

function atom_calc(nucleus::AbstractPotential, grid_type::Symbol, rₘₐₓ, ρ,
                   Δ, δ;
                   config_transform=identity, kwargs...)
    # If we're using a pseudopotential, we don't really need higher
    # order of the FEDVR basis functions in the first finite element,
    # since the orbitals almost vanish there.
    R,r = get_atom_grid(grid_type, rₘₐₓ, ρ, nucleus,
                        amend_order=nucleus isa PointCharge)

    atom = Atom(R, [spin_configurations(config_transform(ground_state(nucleus)))[1]],
                nucleus, eltype(R))

    fock = Fock(atom)

    optimize!(fock; kwargs...)
    energy_errors(fock, exact_energies[nucleus], Δ, δ)
end

@testset "Calculation accuracy" begin
    @testset "Helium" begin
        @testset "$(orb_type) orbitals" for (orb_type,config_transform) in [
            ("Non-relativistic", identity),
            ("Relativistic", relconfigurations)
        ]
            @testset "$(grid_type)" for (grid_type,Δ,δ) in [(:fedvr,6e-9,7e-9),
                                                            (:fd,4e-3,4e-3)]
                atom_calc(pc"He", grid_type, 10.0, 0.1, Δ, δ, ω=0.9,
                          config_transform=config_transform)
            end
        end
    end

    @testset "Beryllium" begin
        @testset "$(grid_type)" for (grid_type,ρ,Δ,δ) in [(:fedvr,0.2,6e-7,2e-6),
                                                          (:fd,0.05,0.02,0.009)]
            atom_calc(pc"Be", grid_type, 15.0, ρ, Δ, δ, ω=0.9, scf_method=:arnoldi)
        end
    end

    @testset "Neon" begin
        @testset "$nucleus" for (nucleus,Δ,δ) in [(pc"Ne",0.005,0.005),
                                                  (PseudoPotentials.NeonHF,0.2,0.05)]
            atom_calc(nucleus, :fedvr, 10, 0.2, Δ, δ,
                      ω=0.999, ωmax=1.0-1e-3, scf_method=:arnoldi)
        end
    end

    @testset "Xenon" begin
        @testset "$nucleus" for (nucleus,grid_type,ρ,Δ,δ,config_transform) in [
            (PseudoPotentials.XenonHF,:fd,0.1,4e-3,4e-3,identity),
            (PseudoPotentials.XenonDF2c,:fd,0.1,0.2,1e-2,relconfigurations)
        ]
            atom_calc(nucleus, grid_type, 7.0, ρ, Δ, δ,
                      ω=0.999, ωmax=1.0-1e-3,
                      config_transform=config_transform, scf_method=:arnoldi)
        end
    end
end


using Base.Iterators: product

"""
    RectilinearGrid{Tl, Tr} <: AbstractVector{Tuple{Tl, Tr}}

A `RectilinearGrid` is parametrised by `AbstractVector`s of points `xl` and `xr`, whose
element types are `Tl` and `Tr` respectively, comprising `length(xl) * length(xr)`
elements. Linear indexing is the same as `product(eachindex(xl), eachindex(xr))` - `xl`
iterates more quickly than `xr`.
"""
struct RectilinearGrid{
    Tl, Tr, Txl<:AbstractVector{Tl}, Txr<:AbstractVector{Tr},
} <: AbstractVector{Tuple{Tl, Tr}}
    xl::Txl
    xr::Txr
end

get_space(x::RectilinearGrid) = Zygote.literal_getfield(x, Val(:xl))

get_times(x::RectilinearGrid) = Zygote.literal_getfield(x, Val(:xr))

Base.size(X::RectilinearGrid) = (length(X.xl) * length(X.xr),)

function Base.collect(X::RectilinearGrid{Tl, Tr}) where{Tl, Tr}
    return vec(
        map(
            ((p, q),) -> (X.xl[p], X.xr[q]),
            product(eachindex(X.xl), eachindex(X.xr)),
        )
    )
end

function Base.getindex(X::RectilinearGrid, n::Integer)
    return (X.xl[mod(n - 1, length(X.xl)) + 1], X.xr[div(n - 1, length(X.xl)) + 1])
end

Base.show(io::IO, x::RectilinearGrid) = Base.show(io::IO, collect(x))

"""
    SpaceTimeGrid{Tr, Tt<:Real}

A `SpaceTimeGrid` is a `RectilinearGrid` in which the left vector corresponds to space, and
the right `time`. The left eltype is arbitrary, but the right must be `Real`.
"""
const SpaceTimeGrid{Tr, Tt<:Real} = RectilinearGrid{
    Tr, Tt, <:AbstractVector{Tr}, <:AbstractVector{Tt},
}

#
# Implement internal API for transforming between "flat" representation, which is useful for
# GPs, an a time-centric representation, which is useful for state-space models.
#

# See docstring elsewhere for context.
function inputs_to_time_form(x::SpaceTimeGrid)
    return Fill(get_space(x), length(get_times(x)))
end

# See docstring elsewhere for context.
function merge_inputs(x1::SpaceTimeGrid, x2::SpaceTimeGrid)
    if get_space(x1) != get_space(x2)
        throw(error("Space coords of inputs not compatible, cannot merge."))
    end
    return RectilinearGrid(get_space(x1), vcat(get_times(x1), get_times(x2)))
end

# See docstring elsewhere for context.
function sort_in_time(x::SpaceTimeGrid)
    idx = sortperm(get_times(x))
    return idx, RectilinearGrid(get_space(x), get_times(x)[idx])
end

# See docstring elsewhere for context.
function observations_to_time_form(x::SpaceTimeGrid, y::AbstractVector{<:Union{Real, Missing}})
    return restructure(y, Fill(length(get_space(x)), length(get_times(x))))
end

function observations_to_time_form(x::SpaceTimeGrid, ::AbstractVector{Missing})
    return fill(missing, length(get_times(x)))
end

function get_zeros(x::SpaceTimeGrid{T}) where {T<:Real}
    return fill(Diagonal(zeros(T, length(get_space(x)))), length(get_times(x)))
end

# See docstring elsewhere for context.
function noise_var_to_time_form(x::RectilinearGrid, S::Diagonal{<:Real})
    vs = restructure(
        diag(S),
        Zygote.ignore() do
            Fill(length(get_space(x)), length(get_times(x)))
        end,
    )
    return zygote_friendly_map(v -> Diagonal(collect(v)), vs)
end

destructure(::RectilinearGrid, y::AbstractVector) = reduce(vcat, y)


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

using Plots

"""
Plots a NeuralFMU (ME).
"""
function fmiPlot(nfmu::ME_NeuralFMU)
    FMI.fmiPlot(nfmu.fmu, nfmu.solution, true)
end

"""
Plots a NeuralFMU (CS).
"""
function fmiPlot(nfmu::CS_NeuralFMU)
    numSig = length(nfmu.valueStack[1])-1
    t = collect(data[1] for data in nfmu.valueStack)

    fig = Plots.plot(xlabel="t [s]")
    for i in 1:numSig
        values = collect(data[i+1] for data in nfmu.valueStack)
        Plots.plot!(fig, t, values, label="NeuralFMU output #$(i)")
    end

    fig
end


using Base.Test
using BandedMatrices
import BandedMatrices: rowstart,
                       rowstop,
                       colstart,
                       colstop,
                       rowlength,
                       collength,
                       diaglength

# rowstart/rowstop business
let
    a = bones(7, 5, 1, 2)
    # 1.0  1.0  1.0  0.0  0.0
    # 1.0  1.0  1.0  1.0  0.0
    # 0.0  1.0  1.0  1.0  1.0
    # 0.0  0.0  1.0  1.0  1.0
    # 0.0  0.0  0.0  1.0  1.0
    # 0.0  0.0  0.0  0.0  1.0
    # 0.0  0.0  0.0  0.0  0.0
    @test (rowstart(a, 1), rowstop(a, 1), rowlength(a, 1)) == (1, 3, 3)
    @test (rowstart(a, 2), rowstop(a, 2), rowlength(a, 2)) == (1, 4, 4)
    @test (rowstart(a, 3), rowstop(a, 3), rowlength(a, 3)) == (2, 5, 4)
    @test (rowstart(a, 4), rowstop(a, 4), rowlength(a, 4)) == (3, 5, 3)
    @test (rowstart(a, 5), rowstop(a, 5), rowlength(a, 5)) == (4, 5, 2)
    @test (rowstart(a, 6), rowstop(a, 6), rowlength(a, 6)) == (5, 5, 1)
    @test (rowstart(a, 7), rowstop(a, 7), rowlength(a, 7)) == (6, 5, 0) # zero length
    @test (colstart(a, 1), colstop(a, 1), collength(a, 1)) == (1, 2, 2)
    @test (colstart(a, 2), colstop(a, 2), collength(a, 2)) == (1, 3, 3)
    @test (colstart(a, 3), colstop(a, 3), collength(a, 3)) == (1, 4, 4)
    @test (colstart(a, 4), colstop(a, 4), collength(a, 4)) == (2, 5, 4)
    @test (colstart(a, 5), colstop(a, 5), collength(a, 5)) == (3, 6, 4)

    a = bones(3, 6, 1, 2)
    # 1.0  1.0  1.0  0.0  0.0  0.0
    # 1.0  1.0  1.0  1.0  0.0  0.0
    # 0.0  1.0  1.0  1.0  1.0  0.0
    @test (rowstart(a, 1), rowstop(a, 1), rowlength(a, 1)) == (1, 3, 3)
    @test (rowstart(a, 2), rowstop(a, 2), rowlength(a, 2)) == (1, 4, 4)
    @test (rowstart(a, 3), rowstop(a, 3), rowlength(a, 3)) == (2, 5, 4)
    @test (colstart(a, 1), colstop(a, 1), collength(a, 1)) == (1, 2, 2)
    @test (colstart(a, 2), colstop(a, 2), collength(a, 2)) == (1, 3, 3)
    @test (colstart(a, 3), colstop(a, 3), collength(a, 3)) == (1, 3, 3)
    @test (colstart(a, 4), colstop(a, 4), collength(a, 4)) == (2, 3, 2)
    @test (colstart(a, 5), colstop(a, 5), collength(a, 5)) == (3, 3, 1)
    @test (colstart(a, 6), colstop(a, 6), collength(a, 6)) == (4, 3, 0) # zero length
end

# test length of diagonal
let
    A = bones(4, 6, 2, 3)
    # 4x6 BandedMatrices.BandedMatrix{Float64}:
    #  1.0  1.0  1.0  1.0
    #  1.0  1.0  1.0  1.0  1.0
    #  1.0  1.0  1.0  1.0  1.0  1.0
    #       1.0  1.0  1.0  1.0  1.0
    @test diaglength(A, -2) == 2
    @test diaglength(A, -1) == 3
    @test diaglength(A,  0) == 4
    @test diaglength(A,  1) == 4
    @test diaglength(A,  2) == 4
    @test diaglength(A,  3) == 3

    A = bones(6, 5, 2, 1)
    # 6x5 BandedMatrices.BandedMatrix{Float64}:
    #  1.0  1.0
    #  1.0  1.0  1.0
    #  1.0  1.0  1.0  1.0
    #       1.0  1.0  1.0  1.0
    #            1.0  1.0  1.0
    #                 1.0  1.0
    @test diaglength(A, -2) == 4
    @test diaglength(A, -1) == 5
    @test diaglength(A,  0) == 5
    @test diaglength(A,  1) == 4

    A = bones(4, 4, 1, 1)
    # 4x4 BandedMatrices.BandedMatrix{Float64}:
    #  1.0  1.0
    #  1.0  1.0  1.0
    #       1.0  1.0  1.0
    #            1.0  1.0
    @test diaglength(A, -1) == 3
    @test diaglength(A,  0) == 4
    @test diaglength(A,  1) == 3

    # test with Band type
    @test diaglength(A,  band(1)) == 3
end

# _firstdiagrow/_firstdiagcol
let
    A = bones(6, 5, 2, 1)
    # 6x5 BandedMatrices.BandedMatrix{Float64}:
    #  1.0  1.0
    #  1.0  1.0  1.0
    #  1.0  1.0  1.0  1.0
    #       1.0  1.0  1.0  1.0
    #            1.0  1.0  1.0
    #                 1.0  1.0
    @test BandedMatrices._firstdiagrow(A, 1) == 2
    @test BandedMatrices._firstdiagrow(A, 2) == 2
    @test BandedMatrices._firstdiagrow(A, 3) == 2
    @test BandedMatrices._firstdiagrow(A, 4) == -3
    @test BandedMatrices._firstdiagrow(A, 5) == -3
    @test BandedMatrices._firstdiagrow(A, 6) == -3

    @test BandedMatrices._firstdiagcol(A, 1) == -3
    @test BandedMatrices._firstdiagcol(A, 2) == -3
    @test BandedMatrices._firstdiagcol(A, 3) == 2
    @test BandedMatrices._firstdiagcol(A, 4) == 2
    @test BandedMatrices._firstdiagcol(A, 5) == 2
end

# scalar - integer - integer
let
    a = bones(5, 5, 1, 1)
    # 1.0  1.0  0.0  0.0  0.0
    # 1.0  1.0  1.0  0.0  0.0
    # 0.0  1.0  1.0  1.0  0.0
    # 0.0  0.0  1.0  1.0  1.0

    # in band
    a[1, 1] = 1
    @test a[1, 1] == 1

    # out of band
    @test_throws BandError a[1, 3] = 1
    @test_throws BandError a[3, 1] = 1

    # out of range
    @test_throws BoundsError a[0, 0] = 1
    @test_throws BoundsError a[0, 1] = 1
    @test_throws BoundsError a[1, 0] = 1
    @test_throws BoundsError a[5, 6] = 1
    @test_throws BoundsError a[6, 5] = 1
    @test_throws BoundsError a[6, 6] = 1
end

# ~ indexing along a column

# scalar - BandRange/Colon - integer
let
    a = bones(5, 5, 1, 1)
    # 1.0  1.0  0.0  0.0  0.0
    # 1.0  1.0  1.0  0.0  0.0
    # 0.0  1.0  1.0  1.0  0.0
    # 0.0  0.0  1.0  1.0  1.0
    # 0.0  0.0  0.0  1.0  1.0

    # in band
    a[BandRange, 1] = 2
    a[BandRange, 2] = 3
    a[BandRange, 3] = 4
    a[BandRange, 4] = 5
    a[BandRange, 5] = 6
    @test a ==  [2  3  0  0  0;
                 2  3  4  0  0;
                 0  3  4  5  0;
                 0  0  4  5  6;
                 0  0  0  5  6]

    @test_throws BandError a[:, 0] = 1
    @test_throws BandError a[:, 1] = 1
    @test_throws BoundsError a[BandRange, 0] = 1
    @test_throws BoundsError a[BandRange, 6] = 1
end

# vector - BandRange/Colon - integer
let
    a = bones(Int, 5, 7, 2, 1)
    # 5x7 BandedMatrices.BandedMatrix{Float64}:
    #  1.0  1.0    0    0    0    0   0   0
    #  1.0  1.0  1.0    0    0    0   0   0
    #  1.0  1.0  1.0  1.0    0    0   0   0
    #    0  1.0  1.0  1.0  1.0    0   0   0
    #    0    0  1.0  1.0  1.0  1.0   0   0

    # in band
    a[BandRange, 1] = [1,   2,  3]
    a[BandRange, 2] = [4,   5,  6,  7]
    a[BandRange, 3] = [8,   9, 10, 11]
    a[BandRange, 4] = [12, 13, 14]
    a[BandRange, 5] = [15, 16]
    a[BandRange, 6] = [17]

    @test a == [ 1  4   0   0   0   0   0;
                 2  5   8   0   0   0   0;
                 3  6   9  12   0   0   0;
                 0  7  10  13  15   0   0;
                 0  0  11  14  16  17   0]

    @test_throws BandError a[:, 0] = [1, 2, 3]
    @test_throws BandError a[:, 1] = [1, 2, 3]
    @test_throws BoundsError a[BandRange, 0] = [1, 2, 3]
    @test_throws BoundsError a[BandRange, 8] = [1, 2, 3]
    @test_throws DimensionMismatch a[BandRange, 1] = [1, 2]
end

# scalar - range - integer
let
    a = bones(3, 4, 2, 1)
    # 1.0  1.0  0.0  0.0
    # 1.0  1.0  1.0  0.0
    # 1.0  1.0  1.0  1.0

    # full column span
    a[1:3, 1] = 1
    a[1:3, 2] = 2
    a[2:3, 3] = 3
    a[3:3, 4] = 4
    @test a == [ 1  2  0  0;
                 1  2  3  0;
                 1  2  3  4]

    # partial span
    a[1:1, 1] = 3
    a[2:3, 2] = 4
    a[3:3, 3] = 5
    @test a == [ 3  2  0  0;
                 1  4  3  0;
                 1  4  5  4]

    # wrong range input
    @test_throws BoundsError   a[0:1, 1] = 3
    @test_throws BoundsError   a[1:4, 1] = 3
    @test_throws BandError a[2:3, 4] = 3
    @test_throws BoundsError   a[3:4, 4] = 3

end

# vector - range - integer
let
    a = bones(5, 4, 1, 2)
    # 1.0  1.0  1.0  0.0
    # 1.0  1.0  1.0  1.0
    # 0.0  1.0  1.0  1.0
    # 0.0  0.0  1.0  1.0
    # 0.0  0.0  0.0  1.0
    # full column span
    a[1:2, 1] = 1:2
    a[1:3, 2] = 3:5
    a[1:4, 3] = 6:9
    a[2:5, 4] = 10:13
    @test a == [1  3  6   0;
                2  4  7  10;
                0  5  8  11;
                0  0  9  12;
                0  0  0  13]

    # partial span
    a[1:1, 1] = [2]
    a[2:3, 2] = 1:2
    a[3:4, 3] = 3:4
    a[4:5, 4] = 3:4
    @test a == [2  3  6   0;
                2  1  7  10;
                0  2  3  11;
                0  0  4   3;
                0  0  0   4]

    # wrong range input
    @test_throws BoundsError a[1:2, 0] = 1:2
    @test_throws BoundsError a[0:1, 1] = 1:2
    @test_throws BoundsError a[1:8, 1] = 1:8
    @test_throws BoundsError a[2:6, 4] = 2:6
    @test_throws BoundsError a[3:4, 5] = 3:4
end

# ~ indexing along a row

# scalar - integer - BandRange/colon
let
    a = bones(5, 5, 1, 1)
    # 1.0  1.0  0.0  0.0  0.0
    # 1.0  1.0  1.0  0.0  0.0
    # 0.0  1.0  1.0  1.0  0.0
    # 0.0  0.0  1.0  1.0  1.0
    # 0.0  0.0  0.0  1.0  1.0

    # in band
    a[1, BandRange] = 2
    a[2, BandRange] = 3
    a[3, BandRange] = 4
    a[4, BandRange] = 5
    a[5, BandRange] = 6
    @test a == [2  2  0  0  0;
                3  3  3  0  0;
                0  4  4  4  0;
                0  0  5  5  5;
                0  0  0  6  6]

    @test_throws BandError a[0, :] = 0
    @test_throws BandError a[1, :] = 0
    @test_throws BoundsError a[0, BandRange] = 0
    @test_throws BoundsError a[6, BandRange] = 0
end

# vector - integer - BandRange/colon
let
    a = bones(7, 5, 1, 2)
    # 7x5 BandedMatrices.BandedMatrix{Float64}:
    # 1.0  1.0  1.0  0.0  0.0
    # 1.0  1.0  1.0  1.0  0.0
    # 0.0  1.0  1.0  1.0  1.0
    # 0.0  0.0  1.0  1.0  1.0
    # 0.0  0.0  0.0  1.0  1.0
    # 0.0  0.0  0.0  0.0  1.0
    # 0.0  0.0  0.0  0.0  0.0

    # in band
    a[1, BandRange] = [1,   2,  3]
    a[2, BandRange] = [4,   5,  6,  7]
    a[3, BandRange] = [8,   9, 10, 11]
    a[4, BandRange] = [12, 13, 14]
    a[5, BandRange] = [15, 16]
    a[6, BandRange] = [17]

    @test a ==  [1  2   3   0   0;
                 4  5   6   7   0;
                 0  8   9  10  11;
                 0  0  12  13  14;
                 0  0   0  15  16;
                 0  0   0   0  17;
                 0  0   0   0   0]

    @test_throws BandError a[0, :] = [1, 2, 3]
    @test_throws BandError a[1, :] = [1, 2, 3]
    @test_throws BoundsError a[0, BandRange] = [1, 2, 3]
    @test_throws BoundsError a[8, BandRange] = [1, 2, 3]
    @test_throws DimensionMismatch a[1, BandRange] = [1, 2]
    @test_throws BoundsError a[7, BandRange] = [1, 2, 3]
    @test_throws BoundsError a[7, BandRange] = 1
end

# scalar - integer - range
let
    a = bones(7, 5, 1, 2)
    # 7x5 BandedMatrices.BandedMatrix{Float64}:
    # 1.0  1.0  1.0  0.0  0.0
    # 1.0  1.0  1.0  1.0  0.0
    # 0.0  1.0  1.0  1.0  1.0
    # 0.0  0.0  1.0  1.0  1.0
    # 0.0  0.0  0.0  1.0  1.0
    # 0.0  0.0  0.0  0.0  1.0
    # 0.0  0.0  0.0  0.0  0.0

    # in band
    a[1, 1:3] = 2
    a[2, 1:4] = 3
    a[3, 2:5] = 4
    a[4, 3:5] = 5
    a[5, 4:5] = 6
    a[6, 5:5] = 7

    @test a ==  [2  2  2  0  0;
                 3  3  3  3  0;
                 0  4  4  4  4;
                 0  0  5  5  5;
                 0  0  0  6  6;
                 0  0  0  0  7;
                 0  0  0  0  0]

    @test_throws BoundsError a[0, 1:3] = 1
    @test_throws BoundsError a[8, 2:3] = 1
    @test_throws BoundsError a[1, 0:3] = 1
    @test_throws BoundsError a[4, 4:6] = 1

    @test_throws BandError a[1, 1:4] = 1
end

# vector - integer - range
let
    a = bones(7, 5, 1, 2)
    # 7x5 BandedMatrices.BandedMatrix{Float64}:
    # 1.0  1.0  1.0  0.0  0.0
    # 1.0  1.0  1.0  1.0  0.0
    # 0.0  1.0  1.0  1.0  1.0
    # 0.0  0.0  1.0  1.0  1.0
    # 0.0  0.0  0.0  1.0  1.0
    # 0.0  0.0  0.0  0.0  1.0
    # 0.0  0.0  0.0  0.0  0.0

    # in band
    a[1, 1:3] = [1, 2, 3]
    a[2, 1:4] = [1, 2, 3, 4]
    a[3, 2:5] = [1, 2, 3, 4]
    a[4, 3:5] = [1, 2, 3]
    a[5, 4:5] = [1, 2]
    a[6, 5:5] = [1]

    @test a ==  [1  2  3  0  0;
                 1  2  3  4  0;
                 0  1  2  3  4;
                 0  0  1  2  3;
                 0  0  0  1  2;
                 0  0  0  0  1;
                 0  0  0  0  0]

    @test_throws BoundsError a[0, 1:3] = [1, 2, 3]
    @test_throws BoundsError a[8, 2:3] = [1, 2]
    @test_throws BoundsError a[1, 0:3] = [1, 2, 3, 4]
    @test_throws BoundsError a[4, 4:6] = [2, 3]

    @test_throws BandError a[1, 1:4] = [1, 2, 3, 4]
    @test_throws BoundsError a[7, 6:7] = [1, 2]
    @test_throws DimensionMismatch a[1, 1:3] = [1, 2]
end

# indexing along a band
let
    a = bzeros(5, 4, 2, 1)
    # 5x4 BandedMatrices.BandedMatrix{Float64}:
    #  0.0  0.0
    #  0.0  0.0  0.0
    #  0.0  0.0  0.0  0.0
    #       0.0  0.0  0.0
    #            0.0  0.0
    a[band(-2)] = 5
    a[band(-1)] = 1
    a[band( 0)] = 2
    a[band( 1)] = 3

    @test full(a) == [ 2  3  0  0;
                       1  2  3  0;
                       5  1  2  3;
                       0  5  1  2;
                       0  0  5  1]

    @test_throws BandError a[band(-3)] = 1

    a[band(-2)] = [4, 4, 4]
    a[band(-1)] = [1, 2, 3, 4]
    a[band( 0)] = [5, 6, 7, 8]
    a[band( 1)] = [9, 10, 11]

    @test full(a) == [ 5  9   0   0;
                       1  6  10   0;
                       4  2   7  11;
                       0  4   3   8;
                       0  0   4   4]

    @test_throws BandError a[band(-3)] = [1, 2, 3]
end

# other special methods
let
    # all elements
    a = bzeros(3, 3, 1, 1)
    a[:] = 1 # this has special meaning in julia, so it should be allowed
    @test a == [1 1 0;
                1 1 1;
                0 1 1]

    # all rows/cols
    a[:, :] = 2 # this has special meaning in julia, so it should be allowed
    @test a == [2 2 0;
                2 2 2;
                0 2 2]
end

# replace a block in the band
let
    a = bzeros(5, 4, 2, 1)
    # 5x4 BandedMatrices.BandedMatrix{Float64}:
    #  0.0  0.0
    #  0.0  0.0  0.0
    #  0.0  0.0  0.0  0.0
    #       0.0  0.0  0.0
    #            0.0  0.0
    a[1:3, 1:2] = 2
    a[3:5, 3:4] = [1 2;
                   3 4;
                   5 6]
    @test full(a) == [2  2  0 0;
                      2  2  0 0;
                      2  2  1 2;
                      0  0  3 4;
                      0  0  5 6]

    @test_throws BoundsError a[0:3, 1:2] = 2
    @test_throws BoundsError a[3:6, 3:4] = 2
    @test_throws BoundsError a[1:3, 0:2] = 2
    @test_throws BoundsError a[3:5, 3:5] = 2
    @test_throws BoundsError a[0:3, 1:2] = [1 2; 3 4]
    @test_throws BoundsError a[3:6, 3:4] = [1 2; 3 4]
    @test_throws BoundsError a[1:3, 0:2] = [1 2; 3 4]
    @test_throws BoundsError a[3:5, 3:5] = [1 2; 3 4]

    @test_throws BandError a[1:3, 1:3] = 2
    @test_throws BandError a[1:3, 1:3] = rand(3, 3)

    @test_throws DimensionMismatch a[1:3, 1:2] = rand(3, 3)
end


using BenchmarkTools
using BSON
using CUDA
using Flux
using ProgressMeter
using Random
using StatsBase
using ValueHistories

using Base.Iterators: partition
using Flux: gpu
using Flux.Optimise: runall, update!, StopException, batchmemaybe
using Flux.Data: DataLoader
using Zygote: Params, gradient

# -------------------------------------------------------------------------------
# Data processing
# -------------------------------------------------------------------------------
function batch_provider(x, y, batchsize)
    neg = findall(vec(y) .== 0)
    pos = findall(vec(y) .== 1)

    n_neg = batchsize ÷ 2
    n_pos = batchsize - n_neg

    x_obs = ndims(x)
    y_obs = ndims(y)

    last_batch = sample(1:batchsize, batchsize; replace = false)

    function make_batch(; buffer = false)
        inds = vcat(
            sample(neg, n_neg; replace = length(neg) < n_neg),
            sample(pos, n_pos; replace = length(pos) < n_pos),
        )
        shuffle!(inds)
        if !ismissing(buffer)
            if buffer
                addind = AccuracyAtTop.BUFFER[].ind
                if 0 < addind <= batchsize
                    inds[end] = last_batch[addind]
                end
                last_batch .= inds
            end
        end
        return (getdim(x, x_obs, inds), getdim(y, y_obs, inds))
    end
    return make_batch
end

# -------------------------------------------------------------------------------
# Custom train!
# -------------------------------------------------------------------------------
function custom_train!(loss, ps, data, opt; cb = (args...) -> ())
  ps = Params(ps)
  cb = runall(cb)

  local loss_val

  for d in data
    try
      gs = gradient(ps) do
        loss_val = loss(batchmemaybe(d)...)
        return loss_val
      end
      update!(opt, ps, gs)
      cb(loss_val, batchmemaybe(d)...)
    catch ex
      if ex isa StopException
        break
      else
        rethrow(ex)
      end
    end
  end
end

# -------------------------------------------------------------------------------
# Callback function
# -------------------------------------------------------------------------------
Base.@kwdef mutable struct CallBack
    iters::Int
    epochlength::Int = iters
    title::String = "Training:"
    bar::Progress = Progress(iters, 5, title)
    showat::Int = 100
    showfunc::Function = (args...) -> []
    saveat::Int = 1000
    savefunc::Function = (args...) -> nothing
    counter::Int = 0
    usershows = []
    loss = History(Float32)
end

function CallBack(iters, epochlength = iters; kwargs...)
    return CallBack(; iters = iters, epochlength = epochlength, kwargs...)
end

function (c::CallBack)(loss_val, x, y)
    c.counter += 1
    push!(c.loss, c.counter, eltype(c.loss.values)(loss_val))

    if mod(c.counter, c.showat) == 0 || c.counter == 1
        c.usershows = c.showfunc(c)
    end
    if mod(c.counter, c.saveat) == 0
        c.savefunc(c, x, y)
    end
    next!(c.bar; showvalues = vcat(
        epochcounter(c),
        epochtimer(c),
        itercounter(c),
        itertimer(c),
        c.usershows
    ))
    return
end

function itercounter(c::CallBack)
    return ("Iteration", string(c.counter, "/", c.iters))
end

function itertimer(c::CallBack)
    tm = round((c.bar.tlast - c.bar.tinit)/c.bar.counter; sigdigits = 2)
    return ("Average time per iteration", string(tm, "s"))
end

function epochcounter(c::CallBack)
    epoch = floor(Int, c.counter/c.epochlength)
    maxepoch = floor(Int, c.iters/c.epochlength)
    return maxepoch <= 1 ? [] : ("Epoch", string(epoch, "/", maxepoch))
end

function epochtimer(c::CallBack)
    tm = round(c.epochlength*(c.bar.tlast - c.bar.tinit)/c.bar.counter; sigdigits = 2)
    return ("Average time per epoch", string(tm, "s"))
end

# -------------------------------------------------------------------------------
# Saving functions
# -------------------------------------------------------------------------------
function save_simulation(
    c::CallBack,
    dataset_settings::Dict,
    train_settings_in::Dict,
    model_settings::Dict,
    model,
    x,
    y,
)

    savedir = modeldir(dataset_settings, train_settings_in, model_settings)

    train_settings = deepcopy(train_settings_in)
    train_settings[:epochs] = floor(Int, c.counter/c.epochlength)
    train_settings[:epochlength] = c.epochlength
    train_settings[:iters] = c.counter

    tm = (c.bar.tlast - c.bar.tinit)/c.bar.counter

    simulation = Dict(
        :dataset_settings => deepcopy(dataset_settings),
        :train_settings => deepcopy(train_settings),
        :model_settings => deepcopy(model_settings),
        :time_per_iter => tm,
        :time_per_epoch => tm * c.epochlength,
        :model => cpu(model),
        :loss => c.loss.values,
        :minibatch => Dict(
            :targets => cpu(vec(y)),
            :scores => cpu(compute_scores(model, x)),
        ),
    )

    # save
    model_dict = deepcopy(model_settings)
    model_dict[:epochs] = simulation[:train_settings][:epochs]
    model_dict[:iters] = simulation[:train_settings][:iters]

    isdir(savedir) || mkpath(savedir)
    bson(joinpath(savedir, simulation_name(train_settings[:epochs])), simulation)
    return
end

function save_simulation_tfco(
    dataset_settings::Dict,
    train_settings_in::Dict,
    model_settings::Dict,
    results,
    y_train,
    y_test,
)

    savedir = modeldir(dataset_settings, train_settings_in, model_settings)

    train_settings = deepcopy(train_settings_in)
    train_settings[:iters] = results["iters"]
    tm = results["tm"]

    simulation = Dict(
        :dataset_settings => deepcopy(dataset_settings),
        :train_settings => deepcopy(train_settings),
        :model_settings => deepcopy(model_settings),
        :time_per_iter => tm / train_settings[:iters],
        :time_per_epoch => tm / train_settings[:epochs],
        :train => Dict(
            :targets => vec(y_train),
            :scores => vec(results["s_train"]),
            :loss => results["loss_train"],
        ),
        :test => Dict(
            :targets => vec(y_test),
            :scores => vec(results["s_test"]),
            :loss => results["loss_test"],
        ),
    )

    # save
    model_dict = deepcopy(model_settings)
    model_dict[:epochs] = simulation[:train_settings][:epochs]
    model_dict[:iters] = simulation[:train_settings][:iters]

    isdir(savedir) || mkpath(savedir)
    bson(joinpath(savedir, simulation_name(model_dict[:epochs])), simulation)
    return
end

# -------------------------------------------------------------------------------
# Runing simulations
# -------------------------------------------------------------------------------
function run_simulations(Dataset_Settings, Train_Settings, Model_Settings; nepochs_save = 1)
    for dataset_settings in dict_list_simple(Dataset_Settings)
        @unpack dataset, posclass = dataset_settings
        @info "Dataset: $(dataset), positive class label: $(posclass)"

        labelmap(y) = y in posclass
        (x_train, y_train), (x_test, y_test) = load(dataset; labelmap = labelmap)

        for train_settings in dict_list_simple(Train_Settings)
            @unpack batchsize, epochs, seed = train_settings
            @info "Batchsize: $(batchsize), seed: $(seed)"

            epochlength = length(y_train) ÷ batchsize
            iters = epochs * epochlength
            make_batch = batch_provider(x_train, y_train, batchsize)

            for model_settings in dict_list_simple(Model_Settings)
                @unpack type, arg, surrogate, reg, buffer = model_settings
                model_settings[:seed] = seed

                # create model
                model = build_network(dataset; seed = seed) |> gpu
                objective = build_loss(type, arg, surrogate, reg)
                pars = params(model)

                loss(x, y) = objective(x, y, model, pars)

                # create callback
                savefunc(c, x, y) = save_simulation(
                    c,
                    dataset_settings,
                    train_settings,
                    model_settings,
                    model,
                    x,
                    y,
                )

                cb = CallBack(
                    title = string(string(type), ": "),
                    iters = iters,
                    epochlength = epochlength;
                    saveat = nepochs_save * epochlength,
                    savefunc = savefunc,
                )

                # training
                @info "Bacth preparation:"
                Random.seed!(seed)
                batches = (gpu(make_batch(; buffer = buffer)) for iter in 1:iters)

                @unpack optimiser, steplength = train_settings
                opt = optimiser(steplength)

                custom_train!(loss, pars, batches, opt; cb = cb)
            end
        end
    end
end

function run_benchmark(Dataset_Settings, Train_Settings, Model_Settings)
    for dataset_settings in dict_list_simple(Dataset_Settings)
        @unpack dataset, posclass = dataset_settings
        @info "Dataset: $(dataset), positive class label: $(posclass)"

        labelmap = (y) -> y == posclass
        (x_train, y_train), ~ = load(dataset; labelmap = labelmap)

        for train_settings in dict_list_simple(Train_Settings)
            @unpack batchsize, iters = train_settings
            @info "Batchsize: $(batchsize), iters: $(iters)"

            epochlength = length(y_train) ÷ batchsize
            make_batch = batch_provider(x_train, y_train, batchsize)

            for model_settings in dict_list_simple(Model_Settings)
                @unpack type, arg, surrogate, reg, buffer = model_settings

                # create model
                model = build_network(dataset; seed = 1) |> gpu
                objective = build_loss(type, arg, surrogate, reg)
                pars = params(model)

                loss(x, y) = objective(x, y, model, pars)

                # training
                @info "Model: $(type)"
                batches = [gpu(make_batch(; buffer = buffer)) for iter in 1:iters] |> gpu

                @unpack optimiser, steplength = train_settings
                opt = optimiser(steplength)

                # precompile
                Flux.train!(loss, pars, (make_batch(; buffer = buffer) for iter in 1:5), opt)

                # run benchmark
                b = @benchmark Flux.train!($loss, $pars, $batches, $opt)
                d = Dict(:times => b.times .* 1e-9, :epochlength => epochlength, :iters_in_run => iters)

                savedir = datadir(
                    "benchmarks",
                    dataset_savename(dataset_settings),
                    train_savename(train_settings),
                )
                file_name = string(model_savename(model_settings), ".bson")

                mkpath(savedir)
                BSON.bson(joinpath(savedir, file_name), d)
            end
        end
    end
end


module SeisSolXDMF

export XDMFFile, timesteps_of, grid_of, data_of

using XMLDict
using HDF5

    struct XDMFFile
        xml         :: XMLDict.XMLDictElement
        base_path   :: AbstractString
        timesteps   :: Array{XMLDict.XMLDictElement,1}
    end

    function expect(expression, error_msg)
        if !expression
            error(error_msg)
        end
    end

    function XDMFFile(file_path::AbstractString)
        xml_file = open(file_path) do file; read(file, String) end
        xml_file = parse_xml(xml_file)

        try
            # The root node "Xdmf" is assumed
            timesteps = xml_file["Domain"]["Grid"]["Grid"]
            num_timesteps = length(timesteps)

            expect(num_timesteps >  0, "No timesteps found in XDMF file.")

            xdmf_base_path = splitdir(file_path)[1]
            return XDMFFile(xml_file, xdmf_base_path, timesteps)
        catch
            error("Input file has no valid SeisSol XDMF structure.")
        end
    end

    function timesteps_of(xdmf::XDMFFile) :: Array{Float64, 1}
        return Array(map(data_item -> parse(Float64, data_item["Time"][:Value]), xdmf.timesteps))
    end

    function grid_of(xdmf::XDMFFile) :: Tuple{AbstractArray{Integer, 2}, AbstractArray{AbstractFloat, 2}}
        # Read geometry (points) and topology (simplices (triangles/tetrahedra)) files.
        # Since the ids of the points referred to in the topology array start at 0 (and Julia counts from 1) we have to add 1.
        simplices = read_dataset(xdmf, xdmf.timesteps[1]["Topology"]["DataItem"]) .+ 1
        expect(ndims(simplices) == 2, "Topology should have 2 dimensions but has $(ndims(simplices)).")
        expect(size(simplices, 1) ∈ (3, 4), "Topology entries must have 3 (triangles) or 4 (tetrahedra) points but have $(size(simplices, 1)).")

        points = read_dataset(xdmf, xdmf.timesteps[1]["Geometry"]["DataItem"])
        expect(ndims(points) == 2, "Geometry should have 2 dimensions but has $(ndims(simplices)).")
        expect(size(points, 1) == 3, "Geometry entries must have 3 coordinates but have $(size(simplices, 1)).")

        return (simplices, points)
    end

    function data_of(xdmf::XDMFFile, timestep::Integer, var_name::AbstractString)
        expect(timestep >= 1, "Timestep $(timestep) is smaller than 1.")
        expect(timestep <= length(xdmf.timesteps), "Timestep $(timestep) exceeds number of timesteps in XDMF file ($(length(xdmf.timesteps))).")

        attr_or_list = xdmf.timesteps[timestep]["Attribute"]
        var_entry = if isa(attr_or_list, AbstractArray)
            var_index = findfirst(attr -> attr[:Name] == var_name, attr_or_list)
            expect(!isnothing(var_index), """Variable $var_name not found in timestep $timestep.
                Found variables: $(join(map(attr -> attr[:Name], attr_or_list), ", ")).""")

            attr_or_list[var_index]["DataItem"]
        else
            attr = attr_or_list
            expect(attr[:Name] == var_name, """Variable $var_name not found in timestep $timestep. Found variable: $(attr[:Name]).""")

            attr["DataItem"]
        end
        return reshape(read_hyperslab(xdmf, var_entry), :)
    end

    function read_dataset(xdmf::XDMFFile, data_item::XMLDict.XMLDictElement; indices=nothing)
        filename = data_item[""]

        file_range = data_item[:Dimensions]
        file_range = split(file_range, ' ')
        file_range = parse.(Int, file_range)
        file_period = file_range[2]

        number_type = get_number_type(data_item)

        if data_item[:Format] == "HDF"
            path_parts = split(filename, ':', limit=2)
            expect(!isempty(path_parts[1]) && !isempty(path_parts[2]), "HDF5 group path is invalid.")

            filename = joinpath(xdmf.base_path, path_parts[1])
            hdf_dataset_path = String(path_parts[2])

            dataset =
                if isnothing(indices)
                    h5read(filename, hdf_dataset_path, (1:file_range[2], 1:file_range[1]))
                else
                    h5read(filename, hdf_dataset_path, (
                        indices[1,2]:indices[2,2]:(indices[1,2] + indices[3,2] - 1),
                        indices[1,1]:indices[2,1]:(indices[1,1] + indices[3,1] - 1)))
                end
        else
            filename = joinpath(xdmf.base_path, filename)

            if !isnothing(indices)
                start_index = (indices[1,1]-1) * file_period + (indices[1,2]-1) # Starting at 0
                stride = (indices[2,1]-1) * file_period + (indices[2,2]-1)
                expect(stride == 0, "Stride not yet supported for POSIX output.") # (no need (?))
                read_size = (indices[3,2], indices[3,1])
            else
                start_index = 0
                read_size = (file_range[2], file_range[1])
            end
            dataset = Array{number_type, 2}(undef, read_size)
            open(filename, "r") do file
                seek(file, start_index * sizeof(number_type))
                read!(file, dataset)
            end
        end

        return dataset
    end

    function read_hyperslab(xdmf::XDMFFile, data_item::XMLDict.XMLDictElement)
        #=
        Example DataItem:
        <DataItem ItemType="HyperSlab" Dimensions="4561037">
            <DataItem NumberType="UInt" Precision="4" Format="XML" Dimensions="3 2">0 0 1 1 1 4561037</DataItem>
            <DataItem NumberType="Float" Precision="8" Format="Binary" Dimensions="1 5242880">out_cell/mesh0/u.bin</DataItem>
        </DataItem>
        =#
        expect(data_item[:ItemType] == "HyperSlab", "Unexpected item type '$(data_item[:ItemType])', expected 'HyperSlab'.")
        range_item = data_item["DataItem"][1]
        binary_item = data_item["DataItem"][2]

        expect(range_item[:Format] == "XML", "Unexpected format '$(range_item[:Format])', expected 'XML'.")
        expect(range_item[:Dimensions] == "3 2", "Unexpected dimensions '$(range_item[:Dimensions])', expected '3 2'.")
        expect(binary_item[:Format] ∈ ["Binary", "HDF"], "Unexpected format '$(binary_item[:Format])', expected 'Binary' or 'HDF'.")

        rng = range_item[""]
        rng = split(rng, ' ')
        rng = parse.(Int, rng)
        rng = reshape(rng, (2, 3))'

        rng[1,:] .+= 1 # Julia start indices

        return read_dataset(xdmf, binary_item, indices=rng)
    end

    function get_number_type(data_item::XMLDict.XMLDictElement) :: Type
        bytes_of_precision = parse(UInt, data_item[:Precision])
        number_type_name = data_item[:NumberType]

        print_error(type_name, precision) =
            error("""Type "$type_name" of size $precision bytes is not recognized by the parser.""")

        if bytes_of_precision == 8
            if number_type_name == "Int"
                return Int64
            elseif number_type_name == "UInt"
                return UInt64
            elseif number_type_name == "Float"
                return Float64
            else
                print_error(number_type_name, bytes_of_precision)
            end
        elseif bytes_of_precision == 4
            if number_type_name == "Int"
                return Int32
            elseif number_type_name == "UInt"
                return UInt32
            elseif number_type_name == "Float"
                return Float32
            else
                print_error(number_type_name, bytes_of_precision)
            end
        else
            print_error(number_type_name, bytes_of_precision)
        end
    end

end # module


using MuZero
import Zygote

include("game.jl")

learner = Learner(Connect4)

# Play a quick game to make sure
# the basic game dynamics are working
#learner |> play

learner |> tests |> svg
learner |> test_results |> svg

selfplay!(learner; limit=10)

learner |> randgame |> svg

# still a fence post problem here
s = learner |> randtrainingsample
learner |> randtrainingsample |> svg

learner |> randtrainingbatch |> show

using Base.Threads: @spawn
for i = 1:num_actors(learner)
    @spawn selfplay!(learner)
end

length(learner.replay_buffer)

@time fit!(learner, steps=1)


# Forward BLAS calls to generic implementation
#
using LinearAlgebra
import LinearAlgebra.BLAS

# generic implementations

@reroute LinearAlgebra.BLAS.dot(x,y) LinearAlgebra.dot(Any, Any)
@reroute LinearAlgebra.BLAS.axpy!(a, x, y) LinearAlgebra.axpy!(Any,
                                                      AbstractArray,
                                                      AbstractArray)

gengemv!(tA, α, A, x, β, y) = LinearAlgebra.generic_matvecmul!(y, tA, A, x, LinearAlgebra.MulAddMul(α, β))

@reroute LinearAlgebra.BLAS.gemv!(tA, α, A, x, β, y) gengemv!(Any, Any, Any, Any, Any, Any)


#=
    NLP
    many unsafe operations here, because we are using JuMP non exposed
    methods (the ones starting with "_")
    - TODO: PIN JUMP
=#

function _has_nlp_data(model)
    return  model.nlp_data !== nothing
end
function _load_nlp_data(model)
    # callen in JuMP.optimize!
    # The nlp_data is not kept in sync, so re-set it here.
    # TODO: Consider how to handle incremental solves.
    if model.nlp_data !== nothing
        MOI.set(model, MOI.NLPBlock(), JuMP._create_nlp_block_data(model))
        empty!(model.nlp_data.nlconstr_duals)
    end
end

function Base.getproperty(m::UpperModel, s::Symbol)
    if s === :nlp_data
        return mylevel_model(m).nlp_data
    else
        getfield(m, s)
    end
end
function Base.getproperty(m::LowerModel, s::Symbol)
    if s === :nlp_data
        return no_nlp_lower()
    else
        getfield(m, s)
    end
end

no_nlp() = error("Non-linear data must be passed to the Upper(.) model")
no_nlp_lower() = error("NLconstraint(s) are not allowed in the lower level")
no_nlp_lower_param() = error("NLparameter(s) are not allowed in the lower level")

JuMP._init_NLP(m::UpperModel) = JuMP._init_NLP(mylevel_model(m))
JuMP._init_NLP(m::LowerModel) = no_nlp_lower()

function JuMP._new_parameter(m::UpperModel, value::Number)
    JuMP._init_NLP(m)
    upper = m.m.upper
    nldata::JuMP._NLPData = upper.nlp_data
    push!(nldata.nlparamvalues, value)
    return JuMP.NonlinearParameter(upper, length(nldata.nlparamvalues))
end

function JuMP._new_parameter(::LowerModel, ::Number)
    no_nlp_lower_param()
end
function JuMP._new_parameter(::BilevelModel, ::Number)
    no_nlp()
end

# function JuMP.set_objective_function(m::UpperModel, func::JuMP._NonlinearExprData)
#     JuMP.set_objective_function(mylevel_model(m), func)
#     return nothing
# end
# function JuMP.set_objective_function(::LowerModel, ::JuMP._NonlinearExprData)
#     no_nlp_lower()
#     return nothing
# end
function JuMP.set_objective(
    m::UpperModel,
    sense::MOI.OptimizationSense,
    ex::JuMP._NonlinearExprData,
)
    JuMP._init_NLP(m)
    JuMP.set_objective_sense(m, sense)
    m.nlp_data.nlobj = ex
    return
end
function JuMP.set_objective(::LowerModel,
    ::MOI.OptimizationSense,
    ::JuMP._NonlinearExprData,
)
    no_nlp_lower()
    return
end

function JuMP._parse_NL_expr_runtime(m::UpperModel, x, tape, parent, values)
    JuMP._parse_NL_expr_runtime(mylevel_model(m), x, tape, parent, values)
    return nothing
end

function JuMP._parse_NL_expr_runtime(m::UpperModel, x::BilevelVariableRef, tape, parent, values)
    if !in_upper(x)
        error(
            "Variable in nonlinear expression does not belong to the " *
            "upper model, it is a LowerOnly variable",
        )
    end
    JuMP._parse_NL_expr_runtime(mylevel_model(m), upper_ref(x), tape, parent, values)
    return nothing
end

# function JuMP._parse_NL_expr_runtime(
#     m::Model,
#     x::NonlinearExpression,
#     tape,
#     parent,
#     values,
# )
#     push!(tape, NodeData(SUBEXPRESSION, x.index, parent))
#     return nothing
# end

# TODO: deal with param
# function JuMP._parse_NL_expr_runtime(
#     m::JuMP.Model,
#     x::JuMP.NonlinearParameter,
#     tape,
#     parent,
#     values,
# )
#     push!(tape, JuMP.NodeData(JuMP.PARAMETER, x.index, parent))
#     return nothing
# end

# function JuMP._parse_NL_expr_runtime(
#     m::Model,
#     x::AbstractArray,
#     tape,
#     parent,
#     values,
# )
#     return error(
#         "Unexpected array $x in nonlinear expression. Nonlinear expressions may contain only scalar expressions.",
#     )
# end

# function JuMP._parse_NL_expr_runtime(
#     m::Model,
#     x::GenericQuadExpr,
#     tape,
#     parent,
#     values,
# )
#     push!(tape, NodeData(CALL, operator_to_id[:+], parent))
#     sum_parent = length(tape)
#     _parse_NL_expr_runtime(m, x.aff, tape, sum_parent, values)
#     for (xy, c) in x.terms
#         push!(tape, NodeData(CALL, operator_to_id[:*], sum_parent))
#         mult_parent = length(tape)
#         _parse_NL_expr_runtime(m, xy.a, tape, mult_parent, values)
#         _parse_NL_expr_runtime(m, xy.b, tape, mult_parent, values)
#         if !isone(c)  # Optimization: no need for * node.
#             _parse_NL_expr_runtime(m, c, tape, mult_parent, values)
#         end
#     end
#     return
# end

# function JuMP._parse_NL_expr_runtime(
#     m::Model,
#     x::GenericAffExpr,
#     tape,
#     parent,
#     values,
# )
#     push!(tape, NodeData(CALL, operator_to_id[:+], parent))
#     sum_parent = length(tape)
#     if !iszero(x.constant)
#         _parse_NL_expr_runtime(m, x.constant, tape, sum_parent, values)
#     end
#     for (v, c) in x.terms
#         if isone(c)  # Optimization: no need for * node.
#             _parse_NL_expr_runtime(m, v, tape, sum_parent, values)
#         else
#             push!(tape, NodeData(CALL, operator_to_id[:*], sum_parent))
#             mult_parent = length(tape)
#             _parse_NL_expr_runtime(m, c, tape, mult_parent, values)
#             _parse_NL_expr_runtime(m, v, tape, mult_parent, values)
#         end
#     end
#     return
# end

# function JuMP._parse_NL_expr_runtime(m::Model, x, tape, parent, values)
#     return error(
#         "Unexpected object $x (of type $(typeof(x)) in nonlinear expression.",
#     )
# end

function JuMP._parse_NL_expr_runtime(m::LowerModel, x, tape, parent, values)
    return no_nlp_lower()
end

#=

function JuMP._Derivatives.expr_to_nodedata(
    ex::VariableRef,
    nd::Vector{NodeData},
    values::Vector{Float64},
    parentid,
    r::_Derivatives.UserOperatorRegistry,
)
    push!(nd, NodeData(MOIVARIABLE, ex.index.value, parentid))
    return nothing
end

function JuMP._Derivatives.expr_to_nodedata(
    ex::NonlinearExpression,
    nd::Vector{NodeData},
    values::Vector{Float64},
    parentid,
    r::_Derivatives.UserOperatorRegistry,
)
    push!(nd, NodeData(SUBEXPRESSION, ex.index, parentid))
    return nothing
end

function JuMP._Derivatives.expr_to_nodedata(
    ex::NonlinearParameter,
    nd::Vector{NodeData},
    values::Vector{Float64},
    parentid,
    r::_Derivatives.UserOperatorRegistry,
)
    push!(nd, NodeData(PARAMETER, ex.index, parentid))
    return nothing
end

=#

#=

# Construct a _NonlinearExprData from a Julia expression.
# VariableRef objects should be spliced into the expression.
function JuMP._NonlinearExprData(m::Model, ex::Expr)
    _init_NLP(m)
    _check_expr(m, ex)
    nd, values = _Derivatives.expr_to_nodedata(ex, m.nlp_data.user_operators)
    return _NonlinearExprData(nd, values)
end
_NonlinearExprData(m::Model, ex) = _NonlinearExprData(m, :($ex + 0))

# Error if:
# 1) Unexpected expression
# 2) VariableRef doesn't match the model
function JuMP._check_expr(m::Model, ex::Expr)
    if ex.head == :ref # if we have x[1] already in there, something is wrong
        error(
            "Unrecognized expression $ex. JuMP variable objects and input coefficients should be spliced directly into expressions.",
        )
    end
    for e in ex.args
        _check_expr(m, e)
    end
    return
end
function JuMP._check_expr(m::Model, v::VariableRef)
    owner_model(v) === m || error("Variable $v does not belong to this model.")
    return
end
_check_expr(m::Model, ex) = nothing

=#

# TODO - convert into bilevel constraint ref
# ConstraintRef(
#     $esc_m,
#     NonlinearConstraintIndex(length($esc_m.nlp_data.nlconstr)),
#     ScalarShape(),
# )