text:
# This implements the Sparse Reconstruction using a Difference Map as described in the paper:
#
#  Will Landecker, Rick Chartrand, and Simon DeDeo. "Robust Compressed
#  Sensing and Sparse Coding with the Difference Map." Submitted to the
#  arXiv on 10/30/2013.
#
# (x, scores) = difference_map_sparse_reconstruction(d, y, s)
#
# Finds a sparse reconstruction 'x' using the Difference Map method.
# 'x' is the solution as in y = d * x.
#
# Outputs:
#  'x' is an array with the sparse reconstruction of y given d using the
#      Difference Map method.
#
#  'scores' is an array indicating the accuracy of 'x' at each iteration.
#
# Inputs:
#  'd' is an N by M matrix.
#
#  'y' is an N by 1 array.
#
#  'x0' is an N by 1 array with the starting point for the iteration. Set to
#       nothing to select a random starting point.
#
#  'max_degree' is the maximum number of non-zero values (the max L0-penalty) in the solution x.
#
#  'time_limit' is max number of time (in seconds) the algorithm is allowed to run.
#
#  'delta_treshold' is the treshold value below which the iteration is deemed to have converged.
function difference_map_sparse_reconstruction(d, y, max_degree;
	beta = -0.14,
	x0 = nothing,
	time_limit = 10.0,
	d_pseudo = nothing,
	delta_treshold = 1e-10,
	verbose = true
	)

	# Check and set up parameters
	M, N = size(d)

	if max_degree > M
		error("Max degree ($(max_degree)) is larger than number of available predictors/covariates ($(M))")
	end

	if x0 == nothing
		x0 = (randn(M) .- 0.5)'
	end

	if d_pseudo == nothing
		d_pseudo = pinv(d)
	end

	# Define the functions used for the update
	num_zero = M - max_degree

	# Init temp variables
	xn = x0
	delta = Inf
	iteration = 0
	start_time = time()

	# We use the stdev in the last max_var_steps iterations as an indication
	# of stability:
	# xs = zeros(Float64, M, max_var_steps)

	# Main loop
	while( delta > delta_treshold )

		iteration += 1
		xo = xn

		# Update x
		xn = xo + beta * ( A( fB(xo, beta, num_zero), d, d_pseudo, y )  -
			B( fA(xo, beta, d, d_pseudo, y), num_zero ) )

		# Change in x in this step
		delta = sum((xo - xn) .^ 2)

		# How well do we reconstruct?
		reconstruction_distance = residual_sum_of_squares(y, B(xn, num_zero) * d)

		# Print some info
		if verbose
      println(@sprintf("DM. iter %s\n  reconstruction distance = %0.5e\n  L2 norm of update distance = %0.5e", iteration, reconstruction_distance, delta))
      show(xn)
      println()
    end

    # Break if over time limit
    delta_t = time() - start_time
    if delta_t > time_limit
    	break
    end

	end

	if verbose
		println("Finished after $(iteration) iterations.")
	end

	# Ensure the solution is sparse.
	x = B(xn, num_zero)

	return( x, iteration, aic(d, y, x) )
end

A(x, d, d_pseudo, y) = x .- ( (x*d .- y) * d_pseudo )

function fA(x, beta, d, d_pseudo, y)
	Ax = A(x, d, d_pseudo, y)
	Ax - (Ax - x)/beta
end

function B(x, num_zero)
	# Set the num_zero smallest values to 0
	res = copy(x)
	res[sortperm(x[:])[1:num_zero]] = 0.0
	res
end

function fB(x, beta, num_zero)
	Bx = B(x, num_zero)
	Bx + (Bx - x)/beta
end

residual_sum_of_squares(y, yhat) = sum( (yhat .- y) .^ 2 )

# AIC assuming independent errors and constant variance:
#   AIC = 2k + n [Ln( 2(pi) RSS/n ) + 1]
function aic_from_rss(rss, d, coefs)
	k = sum(abs(coefs) .> 0)
	m, n = size(d)
	return( 2*k + n * (log(2*pi*rss/n) + 1) )
end

aic(d, y, coefs) = aic_from_rss( residual_sum_of_squares(y, coefs * d), d, coefs)

# Return the reconstruction that gives the lowest AIC assuming independent errors and
# constant variance.
function min_aic_difference_map_reconstruction(d, y;
	time_limit = 2.0,
	max_degree = nothing,
	min_degree = 1
	)

	M, N = size(d)

	if max_degree == nothing
		max_degree = int( 2*sqrt(M) )
	end

	# Pre-calc so we need not do it each time.
	dps = pinv(d)

	degrees = min_degree:max_degree
	num_degrees = Base.length(degrees)
	xs = zeros(Float64, M, num_degrees)
	its = zeros(Int64, num_degrees)
	aics = zeros(Float64, num_degrees)

	time_per_run = time_limit / num_degrees

	for i in 1:num_degrees
		deg = max_degree + 1 - i
		x, its[i], aics[i] = difference_map_sparse_reconstruction(d, y, deg;
			time_limit = time_per_run, d_pseudo = dps, verbose = false)
		xs[:,i] = x'
		println("Min AIC DM, degree = $(deg), aic = $(aics[i]), its = $(its[i])")
	end

	p = sortperm(aics)
	bestindex = p[1]
	bestmodel = xs[:,bestindex]
	degs = reverse(degrees)[p]

	return( bestmodel, degs[1], xs[:,p], its[p], aics[p], degs )

end

# Zero center and unit variance of each row of a matrix.
function zero_center_and_normalize(d)
	d0 = d .- mean(d, 2)
	d0 ./ std(d, 2)
end

# Use stability selection (in the spirit of Meinshausen & Buhlman) with the
# DM sparse reconstruction to find the covariates that are most consistently
# in the selected set.
#
#  inclusion_probability = stability_selection_dm(d, y)
#
#
# Outputs:
#  'inclusion_probability' is the probability that a covariate/feature is
#                          selected when half the samples are used and the features
#                          reweighted by a random weight uniformly sampled in [alpha, 1].
#
function stability_selection_dm(d, y;
	num_bootstrap = 100,
	alpha = 0.20,
	min_degree = nothing,
	time_limit = 1.0,
	verbose = true
	)

	p, n = size(d)

	if min_degree == nothing
		min_degree = int(p/2)
	end

	#dn = zero_center_and_normalize(d)

	halfsize = int(n/2)
	counts = zeros(Float64, 1, p)

	time_per_dm = time_limit / (2*num_bootstrap)

	for i in 1:num_bootstrap

		# Randomly reweight the covariate/predictor values.
		dr = dn .* (alpha .+ (1.0 - alpha) .* rand(p, n))

		# Randomly split the sample in two sets.
		permutation = shuffle(collect(1:n))
		set1 = permutation[1:halfsize]
		set2 = permutation[(halfsize+1):n]

		# Run the difference map sparse reconstructions on each set.
		degree = rand(min_degree:(p-1))
		dm1, its1 = difference_map_sparse_reconstruction(dr[:,set1], y[:,set1], degree; time_limit = time_per_dm, verbose = false)
		dm2, its2 = difference_map_sparse_reconstruction(dr[:,set2], y[:,set2], degree; time_limit = time_per_dm, verbose = false)

		# Update the counts for the selected features
		counts += abs(sign(dm1))
		counts += abs(sign(dm2))

		if verbose
			println("$(i):")
			#println("  dm1 = $(dm1)")
			#println("  dm2 = $(dm2)")
			println("  counts = $(counts/(2*i))")
		end
	end

	# Normalize frequency to a probability
	probs = counts / (2*num_bootstrap)

	# the final stability score is the maximum frequency over the steps
	# result <- apply(freq,2,max)
	return probs
end


text:
n = 9 # rand(1:10)
@test matrixdepot("clement", Float64, n) == matrixdepot("clement", n)

A = matrixdepot("clement", n)
B = matrixdepot("clement", n, 1)

