text:
using Revise
using AdFem
using ADCME
using Random
using ProgressMeter
using PyPlot
using DelimitedFiles
using Distributions

np = pyimport("numpy")
n = 10
m = 15
h = 0.1

xy = Dict(
    1=>(1.95,2.05,0.32,0.38),
    2=>(1.8,2.2,0.2,0.45),
    3=>(1.7,2.2,0.18,0.5),
    4=>(1.5,2.3,0.15,0.5),
    5=>(1.2,3.0,0.0,0.5)
)
for (ii,σ0) in enumerate([0.01 0.05 0.1 0.2 0.5])
xmin,xmax, ymin,ymax = xy[ii]

Random.seed!(233)
obs_id = rand(interior_node("all", m, n, h), 10)
obs_id = [obs_id; obs_id .+ (m+1)*(n+1)]
function solve_elasticity(c)
    c = constant(c)
    E, ν = c[1], c[2]
    H = 1e6*E/(1+ν)/(1-2ν)*tensor([
             1-ν ν 0.0
             ν 1-ν 0.0
             0.0 0.0 (1-2ν)/2
         ])
    K = compute_fem_stiffness_matrix(H, m, n, h)
    bdedge = bcedge("right", m, n, h)
    bdnode = bcnode("left", m, n, h)
    K, _ = fem_impose_Dirichlet_boundary_condition(K, bdnode, m, n, h)
    F = compute_fem_traction_term([zeros(4*m*n) 1e6*ones(4*m*n)], bdedge, m, n, h)
    sol = K\F
    sol = sol[obs_id]
    return sol
end

sol = solve_elasticity([2.0;0.35])
sess = Session(); init(sess)
sol_ = run(sess, sol)
# # visualize_displacement(sol_*1e3, m, n, h)

c = Variable(0.1*rand(2))
sol = solve_elasticity(c)

loss = sum((sol-sol_)^2)
sess = Session(); init(sess)
BFGS!(sess, loss)

y = sol_
hs = run(sess, sol)
H = run(sess, gradients(sol, c))
R = 0.1^2*diagm(0=>ones(length(sol_)))
s = run(sess, c)
Q = diagm(0=>ones(2))

