text:
# Copyright 2021 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# This is a simplified version of primal-dual hybrid gradient implementation, of
# the package FirstOrderLp (see primal_dual_hybrid_gradient.jl), aimed to be
# used for numerical experiments.

import FirstOrderLp
import Printf
import LinearAlgebra
import Random
const norm = LinearAlgebra.norm
const dot = LinearAlgebra.dot
const randn = Random.randn

"""
Information for recording iterates on two dimensional subspace.
"""
struct TwoDimensionalSubspace
  """
  Subspace basis vector 1
  """
  basis_vector1::Vector{Float64}
  """
  Subspace basis vector 2, should be orthognal to vector 1
  """
  basis_vector2::Vector{Float64}
end

"""
A SimplePdhgParameters struct specifies the parameters for solving the saddle
point formulation of an problem using primal-dual hybrid gradient.

It solves a problem of the form (see quadratic_programmming.jl in FirstOrderLp)
minimize objective_vector' * x

s.t. constraint_matrix[1:num_equalities, :] * x =
     right_hand_side[1:num_equalities]

     constraint_matrix[(num_equalities + 1):end, :] * x >=
     right_hand_side[(num_equalities + 1):end, :]

     variable_lower_bound <= x <= variable_upper_bound

We use notation from Chambolle and Pock, "On the ergodic convergence rates of a
first-order primal-dual algorithm"
(http://www.optimization-online.org/DB_FILE/2014/09/4532.pdf).
That paper doesn't explicitly use the terminology "primal-dual hybrid gradient"
but their Theorem 1 is analyzing PDHG. In this file "Theorem 1" without further
reference refers to that paper.

Our problem is equivalent to the saddle point problem:
    min_x max_y L(x, y)
where
    L(x, y) = y' K x + g(x) - h*(y)
    K = -constraint_matrix
    g(x) = objective_vector' x if variable_lower_bound <= x <=
                                                            variable_upper_bound
                               otherwise infinity
    h*(y) = -right_hand_side' y if y[(num_equalities + 1):end] >= 0
                                otherwise infinity

Note that the places where g(x) and h*(y) are infinite effectively limits the
domain of the min and max. Therefore there's no infinity in the code.

We parametrize the primal and dual step sizes (tau and sigma in Chambolle and
Pock) as:
    primal_step_size = step_size / primal_weight
    dual_step_size = step_size * primal_weight.
The algoritm converges if
    primal_stepsize * dual_stepsize * norm(constraint_matrix)^2 < 1.
"""
struct PdhgParameters
  """
  Constant step size used in the algorithm. If nothing is specified, the solver
  computes a provably correct step size.
  """
  step_size::Union{Float64,Nothing}
  """
  Weight relating primal and dual step sizes.
  """
  primal_weight::Float64
  """
  Prints iteration stats at this frequency (in iterations) if verbose is true.
  """
  printing_frequency::Int64
  """
  If true a line of debugging info is printed every printing_frequency
  iterations. Otherwise some info is printed about the final solution.
  """
  verbosity::Bool
  """
  Number of loop iterations to run. Must be postive.
  """
  iteration_limit::Int64
  """
  A subspace where projections of iterates are recorded.
  If the value is nothing then the NaNs are recorded instead of a
  projection of the iterates.
  """
  two_dimensional_subspace::Union{TwoDimensionalSubspace,Nothing}
  """
  Accuracy to determine if a candidate point is a primal infeasibility
  certificate. See FirstOrderLp for documentation.
  """
  eps_primal_infeasible::Float64
  """
  Accuracy to determine if a candidate point is a dual infeasibility
  certificate. See FirstOrderLp for documentation.
  """
  eps_dual_infeasible::Float64
  """
  Determines whether or not to save the iterates of the run
  """
  store_iterates::Bool
  """
  If not Nothing, the solver will output three arrays containing the distance
  between this point and the normalized average, the normalized iterate, and the
  difference of iterates, respectively. The solver will include distances to
  this point in the returned PdhgStats.
  """
  reference_point::Union{Vector{Float64},Nothing}
end

"""
Statistics of the execution.
"""
struct PdhgStats
  """
  Primal objectives of the iterates; the ith entry corresponds to the primal
  objective of the ith iterate.
  """
  primal_objectives::Vector{Float64}
  """
  Primal objectives of the iterates; the ith entry corresponds to the dual
  objective of the ith iterate.
  """
  dual_objectives::Vector{Float64}
  """
  Primal norms of the iterates; the ith entry corresponds to the primal
  norm of the ith iterate.
  """
  primal_solution_norms::Vector{Float64}
  """
  Dual norms of the iterates; the ith entry corresponds to the dual
  norm of the ith iterate.
  """
  dual_solution_norms::Vector{Float64}
  """
  Primal delta norms of the iterates; the ith entry corresponds to the primal
  delta norm of the ith iterate.
  """
  primal_delta_norms::Vector{Float64}
  """
  Dual delta norms of the iterates; the ith entry corresponds to the dual
  delta norm of the ith iterate.
  """
  dual_delta_norms::Vector{Float64}
  """
  First coordinate of a subspace that the current iterates are projected onto.
  """
  current_subspace_coordinate1::Vector{Float64}
  """
  Second coordinate of a subspace that the current iterates are projected onto.
  """
  current_subspace_coordinate2::Vector{Float64}
  """
  First coordinate of a subspace that the average iterates are projected onto.
  """
  average_subspace_coordinate1::Vector{Float64}
  """
  Second coordinate of a subspace that the average iterates are projected onto.
  """
  average_subspace_coordinate2::Vector{Float64}
  """
  Euclidean distance between current_iterate/iteration_number and the
  input reference point at each iteration. If no reference point is given,
  this field is left as Nothing.
  """
  distance_to_normalized_current::Union{Vector{Float64},Nothing}
  """
  Euclidean distance between 2/(iteration_count + 1) * average_iterate and the
  input reference point at each iteration. If no reference point is given, this
  field is left as Nothing.
  """
  distance_to_normalized_average::Union{Vector{Float64},Nothing}
  """
  Euclidean distance between current_iterate - previous_iterate and the input
  reference point at each iteration. If no reference point is given, this field
  is left as Nothing.
  """
  distance_to_difference::Union{Vector{Float64},Nothing}
  """
  Maximum certificate error of the primal infeasibility certificate from
  current_iterate/iteration_number, scaled so that the dual ray objective is 1.
  If the dual ray objective is zero, this will be infinity, and if the dual ray
  objective is negative, this will be negative.
  """
  max_scaled_primal_certificate_error_of_normalized_current::Vector{Float64}
  """
  Maximum certificate error of the primal infeasibility certificate from
  2/(iteration_count + 1) * average_iterate, scaled so that the dual ray
  objective is 1. If the dual ray objective is zero, this will be infinity, and
  if the dual ray objective is negative, this will be negative.
  """
  max_scaled_primal_certificate_error_of_normalized_average::Vector{Float64}
  """
  Maximum certificate error of the primal infeasibility certificate from
  current_iterate - previous_iterate, scaled so that the dual ray objective
  is 1. If the dual ray objective is zero, this will be infinity, and if the
  dual ray objective is negative, this will be negative.
  """
  max_scaled_primal_certificate_error_of_difference::Vector{Float64}
