text:
export Polynomial, coefficients, exponents, nvariables, variables,
    parameters, nparameters, coefficienttype, scale_coefficients!, permutation

"""
    Polynomial{T, SE<:SExponents}

A Polynomial with coefficents in `T` in `NVars` variables and exponents of type `SE`.

    Polynomial(f::MP.AbstractPolynomial, [variables])

Construct a Polynomial from `f`.
"""
struct Polynomial{T, SE, Params}
    coefficients::Vector{T}
    variables::Vector{Symbol}
    parameters::Union{Nothing, Vector{Symbol}}
    perm::Vector{Int}

    function Polynomial{T, SE, P}(coefficients::Vector{T}, variables, parameters, perm::Vector{Int}) where {T, SE, P}
        # @assert length(coefficients) == size(SE, 2) "Coefficients size does not match exponents size"
        new(coefficients, variables, parameters, perm)
    end
end

function Polynomial(coeffs::Vector{T}, SE::SExponents, variables, PE::Union{Nothing,SExponents}, parameters, perm::Vector{Int}) where {T}
    Params = PE isa Nothing ? Nothing : PE
    Polynomial{T, SE, Params}(coeffs, variables, parameters, perm)
end

function Polynomial(p::MP.AbstractPolynomialLike; parameters=nothing, variables=diffvars(MP.variables(p), parameters))
    Polynomial(p, variables, parameters)
end

Polynomial(p::MP.AbstractPolynomialLike, vars) = Polynomial(p, vars, nothing)
@deprecate Polynomial(p, vars) Polynomial(p, variables=vars)

function Polynomial(p::MP.AbstractPolynomialLike, vars, params)
    terms = MP.terms(p)
    nterms = length(terms)
    nvars = length(vars)
    nparams = params === nothing ? 0 : length(params)

    if MP.nvariables(p) > (nvars + nparams)
        error("Number of variables is less than the number of declared variables and parameters.")
    end

    coefficients = MP.coefficient.(MP.terms(p))
    exponents = [MP.degree(term, var) for var in vars, term in terms]
    if params === nothing
        parameters = nothing
        pexponents = nothing
    else
        parameters = Symbol.(params)
        pexponents = [MP.degree(term, param) for param in params, term in terms]
    end

    Polynomial(coefficients, exponents, Symbol.(vars), pexponents, parameters)
end
lengthparams(::Nothing) = 0
lengthparams(params) = length(params)

function Polynomial(coeffs::Vector, E::Matrix{<:Integer}, variables::Vector{Symbol}=defaultvariables(size(E, 1)),
                    PE::Union{Nothing, Matrix{<:Integer}}=nothing, parameters=nothing)
    @assert length(coeffs) == size(E, 2) "Number of coefficients doesn't match number of terms."
    if PE === nothing
        p = revlexicographic_cols_perm(E)
        SPE = nothing
    else
        p = revlexicographic_cols_perm([PE; E])
        SPE = SExponents(PE[:, p])
    end
    return Polynomial(coeffs[p], SExponents(E[:,p]), variables, SPE, parameters, p)
end

defaultvariables(n) = [Symbol("x", i) for i=1:n]

"""
    coefficients(f)

Return the coefficients of `f`.
"""
coefficients(f::Polynomial) = f.coefficients

"""
    exponents(f)

Return the exponents of `f` as an matrix where each column represents
the exponents of a monomial.
"""
exponents(::Polynomial{T, E}) where {T, E} = exponents(E)

"""
    parameter_exponents(f)::Union{Nothing, Matrix{Int}}

Return the exponents of the parameters of `f` as an matrix where each column represents
the exponents of a monomial. If `f` has no parameters `nothing` is returned.
"""
parameter_exponents(::Polynomial{T, E, P}) where {T, E, P} = exponents(P)
parameter_exponents(::Polynomial{T, E, Nothing}) where {T, E} = nothing

sexponents(::Polynomial{T, E}) where {T, E} = E

"""
    variables(f::Polynomial)::Vector{Symbol}

Returns the variables of `f`.
"""
variables(f::Polynomial) = f.variables

"""
    parameters(f::Polynomial)::Union{Nothing, Vector{Symbol}}

Returns the parameters of `f`.
"""
parameters(f::Polynomial) = f.parameters

"""
    nvariables(f::Polynomial)

Return the number of variables of `f`.
"""
nvariables(::Polynomial{T, E}) where {T, E} = size(E)[1]

"""
    nparameters(f::Polynomial)

Return the number of parameters of `f`.
"""
nparameters(f::Polynomial) = f.parameters === nothing ? 0 : length(f.parameters)

"""
    coefficienttype(f::Polynomial)

Return the type of the coefficients of `f`.
"""
coefficienttype(::Polynomial{T}) where {T} = T

"""
    permutation(f::Polynomial)

Return the permutation applied to the support of the polynomial system.
"""
permutation(f::Polynomial) = f.perm

function Base.:(==)(f::Polynomial{T, E1, P1}, g::Polynomial{T, E2, P2}) where {T, E1, E2, P1, P2}
    P1 == P2 && E1 == E2 && coefficients(f) == coefficients(g)
end

"""
    scale_coefficients!(f, λ)

Scale the coefficients of `f` by the factor `λ`.
"""
scale_coefficients!(f::Polynomial, λ) = LinearAlgebra.rmul!(f.coefficients, λ)


text:
include("chainrules.jl")
include("transform_fourier.jl")
include("transform_chebyshev.jl")
include("transform_legendre.jl")
include("operators_convolution.jl")
include("operators_covariance.jl")
include("operators_spectralkernel.jl")

ext: jl


text:
using OptimJulia
using Base.Test
using SymPy

x1,x2 = Sym("x1,x2")

f1 = ((((1 - 8x1) + 2*x1^2) - (10/3)*x1^3) + (1/4)*x1^4 + (4/5)*x1^5) - (1/6)*x1^6

Newton(f1, x1, 3)

vars = [x1,x2]
f2 = -log(1 - x1 - x2) - log(x1) - log(x2)
x0 = [0.85,0.05]

Newton(f2, vars, x0)

for i = 2:10
      QuasiNewton(f1, x1, x0 = [i i+4])
      println("=========================================")
      println("Test $i with QuasiNewton algorithm passed")
end


text:
# module for representing agent
# decide control order to robot
# used for fastslam 2.0
# motion update considering observation

mutable struct FastSlam2Agent
  speed
  yaw_rate
  time_interval
  estimator
  prev_spd
  prev_yr

  # init
  function FastSlam2Agent(speed::Float64, yaw_rate::Float64;
                          time_interval::Float64=0.1,
                          estimator=nothing)
      self = new()
      self.speed = speed
      self.yaw_rate = yaw_rate
      self.time_interval = time_interval
      self.estimator = estimator
      self.prev_spd = 0.0
      self.prev_yr = 0.0
      return self
  end
end

function draw_decision!(self::FastSlam2Agent, observation)
  if self.estimator !== nothing
    motion_update(self.estimator, self.prev_spd, self.prev_yr,
                  self.time_interval, observation)
    self.prev_spd, self.prev_yr = self.speed, self.yaw_rate
    observation_update(self.estimator, observation)
    draw!(self.estimator)
  end
  return self.speed, self.yaw_rate
end

max_stars_repo_name: ShisatoYano/JuliaAutonomy

max_issues_repo_name: ShisatoYano/JuliaAutonomy

max_forks_repo_name: ShisatoYano/JuliaAutonomy


text:
"""
points are in 3 x n matrix
"""
function applyT(T::Matrix{FT1}, points::Matrix{FT2}
    ) where {FT1<:AbstractFloat, FT2<:AbstractFloat}
    return T[1:3, 1:3] * points .+ reshape(T[1:3, 4], (3, 1))
end

"""
points are in 3 x n matrix
function applyT!(T::Matrix{FT}, points::Matrix{FT})
"""
function applyT!(
    T::Matrix{FT1},
    points::Matrix{FT2}
) where {FT1<:AbstractFloat, FT2<:AbstractFloat}
    points .= T[1:3, 1:3] * points .+ reshape(T[1:3, 4], (3, 1))
end

function applyT(T::Matrix{Float64}, point::GeometryTypes.Point3f0)
    return GeometryTypes.Point3f0(T[1:3, 1:3] * point + T[1:3, 4])
end

function applyT(T::Matrix{Float64}, points::Vector{GeometryTypes.Point3f0})
    res = map(pt->applyT(T, pt), points)
    return res
end


text:
export iscompatible

import LinearAlgebraX

## 1. Hypercube and Operation

"""
    iscompatible(hypercube, op::TranslationOperation{<:Integer})

Check whether the translation operation is compatible with the hypercube,
i.e. leaves the superlattice invariant. If the dimensions do not match, throw exception.
Otherwise always return true.

"""
function iscompatible(hypercube::Hypercube, op::TranslationOperation{<:Integer})::Bool
    if dimension(hypercube) != dimension(op)
        throw(DimensionMismatch("hypercube and operation have different dimension"))
    end
    return true
end

"""
    iscompatible(hypercube, op::PointOperation{<:Integer})

Check whether the point operation is compatible with the hypercube,
i.e. leaves the superlattice invariant.
"""
function iscompatible(hypercube::Hypercube, op::PointOperation{<:Integer})::Bool
    R, r = hypercube.wrap(op.matrix * hypercube.shape_matrix)
    return iszero(r) && isone(abs(LinearAlgebraX.detx(R)))
end

## 2. Hypercube and Symmetry

"""
    iscompatible(hypercube, symmetry::FiniteTranslationSymmetry)

Check whether the translation symmetry is compatible with the hypercube,
i.e. has the same superlattice.
"""
function iscompatible(hypercube::Hypercube, tsym::FiniteTranslationSymmetry)::Bool
    R, r = hypercube.wrap(tsym.hypercube.shape_matrix)
    return iszero(r) && isone(abs(LinearAlgebraX.detx(R)))
end

"""
    iscompatible(hypercube, symmetry::PointSymmetry)

Check whether the point symmetry is compatible with the hypercube,
i.e. all of its elements leaves the superlattice invariant.
"""
function iscompatible(hypercube::Hypercube, psym::PointSymmetry)::Bool
    return all(iscompatible(hypercube, pop) for pop in elements(psym))
end

## 3. FiniteTranslationSymmetry and PointOperation/PointSymmetry

"""
    iscompatible(tsym::FiniteTranslationSymmetry, pop::PointOperation{<:Integer})

Check whether the point operation is compatible with the translation symmetry,
i.e. leaves the superlattice invariant.
"""
function iscompatible(tsym::FiniteTranslationSymmetry, pop::PointOperation{<:Integer})::Bool
    return iscompatible(tsym.hypercube, pop)
end

"""
    iscompatible(tsym::FiniteTranslationSymmetry, psym::PointSymmetry)

Check whether the point symmetry is compatible with the translation symmetry,
i.e. all of its elements leaves the superlattice invariant.
"""
function iscompatible(tsym::FiniteTranslationSymmetry, psym::PointSymmetry)::Bool
    return iscompatible(tsym.hypercube, psym)
end

"""
    iscompatible(psym::PointSymmetry, tsym::FiniteTranslationSymmetry)

Check whether the point symmetry is compatible with the translation symmetry,
i.e. all of its elements leaves the superlattice invariant.
"""
function iscompatible(psym::PointSymmetry, tsym::FiniteTranslationSymmetry)::Bool
    return iscompatible(tsym, psym)
end

## 4. Translation Irreps and PointOperation/PointSymmetry

"""
    iscompatible(tsym, tsym_irrep_index, point_operation)

Check whether the point operation is compatible with the given irrep of translation (a.k.a. momentum).
In specific, check whether
(1) the point operation is compatible with the translation symmetry, and
(2) the momentum is invariant under the point operation.
"""
function iscompatible(
    tsym::FiniteTranslationSymmetry,
    tsym_irrep_index::Integer,
    pop::PointOperation{<:Integer},
)::Bool
    !iscompatible(tsym, pop) && return false
    reciprocal_matrix = LinearAlgebraX.invx(transpose(pop.matrix))
    k1 = tsym.fractional_momenta[tsym_irrep_index]
    k2 = (x -> mod(x, 1)).(reciprocal_matrix * k1)
    return k1 == k2
end

"""
    iscompatible(tsym, tsym_irrep_index, psym)

Check whether the point symmetry is compatible with the given irrep of translation (a.k.a. momentum).
In specific, check whether
(1) the point operation is compatible with the translation symmetry, and
(2) the momentum is invariant under all point operations.
"""
function iscompatible(
    tsym::FiniteTranslationSymmetry,
    tsym_irrep_index::Integer,
    psym::PointSymmetry,
)::Bool
    return all(iscompatible(tsym, tsym_irrep_index, pop) for pop in generator_elements(psym))
end

## 5. Lattice and Operation

"""
    iscompatible(lattice, translation_operation)
"""
function iscompatible(lattice::Lattice, op::TranslationOperation{<:Integer})
    if dimension(lattice) != dimension(op)
        throw(DimensionMismatch("lattice and op have different dimension"))
    end
    return true
end

"""
    iscompatible(lattice, point_operation)
"""
function iscompatible(lattice::Lattice, op::PointOperation{<:Integer})
    return iscompatible(lattice.hypercube, op) &&
           !isnothing(findsitemap(lattice.unitcell, op))
end

## 6. Lattice and Symmetry

"""
    iscompatible(lattice, translation_symmetry)

Test whether lattice and the symmetry are compatible.
For translation symmetry, this means that the hypercubic lattice for both are the same.
"""
function iscompatible(lattice::Lattice, tsym::FiniteTranslationSymmetry)
    return iscompatible(lattice.hypercube, tsym)
end

"""
    iscompatible(lattice, point_symmetry)
"""
function iscompatible(lattice::Lattice, psym::PointSymmetry)::Bool
    return (
        iscompatible(lattice.hypercube, psym) &&
        all(pop -> !isnothing(findsitemap(lattice.unitcell, pop)), generator_elements(psym))
    )
end


text:
using Test
using ModelingToolkit, OrdinaryDiffEq

# Basic electric components
@parameters t
@connector function Pin(;name)
    @variables v(t) i(t)
    ODESystem(Equation[], t, [v, i], [], name=name, defaults=[v=>1.0, i=>1.0])
end

function ModelingToolkit.connect(::Type{Pin}, ps...)
    eqs = [
           0 ~ sum(p->p.i, ps) # KCL
          ]
    # KVL
    for i in 1:length(ps)-1
        push!(eqs, ps[i].v ~ ps[i+1].v)
    end

    return eqs