μ, V = uqnlin(y, hs, H, R, s, Q)
d = Distributions.MvNormal(μ, (V+V')/2)

x0 = LinRange(xmin, xmax, 100)
y0 = LinRange(ymin, ymax, 100)
X, Y = np.meshgrid(x0, y0)
Z = hcat([[pdf(d, [X[i, j];Y[i, j]]) for i = 1:100] for j = 1:100]...)
mcmc = readdlm("simdata/sigma$σ0-100000.txt")

close("all")
contour(X, Y, Z, levels=4)
scatter(mcmc[:,1], mcmc[:,2], s=2, c="red", label="RW-MCMC", alpha=0.5)
legend()
xlabel("E")
ylabel("\$\\nu\$")
xlim(xmin, xmax)
ylim(ymin, ymax)
savefig("sigma$σ0.png")
end

max_stars_repo_name: ADCMEMarket/ADCMEImages

max_issues_repo_name: ADCMEMarket/ADCMEImages

max_forks_repo_name: ADCMEMarket/ADCMEImages


text:
using Plots, Test
pgfplotsx()

function create_plot(args...; kwargs...)
   pgfx_plot = plot(args...; kwargs...)
   return pgfx_plot, repr("application/x-tex", pgfx_plot)
end

function create_plot!(args...; kwargs...)
   pgfx_plot = plot!(args...; kwargs...)
   return pgfx_plot, repr("application/x-tex", pgfx_plot)
end

@testset "PGFPlotsX" begin
   pgfx_plot = plot(1:5)
   Plots._update_plot_object(pgfx_plot)
   @test pgfx_plot.o.the_plot isa PGFPlotsX.TikzDocument
   @test pgfx_plot.series_list[1].plotattributes[:quiver] === nothing
   axis = Plots.pgfx_axes(pgfx_plot.o)[1]
   @test count(x -> x isa PGFPlotsX.Plot, axis.contents) == 1
   @test !haskey(axis.contents[1].options.dict, "fill")

   @testset "Legends" begin
      legends_plot =  plot( rand(5,2), lab = ["1" ""] )
      scatter!(legends_plot, rand(5) )
      Plots._update_plot_object(legends_plot)
      axis_contents = Plots.pgfx_axes(legends_plot.o)[1].contents
      leg_entries = filter( x -> x isa PGFPlotsX.LegendEntry, axis_contents )
      series = filter( x -> x isa PGFPlotsX.Plot, axis_contents )
      @test length(leg_entries) == 2
      @test !haskey(series[1].options.dict, "forget plot")
      @test haskey(series[2].options.dict, "forget plot")
      @test !haskey(series[3].options.dict, "forget plot")
   end # testset

   @testset "3D docs example" begin
      n = 100
      ts = range(0, stop = 8π, length = n)
      x = ts .* map(cos, ts)
      y = (0.1ts) .* map(sin, ts)
      z = 1:n
      pl = plot(
         x,
         y,
         z,
         zcolor = reverse(z),
         m = (10, 0.8, :blues, Plots.stroke(0)),
         leg = false,
         cbar = true,
         w = 5,
      )
      pgfx_plot = plot!(pl, zeros(n), zeros(n), 1:n, w = 10)
      Plots._update_plot_object(pgfx_plot)
      if @test_nowarn(
         haskey(Plots.pgfx_axes(pgfx_plot.o)[1].options.dict, "colorbar") == true
      )
         @test Plots.pgfx_axes(pgfx_plot.o)[1]["colorbar"] === nothing
      end
   end # testset
   @testset "Color docs example" begin
      y = rand(100)
      plot(
         0:10:100,
         rand(11, 4),
         lab = "lines",
         w = 3,
         palette = :grays,
         fill = 0,
         α = 0.6,
      )
      pl = scatter!(
         y,
         zcolor = abs.(y .- 0.5),
         m = (:hot, 0.8, Plots.stroke(1, :green)),
         ms = 10 * abs.(y .- 0.5) .+ 4,
         lab = ["grad", "", "ient"],
      )
      Plots._update_plot_object(pl)
      axis = Plots.pgfx_axes(pl.o)[1]
      @test count(x -> x isa PGFPlotsX.LegendEntry, axis.contents) == 6
      @test count(x -> x isa PGFPlotsX.Plot, axis.contents) == 108 # each marker is its own plot, fillranges create 2 plot-objects
      marker = axis.contents[15]
      @test marker isa PGFPlotsX.Plot
      @test marker.options["mark"] == "*"
      @test marker.options["mark options"]["color"] ==
            RGBA{Float64}(colorant"green", 0.8)
      @test marker.options["mark options"]["line width"] == 0.75 # 1px is 0.75pt
   end # testset
   @testset "Plot in pieces" begin
      pic = plot(rand(100) / 3, reg = true, fill = (0, :green))
      scatter!(pic, rand(100), markersize = 6, c = :orange)
      Plots._update_plot_object(pic)
      axis_contents = Plots.pgfx_axes(pic.o)[1].contents
      leg_entries = filter( x -> x isa PGFPlotsX.LegendEntry, axis_contents )
      series = filter( x -> x isa PGFPlotsX.Plot, axis_contents )
      @test length(leg_entries) == 2
      @test length(series) == 4
      @test haskey(series[1].options.dict, "forget plot")
      @test !haskey(series[2].options.dict, "forget plot")
      @test haskey(series[3].options.dict, "forget plot")
      @test !haskey(series[4].options.dict, "forget plot")
   end # testset
   @testset "Marker types" begin
      markers = filter((m -> begin
         m in Plots.supported_markers()
      end), Plots._shape_keys)
      markers = reshape(markers, 1, length(markers))
      n = length(markers)
      x = (range(0, stop = 10, length = n + 2))[2:(end - 1)]
      y = repeat(reshape(reverse(x), 1, :), n, 1)
      scatter(
         x,
         y,
         m = (8, :auto),
         lab = map(string, markers),
         bg = :linen,
         xlim = (0, 10),
         ylim = (0, 10),
      )
   end # testset
   @testset "Layout" begin
      plot(
         Plots.fakedata(100, 10),
         layout = 4,
         palette = [:grays :blues :hot :rainbow],
         bg_inside = [:orange :pink :darkblue :black],
      )
   end # testset
   @testset "Polar plots" begin
      Θ = range(0, stop = 1.5π, length = 100)
      r = abs.(0.1 * randn(100) + sin.(3Θ))
      plot(Θ, r, proj = :polar, m = 2)
   end # testset
   @testset "Drawing shapes" begin
      verts = [
         (-1.0, 1.0),
         (-1.28, 0.6),
         (-0.2, -1.4),
         (0.2, -1.4),
         (1.28, 0.6),
         (1.0, 1.0),
         (-1.0, 1.0),
         (-0.2, -0.6),
         (0.0, -0.2),
         (-0.4, 0.6),
         (1.28, 0.6),
         (0.2, -1.4),
         (-0.2, -1.4),
         (0.6, 0.2),
         (-0.2, 0.2),
         (0.0, -0.2),
         (0.2, 0.2),
         (-0.2, -0.6),
      ]
      x = 0.1:0.2:0.9
      y = 0.7 * rand(5) .+ 0.15
      plot(
         x,
         y,
         line = (3, :dash, :lightblue),
         marker = (Shape(verts), 30, RGBA(0, 0, 0, 0.2)),
         bg = :pink,
         fg = :darkblue,
         xlim = (0, 1),
         ylim = (0, 1),
         leg = false,
      )
   end # testset
   @testset "Histogram 2D" begin
      histogram2d(randn(10000), randn(10000), nbins = 20)
   end # testset
   @testset "Heatmap-like" begin
      xs = [string("x", i) for i = 1:10]
      ys = [string("y", i) for i = 1:4]
      z = float((1:4) * reshape(1:10, 1, :))
      pgfx_plot = heatmap(xs, ys, z, aspect_ratio = 1)
      Plots._update_plot_object(pgfx_plot)
      if @test_nowarn(
         haskey(Plots.pgfx_axes(pgfx_plot.o)[1].options.dict, "colorbar") == true
      )
         @test Plots.pgfx_axes(pgfx_plot.o)[1]["colorbar"] === nothing
         @test Plots.pgfx_axes(pgfx_plot.o)[1]["colormap name"] == "plots1"
      end

      pgfx_plot = wireframe(xs, ys, z, aspect_ratio = 1)
      # TODO: clims are wrong
   end # testset
   @testset "Contours" begin
      x = 1:0.5:20
      y = 1:0.5:10
      f(x, y) = begin
         (3x + y^2) * abs(sin(x) + cos(y))
      end
      X = repeat(reshape(x, 1, :), length(y), 1)
      Y = repeat(y, 1, length(x))
      Z = map(f, X, Y)
      p2 = contour(x, y, Z)
      p1 = contour(x, y, f, fill = true)
      plot(p1, p2)
      # TODO: colorbar for filled contours
   end # testset
   @testset "Varying colors" begin
      t = range(0, stop = 1, length = 100)
      θ = (6π) .* t
      x = t .* cos.(θ)
      y = t .* sin.(θ)
      p1 = plot(x, y, line_z = t, linewidth = 3, legend = false)
      p2 = scatter(
         x,
         y,
         marker_z = ((x, y) -> begin
            x + y
         end),
         color = :bwr,
         legend = false,
      )
      plot(p1, p2)
   end # testset
   @testset "Framestyles" begin
      scatter(
         fill(randn(10), 6),
         fill(randn(10), 6),
         framestyle = [:box :semi :origin :zerolines :grid :none],
         title = [":box" ":semi" ":origin" ":zerolines" ":grid" ":none"],
         color = permutedims(1:6),
         layout = 6,
         label = "",
         markerstrokewidth = 0,
         ticks = -2:2,
      )
      # TODO: support :semi
   end # testset
   @testset "Quiver" begin
      x = (-2pi):0.2:(2 * pi)
      y = sin.(x)

      u = ones(length(x))
      v = cos.(x)
      arrow_plot = plot(x, y, quiver = (u, v), arrow = true)
      # TODO: could adjust limits to fit arrows if too long, but how?
      # TODO: get latex available on CI
      # mktempdir() do path
      #    @test_nowarn savefig(arrow_plot, path*"arrow.pdf")
      # end
   end # testset
   @testset "Annotations" begin
      y = rand(10)
      pgfx_plot = plot(
         y,
         annotations = (3, y[3], Plots.text("this is \\#3", :left)),
         leg = false,
      )
      Plots._update_plot_object(pgfx_plot)
      axis_content = Plots.pgfx_axes(pgfx_plot.o)[1].contents
      nodes = filter(x -> !isa(x, PGFPlotsX.Plot), axis_content)
      @test length(nodes) == 1
      mktempdir() do path
         file_path =joinpath(path,"annotations.tex")
         @test_nowarn savefig(pgfx_plot, file_path)
         open(file_path) do io
            lines = readlines(io)
            @test count(s -> occursin("node", s), lines) == 1
         end
      end
      annotate!([
         (5, y[5], Plots.text("this is \\#5", 16, :red, :center)),
         (10, y[10], Plots.text("this is \\#10", :right, 20, "courier")),
      ])
      Plots._update_plot_object(pgfx_plot)
      axis_content = Plots.pgfx_axes(pgfx_plot.o)[1].contents
      nodes = filter(x -> !isa(x, PGFPlotsX.Plot), axis_content)
      @test length(nodes) == 3
      mktempdir() do path
         file_path =joinpath(path,"annotations.tex")
         @test_nowarn savefig(pgfx_plot, file_path)
         open(file_path) do io
            lines = readlines(io)
            @test count(s -> occursin("node", s), lines) == 3
         end
      end
      annotation_plot = scatter!(
         range(2, stop = 8, length = 6),
         rand(6),
         marker = (50, 0.2, :orange),
         series_annotations = [
            "series",
            "annotations",
            "map",
            "to",
            "series",
            Plots.text("data", :green),
         ],
      )
      Plots._update_plot_object(annotation_plot)
      axis_content = Plots.pgfx_axes(annotation_plot.o)[1].contents
      nodes = filter(x -> !isa(x, PGFPlotsX.Plot), axis_content)
      @test length(nodes) == 9
      mktempdir() do path
         file_path =joinpath(path,"annotations.tex")
         @test_nowarn savefig(annotation_plot, file_path)
         open(file_path) do io
            lines = readlines(io)
            @test count(s -> occursin("node", s), lines) == 9
         end
      end
   end # testset
   @testset "Ribbon" begin
      aa = rand(10)
      bb = rand(10)
      cc = rand(10)
      conf = [aa - cc bb - cc]
      ribbon_plot =
         plot(collect(1:10), fill(1, 10), ribbon = (conf[:, 1], conf[:, 2]))
      Plots._update_plot_object(ribbon_plot)
      axis = Plots.pgfx_axes(ribbon_plot.o)[1]
      plots = filter(x -> x isa PGFPlotsX.Plot, axis.contents)
      @test length(plots) == 4
      @test !haskey(plots[1].options.dict, "fill")
      @test !haskey(plots[2].options.dict, "fill")
      @test !haskey(plots[3].options.dict, "fill")
      @test haskey(plots[4].options.dict, "fill")
      @test ribbon_plot.o !== nothing
      @test ribbon_plot.o.the_plot !== nothing
      #    mktempdir() do path
      #       @test_nowarn savefig(ribbon_plot, path*"ribbon.svg")
      #    end
   end # testset
end # testset

@testset "Extra kwargs" begin
   pl = plot(1:5, test = "me")
   @test pl[1][1].plotattributes[:extra_kwargs][:test] == "me"
   pl = plot(1:5, test = "me", extra_kwargs = :subplot)
   @test pl[1].attr[:extra_kwargs][:test] == "me"
   pl = plot(1:5, test = "me", extra_kwargs = :plot)
   @test pl.attr[:extra_plot_kwargs][:test] == "me"
   pl = plot(1:5, extra_kwargs = Dict(:plot => Dict(:test => "me"), :series => Dict(:and => "me too")))
   @test pl.attr[:extra_plot_kwargs][:test] == "me"
   @test pl[1][1].plotattributes[:extra_kwargs][:and] == "me too"
   pl = plot(
     plot(1:5, title="Line"),
     scatter(1:5, title="Scatter", extra_kwargs=Dict(:subplot=>Dict("axis line shift" => "10pt")))
   )
   Plots._update_plot_object(pl)
   axes = Plots.pgfx_axes(pl.o)
   @test !haskey(axes[1].options.dict, "axis line shift")
   @test haskey(axes[2].options.dict, "axis line shift")
end # testset

ext: jl


text:
include("asap_cluster.jl")
include("asap_colmem.jl")

# Base components for architecture modeling.
#include("components.jl")
#
## Architecture Model Files
#include("generic.jl")
##include("asap3_hex.jl")
#
## Post Creation Transforms
#include("transforms.jl")


text:
using Documenter, TrueSkill

makedocs(
    modules = [TrueSkill],
    format = Documenter.HTML(),
    checkdocs = :exports,
    sitename = "TrueSkill.jl",
    pages = Any["index.md"]
)

deploydocs(
    repo = "github.com/glandfried/TrueSkill.jl.git",
)


text:
# Copyright 2021 The FirstOrderLp Authors
#
# 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 file is for utility code used by both mirror_prox.jl and
# primal_dual_hybrid_gradient.jl .

"""
A solution computed by saddle point mirror prox or primal-dual hybrid
gradient (PDHG).
"""
struct SaddlePointOutput
  """
  The output primal solution vector.
  """
  primal_solution::Vector{Float64}

  """
  The output dual solution vector.
  """
  dual_solution::Vector{Float64}

  """
  One of the possible values from the TerminationReason enum.
  """
  termination_reason::TerminationReason

  """
  Extra information about the termination reason (may be empty).
  """
  termination_string::String

  """
  The total number of algorithmic iterations for the solve.
  """
  iteration_count::Int32

  """
  Detailed statistics about a subset of the iterations. The collection frequency
  is defined by algorithm parameters.
  """
  iteration_stats::Vector{IterationStats}
end

function unscaled_saddle_point_output(
  scaled_problem::ScaledQpProblem,
  primal_solution::AbstractVector{Float64},
  dual_solution::AbstractVector{Float64},
  termination_reason::TerminationReason,
  iterations_completed::Int64,
  iteration_stats::Vector{IterationStats},
)

  # Unscale iterates.
  original_primal_solution =
    primal_solution ./ scaled_problem.variable_rescaling
  original_dual_solution = dual_solution ./ scaled_problem.constraint_rescaling

  return SaddlePointOutput(
    original_primal_solution,
    original_dual_solution,
    termination_reason,
    termination_reason_to_string(termination_reason),
    iterations_completed,
    iteration_stats,
  )
end

"""
Projects the given point onto a set of bounds.
"""
function projection!(
  primal::AbstractVector{Float64},
  variable_lower_bound::AbstractVector{Float64},
  variable_upper_bound::AbstractVector{Float64},
)
  for idx in 1:length(primal)
    primal[idx] = min(
      variable_upper_bound[idx],
      max(variable_lower_bound[idx], primal[idx]),
    )
  end
end

"""Projects the given primal solution onto the feasible set. That is, all
negative duals for inequality constraints are set to zero."""
function project_primal!(
  primal::AbstractVector{Float64},
  problem::QuadraticProgrammingProblem,
)
  projection!(
    primal,
    problem.variable_lower_bound,
    problem.variable_upper_bound,
  )
end

""" Projects the given dual solution onto the feasible set. That is, all
negative duals for inequality constraints are set to zero."""
function project_dual!(
  dual::AbstractVector{Float64},
  problem::QuadraticProgrammingProblem,
)
  for idx in inequality_range(problem)
    dual[idx] = max(dual[idx], 0.0)
  end
end

"""The weighted l2 norm"""
function weighted_norm(
  vec::AbstractVector{Float64},
  weights::AbstractVector{Float64},
)
  sum = 0.0
  for i in eachindex(vec)
    sum += weights[i] * vec[i] * vec[i]
  end
  return sqrt(sum)
end

"""
Computes the local duality gap.
"""
function compute_localized_duality_gap(
  problem::QuadraticProgrammingProblem,
  primal_solution::AbstractVector{Float64},
  dual_solution::AbstractVector{Float64},
  primal_norm_params::AbstractVector{Float64},
  dual_norm_params::AbstractVector{Float64},
  distance_to_optimality::Float64,
  norm::LocalizedDualityGapNorm,
  use_approximate_localized_duality_gap::Bool,
)
  local_duality_gap = bound_optimal_objective(
    problem,
    primal_solution,
    dual_solution,
    primal_norm_params,
    dual_norm_params,
    distance_to_optimality,
    norm,
    solve_approximately = use_approximate_localized_duality_gap,
  )

  return local_duality_gap
end

mutable struct RestartInfo
  """
  The primal_solution recorded at the last restart point.
  """
  primal_solution::Vector{Float64}

  """
  The dual_solution recorded at the last restart point.
  """
  dual_solution::Vector{Float64}

  """
  Localized duality gap at last restart. This has a value of nothing if no
  restart has occurred.
  """
  last_restart_localized_duality_gap::Union{Nothing,OptimalObjectiveBoundResult}

  """
  The length of the last restart interval.
  """
  last_restart_length::Int64

  """
  The primal distance moved from the restart point two restarts ago and the
  average of the iterates across the last restart.
  """
  primal_distance_moved_last_restart_period::Float64

  """
  The dual distance moved from the restart point two restarts ago and the
  average of the iterates across the last restart.
  """
  dual_distance_moved_last_restart_period::Float64

  """
  Reduction in the potential function that was achieved last time we tried to do
  a restart.
  """
  gap_reduction_ratio_last_trial::Float64
end

function create_last_restart_info(
  problem::QuadraticProgrammingProblem,
  primal_solution::AbstractVector,
  dual_solution::AbstractVector,
)
  return RestartInfo(
    copy(primal_solution),
    copy(dual_solution),
    nothing,
    1,
    0.0,
    0.0,
    1.0,
  )
end

mutable struct SolutionWeightedAverage
  sum_primal_solutions::Vector{Float64}
  sum_dual_solutions::Vector{Float64}
  sum_primal_solutions_count::Int64
  sum_dual_solutions_count::Int64
  sum_primal_solution_weights::Float64
  sum_dual_solution_weights::Float64
end

function initialize_solution_weighted_average(
  primal_size::Int64,
  dual_size::Int64,
)
  return SolutionWeightedAverage(
    zeros(primal_size),
    zeros(dual_size),
    0,
    0,
    0.0,
    0.0,
  )
end

function reset_solution_weighted_average(
  solution_weighted_avg::SolutionWeightedAverage,
)
  solution_weighted_avg.sum_primal_solutions =
    zeros(length(solution_weighted_avg.sum_primal_solutions))
  solution_weighted_avg.sum_dual_solutions =
    zeros(length(solution_weighted_avg.sum_dual_solutions))
  solution_weighted_avg.sum_primal_solutions_count = 0
  solution_weighted_avg.sum_dual_solutions_count = 0
  solution_weighted_avg.sum_primal_solution_weights = 0.0
  solution_weighted_avg.sum_dual_solution_weights = 0.0
  return
end

function add_to_primal_solution_weighted_average(
  solution_weighted_avg::SolutionWeightedAverage,
  current_primal_solution::AbstractVector{Float64},
  weight::Float64,
)
  @assert solution_weighted_avg.sum_primal_solutions_count >= 0
  solution_weighted_avg.sum_primal_solutions .+=
    current_primal_solution * weight
  solution_weighted_avg.sum_primal_solutions_count += 1
  solution_weighted_avg.sum_primal_solution_weights += weight
  return
end

function add_to_dual_solution_weighted_average(
  solution_weighted_avg::SolutionWeightedAverage,
  current_dual_solution::AbstractVector{Float64},
  weight::Float64,
)
  @assert solution_weighted_avg.sum_dual_solutions_count >= 0
  solution_weighted_avg.sum_dual_solutions .+= current_dual_solution * weight
  solution_weighted_avg.sum_dual_solutions_count += 1
  solution_weighted_avg.sum_dual_solution_weights += weight
  return
end

function add_to_solution_weighted_average(
  solution_weighted_avg::SolutionWeightedAverage,
  current_primal_solution::AbstractVector{Float64},
  current_dual_solution::AbstractVector{Float64},
  weight::Float64,
)
  add_to_primal_solution_weighted_average(
    solution_weighted_avg,
    current_primal_solution,
    weight,
  )
  add_to_dual_solution_weighted_average(
    solution_weighted_avg,
    current_dual_solution,
    weight,
  )
  return
end

function compute_average(solution_weighted_avg::SolutionWeightedAverage)
  return solution_weighted_avg.sum_primal_solutions /
         solution_weighted_avg.sum_primal_solution_weights,
  solution_weighted_avg.sum_dual_solutions /
  solution_weighted_avg.sum_dual_solution_weights
end

"""
RestartScheme enum

# Values:

-  `NO_RESTARTS`: No restarts are performed.
-  `FIXED_FREQUENCY`: does a restart every [restart_frequency]
    iterations where [restart_frequency] is a user-specified number.
-  `ADAPTIVE_NORMALIZED`: a heuristic based on
    the normalized duality gap to decide when to restart. The general idea is
    is to demand some minimum decrease in the normalized duality gap if that is
    achieved and if either (i) a large reduction in the normalized duality gap
    occurs, or (ii) the normalized duality gap increases from one iteration to
    the other then we restart.
-  `ADAPTIVE_LOCALIZED`: This option restarts when the potential
    function defined by the localized duality gap divided by the restart length
    falls by a factor of necessary_reduction_for_restart. See RestartParameters
    for documentation for necessary_reduction_for_restart.
-  `ADAPTIVE_DISTANCE`: This option restarts when the potential
    function defined by the distance traveled divided by the restart length
    falls by a factor of necessary_reduction_for_restart.
"""
@enum RestartScheme NO_RESTARTS FIXED_FREQUENCY ADAPTIVE_NORMALIZED ADAPTIVE_LOCALIZED ADAPTIVE_DISTANCE

"""
RestartToCurrentMetric enum

# Values:

- `NO_RESTART_TO_CURRENT`: Always reset to the average.
- `GAP_OVER_DISTANCE`: Decide between the average
  current based on which has a smaller normalized duality gap.
- `GAP_OVER_DISTANCE_SQUARED`: Decide between the
  average current based on which has the smaller normalized duality gap divided
  by distance travelled.
"""

@enum RestartToCurrentMetric NO_RESTART_TO_CURRENT GAP_OVER_DISTANCE GAP_OVER_DISTANCE_SQUARED

mutable struct RestartParameters
  """
  Specifies what type of restart scheme is used.
  """
  restart_scheme::RestartScheme

  """
  Specifies how we decide between restarting to the average or current.
  """
  restart_to_current_metric::RestartToCurrentMetric

  """
  If `restart_scheme` = `FIXED_FREQUENCY` then this number determines the
  frequency that the algorithm is restarted.
  """
  restart_frequency_if_fixed::Int64

  """
  If in the past `artificial_restart_threshold` fraction of iterations no
  restart has occurred then a restart will be artificially triggered. The value
  should be between zero and one. Smaller values will have more frequent
  artificial restarts than larger values.
  """
  artificial_restart_threshold::Float64

  """
  Only applies when `restart_scheme` = `ADAPTIVE`. It is the threshold
  improvement in the quality of the current/average iterate compared with that
  of the last restart that will trigger a restart. The value of this parameter
  should be between zero and one. Smaller values make restarts less frequent,
  larger values make restarts more frequent.
  """
  sufficient_reduction_for_restart::Float64

  """
  Only applies when `restart_scheme` = `ADAPTIVE`. It is the threshold
  improvement in the quality of the current/average iterate compared with that
  of the last restart that is neccessary for a restart to be triggered. If this
  thrshold is met and the quality of the iterates appear to be getting worse
  then a restart is triggered. The value of this parameter should be between
  zero and one, and greater than sufficient_reduction_for_restart. Smaller
  values make restarts less frequent, larger values make restarts more frequent.
  """
  necessary_reduction_for_restart::Float64

  """
  Controls the exponential smoothing of log(primal_weight) when the primal
  weight is updated (i.e., on every restart). Must be between 0.0 and 1.0
  inclusive. At 0.0 the primal weight remains frozen at its initial value.
  """
  primal_weight_update_smoothing::Float64

  """
  Use an approximate localized duality gap for restart computations. The
  approximate gap only enforces bounds on variables that are active at the
  center point and objective.
  """
  use_approximate_localized_duality_gap::Bool
end

function construct_restart_parameters(
  restart_scheme::RestartScheme,
  restart_to_current_metric::RestartToCurrentMetric,
  restart_frequency_if_fixed::Int64,
  artificial_restart_threshold::Float64,
  sufficient_reduction_for_restart::Float64,
  necessary_reduction_for_restart::Float64,
  primal_weight_update_smoothing::Float64,
  use_approximate_localized_duality_gap::Bool,
)
  @assert restart_frequency_if_fixed > 1
  @assert 0.0 < artificial_restart_threshold <= 1.0
  @assert 0.0 <
          sufficient_reduction_for_restart <=
          necessary_reduction_for_restart <=
          1.0
  @assert 0.0 <= primal_weight_update_smoothing <= 1.0

  return RestartParameters(
    restart_scheme,
    restart_to_current_metric,
    restart_frequency_if_fixed,
    artificial_restart_threshold,
    sufficient_reduction_for_restart,
    necessary_reduction_for_restart,
    primal_weight_update_smoothing,
    use_approximate_localized_duality_gap,
  )
end

function compute_localized_duality_gaps(
  problem::QuadraticProgrammingProblem,
  current_primal_solution::AbstractVector{Float64},
  current_dual_solution::AbstractVector{Float64},
  avg_primal_solution::Vector{Float64},
  avg_dual_solution::Vector{Float64},
  primal_norm_params::AbstractVector{Float64},
  dual_norm_params::AbstractVector{Float64},
  last_restart_info::RestartInfo,
  use_approximate_localized_duality_gap::Bool,
)

  # Compute bar{r}_i
  distance_traveled_by_average = sqrt(
    weighted_norm(
      avg_primal_solution - last_restart_info.primal_solution,
      primal_norm_params,
    )^2 +
    weighted_norm(
      avg_dual_solution - last_restart_info.dual_solution,
      dual_norm_params,
    )^2,
  )
  # Compute Delta_{bar{r}_i}(bar{w}_{i+1})
  localized_duality_gap_at_average = compute_localized_duality_gap(
    problem,
    avg_primal_solution,
    avg_dual_solution,
    primal_norm_params,
    dual_norm_params,
    distance_traveled_by_average,
    EUCLIDEAN_NORM,
    use_approximate_localized_duality_gap,
  )

  # Compute hat{r}_i
  distance_traveled_by_current = sqrt(
    weighted_norm(
      current_primal_solution - last_restart_info.primal_solution,
      primal_norm_params,
    )^2 +
    weighted_norm(
      current_dual_solution - last_restart_info.dual_solution,
      dual_norm_params,
    )^2,
  )
  # Compute Delta_{hat{r}_i}(hat{w}_{i+1})
  localized_duality_gap_at_current = compute_localized_duality_gap(
    problem,
    current_primal_solution,
    current_dual_solution,
    primal_norm_params,
    dual_norm_params,
    distance_traveled_by_current,
    EUCLIDEAN_NORM,
    use_approximate_localized_duality_gap,
  )

  return (
    gap_at_average = localized_duality_gap_at_average,
    distance_traveled_by_average = distance_traveled_by_average,
    gap_at_current = localized_duality_gap_at_current,
    distance_traveled_by_current = distance_traveled_by_current,
  )
end

# TODO: This should reference the paper when it is done.

"""
Given that a restart was triggered, this function decides if we should restart
to the average or the current iterate.

The general idea is that mirror-prox/pdhg guarrantee:

Delta_{bar{r}_i}(bar{w}_{i+1}) / bar{r}_i^2 <= C / T_i

where i indicates the restart period, bar{r}_i is the distance traveled by the
average iterates across this current restart period, bar{w}_{i+1} is the
iterates averaged across the restart period, T_i is the number of iterations
across the restart period, C is an unknown constant, and Delta_r(w) is the local
duality gap for radius r at the point w. Therefore, we pick the
current iterate hat{w}_{i+1} if

Delta_{hat{r}_i}(hat{w}_{i+1}) / hat{r}_i^2 <
  Delta_{bar{r}_i}(bar{w}_{i+1}) / bar{r}_i^2.

where hat{r}_i is the distance traveled by the current iterate across the
restart period. This choice guarantees that

Delta_{r_i}(w_{i+1}) / r_i^2 <= Delta_{bar{r}_i}(bar{w}_{i+1}) / bar{r}_i^2
                             <= C / T_i

where w_{i+1} = bar{w}_{i+1} and r_i = bar{r}_i if we choose the average
and w_{i+1} = hat{w}_{i+1} and r_i = hat{r}_i if we choose the current
iterate. In other words, this choice guarantees that you do as least as well as
choosing the average in terms of this metric of progress.

"""
function should_reset_to_average(
  current::OptimalObjectiveBoundResult,
  distance_traveled_by_current::Float64,
  average::OptimalObjectiveBoundResult,
  distance_traveled_by_average::Float64,
  restart_to_current_metric::RestartToCurrentMetric,
)
  current_normalized_gap = get_gap(current) / distance_traveled_by_current
  average_normalized_gap = get_gap(average) / distance_traveled_by_average
  if restart_to_current_metric == GAP_OVER_DISTANCE_SQUARED
    return current_normalized_gap / distance_traveled_by_current >=
           average_normalized_gap / distance_traveled_by_average
  elseif restart_to_current_metric == GAP_OVER_DISTANCE
    return current_normalized_gap >= average_normalized_gap
  else
    return true # reset to average
  end
end

function should_do_adaptive_restart_normalized_duality_gap(
  problem::QuadraticProgrammingProblem,
  primal_norm_params::AbstractVector{Float64},
  dual_norm_params::AbstractVector{Float64},
  candidate_localized_gap::OptimalObjectiveBoundResult,
  candidate_distance_traveled::Float64,
  restart_params::RestartParameters,
  last_restart_info::RestartInfo,
  use_approximate_localized_duality_gap::Bool,
  primal_weight::Float64,
)
  lri = last_restart_info
  distance_traveled_last_restart = sqrt(
    lri.primal_distance_moved_last_restart_period^2 * primal_weight +
    lri.dual_distance_moved_last_restart_period^2 / primal_weight,
  )
  # TODO: Cache this after primal weight is updated
  last_restart = compute_localized_duality_gap(
    problem,
    lri.primal_solution,
    lri.dual_solution,
    primal_norm_params,
    dual_norm_params,
    distance_traveled_last_restart,
    EUCLIDEAN_NORM,
    use_approximate_localized_duality_gap,
  )
  do_restart = false
  normalized_candidate_gap =
    get_gap(candidate_localized_gap) / candidate_distance_traveled
  normalized_last_restart_gap =
    get_gap(last_restart) / distance_traveled_last_restart
  gap_reduction_ratio = normalized_candidate_gap / normalized_last_restart_gap
  if gap_reduction_ratio < restart_params.necessary_reduction_for_restart
    if gap_reduction_ratio < restart_params.sufficient_reduction_for_restart
      do_restart = true
    elseif gap_reduction_ratio > lri.gap_reduction_ratio_last_trial
      # Last time we evaluated the results were better.
      do_restart = true
    end
  end
  lri.gap_reduction_ratio_last_trial = gap_reduction_ratio

  return do_restart
end

# TODO: Consider adding ideas from the heuristic restart strategy
# into the theoretical strategies to make them work better.
function should_do_localized_adaptive_restart(
  candidate_localized_gap::OptimalObjectiveBoundResult,
  candidate_restart_length::Int64,
  restart_params::RestartParameters,
  last_restart_info::RestartInfo,
)
  do_restart = false

  lri = last_restart_info
  if candidate_localized_gap == nothing ||
     lri.last_restart_localized_duality_gap == nothing
    return true
  end

  new_potential = get_gap(candidate_localized_gap) / candidate_restart_length
  old_potential =
    get_gap(lri.last_restart_localized_duality_gap) / lri.last_restart_length
  if new_potential / old_potential <
     restart_params.necessary_reduction_for_restart
    do_restart = true
  end

  return do_restart
end

function should_do_distance_based_adaptive_restart(
  candidate_localized_gap::OptimalObjectiveBoundResult,
  candidate_distance_traveled::Float64,
  candidate_restart_length::Int64,
  restart_params::RestartParameters,
  last_restart_info::RestartInfo,
  primal_weight::Float64,
)
  do_restart = false

  lri = last_restart_info

  distance_traveled = candidate_distance_traveled
  distance_traveled_last_restart = sqrt(
    lri.primal_distance_moved_last_restart_period^2 * primal_weight +
    lri.dual_distance_moved_last_restart_period^2 / primal_weight,
  )
  new_potential = distance_traveled / candidate_restart_length
  old_potential = distance_traveled_last_restart / lri.last_restart_length
  if new_potential / old_potential <
     restart_params.necessary_reduction_for_restart
    do_restart = true
  end

  return do_restart
end

"""
  run_restart_scheme(problem, solution_weighted_avg, current_primal_solution,
    current_dual_solution, last_restart_quality, verbosity, restart_params)

This function decides whether to restart and performs the restart.
If it does restart then it updates the solution_weighted_avg,
current_primal_solution, current_dual_solution, and last_restart_quality
accordingly.

# Inputs
- `problem::QuadraticProgrammingProblem`
- `solution_weighted_avg::SolutionWeightedAverage`: If there is a restart then
   this struct will be reset to the restart point.
- `current_primal_solution::AbstractVector`: If there is a restart then this
   vector might be set to the avg_primal_solution. However, it could remain
   unchanged, if the algorithm thinks the current iterate is better than the
   averaged iterate.
- `current_dual_solution::AbstractVector`: If there is a restart then this
   vector might be set to the avg_dual_solution. However, it could remain
   unchanged, if the algorithm thinks the current iterate is better than the
   averaged iterate.
- `last_restart_info::RestartInfo`: Information stored about the last
   restart point.
- `iterations_completed::Int64`: The number of successful iterations completed
   by the algoirthm (i.e., where the step is accepted).
- `primal_norm_params::AbstractVector{Float64}`: The weights of the weighted l2
   norm used to measure distance in the primal and compute steps.
- `dual_norm_params::AbstractVector{Float64}`: The weights of the weighted l2
   norm used to measure distance in the dual and compute steps.
- `primal_weight::Float64`: The current value of the primal weight.
- `verbosity::Int64`: The output level.
- `restart_params::RestartParameters`: Parameters for making restart decisions.

# Output
A RestartChoice enum which tells us if a restart was performed. See the
the definition of RestartChoice for more details.
"""
function run_restart_scheme(
  problem::QuadraticProgrammingProblem,
  solution_weighted_avg::SolutionWeightedAverage,
  current_primal_solution::AbstractVector,
  current_dual_solution::AbstractVector,
  last_restart_info::RestartInfo,
  iterations_completed::Int64,
  primal_norm_params::AbstractVector{Float64},
  dual_norm_params::AbstractVector{Float64},
  primal_weight::Float64,
  verbosity::Int64,
  restart_params::RestartParameters,
)
  # TODO: Add options for the purposes of benchmarking, e.g., use the
  # last iterate instead of the average, or average across all iterates.
  if solution_weighted_avg.sum_primal_solutions_count > 0 &&
     solution_weighted_avg.sum_dual_solutions_count > 0
    avg_primal_solution, avg_dual_solution =
      compute_average(solution_weighted_avg)
  else
    return RESTART_CHOICE_NO_RESTART
  end

  restart_length = solution_weighted_avg.sum_primal_solutions_count
  artificial_restart = false
  do_restart = false
  # If we have not restarted for a very long time (as a fraction of the
  # iterations completed by the algorithm) then force a restart.
  # A restart is always triggered the first time this function is called,
  # because because restart_params.artificial_restart_threshold is less
  # than or equal to one and iterations_completed == restart_length.
  if restart_length >=
     restart_params.artificial_restart_threshold * iterations_completed
    do_restart = true
    artificial_restart = true
  end
  # Decide if we are going to reset to average.
  if restart_params.restart_scheme == NO_RESTARTS
    reset_to_average = false
    candidate_localized_gap = nothing
  else
    localized_duality_gaps = compute_localized_duality_gaps(
      problem,
      current_primal_solution,
      current_dual_solution,
      avg_primal_solution,
      avg_dual_solution,
      primal_norm_params,
      dual_norm_params,
      last_restart_info,
      restart_params.use_approximate_localized_duality_gap,
    )
    reset_to_average = should_reset_to_average(
      localized_duality_gaps.gap_at_current,
      localized_duality_gaps.distance_traveled_by_current,
      localized_duality_gaps.gap_at_average,
      localized_duality_gaps.distance_traveled_by_average,
      restart_params.restart_to_current_metric,
    )

    if reset_to_average
      candidate_localized_gap = localized_duality_gaps.gap_at_average
      candidate_distance_traveled =
        localized_duality_gaps.distance_traveled_by_average
    else
      candidate_localized_gap = localized_duality_gaps.gap_at_current
      candidate_distance_traveled =
        localized_duality_gaps.distance_traveled_by_current
    end
  end

  if !do_restart
    # Decide if we are going to do a restart.
    if restart_params.restart_scheme == ADAPTIVE_NORMALIZED
      do_restart = should_do_adaptive_restart_normalized_duality_gap(
        problem,
        primal_norm_params,
        dual_norm_params,
        candidate_localized_gap,
        candidate_distance_traveled,
        restart_params,
        last_restart_info,
        restart_params.use_approximate_localized_duality_gap,
        primal_weight,
      )

    elseif (
      restart_params.restart_scheme == ADAPTIVE_LOCALIZED ||
      restart_params.restart_scheme == ADAPTIVE_DISTANCE
    ) && last_restart_info.last_restart_localized_duality_gap == nothing
      do_restart = true # automatically restart if no restarts have occurred
    elseif restart_params.restart_scheme == ADAPTIVE_LOCALIZED
      do_restart = should_do_localized_adaptive_restart(
        candidate_localized_gap,
        restart_length,
        restart_params,
        last_restart_info,
      )
    elseif restart_params.restart_scheme == ADAPTIVE_DISTANCE
      do_restart = should_do_distance_based_adaptive_restart(
        candidate_localized_gap,
        candidate_distance_traveled,
        restart_length,
        restart_params,
        last_restart_info,
        primal_weight,
      )
    elseif restart_params.restart_scheme == FIXED_FREQUENCY &&
           restart_params.restart_frequency_if_fixed <= restart_length
      do_restart = true
    end
  end

  if !do_restart
    return RESTART_CHOICE_NO_RESTART
  else
    if reset_to_average
      if verbosity >= 4
        print("  Restarted to average")
      end
      current_primal_solution .= avg_primal_solution
      current_dual_solution .= avg_dual_solution
    else
      # Current point is much better than average point.
      if verbosity >= 4
        print("  Restarted to current")
      end
    end

    if verbosity >= 4
      print(" after ", rpad(restart_length, 4), " iterations")
      if artificial_restart
        println("*")
      else
        println("")
      end
    end
    reset_solution_weighted_average(solution_weighted_avg)

    update_last_restart_info(
      last_restart_info,
      current_primal_solution,
      current_dual_solution,
      avg_primal_solution,
      avg_dual_solution,
      primal_norm_params,
      dual_norm_params,
      primal_weight,
      candidate_localized_gap,
      restart_length,
    )

    if reset_to_average
      return RESTART_CHOICE_RESTART_TO_AVERAGE
    else
      return RESTART_CHOICE_WEIGHTED_AVERAGE_RESET
    end
  end
end

"""
Given the current primal weight and restart information, returns a recomputed
primal weight.

# Inputs:
- `last_restart_info::RestartInfo`: Information on the last restart point.
- `primal_weight::Float64`: The primal weight is the ratio of step sizes
  between the dual and primal.
- `primal_weight_update_smoothing::Float64`: Parameter that decides how smoothed
  the primal weight updates should be.
- `verbosity::Int64`: Output level.
# Output:
The new primal weight.
"""
function compute_new_primal_weight(
  last_restart_info::RestartInfo,
  primal_weight::Float64,
  primal_weight_update_smoothing::Float64,
  verbosity::Int64,
)
  primal_distance = last_restart_info.primal_distance_moved_last_restart_period
  dual_distance = last_restart_info.dual_distance_moved_last_restart_period
  # Note that the choice of eps() is arbitrary. It could have been eps()*100 or
  # eps()/100. It is just to avoid the case when dual_distance or
  # primal_distance equals zero.
  if primal_distance > eps() && dual_distance > eps()
    new_primal_weight_estimate = dual_distance / primal_distance
    # Exponential moving average.
    # If primal_weight_update_smoothing = 1.0 then there is no smoothing.
    # If primal_weight_update_smoothing = 0.0 then the primal_weight is frozen.
    log_primal_weight =
      primal_weight_update_smoothing * log(new_primal_weight_estimate) +
      (1 - primal_weight_update_smoothing) * log(primal_weight)

    primal_weight = exp(log_primal_weight)
    if verbosity >= 4
      Printf.@printf "  New computed primal weight is %.2e\n" primal_weight
    end

    return primal_weight
  else
    return primal_weight
  end
end

function update_last_restart_info(
  last_restart_info::RestartInfo,
  current_primal_solution::AbstractVector{Float64},
  current_dual_solution::AbstractVector{Float64},
  avg_primal_solution::Vector{Float64},
  avg_dual_solution::Vector{Float64},
  primal_norm_params::AbstractVector{Float64},
  dual_norm_params::AbstractVector{Float64},
  primal_weight::Float64,
  candidate_localized_gap::Union{Nothing,OptimalObjectiveBoundResult},
  restart_length::Int64,
)
  # TODO: Define anchor points and only update these anchor points
  # when a nonartifical restart occurs. This should give better primal weight
  # update performance. The reason I add this comment is that I have noticed
  # that it looks like there is sometimes interaction between the primal
  # weights and restart scheme which causes cycling like behaviour.

  last_restart_info.primal_distance_moved_last_restart_period =
    weighted_norm(
      avg_primal_solution - last_restart_info.primal_solution,
      primal_norm_params,
    ) / sqrt(primal_weight)
  last_restart_info.dual_distance_moved_last_restart_period =
    weighted_norm(
      avg_dual_solution - last_restart_info.dual_solution,
      dual_norm_params,
    ) * sqrt(primal_weight)
  last_restart_info.primal_solution .= current_primal_solution
  last_restart_info.dual_solution .= current_dual_solution

  last_restart_info.last_restart_length = restart_length
  last_restart_info.last_restart_localized_duality_gap = candidate_localized_gap

end

"""
A simple string name for a PointType.
"""
function point_type_label(point_type::PointType)
  if point_type == POINT_TYPE_CURRENT_ITERATE
    return "current"
  elseif point_type == POINT_TYPE_AVERAGE_ITERATE
    return "average"
  elseif point_type == POINT_TYPE_ITERATE_DIFFERENCE
    return "difference"
  else
    return "unknown PointType"
  end
end

"""
Logging for when the algorithm terminates.
"""
function generic_final_log(
  problem::QuadraticProgrammingProblem,
  current_primal_solution::AbstractVector{Float64},
  current_dual_solution::AbstractVector{Float64},
  last_iteration_stats::IterationStats,
  verbosity::Int64,
  iteration::Int64,
  termination_reason::TerminationReason,
)
  if verbosity >= 1
    print("Terminated after $iteration iterations: ")
    println(termination_reason_to_string(termination_reason))
  end

  method_specific_stats = last_iteration_stats.method_specific_stats
  if verbosity >= 3
    for convergence_information in last_iteration_stats.convergence_information
      Printf.@printf(
        "For %s candidate:\n",
        point_type_label(convergence_information.candidate_type)
      )
      # Print more decimal places for the primal and dual objective.
      Printf.@printf(
        "Primal objective: %f, ",
        convergence_information.primal_objective
      )
      Printf.@printf(
        "dual objective: %f, ",
        convergence_information.dual_objective
      )
      Printf.@printf(
        "corrected dual objective: %f \n",
        convergence_information.corrected_dual_objective
      )
    end
    if haskey(method_specific_stats, "estimated_lower_bound") &&
       haskey(method_specific_stats, "estimated_upper_bound")
      Printf.@printf(
        "Estimated optimal objective range: [%f, %f] \n",
        method_specific_stats["estimated_lower_bound"],
        method_specific_stats["estimated_upper_bound"],
      )
    end
    Printf.@printf(
      "Lagrangian value: %f \n",
      method_specific_stats["lagrangian_value"],
    )
  end
  if verbosity >= 4
    Printf.@printf(
      "Time (seconds):\n - Basic algorithm: %.2e\n - Full algorithm:  %.2e\n",
      method_specific_stats["time_spent_doing_basic_algorithm"],
      last_iteration_stats.cumulative_time_sec,
    )
  end

  if verbosity >= 7
    for convergence_information in last_iteration_stats.convergence_information
      print_infinity_norms(convergence_information)
    end
    print_variable_and_constraint_hardness(
      problem,
      current_primal_solution,
      current_dual_solution,
    )
  end
end

function update_objective_bound_estimates(
  method_specific_stats::Dict{String,Float64},
  problem::QuadraticProgrammingProblem,
  current_primal_solution::AbstractVector{Float64},
  current_dual_solution::AbstractVector{Float64},
  primal_norm_weights::AbstractVector{Float64},
  dual_norm_weights::AbstractVector{Float64},
)
  # TODO: Use better estimates of the distance to optimality.
  # TODO: Once these better estimates are developed do a study.
  estimated_primal_distance_to_optimality =
    max(1e-8, weighted_norm(current_primal_solution, primal_norm_weights))
  estimated_dual_distance_to_optimality =
    max(1e-8, weighted_norm(current_dual_solution, dual_norm_weights))

  estimated_local_duality_gap = compute_localized_duality_gap(
    problem,
    current_primal_solution,
    current_dual_solution,
    primal_norm_weights / estimated_primal_distance_to_optimality^2,
    dual_norm_weights / estimated_dual_distance_to_optimality^2,
    1.0,
    MAX_NORM,
    false, # don't compute this approximately
  )

  method_specific_stats["lagrangian_value"] =
    estimated_local_duality_gap.lagrangian_value
  method_specific_stats["estimated_lower_bound"] =
    estimated_local_duality_gap.lower_bound_value
  method_specific_stats["estimated_upper_bound"] =
    estimated_local_duality_gap.upper_bound_value
end

function select_initial_primal_weight(
  problem::QuadraticProgrammingProblem,
  primal_norm_params::AbstractVector{Float64},
  dual_norm_params::AbstractVector{Float64},
  primal_importance::Float64,
  verbosity::Int64,
)
  rhs_vec_norm = weighted_norm(problem.right_hand_side, dual_norm_params)
  obj_vec_norm = weighted_norm(problem.objective_vector, primal_norm_params)
  if obj_vec_norm > 0.0 && rhs_vec_norm > 0.0
    # Note that the units of the objective vector are
    # (objective units / x units) and the units of the right hand side are
    # (objective units / y units). Therefore the units of
    # obj_vec_norm / rhs_vec_norm are (y units / x units).
    # Note that this arguement is somewhat handwavy because each component of x
    # could have different physical units. For example, one objective
    # coefficient could be cost per widget and another could be cost per
    # factory.
    primal_weight = primal_importance * (obj_vec_norm / rhs_vec_norm)
  else
    primal_weight = primal_importance
  end
  if verbosity >= 6
    println("Initial primal weight = $primal_weight")
  end
  return primal_weight
end

# The gradient of the Lagrangian with respect to primal variables. This is
# objective_matrix * primal_solution + objective_vector -
# constraint_matrix' * dual_solution, i.e., the value of the reduced costs
# if we wanted dual feasibilty to hold exactly.
function compute_primal_gradient(
  problem::QuadraticProgrammingProblem,
  primal_solution::AbstractVector{Float64},
  dual_solution::AbstractVector{Float64},
)
  return compute_primal_gradient_from_dual_product(
    problem,
    primal_solution,
    problem.constraint_matrix' * dual_solution,
  )
end

function compute_primal_gradient_from_dual_product(
  problem::QuadraticProgrammingProblem,
  primal_solution::AbstractVector{Float64},
  dual_product::AbstractVector{Float64},
)
  return problem.objective_matrix * primal_solution .+
         problem.objective_vector .- dual_product
end

function compute_dual_gradient(
  problem::QuadraticProgrammingProblem,
  primal_solution::AbstractVector{Float64},
)
  return problem.right_hand_side .- problem.constraint_matrix * primal_solution
end

function compute_lagrangian_value(
  problem::QuadraticProgrammingProblem,
  primal_solution::AbstractVector{Float64},
  dual_solution::AbstractVector{Float64},
)
  return 0.5 *
         dot(primal_solution, problem.objective_matrix * primal_solution) +
         dot(primal_solution, problem.objective_vector) -
         dot(primal_solution, problem.constraint_matrix' * dual_solution) +
         dot(dual_solution, problem.right_hand_side) +
         problem.objective_constant
end


text:
mutable struct Variational <: Method
  algorithm::Symbol
  grad_samples::Int64
  elbo_samples::Int64
  #eta_adagrad::Float64
  iter::Int64
  tol_rel_obj::Float64
  eval_elbo::Int64
  output_samples::Int64
end
"""

# Variational type and constructor

Settings for method=Variational() in Stanmodel.

### Method
```julia
Variational(;
  grad_samples=1,
  elbo_samples=100,
  eta_adagrad=0.1,
  iter=10000,
  tol_rel_obj=0.01,
  eval_elbo=100,
  algorithm=:meanfield,
  output_samples=10000
)
```
### Optional arguments
```julia
* `algorithm::Symbol`             : Variational inference algorithm
                                    :meanfiedl;
                                    :fullrank
* `iter::Int64`                   : Maximum number of iterations
* `grad_samples::Int`             : No of samples for MC estimate of gradients
* `elbo_samples::Int`             : No of samples for MC estimate of ELBO
* `eta::Float64`                  : Stepsize weighing parameter for adaptive sequence
* `adapt::Adapt`                  : Warmupadaptations
* `tol_rel_obj::Float64`          : Tolerance on the relative norm of objective
* `eval_elbo::Int`                : Tolerance on the relative norm of objective
* `output_samples::Int`           : Numberof posterior samples to draw and save
```

### Related help
```julia
?Stanmodel                      : Create a StanModel
?Stan.Method                   : Available top level methods
?Stan.Adapt                     : Adaptation settings
```
"""
Variational(;grad_samples=1, elbo_samples=100,
  eta_adagrad=0.1, iter=10000,
  tol_rel_obj=0.01, eval_elbo=100,
  algorithm=:meanfield, output_samples=10000) = Variational(algorithm,
    grad_samples, elbo_samples,
    #eta_adagrad,
    iter, tol_rel_obj, eval_elbo, output_samples)

function variational_show(io::IO, v::Variational, compact::Bool)
  if compact
    println("Variational($(v.algorithm). $(v.grad_samples), $(v.elbo_samples), $(v.eta_adagrad), $(v.iter), $(v.tol_rel_obj), $(v.eval_elbo), $(v.output_samples))")
  else
    println("  method =                     Variational()")
    println("    algorithm =               ", v.algorithm)
    println("    grad_samples =            ", v.grad_samples)
    println("    elbo_samples =            ", v.elbo_samples)
    #println("    eta_adagrad =             ", v.eta_adagrad)
    println("    iter =                    ", v.iter)
    println("    tol_rel_obj =             ", v.tol_rel_obj)
    println("    eval_elbo =               ", v.eval_elbo)
    println("    output_samples =          ", v.output_samples)
  end
end

show(io::IO, v::Variational) = variational_show(io, v, false)


text:
# Note that this script can accept some limited command-line arguments, run
# `julia build_tarballs.jl --help` to see a usage message.
using BinaryBuilder

name = "libusb"
version = v"1.0.23"

# Collection of sources required to complete build
sources = [
    "https://github.com/libusb/libusb.git" =>
    "e782eeb2514266f6738e242cdcb18e3ae1ed06fa",

]

# Bash recipe for building across all platforms
script = raw"""
cd $WORKSPACE/srcdir
cd libusb/
./bootstrap.sh
./configure --prefix=${prefix} --build=${MACHTYPE} --host=${target} --disable-udev
make -j${nproc}
make install
install_license COPYING
"""

# These are the platforms we will build for by default, unless further
# platforms are passed in on the command line
platforms = [p for p in supported_platforms() if !isa(p, FreeBSD)]

# The products that we will ensure are always built
products = [
    LibraryProduct(["libusb", "libusb-1", "libusb-1.0"], :libusb)
]

# Dependencies that must be installed before this package can be built
dependencies = [

]

# Build the tarballs, and possibly a `build.jl` as well.
build_tarballs(ARGS, name, version, sources, script, platforms, products, dependencies)


text:
__precompile__()

module OptionalData

export @OptionalData, OptData, show, push!, isavailable, get

import Base: push!, get, show

mutable struct OptData{T}
    data::Union{T, Nothing}
    name::String
    msg::String
end
OptData{T}(name, msg="") where {T} = OptData{T}(nothing, name, msg)

function show(io::IO, opt::OptData{T}) where T
    val = isavailable(opt) ? get(opt) : ""
    print(io, "OptData{$T}($val)")
end

"""
    push!(opt::OptData{T}, data::T) where T

Push `data` of type `T` to `opt`.
"""
function push!(opt::OptData{T}, data::T) where T
    opt.data = data
    opt
end

"""
    push!(opt::OptData, ::Type{T}, args...) where T

Construct an object of type `T` from `args` and push it to `opt`.
"""
function push!(opt::OptData, ::Type{T}, args...) where {T}
    push!(opt, T(args...))
    opt
end
push!(opt::OptData{T}, args...) where {T} = push!(opt, T, args...)

"""
    get(opt::OptData)

Get data from `opt`. Throw an exception if no data has been pushed to `opt` before.
"""
function get(opt::OptData{T}) where T
    !isavailable(opt) && error(opt.name, " is not available. ", opt.msg)
    opt.data::T
end

"""
    @OptionalData name type msg=""

Initialise a constant `name` with type `OptData{type}`.
An exception with the custom error message `msg` is thrown when `name` is
accessed before data has been pushed to it.

# Example

```julia
@OptionalData OPT_FLOAT Float64
```
"""
macro OptionalData(name, typ, msg="")
   :(const $(esc(name)) = OptData{$(esc(typ))}($(string(name)), $msg))
end

"""
    isavailable(opt::OptData)

Check whether data has been pushed to `opt`.
"""
isavailable(opt::OptData) = opt.data !== nothing

end # module


text:
# Indexing
@propagate_inbounds function Base.getindex(f::T, i::Int) where T <: AbstractPICDataStructure
    getindex(scalarness(T), f, i)
end
@propagate_inbounds function Base.setindex!(f::T, v, i::Int) where T <: AbstractPICDataStructure
    setindex!(scalarness(T), f, v, i)
end

@propagate_inbounds function Base.getindex(f::T; named_idxs...) where T <: AbstractPICDataStructure
    location = named_idxs.data
    dir = only(keys(location))
    val = only(values(location))
    idx = location2idx(domain_discretization(T), f, dir, val)

    dim = dir_to_idx(dir)
    data = selectdim(unwrapdata(f), dim, idx)
    grid = dropdims(getdomain(f); dims=dir)

    newstruct(f, data, grid)
end

Base.firstindex(f::AbstractPICDataStructure) = firstindex(unwrapdata(f))
Base.lastindex(f::AbstractPICDataStructure) = lastindex(unwrapdata(f))

function Base.size(f::AbstractPICDataStructure, dim...)
    size(unwrapdata(f), dim...)
end

Base.LinearIndices(f::AbstractPICDataStructure) = LinearIndices(unwrapdata(f))
Base.IndexStyle(::Type{<:AbstractPICDataStructure}) = Base.IndexLinear()

Base.eltype(::T) where T <: AbstractPICDataStructure = eltype(scalarness(T), T)

# Iteration
Base.iterate(f::AbstractPICDataStructure, state...) = iterate(unwrapdata(f), state...)
Base.length(f::AbstractPICDataStructure) = length(unwrapdata(f))

# Broadcasting
Base.BroadcastStyle(::Type{T}) where T<:AbstractPICDataStructure = Base.BroadcastStyle(scalarness(T), T)

function Base.similar(f::T, typ::Type{S}, dims::Dims) where {S, T <: AbstractPICDataStructure}
    # @debug "similar AbstractPICDataStructure"
    similar(scalarness(T), f, typ, dims)
end

"""
`A = find_filed(Fs)` returns the first `AbstractPICDataStructure` among the arguments.
"""
function find_field(bc::Base.Broadcast.Broadcasted)
    # @debug "Destructuring Broadcasted $(typeof(bc.args))"
    find_field(bc.args)
end
function find_field(args::Tuple)
    # @debug "First argument in broadcast: $(typeof(first(args)))"
    find_field(find_field(first(args)), Base.tail(args))
end
function find_field(f)
    # @debug "Any fallback got $(typeof(f))"
    f
end
find_field(::Tuple{}) = nothing
function find_field(f::AbstractPICDataStructure, rest)
    # @debug "Found AbstractPICDataStructure $(typeof(f))"
    f
end
function find_field(::Any, rest)
    # @debug "Recursing in the second argument $(typeof(rest))"
    find_field(rest)
end

Base.getproperty(f::T, k::Symbol) where T <: AbstractPICDataStructure = getproperty(scalarness(T), f, k)
Base.propertynames(f::T) where T <: AbstractPICDataStructure = propertynames(scalarness(T), f)


text:
``## Motion Planning 3D Benchmark
## =========================
using Sigma
using Lens
using Luxor
import Base.convert
# include("../vis.jl")

mappairs(f, v::Vector) = map(f, v[1:end-1], v[2:end])

# Use for functions which should take a normal or equivalently typed randarray
Point = AbstractVector
Vec = AbstractVector
Mat = AbstractMatrix

"A geometric entity of N dimensions"
abstract type Entity{N} end

"2D Rectangle"
immutable Rectangle <: Entity{2}
  bounds::Mat{Float64}
end

"Circle"
immutable Circle <: Entity{2}
  center::Vec{Float64}
  r
end

"Edge between defined by start and end points"
immutable Edge
  points::Mat
end

"Edge defined by origin and direction vector"
immutable ParametricEdge
  coords::Mat
end

function ParametricEdge(e::Edge)
  origin = e.points[:,1]
  dir = e.points[:,2] - e.points[:,1]
  ParametricEdge(hcat(origin,dir))
end

"Is the `point` in the `box`?"
function ispointinpoly(point::Point, box::Rectangle)
  r = (point[1] >= box.bounds[1,1]) &
      (point[1] <= box.bounds[2,1]) &
      (point[2] >= box.bounds[1,2]) &
      (point[2] <= box.bounds[2,2])
  r
end

"Where - if anywhere - along `p` does it intersect segment"
function intersect_segments(ppos::Point, pdir::Vec, qpos::Point, qdir::Vec)
  @show ppos
  @show qpos
  w = ppos - qpos
  u = pdir
  v = qdir
  (v[2] * w[1] - v[1] * w[2]) / (v[1] * u[2] - v[2] * u[1])
end

"Do edges `e1` and `e2` not intersect?"
function intersects(e1::ParametricEdge, e2::ParametricEdge)
  s = intersect_segments(e1.coords[:,1], e1.coords[:,2],
                         e2.coords[:,1], e2.coords[:,2])
  (s < 0) | (s > 1)
end

"Forward kinematic function"
function kinematic(angles::Vector)
  total = angles[1]
  sin_total = sin(total)
  cos_total = cos(total)
  for i = 2:length(angles)
    total = total + angles[i]
    sin_total += sin(total)
    cos_total += cos(total)
  end
  sin_total, cos_total
end

"Compute vertices from angles"
function vertices(angles::Vector)
  xs = [0.0]
  ys = [0.0]
  total = 0.0
  sin_total = 0.0
  cos_total = 0.0
  for i = 1:length(angles)
    total = total + angles[i]
    sin_total += sin(total)
    xs = vcat(xs, [sin_total])
    cos_total += cos(total)
    ys = vcat(ys, [cos_total])
  end
  println(xs)
  println(ys)
  permutedims(hcat(xs, ys), (2, 1))
end

"Edge directions from points"
directions(points) = points[:, 2:end] - points[:, 1:end-1]

"Parametric Edges from angles"
function anglestoedge(angles::Vector)
  ps = vertices(angles)
  println("PS", ps)
  dirs = directions(ps)
  ps = ps[:, 1:end-1]
  @assert length(ps) == length(dirs)
  ps, dirs
  [ParametricEdge(hcat(ps[:, i], dirs[:, i])) for i = 1:size(ps, 2)]
end

"Does edge `e1` intersect with `circle`?"
function intersects(e1::ParametricEdge, circle::Circle)
  rayorig = e1.coords[:,1]
  raydir = e1.coords[:,2]
  r = circle.r
  f = rayorig - circle.center # Vector from center sphere to ray start
  a = dot(raydir, raydir)
  b = 2.0 * dot(f,raydir)
  c = dot(f,f) - r*r

  # discriminant
  b*b-4*a*c < 0
end

intersects(e1::ParametricEdge, e2::Edge) = intersects(e1, ParametricEdge(e2))

"Do `points` avoid `obs`tacles?"
function pairwisecompare(edges::Vector, obs)
  conditions = [intersects(e, o) for e in edges, o in obs]
  (&)(conditions...)
end

function convert(::Type{Vector{ParametricEdge}}, points::RandArray)
  [ParametricEdge([points[:,i] (points[:,i+1] - points[:,i])])
           for i = 1:size(points,2)-1]
end

"Are `points` valid? starts at `origin`, ends at `dest`, avoids `obstacles`?"
function validpath(points, obstacles, origin, dest)
  is_start_ok = ispointinpoly(points[:, 1], origin)
  is_end_ok = ispointinpoly(points[:,end], dest)
  param_edges = convert(Vector{ParametricEdge}, points)
  avoids_obstacles = pairwisecompare(param_edges, obstacles)
  is_start_ok & avoids_obstacles & is_end_ok
end

function example_data(path_length::Integer)
  # obstacles = [Circle([5.0, 5.0], 3.0)]
  obstacles = [Circle([5.0, 5.0], 1.0)]
  points = mvuniform(0, 10, 2, path_length)
  origin = Rectangle([0.0 0.0
                      0.2 0.2])
  dest = Rectangle([9.9 9.9
                    10.0 10.0])
  points, origin, dest, obstacles
end

function test_mp2d(path_length = 4, nsamples = 1)
  points, origin, dest, obstacles = example_data(path_length)
  good_path = validpath(points, obstacles, origin, dest)
  sample = rand(points, good_path, nsamples; precision = 0.01, parallel = true, ncores = nprocs() - 1) / 10.0
end

function gettiming(results)
  timediffs2 = vcat([get(statsgo, proc_id=i, lensname=:sat_check) for i = 2:nprocs()]...)
  timediffs2 =  float(map(i->float(i)/1e9, timediffs2))
  mean(timediffs2), std(timediffs2)
end

"Draw the target at `x, y`"
function drawtarget(x, y)
  p1 = Luxor.Point(x, y)
  sethue("red")
  circle(p1, 0.1, :fill)
end

"Draw the path"
function drawpath(points)
  setline(3)
  curr = O
  for i = 1:size(points, 2)
    color = randomhue()
    sethue(color)
    x, y = points[1, i], points[2, i]
    println("XY", x, " ", y)
    point = Luxor.Point(x, y)
    line(curr, point, :stroke)
    curr = point
  end
end

function drawobstacles(obstacles)
  for ob in obstacles
    draw(ob)
  end
end

"Draw the path, target and obstacles"
function drawscene(points, obstacles, x, y)
  Drawing(1000, 1000, "scenes.png")
  origin()
  scale(50.0, 50.0)
  background("white")

  drawpath(points)
  drawobstacles(obstacles)
  drawtarget(x, y)

  finish()
  preview()
end

function run_benchmark()
  resultsgo, statsgo = capture(test_mp2d, [:distance, :sat_check, :post_loop])
end

"Draw a circle"
function draw(c::Circle)
  sethue("black")
  circle(Luxor.Point(c.center...), c.r, :fill)
end

function example_data_angles(path_length::Integer)
  angles = mvuniform(0.0, Float64(2pi), path_length)
  obstacles = [Circle([0.5, 1.5], 0.7)
               Circle([2.0, 1.8], 0.7)]
  x_target = 2.4
  y_target = 1.5
  angles, obstacles, x_target, y_target
end

function test_angles(path_length = 5)
  angles, obstacles, x_target, y_target = example_data_angles(path_length)
  x, y = kinematic(angles)
  reach_condition = (x==x_target) & (y==y_target)

  param_edges = anglestoedge(angles)
  obstacle_condition = pairwisecompare(param_edges, obstacles)
  condition = reach_condition & obstacle_condition
  condition = obstacle_condition

  angles_sample = rand(angles, condition; precision=0.01)
  angles_sample = angles_sample/Float64(2pi) # Divide due to bug in DRealRandVar
  println("Actual Position", kinematic(angles_sample))
  points = vertices(angles_sample)
  print("Points", points)
  print("Angles", angles_sample)
  drawscene(points, obstacles, x_target, y_target)
end

function test_rand_angles(path_length = 5)
  angles, obstacles, x_target, y_target = example_data_angles(path_length)
  angles_sample = rand(angles)
  points = vertices(angles_sample)
  drawscene(points, obstacles, x_target, y_target)
end

test_angles()


text:
export Box

# Below, the box axes describe the directions of edges.  For example, if a₁, a₂, a₃ are the
# axes of a 3D box, then the volume inside the box is the collection of points x₁a₁ + x₂a₂
# + x₃a₃, where -rₙ ≤ xₙ ≤ rₙ.  This means aₙ and aₘ with n ≠ m span a face of the box.
#
# The box projection matrix (Box.p) is the matrix that produces the coordinates (x₁, x₂, x₃)
# of a given point.  In other words, for a point x inside the box, the nth entry of Box.p * x
# must have magnitude ≤ rₙ.  This means that Box.p * aₙ = eₙ, because aₙ has only the aₙ-component
# and the component must be 1.  This also means that each row of Box.p is orthogononal to
# two box axes, and therefore normal to the face spanned by the two box axes.  (Note that
# the rows of Box.p are not unit normals to the faces, becuase they are not unit vectors.)
mutable struct Box{N,N²,D,T} <: Shape{N,N²,D,T}
    c::SVector{N,T}  # center of box
    r::SVector{N,T}  # "radii" (semi-axes) in axis directions
    p::SMatrix{N,N,T,N²}  # projection matrix to box coordinates
	l::SVector{N,T}	# lower bounds
	u::SVector{N,T}	# upper bounds
    data::D  # auxiliary data
    #Box{N,N²,D,T}(c,r,p,data) where {N,N²,D,T<:Real} = new(c,r,p,data)  # suppress default outer constructor
end

signmatrix(b::Box{1}) = SMatrix{1,1}(1)
signmatrix(b::Box{2}) = SMatrix{2,2}(1,1, -1,1)
signmatrix(b::Box{3}) = SMatrix{3,4}(1,1,1, -1,1,1, 1,-1,1, 1,1,-1)
signmatrix(v::SVector{1}) = SMatrix{1,1}(1)
signmatrix(v::SVector{2}) = SMatrix{2,2}(1,1, -1,1)
signmatrix(v::SVector{3}) = SMatrix{3,4}(1,1,1, -1,1,1, 1,-1,1, 1,1,-1)
@non_differentiable signmatrix(::Any)

function Box(c::SVector{N,T},
		    d::SVector{N,T},
		    axes::SMatrix{N,N,<:Real}=SMatrix{N,N,T}(I),
		    data::D=nothing) where {N,D,T<:Real}
	r = 0.5d
	@tullio axnorm[i] := axes[i,j]^2 |> sqrt
	@tullio p_inv[i,j] := axes[i,j] / axnorm[j]
	# p_inv = axes ./ sqrt.(sum(abs2,axes,dims=1))
	smr = Zygote.@ignore(signmatrix(r))
	@tullio (max) m[j] := p_inv[i,k] * r[i] * smr[i,j] nograd=smr
    Box{N,N*N,D,T}(c, 0.5d, SMatrix(inv(p_inv)),c-SVector(m),c+SVector(m),data)
end

Box(c::AbstractVector{T},  # center of box
    d::AbstractVector{T},  # size of box in axis directions
    axes::AbstractMatrix{<:Real}=Matrix{T}(I,length(c),length(c)),  # columns are axes vectors (each being parallel to two sets of faces in 3D)
    data=nothing) where T<:Real =
    (N = length(c); Box(SVector{N}(c), SVector{N}(d), SMatrix{N,N}(axes), data))

Base.:(==)(b1::Box, b2::Box) = b1.c==b2.c && b1.r==b2.r && b1.p==b2.p && b1.data==b2.data
Base.isapprox(b1::Box, b2::Box) = b1.c≈b2.c && b1.r≈b2.r && b1.p≈b2.p && b1.data==b2.data
Base.hash(b::Box, h::UInt) = hash(b.c, hash(b.r, hash(b.p, hash(b.data, hash(:Box, h)))))

function bounds(b::Box)
    # A = inv(b.p) .* b.r'
    # m = maximum(abs.(A * signmatrix(b)), dims=2)[:,1] # extrema of all 2^N corners of the box
    # # m = maximum(abs.(A * signmatrix(b)), dims=2)[:,1] # extrema of all 2^N corners of the box
    # # m = maximum(abs.(Array((inv(b.p) .* b.r') * signmatrix(b))), dims=2)[:,1] # extrema of all 2^N corners of the box
	#
	# # A = ( inv(b.p) .* b.r' ) * signmatrix(b)
    # # m = maximum(abs.(A * signmatrix(b)), dims=2)[:,1] # extrema of all 2^N corners of the box
    # return (b.c-m,b.c+m)
	return b.l, b.u
end

function Base.in(x::SVector{N,<:Real}, b::Box{N}) where {N}
    d = b.p * (x - b.c)
    for i = 1:N
        abs(d[i]) > b.r[i] && return false  # boundary is considered inside
    end
    return true
end

function f_onbnd(bin::Box{N,N²,D,T},absdin) where {N,N²,D,T<:Real}
	b = Zygote.@ignore(bin)
	absd = Zygote.@ignore(absdin)
	abs.(b.r.-absd) .≤ Base.rtoldefault(T) .* b.r  # basically b.r .≈ absd but faster
end

f_isout(b,absd) =  (isout = Zygote.@ignore ((b.r.<absd) .| f_onbnd(b,absd) ); isout)
f_signs(d) =  (signs = Zygote.@ignore(sign.(d)'); signs) #(copysign.(1.0,d)'); signs)

@non_differentiable f_onbnd(::Any)
@non_differentiable f_isout(::Any)
@non_differentiable f_signs(::Any)

function surfpt_nearby(x::SVector{N,T}, b::Box{N}) where {N,T<:Real}
    ax = inv(b.p)  # axes: columns are unit vectors
	n0 = @inbounds b.p ./  [ sqrt(b.p[1,1]^2 + b.p[1,2]^2) sqrt(b.p[2,1]^2 + b.p[2,2]^2)  ]
    d= Array(b.p * (x - b.c))
    cosθ = diag(n0*ax)
	n = n0 .* f_signs(d)
    # n = n0 .* copysign.(1.0,d)' #f_signs(d)
    absd = abs.(d)
    ∆ = (b.r .- absd) .* cosθ
    onbnd = f_onbnd(b,absd)
    isout = f_isout(b,absd)
    # projbnd =  all(.!isout .| onbnd)
    if count(isout) == 0  # x strictly inside box; ∆ all positive
        l∆x, i = findmin(∆)  # find closest face
        nout = n[i,:]
        ∆x = l∆x * nout
    else  # x outside box or on boundary in one or multiple directions
        ∆x = n' * (∆ .* isout)  # project out .!isout directions

        # Below, (.!isout .| onbnd) tests if each entry of ∆ (not ∆x) is either projected
        # out or on the boundary.  If the ith entry of ∆ is projected out, ∆x (not ∆) does
        # not have any component in the n[i,:] direction.  Even if the ith entry of ∆ is not
        # projected out, if it is too small then ∆x barely has any component in the n[i,:]
        # direction.  If all the entries of ∆ satisfy one of the two conditions, ∆x ≈ 0 and
        # we cannot use -∆x for nout.  In that case, take n'*onbnd as nout.  When x is
        # outside the box only in one dimension, n'*onbnd is basically the row of n along
        # that dimension, which reduces to normalize(-∆x) even if ∆x is very small.
        # nout = all(.!isout .| onbnd) ? n'*onbnd : -∆x
        # nout = normalize(nout)
		if all(.!isout .| onbnd)
            nout0 = n' * onbnd
        else
            nout0 = -∆x
        end
        nout = nout0 / norm(nout0)
    end

    return x+∆x, nout
end

translate(b::Box{N,N²,D}, ∆::SVector{N,T}) where {N,N²,D,T<:Real} = Box{N,N²,D,T}(b.c+∆, b.r, b.p, b.data)


text:
using DifferentiableNAS
using Flux
using Flux: throttle, logitcrossentropy, onecold, onehotbatch, Optimiser
using Zygote: @nograd
using StatsBase: mean
using Parameters
using CUDA
using Distributions
using BSON
using Dates

include("../CIFAR10.jl")
include("../training_utils.jl")

@show beginscript = now()

@show checkpointed = "test/models/darts_6642942/model29.bson"
BSON.@load checkpointed m_cpu

@show args = parse_commandline()

if args["random_seed"] > -1
    Random.seed!(args["random_seed"])
end

m = m_cpu |> gpu

#worth saving these?
optimiser_α = Optimiser(WeightDecay(1e-3),ADAM(3e-4,(0.5,0.999)))
optimiser_w = Optimiser(WeightDecay(3e-4),CosineAnnealing(args["epochs"]),Momentum(0.025, 0.9))

train, val = get_processed_data(args["val_split"], args["batchsize"], args["trainval_fraction"], args["random_seed"])
test = get_test_data(args["test_fraction"], args["random_seed"])

histepoch = historiessml()
histbatch = historiessml()
losses = [0.0, 0.0]

base_folder = prepare_folder("darts", args)

cbepoch = CbAll(CUDA.reclaim, histepoch, save_progress, CUDA.reclaim)
cbbatch = CbAll(CUDA.reclaim, histbatch, CUDA.reclaim)

for epoch in 1:args["epochs"]
    @show epoch
    display(Dates.format(convert(DateTime,now()-beginscript), "HH:MM:SS"))
    @time DARTStrain1st!(loss, m, train, val, optimiser_α, optimiser_w, losses, epoch; cbepoch = cbepoch, cbbatch = cbbatch)
end
display(("done", Dates.format(convert(DateTime,now()-beginscript), "HH:MM:SS")))

max_stars_repo_name: huangyxi/DifferentiableNAS

max_issues_repo_name: huangyxi/DifferentiableNAS

max_forks_repo_name: huangyxi/DifferentiableNAS


text:
using HDF5
using MapImages
using MAT

function solve_east_up(asc_img, desc_img, asc_los_map, desc_los_map)
    # @show size(asc_los_map), size(desc_los_map), size(asc_img), size(desc_img)

    east = similar(asc_img)
    up = similar(asc_img)
    # TEST: IF ONE IS FLIPPED
    # desc_img *= -1;

    for jj = 1:size(asc_los_map, 2)
        for ii = 1:size(asc_los_map, 1)
            asc_eu = asc_los_map[ii, jj, [1, 3]]
            desc_eu = desc_los_map[ii, jj, [1, 3]]

            A = hcat(asc_eu, desc_eu)'
            b = [asc_img[ii, jj]; desc_img[ii, jj]]
            # @show b

            # b = asc_eu[2] < 0 ? -b : b

            x = A \ b
            east[ii, jj] = x[1]
            up[ii, jj] = x[2]
        end
    end
    @show east[1:3]
    return east, up
end

LOS_MAP = "los_map.h5"

function solve_east_up(
    asc_path::AbstractString,
    desc_path::AbstractString,
    asc_fname::AbstractString,
    desc_fname::AbstractString = asc_fname,
    dset = "velos/1",
)
    asc_img = permutedims(h5read(joinpath(asc_path, asc_fname), dset))
    desc_img = permutedims(h5read(joinpath(desc_path, desc_fname), dset))

    asc_img, desc_img = MapImages._mask_asc_desc(asc_img, desc_img)

    asc_los_map = permutedims(h5read(joinpath(asc_path, LOS_MAP), "stack"), (2, 1, 3))
    desc_los_map = permutedims(h5read(joinpath(desc_path, LOS_MAP), "stack"), (2, 1, 3))
    return solve_east_up(asc_img, desc_img, asc_los_map, desc_los_map)
end

function plot_east_up(
    east,
    up;
    cmap = "seismic_wide_y",
    vm = 20,
    east_scale = 1.0,
    title = "",
    show_cbar = false,
)
    fig, axes = plt.subplots(1, 2, sharex = true, sharey = true)
    axim1 = axes[1].imshow(up, cmap = cmap, vmin = -vm, vmax = vm)
    show_cbar && fig.colorbar(axim1, ax=axes[1])

    vmeast = east_scale * vm
    axim2 = axes[2].imshow(east, cmap = cmap, vmin = -vmeast, vmax = vmeast)
    show_cbar && fig.colorbar(axim2, ax = axes[2])

    axes[1].set_title("up")
    axes[2].set_title("east")
    fig.suptitle(title)
    return fig, axes
end

function demo_east_up(
    fn = "velocities_prune_l1.h5";
    dset = "velos/1",
    vm = 20,
    east_scale = 1.0,
    shifta = 0.0,
    shiftd = 0.0,
    asc_path = "/data1/scott/pecos/path78-bbox2/igrams_looked_18/",
    desc_path = "/data4/scott/path85/stitched/igrams_looked_18/",
    cmap = "seismic_wide_y",
    title = nothing,
    show = true,
    show_cbar = true,
    output_basename=nothing,
)
    # asc_path, desc_path = ("/data1/scott/pecos/path78-bbox2/igrams_looked/", "/data4/scott/path85/stitched/igrams_looked/")
    asc_fname, desc_fname = map(x->joinpath(x, fn), (asc_path, desc_path))
    asc_img, desc_img = MapImages.find_overlaps(asc_fname, desc_fname, dset)

    asc_los_map = MapImage(joinpath(asc_path, LOS_MAP), dset_name = "stack")
    desc_los_map = MapImage(joinpath(desc_path, LOS_MAP), dset_name = "stack")
    #
    asc_idxs, desc_idxs = MapImages.find_overlap_idxs(asc_los_map, desc_los_map)
    # @show size(asc_los_map[asc_idxs..., :]), size(desc_los_map[desc_idxs..., :])

    east, up = solve_east_up(
        asc_img .+ shifta,
        desc_img .+ shiftd,
        asc_los_map[asc_idxs..., :],
        desc_los_map[desc_idxs..., :],
    )

    if show
        title = isnothing(title) ? "$fn: $dset" : title
        fig, axes = plot_east_up(
            east,
            up;
            cmap = cmap,
            vm = vm,
            title = title,
            east_scale = east_scale,
            show_cbar = show_cbar,
        )
    end
    # TODO: merge the geolists here...
    geolist = Sario.load_geolist_from_h5(asc_fname, dset)
    if !isnothing(output_basename)
        save_east_up_tifs(east.image, up.image, output_basename, east.demrsc, geolist=geolist)
    end

    # return east, up, fig, axes
    return east, up
end

function save_east_up_tifs(east_img, up_img, basename, demrsc; namelist=["east", "up"], geolist=nothing)
    for (img, name) in zip((east_img, up_img), namelist)
        if !isnothing(geolist)
            days = (geolist[end] - geolist[1]).value
            @show days
            img = img .* (days / 3650)
        end
        fname = name*basename
        if !endswith(fname, ".h5")
            fname *= ".h5"
        end

        println("Saving image to $fname")
        h5write(fname, name, permutedims(img))
        Sario.save_dem_to_h5(fname, demrsc)
        println("Running `aper hdf5-gtiff $fname`")
        run(`aper hdf5-gtiff $fname`)
    end
end

function eastup_insar(stations, east, up)
    easts, ups = [], []
    for s in stations
        rc = gps.station_rowcol(s, Dict(up.demrsc))
        push!(easts, east[rc...])
        push!(ups, up[rc...])
    end
    return easts, ups
end

function slope_mm_yr(dts, data)
    # Convert to "days since start" for line fitting
    gps_poly = fit_line(dts, data)
    slope = length(gps_poly) == 2 ? Polynomials.coeffs(gps_poly)[2] :
        Polynomials.coeffs(gps_poly)[1]
    return 365 * 10 * slope
end

function eastup_gps(stations, start_date = Date(2014, 11, 1), end_date = Date(2019, 1, 1))

    easts, ups = [], []
    for s in stations
        dts, east, north, up = get_gps_enu(s)
        push!(easts, slope_mm_yr(dts, east))
        push!(ups, slope_mm_yr(dts, up))
    end
    return easts, ups
end

function eastup_diffs(
    stations,
    east,
    up,
    start_date = Date(2014, 11, 1),
    end_date = Date(2019, 1, 1),
)
    return eastup_insar(stations, east, up) .- eastup_gps(stations, start_date, end_date)
end

function save_east_up_decomp(
    fname,
    shifta,
    shiftd;
    vm = 12,
    show = true,
    dset = "velos/1",
    outname = "zoom_pecos_vertical_east.mat",
    latrange = (30.901666673336, 31.90000000028),
    lonrange = (-104.0, -103.001666673056),
)
    # lats = read(ff, "lats");
    # lons = read(ff, "lons");
    east, up = demo_east_up(
        fname;
        full = true,
        dset = dset,
        shifta = shifta,
        shiftd = shiftd,
        vm = vm,
        show = show,
    )
    g = Sario.load_geolist_from_h5(fname, dset)
    days = (g[end] - g[1]).value
    @show days

    if isnothing(latrange) || isnothing(lonrange)
        eastcut = east .* (days / 3650)
        upcut = up .* (days / 3650)
    else
        eastcut = east[latrange, lonrange] .* (days / 3650)
        upcut = up[latrange, lonrange] .* (days / 3650)
    end
    gx, gy = MapImages.grid(eastcut.demrsc)

    println("Lat range: $(extrema(gy))")
    println("Lon range: $(extrema(gx))")
    println("Saving lats, lons, ups, easts to $outname")

    dd = Dict("lats" => gy, "lons" => gx, "up" => upcut.image, "east" => eastcut.image)
    isnothing(outname) && return dd
    matwrite(outname, dd)
    return dd
end

function save_all_east_up(outname = outname)
    dd_all = Dict{String,Any}()

    # p78 = "/data1/scott/pecos/path78-bbox2/igrams_looked/"
    # p85 = "/data4/scott/path85/stitched/igrams_looked/"

    varnames, values = [], []
    for year = 2016:2018
        fname = "velocities_$(year)_current.h5"
        dd = save_east_up_decomp(
            fname,
            0,
            0,
            dset = "velos_shifted/1",
            show = false,
            outname = nothing,
            latrange = nothing,
        )
        dd_all["lats"] = dd["lats"]
        dd_all["lons"] = dd["lons"]
        dd_all["east_$year"] = dd["east"]
        dd_all["up_$year"] = dd["up"]
        dd_all["east_$year"] = dd["east"]
    end
    matwrite(outname, dd_all)
    return dd_all
end

station_overlap = ["TXMH", "TXFS", "TXAD", "TXS3", "NMHB", "TXKM"]
eeups(fname, shifta, shiftd, dset = "velos/1", full = true; show = true) =
    extrema.(eastup_insar(
        station_overlap,
        demo_east_up(
            fname;
            full = full,
            dset = dset,
            shifta = shifta,
            shiftd = shiftd,
            show = show,
        )...,
    ))
eeups_diff(fname, shifta, shiftd, dset = "velos/1", full = true; show = true) =
    eastup_diffs(
        station_overlap,
        demo_east_up(
            fname;
            full = full,
            dset = dset,
            shifta = shifta,
            shiftd = shiftd,
            show = show,
        )...,
    )

# h5write("zoom_pecos_vertical_east.h5", "east_18", permutedims(read(ff3, "east")))


text:
# User-level methods to compute lightcurves

struct Lightcurve{T<:Real}
    dt::T             # Exposure time
    tobs::Vector{T}   # Observed times
    fobs::Vector{T}   # Observed flux
    eobs::Vector{T}   # Measurement errors
    nobs::Int64       # number of flux measurements
    flux::Vector{T}   # Computed model flux
    dfdu::Matrix{T}   # Derivative of flux wrt limbdark coefficients
    dfdk::Matrix{T}   # Derivatives wrt the radius ratios
    dfdq0::Matrix{T}  # Derivatives wrt initial Nbody Cartesian coordinates and masses
    dfdr::Vector{T}   # Derivatives wrt stellar radius

    # Transit parameters
    u_n::Vector{T}   # Limbdark coefficients
    k::Vector{T}     # radius ratios
    rstar::Vector{T} # Stellar radius

    # Interal arrays/values
    dtinv::T          # Inverse of the exposure time
    dbdq0::Vector{T}  # derivative of the impact parameter wrt the Nbody initial conditions
    n_params::Int64   # Number of nbody model parameters
    do_grad::Bool     # Compute gradients or not

    function Lightcurve(dt::T, tobs::Vector{T}, fobs::Vector{T}, eobs::Vector{T}, u_n::Vector{T}, k::Vector{T}, rstar::T, n_params::Int64=0) where T<:Real
        @assert (length(tobs) == length(fobs)) && (length(tobs) == length(eobs)) "Data arrays are different sizes"
        @assert n_params >= 0 "Number of model parameters must be a positive integer"
        n_params > 0 ? do_grad = true : do_grad = false
        nobs = length(tobs)
        flux = zeros(T,nobs)
        dtinv = dt == 0.0 ? 0.0 : inv(dt)
        dfdu = zeros(T, nobs, length(u_n))
        dfdk = zeros(T, nobs, length(k))
        dfdq0 = zeros(T, nobs, n_params)
        dbdq0 = zeros(T, n_params)
        dfdr = zeros(T, nobs)
        return new{T}(dt,tobs,fobs,eobs,nobs,flux,dfdu,dfdk,dfdq0,dfdr,u_n,k,[rstar],dtinv,dbdq0,n_params,do_grad)
    end
end

"""Zero out the model arrays"""
function zero_out!(lc::Lightcurve{T}) where T<:Real
    lc.dfdu .= 0.0
    lc.dfdk .= 0.0
    lc.dfdq0 .= 0.0
    lc.dfdr .= 0.0
    lc.flux .= 0.0
end

# Normalize the exposure integration by the exposure time.
function normalize!(lc::Lightcurve{T}, ia::IntegralArrays{T}) where T<:Real
    lc.flux .*= lc.dtinv # Divide by exposure time to get average flux
    if lc.do_grad
        # Do the same for derivatives
        lc.dfdk .*= lc.dtinv
        lc.dfdu .*= lc.dtinv
        lc.dfdq0 .*= lc.dtinv
        lc.dfdr .*= lc.dtinv
    end
end

# If non-integrated exposure, do nothing
@inline normalize!(lc::Lightcurve{T}, ia::T) where T<:Real = nothing

function find_transit_time(t0::T, h::T, points::AbstractMatrix{T}) where T<:Real
    tt = find_zero(t->compute_impact_parameter(t, t0, h, points), t0)
    return tt
end

function points_of_contact_4(tt::T,t0::T,h::T,points::AbstractMatrix{T},k::T) where T<:Real
    t1 = find_zero(t -> (1.0+k-compute_impact_parameter(t,t0,h,points)), tt-h)
    t2 = find_zero(t -> (1.0-k-compute_impact_parameter(t,t0,h,points)), tt-h)
    t3 = find_zero(t -> (1.0-k-compute_impact_parameter(t,t0,h,points)), tt+h)
    t4 = find_zero(t -> (1.0+k-compute_impact_parameter(t,t0,h,points)), tt+h)
    return SVector(t1,t2,t3,t4)
end

function points_of_contact_2(t0::T,tt::T,h::T,points::AbstractMatrix{T},k::T) where T<:Real
    t1 = find_zero(t -> (1.0+k-compute_impact_parameter(t,t0,h,points)), tt-h)
    t4 = find_zero(t -> (1.0+k-compute_impact_parameter(t,t0,h,points)), tt+h)
    return SVector(t1,t4)
end

function compute_lightcurve!(lc::Lightcurve{T}, ts::TransitSeries{T, ProvidedTimes}; tol::T=1e-6, maxdepth::Int64=6) where T<:Real

    zero_out!(lc) # Zero out model arrays

    # Check if we're doing an integrated lightcurve
    ia = lc.dt == 0.0 ? 0.0 : IntegralArrays(lc.do_grad ? (lc.n_params + length(lc.u_n) + length(lc.k) + 2) : 1, maxdepth, tol) # Plus 2 for flux and rstar

    # Make transit structure (will be updated with proper r and b later)
    trans = transit_init(lc.k[1], 0.0, lc.u_n, lc.do_grad)

    # Iterate over each transit time and sum Lightcurve
    rstar = lc.rstar[1]
    for it in eachindex(ts.times)
        # check for transit
        t0 = ts.times[it]   # Get "data" transit time (expansion point)
        ib = ts.bodies[it]  # Get transiting body

        tt = try
            # Compute the transit time by finding where the x component of
            # the impact parameter is equal to zero
            # This assumes the body crosses the mid plane during the transit
            find_transit_time(t0, ts.h, @views(ts.points[ib,it,:,:]./rstar))
        catch e
            # If no convergence, just use the supplied transit time
            t0
        end

        # Get the impact parameter at the transit midpoint (computed above)
        b0 = compute_impact_parameter(t0, tt, ts.h, @views(ts.points[ib,it,:,:]./rstar))
        if b0 > 1.0+lc.k[ib-1]; continue; end

        # Compute points of contact
        # If grazing transit, only two points
        if (b0 + lc.k[ib-1]) >= 1.0
            tc = points_of_contact_2(t0, tt, ts.h, @views(ts.points[ib,it,:,:]./rstar), lc.k[ib-1])
        else
            tc = points_of_contact_4(t0, tt, ts.h, @views(ts.points[ib,it,:,:]./rstar), lc.k[ib-1])
        end

        # Integrate lightcurve
        trans.r = lc.k[ib-1]
        integrate_transit!(ib,it,t0,tc,trans,lc,ts,ia)
    end

    # Normalize by exposure time, if needed.
    normalize!(lc, ia)
    return
end

function integrate_transit!(ib::Int64,it::Int64,t0::T,tc::SVector{N,T},trans::Transit_Struct{T},lc::Lightcurve{T},ts::TransitSeries{T},ia::IntegralArrays{T}) where {N, T<:Real}
    nc = length(tc) # Number of points of contact
    dt = lc.dt
    inv_rstar = inv(lc.rstar[1])

    # Integrate over each Exposure
    for i in 1:lc.nobs
        tstart = lc.tobs[i] - 0.5*dt
        tend = lc.tobs[i] + 0.5*dt

        # Check if remaining exposures are inside transit
        if tstart > tc[end]; break; end # Don't need to continue loop if passed transit
        if tend < tc[1]; continue; end

        # Check if points of contact are within exposure
        tlim = [tstart]
        for j in 1:nc
            if tstart < tc[j] && tc[j] < tend
                push!(tlim, tc[j])
            end
        end
        push!(tlim, tend)

        # Get series expansion components
        xc = components(@views(ts.points[ib,it,:,1].*inv_rstar), ts.h)
        yc = components(@views(ts.points[ib,it,:,2].*inv_rstar), ts.h)

        if lc.do_grad
            n_bodies = length(ts.count)
            dxc = [components(ts.dpoints[ib,it,:,1,k,i].*inv_rstar, ts.h) for i in 1:7, k in 1:n_bodies][:]
            dyc = [components(ts.dpoints[ib,it,:,2,k,i].*inv_rstar, ts.h) for i in 1:7, k in 1:n_bodies][:]

            # integrate over exposure
            for j in 1:length(tlim)-1
                integrate_timestep!(t0, tlim[j], tlim[j+1], xc, yc, dxc, dyc, trans, ia, lc.dbdq0, ib-1)
                lc.flux[i] += ia.I_of_f[1]
                lc.dfdq0[i,:] .+= ia.I_of_f[2:1+n_bodies*7]
                lc.dfdk[i,:] .+= ia.I_of_f[2+n_bodies*7:n_bodies*8]
                lc.dfdu[i,:] .+= trans.dgdu' * ia.I_of_f[end-trans.n-1:end-1]
                lc.dfdr[i] += ia.I_of_f[end] * inv_rstar
            end
        else
            # Integrate over exposure
            for j in 1:length(tlim)-1
                integrate_timestep!(t0, tlim[j], tlim[j+1], xc, yc, trans, ia)
                lc.flux[i] += ia.I_of_f[1]
            end
        end
    end
    return
end

function integrate_transit!(ib::Int64,it::Int64,t0::T,tc::SVector{N,T},trans::Transit_Struct{T},lc::Lightcurve{T},ts::TransitSeries{T},ia::T) where {N, T<:Real}
    inv_rstar = inv(lc.rstar[1])

    # Compute the flux at each point
    for i in 1:lc.nobs
        # Check if observation is outside of a transit
        if lc.tobs[i] > tc[end]; break; end
        if lc.tobs[i] < tc[1]; continue; end

        # Get series expansion components
        xc = components(@views(ts.points[ib,it,:,1].*inv_rstar), ts.h)
        yc = components(@views(ts.points[ib,it,:,2].*inv_rstar), ts.h)

        if lc.do_grad
            n_bodies = length(ts.count)
            dxc = [components(ts.dpoints[ib,it,:,1,k,i].*inv_rstar, ts.h) for i in 1:7, k in 1:n_bodies][:]
            dyc = [components(ts.dpoints[ib,it,:,2,k,i].*inv_rstar, ts.h) for i in 1:7, k in 1:n_bodies][:]

            compute_flux!(lc.tobs[i], t0, xc, yc, dxc, dyc, lc, trans, i, ki, inv_rstar)
        else
            lc.flux[i] += compute_flux(lc.tobs[i], t0, xc, yc, trans)
        end
    end
    return
end

function integrate_timestep!(t0::T, a::T, b::T, xc::SVector{N,T}, yc::SVector{N,T}, trans::Transit_Struct{T}, ia::IntegralArrays{T}) where {N, T<:Real}
    # Computes the flux as function of time.
    # Closure to be passed to integrator function
    transit_flux! = let trans=trans, t0=t0, xc=xc, yc=yc
        (time::T, flux::Vector{T}) -> begin
            trans.b = compute_impact_parameter(time, t0, xc, yc)
            flux[1] = transit_poly_g(trans)-1
        end
    end

    # Integrate transit_flux! over interval [a,b]
    integrate_simpson!(a,b,transit_flux!,ia)
end

function compute_flux(tc::T, t0::T, xc::SVector{N,T}, yc::SVector{N,T}, trans::Transit_Struct{T}) where {N, T<:Real}
    # Compute flux at a particular time
    trans.b = compute_impact_parameter(tc, t0, xc, yc)
    flux = transit_poly_g(trans) - 1
    return flux
end

function integrate_timestep!(t0::T, a::T, b::T, xc::SVector{N,T}, yc::SVector{N,T}, dxc, dyc, trans::Transit_Struct{T}, ia, dbdq0, ki) where {N,T<:Real}
    n_coords = length(dbdq0)

    # Map radius ratio index to the gradient array index.
    k_ind = ki + 1 + n_coords

    # Compute the flux and derivatives
    transit_flux_grad! = let trans=trans, t0=t0, xc=xc, yc=yc, dbdq0=dbdq0, n_coords = n_coords, k_ind=k_ind
        (time::T, dflux::Vector{T}) -> begin
            trans.b = compute_impact_parameter!(time, t0, xc, yc, dxc, dyc, dbdq0)
            dflux .= 0.0
            dflux[1] = transit_poly_g!(trans)-1            # Flux at time
            dflux[2:1+n_coords] .= trans.dfdrb[2] .* dbdq0 # Gradient wrt the initial conditions
            dflux[k_ind] = trans.dfdrb[1]                  # radius ratio (derivative of others are zero)
            dflux[end-trans.n-1:end-1] .= trans.dfdg       # dfdg (greens basis)
            dflux[end] = -trans.dfdrb[2]*trans.b           # dfdrstar * rstar (divide by rstar outside)
        end
    end

    # Integrate flux and derivatives over exposure interval [a,b]
    integrate_simpson!(a,b,transit_flux_grad!,ia)
end

function compute_flux!(tc::T, t0::T, xc, yc, dxc, dyc, lc, trans, i, ki, inv_rstar) where T<:Real
    # Compute flux and derivatives
    trans.b = compute_impact_parameter!(tc, t0, xc, yc, dxc, dyc, lc.dbdq0)
    lc.flux[i] += transit_poly_g!(trans) - 1            # Flux
    lc.dfdq0[i,:] .+= trans.dfdrb[2] .* lc.dbdq0        # Derivative wrt initial conditions
    lc.dfdk[i,ki] += trans.dfdrb[1]                     # Radius ratio
    lc.dfdu[i,:] .+= trans.dgdu' * trans.dfdg           # Limbdarkening params
    lc.dfdr[i] += -trans.dfdrb[2]*trans.b*inv_rstar     # Stellar radius
    return
end


text:
using BinaryBuilder

name = "Sundials"
version = v"3.1.2"

# Collection of sources required to build Sundials
sources = [
    GitSource("https://github.com/LLNL/sundials.git", "3bea69fccb5d5dc35ad030a1c05f27f316461530"),
    DirectorySource("./bundled"),
]

# Bash recipe for building across all platforms
script = raw"""
cd $WORKSPACE/srcdir/sundials*
atomic_patch -p0 $WORKSPACE/srcdir/patches/Sundials_windows.patch

CMAKE_FLAGS=(-DCMAKE_INSTALL_PREFIX=${prefix} -DCMAKE_TOOLCHAIN_FILE="${CMAKE_TARGET_TOOLCHAIN}")
CMAKE_FLAGS+=(-DCMAKE_BUILD_TYPE=Release -DEXAMPLES_ENABLE_C=OFF)
CMAKE_FLAGS+=(-DKLU_ENABLE=ON -DKLU_INCLUDE_DIR="$prefix/include" -DKLU_LIBRARY_DIR="$libdir")

### Disable BLAS and LAPACK because sundials 3.1 cannot use 64-bit ints with these libraries
#CMAKE_FLAGS+=(-DBLAS_ENABLE=ON -DLAPACK_ENABLE=ON)
#if [[ ${nbits} == 64 ]] && [[ ${target} != aarch64* ]]; then
#    patch -p0 < $WORKSPACE/srcdir/patches/Sundials_ilp64.patch
#    CMAKE_FLAGS+=(-DBLAS_LIBRARIES="-L${libdir} -lopenblas64_" -DLAPACK_LIBRARIES="-L${libdir} -lopenblas64_")
#else
#    CMAKE_FLAGS+=(-DBLAS_LIBRARIES="-L${libdir} -lopenblas" -DLAPACK_LIBRARIES="-L${libdir} -lopenblas")
#fi

if [[ ${target} != *darwin* ]]; then
    # Needed to find libgfortran for OpenBLAS.
    export CFLAGS="-Wl,-rpath-link,/opt/${target}/${target}/lib -Wl,-rpath-link,/opt/${target}/${target}/lib64"
fi

mkdir build
cd build
cmake "${CMAKE_FLAGS[@]}" ..
make -j${nproc}
make install

# Move libraries to ${libdir} on Windows
if [[ "${target}" == *-mingw* ]]; then
    mv ${prefix}/lib/libsundials_*.${dlext} "${libdir}"
fi
"""

# We attempt to build for all defined platforms
platforms = supported_platforms()

products = [
    LibraryProduct("libsundials_arkode", :libsundials_arkode),
    LibraryProduct("libsundials_cvode", :libsundials_cvode),
    LibraryProduct("libsundials_cvodes", :libsundials_cvodes),
    LibraryProduct("libsundials_ida", :libsundials_ida),
    LibraryProduct("libsundials_idas", :libsundials_idas),
    LibraryProduct("libsundials_kinsol", :libsundials_kinsol),
    LibraryProduct("libsundials_nvecserial", :libsundials_nvecserial),
    LibraryProduct("libsundials_sunlinsolband", :libsundials_sunlinsolband),
    LibraryProduct("libsundials_sunlinsoldense", :libsundials_sunlinsoldense),
    LibraryProduct("libsundials_sunlinsolklu", :libsundials_sunlinsolklu),
    LibraryProduct("libsundials_sunlinsolpcg", :libsundials_sunlinsolpcg),
    LibraryProduct("libsundials_sunlinsolspbcgs", :libsundials_sunlinsolspbcgs),
    LibraryProduct("libsundials_sunlinsolspfgmr", :libsundials_sunlinsolspfgmr),
    LibraryProduct("libsundials_sunlinsolspgmr", :libsundials_sunlinsolspgmr),
    LibraryProduct("libsundials_sunlinsolsptfqmr", :libsundials_sunlinsolsptfqmr),
    LibraryProduct("libsundials_sunmatrixband", :libsundials_sunmatrixband),
    LibraryProduct("libsundials_sunmatrixdense", :libsundials_sunmatrixdense),
    LibraryProduct("libsundials_sunmatrixsparse", :libsundials_sunmatrixsparse),
]

dependencies = [
    Dependency("OpenBLAS_jll"),
    Dependency("SuiteSparse_jll"),
]

# Build the tarballs, and possibly a `build.jl` as well.
build_tarballs(ARGS, name, version, sources, script, platforms, products, dependencies)

max_stars_repo_name: sharanry/Yggdrasil

max_issues_repo_name: sharanry/Yggdrasil

max_forks_repo_name: sharanry/Yggdrasil