end

"""
Output of the solver.
"""
struct PdhgOutput
  primal_solution::Vector{Float64}
  dual_solution::Vector{Float64}
  primal_delta::Vector{Float64}
  dual_delta::Vector{Float64}
  iteration_stats::PdhgStats
  """
  Records the first iteration in which a PointType detects infeasibility.
  """
  first_infeasibility_detection::Dict{FirstOrderLp.PointType,Int64}
  """
  Vector with all the primal iterates generated in the run.
  """
  primal_iterates::Union{Vector{Vector{Float64}},Nothing}
  """
  Vector with all the dual iterates generated in the run.
  """
  dual_iterates::Union{Vector{Vector{Float64}},Nothing}
  """
  The last iteration in which the active set of the primal and dual variables change.
  """
  last_active_set_change::Int64
end

"""
Computes statistics for the current iteration. The arguments primal_delta and
dual_delta correspond to the difference between the last two primal and dual
iterates, respectively.
"""
function compute_stats(
  problem::FirstOrderLp.QuadraticProgrammingProblem,
  iteration::Int64,
  two_dimensional_subspace::Union{TwoDimensionalSubspace,Nothing},
  stats::PdhgStats,
  primal_solution::Vector{Float64},
  dual_solution::Vector{Float64},
  primal_delta::Vector{Float64},
  dual_delta::Vector{Float64},
  solution_weighted_avg::FirstOrderLp.SolutionWeightedAverage,
)
  dual_stats =
    FirstOrderLp.compute_dual_stats(problem, primal_solution, dual_solution)
  stats.primal_objectives[iteration] =
    FirstOrderLp.primal_obj(problem, primal_solution)
  stats.dual_objectives[iteration] = dual_stats.dual_objective
  stats.primal_solution_norms[iteration] = norm(primal_solution)
  stats.dual_solution_norms[iteration] = norm(dual_solution)
  stats.primal_delta_norms[iteration] = norm(primal_delta)
  stats.dual_delta_norms[iteration] = norm(dual_delta)
  if nothing != two_dimensional_subspace
    z = [primal_solution; dual_solution]
    stats.current_subspace_coordinate1[iteration] =
      dot(z, two_dimensional_subspace.basis_vector1)
    stats.current_subspace_coordinate2[iteration] =
      dot(z, two_dimensional_subspace.basis_vector2)
    primal_avg, dual_avg = FirstOrderLp.compute_average(solution_weighted_avg)
    z_avg = [primal_avg; dual_avg]
    stats.average_subspace_coordinate1[iteration] =
      dot(z_avg, two_dimensional_subspace.basis_vector1)
    stats.average_subspace_coordinate2[iteration] =
      dot(z_avg, two_dimensional_subspace.basis_vector2)
  else
    stats.current_subspace_coordinate1[iteration] = NaN
    stats.current_subspace_coordinate2[iteration] = NaN
    stats.average_subspace_coordinate1[iteration] = NaN
    stats.average_subspace_coordinate2[iteration] = NaN
  end
end

function compute_infeasibility_distance_stats(
  iteration::Int64,
  stats::PdhgStats,
  reference_point::Vector{Float64},
  primal_solution::Vector{Float64},
  dual_solution::Vector{Float64},
  solution_weighted_avg::FirstOrderLp.SolutionWeightedAverage,
  primal_delta::Vector{Float64},
  dual_delta::Vector{Float64},
)
  primal_avg, dual_avg = FirstOrderLp.compute_average(solution_weighted_avg)
  stats.distance_to_normalized_current[iteration] =
    norm(reference_point - (1 / iteration) * [primal_solution; dual_solution])
  stats.distance_to_normalized_average[iteration] =
    norm(reference_point - (2 / (iteration + 1)) * [primal_avg; dual_avg])
  stats.distance_to_difference[iteration] =
    norm(reference_point - [primal_delta; dual_delta])
end

"""
Logging while the algorithm is running.
"""
function pdhg_log(
  problem::FirstOrderLp.QuadraticProgrammingProblem,
  iteration::Int64,
  stats::PdhgStats,
)

  Printf.@printf(
    "   %5d objectives=(%9g, %9g) norms=(%9g, %9g) res_norm=(%9g, %9g)\n",
    iteration,
    stats.primal_objectives[iteration],
    stats.dual_objectives[iteration],
    stats.primal_solution_norms[iteration],
    stats.dual_solution_norms[iteration],
    stats.primal_delta_norms[iteration],
    stats.dual_delta_norms[iteration],
  )
end

"""
Prints infeasibility information
"""
function infeasibility_log(
  iteration::Int64,
  infeasibility_information::FirstOrderLp.InfeasibilityInformation,
  point_type::FirstOrderLp.PointType,
)
  Printf.@printf(
    "   %5d %s max_dual: %9g dual_obj: %9g max_primal: %9g primal_obj %9g \n",
    iteration,
    string(point_type),
    infeasibility_information.max_dual_ray_infeasibility,
    infeasibility_information.dual_ray_objective,
    infeasibility_information.max_primal_ray_infeasibility,
    infeasibility_information.primal_ray_linear_objective,
  )
end