end

function Ground(;name)
    @named g = Pin()
    eqs = [g.v ~ 0]
    ODESystem(eqs, t, [], [], systems=[g], name=name)
end

function ConstantVoltage(;name, V = 1.0)
    val = V
    @named p = Pin()
    @named n = Pin()
    @parameters V
    eqs = [
           V ~ p.v - n.v
           0 ~ p.i + n.i
          ]
    ODESystem(eqs, t, [], [V], systems=[p, n], defaults=Dict(V => val), name=name)
end

function Resistor(;name, R = 1.0)
    val = R
    @named p = Pin()
    @named n = Pin()
    @variables v(t)
    @parameters R
    eqs = [
           v ~ p.v - n.v
           0 ~ p.i + n.i
           v ~ p.i * R
          ]
    ODESystem(eqs, t, [v], [R], systems=[p, n], defaults=Dict(R => val), name=name)
end

function Capacitor(;name, C = 1.0)
    val = C
    @named p = Pin()
    @named n = Pin()
    @variables v(t)
    @parameters C
    D = Differential(t)
    eqs = [
           v ~ p.v - n.v
           0 ~ p.i + n.i
           D(v) ~ p.i / C
          ]
    ODESystem(eqs, t, [v], [C], systems=[p, n], defaults=Dict(C => val), name=name)
end

function Inductor(; name, L = 1.0)
    val = L
    @named p = Pin()
    @named n = Pin()
    @variables v(t) i(t)
    @parameters L
    D = Differential(t)
    eqs = [
           v ~ p.v - n.v
           0 ~ p.i + n.i
           i ~ p.i
           D(i) ~ v / L
          ]
    ODESystem(eqs, t, [v, i], [L], systems=[p, n], defaults=Dict(L => val), name=name)
end


text:
using Test
using Dash
using Dash:dash_env, @env_default!

@testset "env" begin
    @test isnothing(dash_env("test"))
    @test dash_env("test", "aaaa") == "aaaa"
    @test isnothing(dash_env(Int, "test"))
    @test dash_env(Int, "test", 10) == 10
    ENV["DASH_STRING_TEST"] = "test_path"
    @test dash_env("string_test", "aaaa") == "test_path"
    @test_throws ArgumentError dash_env(Int, "string_test", "aaaa") == "test_path"

    string_test = nothing
    @env_default! string_test
    @test string_test == "test_path"
    string_test = "aaaa"
    @env_default! string_test
    @test string_test == "aaaa"

    ENV["DASH_INT_TEST"] = "100"
    @test dash_env("int_test", "aaaa") == "100"
    @test dash_env(Int, "int_test", 50) == 100

    int_test = nothing
    @env_default! int_test Int
    @test int_test == 100
    int_test = 50
    @env_default! int_test Int
    @test int_test == 50
    int_test2 = nothing
    @env_default! int_test2 Int 40
    @test int_test2 == 40

    ENV["DASH_BOOL_TEST"] = "1"
    @test dash_env(Bool, "bool_test", 50) == true

    ENV["DASH_BOOL_TEST"] = "0"
    @test dash_env(Bool, "bool_test", 50) == false

    ENV["DASH_BOOL_TEST"] = "TRUE"
    @test dash_env(Bool, "bool_test", 50) == true

    ENV["DASH_BOOL_TEST"] = "FALSE"
    @test dash_env(Bool, "bool_test", 50) == false
end


text:
@recipe function f(dp::EHDriftPath{T}; showlabel = true) where {T <: Real}
    legendfontsize --> 15
    linewidth --> 3
    seriestype --> :path3d
    @series begin
        showlabel ? label --> "e-" : label --> ""
        seriescolor --> :red
        [CartesianPoint{T}(path) for path in dp.e_path]
    end
    @series begin
        showlabel ? label --> "h+" : label --> ""
        seriescolor --> :green
        [CartesianPoint{T}(path) for path in dp.h_path]
    end
    @series begin
        label := ""
        seriestype := :scatter
        seriescolor --> :red
        CartesianPoint{T}(dp.e_path[end])
    end
    @series begin
        label := ""
        seriestype := :scatter
        seriescolor --> :green
        CartesianPoint{T}(dp.h_path[end])
    end
end
@recipe function f(dps::Vector{EHDriftPath{T}}) where {T <: Real}
    for i in eachindex(dps)
        @series begin
            showlabel --> i == 1
            dps[i]
        end
    end
end


text:
module TestGradfunc
    using Calculus2
    using Base.Test

    f(x) = sin(x)
    g(x) = @gradfunc(sin(x), x)
    h(x) = @hessfunc(sin(x), x)

    f(1.0)
    g(1.0)
    h(1.0)

    f(z) = 1 / (1 + exp(-z))
    g(z) = @gradfunc(1 / (1 + exp(-z)), z)

    @test isapprox(
        g(0.0),
        grad(f, 0.0)
    )

    @test isapprox(
        g(1.0),
        grad(f, 1.0)
    )

    f(x) = sin(x) + cos(x)^2 + sin(x)^3
    g(x) = @gradfunc(sin(x) + cos(x)^2 + sin(x)^3, x)

    @test isapprox(
        g(0.0),
        grad(f, 0.0)
    )

    @test isapprox(
        g(1.0),
        grad(f, 1.0)
    )

    # NB: This won't work for multivariate functions
    f(x) = sin(x) + cos(x)^2 + sin(x)^3
    grad(f, 1.0, method = :finite)
    grad(f, 1.0, method = :ad)
    gradfunc(f, 1.0, method = :finite)(1.0)
    gradfunc(f, 1.0, method = :ad)(1.0)
    g(x) = @gradfunc(sin(x) + cos(x)^2 + sin(x)^3, x)
    g(1.0)
    hess(f, 1.0, method = :finite)
    hessfunc(f, 1.0, method = :finite)(1.0)
    h(x) = @hessfunc(sin(x) + cos(x)^2 + sin(x)^3, x)
    h(1.0)
end


text:
using CSV, StatsBase

jp(csv) = joinpath(@__DIR__, csv)

datasets = Dict(
:toy1 => (jp("./data/toy/toy1.csv"), 90, :Y),
:toy2 => (jp("./data/toy/toy2.csv"), 90, :Y),
:wine => (jp("./data/wine/wine5K_f32.csv"), 3600, :quality),
:garnet => (jp("./data/garnet/garnet_norm_f32.csv"), 9000, :TiO2),
)

type RealData{T}
  X::AbstractMatrix{T}
  signal::AbstractVector{T}
  noise::AbstractVector{T}
  y::AbstractVector{T}
  n::Int
end

RealData(X, y) = RealData(X, y, y*0, y, length(y))

loadrealdata(datasetname::Symbol) = loadrealdata(datasets[datasetname]...)

function loadrealdata(csvfilename::String, ntrain::Int, yname::Symbol)
    data = CSV.read(csvfilename)
    disallowmissing!(data)
    ntotal = nrow(data)

    srand(20181103)
    testindices = sample(1:ntotal, ntotal-ntrain, replace=false)
    trainindices = trues(ntotal)
    trainindices[testindices] = false

    test = data[testindices, :]
    train = data[trainindices, :]

    xnames = [name for name in names(data) if name != yname]
    trainx = Matrix(train[:, xnames])
    trainy = Vector(train[yname])
    testx = Matrix(test[:, xnames])
    testy = Vector(test[yname])

    RealData(trainx, trainy), RealData(testx, testy)
end


text:
module CommunalHelperDreamSwapBlock

using ..Ahorn, Maple
using Ahorn.CommunalHelper
using Ahorn.CommunalHelperEntityPresets: CustomDreamBlockData

const entityData = appendkwargs(CustomDreamBlockData, :(
    noReturn::Bool=false,
))
@mapdefdata Entity "CommunalHelper/DreamSwapBlock" DreamSwapBlock entityData

const placements = Ahorn.PlacementDict(
    "Dream Swap Block (Communal Helper)" => Ahorn.EntityPlacement(
        DreamSwapBlock,
        "rectangle",
        Dict{String,Any}(),
        Ahorn.SwapBlock.swapFinalizer,
    ),
)

Ahorn.nodeLimits(entity::DreamSwapBlock) = 1, 1

Ahorn.minimumSize(entity::DreamSwapBlock) = 16, 16
Ahorn.resizable(entity::DreamSwapBlock) = true, true

const crossSprite = "objects/CommunalHelper/dreamMoveBlock/x"

function Ahorn.selection(entity::DreamSwapBlock)
    x, y = Ahorn.position(entity)
    stopX, stopY = Int.(entity.data["nodes"][1])

    width = Int(get(entity.data, "width", 8))
    height = Int(get(entity.data, "height", 8))

    return [
        Ahorn.Rectangle(x, y, width, height),
        Ahorn.Rectangle(stopX, stopY, width, height),
    ]
end

function Ahorn.renderSelectedAbs(ctx::Ahorn.Cairo.CairoContext, entity::DreamSwapBlock)
    startX, startY = Int(entity.data["x"]), Int(entity.data["y"])
    stopX, stopY = Int.(entity.data["nodes"][1])

    width = Int(get(entity.data, "width", 32))
    height = Int(get(entity.data, "height", 32))

    renderDreamBlock(ctx, stopX, stopY, width, height, entity.data)

    Ahorn.drawArrow(ctx, startX + width / 2, startY + height / 2, stopX + width / 2, stopY + height / 2, Ahorn.colors.selection_selected_fc, headLength=6)
end

function Ahorn.renderAbs(ctx::Ahorn.Cairo.CairoContext, entity::DreamSwapBlock)

    startX, startY = Int(entity.data["x"]), Int(entity.data["y"])
    stopX, stopY = Int.(entity.data["nodes"][1])

    width = Int(get(entity.data, "width", 32))
    height = Int(get(entity.data, "height", 32))

    Ahorn.SwapBlock.renderTrail(ctx, min(startX, stopX), min(startY, stopY), abs(startX - stopX) + width, abs(startY - stopY) + height, "objects/swapblock/target")

    renderDreamBlock(ctx, startX, startY, width, height, entity.data)

    if Bool(get(entity.data, "noReturn", false))
        noReturnSprite = Ahorn.getSprite(crossSprite, "Gameplay")
        Ahorn.drawImage(ctx, noReturnSprite, startX + div(width - noReturnSprite.width, 2), startY + div(height - noReturnSprite.height, 2))
    end
end

end


text:
#=
Test ROI Align operations:
=#

@testset "ROI Align" begin
end


text:
using Compat
using Base.Meta
norm_ast(a::Any) = begin
    if isa(a, Expr)
        for (i, arg) in enumerate(a.args)
            a.args[i] = norm_ast(arg)
        end
        if a.head === :line
            return Expr(:line, a.args[1], :none)
        end
        if a.head === :macrocall
            fa = a.args[1]
            if fa === Symbol("@int128_str")
                return @compat parse(Int128,a.args[2])
            elseif fa === Symbol("@uint128_str")
                return @compat parse(UInt128,a.args[2])
            elseif fa === Symbol("@bigint_str")
                return @compat parse(BigInt,a.args[2])
            elseif fa == Symbol("@big_str")
                s = a.args[2]
                n = tryparse(BigInt,s)
                if !isnull(n)
                    return get(n)
                end
                n = tryparse(BigFloat,s)
                if !isnull(n)
                    return isnan(get(n)) ? :NaN : get(n)
                end
                return s
            end
        elseif isexpr(a, :call) && a.args[1] == :- && isa(a.args[2], Number)
            return -a.args[2]
        end
        return a
    elseif isa(a, QuoteNode)
        return Expr(:quote, norm_ast(a.value))
    elseif isa(a, AbstractFloat) && isnan(a)
        return :NaN
    end
    return a
end


text:
using Flux
using Optimisers
using PyCallChainRules.Torch: TorchModuleWrapper, torch

input_dim = 4
output_dim = 2
hiddendim = 16

batchsize = 8

torch_module = torch.nn.Sequential(
                            torch.nn.Linear(input_dim, hiddendim), torch.nn.ReLU(),
                            torch.nn.Linear(hiddendim, hiddendim), torch.nn.ReLU(),
                            torch.nn.Linear(hiddendim, output_dim)
                        )
jlmodel = TorchModuleWrapper(torch_module)

input = randn(Float32, input_dim, batchsize)
target = randn(Float32, output_dim, batchsize)

loss(model, x, y) = Flux.Losses.mse(model(x), y)

@info "before" loss(jlmodel, input, target)

function train(model;nsteps=100)
    opt = Optimisers.ADAM(0.01)
    state = Optimisers.setup(opt, model)

    for i in 1:nsteps
        gs, _ = Flux.gradient(model, input, target) do m, x, y
            loss(m, x, y)
        end
        state, model = Optimisers.update(state, model, gs)
    end
    return model
end

jlmodel = train(jlmodel)

@info "after" loss(jlmodel, input, target)


text: Dict{String, Any}("thevoid" => Any[Any[Any[Any[Any[]]]]])


text:
####
#### NNLS submodule
####

# This NNLS submodule is modified version of the corresponding NNLS module from
# the forked NonNegLeastSquares.jl repository:
#
#   https://github.com/jondeuce/NonNegLeastSquares.jl/blob/a122bf7acb498efcaf140b719133691e7c4cd03d/src/nnls.jl#L1
#
# The original MIT licence from NonNegLeastSquares.jl is included below:
#
# -----------------------------------------------------------------------------
#
# The MIT License (MIT)
#
# Copyright (c) 2015 Alex Williams
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
# -----------------------------------------------------------------------------

module NNLS

using LinearAlgebra
using UnPack: @unpack
using UnsafeArrays: @uviews, uviews, uview
using MuladdMacro: @muladd
# using LoopVectorization: @avx

export nnls, nnls!, NNLSWorkspace, load!

@muladd begin

struct NNLSWorkspace{T}
    A::Matrix{T}
    b::Vector{T}
    x::Vector{T}
    w::Vector{T}
    zz::Vector{T}
    idx::Vector{Int}
    rnorm::Base.RefValue{T}
    mode::Base.RefValue{Int}
    nsetp::Base.RefValue{Int}