@test diag(A+A', 1) == n*ones(n-1)
@test issymmetric(Array(B))
θ = matrixdepot("clement", 1)
println("'clement' passed test...")

ext: jl


text:
import ElectricFields: @namespace!

@testset "Namespace macro" begin
    params = Dict{Symbol,Rational}()

    @namespace!(params) do
        x = 5
        y = 3x
        z = x//3
        if x
            w = x^2
        end
    end

    @test params[:x] == 5
    @test params[:y] == 15
    @test params[:z] == 5//3
    @test params[:w] == 25

    params2 = Dict{Symbol,Quantity}(:λ => 1u"m")

    @namespace!(params2) do
        if λ
            ν = 1/λ
        end
    end

    @test params2[:ν] == 1u"m"^-1

    params3 = Dict{Symbol,Rational}()

    factor = 4
    @namespace!(params3) do
        x = 5
        y = factor*x
    end

    @test params3[:y] == 20

    for i in 1:8
        params4 = Dict{Symbol,Rational}()
        @namespace!(params4) do
            x = 5
            y = i*x
        end

        @test params4[:y] == 5i
    end
end


text:
using DSGE, JLD
using Base.Test

path = dirname(@__FILE__)

# Set up arguments
m = AnSchorfheide(testing = true)

system = jldopen("$path/../reference/forecast_args.jld","r") do file
    read(file, "system")
end

# Run impulse responses
states, obs, pseudo = impulse_responses(m, system)

# Compare to expected output
exp_states, exp_obs, exp_pseudo =
    jldopen("$path/../reference/impulse_responses_out.jld", "r") do file
        read(file, "exp_states"), read(file, "exp_obs"), read(file, "exp_pseudo")
    end

@testset "Compare irfs to expected output" begin
    @test @test_matrix_approx_eq exp_states states
    @test @test_matrix_approx_eq exp_obs    obs
    @test @test_matrix_approx_eq exp_pseudo pseudo
end

# Test impulse response methods for:
# (1) Choosing a set of shocks and their shock values
# (2) Choosing a shock, a state, and the value the state should hit in period 1 based on
# an impulse from that shock at time 1.
# (3) Same as (2) but with an observable.

# Setup
endo   = m.endogenous_states
obsvs  = m.observables
exo    = m.exogenous_shocks
horizon = 10

# Testing both for exact output match and for certain values to match to be
# explicit about the expected outcome of the impulse response

# Shock
shock_names  = [:z_sh, :g_sh]
shock_values = [1., 1.]
states, obs, pseudo = impulse_responses(m, system, horizon, shock_names, shock_values)

# Explicit checks
@testset "Compare shock-set irfs to the states they affect 1-for-1" begin
    @test states[endo[:z_t], 1, exo[:z_sh]] ≈ 1.
    @test states[endo[:y_t], 1, exo[:g_sh]] ≈ 1.
    @test states[endo[:g_t], 1, exo[:g_sh]] ≈ 1.
end

exp_states_shockset, exp_obs_shockset, exp_pseudo_shockset =
    jldopen("$path/../reference/impulse_responses_out.jld", "r") do file
        read(file, "exp_states_shockset"), read(file, "exp_obs_shockset"), read(file, "exp_pseudo_shockset")
    end

@testset "Checking shockset irfs" begin
    @test states ≈ exp_states_shockset
    @test obs    ≈ exp_obs_shockset
    @test pseudo ≈ exp_pseudo_shockset
end

# State
shock_name = :z_sh
state_name   = :y_t
state_value  = 1.
states, obs, pseudo = impulse_responses(m, system, horizon, shock_name, state_name, state_value)

# Explicit checks
@testset "Check state irf initial period is exactly the specified state value" begin
    @test states[endo[:y_t], 1, exo[:z_sh]] ≈ 1.
end

exp_states_shockstates, exp_obs_shockstates, exp_pseudo_shockstates =
    jldopen("$path/../reference/impulse_responses_out.jld", "r") do file
        read(file, "exp_states_shockstate"), read(file, "exp_obs_shockstate"), read(file, "exp_pseudo_shockstate")
    end

@testset "Checking shockstate irfs" begin
    @test states ≈ exp_states_shockstates
    @test obs    ≈ exp_obs_shockstates
    @test pseudo ≈ exp_pseudo_shockstates
end

# Observable
shock_name = :z_sh
obs_name   = :obs_gdp
obs_value  = 1.
states, obs, pseudo = impulse_responses(m, system, horizon, shock_name, obs_name, obs_value)

# Explicit checks
@testset "Check obs irf initial period is exactly the specified obs value" begin
    @test obs[obsvs[:obs_gdp], 1, exo[:z_sh]] ≈ 1.
end

exp_states_shockobs, exp_obs_shockobs, exp_pseudo_shockobs =
    jldopen("$path/../reference/impulse_responses_out.jld", "r") do file
        read(file, "exp_states_shockobs"), read(file, "exp_obs_shockobs"), read(file, "exp_pseudo_shockobs")
    end

@testset "Checking shockobs irfs" begin
    @test states ≈ exp_states_shockobs
    @test obs    ≈ exp_obs_shockobs
    @test pseudo ≈ exp_pseudo_shockobs
end

nothing


text:
# Main DG data structure that contains all relevant data for the DG solver
mutable struct Dg1D{Eqn<:AbstractEquation, MeshType, NVARS, POLYDEG,
                  SurfaceFlux, VolumeFlux, InitialConditions, SourceTerms, BoundaryConditions,
                  VolumeIntegralType, ShockIndicatorVariable,
                  VectorNnodes, MatrixNnodes, MatrixNnodes2,
                  InverseVandermondeLegendre, MortarMatrix,
                  VectorAnalysisNnodes, AnalysisVandermonde} <: AbstractDg{1, POLYDEG, MeshType}
  equations::Eqn

  surface_flux_function::SurfaceFlux
  volume_flux_function::VolumeFlux

  initial_conditions::InitialConditions
  source_terms::SourceTerms

  elements::ElementContainer1D{NVARS, POLYDEG}
  n_elements::Int

  interfaces::InterfaceContainer1D{NVARS, POLYDEG}
  n_interfaces::Int

  boundaries::BoundaryContainer1D{NVARS, POLYDEG}
  n_boundaries::Int
  n_boundaries_per_direction::SVector{2, Int}

  n_l2mortars::Int # TODO: Taal. Only needed for simulation summary output -> fix me when Taal is alive

  boundary_conditions::BoundaryConditions

  nodes::VectorNnodes
  weights::VectorNnodes
  inverse_weights::VectorNnodes
  inverse_vandermonde_legendre::InverseVandermondeLegendre
  lhat::MatrixNnodes2

  volume_integral_type::VolumeIntegralType
  dhat::MatrixNnodes
  dsplit::MatrixNnodes
  dsplit_transposed::MatrixNnodes

  amr_refine_right::MortarMatrix
  amr_refine_left::MortarMatrix
  amr_coarsen_right::MortarMatrix
  amr_coarsen_left::MortarMatrix

  analysis_nodes::VectorAnalysisNnodes
  analysis_weights::VectorAnalysisNnodes
  analysis_weights_volume::VectorAnalysisNnodes
  analysis_vandermonde::AnalysisVandermonde
  analysis_total_volume::Float64
  analysis_quantities::Vector{Symbol}
  save_analysis::Bool
  analysis_filename::String

  shock_indicator_variable::ShockIndicatorVariable
  shock_alpha_max::Float64
  shock_alpha_min::Float64
  shock_alpha_smooth::Bool
  amr_indicator::Symbol
  amr_alpha_max::Float64
  amr_alpha_min::Float64
  amr_alpha_smooth::Bool

  n_elements_global::Int

  element_variables::Dict{Symbol, Union{Vector{Float64}, Vector{Int}}}
  cache::Dict{Symbol, Any}
  thread_cache::Any # to make fully-typed output more readable
  initial_state_integrals::Vector{Float64}
end

# Convenience constructor to create DG solver instance
function Dg1D(equation::AbstractEquation{NDIMS, NVARS}, surface_flux_function, volume_flux_function, initial_conditions, source_terms, mesh::TreeMesh1D, POLYDEG) where {NDIMS, NVARS}
  # Get cells for which an element needs to be created (i.e., all leaf cells)
  leaf_cell_ids = leaf_cells(mesh.tree)

  # Initialize element container
  elements = init_elements(leaf_cell_ids, mesh, Val(NVARS), Val(POLYDEG))
  n_elements = nelements(elements)

  # Initialize interface container
  interfaces = init_interfaces(leaf_cell_ids, mesh, Val(NVARS), Val(POLYDEG), elements)
  n_interfaces = ninterfaces(interfaces)

  # Initialize boundaries
  boundaries, n_boundaries_per_direction = init_boundaries(leaf_cell_ids, mesh, Val(NVARS), Val(POLYDEG), elements)
  n_boundaries = nboundaries(boundaries)

  n_l2mortars = -1 # TODO: Taal. Only needed for simulation summary output -> fix me when Taal is alive

  # Sanity checks
  if isperiodic(mesh.tree)
    @assert n_interfaces == 1*n_elements ("For 1D and periodic domains, n_surf must be the same as 1*n_elem")
  end

  # Initialize boundary conditions
  boundary_conditions = init_boundary_conditions(n_boundaries_per_direction, mesh)

  # Initialize interpolation data structures
  n_nodes = POLYDEG + 1
  nodes, weights = gauss_lobatto_nodes_weights(n_nodes)
  inverse_weights = 1 ./ weights
  _, inverse_vandermonde_legendre = vandermonde_legendre(nodes)
  lhat = zeros(n_nodes, 2)
  lhat[:, 1] = calc_lhat(-1.0, nodes, weights)
  lhat[:, 2] = calc_lhat( 1.0, nodes, weights)

  # Initialize differentiation operator
  volume_integral_type = Val(Symbol(parameter("volume_integral_type", "weak_form",
                                              valid=["weak_form", "split_form", "shock_capturing"])))
  dhat = calc_dhat(nodes, weights)
  dsplit = calc_dsplit(nodes, weights)
  dsplit_transposed = transpose(calc_dsplit(nodes, weights))

  # Initialize L2 mortar projection operators
  amr_refine_right  = calc_forward_upper(n_nodes)
  amr_refine_left   = calc_forward_lower(n_nodes)
  amr_coarsen_right = calc_reverse_upper(n_nodes, Val(:gauss))
  amr_coarsen_left  = calc_reverse_lower(n_nodes, Val(:gauss))

  # Initialize data structures for error analysis (by default, we use twice the
  # number of analysis nodes as the normal solution)
  analysis_polydeg = 2 * (n_nodes) - 1
  analysis_nodes, analysis_weights = gauss_lobatto_nodes_weights(analysis_polydeg + 1)
  analysis_weights_volume = analysis_weights
  analysis_vandermonde = polynomial_interpolation_matrix(nodes, analysis_nodes)
  analysis_total_volume = mesh.tree.length_level_0^ndims(mesh)

  # Store which quantities should be analyzed in `analyze_solution`
  if parameter_exists("extra_analysis_quantities")
    extra_analysis_quantities = Symbol.(parameter("extra_analysis_quantities"))
  else
    extra_analysis_quantities = Symbol[]
  end
  analysis_quantities = vcat(collect(Symbol.(default_analysis_quantities(equation))),
                             extra_analysis_quantities)

  # If analysis should be saved to file, create file with header
  save_analysis = parameter("save_analysis", false)
  if save_analysis
    # Create output directory (if it does not exist)
    output_directory = parameter("output_directory", "out")
    mkpath(output_directory)

    # Determine filename
    analysis_filename = joinpath(output_directory, "analysis.dat")

    # Open file and write header
    save_analysis_header(analysis_filename, analysis_quantities, equation)
  else
    analysis_filename = ""
  end

  # Initialize AMR
  amr_indicator = Symbol(parameter("amr_indicator", "n/a",
                                   valid=["n/a", "gauss", "blast_wave"]))

  # Set global number of elements
  n_elements_global = n_elements

  # Initialize storage for element variables
  element_variables = Dict{Symbol, Union{Vector{Float64}, Vector{Int}}}()

  # maximum and minimum alpha for shock capturing
  shock_alpha_max = parameter("shock_alpha_max", 0.5)
  shock_alpha_min = parameter("shock_alpha_min", 0.001)
  shock_alpha_smooth = parameter("shock_alpha_smooth", true)

  # variable used to compute the shock capturing indicator
  # "eval is evil"
  # This is a temporary hack until we have switched to a library based approach
  # with pure Julia code instead of parameter files.
  shock_indicator_variable = eval(Symbol(parameter("shock_indicator_variable", "density_pressure")))

  # maximum and minimum alpha for amr control
  amr_alpha_max = parameter("amr_alpha_max", 0.5)
  amr_alpha_min = parameter("amr_alpha_min", 0.001)
  amr_alpha_smooth = parameter("amr_alpha_smooth", false)

  # Initialize element variables such that they are available in the first solution file
  if volume_integral_type === Val(:shock_capturing)
    element_variables[:blending_factor] = zeros(n_elements)
  end

  # Initialize storage for the cache
  cache = Dict{Symbol, Any}()
  thread_cache = create_thread_cache_1d(NVARS, POLYDEG+1)

  # Store initial state integrals for conservation error calculation
  initial_state_integrals = Vector{Float64}()

  # Convert all performance-critical fields to StaticArrays types
  nodes           = SVector{POLYDEG+1}(nodes)
  weights         = SVector{POLYDEG+1}(weights)
  inverse_weights = SVector{POLYDEG+1}(inverse_weights)
  lhat = SMatrix{POLYDEG+1,2}(lhat)
  dhat              = SMatrix{POLYDEG+1,POLYDEG+1}(dhat)
  dsplit            = SMatrix{POLYDEG+1,POLYDEG+1}(dsplit)
  dsplit_transposed = SMatrix{POLYDEG+1,POLYDEG+1}(dsplit_transposed)
  amr_refine_right   = SMatrix{POLYDEG+1,POLYDEG+1}(amr_refine_right)
  amr_refine_left    = SMatrix{POLYDEG+1,POLYDEG+1}(amr_refine_left)
  amr_coarsen_right  = SMatrix{POLYDEG+1,POLYDEG+1}(amr_coarsen_right)
  amr_coarsen_left   = SMatrix{POLYDEG+1,POLYDEG+1}(amr_coarsen_left)
  analysis_nodes          = SVector{analysis_polydeg+1}(analysis_nodes)
  analysis_weights        = SVector{analysis_polydeg+1}(analysis_weights)
  analysis_weights_volume = SVector{analysis_polydeg+1}(analysis_weights_volume)

  # Create actual DG solver instance
  dg = Dg1D{typeof(equation), typeof(mesh), NVARS, POLYDEG,
            typeof(surface_flux_function), typeof(volume_flux_function), typeof(initial_conditions),
            typeof(source_terms), typeof(boundary_conditions),
            typeof(volume_integral_type), typeof(shock_indicator_variable),
            typeof(nodes), typeof(dhat), typeof(lhat), typeof(inverse_vandermonde_legendre),
            typeof(amr_refine_right), typeof(analysis_nodes), typeof(analysis_vandermonde)}(
      equation,
      surface_flux_function, volume_flux_function,
      initial_conditions, source_terms,
      elements, n_elements,
      interfaces, n_interfaces,
      boundaries, n_boundaries, n_boundaries_per_direction,
      n_l2mortars,
      boundary_conditions,
      nodes, weights, inverse_weights,
      inverse_vandermonde_legendre, lhat,
      volume_integral_type,
      dhat, dsplit, dsplit_transposed,
      amr_refine_right, amr_refine_left,
      amr_coarsen_right, amr_coarsen_left,
      analysis_nodes, analysis_weights, analysis_weights_volume,
      analysis_vandermonde, analysis_total_volume,
      analysis_quantities, save_analysis, analysis_filename,
      shock_indicator_variable, shock_alpha_max, shock_alpha_min, shock_alpha_smooth,
      amr_indicator, amr_alpha_max, amr_alpha_min, amr_alpha_smooth,
      n_elements_global,
      element_variables, cache, thread_cache,
      initial_state_integrals)

  return dg
end

function create_thread_cache_1d(n_variables, n_nodes)
  # Type alias only for convenience
  A3d     = Array{Float64, 3}
  A2d     = Array{Float64, 2}
  A2dp1_x = Array{Float64, 2}

  MA1d    = MArray{Tuple{n_variables, n_nodes}, Float64}

  # Pre-allocate data structures to speed up computation (thread-safe)
  f1_threaded     = A3d[A3d(undef, n_variables, n_nodes, n_nodes) for _ in 1:Threads.nthreads()]
  fstar1_threaded = A2dp1_x[A2dp1_x(undef, n_variables, n_nodes+1) for _ in 1:Threads.nthreads()]

  indicator_threaded  = [A2d(undef, 1, n_nodes) for _ in 1:Threads.nthreads()]
  modal_threaded      = [A2d(undef, 1, n_nodes) for _ in 1:Threads.nthreads()]

  return (; f1_threaded,
            fstar1_threaded,
            indicator_threaded, modal_threaded)
end

# Count the number of interfaces that need to be created
function count_required_interfaces(mesh::TreeMesh1D, cell_ids)
  count = 0

  # Iterate over all cells
  for cell_id in cell_ids
    for direction in 1:n_directions(mesh.tree)
      # Only count interfaces in positive direction to avoid double counting
      if direction == 1
        continue
      end

      # Skip if no neighbor exists
      if !has_any_neighbor(mesh.tree, cell_id, direction)
        continue
      end

      count += 1
    end
  end

  return count
end

# Count the number of boundaries that need to be created
function count_required_boundaries(mesh::TreeMesh1D, cell_ids)
  count = 0

  # Iterate over all cells
  for cell_id in cell_ids
    for direction in 1:n_directions(mesh.tree)
      # If neighbor exists, current cell is not at a boundary
      if has_neighbor(mesh.tree, cell_id, direction)
        continue
      end

      # If coarse neighbor exists, current cell is not at a boundary
      if has_coarse_neighbor(mesh.tree, cell_id, direction)
        continue
      end

      # No neighbor exists in this direction -> must be a boundary
      count += 1
    end
  end

  return count
end

# Create element container, initialize element data, and return element container for further use
#
# NVARS: number of variables
# POLYDEG: polynomial degree
function init_elements(cell_ids, mesh::TreeMesh1D, ::Val{NVARS}, ::Val{POLYDEG}) where {NVARS, POLYDEG}
  # Initialize container
  n_elements = length(cell_ids)
  elements = ElementContainer1D{NVARS, POLYDEG}(n_elements)

  # Store cell ids
  elements.cell_ids .= cell_ids

  # Determine node locations
  n_nodes = POLYDEG + 1
  nodes, _ = gauss_lobatto_nodes_weights(n_nodes)

  # Calculate inverse Jacobian and node coordinates
  for element_id in 1:nelements(elements)
    # Get cell id
    cell_id = cell_ids[element_id]

    # Get cell length
    dx = length_at_cell(mesh.tree, cell_id)

    # Calculate inverse Jacobian as 1/(h/2)
    elements.inverse_jacobian[element_id] = 2/dx

    # Calculate node coordinates
      for i in 1:n_nodes
        elements.node_coordinates[1, i, element_id] = (
            mesh.tree.coordinates[1, cell_id] + dx/2 * nodes[i])
      end
  end

  return elements
end

# Create interface container, initialize interface data, and return interface container for further use
#
# NVARS: number of variables
# POLYDEG: polynomial degree
function init_interfaces(cell_ids, mesh::TreeMesh1D, ::Val{NVARS}, ::Val{POLYDEG}, elements) where {NVARS, POLYDEG}
  # Initialize container
  n_interfaces = count_required_interfaces(mesh, cell_ids)
  interfaces = InterfaceContainer1D{NVARS, POLYDEG}(n_interfaces)

  # Connect elements with interfaces
  init_interface_connectivity!(elements, interfaces, mesh)

  return interfaces
end

# Create boundaries container, initialize boundary data, and return boundaries container
#
# NVARS: number of variables
# POLYDEG: polynomial degree
function init_boundaries(cell_ids, mesh::TreeMesh1D, ::Val{NVARS}, ::Val{POLYDEG}, elements) where {NVARS, POLYDEG}
  # Initialize container
  n_boundaries = count_required_boundaries(mesh, cell_ids)
  boundaries = BoundaryContainer1D{NVARS, POLYDEG}(n_boundaries)

  # Connect elements with boundaries
  n_boundaries_per_direction = init_boundary_connectivity!(elements, boundaries, mesh)

  return boundaries, n_boundaries_per_direction
end

# Initialize connectivity between elements and interfaces
function init_interface_connectivity!(elements, interfaces, mesh::TreeMesh1D)
  # Construct cell -> element mapping for easier algorithm implementation
  tree = mesh.tree
  c2e = zeros(Int, length(tree))
  for element_id in 1:nelements(elements)
    c2e[elements.cell_ids[element_id]] = element_id
  end

  # Reset interface count
  count = 0

  # Iterate over all elements to find neighbors and to connect via interfaces
  for element_id in 1:nelements(elements)
    # Get cell id
    cell_id = elements.cell_ids[element_id]

    # Loop over directions
    for direction in 1:n_directions(mesh.tree)
      # Only create interfaces in positive direction
      if direction == 1
        continue
      end

      # Skip if no neighbor exists and current cell is not small
      if !has_any_neighbor(mesh.tree, cell_id, direction)
        continue
      end

      count += 1

      if has_neighbor(mesh.tree, cell_id, direction)
        neighbor_cell_id = mesh.tree.neighbor_ids[direction, cell_id]
        if has_children(mesh.tree, neighbor_cell_id) # Cell has small neighbor
          interfaces.neighbor_ids[2, count] = c2e[mesh.tree.child_ids[1, neighbor_cell_id]]
        else # Cell has same refinement level neighbor
          interfaces.neighbor_ids[2, count] = c2e[neighbor_cell_id]
        end
      else # Cell is small and has large neighbor
        parent_id = mesh.tree.parent_ids[cell_id]
        neighbor_cell_id = mesh.tree.neighbor_ids[direction, parent_id]
        interfaces.neighbor_ids[2, count] = c2e[neighbor_cell_id]
      end

      interfaces.neighbor_ids[1, count] = element_id
      # Set orientation (x -> 1)
      interfaces.orientations[count] = 1
    end
  end

  @assert count == ninterfaces(interfaces) ("Actual interface count ($count) does not match " *
                                            "expectations $(ninterfaces(interfaces))")
end

# Initialize connectivity between elements and boundaries
function init_boundary_connectivity!(elements, boundaries, mesh::TreeMesh1D)
  # Reset boundaries count
  count = 0

  # Initialize boundary counts
  counts_per_direction = MVector(0, 0)

  # OBS! Iterate over directions first, then over elements, and count boundaries in each direction
  # Rationale: This way the boundaries are internally sorted by the directions -x, +x, -y etc.,
  #            obviating the need to store the boundary condition to be applied explicitly.
  # Loop over directions
  for direction in 1:n_directions(mesh.tree)
    # Iterate over all elements to find missing neighbors and to connect to boundaries
    for element_id in 1:nelements(elements)
      # Get cell id
      cell_id = elements.cell_ids[element_id]

      # If neighbor exists, current cell is not at a boundary
      if has_neighbor(mesh.tree, cell_id, direction)
        continue
      end

      # If coarse neighbor exists, current cell is not at a boundary
      if has_coarse_neighbor(mesh.tree, cell_id, direction)
        continue
      end

      # Create boundary
      count += 1
      counts_per_direction[direction] += 1

      # Set neighbor element id
      boundaries.neighbor_ids[count] = element_id

      # Set neighbor side, which denotes the direction (1 -> negative, 2 -> positive) of the element
      if direction == 2
        boundaries.neighbor_sides[count] = 1
      else
        boundaries.neighbor_sides[count] = 2
      end

      # Set orientation (x -> 1)
      boundaries.orientations[count] = 1

      # Store node coordinates
      enc = elements.node_coordinates
      if direction == 1 # -x direction
        boundaries.node_coordinates[:, count] .= enc[:, 1,  element_id]
      elseif direction == 2 # +x direction
        boundaries.node_coordinates[:, count] .= enc[:, end, element_id]
      else
        error("should not happen")
      end
    end
  end

  @assert count == nboundaries(boundaries) ("Actual boundaries count ($count) does not match " *
                                            "expectations $(nboundaries(boundaries))")

  @assert sum(counts_per_direction) == count

  return SVector(counts_per_direction)
end

function init_boundary_conditions(n_boundaries_per_direction, mesh::TreeMesh1D)
  # "eval is evil"
  # This is a temporary hack until we have switched to a library based approach
  # with pure Julia code instead of parameter files.
  bcs = parameter("boundary_conditions", ["nothing", "nothing"])
  if bcs isa AbstractArray
    boundary_conditions = eval_if_not_function.(bcs)
  else
    # This adds support for using a scalar boundary condition (like 'periodicity = "false"')
    boundary_conditions = eval_if_not_function.([bcs for _ in 1:n_directions(mesh.tree)])
  end

  # Sanity check about specifying boundary conditions
  for direction in 1:n_directions(mesh.tree)
    bc = boundary_conditions[direction]
    count = n_boundaries_per_direction[direction]
    if direction == 1
      dir = "-x"
    elseif direction == 2
      dir = "+x"
    end

    # All directions with boundaries should have a boundary condition
    if count > 0 && isnothing(bc)
      error("Found $(count) boundaries in the $(dir)-direction, but corresponding boundary " *
            "condition is '$(get_name(bc))'")
    end
  end

  return Tuple(boundary_conditions)
end

"""
    integrate(func, dg::Dg1D, args...; normalize=true)

Call function `func` for each DG node and integrate the result over the computational domain.

The function `func` is called as `func(i, j, element_id, dg, args...)` for each
volume node `(i, j)` and each `element_id`. Additional positional
arguments `args...` are passed along as well. If `normalize` is true, the result
is divided by the total volume of the computational domain.

# Examples
Calculate the integral of the time derivative of the entropy, i.e.,
∫(∂S/∂t)dΩ = ∫(∂S/∂u ⋅ ∂u/∂t)dΩ:
```julia
# Calculate integral of entropy time derivative
dsdu_ut = integrate(dg, dg.elements.u, dg.elements.u_t) do i, element_id, dg, u, u_t
  u_node   = get_node_vars(u,   dg, i, element_id)
  u_t_node = get_node_vars(u_t, dg, i, element_id)
  dot(cons2entropy(u_node, equations(dg)), u_t_node)
end
```
"""
function integrate(func, dg::Dg1D, args...; normalize=true)
  # Initialize integral with zeros of the right shape
  integral = zero(func(1, 1, dg, args...))

  # Use quadrature to numerically integrate over entire domain
  for element_id in 1:dg.n_elements
    jacobian_volume = inv(dg.elements.inverse_jacobian[element_id])^ndims(dg)
      for i in 1:nnodes(dg)
        integral += jacobian_volume * dg.weights[i] * func(i, element_id, dg, args...)
      end
  end

  # Normalize with total volume
  if normalize
    integral = integral/dg.analysis_total_volume
  end

  return integral
end

"""
    integrate(func, u, dg::Dg1D; normalize=true)
    integrate(u, dg::Dg1D; normalize=true)

Call function `func` for each DG node and integrate the result over the computational domain.

The function `func` is called as `func(u_local)` for each volume node `(i, j)`
and each `element_id`, where `u_local` is an `SVector`ized copy of
`u[:, i, j, element_id]`. If `normalize` is true, the result is divided by the
total volume of the computational domain. If `func` is omitted, it defaults to
`identity`.

# Examples
Calculate the integral over all conservative variables:
```julia
state_integrals = integrate(dg.elements.u, dg)
```
"""
function integrate(func, u, dg::Dg1D; normalize=true)
  func_wrapped = function(i, element_id, dg, u)
    u_local = get_node_vars(u, dg, i, element_id)
    return func(u_local)
  end
  return integrate(func_wrapped, dg, u; normalize=normalize)
end
integrate(u, dg::Dg1D; normalize=true) = integrate(identity, u, dg; normalize=normalize)

# Calculate L2/Linf error norms based on "exact solution"
function calc_error_norms(func, dg::Dg1D, t)
  # Gather necessary information
  equation = equations(dg)
  n_nodes_analysis = size(dg.analysis_vandermonde, 1)

  # pre-allocate buffers
  u = zeros(eltype(dg.elements.u),
            nvariables(dg), size(dg.analysis_vandermonde, 1))

  x = zeros(eltype(dg.elements.node_coordinates),
            1, size(dg.analysis_vandermonde, 1))

  # Set up data structures
  l2_error   = zero(func(get_node_vars(dg.elements.u, dg, 1, 1), equation))
  linf_error = zero(func(get_node_vars(dg.elements.u, dg, 1, 1), equation))

  # Iterate over all elements for error calculations
  for element_id in 1:dg.n_elements
    # Interpolate solution and node locations to analysis nodes
    multiply_dimensionwise!(u, dg.analysis_vandermonde, view(dg.elements.u,                :, :, element_id))
    multiply_dimensionwise!(x, dg.analysis_vandermonde, view(dg.elements.node_coordinates, :, :, element_id))

    # Calculate errors at each analysis node
    weights = dg.analysis_weights_volume
    jacobian_volume = inv(dg.elements.inverse_jacobian[element_id])^ndims(dg)
    for i in 1:n_nodes_analysis
      u_exact = dg.initial_conditions(get_node_coords(x, dg, i), t, equation)
      diff = func(u_exact, equation) - func(get_node_vars(u, dg, i), equation)
      l2_error += diff.^2 * (weights[i] *  jacobian_volume)
      linf_error = @. max(linf_error, abs(diff))
    end
  end

  # For L2 error, divide by total volume
  l2_error = @. sqrt(l2_error / dg.analysis_total_volume)

  return l2_error, linf_error
end

# Integrate ∂S/∂u ⋅ ∂u/∂t over the entire domain
function calc_entropy_timederivative(dg::Dg1D, t)
  # Compute ut = rhs(u) with current solution u
  @notimeit timer() rhs!(dg, t)

  # Calculate ∫(∂S/∂u ⋅ ∂u/∂t)dΩ
  dsdu_ut = integrate(dg, dg.elements.u, dg.elements.u_t) do i, element_id, dg, u, u_t
    u_node   = get_node_vars(u,   dg, i, element_id)
    u_t_node = get_node_vars(u_t, dg, i, element_id)
    dot(cons2entropy(u_node, equations(dg)), u_t_node)
  end

  return dsdu_ut
end

"""
    analyze_solution(dg::Dg1D, mesh::TreeMesh, time, dt, step, runtime_absolute, runtime_relative)

Calculate error norms and other analysis quantities to analyze the solution
during a simulation, and return the L2 and Linf errors. `dg` and `mesh` are the
DG and the mesh instance, respectively. `time`, `dt`, and `step` refer to the
current simulation time, the last time step size, and the current time step
count. The run time (in seconds) is given in `runtime_absolute`, while the
performance index is specified in `runtime_relative`.

**Note:** Keep order of analysis quantities in sync with
          [`save_analysis_header`](@ref) when adding or changing quantities.
"""
function analyze_solution(dg::Dg1D, mesh::TreeMesh, time::Real, dt::Real, step::Integer,
                          runtime_absolute::Real, runtime_relative::Real; solver_gravity=nothing)
  equation = equations(dg)

  # General information
  println()
  println("-"^80)
  println(" Simulation running '", get_name(equation), "' with POLYDEG = ", polydeg(dg))
  println("-"^80)
  println(" #timesteps:     " * @sprintf("% 14d", step) *
          "               " *
          " run time:       " * @sprintf("%10.8e s", runtime_absolute))
  println(" dt:             " * @sprintf("%10.8e", dt) *
          "               " *
          " Time/DOF/step:  " * @sprintf("%10.8e s", runtime_relative))
  println(" sim. time:      " * @sprintf("%10.8e", time))

  # Level information (only show for AMR)
  if parameter("amr_interval", 0)::Int > 0
    levels = Vector{Int}(undef, dg.n_elements)
    for element_id in 1:dg.n_elements
      levels[element_id] = mesh.tree.levels[dg.elements.cell_ids[element_id]]
    end
    min_level = minimum(levels)
    max_level = maximum(levels)

    println(" #elements:      " * @sprintf("% 14d", dg.n_elements))
    for level = max_level:-1:min_level+1
      println(" ├── level $level:    " * @sprintf("% 14d", count(x->x==level, levels)))
    end
    println(" └── level $min_level:    " * @sprintf("% 14d", count(x->x==min_level, levels)))
  end
  println()

  # Open file for appending and store time step and time information
  if dg.save_analysis
    f = open(dg.analysis_filename, "a")
    @printf(f, "% 9d", step)
    @printf(f, "  %10.8e", time)
    @printf(f, "  %10.8e", dt)
  end

  # Calculate and print derived quantities (error norms, entropy etc.)
  # Variable names required for L2 error, Linf error, and conservation error
  if any(q in dg.analysis_quantities for q in
         (:l2_error, :linf_error, :conservation_error, :residual))
    print(" Variable:    ")
    for v in 1:nvariables(equation)
      @printf("   %-14s", varnames_cons(equation)[v])
    end
    println()
  end

  # Calculate L2/Linf errors, which are also returned by analyze_solution
  l2_error, linf_error = calc_error_norms(dg, time)

  # L2 error
  if :l2_error in dg.analysis_quantities
    print(" L2 error:    ")
    for v in 1:nvariables(equation)
      @printf("  % 10.8e", l2_error[v])
      dg.save_analysis && @printf(f, "  % 10.8e", l2_error[v])
    end
    println()
  end

  # Linf error
  if :linf_error in dg.analysis_quantities
    print(" Linf error:  ")
    for v in 1:nvariables(equation)
      @printf("  % 10.8e", linf_error[v])
      dg.save_analysis && @printf(f, "  % 10.8e", linf_error[v])
    end
    println()
  end

  # Conservation errror
  if :conservation_error in dg.analysis_quantities
    # Calculate state integrals
    state_integrals = integrate(dg.elements.u, dg)

    # Store initial state integrals at first invocation
    if isempty(dg.initial_state_integrals)
      dg.initial_state_integrals = zeros(nvariables(equation))
      dg.initial_state_integrals .= state_integrals
    end

    print(" |∑U - ∑U₀|:  ")
    for v in 1:nvariables(equation)
      err = abs(state_integrals[v] - dg.initial_state_integrals[v])
      @printf("  % 10.8e", err)
      dg.save_analysis && @printf(f, "  % 10.8e", err)
    end
    println()
  end

  # Residual (defined here as the vector maximum of the absolute values of the time derivatives)
  if :residual in dg.analysis_quantities
    print(" max(|Uₜ|):   ")
    for v in 1:nvariables(equation)
      # Calculate maximum absolute value of Uₜ
      @views res = maximum(abs, view(dg.elements.u_t, v, :, :))
      @printf("  % 10.8e", res)
      dg.save_analysis && @printf(f, "  % 10.8e", res)
    end
    println()
  end

  # L2/L∞ errors of the primitive variables
  if :l2_error_primitive in dg.analysis_quantities || :linf_error_primitive in dg.analysis_quantities
    l2_error_prim, linf_error_prim = calc_error_norms(cons2prim, dg, time)

    print(" Variable:    ")
    for v in 1:nvariables(equation)
      @printf("   %-14s", varnames_prim(equation)[v])
    end
    println()

    # L2 error
    if :l2_error_primitive in dg.analysis_quantities
      print(" L2 error prim.: ")
      for v in 1:nvariables(equation)
        @printf("%10.8e   ", l2_error_prim[v])
        dg.save_analysis && @printf(f, "  % 10.8e", l2_error_prim[v])
      end
      println()
    end

    # L∞ error
    if :linf_error_primitive in dg.analysis_quantities
      print(" Linf error pri.:")
      for v in 1:nvariables(equation)
        @printf("%10.8e   ", linf_error_prim[v])
        dg.save_analysis && @printf(f, "  % 10.8e", linf_error_prim[v])
      end
      println()
    end
  end

  # Entropy time derivative
  if :dsdu_ut in dg.analysis_quantities
    duds_ut = calc_entropy_timederivative(dg, time)
    print(" ∑∂S/∂U ⋅ Uₜ: ")
    @printf("  % 10.8e", duds_ut)
    dg.save_analysis && @printf(f, "  % 10.8e", duds_ut)
    println()
  end

  # Entropy
  if :entropy in dg.analysis_quantities
    s = integrate(dg, dg.elements.u) do i, element_id, dg, u
      cons = get_node_vars(u, dg, i, element_id)
      return entropy(cons, equations(dg))
    end
    print(" ∑S:          ")
    @printf("  % 10.8e", s)
    dg.save_analysis && @printf(f, "  % 10.8e", s)
    println()
  end

  # Total energy
  if :energy_total in dg.analysis_quantities
    e_total = integrate(dg, dg.elements.u) do i, element_id, dg, u
      cons = get_node_vars(u, dg, i, element_id)
      return energy_total(cons, equations(dg))
    end
    print(" ∑e_total:    ")
    @printf("  % 10.8e", e_total)
    dg.save_analysis && @printf(f, "  % 10.8e", e_total)
    println()
  end

  # Kinetic energy
  if :energy_kinetic in dg.analysis_quantities
    e_kinetic = integrate(dg, dg.elements.u) do i, element_id, dg, u
      cons = get_node_vars(u, dg, i, element_id)
      return energy_kinetic(cons, equations(dg))
    end
    print(" ∑e_kinetic:  ")
    @printf("  % 10.8e", e_kinetic)
    dg.save_analysis && @printf(f, "  % 10.8e", e_kinetic)
    println()
  end

  # Internal energy
  if :energy_internal in dg.analysis_quantities
    e_internal = integrate(dg, dg.elements.u) do i, element_id, dg, u
      cons = get_node_vars(u, dg, i, element_id)
      return energy_internal(cons, equations(dg))
    end
    print(" ∑e_internal: ")
    @printf("  % 10.8e", e_internal)
    dg.save_analysis && @printf(f, "  % 10.8e", e_internal)
    println()
  end

  println("-"^80)
  println()

  # Add line break and close analysis file if it was opened
  if dg.save_analysis
    println(f)
    close(f)
  end

  # Return errors for EOC analysis
  return l2_error, linf_error
end

"""
    save_analysis_header(filename, quantities, equation)

Truncate file `filename` and save a header with the names of the quantities
`quantities` that will subsequently written to `filename` by
[`analyze_solution`](@ref). Since some quantities are equation-specific, the
system of equations instance is passed in `equation`.

**Note:** Keep order of analysis quantities in sync with
          [`analyze_solution`](@ref) when adding or changing quantities.
"""
function save_analysis_header(filename, quantities, equation::AbstractEquation{1})
  open(filename, "w") do f
    @printf(f, "#%-8s", "timestep")
    @printf(f, "  %-14s", "time")
    @printf(f, "  %-14s", "dt")
    if :l2_error in quantities
      for v in varnames_cons(equation)
        @printf(f, "   %-14s", "l2_" * v)
      end
    end
    if :linf_error in quantities
      for v in varnames_cons(equation)
        @printf(f, "   %-14s", "linf_" * v)
      end
    end
    if :conservation_error in quantities
      for v in varnames_cons(equation)
        @printf(f, "   %-14s", "cons_" * v)
      end
    end
    if :residual in quantities
      for v in varnames_cons(equation)
        @printf(f, "   %-14s", "res_" * v)
      end
    end
    if :l2_error_primitive in quantities
      for v in varnames_prim(equation)
        @printf(f, "   %-14s", "l2_" * v)
      end
    end
    if :linf_error_primitive in quantities
      for v in varnames_prim(equation)
        @printf(f, "   %-14s", "linf_" * v)
      end
    end
    if :dsdu_ut in quantities
      @printf(f, "   %-14s", "dsdu_ut")
    end
    if :entropy in quantities
      @printf(f, "   %-14s", "entropy")
    end
    if :energy_total in quantities
      @printf(f, "   %-14s", "e_total")
    end
    if :energy_kinetic in quantities
      @printf(f, "   %-14s", "e_kinetic")
    end
    if :energy_internal in quantities
      @printf(f, "   %-14s", "e_internal")
    end
    println(f)
  end
end

# Call equation-specific initial conditions functions and apply to all elements
function set_initial_conditions!(dg::Dg1D, time)
  equation = equations(dg)
  # make sure that the random number generator is reseted and the ICs are reproducible in the julia REPL/interactive mode
  seed!(0)
  for element_id in 1:dg.n_elements
      for i in 1:nnodes(dg)
        dg.elements.u[:, i, element_id] .= dg.initial_conditions(
            dg.elements.node_coordinates[:, i, element_id], time, equation)
      end
  end
end

# Calculate time derivative
function rhs!(dg::Dg1D, t_stage)
  # Reset u_t
  @timeit timer() "reset ∂u/∂t" dg.elements.u_t .= 0

  # Calculate volume integral
  @timeit timer() "volume integral" calc_volume_integral!(dg)

  # Prolong solution to interfaces
  @timeit timer() "prolong2interfaces" prolong2interfaces!(dg)

  # Calculate interface fluxes
  @timeit timer() "interface flux" calc_interface_flux!(dg)

  # Prolong solution to boundaries
  @timeit timer() "prolong2boundaries" prolong2boundaries!(dg)

  # Calculate boundary fluxes
  @timeit timer() "boundary flux" calc_boundary_flux!(dg, t_stage)

  # Calculate surface integrals
  @timeit timer() "surface integral" calc_surface_integral!(dg)

  # Apply Jacobian from mapping to reference element
  @timeit timer() "Jacobian" apply_jacobian!(dg)

  # Calculate source terms
  @timeit timer() "source terms" calc_sources!(dg, dg.source_terms, t_stage)
end

# TODO: implement 1D!!!
# Apply positivity limiter of Zhang and Shu to nodal values elements.u
function apply_positivity_preserving_limiter!(dg::Dg1D)
end

# Calculate volume integral and update u_t
calc_volume_integral!(dg::Dg1D) = calc_volume_integral!(dg.elements.u_t, dg.volume_integral_type, dg)

# Calculate volume integral (DGSEM in weak form)
function calc_volume_integral!(u_t, ::Val{:weak_form}, dg::Dg1D)
  @unpack dhat = dg

  Threads.@threads for element_id in 1:dg.n_elements
    # Calculate volume integral
      for i in 1:nnodes(dg)
        u_node = get_node_vars(dg.elements.u, dg, i, element_id)

        flux1 = calcflux(u_node, 1, equations(dg))
        for l in 1:nnodes(dg)
          integral_contribution = dhat[l, i] * flux1
          add_to_node_vars!(u_t, integral_contribution, dg, l, element_id)
        end
      end
  end
end

# Calculate volume integral (DGSEM in split form)
@inline function calc_volume_integral!(u_t, volume_integral_type::Val{:split_form}, dg::Dg1D)
  calc_volume_integral!(u_t, volume_integral_type, dg.thread_cache, dg)
end

function calc_volume_integral!(u_t, ::Val{:split_form}, cache, dg::Dg1D)
  Threads.@threads for element_id in 1:dg.n_elements
    split_form_kernel!(u_t, element_id, cache, dg)
  end
end

@inline function split_form_kernel!(u_t, element_id, cache, dg::Dg1D, alpha=true)
  # true * [some floating point value] == [exactly the same floating point value]
  # This can get optimized away due to constant propagation.
  @unpack volume_flux_function, dsplit = dg

  # Calculate volume integral in one element
  for i in 1:nnodes(dg)
    u_node = get_node_vars(dg.elements.u, dg, i, element_id)

    # x direction
    # use consistency of the volume flux to make this evaluation cheaper
    flux = calcflux(u_node, 1, equations(dg))
    integral_contribution = alpha * dsplit[i, i] * flux
    add_to_node_vars!(u_t, integral_contribution, dg, i, element_id)
    # use symmetry of the volume flux for the remaining terms
    for ii in (i+1):nnodes(dg)
      u_node_ii = get_node_vars(dg.elements.u, dg, ii, element_id)
      flux = volume_flux_function(u_node, u_node_ii, 1, equations(dg))
      integral_contribution = alpha * dsplit[i, ii] * flux
      add_to_node_vars!(u_t, integral_contribution, dg, i, element_id)
      integral_contribution = alpha * dsplit[ii, i] * flux
      add_to_node_vars!(u_t, integral_contribution, dg, ii, element_id)
    end

  end
end

# Calculate volume integral (DGSEM in split form with shock capturing)
function calc_volume_integral!(u_t, ::Val{:shock_capturing}, dg::Dg1D)
  # (Re-)initialize element variable storage for blending factor
  if (!haskey(dg.element_variables, :blending_factor) ||
      length(dg.element_variables[:blending_factor]) != dg.n_elements)
    dg.element_variables[:blending_factor] = Vector{Float64}(undef, dg.n_elements)
  end
  if (!haskey(dg.element_variables, :blending_factor_tmp) ||
      length(dg.element_variables[:blending_factor_tmp]) != dg.n_elements)
    dg.element_variables[:blending_factor_tmp] = Vector{Float64}(undef, dg.n_elements)
  end

  # Initialize element variable storage for the cache
  if (!haskey(dg.cache, :element_ids_dg))
    dg.cache[:element_ids_dg] = Int[]
    sizehint!(dg.cache[:element_ids_dg], dg.n_elements)
  end
  if (!haskey(dg.cache, :element_ids_dgfv))
    dg.cache[:element_ids_dgfv] = Int[]
    sizehint!(dg.cache[:element_ids_dgfv], dg.n_elements)
  end

  calc_volume_integral!(u_t, Val(:shock_capturing),
                        dg.element_variables[:blending_factor], dg.element_variables[:blending_factor_tmp],
                        dg.cache[:element_ids_dg], dg.cache[:element_ids_dgfv],
                        dg.thread_cache,
                        dg)
end

function calc_volume_integral!(u_t, ::Val{:shock_capturing}, alpha, alpha_tmp,
                               element_ids_dg, element_ids_dgfv, thread_cache, dg::Dg1D)
  @unpack dsplit_transposed, inverse_weights = dg
  @unpack fstar1_threaded = thread_cache

  # Calculate blending factors α: u = u_DG * (1 - α) + u_FV * α
  @timeit timer() "blending factors" calc_blending_factors!(alpha, alpha_tmp, dg.elements.u,
    dg.shock_alpha_max,
    dg.shock_alpha_min,
    dg.shock_alpha_smooth,
    dg.shock_indicator_variable, thread_cache, dg)

  # Determine element ids for DG-only and blended DG-FV volume integral
  pure_and_blended_element_ids!(element_ids_dg, element_ids_dgfv, alpha, dg)

  # Loop over pure DG elements
  @timeit timer() "pure DG" Threads.@threads for element_id in element_ids_dg
    split_form_kernel!(u_t, element_id, thread_cache, dg)
  end

  # Loop over blended DG-FV elements
  @timeit timer() "blended DG-FV" Threads.@threads for element_id in element_ids_dgfv
    # Calculate DG volume integral contribution
    split_form_kernel!(u_t, element_id, thread_cache, dg, 1 - alpha[element_id])

    # Calculate FV two-point fluxes
    fstar1 = fstar1_threaded[Threads.threadid()]
    calcflux_fv!(fstar1, dg.elements.u, element_id, dg)

    # Calculate FV volume integral contribution
    for i in 1:nnodes(dg)
      for v in 1:nvariables(dg)
        u_t[v, i, element_id] += ( alpha[element_id] * (inverse_weights[i] * (fstar1[v, i+1] - fstar1[v, i]) ) )

      end
    end
  end
end

"""
    calcflux_fv!(fstar1, u_leftright, u, element_id, dg::Dg1D)

Calculate the finite volume fluxes inside the elements.

# Arguments
- `fstar1::AbstractArray{T} where T<:Real`:
- `dg::Dg1D`
- `u::AbstractArray{T} where T<:Real`
- `element_id::Integer`
"""
@inline function calcflux_fv!(fstar1, u, element_id, dg::Dg1D)
  @unpack surface_flux_function = dg

  fstar1[:, 1           ] .= zero(eltype(fstar1))
  fstar1[:, nnodes(dg)+1] .= zero(eltype(fstar1))

    for i in 2:nnodes(dg)
      u_ll = get_node_vars(u, dg, i-1, element_id)
      u_rr = get_node_vars(u, dg, i,   element_id)
      flux = surface_flux_function(u_ll, u_rr, 1, equations(dg)) # orientation 1: x direction
      set_node_vars!(fstar1, flux, dg, i)
    end
end

# Prolong solution to interfaces (for GL nodes: just a copy)
function prolong2interfaces!(dg::Dg1D)
  equation = equations(dg)

  Threads.@threads for s in 1:dg.n_interfaces
    left_element_id = dg.interfaces.neighbor_ids[1, s]
    right_element_id = dg.interfaces.neighbor_ids[2, s]
    # interface in x-direction
    for v in 1:nvariables(dg)
      dg.interfaces.u[1, v, s] = dg.elements.u[v, nnodes(dg), left_element_id]
      dg.interfaces.u[2, v, s] = dg.elements.u[v,          1, right_element_id]
    end
  end
end

# Prolong solution to boundaries (for GL nodes: just a copy)
function prolong2boundaries!(dg::Dg1D)
  equation = equations(dg)

  for b in 1:dg.n_boundaries
    element_id = dg.boundaries.neighbor_ids[b]
    if dg.boundaries.neighbor_sides[b] == 1 # Element in -x direction of boundary
      for v in 1:nvariables(dg)
        dg.boundaries.u[1, v, b] = dg.elements.u[v, nnodes(dg), element_id]
      end
    else # Element in +x direction of boundary
      for v in 1:nvariables(dg)
        dg.boundaries.u[2, v, b] = dg.elements.u[v, 1,          element_id]
      end
    end
  end
end

# Calculate and store the surface fluxes (standard Riemann and nonconservative parts) at an interface
# OBS! Regarding the nonconservative terms: 1) currently only needed for the MHD equations
#                                           2) not implemented for boundaries
calc_interface_flux!(dg::Dg1D) = calc_interface_flux!(dg.elements.surface_flux_values, dg)

function calc_interface_flux!(surface_flux_values,  dg::Dg1D)
  @unpack surface_flux_function = dg
  @unpack u, neighbor_ids, orientations = dg.interfaces

  Threads.@threads for s in 1:dg.n_interfaces
    # Get neighboring elements
    left_id  = neighbor_ids[1, s]
    right_id = neighbor_ids[2, s]

    # Determine interface direction with respect to elements:
    # orientation = 1: left -> 2, right -> 1
    left_direction  = 2 * orientations[s]
    right_direction = 2 * orientations[s] - 1

    # Call pointwise Riemann solver
    u_ll, u_rr = get_surface_node_vars(u, dg, s)
    flux = surface_flux_function(u_ll, u_rr, orientations[s], equations(dg))

    # Copy flux to left and right element storage
    for v in 1:nvariables(dg)
      surface_flux_values[v, left_direction, left_id]  = surface_flux_values[v, right_direction, right_id] = flux[v]
    end
  end
end

# Calculate and store boundary flux across domain boundaries
#NOTE: Do we need to dispatch on have_nonconservative_terms(dg.equations)?
calc_boundary_flux!(dg::Dg1D, time) = calc_boundary_flux!(dg.elements.surface_flux_values, dg, time)

function calc_boundary_flux!(surface_flux_values, dg::Dg1D, time)
  @unpack n_boundaries_per_direction, boundary_conditions = dg

  # Calculate indices
  lasts = accumulate(+, n_boundaries_per_direction)
  firsts = lasts - n_boundaries_per_direction .+ 1

  # Calc boundary fluxes in each direction
  calc_boundary_flux_by_direction!(surface_flux_values, dg, time,
                                   boundary_conditions[1], 1, firsts[1], lasts[1])
  calc_boundary_flux_by_direction!(surface_flux_values, dg, time,
                                   boundary_conditions[2], 2, firsts[2], lasts[2])
end

function calc_boundary_flux_by_direction!(surface_flux_values, dg::Dg1D, time, boundary_condition,
                                          direction, first_boundary_id, last_boundary_id)
  @unpack surface_flux_function = dg
  @unpack u, neighbor_ids, neighbor_sides, node_coordinates, orientations = dg.boundaries

  Threads.@threads for b in first_boundary_id:last_boundary_id
    # Get neighboring element
    neighbor_id = neighbor_ids[b]

    # Get boundary flux
    u_ll, u_rr = get_surface_node_vars(u, dg, b)
    if neighbor_sides[b] == 1 # Element is on the left, boundary on the right
      u_inner = u_ll
    else # Element is on the right, boundary on the left
      u_inner = u_rr
    end
    x = get_node_coords(node_coordinates, dg, b)
    flux = boundary_condition(u_inner, orientations[b], direction, x, time, surface_flux_function,
                                equations(dg))

    # Copy flux to left and right element storage
    for v in 1:nvariables(dg)
      surface_flux_values[v, direction, neighbor_id] = flux[v]
    end
  end
end

# Calculate surface integrals and update u_t
calc_surface_integral!(dg::Dg1D) = calc_surface_integral!(dg.elements.u_t, dg.elements.surface_flux_values, dg)
function calc_surface_integral!(u_t, surface_flux_values, dg::Dg1D)
  @unpack lhat = dg

  Threads.@threads for element_id in 1:dg.n_elements
      for v in 1:nvariables(dg)
        # surface at -x
        u_t[v, 1,          element_id] -= surface_flux_values[v, 1, element_id] * lhat[1,          1]
        # surface at +x
        u_t[v, nnodes(dg), element_id] += surface_flux_values[v, 2, element_id] * lhat[nnodes(dg), 2]
      end
  end
end

# Apply Jacobian from mapping to reference element
function apply_jacobian!(dg::Dg1D)
  Threads.@threads for element_id in 1:dg.n_elements
    factor = -dg.elements.inverse_jacobian[element_id]
      for i in 1:nnodes(dg)
        for v in 1:nvariables(dg)
          dg.elements.u_t[v, i, element_id] *= factor
        end
      end
  end
end

# Calculate source terms and apply them to u_t
function calc_sources!(dg::Dg1D, source_terms::Nothing, t)
  return nothing
end

function calc_sources!(dg::Dg1D, source_terms, t)
  Threads.@threads for element_id in 1:dg.n_elements
    source_terms(dg.elements.u_t, dg.elements.u,
                 dg.elements.node_coordinates, element_id, t, nnodes(dg), equations(dg))
  end
end

# Calculate stable time step size
function calc_dt(dg::Dg1D, cfl)
  min_dt = Inf
  for element_id in 1:dg.n_elements
    dt = calc_max_dt(dg.elements.u, element_id,
                     dg.elements.inverse_jacobian[element_id], cfl, equations(dg), dg)
    min_dt = min(min_dt, dt)
  end

  return min_dt
end

# Calculate blending factors used for shock capturing, or amr control
function calc_blending_factors!(alpha, alpha_pre_smooth, u,
                                alpha_max, alpha_min, do_smoothing,
                                indicator_variable, thread_cache, dg::Dg1D)
  # temporary buffers
  @unpack indicator_threaded, modal_threaded = thread_cache
  # magic parameters
  threshold = 0.5 * 10^(-1.8 * (nnodes(dg))^0.25)
  parameter_s = log((1 - 0.0001)/0.0001)

  Threads.@threads for element_id in 1:dg.n_elements
    indicator  = indicator_threaded[Threads.threadid()]
    modal      = modal_threaded[Threads.threadid()]

    # Calculate indicator variables at Gauss-Lobatto nodes
    cons2indicator!(indicator, u, element_id, indicator_variable, dg)

    # Convert to modal representation
    multiply_dimensionwise!(modal, dg.inverse_vandermonde_legendre, indicator)

    # Calculate total energies for all modes, without highest, without two highest
    total_energy = 0.0
    for i in 1:nnodes(dg)
      total_energy += modal[1, i]^2
    end
    total_energy_clip1 = 0.0
    for i in 1:(nnodes(dg)-1)
      total_energy_clip1 += modal[1, i]^2
    end
    total_energy_clip2 = 0.0
    for  i in 1:(nnodes(dg)-2)
      total_energy_clip2 += modal[1, i]^2
    end

    # Calculate energy in lower modes
    energy = max((total_energy - total_energy_clip1)/total_energy,
                 (total_energy_clip1 - total_energy_clip2)/total_energy_clip1)

    alpha[element_id] = 1 / (1 + exp(-parameter_s/threshold * (energy - threshold)))

    # Take care of the case close to pure DG
    if (alpha[element_id] < alpha_min)
      alpha[element_id] = 0.
    end

    # Take care of the case close to pure FV
    if (alpha[element_id] > 1-alpha_min)
      alpha[element_id] = 1.
    end

    # Clip the maximum amount of FV allowed
    alpha[element_id] = min(alpha_max, alpha[element_id])
  end

  if (do_smoothing)
    # Diffuse alpha values by setting each alpha to at least 50% of neighboring elements' alpha
    # Copy alpha values such that smoothing is indpedenent of the element access order
    alpha_pre_smooth .= alpha

    # Loop over interfaces
    for interface_id in 1:dg.n_interfaces
      # Get neighboring element ids
      left  = dg.interfaces.neighbor_ids[1, interface_id]
      right = dg.interfaces.neighbor_ids[2, interface_id]

      # Apply smoothing
      alpha[left]  = max(alpha_pre_smooth[left],  0.5 * alpha_pre_smooth[right], alpha[left])
      alpha[right] = max(alpha_pre_smooth[right], 0.5 * alpha_pre_smooth[left],  alpha[right])
    end
  end
end

# Convert conservative variables to indicator variable for discontinuities (elementwise version)
@inline function cons2indicator!(indicator, u, element_id, indicator_variable, dg::Dg1D)
  eqs = equations(dg)

  for i in 1:nnodes(dg)
    u_node = get_node_vars(u, dg, i, element_id)
    indicator[1, i] = indicator_variable(u_node, eqs)
  end
end

"""
    pure_and_blended_element_ids!(element_ids_dg, element_ids_dgfv, alpha, dg)

Given blending factors `alpha` and the solver `dg`, fill
`element_ids_dg` with the IDs of elements using a pure DG scheme and
`element_ids_dgfv` with the IDs of elements using a blended DG-FV scheme.
"""
function pure_and_blended_element_ids!(element_ids_dg, element_ids_dgfv, alpha, dg::Dg1D)
  empty!(element_ids_dg)
  empty!(element_ids_dgfv)

  for element_id in 1:dg.n_elements
    # Clip blending factor for values close to zero (-> pure DG)
    dg_only = isapprox(alpha[element_id], 0, atol=1e-12)
    if dg_only
      push!(element_ids_dg, element_id)
    else
      push!(element_ids_dgfv, element_id)
    end
  end
end


text:
abstract type AbstractNode end

"""
    Node
This data type holds the basic Node structure. The type T is used to specify the type of the data
stored in the node.
* If `nchild` is `0` the Node is a leaf node.
* If `root` is `False` the Node is a child of another node.
* `inc_length` specifies the length of the incomming branch.
* `binary` specifies the path from the root to the Node. `1` and `0` represent left and right turns respectively.
"""
mutable struct GeneralNode{R<:Real, I<:Integer} <: AbstractNode
    name::String
    mother::Union{GeneralNode{R,I}, Missing}
    children::Vector{GeneralNode{R,I}}
    nchild::I
    root::Bool
    inc_length::R
    binary::String
    num::I
    height::R
    IntExtMap::Vector{I}
    blv::Vector{R}
    stats::Dict{String, Float64}
end # struct Node

"""
    function Node()::GeneralNode{Float64, Int64}
This function will initialize an empty node.
"""
function Node()::GeneralNode{Float64, Int64}
        GeneralNode("no_name", missing, GeneralNode{Float64, Int64}[], 0, true, 1.0, "0",1, 1.0,Int64[], Float64[], Dict{String, Float64}())
end

function Node(name::String)::GeneralNode{Float64, Int64}
        GeneralNode(name, missing, GeneralNode{Float64, Int64}[], 0, true, 1.0, "0", 1, 1.0, Int64[], Float64[], Dict{String, Float64}())
end

function Node(name::String, inc_len::T)::GeneralNode{T, Int64} where T<:Real
    GeneralNode(name, missing, GeneralNode{T, Int64}[], 0, true, inc_len, "0", 1, 1.0, Int64[], Float64[], Dict{String, Float64}())
end

#################### Base functionality ####################

Base.:(==)(x::T, y::T) where T<:AbstractNode = x.num == y.num
Base.size(x::T) where T<:AbstractNode = size(post_order(x))
Base.length(x::T) where T<:AbstractNode = x.nchild

function Base.summary(io::IO, d::N) where N <: GeneralNode
    summary(io, d.name)
end

function Base.show(io::IO, d::N) where N <: GeneralNode
    print(io, "Tree with root:\n")
    show(io, d.name)
    print(io, "\nLength:\n")
    show(io, "text/plain", tree_length(d))
    print(io, "\nHeight:\n")
    show(io, "text/plain", tree_height(d))

end

function showall(io::IO, d::N) where N <: GeneralNode
  show(io, d)
  print(io, "\nNode:\n")
  show(io, "text/plain", d.name)
  print(io, "\n#children:\n")
  show(io, d.nchild)
  print(io, "\nbinary:\n")
  show(io, d.binary)
end


text:
import Knet.Ops21: linear
using Knet.KnetArrays: DevArray
using CUDA.CUBLAS: CUBLAS, cublasDgemm_v2, cublasSgemm_v2, CUBLAS_OP_T, CUBLAS_OP_N
using AutoGrad: AutoGrad, @primitive1

# op(A) m × k , op(B) k × n and C m × n; lda, ldb, ldc leading dimensions

for (T,F) in ((Float32,cublasSgemm_v2), (Float64,cublasDgemm_v2)); @eval begin

    function linear(A::DevArray{$T}, B::DevArray{$T}; dims=1)
        @assert ndims(A) > dims "ndims(A) must be > dims"
        ntuple(i->size(A,ndims(A)-dims+i),dims) === ntuple(i->size(B,i),dims) || throw(DimensionMismatch("linear: w=$(size(A)) x=$(size(B)) dims=$dims"))
        transa,transb = (CUBLAS_OP_N,CUBLAS_OP_N)
        msize = (size(A,i) for i in 1:ndims(A)-dims)
        ksize = (size(B,i) for i in 1:dims)
        nsize = (size(B,i) for i in dims+1:ndims(B))
        m,n,k = prod.((msize, nsize, ksize))
        lda,ldb,ldc = m,k,m
        C = similar(A, (msize..., nsize...))
        $F(CUBLAS.handle(), transa, transb, m, n, k, 1, A, lda, B, ldb, 0, C, ldc)
        return C
    end

    function _linear1(A::DevArray{$T}, B::DevArray{$T}; dims)  # A transposed
        @assert ntuple(i->size(A,i),dims) === ntuple(i->size(B,i),dims)
        transa,transb = (CUBLAS_OP_T,CUBLAS_OP_N)
        msize = (size(A,i) for i in dims+1:ndims(A))
        ksize = (size(B,i) for i in 1:dims)
        nsize = (size(B,i) for i in dims+1:ndims(B))
        m,n,k = prod.((msize, nsize, ksize))
        lda,ldb,ldc = k,k,m
        C = similar(A, (msize..., nsize...))
        $F(CUBLAS.handle(), transa, transb, m, n, k, 1, A, lda, B, ldb, 0, C, ldc)
        return C
    end

    function _linear2(A::DevArray{$T}, B::DevArray{$T}; dims)  # B transposed
        @assert ndims(A) > dims && ndims(B) > dims
        @assert ntuple(i->size(A,ndims(A)-dims+i),dims) === ntuple(i->size(B,ndims(B)-dims+i),dims)
        transa,transb = (CUBLAS_OP_N,CUBLAS_OP_T)
        msize = (size(A,i) for i in 1:ndims(A)-dims)
        ksize = (size(A,i) for i in 1+ndims(A)-dims:ndims(A))
        nsize = (size(B,i) for i in 1:ndims(B)-dims)
        m,n,k = prod.((msize, nsize, ksize))
        lda,ldb,ldc = m,n,m
        C = similar(A, (msize..., nsize...))
        $F(CUBLAS.handle(), transa, transb, m, n, k, 1, A, lda, B, ldb, 0, C, ldc)
        return C
    end

    function _linear3(A::DevArray{$T}, B::DevArray{$T}; dims)  # A,B transposed
        @assert ndims(B) > dims
        @assert ntuple(i->size(A,i),dims) === ntuple(i->size(B,ndims(B)-dims+i),dims)
        transa,transb = (CUBLAS_OP_T,CUBLAS_OP_T)
        msize = (size(A,i) for i in dims+1:ndims(A))
        ksize = (size(A,i) for i in 1:dims)
        nsize = (size(B,i) for i in 1:ndims(B)-dims)
        m,n,k = prod.((msize, nsize, ksize))
        lda,ldb,ldc = k,n,m
        C = similar(A, (msize..., nsize...))
        $F(CUBLAS.handle(), transa, transb, m, n, k, 1, A, lda, B, ldb, 0, C, ldc)
        return C
    end
end; end

@primitive1   linear(A::DevArray,B::DevArray;dims=1),dy  _linear2(dy,B; dims=ndims(B)-dims)  _linear1(A,dy; dims=ndims(A)-dims)
@primitive1 _linear1(A::DevArray,B::DevArray;dims=1),dy  _linear2(B,dy; dims=ndims(B)-dims)    linear(A,dy; dims=ndims(A)-dims)
@primitive1 _linear2(A::DevArray,B::DevArray;dims=1),dy    linear(dy,B; dims=ndims(B)-dims)  _linear1(dy,A; dims=ndims(A)-dims)
@primitive1 _linear3(A::DevArray,B::DevArray;dims=1),dy  _linear3(B,dy; dims=ndims(B)-dims)  _linear3(dy,A; dims=ndims(A)-dims)


text:
module AIBECS

using Reexport
@reexport using SparseArrays
using SuiteSparse
using LinearAlgebra
using Dates
using Printf
using ForwardDiff
@reexport using UnPack
using FieldMetadata
using Flatten
using DataFrames
using Distributions
using JLD2
using CodecZlib
using NCDatasets
@reexport using OceanGrids
@reexport using Unitful
using Unitful: °
using RecipesBase
using UnitfulRecipes
using Interpolations
using Distances
using NearestNeighbors
@reexport using MetadataArrays
using Shapefile
using Bijectors
using ImageFiltering
using SciMLBase

include("CirculationGeneration.jl")
include("SinkingParticles.jl")
include("newTypes.jl")
include("CTKsolvers.jl")
include("Parameters.jl")
include("overload_solve.jl")
include("multiTracer.jl")

# Available circulations
include("OCIM0.jl")
include("OCIM1.jl")
include("OCIM2.jl")
include("OCIM2_48L.jl")
include("OCCA.jl")
include("Archer_etal_2000.jl")
include("TwoBoxModel.jl")
include("Primeau_2x2x2.jl")

# AWESOME OCIM data
include("AO.jl") # TODO talk about it to Seth

# Aeolian source data
include("aeolian_sources.jl")

# River source data
include("Rivers.jl")

# Groundwater source data
include("GroundWaters.jl")

# ETOPO data
include("ETOPO.jl")

# Recipes for plotting
include("plot_recipes.jl")

# Diagnostics
include("diagnostics.jl")

end # module


text:
###############################################
#                                             #
# Main file for starting solving functions    #
# and generating figures,                     #
# parallelized with 'Distributed'             #
#                                             #
###############################################

# Load parallel computing package
using Distributed
using SharedArrays
using NLopt
using FiniteDifferences
using GridInterpolations
using JLD2
using Plots

# Include main functions
include("nopolicy.jl")
include("withpolicy_parallel.jl")
include("figures_paper.jl")

# Include helper functions
include("helperfunctions.jl")

# Tell Julia to use several processes for distributed computing
nprocs = 8;
addprocs(nprocs)

@everywhere begin
    # Load packages
    using Pkg; Pkg.activate(".") # Activate environment so that workers "see" packages
    using Distributed
    using SharedArrays
    using NLopt
    using FiniteDifferences
    using GridInterpolations

    # Include helper functions
    include("helperfunctions.jl")
end

# Solve model with standard parameters
@time noint    = nopolicy();                          # no intervention at all
@time notest   = withpolicy("notest", nprocs;);        # no testing
@time untarget = withpolicy("untargettest", nprocs;);  # untargeted testing
@time target   = withpolicy("targettest", nprocs;);    # targeted testing
@time isolate  = withpolicy("isolate", nprocs;);       # isolation

# Save important objects
#save("src/data/calib1_nointervention.jld2", noint);
#save("src/data/calib1_notest_small.jld2", notest);
#save("src/data/calib1_untarget_small.jld2", untarget);
#save("src/data/calib1_target_small.jld2", target);
#save("src/data/calib1_isolate_small.jld2", isolate);

# Load objects
#noint = load("src/data/calib1_nointervention.jld2");
#notest = load("src/data/calib1_notest_small.jld2");
#untarget = load("src/data/calib1_untarget_small.jld2");
#target = load("src/data/calib1_target_small.jld2");
#isolate = load("src/data/calib1_isolate_small.jld2");

# Figures (using freshly-generated data)
createFig5(noint, notest, untarget)
#savefig("src/figures/Fig5_dim20")
createFig6(notest, untarget, target)
#savefig("src/figures/Fig6_dim20")
createFig7(untarget, target, isolate)
#savefig("src/figures/Fig7_dim20")

# Figures (using stored data)
#createFig5()
#createFig6()
#createFig7()


text:
using Test
using NDTensors

@testset "Test non-exported tuple tools" begin

  @test NDTensors.diff((1, 3, 6, 4)) == (2, 3, -2)
  @test NDTensors.diff((1, 2, 3)) == (1, 1)

end

nothing


text:
###############################################################################
#
# Compute leaf inclination angle distribution?
#
###############################################################################
"""
    dcum(a::FT, b::FT, t::FT) where {FT<:AbstractFloat}

TODO Add function description
"""
function dcum(a::FT, b::FT, t::FT) where {FT<:AbstractFloat}
    if a >= 1
        f = 1 - cosd(t);
    else
        y    = 0;
        epsi = FT(1e-8);
        delx = 1;
        x    = 2 * deg2rad(t);
        p    = x;
        # add a iteration number (n) control to avoid presicion issues
        n    = 0;
        while (delx >= epsi) && (n<50)
            n   += 1;
            y    = a*sin(x) + b*sin(2x)/2;
            dx   = (y - x + p) / 2;
            x   += dx;
            delx = abs(dx);
        end
        f = (2*y + p) / FT(pi);
    end

    return f
end

"""
    dladgen(a::FT, b::FT, litab_bnd::Array{FT,2}) where {FT<:AbstractFloat}

TODO Calculate the freqency of WHAT?
"""
function dladgen(
            a::FT,
            b::FT,
            litab_bnd::Array{FT,2}
) where {FT<:AbstractFloat}
    @assert a + b <= 1;
    freq = similar(litab_bnd[:,1]);
    for i in eachindex(freq)
        freq[i] = dcum(a, b, litab_bnd[i,2]) - dcum(a, b, litab_bnd[i,1]);
    end

    return freq
end

max_stars_repo_name: bbuman/Land

max_issues_repo_name: bbuman/Land

max_forks_repo_name: bbuman/Land


text:
const DATA = readlines("day05.txt")

function input_to_int(row::AbstractString, pair1::Pair, pair2::Pair, pair3::Pair, pair4::Pair)
    row_bin = replace(replace(replace(replace(row, pair1), pair2), pair3), pair4)
    return parse(Int, row_bin, base = 2)
end

function compute(seat::AbstractString)
    return input_to_int(seat, 'R' => '1', 'L' => '0', 'F' => '0', 'B' => '1')
end

const seats = Set(compute(seat) for seat in DATA)

println("Part 1 : $(maximum(seats))")

part2 = setdiff(minimum(seats):maximum(seats), seats)

println("Part 2 : $(part2[1])")


text:
using OverconfidentTraders

ot = OverconfidentTraders

par  = ot.Pars()
beta = ot.findbeta(par)
bpar = ot.CstrBPar(beta, par)

a    = 1
ϵ    = 0

z    = a + ϵ/beta
p    = ot.price(z, par, bpar)


text:
# assuming
# (A d_uu + B d_u + C Id) f = -Source

function phid_eq_coeff!(ABCS::Vector, vars)
    (potential, u, Sd, phi, phi_u, phi_xx, phi_yy) = vars

    u2 = u * u
    u3 = u2 * u
    u4 = u2 * u2
    u5 = u3 * u2
    u6 = u4 * u2
    u9 = u6 * u3

    phi3 = phi * phi * phi
    phi4 = phi * phi3

    ABCS[1] = 0
    ABCS[2] = 9 * u
    ABCS[3] = 9//2

    ABCS[4] = -9//2 * phi - 27 * u4 * phi * Sd -
        4 * u6 * phi3 * UU(phi,potential) - u9 * phi4 * UUp(phi,potential) +
        3 * u2 * (phi_xx + phi_yy) -
        9//2 * u * phi_u - 9 * u5 * Sd * phi_u

    nothing
end


text:
using Interpolations: Flag, Quadratic, InPlace

export ItpSharedEV, set_up_dcdp_workers

struct ItpSharedEV{T,A1<:Interpolations.AbstractInterpolation{T},A2<:Interpolations.AbstractInterpolation{T},A3<:Interpolations.AbstractInterpolation{T},TT<:Tuple}
  EV::A1
  dEV::A2
  dEVσ::A3
  itypes::TT
end

function ItpSharedEV(σ::Real, ψflag::Flag=Quadratic(InPlace()) )
  global g_SharedEV
  global g_dcdp_primitives
  return ItpSharedEV(g_SharedEV, g_dcdp_primitives, σ, ψflag)
end

@GenGlobal g_ItpSharedEV

function ItpSharedEV(sev::SharedEV{T,N,N2,TT}, p::dcdp_primitives, σ::Real=1.0; ψflag::Flag=Quadratic(InPlace())) where {T,N,N2,TT}

    # dimensions of things
    ntyp = length(sev.itypes)
    nθ = _nθt(p)
    nS = _nS(p)
    nSexp1 = _nSexp(p)

    # form spline specifications
    splinetype(r::Union{StepRange{<:AbstractFloat},StepRangeLen{<:AbstractFloat}}) = BSpline(Quadratic(InPlace()))
    splinetype(r::AbstractVector) = NoInterp()
    splinetype(r::LinRange{<:AbstractFloat}) = BSpline(Constant())

    it_z     = splinetype.(p.zspace)
    it_itype = splinetype.(sev.itypes)

    it_EV   = (it_z..., BSpline(ψflag), NoInterp(),             it_itype...)
    it_dEV  = (it_z..., BSpline(ψflag), NoInterp(), NoInterp(), it_itype...)
    it_dEVσ = (it_z..., BSpline(ψflag), NoInterp(),             it_itype...)

    # form UNFILTERED spline objects.
    # These specify the SharedArray as the coef field and add type information about spline interpolation
    # if quadtratic, we will have to solve a system of linear equations later to interpolate
    itp_EV    = BSplineInterpolation(tweight(sev.EV  ), sev.EV  , it_EV  , OnCell(), Val{0}())
    itp_dEV   = BSplineInterpolation(tweight(sev.dEV ), sev.dEV , it_dEV , OnCell(), Val{0}())
    itp_dEVσ  = BSplineInterpolation(tweight(sev.dEVσ), sev.dEVσ, it_dEVσ, OnCell(), Val{0}())

    # information for how to scale the interpolation object
    scalegrid(x::AbstractRange{S}) where {S<:AbstractFloat} = StepRangeLen(x)
    scalegrid(x::Integer) = Base.OneTo(x)
    scalegrid(x::AbstractVector) = Base.OneTo(length(x))

    # scaled interpolation
    scl_EV   = scalegrid.((p.zspace..., _ψspace(p),         nS, sev.itypes...))
    scl_dEV  = scalegrid.((p.zspace..., _ψspace(p), nθ,     nS, sev.itypes...))
    scl_dEVσ = scalegrid.((p.zspace..., _ψspace(p),     nSexp1, sev.itypes...))

    sit_EV   = scale(itp_EV   , scl_EV...)
    sit_dEV  = scale(itp_dEV  , scl_dEV...)
    sit_dEVσ = scale(itp_dEVσ , scl_dEVσ...)
    # sit_dEVψ = Interpolations.scale(it_dEVψ , scl_dEVσ...)

    return ItpSharedEV{T,typeof(sit_EV),typeof(sit_dEV),typeof(sit_dEVσ),TT}(sit_EV, sit_dEV, sit_dEVσ, sev.itypes)
end

for typ in (:parallel, :serial)
    fun = Symbol("$(typ)_prefilterByView!")
    @eval begin
        function ($fun)(sev::SharedEV{T,N,N2}, isev::ItpSharedEV, dograd::Bool=true) where {T,N,N2}
            Ntyps = (length(sev.itypes[end]), )
            ($fun)(sev.EV, isev.EV.itp,         Ntyps...)
            if dograd
                ($fun)(sev.dEV, isev.dEV.itp,   Ntyps...)
                ($fun)(sev.dEVσ, isev.dEVσ.itp, Ntyps...)
            end
        end
    end
end

# function parallel_prefilterByView!(sev::SharedEV{T,N,N2}, isev::ItpSharedEV)
#     Ntyps = length.(sev.itypes)
#     parallel_prefilterByView!(sev.EV, isev.EV,     N- Ntyps+1:N...)
#     parallel_prefilterByView!(sev.dEV, isev.dEV,   N2-Ntyps+1:N...)
#     parallel_prefilterByView!(sev.dEVσ, isev.dEVσ, N- Ntyps+1:N...)
# end
#
# function parallel_prefilterByView!(sev::SharedEV{T,N,N2}, isev::ItpSharedEV)
#     Ntyps = length.(sev.itypes)
#     parallel_prefilterByView!(sev.EV, isev.EV,     N- Ntyps+1:N...)
#     parallel_prefilterByView!(sev.dEV, isev.dEV,   N2-Ntyps+1:N...)
#     parallel_prefilterByView!(sev.dEVσ, isev.dEVσ, N- Ntyps+1:N...)
# end

# TODO: MAKE SURE to check if data is within the indices

unsafegetityp(grid::AbstractArray, idx::Real) = idx
unsafegetityp(grid::Vector, idx::Integer) = grid[idx]

getitype(grid::Vector, idx::Integer) = unsafegetityp(grid, idx)

function getitype(grid::AbstractArray, idx::Real)
    minimum(grid) <= idx <= maximum(grid) || throw(DomainError(idx))
    unsafegetityp(grid, idx)
end

function unsafegetityp(sev::Union{SharedEV,ItpSharedEV}, itypidx::Tuple)
    length(sev.itypes) == length(itypidx) || throw(DimensionMismatch())
    return unsafegetityp.(sev.itypes, itypidx)
end

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

function set_up_dcdp_workers(pids::AbstractVector{<:Integer}, prim::dcdp_primitives, typegrids::AbstractVector...; σ0::Real=1.0, ψflag::Interpolations.Flag=Quadratic(InPlace()), kwargs...)

    println("initialize shared arrays")
    sev = SharedEV(pids, prim, typegrids...)

    println("setting up tmpvars")
    @eval @everywhere begin
        # Inner loop setup
        set_g_dcdp_primitives($prim)
        set_g_dcdp_tmpvars( dcdp_tmpvars(get_g_dcdp_primitives()) )
        set_g_SharedEV($sev)

        set_g_ItpSharedEV( ItpSharedEV(get_g_SharedEV(), get_g_dcdp_primitives(), $σ0; ψflag=$ψflag) )
    end

    println("worker dcdp problems setup")

    # add a few checks
    fetch(@spawn size(get_g_SharedEV().EV))   == size(sev.EV)               ||  throw(error("remote sharedEV not created (?)"))
    fetch(@spawn get_g_SharedEV().EV  === get_g_ItpSharedEV().EV.itp.coefs) ||  throw(error("remote sharedEV not linked to ItpSharedEV"))
    fetch(@spawn get_g_SharedEV().EV) === sev.EV                            ||  throw(error("remote sharedEV not same object as local"))

    return sev, get_g_ItpSharedEV()
end