"""
Checks if the normalized current iterate, normalized average iterate, and
iterate difference provide (approximate) certificates of infeasibility.

# Returns
A dictionary from point type to the maximum infeasibility of the corresponding
scaled dual ray, scaled so that the dual ray objective is 1.0. This will be
infinity if the dual ray objective is 0.0, and negative is the dual ray
objective is negative.
"""
function check_infeasibility(
  problem::FirstOrderLp.QuadraticProgrammingProblem,
  iteration::Int64,
  first_infeasibility_detection::Dict{FirstOrderLp.PointType,Int64},
  primal_solution::Vector{Float64},
  primal_delta::Vector{Float64},
  dual_solution::Vector{Float64},
  dual_delta::Vector{Float64},
  solution_weighted_avg::FirstOrderLp.SolutionWeightedAverage,
  eps_primal_infeasible::Float64,
  eps_dual_infeasible::Float64,
  params::PdhgParameters,
)
  scaled_infeasibility = Dict{FirstOrderLp.PointType,Float64}()
  for point_type in instances(FirstOrderLp.PointType)
    primal_candidate = nothing
    dual_candidate = nothing
    if point_type == FirstOrderLp.POINT_TYPE_AVERAGE_ITERATE
      primal_candidate, dual_candidate =
        FirstOrderLp.compute_average(solution_weighted_avg)
    elseif point_type == FirstOrderLp.POINT_TYPE_CURRENT_ITERATE
      primal_candidate = primal_solution
      dual_candidate = dual_solution
    elseif point_type == FirstOrderLp.POINT_TYPE_ITERATE_DIFFERENCE
      primal_candidate = primal_delta
      dual_candidate = dual_delta
    else
      continue
    end
    infeasibility_information = FirstOrderLp.compute_infeasibility_information(
      problem,
      primal_candidate,
      dual_candidate,
      point_type,
    )
    scaled_infeasibility[point_type] =
      infeasibility_information.max_dual_ray_infeasibility /
      infeasibility_information.dual_ray_objective
    if params.verbosity && mod(iteration - 1, params.printing_frequency) == 0
      infeasibility_log(iteration, infeasibility_information, point_type)
    end
    if haskey(first_infeasibility_detection, point_type)
      continue
    end
    if FirstOrderLp.primal_infeasibility_criteria_met(
      eps_primal_infeasible,
      infeasibility_information,
    ) || FirstOrderLp.dual_infeasibility_criteria_met(
      eps_dual_infeasible,
      infeasibility_information,
    )
      infeasibility_log(iteration, infeasibility_information, point_type)
      first_infeasibility_detection[point_type] = iteration
      Printf.@printf(
        "   %5d Infeasibility detected using %s\n",
        iteration,
        string(point_type)
      )
    end
  end
  return scaled_infeasibility
end

"""
Counts the changes of primal active constraints for a given primal vector.
"""
function count_primal_active_constraint_changes(
  primal::Vector{Float64},
  next_primal::Vector{Float64},
  problem::FirstOrderLp.QuadraticProgrammingProblem,
)
  counter = 0
  for idx in 1:length(primal)
    counter +=
      (
        (primal[idx] == problem.variable_upper_bound[idx]) !=
        (next_primal[idx] == problem.variable_upper_bound[idx])
      ) || (
        (primal[idx] == problem.variable_lower_bound[idx]) !=
        (next_primal[idx] == problem.variable_lower_bound[idx])
      )
  end
  return counter
end

""" Counts the number of dual active constraints for a given dual vector."""
function count_dual_active_constraint_changes(
  dual::Vector{Float64},
  next_dual::Vector{Float64},
  problem::FirstOrderLp.QuadraticProgrammingProblem,
)
  counter = 0
  for idx in FirstOrderLp.inequality_range(problem)
    dual_active = 0
    next_dual_active = 0
    if dual[idx] == 0.0
      dual_active = 1
    end
    if next_dual[idx] == 0.0
      next_dual_active = 1
    end
    counter += abs(dual_active - next_dual_active)
  end
  return counter
end

function take_pdhg_step(
  problem::FirstOrderLp.QuadraticProgrammingProblem,
  current_primal_solution::Vector{Float64},
  current_dual_solution::Vector{Float64},
  primal_weight::Float64,
  step_size::Float64,
)
  # The next lines compute the primal portion of the PDHG algorithm:
  # argmin_x [g(x) + current_dual_solution' K x
  #          + (0.5 * primal_weight / step_size)
  #             * norm(x - current_primal_solution)^2]
  # See Sections 2-3 of Chambolle and Pock and the comment above
  # SimplePdhgParameters.
  # This minimization is easy to do in closed form since it can be separated
  # into independent problems for each of the primal variables. The
  # projection onto the primal feasibility set comes from the closed form
  # for the above minimization and the cases where g(x) is infinite - there
  # isn't officially any projection step in the algorithm.
  primal_gradient = FirstOrderLp.compute_primal_gradient(
    problem,
    current_primal_solution,
    current_dual_solution,
  )
  next_primal =
    current_primal_solution .- primal_gradient * (step_size / primal_weight)
  FirstOrderLp.project_primal!(next_primal, problem)

  # The next two lines compute the dual portion:
  # argmin_y [H*(y) - y' K (2.0*next_primal - current_primal_solution)
  #           + (0.5 / (primal_weight * step_size))
  #.             * norm(y-current_dual_solution)^2]
  dual_gradient = FirstOrderLp.compute_dual_gradient(
    problem,
    2.0 * next_primal - current_primal_solution,
  )
  next_dual =
    current_dual_solution .+ dual_gradient * (step_size * primal_weight)
  FirstOrderLp.project_dual!(next_dual, problem)

  return next_primal, next_dual
end