end
@inline solution(work::NNLSWorkspace) = work.x
@inline dual(work::NNLSWorkspace) = work.w
@inline residualnorm(work::NNLSWorkspace) = work.rnorm[]
@inline ncomponents(work::NNLSWorkspace) = work.nsetp[]
@inline components(work::NNLSWorkspace) = uview(work.idx, 1:ncomponents(work))
@inline choleskyfactor(work::NNLSWorkspace, ::Val{:L}) = uview(work.A, 1:ncomponents(work), components(work))
@inline choleskyfactor(work::NNLSWorkspace, ::Val{:U}) = choleskyfactor(work, Val(:L))'

function NNLSWorkspace{T}(m, n) where {T}
    NNLSWorkspace{T}(
        zeros(T, m, n), # A
        zeros(T, m),    # b
        zeros(T, n),    # x
        zeros(T, n),    # w
        zeros(T, m),    # zz
        zeros(Int, n),  # idx
        Ref(zero(T)),   # rnorm
        Ref(0),         # mode
        Ref(0),         # nsetp
    )
end

function NNLSWorkspace(A::AbstractMatrix{T}, b::AbstractVector{T}) where {T}
    m, n = size(A)
    @assert size(b) == (m,)
    work = NNLSWorkspace{T}(m, n)
    load!(work, A, b)
    return work
end

function NNLSWorkspace(m::Int, n::Int, ::Type{T} = Float64) where {T}
    return NNLSWorkspace{T}(m, n)
end

function load!(work::NNLSWorkspace{T}, A::AbstractMatrix{T}, b::AbstractVector{T}) where {T}
    @assert size(A) == size(work.A)
    @assert size(b) == size(work.b)
    copyto!(work.A, A)
    copyto!(work.b, b)
    return work
end

@noinline function checkargs(work::NNLSWorkspace)
    m, n = size(work.A)
    @assert size(work.b) == (m,)
    @assert size(work.x) == (n,)
    @assert size(work.w) == (n,)
    @assert size(work.zz) == (m,)
    @assert size(work.idx) == (n,)
end

"""
x = nnls(A, b; ...)

Solves non-negative least-squares problem by the active set method
of Lawson & Hanson (1974).

Optional arguments:
    max_iter: maximum number of iterations (counts inner loop iterations)

References:
    Lawson, C.L. and R.J. Hanson, Solving Least-Squares Problems,
    Prentice-Hall, Chapter 23, p. 161, 1974.
"""
function nnls(
        A,
        b::AbstractVector{T};
        max_iter::Int = 3*size(A, 2)
    ) where {T}
    work = NNLSWorkspace(A, b)
    nnls!(work, max_iter)
    return solution(work)
end

function nnls!(
        work::NNLSWorkspace{T},
        A::AbstractMatrix{T},
        b::AbstractVector{T},
        max_iter::Int = 3*size(A, 2)
    ) where {T}
    load!(work, A, b)
    nnls!(work, max_iter)
    return solution(work)
end

"""
CONSTRUCTION AND/OR APPLICATION OF A SINGLE
HOUSEHOLDER TRANSFORMATION..     Q = I + U*(U**T)/B

The original version of this code was developed by
Charles L. Lawson and Richard J. Hanson at Jet Propulsion Laboratory
1973 JUN 12, and published in the book
"SOLVING LEAST SQUARES PROBLEMS", Prentice-HalL, 1974.
Revised FEB 1995 to accompany reprinting of the book by SIAM.
"""
@inline function construct_householder!(u::AbstractVector{T}, up::T) where {T}
    if length(u) <= 1
        return up
    end

    sm = zero(T)
    @inbounds @simd for ui in u #@avx
        sm = sm + ui * ui
    end
    cl = sqrt(sm)

    @inbounds u1 = u[1]
    cl = ifelse(u1 > 0, -cl, cl)
    up = u1 - cl
    @inbounds u[1] = cl

    return up
end

"""
CONSTRUCTION AND/OR APPLICATION OF A SINGLE
HOUSEHOLDER TRANSFORMATION..     Q = I + U*(U**T)/B

The original version of this code was developed by
Charles L. Lawson and Richard J. Hanson at Jet Propulsion Laboratory
1973 JUN 12, and published in the book
"SOLVING LEAST SQUARES PROBLEMS", Prentice-HalL, 1974.
Revised FEB 1995 to accompany reprinting of the book by SIAM.
"""
@inline function apply_householder!(u::AbstractVector{T}, up::T, c::AbstractVector{T}) where {T}
    m = length(u)
    if m <= 1
        return
    end

    @inbounds u1 = u[1]
    cl = abs(u1)
    @assert cl > 0
    b = up * u1
    if b >= 0
        return
    end

    @inbounds c1 = c[1]
    sm = c1 * up
    @inbounds @simd for i in 2:m #@avx
        sm = sm + c[i] * u[i]
    end

    if sm != 0
        sm /= b
        @inbounds c[1] = c[1] + sm * up
        @inbounds @simd ivdep for i in 2:m #@avx
            c[i] = c[i] + sm * u[i]
        end
    end
end

@inline function apply_householder!(A::AbstractMatrix{T}, up::T, idx::AbstractVector{Int}, nsetp::Int, jup::Int) where {T}
    m, n = size(A)
    if m - nsetp <= 0
        return
    end

    @inbounds u1 = A[nsetp, jup]
    cl = abs(u1)
    @assert cl > 0
    b = up * u1
    if b >= 0
        return
    end

    @inbounds A[nsetp, jup] = up
    @inbounds for ip in nsetp+1:n
        j = idx[ip]
        sm = zero(T)
        @inbounds @simd for l in nsetp:m #@avx
            sm = sm + A[l, j] * A[l, jup]
        end
        @inbounds if sm != 0
            sm /= b
            @inbounds @simd ivdep for l in nsetp:m #@avx
                A[l, j] = A[l, j] + sm * A[l, jup]
            end
        end
    end
    @inbounds A[nsetp, jup] = u1

    return
end

"""
   COMPUTE ORTHOGONAL ROTATION MATRIX..
The original version of this code was developed by
Charles L. Lawson and Richard J. Hanson at Jet Propulsion Laboratory
1973 JUN 12, and published in the book
"SOLVING LEAST SQUARES PROBLEMS", Prentice-HalL, 1974.
Revised FEB 1995 to accompany reprinting of the book by SIAM.

   COMPUTE.. MATRIX   (C, S) SO THAT (C, S)(A) = (SQRT(A**2+B**2))
                      (-S,C)         (-S,C)(B)   (   0          )
   COMPUTE SIG = SQRT(A**2+B**2)
      SIG IS COMPUTED LAST TO ALLOW FOR THE POSSIBILITY THAT
      SIG MAY BE IN THE SAME LOCATION AS A OR B .
"""
@inline function orthogonal_rotmat(a::T, b::T) where {T}
    if abs(a) > abs(b)
        xr = b / a
        yr = sqrt(one(T) + xr * xr)
        c = inv(yr) * sign(a)
        s = c * xr
        sig = abs(a) * yr
    elseif b != 0
        xr = a / b
        yr = sqrt(one(T) + xr * xr)
        s = inv(yr) * sign(b)
        c = s * xr
        sig = abs(b) * yr
    else
        sig = zero(T)
        c = zero(T)
        s = one(T)
    end
    return c, s, sig
end

"""
The original version of this code was developed by
Charles L. Lawson and Richard J. Hanson at Jet Propulsion Laboratory
1973 JUN 15, and published in the book
"SOLVING LEAST SQUARES PROBLEMS", Prentice-HalL, 1974.
Revised FEB 1995 to accompany reprinting of the book by SIAM.
"""
@inline function solve_triangular_system!(zz::AbstractVector{T}, A::AbstractMatrix{T}, idx::AbstractVector{Int}, nsetp::Int) where {T}
    if nsetp <= 0
        return
    end
    @inbounds j = idx[nsetp]
    @inbounds zz[nsetp] /= A[nsetp, j]
    @inbounds for ip in nsetp-1:-1:1
        zz1 = zz[ip + 1]
        @inbounds @simd ivdep for l in 1:ip #@avx
            zz[l] = zz[l] - A[l, j] * zz1
        end
        j = idx[ip]
        zz[ip] /= A[ip, j]
    end
    return
end

@inline function largest_positive_dual(w::AbstractVector{T}, idx::AbstractVector{Int}, nsetp::Int) where {T}
    n = length(w)
    wmax = zero(T)
    i_wmax = 0
    @inbounds for ip in nsetp+1:n
        j = idx[ip]
        if w[j] > wmax
            wmax = w[j]
            i_wmax = ip
        end
    end
    wmax, i_wmax
end

"""
Algorithm NNLS: NONNEGATIVE LEAST SQUARES

The original version of this code was developed by
Charles L. Lawson and Richard J. Hanson at Jet Propulsion Laboratory
1973 JUN 15, and published in the book
"SOLVING LEAST SQUARES PROBLEMS", Prentice-HalL, 1974.
Revised FEB 1995 to accompany reprinting of the book by SIAM.

GIVEN AN M BY N MATRIX, A, AND AN M-VECTOR, B,  COMPUTE AN
N-VECTOR, X, THAT SOLVES THE LEAST SQUARES PROBLEM
                 A * X = B  SUBJECT TO X .GE. 0
"""
function nnls!(
        work::NNLSWorkspace{T},
        max_iter::Int = 3*size(work.A, 2)
    ) where {T}

    checkargs(work)
    @unpack A, b, x, w, zz, idx = work
    m, n = size(A)

    iter = 0
    nsetp = 0
    up = zero(T)
    fill!(x, zero(T))
    copyto!(idx, 1:n)

    work.mode[] = 1
    terminated = false

    # ******  MAIN LOOP BEGINS HERE  ******
    @inbounds while true
        local i_curr, j_curr, i_maxdual, j_maxdual

        # QUIT IF ALL COEFFICIENTS ARE ALREADY IN THE SOLUTION.
        # OR IF M COLS OF A HAVE BEEN TRIANGULARIZED.
        if (nsetp + 1 > n || nsetp >= m)
            terminated = true
            break
        end

        # COMPUTE COMPONENTS OF THE DUAL (NEGATIVE GRADIENT) VECTOR W().
        @inbounds for ip in nsetp+1:n
            j = idx[ip]
            sm = zero(T)
            @inbounds @simd for l in nsetp+1:m #@avx
                sm = sm + A[l, j] * b[l]
            end
            w[j] = sm
        end

        @inbounds while true
            # FIND LARGEST POSITIVE W(J).
            wmax, i_wmax = largest_positive_dual(w, idx, nsetp)

            # IF WMAX .LE. 0. GO TO TERMINATION.
            # THIS INDICATES SATISFACTION OF THE KUHN-TUCKER CONDITIONS.
            if wmax <= 0
                terminated = true
                break
            end

            i_maxdual = i_wmax
            j_maxdual = idx[i_maxdual]

            # THE SIGN OF W(J) IS OK FOR J TO BE MOVED TO SET P.
            # BEGIN THE TRANSFORMATION AND CHECK NEW DIAGONAL ELEMENT TO AVOID
            # NEAR LINEAR DEPENDENCE.
            Asave = A[nsetp + 1, j_maxdual]
            up = construct_householder!(uview(A, nsetp+1:m, j_maxdual), up)

            if abs(A[nsetp + 1, j_maxdual]) > 0
                # COL J IS SUFFICIENTLY INDEPENDENT.  COPY B INTO ZZ, UPDATE ZZ
                # AND SOLVE FOR ZTEST ( = PROPOSED NEW VALUE FOR X(J) ).
                copyto!(zz, b)
                apply_householder!(uview(A, nsetp+1:m, j_maxdual), up, uview(zz, nsetp+1:m))
                ztest = zz[nsetp + 1] / A[nsetp + 1, j_maxdual]

                # SEE IF ZTEST IS POSITIVE
                if ztest > 0
                    break
                end
            end

            # REJECT J AS A CANDIDATE TO BE MOVED FROM SET Z TO SET P.
            # RESTORE A(NPP1,J), SET W(J)=0., AND LOOP BACK TO TEST DUAL
            # COEFFS AGAIN.
            A[nsetp + 1, j_maxdual] = Asave
            w[j_maxdual] = zero(T)
        end
        if terminated
            break
        end

        # THE INDEX  J=INDEX(IZ)  HAS BEEN SELECTED TO BE MOVED FROM
        # SET Z TO SET P.    UPDATE B,  UPDATE INDICES,  APPLY HOUSEHOLDER
        # TRANSFORMATIONS TO COLS IN NEW SET Z,  ZERO SUBDIAGONAL ELTS IN
        # COL J,  SET W(J)=0.
        copyto!(b, zz)

        idx[i_maxdual] = idx[nsetp + 1]
        idx[nsetp + 1] = j_maxdual
        nsetp += 1

        if nsetp + 1 <= n
            apply_householder!(A, up, idx, nsetp, j_maxdual)
        end

        if nsetp != m
            @inbounds @simd ivdep for l in nsetp+1:m #@avx
                A[l, j_maxdual] = zero(T)
            end
        end

        w[j_maxdual] = zero(T)

        # SOLVE THE TRIANGULAR SYSTEM.
        # STORE THE SOLUTION TEMPORARILY IN ZZ().
        solve_triangular_system!(zz, A, idx, nsetp)
        if nsetp > 0
            i_curr = first(idx)
        end

        # ******  SECONDARY LOOP BEGINS HERE ******
        #
        # ITERATION COUNTER.
        @inbounds while true
            iter += 1
            if iter > max_iter
                work.mode[] = 3
                terminated = true
                break
            end

            # SEE IF ALL NEW CONSTRAINED COEFFS ARE FEASIBLE.
            # IF NOT COMPUTE ALPHA.
            alpha = convert(T, 2)
            @inbounds for ip in 1:nsetp
                j = idx[ip]
                if zz[ip] <= 0
                    t = -x[j] / (zz[ip] - x[j])
                    if alpha > t
                        alpha = t
                        i_curr = ip
                    end
                end
            end

            # IF ALL NEW CONSTRAINED COEFFS ARE FEASIBLE THEN ALPHA WILL
            # STILL = 2.    IF SO EXIT FROM SECONDARY LOOP TO MAIN LOOP.
            if alpha == 2
                break
            end

            # OTHERWISE USE ALPHA WHICH WILL BE BETWEEN 0 AND 1 TO
            # INTERPOLATE BETWEEN THE OLD X AND THE NEW ZZ.
            @inbounds @simd for ip in 1:nsetp
                j = idx[ip]
                x[j] = x[j] + alpha * (zz[ip] - x[j])
            end

            # MODIFY A AND B AND THE INDEX ARRAYS TO MOVE COEFFICIENT I
            # FROM SET P TO SET Z.
            j_curr = idx[i_curr]
            @inbounds while true
                x[j_curr] = zero(T)

                if i_curr != nsetp
                    i_curr += 1
                    @inbounds for ip in i_curr:nsetp
                        j = idx[ip]
                        idx[ip-1] = j

                        cc, ss, sig = orthogonal_rotmat(A[ip-1, j], A[ip, j])
                        A[ip-1, j] = sig
                        A[ip, j]   = zero(T)

                        # Apply procedure G2 (CC,SS,A(J-1,L),A(J,L))
                        @inbounds @simd for l in 1:j-1
                            tmp = A[ip-1, l]
                            A[ip-1, l] =  cc * tmp + ss * A[ip, l]
                            A[ip,   l] = -ss * tmp + cc * A[ip, l]
                        end

                        @inbounds @simd for l in j+1:n
                            tmp = A[ip-1, l]
                            A[ip-1, l] =  cc * tmp + ss * A[ip, l]
                            A[ip,   l] = -ss * tmp + cc * A[ip, l]
                        end

                        # Apply procedure G2 (CC,SS,B(J-1),B(J))
                        tmp = b[ip-1]
                        b[ip-1] =  cc * tmp + ss * b[ip]
                        b[ip]   = -ss * tmp + cc * b[ip]
                    end
                end

                nsetp -= 1
                idx[nsetp + 1] = j_curr

                # SEE IF THE REMAINING COEFFS IN SET P ARE FEASIBLE.  THEY SHOULD
                # BE BECAUSE OF THE WAY ALPHA WAS DETERMINED.
                # IF ANY ARE INFEASIBLE IT IS DUE TO ROUND-OFF ERROR.  ANY
                # THAT ARE NONPOSITIVE WILL BE SET TO ZERO
                # AND MOVED FROM SET P TO SET Z.
                allfeasible = true
                @inbounds for ip in 1:nsetp
                    j_curr = idx[ip]
                    if x[j_curr] <= 0
                        allfeasible = false
                        i_curr = ip
                        break
                    end
                end
                if allfeasible
                    break
                end
            end

            # COPY B( ) INTO ZZ( ).  THEN SOLVE AGAIN AND LOOP BACK.
            copyto!(zz, b)
            solve_triangular_system!(zz, A, idx, nsetp)
            if nsetp > 0
                i_curr = first(idx)
            end
        end
        if terminated
            break
        end
        # ******  END OF SECONDARY LOOP  ******

        @inbounds @simd for ip in 1:nsetp
            x[idx[ip]] = zz[ip]
        end
        # ALL NEW COEFFS ARE POSITIVE.  LOOP BACK TO BEGINNING.
    end

    # ******  END OF MAIN LOOP  ******
    # COME TO HERE FOR TERMINATION.
    # COMPUTE THE NORM OF THE FINAL RESIDUAL VECTOR.

    sm = zero(T)
    if nsetp < m
        @inbounds @simd for ip in nsetp+1:m #@avx
            bi = b[ip]
            sm = sm + bi * bi
        end
    else
        fill!(w, zero(T))
    end
    work.rnorm[] = sqrt(sm)
    work.nsetp[] = nsetp
    return work.x