"""
`optimize(params::PdhgParameters,
          problem::QuadraticProgrammingProblem)`

Solves a linear program using primal-dual hybrid gradient. If the step_size
specified in params is negative, picks a step size that ensures
step_size^2 * norm(constraint_matrix)^2 < 1,
a condition that guarantees provable convergence.

# Arguments
- `params::PdhgParameters`: parameters.
- `original_problem::QuadraticProgrammingProblem`: the QP to solve.

# Returns
A SaddlePointOutput struct containing the solution found.
"""
function optimize(
  params::PdhgParameters,
  problem::FirstOrderLp.QuadraticProgrammingProblem,
  initial_primal_solution::Vector{Float64},
  initial_dual_solution::Vector{Float64},
)
  primal_size = length(problem.variable_lower_bound)
  dual_size = length(problem.right_hand_side)

  primal_weight = params.primal_weight
  reference_point = params.reference_point

  if isnothing(params.step_size)
    desired_relative_error = 0.2
    maximum_singular_value, number_of_power_iterations =
      FirstOrderLp.estimate_maximum_singular_value(
        problem.constraint_matrix,
        probability_of_failure = 0.001,
        desired_relative_error = desired_relative_error,
      )
    step_size = (1 - desired_relative_error) / maximum_singular_value
  else
    step_size = params.step_size
  end
  Printf.@printf("Step size %9g Primal weight %9g\n", step_size, primal_weight)

  solution_weighted_avg =
    FirstOrderLp.initialize_solution_weighted_average(primal_size, dual_size)
  # Difference between current_primal_solution and last primal iterate, i.e.
  # x_k - x_{k-1}.
  primal_delta = zeros(primal_size)
  # Difference between current_dual_solution and last dual iterate, i.e. y_k -
  # y_{k-1}
  dual_delta = zeros(dual_size)
  iterations_completed = 0
  iteration_limit = params.iteration_limit
  printing_frequency = params.printing_frequency
  first_infeasibility_detection = Dict{FirstOrderLp.PointType,Int64}()

  stats = PdhgStats(
    Vector{Float64}(undef, iteration_limit),
    Vector{Float64}(undef, iteration_limit),
    Vector{Float64}(undef, iteration_limit),
    Vector{Float64}(undef, iteration_limit),
    Vector{Float64}(undef, iteration_limit),
    Vector{Float64}(undef, iteration_limit),
    Vector{Float64}(undef, iteration_limit),
    Vector{Float64}(undef, iteration_limit),
    Vector{Float64}(undef, iteration_limit),
    Vector{Float64}(undef, iteration_limit),
    (
      isnothing(reference_point) ? nothing :
      Vector{Float64}(undef, iteration_limit)
    ), # distance_to_normalized_current
    (
      isnothing(reference_point) ? nothing :
      Vector{Float64}(undef, iteration_limit)
    ), # distance_to_normalized_average
    (
      isnothing(reference_point) ? nothing :
      Vector{Float64}(undef, iteration_limit)
    ), # distance_to_difference
    Vector{Float64}(undef, iteration_limit),
    Vector{Float64}(undef, iteration_limit),
    Vector{Float64}(undef, iteration_limit),
  )
  primal_iterates = nothing
  dual_iterates = nothing
  if params.store_iterates
    primal_iterates = Vector{Vector{Float64}}(undef, iteration_limit)
    dual_iterates = Vector{Vector{Float64}}(undef, iteration_limit)
  end

  current_primal_solution = initial_primal_solution
  current_dual_solution = initial_dual_solution
  last_active_set_change = 0
  iteration = 0
  while iteration < iteration_limit
    iteration += 1

    if params.store_iterates
      primal_iterates[iteration] = current_primal_solution
      dual_iterates[iteration] = current_dual_solution
    end
    time_spent_doing_basic_algorithm_checkpoint = time()

    next_primal, next_dual = take_pdhg_step(
      problem,
      current_primal_solution,
      current_dual_solution,
      primal_weight,
      step_size,
    )

    # Update deltas and iterates
    primal_delta = next_primal .- current_primal_solution
    dual_delta = next_dual .- current_dual_solution

    FirstOrderLp.add_to_solution_weighted_average(
      solution_weighted_avg,
      next_primal,
      next_dual,
      1.0,
    )
    # Compute stats and log.
    compute_stats(
      problem,
      iteration,
      params.two_dimensional_subspace,
      stats,
      current_primal_solution,
      current_dual_solution,
      primal_delta,
      dual_delta,
      solution_weighted_avg,
    )

    # Check infeasibility. This doesn't stop the run, just reports detection.
    scaled_infeasibility = check_infeasibility(
      problem,
      iteration,
      first_infeasibility_detection,
      current_primal_solution,
      primal_delta,
      current_dual_solution,
      dual_delta,
      solution_weighted_avg,
      params.eps_primal_infeasible,
      params.eps_dual_infeasible,
      params,
    )

    stats.max_scaled_primal_certificate_error_of_normalized_current[iteration] =
      get(scaled_infeasibility, FirstOrderLp.POINT_TYPE_CURRENT_ITERATE, Inf)
    stats.max_scaled_primal_certificate_error_of_normalized_average[iteration] =
      get(scaled_infeasibility, FirstOrderLp.POINT_TYPE_AVERAGE_ITERATE, Inf)
    stats.max_scaled_primal_certificate_error_of_difference[iteration] =
      get(scaled_infeasibility, FirstOrderLp.POINT_TYPE_ITERATE_DIFFERENCE, Inf)

    if !isnothing(reference_point)
      compute_infeasibility_distance_stats(
        iteration,
        stats,
        reference_point,
        current_primal_solution,
        current_dual_solution,
        solution_weighted_avg,
        primal_delta,
        dual_delta,
      )
    end

    if count_primal_active_constraint_changes(
      current_primal_solution,
      next_primal,
      problem,
    ) + count_dual_active_constraint_changes(
      current_dual_solution,
      next_dual,
      problem,
    ) > 0
      last_active_set_change = iteration
    end
    if params.verbosity && mod(iteration - 1, printing_frequency) == 0
      pdhg_log(problem, iteration, stats)
    end

    # update iterates
    current_primal_solution = next_primal
    current_dual_solution = next_dual

    iterations_completed += 1
  end
  return PdhgOutput(
    current_primal_solution,
    current_dual_solution,
    primal_delta,
    dual_delta,
    stats,
    first_infeasibility_detection,
    primal_iterates,
    dual_iterates,
    last_active_set_change,
  )
end

function optimize(
  params::PdhgParameters,
  problem::FirstOrderLp.QuadraticProgrammingProblem,
)
  current_primal_solution = zeros(length(problem.variable_lower_bound))
  current_dual_solution = zeros(length(problem.right_hand_side))

  return optimize(
    params,
    problem,
    current_primal_solution,
    current_dual_solution,
  )
end


text:
function handle_macro(@nospecialize(x), state) end
function handle_macro(x::EXPR, state)
    !is_macro_call(x) && return
    if is_macroname(x[1])
        state(x[1])
        if _points_to_Base_macro(x[1], :deprecate, state) && length(x) == 3
            if bindingof(x[2]) !== nothing
                return
            elseif CSTParser.is_func_call(x[2])
                # add deprecated method
                # add deprecated function binding and args in new scope
                mark_binding!(x[2], x)
                mark_sig_args!(x[2])
                s0 = state.scope # store previous scope
                state.scope = Scope(s0, x, Dict(), nothing, nothing)
                setscope!(x, state.scope) # tag new scope to generating expression
                state(x[2])
                state(x[3])
                state.scope = s0
            elseif isidentifier(x[2])
                mark_binding!(x[2], x)
            end
        elseif _points_to_Base_macro(x[1], :deprecate_binding, state) && length(x) == 3 && isidentifier(x[2]) && isidentifier(x[3])
            setref!(x[2], refof(x[3]))
        elseif _points_to_Base_macro(x[1], :eval, state) && length(x) == 2 && state isa Toplevel
            # Create scope around eval'ed expression. This ensures anybindings are
            # correctly hoisted to the top-level scope.
            setscope!(x, Scope(x))
            setparent!(scopeof(x), state.scope)
            s0 = state.scope
            state.scope = scopeof(x)
            interpret_eval(x[2], state)
            state.scope = s0
        elseif _points_to_Base_macro(x[1], :irrational, state) && length(x) == 4
            mark_binding!(x[2], x)
        elseif _points_to_Base_macro(x[1], :enum, state)
            for i = 2:length(x)
                if !ispunctuation(x[i])
                    if bindingof(x[i]) !== nothing
                        break
                    end
                    if i == 3 && typof(x[i]) === CSTParser.Begin
                        for j in 1:length(x[i][2])
                            mark_binding!(x[i][2][j], x)
                        end
                        break
                    end
                    mark_binding!(x[i], x)
                end
            end
        elseif _points_to_Base_macro(x[1], :goto, state)
            if length(x) == 2 && isidentifier(x[2])
                setref!(x[2], Binding(noname, nothing, nothing, EXPR[], nothing, nothing))
            end
        elseif _points_to_Base_macro(x[1], :label, state)
            if length(x) == 2 && isidentifier(x[2])
                mark_binding!(x[2])
            end
        elseif is_nospecialize(x[1])
            for i = 2:length(x)
                if !ispunctuation(x[i])
                    if bindingof(x[i]) !== nothing
                        break
                    end
                    mark_binding!(x[i], x)
                end
            end
        elseif _points_to_arbitrary_macro(x[1], :Turing, :model, state) && length(x) == 2 &&
            is_binary_call(x[2], CSTParser.Tokens.EQ) &&
            typof(x[2][3]) === CSTParser.Begin && length(x[2][3]) == 3 && typof(x[2][3][2]) === CSTParser.Block
            for i = 1:length(x[2][3][2])
                ex = x[2][3][2][i]
                if is_binary_call(ex, CSTParser.Tokens.APPROX)
                    mark_binding!(ex)
                end
            end
        elseif _points_to_arbitrary_macro(x[1], :JuMP, :variable, state)
            if length(x) < 3
                return
            elseif length(x) >= 5 && ispunctuation(x[2])
                _mark_JuMP_binding(x[5])
            else
                _mark_JuMP_binding(x[3])
            end
        elseif (_points_to_arbitrary_macro(x[1], :JuMP, :expression, state) ||
            _points_to_arbitrary_macro(x[1], :JuMP, :NLexpression, state) ||
            _points_to_arbitrary_macro(x[1], :JuMP, :constraint, state) || _points_to_arbitrary_macro(x[1], :JuMP, :NLconstraint, state)) && length(x) > 1
            if ispunctuation(x[2])
                if length(x) == 8
                    _mark_JuMP_binding(x[5])
                end
            else
                if length(x) == 4
                    _mark_JuMP_binding(x[3])
                end
            end
        end
    elseif typof(x[1]) == CSTParser.GlobalRefDoc && length(x) == 3 && CSTParser.isliteral(x[2]) && kindof(x[2]) === CSTParser.Tokens.TRIPLE_STRING && isidentifier(x[3])
        mark_binding!(x[3])
        setref!(x[3], bindingof(x[3]))
    end
end

function is_nospecialize(x)
    return length(x) == 2 &&
        isidentifier(x[2]) &&
        valofid(x[2]) == "nospecialize"
end

function _rem_ref(x::EXPR)
    if typof(x) === CSTParser.Ref && length(x) > 0
        return x[1]
    end
    return x
end

function _mark_JuMP_binding(arg)
    if isidentifier(arg) || typof(arg) === CSTParser.Ref
        mark_binding!(_rem_ref(arg))
    elseif is_binary_call(arg, CSTParser.Tokens.EQEQ) || is_binary_call(arg, CSTParser.Tokens.LESS_EQ)  || is_binary_call(arg, CSTParser.Tokens.GREATER_EQ)
        if isidentifier(arg[1]) || typof(arg[1]) === CSTParser.Ref
            mark_binding!(_rem_ref(arg[1]))
        else
            mark_binding!(_rem_ref(arg[3]))
        end
    elseif typof(arg) === CSTParser.Comparison && length(arg) == 5
        mark_binding!(_rem_ref(arg[3]))
    end
end

function _points_to_Base_macro(x::EXPR, name, state)
    is_getfield_w_quotenode(x) && return _points_to_Base_macro(x[3][1], name, state)
    targetmacro =  maybe_lookup(getsymbolserver(state.server)[:Base][Symbol("@", name)], state.server)
    length(x) == 2 && isidentifier(x[2]) && Symbol(valofid(x[2])) == name && (ref = refof(x[2])) !== nothing &&
    (ref == targetmacro || (ref isa Binding && ref.val == targetmacro))
end