end

end # @muladd

end # module


text:
using IAEAPhsp
using Base.Test

@testset "write read" begin
    path = "some_file"
    s = Source(path, Write)
    u,v,w = normalize!(randn(Float32,3))
    p = Particle(photon, 1f0,2f0,3f0,4f0,5f0,u,v,w, (), (Int32(13),))
    writeparticle(s, p)
    destroy(s)

    s2 = Source(path, Read)

    @test length(s2) == 1
    p2 = first(s2)
    @test p2 ≈ p
    rm(path*".IAEAphsp", force=true)
    rm(path*".IAEAheader", force=true)
end


text:
module Boltzmann

export RBM,
       BernoulliRBM,
       Monitor,
       Update!,
       GRBM,
       DBN,
       DAE,
       fit,
       transform,
       sample_hiddens,
       sample_visibles,
       mag_vis_tap2,
       mag_hid_tap2,
       generate,
       components,
       features,
       unroll,
       save_params,
       load_params,
       chart_weights,
       chart_weights_distribution,
       chart_activation_distribution,
       binarize,
       binarize!,
       normalize,
       normalize!,
       normalize_samples,
       normalize_samples!,
       ShowMonitor,
       WriteMonitorChartPDF,
       SaveMonitorHDF5,
       plot_scores,
       plot_evolution,
       plot_rf,
       plot_chain,
       plot_vbias,
       plot_weightdist,
       chart_likelihood_evolution,
       PinField!

include("core.jl")

end


text:
module MVDA

using DataFrames: copy, copyto!
using DataDeps, CSV, DataFrames, CodecZlib
using Parameters, Printf, MLDataUtils, ProgressMeter
using LinearAlgebra, Random, Statistics, StatsBase, StableRNGs
using KernelFunctions
using DelimitedFiles, Plots
using Polyester, MLDataUtils

import Base: show, iterate

##### DATA #####

#=
Uses DataDeps to download data as needed.
Inspired by UCIData.jl: https://github.com/JackDunnNZ/UCIData.jl
=#

const DATA_DIR = joinpath(@__DIR__, "data")

"""
`list_datasets()`

List available datasets in MVDA.
"""
list_datasets() = map(x -> splitext(x)[1], readdir(DATA_DIR))

function __init__()
    for dataset in list_datasets()
        include(joinpath(DATA_DIR, dataset * ".jl"))
    end
end

"""
`dataset(str)`

Load a dataset named `str`, if available. Returns data as a `DataFrame` where
the first column contains labels/targets and the remaining columns correspond to
distinct features.
"""
function dataset(str)
    # Locate dataset file.
    dataset_path = @datadep_str str
    file = readdir(dataset_path)
    index = findfirst(x -> occursin("data.", x), file)
    if index isa Int
        dataset_file = joinpath(dataset_path, file[index])
    else # is this unreachable?
        error("Failed to locate a data.* file in $(dataset_path)")
    end

    # Read dataset file as a DataFrame.
    df = if splitext(dataset_file)[2] == ".csv"
        CSV.read(dataset_file, DataFrame)
    else # assume .csv.gz
        open(GzipDecompressorStream, dataset_file, "r") do stream
            CSV.read(stream, DataFrame)
        end
    end
    return df
end

"""
Process the dataset located at the given `path`.

This is an extra step to give fine-grain control in generating files with DataDeps.jl.
"""
function process_dataset(path::AbstractString; header=false, missingstrings="", kwargs...)
    input_df = CSV.read(path, DataFrame, header=header, missingstrings=missingstrings)
    process_dataset(input_df; kwargs...)
    rm(path)
end

"""
Final step in processing the given dataset `input_df`.

This standardizes cached files that live in ~/.julia/datadeps so that labels/targets appear in first column
followed by features in the remaining columns.
We also check for uniqueness in features.
"""
function process_dataset(input_df::DataFrame;
    target_index=-1,
    feature_indices=1:0,
    ext=".csv")
    # Build output DataFrame.
    output_df = DataFrame()
    output_df.target = input_df[!, target_index]
    output_df = hcat(output_df, input_df[!, feature_indices], makeunique=true)
    output_cols = [ :target; [Symbol("x", n) for n in eachindex(feature_indices)] ]
    rename!(output_df, output_cols)
    dropmissing!(output_df)

    # Write to disk.
    output_path = "data" * ext
    if ext == ".csv"
        CSV.write(output_path, output_df, delim=',', writeheader=true)
    elseif ext == ".csv.gz"
        open(GzipCompressorStream, output_path, "w") do stream
            CSV.write(stream, output_df, delim=",", writeheader=true)
        end
    else
        error("Unknown file extension option '$(ext)'")
    end
end

##### IMPLEMENTATION #####

include("problem.jl")
include("utilities.jl")
include("projections.jl")
include("simulation.jl")

abstract type AbstractMMAlg end

include(joinpath("algorithms", "SD.jl"))
include(joinpath("algorithms", "MMSVD.jl"))
include(joinpath("algorithms", "CyclicVDA.jl"))

const DEFAULT_ANNEALING = geometric_progression
const DEFAULT_CALLBACK = __do_nothing_callback__
const DEFAULT_SCORE_FUNCTION = prediction_error

"""
    fit(algorithm, problem, ϵ, s; kwargs...)

Solve optimization problem at sparsity level `s` using a deadzone of size `ϵ`.

The solution is obtained via a proximal distance `algorithm` that gradually anneals parameter estimates
toward the target sparsity set.
"""
function fit(algorithm::AbstractMMAlg, problem::MVDAProblem, ϵ::Real, s::Real; kwargs...)
    extras = __mm_init__(algorithm, problem, nothing) # initialize extra data structures
    MVDA.fit!(algorithm, problem, ϵ, s, extras, (true,false,); kwargs...)
end

"""
    fit!(algorithm, problem, ϵ, s, [extras], [update_extras]; kwargs...)

Same as `fit_MVDA(algorithm, problem, ϵ, s)`, but with preallocated data structures in `extras`.

!!! Note
    The caller should specify whether to update data structures depending on `s` and `ρ` using `update_extras[1]` and `update_extras[2]`, respectively.

    Convergence is determined based on the rule `dist < dtol || abs(dist - old) < rtol * (1 + old)`, where `dist` is the squared distance and `dtol` and `rtol` are tolerance parameters.

!!! Tip
    The `extras` argument can be constructed using `extras = __mm_init__(algorithm, problem, nothing)`.

# Keyword Arguments

- `nouter`: The number of outer iterations; i.e. the maximum number of `ρ` values to use in annealing (default=`100`).
- `dtol`: An absolute tolerance parameter for the squared distance (default=`1e-6`).
- `rtol`: A relative tolerance parameter for the squared distance (default=`1e-6`).
- `rho_init`: The initial value for `ρ` (default=1.0).
- `rho_max`: The maximum value for `ρ` (default=1e8).
- `rhof`: A function `rhof(ρ, iter, rho_max)` used to determine the next value for `ρ` in the annealing sequence. The default multiplies `ρ` by `1.2`.
- `verbose`: Print convergence information (default=`false`).
- `cb`: A callback function for extending functionality.

See also: [`MVDA.anneal!`](@ref) for additional keyword arguments applied at the annealing step.
"""
function fit!(algorithm::AbstractMMAlg, problem::MVDAProblem, ϵ::Real, s::Real,
    extras=nothing,
    update_extras::NTuple{2,Bool}=(true,false,);
    nouter::Int=100,
    dtol::Real=1e-6,
    rtol::Real=1e-6,
    rho_init::Real=1.0,
    rho_max::Real=1e8,
    rhof::Function=DEFAULT_ANNEALING,
    verbose::Bool=false,
    cb::Function=DEFAULT_CALLBACK,
    kwargs...)
    # Check for missing data structures.
    if extras isa Nothing
        error("Detected missing data structures for algorithm $(algorithm).")
    end

    # Get problem info and extra data structures.
    @unpack intercept, coeff, coeff_prev, proj = problem
    @unpack projection = extras

    # Fix model size(s).
    k = sparsity_to_k(problem, s)

    # Initialize ρ and iteration count.
    ρ, iters = rho_init, 0

    # Update data structures due to hyperparameters.
    update_extras[1] && __mm_update_sparsity__(algorithm, problem, ϵ, ρ, k, extras)
    update_extras[2] && __mm_update_rho__(algorithm, problem, ϵ, ρ, k, extras)

    # Check initial values for loss, objective, distance, and norm of gradient.
    apply_projection(projection, problem, k)
    init_result = __evaluate_objective__(problem, ϵ, ρ, extras)
    result = SubproblemResult(0, init_result)
    cb(0, problem, ϵ, ρ, k, result)
    old = sqrt(result.distance)

    for iter in 1:nouter
        # Solve minimization problem for fixed rho.
        verbose && print("\n",iter,"  ρ = ",ρ)
        result = MVDA.anneal!(algorithm, problem, ϵ, ρ, s, extras, (false,true,); verbose=verbose, cb=cb, kwargs...)

        # Update total iteration count.
        iters += result.iters

        cb(iter, problem, ϵ, ρ, k, result)

        # Check for convergence to constrained solution.
        dist = sqrt(result.distance)
        if dist < dtol || abs(dist - old) < rtol * (1 + old)
            break
        else
          old = dist
        end

        # Update according to annealing schedule.
        ρ = ifelse(iter < nouter, rhof(ρ, iter, rho_max), ρ)
    end

    # Project solution to the constraint set.
    apply_projection(projection, problem, k)
    loss, obj, dist, gradsq = __evaluate_objective__(problem, ϵ, ρ, extras)

    if verbose
        print("\n\niters = ", iters)
        print("\n∑ᵢ max{0, |yᵢ-Bᵀxᵢ|₂-ϵ}² = ", loss)
        print("\nobjective  = ", obj)
        print("\ndistance   = ", sqrt(dist))
        println("\n|gradient| = ", sqrt(gradsq))
    end

    return SubproblemResult(iters, loss, obj, dist, gradsq)
end

"""
    anneal(algorithm, problem, ϵ, ρ, s; kwargs...)

Solve the `ρ`-penalized optimization problem at sparsity level `s` with deadzone `ϵ`.
"""
function anneal(algorithm::AbstractMMAlg, problem::MVDAProblem, ϵ::Real, ρ::Real, s::Real; kwargs...)
    extras = __mm_init__(algorithm, problem, nothing)
    MVDA.anneal!(algorithm, problem, ϵ, ρ, s, extras, (true,true,); kwargs...)
end