function _points_to_arbitrary_macro(x::EXPR, module_name, name, state)
    length(x) == 2 && isidentifier(x[2]) && valof(x[2]) == name && haskey(getsymbolserver(state.server), Symbol(module_name)) && haskey(getsymbolserver(state.server)[Symbol(module_name)], Symbol("@", name)) && (refof(x[2]) == maybe_lookup(getsymbolserver(state.server)[Symbol(module_name)][Symbol("@", name)], state.server) ||
    (refof(x[2]) isa Binding && refof(x[2]).val == maybe_lookup(getsymbolserver(state.server)[Symbol(module_name)][Symbol("@", name)], state.server)))
end

maybe_lookup(x, server) = x isa SymbolServer.VarRef ? SymbolServer._lookup(x, getsymbolserver(server), true) : x

function maybe_eventually_get_id(x::EXPR)
    if isidentifier(x)
        return x
    elseif is_invis_brackets(x)
        return maybe_eventually_get_id(x[2])
    end
    return nothing
end

is_eventually_interpolated(x::EXPR) = is_invis_brackets(x) ? is_eventually_interpolated(x[2]) : is_unary_call(x) && CSTParser.is_exor(x[1])
isquoted(x::EXPR) = typof(x) === CSTParser.Quotenode && length(x) == 2 && kindof(x[1]) === CSTParser.Tokens.COLON
maybeget_quotedsymbol(x::EXPR) = isquoted(x) ? maybe_eventually_get_id(x[2]) : nothing

function is_loop_iterator(x::EXPR)
    CSTParser.is_range(x) &&
    ((parentof(x) isa EXPR && typof(parentof(x)) === CSTParser.For) ||
    (parentof(x) isa EXPR && parentof(parentof(x)) isa EXPR && typof(parentof(parentof(x))) === CSTParser.For))
end

"""
    maybe_quoted_list(x::EXPR)

Try and get a list of quoted symbols from x. Return nothing if not possible.
"""
function maybe_quoted_list(x::EXPR)
    names = EXPR[]
    if typof(x) === CSTParser.Vect || typof(x) === CSTParser.TupleH
        for i = 1:length(x)
            CSTParser.ispunctuation(x[i]) && continue
            name = maybeget_quotedsymbol(x[i])
            if name !== nothing
                push!(names, name)
            else
                return nothing
            end
        end
        return names
    end
end

"""
interpret_eval(x::EXPR, state)

Naive attempt to interpret `x` as though it has been eval'ed. Lifts
any bindings made within the scope of `x` to the toplevel and replaces
(some) interpolated binding names with the value where possible.
"""
function interpret_eval(x::EXPR, state)
    # make sure we have bindings etc
    state(x)
    tls = retrieve_toplevel_scope(x)
    for ex in collect_expr_with_bindings(x)
        b = bindingof(ex)
        if isidentifier(b.name)
            # The name of the binding is fixed
            add_binding(ex, state, tls)
        elseif CSTParser.is_exor(b.name)
            # The name of the binding is variable, we need to work out what the
            # interpolated symbol points to.
            variable_name = parentof(b.name)[2]
            resolve_ref(variable_name, state.scope, state)
            if (ref = refof(variable_name)) isa Binding
                if is_assignment(ref.val) && (rhs = maybeget_quotedsymbol(ref.val[3])) !== nothing
                    # `name = :something`
                    toplevel_binding = Binding(rhs, b.val, b.type, [], nothing, nothing)
                    if scopehasbinding(tls, valofid(toplevel_binding.name))
                        toplevel_binding.prev = tls.names[valofid(toplevel_binding.name)]
                        tls.names[valofid(toplevel_binding.name)] = toplevel_binding
                        toplevel_binding.prev.next = toplevel_binding
                    else
                        tls.names[valofid(toplevel_binding.name)] = toplevel_binding
                    end
                elseif is_loop_iterator(ref.val) && (names = maybe_quoted_list(ref.val[3])) !== nothing
                    # name is of a collection of quoted symbols
                    for name in names
                        toplevel_binding = Binding(name, b.val, b.type, [], nothing, nothing)
                        if scopehasbinding(tls, valofid(toplevel_binding.name))
                            toplevel_binding.prev = tls.names[valofid(toplevel_binding.name)]
                            tls.names[valofid(toplevel_binding.name)] = toplevel_binding
                            toplevel_binding.prev.next = toplevel_binding
                        else
                            tls.names[valofid(toplevel_binding.name)] = toplevel_binding
                        end
                    end
                end
            end
        end
    end
end

function collect_expr_with_bindings(x, bound_exprs=EXPR[])
    if hasbinding(x)
        push!(bound_exprs, x)
        # Assuming here that if an expression has a binding we don't want anything bound to chlid nodes.
    elseif !((CSTParser.defines_function(x) && !(is_eventually_interpolated(x[1]))) || CSTParser.defines_macro(x) || typof(x) === CSTParser.Export)
        for a in x
            collect_expr_with_bindings(a, bound_exprs)
        end
    end
    return bound_exprs
end

ext: jl


text:
"A `ProductGrid` represents the tensor product of other grids."
abstract type ProductGrid{T,N} <: CompositeGrid{T,N}
end

size(g::ProductGrid) = map(length, g.grids)
size(g::ProductGrid, j::Int) = length(g.grids[j])

ProductGrid(grids...) = VcatGrid(grids...)
ProductGrid(grids::AbstractGrid1d...) = FlatProductGrid(grids...)

ProductGrid{T}(grids...) where {N,S,T <: SVector{N,S}} = VcatGrid{N,S}(grids...)
ProductGrid{T}(grids::AbstractGrid1d...) where {N,S,T <: SVector{N,S}} =
	FlatProductGrid{N,S}(grids...)

similargrid(grid::ProductGrid, ::Type{T}, dims...) where T = error()#ProductGrid([similargrid(g, eltype(T), dims[i]) for (i,g) in enumerate(components(grid))]...)
resize(grid::ProductGrid, n) = ProductGrid(map(resize, components(grid), n)...)

covering(g::ProductGrid) = productdomain(map(covering, components(g))...)

isperiodic(g::ProductGrid) = mapreduce(isperiodic, &, components(g))

factors(g::ProductGrid) = components(g)

unsafe_grid_getindex(g::ProductGrid{T,N}, I::Vararg{Int,N}) where {T,N} =
	convert(T, map(getindex, components(g), I))

productgrid() = ()
productgrid(d) = d
productgrid(d1, d2, d3...) = productgrid(productgrid(d1, d2), d3...)

productgrid(d1, d2) = productgrid1(d1, d2)
productgrid1(d1, d2) = productgrid2(d1, d2)
productgrid2(d1, d2) = ProductGrid(d1, d2)

productgrid(d1::ProductGrid, d2::ProductGrid) =
	ProductGrid(components(d1)..., components(d2)...)
productgrid1(d1::ProductGrid, d2) = ProductGrid(components(d1)..., d2)
productgrid2(d1, d2::ProductGrid) = ProductGrid(d1, components(d2)...)

cross(grids::AbstractGrid...) = productgrid(grids...)
^(grid::AbstractGrid, n::Int) = productgrid(ntuple(i->grid, n)...)

canonicalgrid(g::ProductGrid) = ProductGrid(map(canonicalgrid, components(g)))

mapto_canonical(g::ProductGrid) = ProductMap(map(mapto_canonical, components(g)))
mapfrom_canonical(g::ProductGrid) = ProductMap(map(mapfrom_canonical, components(g)))

Display.combinationsymbol(d::ProductGrid) = Display.Times()
Display.displaystencil(d::ProductGrid) = composite_displaystencil(d)
Base.show(io::IO, mime::MIME"text/plain", d::ProductGrid) = composite_show(io, mime, d)
Base.show(io::IO, d::ProductGrid) = composite_show_compact(io, d)

"A `FlatProductGrid` is a product grid of `N` 1-D grids."
struct FlatProductGrid{N,T,GG} <: ProductGrid{SVector{N,T},N}
	grids	::	GG
end

FlatProductGrid(grids::Vararg{AbstractGrid,N}) where {N} =
	FlatProductGrid{N}(grids...)
function FlatProductGrid{N}(grids::Vararg{AbstractGrid,N}) where {N}
	T = mapreduce(numtype, promote_type, grids)
	FlatProductGrid{N,T}(grids...)
end
FlatProductGrid{N,T}(grids::Vararg{Any,N}) where {N,T} =
	FlatProductGrid{N,T,typeof(grids)}(grids)

"""
A `VcatGrid` is a product grid that concatenates the elements of all composing
grids into a flat vector, like `vcat` does.
"""
struct VcatGrid{N,T,M,DIMS,GG} <: ProductGrid{SVector{N,T},M}
	grids	::	GG
end

function VcatGrid(grids::Vararg{AbstractGrid{<:Any,<:Any},M}) where {M}
	T = mapreduce(numtype, promote_type, grids)
	DIMS = map(dimension, grids)
	N = sum(DIMS)
	VcatGrid{N,T,M,DIMS}(grids...)
end
function VcatGrid{N,T}(grids::Vararg{AbstractGrid{<:Any,<:Any},M}) where {N,T,M}
	DIMS = map(dimension, grids)
	VcatGrid{N,T,M,DIMS}(grids...)
end

VcatGrid{N,T,M,DIMS}(grids...) where {N,T,M,DIMS} =
	VcatGrid{N,T,M,DIMS,typeof(grids)}(grids)

unsafe_grid_getindex(g::VcatGrid{N,T,M,DIMS}, I::Vararg{Int,M}) where {N,T,M,DIMS} =
	DomainSets.convert_tocartesian(map(getindex, components(g), I), Val(DIMS))


text:
##############################################################################
#
# Split string into tokens on whitespace
#
##############################################################################

tokenize(lang::S, s::T) where {S <: Language, T <: AbstractString} = WordTokenizers.tokenize(s)

sentence_tokenize(lang::S, s::T) where {S <: Language, T<:AbstractString} = WordTokenizers.split_sentences(s)


text:
module MappingVariableMethods

using NQCDynamics:
    NQCDynamics,
    AbstractSimulation,
    Simulation,
    RingPolymerSimulation,
    DynamicsMethods,
    DynamicsUtils,
    Calculators,
    Estimators,
    NonadiabaticDistributions,
    TimeCorrelationFunctions,
    ndofs
using NQCModels: NQCModels, Model
using NQCBase: Atoms
using .DynamicsUtils: get_mapping_momenta, get_mapping_positions

include("nrpmd.jl")
export NRPMD

include("cmm.jl")
export eCMM
include("rpecmm.jl")

end # module


text:
insts = OrderedDict()
insts["minimal"] = [
    ((1, 2),),
    ((2, 1),),
    ]
insts["fast"] = [
    ((2, 3),),
    ((3, 3),),
    ]
insts["various"] = [
    ((2, 2),),
    ((4, 2),),
    ((6, 2),),
    ((2, 4),),
    ((3, 4),),
    ((4, 4),),
    ]
return (RelEntrEntanglementJuMP, insts)


text:
import DataStructures
import AVLTrees
import StaticArrays
using BenchmarkTools
using Profile

include("../src/soalight.jl")
include("../src/allocator.jl")
include("../src/tree_iter.jl")

function testbase1(n, p)
    tree = Dict{Int, typeof(p)}()
    for j in 1:n[]
        tree[j] = p
        tree[-j] = p
    end
    @assert !haskey(tree, 0)
    for j in 1:n[]
        @assert tree[j] === p
        @assert tree[-j] === p
    end
    @assert !haskey(tree, n[] + 1)
    @assert !haskey(tree, -n[] - 1)
    for j in 1:n[]
        @assert haskey(tree, j)
        @assert haskey(tree, -j)
        pop!(tree, j)
        pop!(tree, -j)
        @assert !haskey(tree, j)
        @assert !haskey(tree, -j)
    end
end

function testbase2(n, p)
    tree = DataStructures.SortedDict{Int, typeof(p)}()
    for j in 1:n[]
        tree[j] = p
        tree[-j] = p
    end
    @assert !haskey(tree, 0)
    for j in 1:n[]
        @assert tree[j] === p
        @assert tree[-j] === p
    end
    @assert !haskey(tree, n[] + 1)
    @assert !haskey(tree, -n[] - 1)
    for j in 1:n[]
        @assert haskey(tree, j)
        @assert haskey(tree, -j)
        pop!(tree, j)
        pop!(tree, -j)
        @assert !haskey(tree, j)
        @assert !haskey(tree, -j)
    end