"""
    anneal!(algorithm, problem, ϵ, ρ, s, [extras], [update_extras]; kwargs...)

Same as `anneal(algorithm, problem, ϵ, ρ, s)`, but with preallocated data structures in `extras`.

!!! Note
    The caller should specify whether to update data structures depending on `s` and `ρ` using `update_extras[1]` and `update_extras[2]`, respectively.

    Convergence is determined based on the rule `gradsq < gtol`, where `gradsq` is squared Euclidean norm of the gradient and `gtol` is a tolerance parameter.

!!! Tip
    The `extras` argument can be constructed using `extras = __mm_init__(algorithm, problem, nothing)`.

# Keyword Arguments

- `ninner`: The maximum number of iterations (default=`10^4`).
- `gtol`: An absoluate tolerance parameter on the squared Euclidean norm of the gradient (default=`1e-6`).
- `nesterov_threshold`: The number of early iterations before applying Nesterov acceleration (default=`10`).
- `verbose`: Print convergence information (default=`false`).
- `cb`: A callback function for extending functionality.
"""
function anneal!(algorithm::AbstractMMAlg, problem::MVDAProblem, ϵ::Real, ρ::Real, s::Real,
    extras=nothing,
    update_extras::NTuple{2,Bool}=(true,true);
    ninner::Int=10^4,
    gtol::Real=1e-6,
    nesterov_threshold::Int=10,
    verbose::Bool=false,
    cb::Function=DEFAULT_CALLBACK,
    kwargs...
    )
    # Check for missing data structures.
    if extras isa Nothing
        error("Detected missing data structures for algorithm $(algorithm).")
    end

    # Get problem info and extra data structures.
    @unpack intercept, coeff, coeff_prev, proj = problem
    @unpack projection = extras

    # Fix model size(s).
    k = sparsity_to_k(problem, s)

    # Update data structures due to hyperparameters.
    update_extras[1] && __mm_update_sparsity__(algorithm, problem, ϵ, ρ, k, extras)
    update_extras[2] && __mm_update_rho__(algorithm, problem, ϵ, ρ, k, extras)

    # Check initial values for loss, objective, distance, and norm of gradient.
    apply_projection(projection, problem, k)
    result = __evaluate_objective__(problem, ϵ, ρ, extras)
    cb(0, problem, ϵ, ρ, k, result)
    old = result.objective

    if sqrt(result.gradient) < gtol
        return SubproblemResult(0, result)
    end

    # Use previous estimates in case of warm start.
    copyto!(coeff.all, coeff_prev.all)

    # Initialize iteration counts.
    iters = 0
    nesterov_iter = 1
    verbose && @printf("\n%-5s\t%-8s\t%-8s\t%-8s\t%-8s", "iter.", "loss", "objective", "distance", "|gradient|")
    for iter in 1:ninner
        iters += 1

        # Apply the algorithm map to minimize the quadratic surrogate.
        __mm_iterate__(algorithm, problem, ϵ, ρ, k, extras)

        # Update loss, objective, distance, and gradient.
        apply_projection(projection, problem, k)
        result = __evaluate_objective__(problem, ϵ, ρ, extras)

        cb(iter, problem, ϵ, ρ, k, result)

        if verbose
            @printf("\n%4d\t%4.3e\t%4.3e\t%4.3e\t%4.3e", iter, result.loss, result.objective, sqrt(result.distance), sqrt(result.gradient))
        end

        # Assess convergence.
        obj = result.objective
        gradsq = sqrt(result.gradient)
        if gradsq < gtol
            break
        elseif iter < ninner
            needs_reset = iter < nesterov_threshold || obj > old
            nesterov_iter = __apply_nesterov__!(coeff.all, coeff_prev.all, nesterov_iter, needs_reset)
            old = obj
        end
    end
    # Save parameter estimates in case of warm start.
    copyto!(coeff_prev.all, coeff.all)

    return SubproblemResult(iters, result)
end

"""
    cv(algorithm, problem, grids; [at], [kwargs...])

Split data in `problem` into cross-validation and a test sets, then run cross-validation over the `grids`.

# Keyword Arguments

- `at`: A value between `0` and `1` indicating the proportion of samples/instances used for cross-validation, with remaining samples used for a test set (default=`0.8`).

See also: [`MVDA.cv(algorithm::AbstractMMAlg, problem::MVDAProblem, grids::Tuple{E,S}, dataset_split::Tuple{Any,Any})`](@ref)
"""
function cv(algorithm::AbstractMMAlg, problem::MVDAProblem, grids::Tuple{E,S}; at::Real=0.8, kwargs...) where {E,S}
    # Split data into cross-validation and test sets.
    @unpack p, Y, X, intercept = problem
    dataset_split = splitobs((Y, view(X, :, 1:p)), at=at, obsdim=1)
    MVDA.cv(algorithm, problem, grids, dataset_split; kwargs...)
end

"""
    cv(algorithm, problem, grids, dataset_split; [kwargs...])

Run k-fold cross-validation over hyperparameters `(ϵ, s)` for deadzone radius and sparsity level, respectively.

The given `problem` should enter with initial model parameters in `problem.coeff.all`.
Hyperparameters are specified in `grids = (ϵ_grid, s_grid)`, and data subsets are given as `dataset_split = (cv_set, test_set)`.

# Keyword Arguments

- `nfolds`: The number of folds to run in cross-validation.
- `scoref`: A function that evaluates a classifier over training, validation, and testing sets (default uses misclassification error).
- `show_progress`: Toggles progress bar.

Additional arguments are propagated to `fit` and `anneal`. See also [`MVDA.fit`](@ref) and [`MVDA.anneal`](@ref).
"""
function cv(algorithm::AbstractMMAlg, problem::MVDAProblem, grids::Tuple{E,S}, dataset_split::Tuple{Any,Any};
    lambda::Real=1e-3,
    maxiter::Int=10^4,
    tol::Real=1e-4,
    nfolds::Int=5,
    scoref::Function=DEFAULT_SCORE_FUNCTION,
    cb::Function=DEFAULT_CALLBACK,
    show_progress::Bool=true,
    kwargs...) where {E,S}
    # Initialize the output.
    cv_set, test_set = dataset_split
    ϵ_grid, s_grid = grids
    nϵ, ns = length(ϵ_grid), length(s_grid)
    alloc_score_arrays(a, b, c) = [Matrix{Float64}(undef, a, b) for _ in 1:c]
    result = (;
        train=alloc_score_arrays(ns, nϵ, nfolds),
        validation=alloc_score_arrays(ns, nϵ, nfolds),
        test=alloc_score_arrays(ns, nϵ, nfolds),
        time=alloc_score_arrays(ns, nϵ, nfolds),
    )

    # Run cross-validation.
    if show_progress
        progress_bar = Progress(nfolds * nϵ * ns, 1, "Running CV w/ $(algorithm)... ")
    end

    for (k, fold) in enumerate(kfolds(cv_set, k=nfolds, obsdim=1))
        # Retrieve the training set and validation set.
        # TODO: Does this guarantee copies?
        train_set, validation_set = fold
        train_Y, train_X = getobs(train_set, obsdim=1)
        val_Y, val_X = getobs(validation_set, obsdim=1)
        test_Y, test_X = getobs(test_set, obsdim=1)

        # Standardize ALL data based on the training set.
        F = StatsBase.fit(ZScoreTransform, train_X, dims=1)
        has_nan = any(isnan, F.scale) || any(isnan, F.mean)
        has_inf = any(isinf, F.scale) || any(isinf, F.mean)
        has_zero = any(iszero, F.scale)
        if has_nan
            error("Detected NaN in z-score.")
        elseif has_inf
            error("Detected Inf in z-score.")
        elseif has_zero
            for idx in eachindex(F.scale)
                x = F.scale[idx]
                F.scale[idx] = ifelse(iszero(x), one(x), x)
            end
        end

        foreach(X -> StatsBase.transform!(F, X), (train_X, val_X, test_X))

        # Create a problem object for the training set.
        train_idx, _ = parentindices(train_set[1])
        train_problem = change_data(problem, train_Y, train_X)
        extras = __mm_init__(algorithm, train_problem, nothing)

        for (j, ϵ) in enumerate(ϵ_grid)
            # Set initial model parameters.
            set_initial_coefficients!(train_problem, problem, train_idx)

            for (i, s) in enumerate(s_grid)
                # Obtain solution as function of (ϵ, s).
                if s != 0.0
                    result.time[k][i,j] = @elapsed MVDA.fit!(algorithm, train_problem, ϵ, s, extras, (true, false,);
                        cb=cb, kwargs...
                    )
                else# s == 0
                    result.time[k][i,j] = @elapsed MVDA.init!(algorithm, train_problem, ϵ, lambda, extras;
                        maxiter=maxiter, gtol=tol, nesterov_threshold=0,
                    )
                end
                copyto!(train_problem.coeff.all, train_problem.proj.all)

                # Evaluate the solution.
                r = scoref(train_problem, (train_Y, train_X), (val_Y, val_X), (test_Y, test_X))
                for (arr, val) in zip(result, r)
                    arr[k][i,j] = val
                end

                # Update the progress bar.
                if show_progress
                    spercent = string(round(100*s, digits=6), '%')
                    next!(progress_bar, showvalues=[(:fold, k), (:sparsity, spercent), (:ϵ, ϵ)])
                end
            end
        end
    end

    return result
end

function cv_estimation(algorithm::AbstractMMAlg, problem::MVDAProblem, grids::Tuple{E,S}; at::Real=0.8, kwargs...) where {E,S}
    # Split data into cross-validation and test sets.
    @unpack p, Y, X, intercept = problem
    dataset_split = splitobs((Y, view(X, :, 1:p)), at=at, obsdim=1)
    MVDA.cv_estimation(algorithm, problem, grids, dataset_split; kwargs...)
end

function cv_estimation(algorithm::AbstractMMAlg, problem::MVDAProblem, grids::Tuple{E,S}, dataset_split::Tuple{Any,Any};
    nreplicates::Int=10,
    show_progress::Bool=true,
    rng::AbstractRNG=StableRNG(1903),
    kwargs...) where {E,S}
    # Retrieve subsets and create index set into cross-validation set.
    cv_set, test_set = dataset_split

    if show_progress
        progress_bar = Progress(nreplicates, 1, "Running CV w/ $(algorithm)... ")
    end

    # Replicate CV procedure several times.
    replicate = NamedTuple[]
    for r in 1:nreplicates
        # Shuffle cross-validation data.
        cv_shuffled = shuffleobs(cv_set, obsdim=1, rng=rng)

        # Run k-fold cross-validation and store results.
        result = MVDA.cv(algorithm, problem, grids, (cv_shuffled, test_set); show_progress=false, kwargs...)
        push!(replicate, result)

        # Update the progress bar.
        if show_progress
            next!(progress_bar, showvalues=[(:replicate, r),])
        end
    end

    return replicate
end

"""
```init!(algorithm, problem, ϵ, λ, [_extras_]; [maxiter=10^3], [gtol=1e-6], [nesterov_threshold=10], [verbose=false])```

Initialize a `problem` with its `λ`-regularized solution.
"""
function init!(algorithm::AbstractMMAlg, problem::MVDAProblem, ϵ, λ, _extras_=nothing;
    maxiter::Int=10^3,
    gtol::Real=1e-6,
    nesterov_threshold::Int=10,
    verbose::Bool=false,
    )
    # Check for missing data structures.
    extras = __mm_init__(algorithm, problem, _extras_)

    # Get problem info and extra data structures.
    @unpack coeff, coeff_prev, proj = problem

    # Update data structures due to hyperparameters.
    __mm_update_lambda__(algorithm, problem, ϵ, λ, extras)

    # Initialize coefficients.
    randn!(coeff.all)
    copyto!(coeff_prev.all, coeff.all)

    # Check initial values for loss, objective, distance, and norm of gradient.
    result = __evaluate_reg_objective__(problem, ϵ, λ, extras)
    old = result.objective

    if sqrt(result.gradient) < gtol
        return SubproblemResult(0, result)
    end

    # Initialize iteration counts.
    iters = 0
    nesterov_iter = 1
    verbose && @printf("\n%-5s\t%-8s\t%-8s\t%-8s", "iter.", "loss", "objective", "|gradient|")
    for iter in 1:maxiter
        iters += 1

        # Apply the algorithm map to minimize the quadratic surrogate.
        __reg_iterate__(algorithm, problem, ϵ, λ, extras)

        # Update loss, objective, and gradient.
        result = __evaluate_reg_objective__(problem, ϵ, λ, extras)

        if verbose
            @printf("\n%4d\t%4.3e\t%4.3e\t%4.3e", iter, result.loss, result.objective, sqrt(result.gradient))
        end

        # Assess convergence.
        obj = result.objective
        gradsq = sqrt(result.gradient)
        if gradsq < gtol
            break
        elseif iter < maxiter
            needs_reset = iter < nesterov_threshold || obj > old
            nesterov_iter = __apply_nesterov__!(coeff.all, coeff_prev.all, nesterov_iter, needs_reset)
            old = obj
        end
    end
    # Save parameter estimates in case of warm start.
    copyto!(coeff_prev.all, coeff.all)
    copyto!(proj.all, coeff.all)

    return SubproblemResult(iters, result)
end

# function fit_MVDA(algorithm::CyclicVDA, problem, ϵ, δ, λ₁, λ₂;
#         niter::Int=10^3,
#         atol=1e-4,
#     )
#     @unpack Y, X, res, coeff = problem
#     # δ = 1 / 20
#     # ϵ = 1//2 * sqrt(2*c/(c-1))
#     n, p, c = probdims(problem)
#     μ₁ = n * λ₁
#     μ₂ = n * λ₂

#     # initialize residuals
#     mul!(res.main.all, X, coeff.all)
#     axpby!(1.0, Y, -1.0, res.main.all)
#     extras = nothing

#     full_objective, _, _ = fetch_objective(problem, p+1, 1, ϵ, δ, λ₁, λ₂)
#     penalty1 = 0.0
#     penalty2 = 0.0
#     for j in 1:p # does not include intercept here
#         β = view(problem.coeff.all, j, :)
#         penalty1 = penalty1 + μ₁ * norm(β, 1)
#         penalty2 = penalty2 + μ₂ * norm(β, 2)
#     end
#     full_objective = full_objective + penalty2 + penalty1

#     iters = 0
#     for iter in 1:niter
#         iters += 1

#         __mm_iterate__(algorithm, problem, ϵ, δ, λ₁, λ₂, extras)
#         loss, _, _ = fetch_objective(problem, p+1, 1, ϵ, δ, λ₁, λ₂)
#         penalty1 = 0.0
#         penalty2 = 0.0
#         for j in 1:p # does not include intercept here
#             β = view(problem.coeff.all, j, :)
#             penalty1 = penalty1 + μ₁ * norm(β, 1)
#             penalty2 = penalty2 + μ₂ * norm(β, 2)
#         end
#         objective = loss + penalty2 + penalty1

#         if objective > full_objective error("Descent failure") end

#         if full_objective - objective < atol break end
#         full_objective = objective
#     end

#     copyto!(problem.proj.all, coeff.all)

#     return full_objective, penalty1, penalty2
# end

export IterationResult, SubproblemResult
export MVDAProblem, SD, MMSVD # CyclicVDA

end


text:
## bool