end

# For later use
#= function testbase3(n, p)
    tree = AVLTrees.AVLDict{Int, typeof(p)}()
    for j in 1:n[]
        tree[j] = p
        tree[-j] = p
    end
    @assert !(0 in tree)
    for j in 1:n[]
        @assert tree[j] === p
        @assert tree[-j] === p
    end
    @assert !(n[] + 1 in tree)
    @assert !(-n[] - 1 in tree)
    for j in 1:n[]
        @assert j in tree
        @assert -j in tree
        delete!(tree, j)
        delete!(tree, -j)
        @assert !(j in tree)
        @assert !(-j in tree)
    end
end =#

function testbase4(n, p)
    tree = nil()
    for j in 1:n[]
        tree = insert(tree, j, p)
        tree = insert(tree, -j, p)
    end
    @assert !haskey(tree, 0)
    for j in 1:n[]
        @assert tree[j] === p
        @assert tree[-j] === p
    end
    @assert !haskey(tree, n[] + 1)
    @assert !haskey(tree, -n[] - 1)
    for j in 1:n[]
        @assert haskey(tree, j)
        @assert haskey(tree, -j)
        tree = delete(tree, j)
        tree = delete(tree, -j)
        @assert !haskey(tree, j)
        @assert !haskey(tree, -j)
    end
end

function testalloc(n, p, alloc)
    tree = nil(alloc)
    for j in 1:n[]
        tree = insert(tree, j, p, alloc)
        tree = insert(tree, -j, p, alloc)
    end
    @assert !haskey(tree, 0, alloc)
    for j in 1:n[]
        @assert getindex(tree, j, alloc) === p
        @assert getindex(tree, -j, alloc) === p
    end
    @assert !haskey(tree, n[] + 1, alloc)
    @assert !haskey(tree, -n[] - 1, alloc)
    for j in 1:n[]
        @assert haskey(tree, j, alloc)
        @assert haskey(tree, -j, alloc)
        tree = delete(tree, j, alloc)
        tree = delete(tree, -j, alloc)
        @assert !haskey(tree, j, alloc)
        @assert !haskey(tree, -j, alloc)
    end
    emptyend!(alloc)
end

function testfree(n, p, alloc)
    tree = nil(alloc)
    for j in 1:n[]
        tree = insert(tree, j, p, alloc)
        tree = insert(tree, -j, p, alloc)
    end
    @assert !haskey(tree, 0, alloc)
    for j in 1:n[]
        @assert getindex(tree, j, alloc) === p
        @assert getindex(tree, -j, alloc) === p
    end
    @assert !haskey(tree, n[] + 1, alloc)
    @assert !haskey(tree, -n[] - 1, alloc)
    for j in 1:n[]
        @assert haskey(tree, j, alloc)
        @assert haskey(tree, -j, alloc)
        tree = delete(tree, j, alloc)
        tree = delete(tree, -j, alloc)
        @assert !haskey(tree, j, alloc)
        @assert !haskey(tree, -j, alloc)
    end
    @assert isempty(alloc)
end

ans = nothing
for n in [100000, 10]
    println(n, " elements")

    for p in [1, 10, 100]
        println(" payload of ", p, " float(s)")

        Payload = StaticArrays.SVector{p ,Float64}

        GC.gc()

        print("  Dict                                   ")
        @btime ($ans = testbase1(Ref($n), zeros($Payload)))

        GC.gc()

        print("  DataStructures.SortedDict              ")
        @btime ($ans = testbase2(Ref($n), zeros($Payload)))

        #= GC.gc()

        print("  AVLTrees.AVLDict                       ")
        @btime ($ans = testbase3(Ref($n), zeros($Payload))) =#

        GC.gc()

        print("  Tree without allocator                 ")
        @btime ($ans = testbase4(Ref($n), zeros($Payload)))

        GC.gc()

        print("  with fixed allocator                   ")
        alloc = Allocator{TreeNode{Int, Int, Payload}, Int}(2 * n)
        @btime ($ans = testalloc(Ref($n), zeros($Payload), $alloc))

        GC.gc()

        print("  with resizable allocator               ")
        alloc = Allocator{TreeNode{Int, Int, Payload}, Int}(nothing)
        @btime ($ans = testalloc(Ref($n), zeros($Payload), $alloc))

        GC.gc()

        print("  with fixed free list allocator         ")
        alloc = FreeListAllocator{TreeNode{Int, Int, Payload}, Int}(2 * n)
        @btime ($ans = testfree(Ref($n), zeros($Payload), $alloc))

        GC.gc()

        print("  with resizable free list allocator     ")
        alloc = FreeListAllocator{TreeNode{Int, Int, Payload}, Int}(nothing)
        @btime ($ans = testfree(Ref($n), zeros($Payload), $alloc))

        GC.gc()

        print("  with fixed SOA allocator               ")
        store = TupleVector{TreeNode{Int, Int, Payload}}(2 * n)
        alloc = SOAllocator{TreeNode{Int, Int, Payload}, Int, typeof(store)}(store, 2 * n)
        @btime ($ans = testalloc(Ref($n), zeros($Payload), $alloc))

        GC.gc()

        print("  with resizable SOA allocator           ")
        store = TupleVector{TreeNode{Int, Int, Payload}}(N)
        alloc = SOAllocator{TreeNode{Int, Int, Payload}, Int, typeof(store)}(store, nothing)
        @btime ($ans = testalloc(Ref($n), zeros($Payload), $alloc))

        GC.gc()

        print("  with fixed free list SOA allocator     ")
        store = TupleVector{TreeNode{Int, Int, Payload}}(2 * n)
        alloc = FreeListSOAllocator{TreeNode{Int, Int, Payload}, Int, typeof(store)}(store, 2 * n)
        @btime ($ans = testfree(Ref($n), zeros($Payload), $alloc))

        GC.gc()

        print("  with resizable free list SOA allocator ")
        store = TupleVector{TreeNode{Int, Int, Payload}}(N)
        alloc = FreeListSOAllocator{TreeNode{Int, Int, Payload}, Int, typeof(store)}(store, nothing)
        @btime ($ans = testfree(Ref($n), zeros($Payload), $alloc))
    end
end