const True = API.LLVMBool(1)
const False = API.LLVMBool(0)

struct Bool
    bool::Core.Bool
end

Base.convert(::Type{Bool}, bool::Core.Bool) = bool ? True : False

function Base.convert(::Type{Core.Bool}, bool::API.LLVMBool)
    if bool == True
        return true
    elseif bool == False
        return false
    else
        throw(ArgumentError("Invalid LLVMBool value $(bool)"))
    end
end


text:
function retune_params(suite)
    tune!(suite)
    BenchmarkTools.save(paramsfile, params(suite))
end

function run_benchmark(;verbose=true)
    disable_logging(Logging.Info);
    res = run(SUITE,verbose=verbose);
    disable_logging(Logging.Debug)
    return res
end

function save_benchmark(results, tags::Vector{<:Any}=[])
    # Add any custom tags (includes date by default)
    append!(SUITE.tags, tags)

    # Load in history
    suite_history = load_benchmark_history()

    # Add current to history
    key = string(today())
    if haskey(suite_history, key)
        print("A benchmark has already been run today. Enter:\n  0 - to Cancel\n  1 - to Overwrite\n  \t Another key name\n INPUT > ")
        input = readline()
        if input == "0"
            return nothing
        elseif input != "1"
            key = input
        end
    end

    suite_history[key] = results

    # Save to file
    printstyled("Saving to file",color=:blue)
    BenchmarkTools.save(histfile, suite_history)
    return nothing
end

function set_baseline_benchmark(tag::String)
    suite_history = load_benchmark_history()
    results = suite_history[tag]
    set_baseline_benchmark(results)
end

function set_baseline_benchmark(results::BenchmarkGroup)
    suite_history = load_benchmark_history()
    suite_history["baseline"] = results
    BenchmarkTools.save(histfile, suite_history)
    return nothing
end

function get_last_benchmark(;offset=0)
    suite_history = load_benchmark_history()
    tags = collect(keys(suite_history))
    tags = tags[tags .!= "baseline"]
    last = sort(tags)[end-offset]
    suite_history[last]
end

function get_baseline_benchmark()
    suite_history = load_benchmark_history()
    suite_history["baseline"]
end

function load_benchmark_history()
    suite_history = BenchmarkTools.load(histfile)[1]
end

function prev_benchmark_comparison(results; f=BenchmarkTools.median, styled=true, offset=0)
    base_res = get_last_benchmark(offset=offset)
    benchmark_comparison(results, base_res)
end

function prev_benchmark_comparison(;verbose=false, offset=0)
    results = run(suite,verbose=verbose)
    prev_benchmark_comparison(results,offset=offset)
end

function baseline_comparison(results; f=BenchmarkTools.median, styled=true)
    base_res = get_baseline_benchmark()
    benchmark_comparison(results, base_res)
end

function baseline_comparison(;verbose=false)
    results = run(suite,verbose=verbose)
    baseline_comparison(results)
    results
end

function benchmark_comparison(r1::BenchmarkGroup,r2::BenchmarkGroup;  f=BenchmarkTools.median)
    printstyled("Kuka Arm",bold=true,color=:yellow); println()
    compare_system(r1["kuka"],r2["kuka"],"constrained")
    compare_system(r1["kuka"],r2["kuka"],"unconstrained")
    printstyled("Dubins Car",bold=true,color=:yellow); println()
    compare_system(r1["dubinscar"],r2["dubinscar"],"constrained")
    compare_system(r1["dubinscar"],r2["dubinscar"],"unconstrained")
    return nothing
end

function benchmark_comparison(r1,s1,r2,s2; f=BenchmarkTools.median, styled=true)
    r_comp = compare_results(r1, r2,f=f)
    s_comp = compare_stats(s1, s2)
    if [r[1] for r in r_comp] == [s[1] for s in s_comp]
    end
    for (r,s) in zip(r_comp,s_comp)
        if r[1] == s[1]  # Check to make sure the runs match
            # Merge the Dictionaries
            merge!(r[2], s[2])
        end
    end
    print_comparison_table(r_comp, styled=styled)
end

function compare_results(r1::BenchmarkGroup, r2::BenchmarkGroup; f=BenchmarkTools.median)
    tmp = leaves(r1)
    run_list1 = [t[1] for t in tmp]
    tmp = leaves(r2)
    run_list2 = [t[1] for t in tmp]
    run_list = intersect(run_list1, run_list2)

    j = judge(f(r1),f(r2))

    comp = []
    for run in run_list
        vals = [f(r1[run]), f(r2[run])]
        t = j[run]
        stats_comp = Dict()
        stats_comp["time"] =  (join([prettytime(val.time) for val in vals], " ← ") * " (" * prettydiff(time(ratio(t))) * ")", t.time)
        stats_comp["memory"] = (join([prettymemory(val.memory) for val in vals], " ← ") * " (" * prettydiff(memory(ratio(t))) * ")", t.memory)
        stats_comp["gctime"] = (join([prettytime(val.gctime) for val in vals], " ← ") * " (" * prettytime(vals[1].gctime - vals[2].gctime) * ")",
            check_improvement(vals[1].gctime, vals[2].gctime, Inf))
        stats_comp["allocs"] = (join([format(val.allocs, autoscale=:metric, precision=2) for val in vals], " ← ") * " (" * format(vals[1].allocs - vals[2].allocs, signed=true, autoscale=:metric, precision=3) * ")",
            check_improvement(vals[1].allocs, vals[2].allocs))
        push!(comp, (run, stats_comp))
    end
    return sort(comp)
end

function check_improvement(v1,v2,tol=0.05)
    diff = v1-v2
    if diff > tol*v2
        return :regression
    elseif diff < -tol*v2
        return :improvement
    else
        return :invariant
    end
end

prettytimesigned(x::Real) = (x > 0 ? "+" : "") * prettytime(x)
prettyjudgetime(t::BenchmarkTools.TrialJudgement) = prettydiff(time(ratio(t))) * " => " * string(time(t)) * " (" * prettypercent(params(t).time_tolerance) * " tolerance)"
prettyjudgememory(t::BenchmarkTools.TrialJudgement) = prettydiff(memory(ratio(t))) * " => " * string(memory(t)) * " (" * prettypercent(params(t).memory_tolerance) * " tolerance)"

function compare_stats(s1::BenchmarkGroup, s2::BenchmarkGroup)
    stat_list = ["iterations","c_max","cost","setup_time","major iterations"]
    tmp = leaves(s1[@tagged "iterations"])
    run_list1 = [t[1] for t in tmp]
    tmp = leaves(s2[@tagged "iterations"])
    run_list2 = [t[1] for t in tmp]
    run_list = intersect(run_list1, run_list2)
    run_list = [run[1:end-1] for run in run_list]

    width = 8
    comp = []
    for run in run_list
        stat_comp = Dict()
        for stat in stat_list
            v1 = s1[run][stat]
            v2 = s2[run][stat]
            diff = v1 - v2
            str = trim_entry(v1, width, pad=false, positivespace=false) * " ← " * trim_entry(v2, width, pad=false, positivespace=false) * " (" * trim_entry(diff, width, pad=false, positivespace=false, signed=true) * ")"
            stat_comp[stat] = (str, check_improvement(v1,v2))
        end
        push!(comp,(run, stat_comp))
    end
    return sort(comp)
end

function loop_groups(s1::BenchmarkGroup, s2::BenchmarkGroup)
    k1 = keys(s1)
    k2 = keys(s2)
    both_keys = intersect(k1,k2)
    for key in both_keys
        if s1[key] isa BenchmarkGroup && s2[key] isa BenchmarkGroup
            loop_groups(s1[key], s2[key])
        else
            @show key, s1[key], s2[key]
        end
    end
end

function print_comparison_table(r_comp::Vector{Any}; styled=true)
    cols = ["time","memory","allocs","iterations","major iterations","c_max","cost"]

    ids = [join(run[1], ", ") for run in r_comp]
    change = [[run[2][key][2] for key in cols] for run in r_comp]
    change = hcat(change...)
    table = [[rpad(run[2][key][1],20) for key in cols] for run in r_comp]
    table = hcat(table...)
    table = vcat(reshape(ids,1,length(ids)),table)

    table = hcat(["IDs"; cols], table)
    col_widths = maximum(length.(table), dims=2)
    n_stats,n_runs = size(table)
    for j = 1:n_runs, i=1:n_stats
        table[i,j] = rpad(table[i,j],col_widths[i])
    end

    header = "| " * join(table[:,1]," | ") * " |";
    hline = "|" * join(["-"^(width+2) for width in col_widths],"|") * "|";

    if styled
        printstyled(header,bold=true); println();
        println(hline)
        colormap = Dict(:invariant=>:grey, :regression=>:red, :improvement=>:green)
        for i = 2:n_runs
            print("| ")
            printstyled(table[1,i],bold=true)
            for j = 2:n_stats
                print(" | ")
                printstyled(table[j,i],color=colormap[change[j-1,i-1]])
            end
            println(" |")
        end
    else
        println(header)
        println(hline)
        for i = 2:n_runs
            println("| " * join(table[:,i]," | ") * " |")
        end
    end
end


text:
using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

                     @variable(m, -1 <= x[i=1:3] <= 1)

                     @variable(m, -19.736879506320445 <= q <= 22.217442927624916)

                     add_NL_constraint(m, :(sigmoid(-0.3136061325511519 + 0.2821341474705279*sigmoid(0.7833824073251248 + 0.1390648465658577*$(x[1]) + -0.6617731040629309*$(x[2]) + -0.5974574624055107*$(x[3])) + 0.4396178028948987*sigmoid(-0.7125949557702831 + 0.9109805514170546*$(x[1]) + 0.42115237771853087*$(x[2]) + 0.7134573642319126*$(x[3])) + -0.24849525966922226*sigmoid(-0.36270963442709325 + 0.07407412170431016*$(x[1]) + -0.6239570881148437*$(x[2]) + -0.6542866380333012*$(x[3])) + 0.5972646324288218*sigmoid(-0.5451890615285553 + -0.19626106399785304*$(x[1]) + 0.5284606976702149*$(x[2]) + 0.2694706847237054*$(x[3])) + -0.8620507270647968*sigmoid(0.35158020081654806 + 0.8300503955389886*$(x[1]) + -0.8477304064925764*$(x[2]) + 0.1840810006485709*$(x[3]))) + sigmoid(-0.9705953045797062 + -0.6259685253929317*sigmoid(0.7833824073251248 + 0.1390648465658577*$(x[1]) + -0.6617731040629309*$(x[2]) + -0.5974574624055107*$(x[3])) + 0.41918630672941015*sigmoid(-0.7125949557702831 + 0.9109805514170546*$(x[1]) + 0.42115237771853087*$(x[2]) + 0.7134573642319126*$(x[3])) + -0.9388772793667424*sigmoid(-0.36270963442709325 + 0.07407412170431016*$(x[1]) + -0.6239570881148437*$(x[2]) + -0.6542866380333012*$(x[3])) + -0.18849282727638839*sigmoid(-0.5451890615285553 + -0.19626106399785304*$(x[1]) + 0.5284606976702149*$(x[2]) + 0.2694706847237054*$(x[3])) + -0.7819718693068158*sigmoid(0.35158020081654806 + 0.8300503955389886*$(x[1]) + -0.8477304064925764*$(x[2]) + 0.1840810006485709*$(x[3]))) + sigmoid(0.8285261526161638 + 0.6863822914423507*sigmoid(0.7833824073251248 + 0.1390648465658577*$(x[1]) + -0.6617731040629309*$(x[2]) + -0.5974574624055107*$(x[3])) + -0.7060930690780838*sigmoid(-0.7125949557702831 + 0.9109805514170546*$(x[1]) + 0.42115237771853087*$(x[2]) + 0.7134573642319126*$(x[3])) + -0.2956343589188948*sigmoid(-0.36270963442709325 + 0.07407412170431016*$(x[1]) + -0.6239570881148437*$(x[2]) + -0.6542866380333012*$(x[3])) + 0.32353098534638214*sigmoid(-0.5451890615285553 + -0.19626106399785304*$(x[1]) + 0.5284606976702149*$(x[2]) + 0.2694706847237054*$(x[3])) + 0.6963022628505415*sigmoid(0.35158020081654806 + 0.8300503955389886*$(x[1]) + -0.8477304064925764*$(x[2]) + 0.1840810006485709*$(x[3]))) + sigmoid(-0.2246585184074279 + 0.8591320330102441*sigmoid(0.7833824073251248 + 0.1390648465658577*$(x[1]) + -0.6617731040629309*$(x[2]) + -0.5974574624055107*$(x[3])) + -0.5037522661683527*sigmoid(-0.7125949557702831 + 0.9109805514170546*$(x[1]) + 0.42115237771853087*$(x[2]) + 0.7134573642319126*$(x[3])) + 0.9621050832972355*sigmoid(-0.36270963442709325 + 0.07407412170431016*$(x[1]) + -0.6239570881148437*$(x[2]) + -0.6542866380333012*$(x[3])) + -0.1995700954953512*sigmoid(-0.5451890615285553 + -0.19626106399785304*$(x[1]) + 0.5284606976702149*$(x[2]) + 0.2694706847237054*$(x[3])) + 0.442251462608958*sigmoid(0.35158020081654806 + 0.8300503955389886*$(x[1]) + -0.8477304064925764*$(x[2]) + 0.1840810006485709*$(x[3]))) + sigmoid(0.5761249271362421 + -0.042724357144423575*sigmoid(0.7833824073251248 + 0.1390648465658577*$(x[1]) + -0.6617731040629309*$(x[2]) + -0.5974574624055107*$(x[3])) + -0.4537006876708345*sigmoid(-0.7125949557702831 + 0.9109805514170546*$(x[1]) + 0.42115237771853087*$(x[2]) + 0.7134573642319126*$(x[3])) + -0.12519521024691826*sigmoid(-0.36270963442709325 + 0.07407412170431016*$(x[1]) + -0.6239570881148437*$(x[2]) + -0.6542866380333012*$(x[3])) + -0.7426009177864223*sigmoid(-0.5451890615285553 + -0.19626106399785304*$(x[1]) + 0.5284606976702149*$(x[2]) + 0.2694706847237054*$(x[3])) + -0.8758186284143128*sigmoid(0.35158020081654806 + 0.8300503955389886*$(x[1]) + -0.8477304064925764*$(x[2]) + 0.1840810006485709*$(x[3]))) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m

max_stars_repo_name: PSORLab/RSActivationFunctions

max_issues_repo_name: PSORLab/RSActivationFunctions

max_forks_repo_name: PSORLab/RSActivationFunctions


text:
"""
    RSOCBridge{T, F, G}

The `RotatedSecondOrderCone` is `SecondOrderCone` representable; see [1, p. 104].
Indeed, we have ``2tu = (t/√2 + u/√2)^2 - (t/√2 - u/√2)^2`` hence
```math
2tu \\ge || x ||_2^2
```
is equivalent to
```math
(t/√2 + u/√2)^2 \\ge || x ||_2^2 + (t/√2 - u/√2)^2.
```
We can therefore use the transformation ``(t, u, x) \\mapsto (t/√2+u/√2, t/√2-u/√2, x)``.
Note that the linear transformation is a symmetric involution (i.e. it is its own transpose and its own inverse).
That means in particular that the norm is of constraint primal and duals are preserved by the tranformation.

[1] Ben-Tal, Aharon, and Arkadi Nemirovski. *Lectures on modern convex optimization: analysis, algorithms, and engineering applications*. Society for Industrial and Applied Mathematics, 2001.
"""
struct RSOCBridge{T, F, G} <: AbstractBridge
    soc::CI{F, MOI.SecondOrderCone}
end
function rotate_function(f::MOI.AbstractVectorFunction, T::Type)
    d = MOI.output_dimension(f)
    f_scalars = MOIU.eachscalar(f)
    t = f_scalars[1]
    u = f_scalars[2]
    x = f_scalars[3:d]
    s2 = √T(2)
    ts = MOIU.operate!(/, T, t, s2)
    us = MOIU.operate!(/, T, u, s2)
    # Cannot use `operate!` here since `ts` and `us` are needed for the next
    # line
    y  = ts - us
    z  = MOIU.operate!(+, T, ts, us)
    return MOIU.operate(vcat, T, z, y, x)
end
function bridge_constraint(::Type{RSOCBridge{T, F, G}}, model,
                           f::MOI.AbstractVectorFunction,
                           s::MOI.RotatedSecondOrderCone) where {T, F, G}
    soc = MOI.add_constraint(model, rotate_function(f, T),
                             MOI.SecondOrderCone(MOI.dimension(s)))
    return RSOCBridge{T, F, G}(soc)
end

function MOI.supports_constraint(::Type{RSOCBridge{T}},
                                ::Type{<:MOI.AbstractVectorFunction},
                                ::Type{MOI.RotatedSecondOrderCone}) where T
    return true
end
MOIB.added_constrained_variable_types(::Type{<:RSOCBridge}) = Tuple{DataType}[]
function MOIB.added_constraint_types(::Type{<:RSOCBridge{T, F}}) where {T, F}
    return [(F, MOI.SecondOrderCone)]
end
function concrete_bridge_type(::Type{<:RSOCBridge{T}},
                              G::Type{<:MOI.AbstractVectorFunction},
                              ::Type{MOI.RotatedSecondOrderCone}) where T
    S = MOIU.promote_operation(/, T, MOIU.scalar_type(G), T)
    Y = MOIU.promote_operation(-, T, S, S)
    Z = MOIU.promote_operation(+, T, S, S)
    F = MOIU.promote_operation(vcat, T, Z, Y, G)
    return RSOCBridge{T, F, G}
end

# Attributes, Bridge acting as a model
function MOI.get(b::RSOCBridge{T, F},
                 ::MOI.NumberOfConstraints{F, MOI.SecondOrderCone}) where {T, F}
    return 1
end
function MOI.get(b::RSOCBridge{T, F},
                 ::MOI.ListOfConstraintIndices{F, MOI.SecondOrderCone}) where {T, F}
    return [b.soc]
end

# References
function MOI.delete(model::MOI.ModelLike, bridge::RSOCBridge)
    MOI.delete(model, bridge.soc)
end

# Attributes, Bridge acting as a constraint
function MOI.get(model::MOI.ModelLike, attr::MOI.ConstraintFunction,
                 bridge::RSOCBridge{T, F, G}) where {T, F, G}
    # As it is an involution, we can just reapply the same transformation
    func = MOI.get(model, attr, bridge.soc)
    return MOIU.convert_approx(G, rotate_function(func, T))
end
function MOI.get(model::MOI.ModelLike, attr::MOI.ConstraintSet, bridge::RSOCBridge)
    set = MOI.get(model, attr, bridge.soc)
    return MOI.RotatedSecondOrderCone(MOI.dimension(set))
end
# As the linear transformation is a symmetric involution,
# the constraint primal and dual both need to be processed by reapplying the
# same transformation
function rotate_result(model,
                       attr::Union{MOI.ConstraintPrimal, MOI.ConstraintDual},
                       ci::MOI.ConstraintIndex)
    x = MOI.get(model, attr, ci)
    T = eltype(x)
    s2 = √T(2)
    return [x[1]/s2 + x[2]/s2; x[1]/s2 - x[2]/s2; x[3:end]]
end
# Need to define both `get` methods and redirect to `rotate_result` to avoid
# ambiguity in dispatch
function MOI.get(model::MOI.ModelLike, attr::MOI.ConstraintPrimal,
                 bridge::RSOCBridge)
    return rotate_result(model, attr, bridge.soc)
end
function MOI.get(model::MOI.ModelLike, attr::MOI.ConstraintDual,
                 bridge::RSOCBridge)
    return rotate_result(model, attr, bridge.soc)
end


text:
using Oceananigans.Solvers
using Oceananigans.Operators
using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid, GridFittedBottom
using Oceananigans.Architectures
using Oceananigans.Grids: with_halo, isrectilinear
using Oceananigans.Fields: Field, ZReducedField
using Oceananigans.Architectures: device

import Oceananigans.Solvers: solve!, precondition!
import Oceananigans.Architectures: architecture

"""
    struct PCGImplicitFreeSurfaceSolver{V, S, R}

The preconditioned conjugate gradient iterative implicit free-surface solver.

$(TYPEDFIELDS)
"""
struct PCGImplicitFreeSurfaceSolver{V, S, R}
    "The vertically-integrated lateral areas"
    vertically_integrated_lateral_areas :: V
    "The preconditioned conjugate gradient solver"
    preconditioned_conjugate_gradient_solver :: S
    "The right hand side of the free surface evolution equation"
    right_hand_side :: R
end

architecture(solver::PCGImplicitFreeSurfaceSolver) =
    architecture(solver.preconditioned_conjugate_gradient_solver)

"""
    PCGImplicitFreeSurfaceSolver(grid, settings)

Return a solver based on a preconditioned conjugate gradient method for
the elliptic equation

```math
[∇ ⋅ H ∇ - 1 / (g Δt²)] ηⁿ⁺¹ = (∇ʰ ⋅ Q★ - ηⁿ / Δt) / (g Δt)
```

representing an implicit time discretization of the linear free surface evolution equation
for a fluid with variable depth `H`, horizontal areas `Az`, barotropic volume flux `Q★`, time
step `Δt`, gravitational acceleration `g`, and free surface at time-step `n` `ηⁿ`.
"""
function PCGImplicitFreeSurfaceSolver(grid::AbstractGrid, settings, gravitational_acceleration=nothing)
    # Initialize vertically integrated lateral face areas
    ∫ᶻ_Axᶠᶜᶜ = Field{Face, Center, Nothing}(with_halo((3, 3, 1), grid))
    ∫ᶻ_Ayᶜᶠᶜ = Field{Center, Face, Nothing}(with_halo((3, 3, 1), grid))

    vertically_integrated_lateral_areas = (xᶠᶜᶜ = ∫ᶻ_Axᶠᶜᶜ, yᶜᶠᶜ = ∫ᶻ_Ayᶜᶠᶜ)

    compute_vertically_integrated_lateral_areas!(vertically_integrated_lateral_areas)

    # Set some defaults
    settings = Dict{Symbol, Any}(settings)
    settings[:maxiter] = get(settings, :maxiter, grid.Nx * grid.Ny)
    settings[:reltol] = get(settings, :reltol, min(1e-7, 10 * sqrt(eps(eltype(grid)))))

    # FFT preconditioner for rectilinear grids, nothing otherwise.
    settings[:preconditioner] = isrectilinear(grid) ?
        get(settings, :preconditioner, FFTImplicitFreeSurfaceSolver(grid)) :
        get(settings, :preconditioner, nothing)

    # TODO: reuse solver.storage for rhs when preconditioner isa FFTImplicitFreeSurfaceSolver?
    right_hand_side = Field{Center, Center, Nothing}(grid)

    solver = PreconditionedConjugateGradientSolver(implicit_free_surface_linear_operation!;
                                                   template_field = right_hand_side,
                                                   settings...)

    return PCGImplicitFreeSurfaceSolver(vertically_integrated_lateral_areas, solver, right_hand_side)
end

build_implicit_step_solver(::Val{:PreconditionedConjugateGradient}, grid, settings, gravitational_acceleration) =
    PCGImplicitFreeSurfaceSolver(grid, settings, gravitational_acceleration)

#####
##### Solve...
#####

function solve!(η, implicit_free_surface_solver::PCGImplicitFreeSurfaceSolver, rhs, g, Δt)
    # Take explicit step first? We haven't found improvement from this yet, but perhaps it will
    # help eventually.
    #event = explicit_ab2_step_free_surface!(free_surface, model, Δt, χ)
    #wait(device(model.architecture), event)

    ∫ᶻA = implicit_free_surface_solver.vertically_integrated_lateral_areas
    solver = implicit_free_surface_solver.preconditioned_conjugate_gradient_solver

    # solve!(x, solver, b, args...) solves A*x = b for x.
    return solve!(η, solver, rhs, ∫ᶻA.xᶠᶜᶜ, ∫ᶻA.yᶜᶠᶜ, g, Δt)
end

function compute_implicit_free_surface_right_hand_side!(rhs, implicit_solver::PCGImplicitFreeSurfaceSolver,
                                                        g, Δt, ∫ᶻQ, η)

    solver = implicit_solver.preconditioned_conjugate_gradient_solver
    arch = architecture(solver)
    grid = solver.grid

    event = launch!(arch, grid, :xy,
                    implicit_free_surface_right_hand_side!,
                    rhs, grid, g, Δt, ∫ᶻQ, η,
		            dependencies = device_event(arch))

    return event
end

""" Compute the divergence of fluxes Qu and Qv. """
@inline flux_div_xyᶜᶜᶜ(i, j, k, grid, Qu, Qv) = δxᶜᵃᵃ(i, j, k, grid, Qu) + δyᵃᶜᵃ(i, j, k, grid, Qv)

@kernel function implicit_free_surface_right_hand_side!(rhs, grid, g, Δt, ∫ᶻQ, η)
    i, j = @index(Global, NTuple)
    Az = Azᶜᶜᶜ(i, j, 1, grid)
    δ_Q = flux_div_xyᶜᶜᶜ(i, j, 1, grid, ∫ᶻQ.u, ∫ᶻQ.v)
    @inbounds rhs[i, j, 1] = (δ_Q - Az * η[i, j, 1] / Δt) / (g * Δt)
end

"""
Returns `L(ηⁿ)`, where `ηⁿ` is the free surface displacement at time step `n`
and `L` is the linear operator that arises
in an implicit time step for the free surface displacement `η`.

(See the docs section on implicit time stepping.)
"""
function implicit_free_surface_linear_operation!(L_ηⁿ⁺¹, ηⁿ⁺¹, ∫ᶻ_Axᶠᶜᶜ, ∫ᶻ_Ayᶜᶠᶜ, g, Δt)
    grid = L_ηⁿ⁺¹.grid
    arch = architecture(L_ηⁿ⁺¹)
    fill_halo_regions!(ηⁿ⁺¹)

    event = launch!(arch, grid, :xy, _implicit_free_surface_linear_operation!,
                    L_ηⁿ⁺¹, grid,  ηⁿ⁺¹, ∫ᶻ_Axᶠᶜᶜ, ∫ᶻ_Ayᶜᶠᶜ, g, Δt,
                    dependencies = device_event(arch))

    wait(device(arch), event)

    return nothing
end

# Kernels that act on vertically integrated / surface quantities
@inline ∫ᶻ_Ax_∂x_ηᶠᶜᶜ(i, j, k, grid, ∫ᶻ_Axᶠᶜᶜ, η) = @inbounds ∫ᶻ_Axᶠᶜᶜ[i, j, k] * ∂xᶠᶜᶜ(i, j, k, grid, η)
@inline ∫ᶻ_Ay_∂y_ηᶜᶠᶜ(i, j, k, grid, ∫ᶻ_Ayᶜᶠᶜ, η) = @inbounds ∫ᶻ_Ayᶜᶠᶜ[i, j, k] * ∂yᶜᶠᶜ(i, j, k, grid, η)

"""
Compute the horizontal divergence of vertically-uniform quantity using
vertically-integrated face areas `∫ᶻ_Axᶠᶜᶜ` and `∫ᶻ_Ayᶜᶠᶜ`.
"""
@inline Az_∇h²ᶜᶜᶜ(i, j, k, grid, ∫ᶻ_Axᶠᶜᶜ, ∫ᶻ_Ayᶜᶠᶜ, η::ZReducedField) =
    (δxᶜᵃᵃ(i, j, k, grid, ∫ᶻ_Ax_∂x_ηᶠᶜᶜ, ∫ᶻ_Axᶠᶜᶜ, η) +
     δyᵃᶜᵃ(i, j, k, grid, ∫ᶻ_Ay_∂y_ηᶜᶠᶜ, ∫ᶻ_Ayᶜᶠᶜ, η))

"""
    _implicit_free_surface_linear_operation!(L_ηⁿ⁺¹, grid, ηⁿ⁺¹, ∫ᶻ_Axᶠᶜᶜ, ∫ᶻ_Ayᶜᶠᶜ, g, Δt)

Return the left side of the "implicit η equation"

```math
(∇ʰ⋅ H ∇ʰ - 1 / (g Δt²)) ηⁿ⁺¹ = 1 / (g Δt) ∇ʰ ⋅ Q★ - 1 / (g Δt²) ηⁿ
----------------------
        ≡ L_ηⁿ⁺¹
```

which is derived from the discretely summed barotropic mass conservation equation,
and arranged in a symmetric form by multiplying by horizontal areas Az:

```
δⁱÂʷ∂ˣηⁿ⁺¹ + δʲÂˢ∂ʸηⁿ⁺¹ - Az ηⁿ⁺¹ / (g Δt²) = 1 / (g Δt) (δⁱÂʷu̅ˢᵗᵃʳ + δʲÂˢv̅ˢᵗᵃʳ) - Az ηⁿ / (g Δt²)
```

where  ̂ indicates a vertical integral, and
       ̅ indicates a vertical average
"""
@kernel function _implicit_free_surface_linear_operation!(L_ηⁿ⁺¹, grid, ηⁿ⁺¹, ∫ᶻ_Axᶠᶜᶜ, ∫ᶻ_Ayᶜᶠᶜ, g, Δt)
    i, j = @index(Global, NTuple)
    Az = Azᶜᶜᶜ(i, j, 1, grid)
    @inbounds L_ηⁿ⁺¹[i, j, 1] = Az_∇h²ᶜᶜᶜ(i, j, 1, grid, ∫ᶻ_Axᶠᶜᶜ, ∫ᶻ_Ayᶜᶠᶜ, ηⁿ⁺¹) - Az * ηⁿ⁺¹[i, j, 1] / (g * Δt^2)
end

#####
##### Preconditioners
#####

"""
add to the rhs - H⁻¹ ∇H ⋅ ∇ηⁿ to the rhs...
"""
@inline function precondition!(P_r, preconditioner::FFTImplicitFreeSurfaceSolver, r, η, ∫ᶻ_Axᶠᶜᶜ, ∫ᶻ_Ayᶜᶠᶜ, g, Δt)
    poisson_solver = preconditioner.fft_poisson_solver
    arch = architecture(poisson_solver)
    grid = preconditioner.three_dimensional_grid
    Az = grid.Δxᶜᵃᵃ * grid.Δyᵃᶜᵃ # assume horizontal regularity
    Lz = grid.Lz

    event = launch!(arch, grid, :xy,
                    fft_preconditioner_right_hand_side!,
                    poisson_solver.storage, r, η, grid, Az, Lz,
                    dependencies = device_event(arch))

    wait(device(arch), event)

    return solve!(P_r, preconditioner, poisson_solver.storage, g, Δt)
end

@kernel function fft_preconditioner_right_hand_side!(fft_rhs, pcg_rhs, η, grid, Az, Lz)
    i, j = @index(Global, NTuple)
    @inbounds fft_rhs[i, j, 1] = pcg_rhs[i, j, 1] / (Lz * Az)
end

# TODO: make it so adding this term:
#
#   - ∇H_∇η(i, j, 1, grid, η) / H
#
# speeds up the convergence.
#=
@inline ∇H_∇η(i, j, k, grid, η) = zero(eltype(grid)) # fallback
@inline depth(i, j, k, grid) = grid.Lz

const GFBIBG = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:GridFittedBottom}

# Assumes surface is z=0:
@inline depth(i, j, k, ibg::GFBIBG) = @inbounds max(zero(eltype(ibg)), min(ibg.Lz, -ibg.immersed_boundary.bottom[i, j]))
@inline ∂x_H_∂x_η(i, j, k, ibg, η) = ∂xᶠᶜᶜ(i, j, k, ibg, depth) * ∂xᶠᶜᶜ(i, j, k, ibg, η)
@inline ∂y_H_∂y_η(i, j, k, ibg, η) = ∂yᶜᶠᶜ(i, j, k, ibg, depth) * ∂yᶜᶠᶜ(i, j, k, ibg, η)
@inline ∇H_∇η(i, j, k, ibg::GFBIBG, η) = ℑxᶜᵃᵃ(i, j, k, ibg, ∂x_H_∂x_η, η) + ℑyᵃᶜᵃ(i, j, k, ibg, ∂y_H_∂y_η, η)

@inline function H⁻¹_∇H_∇η(i, j, k, ibg::GFBIBG, η)
    H = depth(i, j, k, ibg)
    return ifelse(H == 0, zero(eltype(ibg)), ∇H_∇η(i, j, k, ibg, η) / H)
end

# The rhs below becomes pcg_rhs[i, j, 1] / (H * Az) - ∇H_∇η(i, j, 1, grid, η) / H
=#

#####
##### "Asymptotically diagonally-dominant" preconditioner
#####

struct DiagonallyDominantInversePreconditioner end

@inline precondition!(P_r, ::DiagonallyDominantInversePreconditioner, r, ∫ᶻ_Axᶠᶜᶜ, ∫ᶻ_Ayᶜᶠᶜ, g, Δt) =
    diagonally_dominant_precondition!(P_r, r, ∫ᶻ_Axᶠᶜᶜ, ∫ᶻ_Ayᶜᶠᶜ, g, Δt)

"""
    _diagonally_dominant_precondition!(P_r, grid, r, ∫ᶻ_Axᶠᶜᶜ, ∫ᶻ_Ayᶜᶠᶜ, g, Δt)

Return the diagonally dominant inverse preconditioner applied to the residuals consistently with
 `M = D⁻¹(I - (A - D)D⁻¹) ≈ A⁻¹` where `I` is the Identity matrix,
A is the linear operator applied to η, and D is the diagonal of A.

```math
P_r = M * r
```
which expanded in components is
```math
P_rᵢⱼ = rᵢⱼ / Acᵢⱼ - 1 / Acᵢⱼ ( Ax⁻ / Acᵢ₋₁ rᵢ₋₁ⱼ + Ax⁺ / Acᵢ₊₁ rᵢ₊₁ⱼ + Ay⁻ / Acⱼ₋₁ rᵢⱼ₋₁+ Ay⁺ / Acⱼ₊₁ rᵢⱼ₊₁ )
```

where `Ac`, `Ax⁻`, `Ax⁺`, `Ay⁻` and `Ay⁺` are the coefficients of
`ηᵢⱼ`, `ηᵢ₋₁ⱼ`, `ηᵢ₊₁ⱼ`, `ηᵢⱼ₋₁` and `ηᵢⱼ₊₁` in `_implicit_free_surface_linear_operation!`
"""
function diagonally_dominant_precondition!(P_r, r, ∫ᶻ_Axᶠᶜᶜ, ∫ᶻ_Ayᶜᶠᶜ, g, Δt)
    grid = ∫ᶻ_Axᶠᶜᶜ.grid
    arch = architecture(P_r)

    fill_halo_regions!(r)

    event = launch!(arch, grid, :xy, _diagonally_dominant_precondition!,
                    P_r, grid, r, ∫ᶻ_Axᶠᶜᶜ, ∫ᶻ_Ayᶜᶠᶜ, g, Δt,
                    dependencies = device_event(arch))

    wait(device(arch), event)

    return nothing
end

# Kernels that calculate coefficients for the preconditioner
@inline Ax⁻(i, j, grid, ax) = @inbounds   ax[i, j, 1] / Δxᶠᶜᶜ(i, j, 1, grid)
@inline Ay⁻(i, j, grid, ay) = @inbounds   ay[i, j, 1] / Δyᶜᶠᶜ(i, j, 1, grid)
@inline Ax⁺(i, j, grid, ax) = @inbounds ax[i+1, j, 1] / Δxᶠᶜᶜ(i+1, j, 1, grid)
@inline Ay⁺(i, j, grid, ay) = @inbounds ay[i, j+1, 1] / Δyᶜᶠᶜ(i, j+1, 1, grid)

@inline Ac(i, j, grid, g, Δt, ax, ay) = - Ax⁻(i, j, grid, ax) -
                                          Ax⁺(i, j, grid, ax) -
                                          Ay⁻(i, j, grid, ay) -
                                          Ay⁺(i, j, grid, ay) -
                                          Azᶜᶜᶜ(i, j, 1, grid) / (g * Δt^2)

@inline heuristic_inverse_times_residuals(i, j, r, grid, g, Δt, ax, ay) =
    @inbounds 1 / Ac(i, j, grid, g, Δt, ax, ay) * (r[i, j, 1] - Ax⁻(i, j, grid, ax) / Ac(i-1, j, grid, g, Δt, ax, ay) * r[i-1, j, 1] -
                                                                Ax⁺(i, j, grid, ax) / Ac(i+1, j, grid, g, Δt, ax, ay) * r[i+1, j, 1] -
                                                                Ay⁻(i, j, grid, ay) / Ac(i, j-1, grid, g, Δt, ax, ay) * r[i, j-1, 1] -
                                                                Ay⁺(i, j, grid, ay) / Ac(i, j+1, grid, g, Δt, ax, ay) * r[i, j+1, 1])

@kernel function _diagonally_dominant_precondition!(P_r, grid, r, ∫ᶻ_Axᶠᶜᶜ, ∫ᶻ_Ayᶜᶠᶜ, g, Δt)
    i, j = @index(Global, NTuple)
    @inbounds P_r[i, j, 1] = heuristic_inverse_times_residuals(i, j, r, grid, g, Δt, ∫ᶻ_Axᶠᶜᶜ, ∫ᶻ_Ayᶜᶠᶜ)
end


text:
struct SquareRange
  xRange :: Tuple{Int,Int}
  yRange :: Tuple{Int,Int}
end

mutable struct Probe
  pos :: Tuple{Int,Int}
  vel :: Tuple{Int,Int}
end

Probe(v::Tuple{Int,Int}) = Probe((0,0),v)

function step!(p::Probe)
  p.pos = p.pos .+ p.vel
  p.vel = p.vel .+ (-sign(p.vel[1]),-1)
end

function inRange(p::Probe, r::SquareRange) :: Bool
  return (p.pos[1] >= r.xRange[1] && p.pos[1] <= r.xRange[2] && p.pos[2] >= r.yRange[1] && p.pos[2] <= r.yRange[2])
end

function checkTrajectory(p::Probe, r::SquareRange) :: Int
  yMax :: Int = 0
  while(p.pos[2] > r.yRange[1])
    step!(p)
    yMax = max(yMax,p.pos[2])
    inRange(p,r) && return yMax
  end
  return -1
end

function readInput()
  xRange = (135,155)
  yRange = (-102,-78)
  return SquareRange(xRange, yRange)
end

# For the first problem, the 'x' does not matter (since it does not affect y or vy in any way).
# This means we can set vx = 16 or 17 (velocities for which we will drop straight down inside the range)
# and just maximize for vy.
function problem1(r::SquareRange) :: Tuple{Int,Int}
  vmax :: Int = -1
  ymax :: Int = 0
  for v in 0:1000
    y = checkTrajectory(Probe((17,v)),r)
    y > ymax && begin vmax = v; ymax = y; end
  end
  return (vmax,ymax)
end

# For the second problem, we need to consider the velocity range that is possibe. In order
# to do so, we can try to set some approppriate ranges for both directions:
#  - If vx < 16, the probe stops before reaching the range.
#  - If vx > 156, the probe will pass right through.
# For vy, we have that vy = -|vy_0| at y=0 (either at the start of when it crosses over) so
# we have that
#  - If vy > 103, the probe goes through the range without stoping in it.
#  - If vy < -103, same thing happens
function problem2(r::SquareRange) :: Int
  cnter :: Int = 0
  for vx in 16:156
    for vy in -103:103
      (checkTrajectory(Probe((vx,vy)),r) != -1) && (cnter += 1)
    end
  end
  return cnter
end

range = readInput()
res = problem1(range)
println(res)
res = problem2(range)
println(res)


text:
function all_same_type(::Type{T}) where T
    r = Vector{Any}()
    for s in fieldnames(T)
        push!(r,fieldtype(T,s))
    end
    re = all(y->y===r[1],r)
    return re
end

function is_tuple_wrapper(::Type{ST}) where ST
    return fieldcount(ST) == 1 && fieldtype(ST,1) <: NTuple
end

@generated function check_parent_and_struct_match(::Type{ST},::Type{T}) where {ST,T}

    t = ST
    allsame = false
    isstruct = false
    NF = 1
    if fieldcount(ST) != 0
        isstruct = true
        if is_tuple_wrapper(ST)
            IST = fieldtype(ST,1)
            NF = fieldcount(IST)
            t = fieldtype(IST,1)
            allsame = all_same_type(IST)
        else
            NF = fieldcount(ST)
            t = fieldtype(ST,1)
            allsame = all_same_type(ST)
        end
    end

    correctArrayType = t === T

    return quote
        $isstruct || throw(ArgumentError("Type $ST is not supported, only composite types."))

        $allsame || throw(ArgumentError("Struct $ST not supported. All struct elements must have the same type."))

        $correctArrayType || throw(ArgumentError("Array type does not match struct fields type."))

        return $NF
    end
end

myfieldcount(::Type{ST}) where {ST} = is_tuple_wrapper(ST) ? fieldcount(fieldtype(ST,1)) : fieldcount(ST)


text:
module TestStatic

using Test, MLJBase
using StableRNGs

rng = StableRNG(5312515)

## SIMPLE UNIVARIATE FUNCTION

mutable struct Scale <: MLJBase.Static
    scaling::Float64
end

function MLJBase.transform(s::Scale, _, X)
    X isa AbstractVecOrMat && return X * s.scaling
    MLJBase.table(s.scaling * MLJBase.matrix(X), prototype=X)
end

function MLJBase.inverse_transform(s::Scale, _, X)
    X isa AbstractVecOrMat && return X / s.scaling
    MLJBase.table(MLJBase.matrix(X) / s.scaling, prototype=X)
end

s = Scale(2)
X = randn(rng, 2, 3)
Xt = MLJBase.table(X)

R = transform(s, nothing, X)
IR = inverse_transform(s, nothing, R)
@test IR ≈ X

R = transform(s, nothing, Xt)
IR = inverse_transform(s, nothing, R)
@test MLJBase.matrix(IR) ≈ X

## MULTIVARIATE FUNCTION

mutable struct PermuteArgs <: MLJBase.Static
    permutation::NTuple{N,Int} where N
end

MLJBase.transform(p::PermuteArgs, _, args...) =
    Tuple([args[i] for i in p.permutation])
MLJBase.inverse_transform(p::PermuteArgs, _, args...) =
    Tuple([args[i] for i in sortperm(p.permutation |> collect)])

p = PermuteArgs((2, 3, 1))
@test transform(p, nothing, 10, 20, 30) == (20, 30, 10)
@test inverse_transform(p, nothing, 20, 30, 10) == (10, 20, 30)

# no-op
fitresult, _, _ = MLJBase.fit(p, 1, (1, 2, 3))

@test transform(p, fitresult, 10, 20, 30) == (20, 30, 10)
@test inverse_transform(p, fitresult, 20, 30, 10) == (10, 20, 30)

end
true