using LinearAlgebra
using Statistics
using Random

function abs_val(d)
	return abs(d)
end

function obj(X, y, w, lam, tau)
    n, ncol = size(X);
    pen = 0;
    obj_fun = 0;
    r = y - X * w;
    rw = zeros(n);
    for q_i in 1:n
        if r[q_i] < 0
            rw[q_i] = r[q_i] * r[q_i] * (1 - tau);
        else
            rw[q_i] = r[q_i] * r[q_i] * tau;
        end
    end
    pen = lam * sum(broadcast(abs, w));
    obj_fun = sum(rw) + pen;
    return obj_fun;
end

function median_index(vector)
    med = median(vector);
    abs_demed_vector = broadcast(abs, vector .- med);
    idx = argmin(abs_demed_vector);
    return idx;
end

function buckets(K_buckets, nrow)
    tmp_buckets_id = range(0, stop=K_buckets-1, length=K_buckets);
    tmp2buckets_id = repeat(tmp_buckets_id, 1, Int(ceil(nrow/K_buckets)));
    buckets_id = sort(vec(tmp2buckets_id)[1:nrow]);
    return buckets_id;
end

function cvfolds(nfolds, nrow)
    tmp_foldsid = range(0, stop=nfolds-1, length=nfolds);
    tmp2_foldsid = repeat(tmp_foldsid, 1, Int(ceil(nrow/nfolds)));
    foldsid = sort(vec(tmp2_foldsid)[1:nrow]);
    return foldsid;
end

function block_X(X, which_idx)
    return X[which_idx,:];
end

function block_y(y, which_idx)
    return y[which_idx];
end

function mom_obj(X, y, w, w_prime, lam, tau)
    loss = obj(X, y, w, lam, tau);
    loss_prime = obj(X, y, w_prime, lam, tau);
    loss_mom = loss - loss_prime;
    return loss_mom;
end

function f_grad(X, y, w, tau)
    n, ncol = size(X);
    r = y - X * w;
    grad_tmp = zeros(n);
    for q_i in 1:n
        if r[q_i] < 0
            grad_tmp[q_i] = 2 * r[q_i] * (1 - tau);
        else
            grad_tmp[q_i] = 2 * r[q_i] * tau;
        end
    end
    grad = -transpose(X) * grad_tmp;
    return grad;
end

function block(X, y, w, w_prime, K, tau, is_random)
	vect_means = zeros(K);
	n = size(y)[1];
	buckets_id = buckets(K, n);
	if is_random == true
		buckets_id = shuffle(buckets_id);
	end
	for k in 1:K
		Xk = X[findall(x -> x==(k-1), buckets_id), :];
		yk = y[findall(x -> x==(k-1), buckets_id)];
		excess_loss_k = obj(Xk, yk, w, 0, tau) - obj(Xk, yk, w_prime, 0, tau);
		vect_means[k] = excess_loss_k;
	end
	idx = median_index(vect_means);
	which_idx = findall(x -> x==(idx-1), buckets_id);
	return which_idx;
end

function block_cv(X_test, y_test, w, K_prime, tau)
	vect_means = zeros(K_prime);
	n = size(y_test)[1];
	buckets_id = buckets(K_prime, n);
	for k in 1:K_prime
		Xk = X_test[findall(x -> x==(k-1), buckets_id), :];
		yk = y_test[findall(x -> x==(k-1), buckets_id)];
		excess_loss_k = obj(Xk, yk, w, 0, tau);
		vect_means[k] = excess_loss_k;
	end
  idx = median_index(vect_means);
  which_idx = findall(x -> x==(idx-1), buckets_id);
  return which_idx;
end

function soft_threshold(w, mu)
	p = size(w)[1];
	max_tmp = zeros(p);
	for i in 1:p
		if abs_val(w[i]) > mu
			max_tmp[i] = abs_val(w[i]) - mu;
		else
			max_tmp[i] = 0.0;
		end
	end
	st = sign.(w) .* max_tmp;
	return st;
end

function pg_als(X, y, beta_0, tau, max_iter, lam)
	w = beta_0;
	w_prev = w;
	obj_pg_als = zeros(max_iter+1);
	tmp = obj(X, y, w, lam, tau);
	obj_pg_als[1] = tmp;

	beta = 1.2;
	delta = 1.0;
	LL = 1.0;
	for j in 1:max_iter
		w_prev = w;
		delta = 1.0;
		obj_val = obj(X, y, w, lam, tau);
		while (delta > 0.001)
			gamma = 1.0 / LL;
			w = w_prev - gamma * f_grad(X, y, w_prev, tau);
			w = soft_threshold(w, lam * gamma);
			obj_val_new = obj(X, y, w, lam, tau);
			f_grad_w = sum(transpose(f_grad(X, y, w_prev, tau)) * (w - w_prev));
			delta = obj_val_new - obj_val - f_grad_w - (LL / 2) * norm(w-w_prev, 2) ^ 2;
			LL = LL * beta;
		end
		LL = LL / beta;
		tmp = obj(X, y, w, lam, tau);
		obj_pg_als[j+1] = tmp;
	end
	out = Dict("fitted_values" => obj_pg_als, "beta" => w);
	return out;
end

function pg_mom_als(X, y, beta_0, tau, K, max_iter, lam, epsil, is_random)
	w = beta_0;
	w_prime = beta_0;
	mom_obj_pg_als = zeros(max_iter+1);
	tmp = mom_obj(X, y, w, w_prime, lam, tau);
	mom_obj_pg_als[1] = tmp;
	beta = 1.2;
	stopping = false;
	counter = 0.0;
	tmp_obj = 0.0;
	tmp_obj_prev = 0.0;
	while (counter < max_iter && stopping == false)
		counter = counter + 1.0;
		LL = 1.0;
		w_prev = w;
		delta = ones(1);
		idx_k = block(X, y, w, w_prime, K, tau, is_random);
		Xk = block_X(X, idx_k);
		yk = block_y(y, idx_k);
		obj_val = obj(Xk, yk, w, lam, tau);
		while (delta[1] > 0.001)
			gamma = 1 / LL;
			w = w_prev - gamma * f_grad(Xk, yk, w_prev, tau);
			w = soft_threshold(w, lam * gamma);
			delta = obj(Xk, yk, w, lam, tau) - obj_val - transpose(Array(f_grad(Xk, yk, w_prev, tau))) * (w - w_prev) - (LL / 2) * norm(w-w_prev, 2) ^ 2;
			LL = LL * beta;
		end
		LL_prime = 1.0;
		w_prime_prev = w_prime;
		delta_prime = ones(1);
		idx_k = block(X, y, w, w_prime, K, tau, is_random);
		Xk = block_X(X, idx_k);
		yk = block_y(y, idx_k);
		obj_val = obj(Xk, yk, w_prime, lam, tau);
		while (delta_prime[1] > 0.001)
			gamma = 1 / LL_prime;
			w_prime = w_prime_prev - gamma * f_grad(Xk, yk, w_prime_prev, tau);
			w_prime = soft_threshold(w_prime, lam * gamma);
			delta_prime = obj(Xk, yk, w_prime, lam, tau) - obj_val - transpose(Array(f_grad(Xk, yk, w_prime_prev, tau))) * (w_prime - w_prime_prev) - (LL_prime / 2) * norm(w_prime - w_prime_prev, 2) ^ 2;
			LL_prime = LL_prime * beta;
		end
		tmp_obj_prev = tmp_obj;
		tmp_obj = mom_obj(X, y, w, w_prime, lam, tau);
		mom_obj_pg_als[Int(counter)+1] = tmp_obj;
		#if (abs_val(tmp_obj) <= epsil)
		#	if (abs_val(tmp_obj) + abs_val(tmp_obj_prev) <= 2 * epsil)
		#		stopping = true;
		#		counter = counter + 1.0;
		#	end
		#end
	end
	out = Dict("fitted_values" => mom_obj_pg_als, "beta" => w, "beta_prime" => w_prime, "counter" => counter, "stopping" => stopping);
	return out
end

function pg_mom_als_w(X, y, beta_0, tau, K, max_iter, lam, epsil, is_random)
	w = beta_0;
	w_prime = beta_0;
	mom_obj_pg_als = zeros(max_iter+1);
	tmp = mom_obj(X, y, w, w_prime, lam, tau);
	mom_obj_pg_als[1] = tmp;
	beta = 1.2;
    stopping = false;
	counter = 0.0;
	tmp_obj = 0.0;
	tmp_obj_prev = 0.0;
	while (counter < max_iter && stopping == false)
		counter = counter + 1.0;
		LL = 1.0;
		w_prev = w;
		delta = ones(1);
		idx_k = block(X, y, w, w_prime, K, tau, is_random);
		Xk = block_X(X, idx_k);
		yk = block_y(y, idx_k);
		obj_val = obj(Xk, yk, w, lam, tau);
		while (delta[1] > 0.001)
			gamma = 1 / LL;
			w = w_prev - gamma * f_grad(Xk, yk, w_prev, tau);
			w = soft_threshold(w, lam * gamma);
			delta = obj(Xk, yk, w, lam, tau) - obj_val - transpose(Array(f_grad(Xk, yk, w_prev, tau)))*(w - w_prev)-(LL/2)* norm(w-w_prev, 2) ^ 2;
			LL = LL*beta;
		end
		LL_prime = 1.0;
		w_prime_prev = w_prime;
		delta_prime = ones(1);
		idx_k = block(X, y, w, w_prime, K, tau, is_random);
		Xk = block_X(X, idx_k);
		yk = block_y(y, idx_k);
		obj_val = obj(Xk, yk, w_prime, lam, tau);
		while (delta_prime[1] > 0.001)
			gamma = 1 / LL_prime;
			w_prime = w_prime_prev - gamma * f_grad(Xk, yk, w_prime_prev, tau);
			w_prime = soft_threshold(w_prime, lam * gamma);
			delta_prime = obj(Xk, yk, w_prime, lam, tau) - obj_val - transpose(Array(f_grad(Xk, yk, w_prime_prev, tau)))*(w_prime - w_prime_prev)-(LL_prime/2)*norm(w_prime - w_prime_prev, 2) ^ 2;
			LL_prime = LL_prime*beta;
		end
		tmp_obj_prev = tmp_obj;
		tmp_obj = mom_obj(X, y, w, w_prime, lam, tau);
		mom_obj_pg_als[Int(counter)+1] = tmp_obj;
		if (abs_val(tmp_obj) <= epsil)
			if (abs_val(tmp_obj)+abs_val(tmp_obj_prev) <= 2*epsil)
				stopping = true;
				counter = counter + 1.0;
			end
		end
	end
	return w;
end

"""

X = [1 2 3; 4 5 6; 7 8 9];
y = [2, -4, 6];
w = [1, 2, 5];
w_prime = [1, 2, 3];
tau = 0.5;
lam = 2.0;
K = 3;
K_buckets = 3;
nrow = 10;
max_iter = 10;
epsil = 0.001;
is_random = false;

println(obj(X, y, w, lam, tau))
println(median_index(y));
println(buckets(K_buckets, nrow));
println(cvfolds(K_buckets, nrow));
println(block_X(X, [1,2]));
println(block_y(X, [1,2]));
println(mom_obj(X, y, w, w_prime, lam, tau));
println(f_grad(X, y, w, tau));
println(block(X, y, w, w_prime, K, tau, is_random))
println(pg_als(X, y, w, tau, max_iter, lam))
println(pg_mom_als(X, y, w, tau, K, max_iter, lam, epsil, is_random))
println(pg_mom_als_w(X, y, w, tau, K, max_iter, lam, epsil, is_random))

"""


using LazySets, Plots

include("SparsePolynomialZonotope.jl")

# ==========
# Slide 4
# ==========
Z = Zonotope([0, 0.], [2 1 2; 0 2 2.])
fig = plot(Z)

# equivalent:
# Z = Zonotope([0, 0.], [[2, 0], [1, 2.], [2, 2.]])

# plot generators
plot!(fig, [0, 2.], [0, 0.], arrow=:arrow, c=:black, lw=2.0, lab="")
plot!(fig, [0, 1.], [0, 2.], arrow=:arrow, c=:black, lw=2.0, lab="")
plot!(fig, [0, 2.], [0, 2.], arrow=:arrow, c=:black, lw=2.0, lab="")
savefig(fig, "slide_4.png")

# ==========
# Slide 5
# ==========
SP = SparsePolynomialZonotope([0, 0.], [2 1 2; 0 2 2.], [1 0 3; 0 1 1.])
#fig = plot(SP)

# ==========
# Slide 10
# ==========

# ==========
# Slide 12
# ==========

# ==========
# Slide 13
# ==========


"""
    ThermistorHW

Cálculo de propriedades de termistores e sua aplicação em
termo-anemometria.

"""
module ThermistorHW

# package code goes here

export Thermistor, Resistor, resistance, temperature
export myhorner, IdealGas, AirType, Air, heatcond, viscosity, density, kinviscosity, specheat
export thermaldiff, prandtl, volthermalexpansion
export reynolds, nusphere, nucylinder, grashof, rayleigh, nuspherefree, hconvect, hconvectcyl
export Solid, Fe2O3, Fe₂O₃, const_current, const_temperature, thermalcond
export HWSensor

abstract type AbstractSubstance end
abstract type AbstractFluid <: AbstractSubstance end
abstract type AbstractGas <: AbstractFluid end
abstract type AbstractIdealGas <: AbstractGas end

include("thermistor.jl")
include("air.jl")
include("solid.jl")
include("heattrans.jl")
include("anemometer.jl")

end # module


using NumCME
using Test
using Sundials: CVODE_BDF

𝕊 = [[-1, 1, 0] [1, -1, 0] [0, 0, 1] [0, 0, -1]]
x₀ = [1, 0, 0]
k₀₁ = 0.05
k₁₀ = 0.1
λ = 5.0
γ = 0.5
L = 20.0
θ = [k₀₁, k₁₀, λ, γ, L]
a1 = propensity() do x, p
    p[1] * x[1]
end
a2 = propensity() do x, p
    p[2] * x[2]
end
a3 = propensity() do x, p
    p[3] * x[2]
end
a4 = propensity((t, p) -> max(0.0, 1.0 - sin(π * t / p[5]))) do x, p
    p[4] * x[3]
end
propensities = [a1, a2, a3, a4]
sensmodel = CmeModelWithSensitivity(CmeModel(𝕊, propensities, θ))
init_cond = forwardsens_initial_condition([x₀], [1.0], [[0.0] for i in 1:5])
sensfspalgorithm = AdaptiveForwardSensFspSparse(
    ode_method = CVODE_BDF(linear_solver = :GMRES),
    space_adapter = ForwardSensRStepAdapter(10, 10, true)
)
senssol = solve(sensmodel,
    init_cond,
    (0.0, 400.0),
    sensfspalgorithm;
    saveat = [],
    fsptol = 1.0E-6,
    odeatol = 1.0E-10,
    odertol = 1.0E-4)


@doc raw"""
    AbstractKernel

Integral kernel `K(x, y)`.

Abstract base type for an integral kernel, i.e. a AbstractFloat binary function
``K(x, y)`` used in a Fredhold integral equation of the first kind:
```math
    u(x) = ∫ K(x, y) v(y) dy
```
where ``x ∈ [x_\mathrm{min}, x_\mathrm{max}]`` and
``y ∈ [y_\mathrm{min}, y_\mathrm{max}]``.  For its SVE to exist,
the kernel must be square-integrable, for its singular values to decay
exponentially, it must be smooth.

In general, the kernel is applied to a scaled spectral function ``ρ'(y)`` as:
```math
    ∫ K(x, y) ρ'(y) dy,
```
where ``ρ'(y) = w(y) ρ(y)``.
"""
abstract type AbstractKernel end

@doc raw"""
    LogisticKernel <: AbstractKernel

Fermionic/bosonic analytical continuation kernel.

In dimensionless variables ``x = 2 τ/β - 1``, ``y = β ω/Λ``,
the integral kernel is a function on ``[-1, 1] × [-1, 1]``:
```math
    K(x, y) = \frac{e^{-Λ y (x + 1) / 2}}{1 + e^{-Λ y}}
```
LogisticKernel is a fermionic analytic continuation kernel.
Nevertheless, one can model the ``τ`` dependence of
a bosonic correlation function as follows:
```math
    ∫ \frac{e^{-Λ y (x + 1) / 2}}{1 - e^{-Λ y}} ρ(y) dy = ∫ K(x, y) ρ'(y) dy,
```
with
```math
    ρ'(y) = w(y) ρ(y),
```
where the weight function is given by
```math
    w(y) = \frac{1}{\tanh(Λ y/2)}.
```
"""
struct LogisticKernel <: AbstractKernel
    Λ::Float64
    function LogisticKernel(Λ)
        Λ ≥ 0 || throw(DomainError(Λ, "Kernel cutoff Λ must be non-negative"))
        return new(Λ)
    end
end

@doc raw"""
    RegularizedBoseKernel <: AbstractKernel

Regularized bosonic analytical continuation kernel.

In dimensionless variables ``x = 2 τ/β - 1``, ``y = β ω/Λ``, the fermionic
integral kernel is a function on ``[-1, 1] × [-1, 1]``:
```math
    K(x, y) = y \frac{e^{-Λ y (x + 1) / 2}}{e^{-Λ y} - 1}
```
Care has to be taken in evaluating this expression around ``y = 0``.
"""
struct RegularizedBoseKernel <: AbstractKernel
    Λ::Float64
    function RegularizedBoseKernel(Λ)
        Λ ≥ 0 || throw(DomainError(Λ, "Kernel cutoff Λ must be non-negative"))
        return new(Λ)
    end
end

"""
    AbstractSVEHints

Discretization hints for singular value expansion of a given kernel.
"""
abstract type AbstractSVEHints end

struct SVEHintsLogistic{T} <: AbstractSVEHints
    kernel::LogisticKernel
    ε::T
end

struct SVEHintsRegularizedBose{T} <: AbstractSVEHints
    kernel::RegularizedBoseKernel
    ε::T
end

struct SVEHintsReduced{T<:AbstractSVEHints} <: AbstractSVEHints
    inner_hints::T
end

abstract type AbstractReducedKernel <: AbstractKernel end

@doc raw"""
    ReducedKernel

Restriction of centrosymmetric kernel to positive interval.

For a kernel ``K`` on ``[-1, 1] × [-1, 1]`` that is centrosymmetric, i.e.
``K(x, y) = K(-x, -y)``, it is straight-forward to show that the left/right
singular vectors can be chosen as either odd or even functions.

Consequentially, they are singular functions of a reduced kernel ``K_\mathrm{red}``
on ``[0, 1] × [0, 1]`` that is given as either:
```math
    K_\mathrm{red}(x, y) = K(x, y) \pm K(x, -y)
```
This kernel is what this type represents.  The full singular functions can
be reconstructed by (anti-)symmetrically continuing them to the negative
axis.
"""
struct ReducedKernel{K<:AbstractKernel} <: AbstractReducedKernel
    inner::K
    sign::Int
end

@doc raw"""
    LogisticKernelOdd <: AbstractReducedKernel

Fermionic analytical continuation kernel, odd.

In dimensionless variables ``x = 2τ/β - 1``, ``y = βω/Λ``, the fermionic
integral kernel is a function on ``[-1, 1] × [-1, 1]``:
```math
    K(x, y) = -\frac{\sinh(Λ x y / 2)}{\cosh(Λ y / 2)}
```
"""
struct LogisticKernelOdd <: AbstractReducedKernel
    inner::LogisticKernel
    sign::Int

    function LogisticKernelOdd(inner::LogisticKernel, sign)
        iscentrosymmetric(inner) || error("inner kernel must be centrosymmetric")
        abs(sign) == 1 || throw(DomainError(sign, "sign must be -1 or 1"))
        return new(inner, sign)
    end
end

@doc raw"""
    RegularizedBoseKernelOdd <: AbstractReducedKernel

Bosonic analytical continuation kernel, odd.

In dimensionless variables ``x = 2 τ / β - 1``, ``y = β ω / Λ``, the fermionic
integral kernel is a function on ``[-1, 1] × [-1, 1]``:
```math
    K(x, y) = -y \frac{\sinh(Λ x y / 2)}{\sinh(Λ y / 2)}
```
"""
struct RegularizedBoseKernelOdd <: AbstractReducedKernel
    inner::RegularizedBoseKernel
    sign::Int

    function RegularizedBoseKernelOdd(inner::RegularizedBoseKernel, sign)
        iscentrosymmetric(inner) || error("inner kernel must be centrosymmetric")
        abs(sign) == 1 || throw(DomainError(sign, "sign must be -1 or 1"))
        return new(inner, sign)
    end
end

@doc raw"""
    xrange(kernel)

Return a tuple ``(x_\mathrm{min}, x_\mathrm{max})`` delimiting the range
of allowed `x` values.
"""
xrange(::AbstractKernel) = (-1, 1)
xrange(kernel::AbstractReducedKernel) = (0, last(xrange(kernel.inner)))

@doc raw"""
    yrange(kernel)

Return a tuple ``(y_\mathrm{min}, y_\mathrm{max})`` delimiting the range
 of allowed `y` values.
"""
yrange(::AbstractKernel) = (-1, 1)
yrange(kernel::AbstractReducedKernel) = (0, last(yrange(kernel.inner)))

function compute(::LogisticKernel, u₊, u₋, v)
    # By introducing u_± = (1 ± x)/2 and v = Λ * y, we can write
    # the kernel in the following two ways:
    #
    #    k = exp(-u₊ * v) / (exp(-v) + 1)
    #      = exp(-u₋ * -v) / (exp(v) + 1)
    #
    # We need to use the upper equation for v ≥ 0 and the lower one for
    # v < 0 to avoid overflowing both numerator and denominator

    enum = exp(-abs(v) * (v ≥ 0 ? u₊ : u₋))
    denom = 1 + exp(-abs(v))
    return enum / denom
end

function compute(kernel::RegularizedBoseKernel, u₊, u₋, v)
    # With "reduced variables" u, v we have:
    #
    #   K = -1/lambda * exp(-u_+ * v) * v / (exp(-v) - 1)
    #     = -1/lambda * exp(-u_- * -v) * (-v) / (exp(v) - 1)
    #
    # where we again need to use the upper equation for v ≥ 0 and the
    # lower one for v < 0 to avoid overflow.
    absv = abs(v)
    enum = exp(-absv * (v ≥ 0 ? u₊ : u₋))
    T = eltype(v)

    # The expression ``v / (exp(v) - 1)`` is tricky to evaluate: firstly,
    # it has a singularity at v=0, which can be cured by treating that case
    # separately.  Secondly, the denominator loses precision around 0 since
    # exp(v) = 1 + v + ..., which can be avoided using expm1(...)
    denom = absv ≥ 1e-200 ? absv / expm1(-absv) : -one(absv)
    return -1 / T(kernel.Λ) * enum * denom
end

"""
    segments_x(kernel)

Segments for piecewise polynomials on the ``x`` axis.

List of segments on the ``x`` axis for the associated piecewise
polynomial. Should reflect the approximate position of roots of a
high-order singular function in ``x``.
"""
function segments_x end

"""
    segments_y(kernel)

Segments for piecewise polynomials on the ``y`` axis.

List of segments on the ``y`` axis for the associated piecewise
polynomial. Should reflect the approximate position of roots of a
high-order singular function in ``y``.
"""
function segments_y end

function segments_x(hints::SVEHintsLogistic)
    nzeros = max(round(Int, 15 * log10(hints.kernel.Λ)), 1)
    diffs = 1 ./ cosh.(0.143 * range(0; length=nzeros))
    zeros = cumsum(diffs)
    zeros ./= last(zeros)
    return [-reverse(zeros); 0; zeros]
end

function segments_y(hints::SVEHintsLogistic)
    nzeros = max(round(Int, 20 * log10(hints.kernel.Λ)), 2)

    # Zeros around -1 and 1 are distributed asymptotically identically
    leading_diffs = [0.01523, 0.03314, 0.04848, 0.05987, 0.06703, 0.07028, 0.07030, 0.06791,
        0.06391, 0.05896, 0.05358, 0.04814, 0.04288, 0.03795, 0.03342, 0.02932, 0.02565,
        0.02239, 0.01951, 0.01699][1:min(nzeros, 20)]

    diffs = [leading_diffs; 0.25 ./ exp.(0.141 * (20:(nzeros - 1)))]
    zeros = cumsum(diffs)
    zeros ./= pop!(zeros)
    zeros .-= 1
    return [-1; zeros; 0; -reverse(zeros); 1]
end

function segments_x(hints::SVEHintsRegularizedBose)
    # Somewhat less accurate...
    nzeros = max(round(Int, 15 * log10(hints.kernel.Λ)), 15)
    diffs = 1 ./ cosh.(0.18 * range(0; length=nzeros))
    zeros = cumsum(diffs)
    zeros ./= last(zeros)
    return [-reverse(zeros); 0; zeros]
end

function segments_y(hints::SVEHintsRegularizedBose)
    nzeros = max(round(Int, 20 * log10(hints.kernel.Λ)), 20)
    i = range(0; length=nzeros)
    diffs = @. 0.12 / exp(0.0337 * i * log(i + 1))
    zeros = cumsum(diffs)
    zeros ./= pop!(zeros)
    zeros .-= 1
    return [-1; zeros; 0; -reverse(zeros); 1]
end

"""
    matrix_from_gauss(kernel, gauss_x, gauss_y)

Compute matrix for kernel from Gauss rules.
"""
function matrix_from_gauss(kernel, gauss_x, gauss_y)
    # (1 ± x) is problematic around x = -1 and x = 1, where the quadrature
    # nodes are clustered most tightly.  Thus we have the need for the
    # matrix method.
    return @inbounds kernel.(
        gauss_x.x, transpose(gauss_y.x), gauss_x.x_forward, gauss_x.x_backward
    )
end

"""
    check_domain(kernel, x, y)

Check that `(x, y)` lies within `kernel`'s domain and return it.
"""
function check_domain(kernel, x, y)
    xmin, xmax = xrange(kernel)
    xmin ≤ x ≤ xmax || throw(DomainError(x, "x value not in range [$xmin, $xmax]"))

    ymin, ymax = yrange(kernel)
    ymin ≤ y ≤ ymax || throw(DomainError(y, "y value not in range [$ymin, $ymax]"))

    return x, y
end

function compute_uv(Λ, x, y, x₊=1 + x, x₋=1 - x)
    u₊ = x₊ / 2
    u₋ = x₋ / 2
    v = Λ * y
    return u₊, u₋, v
end

"""
    get_symmetrized(kernel, sign)

Construct a symmetrized version of `kernel`, i.e. `kernel(x, y) + sign * kernel(x, -y)`.

!!! warning "Beware!"

    By default, this returns a simple wrapper over the current instance
    which naively performs the sum.  You may want to override this
    to avoid cancellation.
"""
get_symmetrized(kernel::AbstractKernel, sign) = ReducedKernel(kernel, sign)

function get_symmetrized(kernel::LogisticKernel, sign)
    sign == -1 && return LogisticKernelOdd(kernel, sign)
    return Base.invoke(get_symmetrized, Tuple{AbstractKernel,typeof(sign)}, kernel, sign)
end

function get_symmetrized(kernel::RegularizedBoseKernel, sign)
    sign == -1 && return RegularizedBoseKernelOdd(kernel, sign)
    return Base.invoke(get_symmetrized, Tuple{AbstractKernel,typeof(sign)}, kernel, sign)
end

get_symmetrized(::AbstractReducedKernel, sign) = error("cannot symmetrize twice")

function callreduced(kernel::AbstractReducedKernel, x, y, x₊, x₋)
    x, y = check_domain(kernel, x, y)

    # The reduced kernel is defined only over the interval [0, 1], which
    # means we must add one to get the x_plus for the inner kernels.  We
    # can compute this as 1 + x, since we are away from -1.
    x₊ = 1 + x₊

    K₊ = kernel.inner(x, +y, x₊, x₋)
    K₋ = kernel.inner(x, -y, x₊, x₋)
    return K₊ + kernel.sign * K₋
end

(kernel::ReducedKernel)(x, y, x₊, x₋) = callreduced(kernel, x, y, x₊, x₋)

"""
    is_centrosymmetric(kernel)

Return `true` if `kernel(x, y) == kernel(-x, -y)` for all values of `x` and `y`
in range. This allows the kernel to be block-diagonalized,
speeding up the singular value expansion by a factor of 4.  Defaults
to `false`.
"""
iscentrosymmetric(::AbstractKernel) = false
iscentrosymmetric(::LogisticKernel) = true
iscentrosymmetric(::RegularizedBoseKernel) = true
iscentrosymmetric(::AbstractReducedKernel) = false

"""
    (kernel::AbstractKernel)(x, y[, x₊, x₋])

Evaluate `kernel` at point `(x, y)`.

The parameters `x₊` and `x₋`, if given, shall contain the
values of `x - xₘᵢₙ` and `xₘₐₓ - x`, respectively.  This is useful
if either difference is to be formed and cancellation expected.
"""
function (kernel::AbstractKernel)(
    x, y, x₊=x - first(xrange(kernel)), x₋=last(xrange(kernel)) - x
)
    @boundscheck x, y = check_domain(kernel, x, y)
    u₊, u₋, v = compute_uv(kernel.Λ, x, y, x₊, x₋)
    return compute(kernel, u₊, u₋, v)
end

function (kernel::LogisticKernelOdd)(
    x, y, x₊=x - first(xrange(kernel)), x₋=last(xrange(kernel)) - x
)
    result = callreduced(kernel, x, y, x₊, x₋)

    # For x * y around 0, antisymmetrization introduces cancellation, which
    # reduces the relative precision. To combat this, we replace the
    # values with the explicit form
    v_half = kernel.inner.Λ / 2 * y
    xy_small = x * v_half < 1
    cosh_finite = v_half < 85
    return (xy_small && cosh_finite) ? -sinh(v_half * x) / cosh(v_half) : result
end

function (kernel::RegularizedBoseKernelOdd)(
    x, y, x₊=x - first(xrange(kernel)), x₋=last(xrange(kernel)) - x
)
    result = callreduced(kernel, x, y, x₊, x₋)

    # For x * y around 0, antisymmetrization introduces cancellation, which
    # reduces the relative precision.  To combat this, we replace the
    # values with the explicit form
    v_half = kernel.inner.Λ / 2 * y
    xv_half = x * v_half
    xy_small = xv_half < 1
    sinh_range = 1e-200 < v_half < 85
    return xy_small && sinh_range ? -y * sinh(xv_half) / sinh(v_half) : result
end

segments_x(hints::SVEHintsReduced) = symm_segments(segments_x(hints.inner_hints))
segments_y(hints::SVEHintsReduced) = symm_segments(segments_y(hints.inner_hints))

function symm_segments(x)
    x ≈ -reverse(x) || error("segments must be symmetric")
    xpos = x[(begin + length(x) ÷ 2):end]
    iszero(first(xpos)) || (xpos = [0; xpos])
    return xpos
end

"""
    sve_hints(kernel, ε)

Provide discretisation hints for the SVE routines.

Advises the SVE routines of discretisation parameters suitable in
tranforming the (infinite) SVE into an (finite) SVD problem.

See also [`AbstractSVEHints`](@ref).
"""
function sve_hints end

sve_hints(kernel::LogisticKernel, ε) = SVEHintsLogistic(kernel, ε)
sve_hints(kernel::RegularizedBoseKernel, ε) = SVEHintsRegularizedBose(kernel, ε)
sve_hints(kernel::AbstractReducedKernel, ε) = SVEHintsReduced(sve_hints(kernel.inner, ε))

"""
    nsvals(hints)

Upper bound for number of singular values.

Upper bound on the number of singular values above the given threshold, i.e. where
`s[l] ≥ ε * first(s)`.
"""
function nsvals end
function nsvals(hints::SVEHintsLogistic)
    log10_Λ = max(1, log10(hints.kernel.Λ))
    return round(Int, (25 + log10_Λ) * log10_Λ)
end
function nsvals(hints::SVEHintsRegularizedBose)
    log10_Λ = max(1, log10(hints.kernel.Λ))
    return round(Int, 28 * log10_Λ)
end
function nsvals(hints::SVEHintsReduced)
    return (nsvals(hints.inner_hints) + 1) ÷ 2
end

"""
    ngauss(hints)

Gauss-Legendre order to use to guarantee accuracy.
"""
ngauss(hints::SVEHintsLogistic) = hints.ε ≥ 1e-8 ? 10 : 16
ngauss(hints::SVEHintsRegularizedBose) = hints.ε ≥ 1e-8 ? 10 : 16
ngauss(hints::SVEHintsReduced) = ngauss(hints.inner_hints)

"""
    ypower(kernel)

Power with which the ``y`` coordinate scales.
"""
ypower(::AbstractKernel) = 0
ypower(::RegularizedBoseKernel) = 1
ypower(kernel::AbstractReducedKernel) = ypower(kernel.inner)

"""
    conv_radius(kernel)

Convergence radius of the Matsubara basis asymptotic model.

For improved relative numerical accuracy, the IR basis functions on the
Matsubara axis `uhat(basis, n)` can be evaluated from an asymptotic
expression for `abs(n) > conv_radius`.  If `isinf(conv_radius)`, then
the asymptotics are unused (the default).
"""
conv_radius(::AbstractKernel) = Inf
conv_radius(kernel::LogisticKernel) = 40 * kernel.Λ
conv_radius(kernel::RegularizedBoseKernel) = 40 * kernel.Λ
conv_radius(kernel::AbstractReducedKernel) = conv_radius(kernel.inner)

"""
    weight_func(kernel, statistics::Statistics)

Return the weight function for the given statistics.
"""
function weight_func(::AbstractKernel, ::Statistics)
    return x -> ones(eltype(x), size(x))
end
function weight_func(kernel::LogisticKernel, statistics::Statistics)
    if statistics == fermion
        return y -> ones(eltype(y), size(y))
    else
        return y -> 1 ./ tanh.(0.5 * kernel.Λ * y)
    end
end
function weight_func(::RegularizedBoseKernel, statistics::Statistics)
    statistics == boson || error("Kernel is designed for bosonic functions")
    return y -> 1 / y
end


using SparseInverseProblems
import SparseInverseProblems: lmo, phi, solveFiniteDimProblem, localDescent,
getStartingPoint, parameterBounds, computeGradient
using SparseInverseProblems.Util
const min_theta = 0.001
const max_theta = pi
const min_r = 1E-3
const max_r = 0.999
const max_x = 1.0
const max_b = 1.0
struct LinearSysID <: BoxConstrainedDifferentiableModel
  u :: Vector{Float64}  #u is the input time series
  grid
  normalization :: Function # (r,theta) -> n(r,theta)
  dndr :: Function # (r,theta) -> dn/dr
  dndt :: Function # (r,theta) -> dn/dt
  LinearSysID(u, gr,gt,n,dndr,dndt) = new(u, grid(gr,gt,u),n,dndr,dndt)
end

function getStartingPoint(model :: LinearSysID, v :: Vector{Float64})
  (gridR, gridT, l_x, l_b)  = model.grid

  scores = zeros(length(gridR), length(gridT))
  for i = 1:length(gridR)
    for j = 1:length(gridT)
      r,theta = gridR[i], gridT[j]
      nf = model.normalization(r,theta)
      scores[i,j] = nf*(max_x*sumabs(v'*l_x[i,j]) + max_b*sumabs(v'*l_b[i,j]))
    end
  end
  ind = indmax(scores)
  (rind,tind) = ind2sub((length(gridR), length(gridT)),ind)
  best_x = -max_x * sign(v'*l_x[ind])
  best_b = -max_b * sign(v'*l_b[ind])
  gridP = [gridR[rind]; gridT[tind]; vec(best_x); vec(best_b)]

  return gridP
end

parameterBounds(model :: LinearSysID) =
  ([min_r, min_theta,-max_x,-max_x, -max_b, -max_b],
  [max_r, max_theta,max_x,max_x, max_b, max_b])

function computeGradient(model :: LinearSysID, weights :: Vector{Float64},
  thetas :: Matrix{Float64}, v :: Vector{Float64})
  gradient = zeros(thetas)
  for i = 1:size(thetas,2)
    gradient[:,i] = weights[i]*jacobian(model, vec(thetas[:,i]))* v
  end
  return gradient
end

function phi(sim :: LinearSysID,parameters :: Matrix{Float64}, weights :: Vector{Float64})
  return runSim(parameters,sim.u,sim.normalization)*weights
end

function runSim(atoms,u,n)
  nAtoms = size(atoms,2)
  results = zeros(length(u),nAtoms)
  for i = 1:nAtoms
    w = vec(atoms[:,i])
    r,t,x,b = unpackParameters(w)
    l_x,l_b = getLinearOps(r,t,u)
    results[:,i] = (l_x*x + l_b*b)*n(r,t)
  end
  return results
end

function getLinearOps(r,theta,u; l_x :: Matrix{Float64} = zeros(length(u),2),l_b :: Matrix{Float64} = zeros(length(u),2) )
  n = length(u)
  rt = r.^(1:n)
  ttheta = (1:n)*theta
  ct = cos(ttheta)
  st = sin(ttheta)
  @inbounds for t = 1:n
    l_x[t,1] = rt[t]*ct[t]
    l_x[t,2] = -rt[t]*st[t]
     @inbounds for tau = 1:t-2 #u[0] = 0
      flipped_t = t-1-tau
      c = u[tau]*rt[flipped_t]
      l_b[t,1] += c*ct[flipped_t]
      l_b[t,2] -= c*st[flipped_t]
    end
  end
  @inbounds for t = 2:n
    l_b[t,1] += u[t-1] # flipped_t = 0
  end
  return (l_x, l_b)
end

function jacobian(s :: LinearSysID, w :: Vector{Float64}; j_r :: Vector{Float64} = zeros(length(s.u)), j_theta :: Vector{Float64} = zeros(length(s.u)), l_x :: Matrix{Float64} = zeros(length(s.u),2), l_b :: Matrix{Float64} = zeros(length(s.u),2))
  r,theta,x,b = unpackParameters(w)
  n = length(s.u)
  u :: Vector{Float64} = s.u
  rt = r.^(1:n)
  ttheta = (1:n)*theta
  ct = cos(ttheta)
  st = sin(ttheta)

  @inbounds for t = 1:n
    l_x[t,1] = rt[t]*ct[t]
    l_x[t,2] = -rt[t]*st[t]
    j_theta[t] -= t*rt[t]*st[t]*x[1]
    j_theta[t] -= t*rt[t]*ct[t]*x[2]
     for tau = 1:t-2 #u[0] = 0
      flipped_t = t-1-tau
      c = u[tau]*rt[flipped_t]
      l_b[t,1] += c*ct[flipped_t]
      l_b[t,2] -= c*st[flipped_t]
      j_theta[t] -= u[tau]*flipped_t*rt[flipped_t]*st[flipped_t]*b[1]
      j_theta[t] -= u[tau]*flipped_t*rt[flipped_t]*ct[flipped_t]*b[2]
    end
     for tau = 1:t-3 #u[0] = 0
      flipped_t = t-1-tau
      j_r[t] += u[tau]*flipped_t*rt[flipped_t-1]*ct[flipped_t]*b[1]
      j_r[t] -= u[tau]*flipped_t*rt[flipped_t-1]*st[flipped_t]*b[2]
    end
  end
  @inbounds for t = 2:n
    j_r[t] += t*rt[t-1]*ct[t]*x[1]
    j_r[t] -= t*rt[t-1]*st[t]*x[2]
    l_b[t,1] += u[t-1] # flipped_t = 0
  end
  #normalization!
  nf = s.normalization(r,theta)
  o = l_x*x + l_b*b
  j_r = j_r*nf + o*s.dndr(r,theta)
  j_theta = j_theta*nf + o*s.dndt(r,theta)
  return [j_r j_theta l_x*nf l_b*nf]'
end

function grid(nr,nt,u)
  gridR = max_r * sqrt(linspace(min_r,1.0,nr))
  gridT = linspace(min_theta,max_theta,nt)
  l_x = Array(Matrix{Float64},nr,nt)
  l_b = Array(Matrix{Float64},nr,nt)
  for i = 1:nr, j = 1:nt
    r = gridR[i]
    theta = gridT[j]
    (l_x[i,j],l_b[i,j]) = getLinearOps(r,theta,u)
  end
  return (gridR, gridT, l_x,l_b)
end

unpackParameters(p) = (p[1],p[2],p[3:4],p[5:6]) #r, t, x_0, B

function solveFiniteDimProblem(sim :: LinearSysID, lossFn :: LSLoss, thetas :: Matrix{Float64}, y :: Vector{Float64}, tau :: Float64)
  A = runSim(thetas,sim.u,sim.normalization)
  return lasso_kernel(A'*A,A'*y,tau)
end


using OmegaCore, Test, Distributions

function test()
  function f(ω)
    a = (1 ~ Normal(0, 1))(ω)
    b = (2 ~ Normal(0, 1))(ω)
    a + b
  end

  function g(ω)
    x = 0.0
    for i = 1:100
      x += (i <| f)(ω)
    end
    x
  end

  function h(ω)
    α = 1.0
    for i = 1:10
      α *= (i <| f)(ω)
    end
    α
  end

  i(ω) = (1 <| g)(ω) + (2 <| h)(ω)

  ω = defω()
  i(ω)
  @test length(keys(ω)) == 220
end

function normal_normals_test()
  μ = 1 ~ Normal(0, 1)
  x = 2 ~ Normal(μ, 1)
end

function blr()
  m = 1 ~ Normal(0, 1)
  c = 2 ~ Normal(0, 1)
  n = 100
  data = rand(n)
  f(x, m_, c_) = m_*x + c_
  obs(id, ω) = f(x, m_(ω), c_(ω)) + id ~ Normal(ω, 0, 0.1)
  obs(ω) = [f(x, m_(ω), c_(ω)) + i ~ Normal(ω, 0, 0.1) for i in data]
end


@test array_matches(binedges(DiscretizeUniformCount(2), [1.0,2.0,3.0,4.0,5.0,6.0]), [1.0, 3.5, 6.0], 1e-8)
@test array_matches(binedges(DiscretizeUniformCount(2), [6.0,2.0,30.0,4.0,1.0,1.0]), [1.0, 3.0, 30.0], 1e-8)
@test array_matches(binedges(DiscretizeUniformCount(2), [1.0,2.0,3.0,4.0,5.0]), [1.0, 3.5, 5.0], 1e-8)
@test array_matches(binedges(DiscretizeUniformCount(3), [1.0,2.0,3.0,4.0,5.0]), [1.0, 2.5, 4.5, 5.0], 1e-8)
@test_throws ErrorException binedges(DiscretizeUniformCount(3), [1.0, 1.0])

@test binedges(DiscretizeUniformCount(2), [1,2]) == [1, 2, 2]
@test binedges(DiscretizeUniformCount(3), [1,2,2,2,3]) == [1, 2, 3, 3]
@test binedges(DiscretizeUniformCount(3), [1,2,2,2,50]) == [1, 2, 26, 50]
@test binedges(DiscretizeUniformCount(3), [1,2,2,2,51]) == [1, 2, 27, 51]
@test binedges(DiscretizeUniformCount(3), [1,1,1,2,3,4,50,51,52]) == [1, 2, 27, 52]


import Base.isempty

export ConvexHull, CH,
       convex_hull,
       convex_hull!,
       ConvexHull!,
       swap

"""
    ConvexHull{N<:Real, S1<:LazySet{N}, S2<:LazySet{N}} <: LazySet{N}

Type that represents the convex hull of the union of two convex sets.

### Fields

- `X` -- convex set
- `Y` -- convex set

### Notes

The `EmptySet` is the neutral element for `ConvexHull`.

### Examples

Convex hull of two 100-dimensional Euclidean balls:

```jldoctest
julia> b1, b2 = Ball2(zeros(100), 0.1), Ball2(4*ones(100), 0.2);

julia> c = ConvexHull(b1, b2);

julia> typeof(c)
ConvexHull{Float64,Ball2{Float64,Array{Float64,1}},Ball2{Float64,Array{Float64,1}}}
```
"""
struct ConvexHull{N<:Real, S1<:LazySet{N}, S2<:LazySet{N}} <: LazySet{N}
    X::S1
    Y::S2

    # default constructor with dimension check
    function ConvexHull(X::S1, Y::S2) where {N<:Real, S1<:LazySet{N},
                                             S2<:LazySet{N}}
        @assert dim(X) == dim(Y) "sets in a convex hull must have the same " *
            "dimension"
        return new{N, S1, S2}(X, Y)
    end
end

isoperationtype(::Type{<:ConvexHull}) = true
isconvextype(::Type{<:ConvexHull}) = true

# EmptySet is the neutral element for ConvexHull
@neutral(ConvexHull, EmptySet)

# Universe is the absorbing element for ConvexHull
@absorbing(ConvexHull, Universe)

"""
    CH

Alias for `ConvexHull`.
"""
const CH = ConvexHull

"""
    swap(ch::ConvexHull)

Return a new `ConvexHull` object with the arguments swapped.

### Input

- `ch` -- convex hull of two convex sets

### Output

A new `ConvexHull` object with the arguments swapped.
"""
function swap(ch::ConvexHull)
    return ConvexHull(ch.Y, ch.X)
end

"""
    dim(ch::ConvexHull)

Return the dimension of a convex hull of two convex sets.

### Input

- `ch` -- convex hull of two convex sets

### Output

The ambient dimension of the convex hull of two convex sets.
"""
function dim(ch::ConvexHull)
    return dim(ch.X)
end

"""
    σ(d::AbstractVector{N}, ch::ConvexHull{N}) where {N<:Real}

Return the support vector of a convex hull of two convex sets in a given
direction.

### Input

- `d`  -- direction
- `ch` -- convex hull of two convex sets

### Output

The support vector of the convex hull in the given direction.
"""
function σ(d::AbstractVector{N}, ch::ConvexHull{N}) where {N<:Real}
    σ1 = σ(d, ch.X)
    σ2 = σ(d, ch.Y)
    ρ1 = dot(d, σ1)
    ρ2 = dot(d, σ2)
    return ρ1 >= ρ2 ? σ1 : σ2
end

"""
    ρ(d::AbstractVector{N}, ch::ConvexHull{N}) where {N<:Real}

Return the support function of a convex hull of two convex sets in a given
direction.

### Input

- `d`  -- direction
- `ch` -- convex hull of two convex sets

### Output

The support function of the convex hull in the given direction.
"""
function ρ(d::AbstractVector{N}, ch::ConvexHull{N}) where {N<:Real}
    return max(ρ(d, ch.X), ρ(d, ch.Y))
end

"""
    isbounded(ch::ConvexHull)

Determine whether a convex hull of two convex sets is bounded.

### Input

- `ch` -- convex hull of two convex sets

### Output

`true` iff both wrapped sets are bounded.
"""
function isbounded(ch::ConvexHull)
    return isbounded(ch.X) && isbounded(ch.Y)
end

"""
    isempty(ch::ConvexHull)

Return if a convex hull of two convex sets is empty or not.

### Input

- `ch` -- convex hull

### Output

`true` iff both wrapped sets are empty.
"""
function isempty(ch::ConvexHull)
    return isempty(ch.X) && isempty(ch.Y)
end


using Unums
using Base.Test

import Unums: Unum00, Ubound00, Ubound22, Ubound34

for U in (Ubound00, Ubound22, Ubound34, Unum00)
    @test U(0) + U(1) == U(1)
    @test U(0) - U(1) == U(-1)
    @test U(0) + U(2) == U(2)
    @test U(0) - U(2) == U(-2)
    @test U(1) + U(1) == U(2)
    @test U(1) - U(1) == U(0)
    @test U(1) + U(2) == U(3)
    @test U(1) - U(2) == U(-1)
    @test U(0) + U(0.5) == U(0.5)
    @test U(0) - U(0.5) == U(-0.5)
    @test U(1) + U(0.5) == U(1.5)
    @test U(1) - U(0.5) == U(0.5)

    @test U(1)*U(1) == U(1)
    @test U(1)*U(2) == U(2)
    @test U(1)/U(2) == U(0.5)
    @test U(1)/U(-2) == U(-0.5)
    @test sqrt(U(1)) == U(1)
    @test sqrt(U(2)) == U(sqrt(2))
    @test U(2)^4 == U(2^4)

    if U <: Unums.Ubound
        @test U(-1,1)^2 == U(0,1)
    end
end

for U in (Unum22, Ubound22)
    u = convert(U, 2.1)
    lo,hi,olo,ohi = limits(u)
    @test lo == 2
    @test hi == 2.125
    @test olo == true
    @test ohi == true
end

@test Unum22(2) in Unum22(2)
@test Unum22(2) in Ubound22(2,3)
@test !(Unum22(2) in Unum22(2.1))

@test isfinite(Unum22(2))
@test isfinite(Unum22(2.1))
@test !isfinite(Unum22(1e300))


export project_l1_Duchi!

"""
    project_l1_Duchi!(v::Union{Vector{TF},Vector{Complex{TF}}}, b::TF)

Projects ector onto L1 ball of specified radius.

w = ProjectOntoL1Ball(v, b) returns the vector w which is the solution
  to the following constrained minimization problem:

   min   ||w - v||_2
   s.t.  ||w||_1 <= b.

  That is, performs Euclidean projection of v to the 1-norm ball of radius
  b.

Author: John Duchi (jduchi@cs.berkeley.edu)
Translated (with some modification) to Julia 1.1 by Bas Peters
"""

function project_l1_Duchi!(v::Union{Vector{TF},Vector{Complex{TF}}}, b::TF) where {TF<:Real}
  b <= TF(0) && error("Radius of L1 ball is negative")
  norm(v, 1) <= b && return v

  lv = length(v)
  u  = similar(v)
  sv = Vector{TF}(undef, lv)

  #use RadixSort for Float32 (short keywords)
  copyto!(u, v)
  u .= abs.(u)
  u  = convert(Vector{TF},u)
  if TF==Float32
    sort!(u, rev=true, alg=RadixSort)
  else
    u = sort!(u, rev=true, alg=QuickSort)
  end

  # if TF==Float32
  #   u = sort!(abs.(u), rev=true, alg=RadixSort)
  # else
  #   u = sort!(abs.(u), rev=true, alg=QuickSort)
  # end

  cumsum!(sv, u)

  # Thresholding level
  temp  = TF(1.0):TF(1.0):TF(lv)
  rho   = max(1, min(lv, findlast(u .> ((sv.-b) ./ temp ) ) ))::Int
  theta = max.(TF(0) , (sv[rho] .- b) ./ rho)::TF

  # Projection as soft thresholding
  v .= sign.(v) .* max.(abs.(v) .- theta, TF(0))

  return v
end


using LabelledArrays

export rossler_rossler_bidir

@inline @inbounds function eom_rossler_rossler_bidir(u, p, t)
    ω₁, ω₂, c_xy, c_yx, a₁, a₂, a₃, b₁, b₂, b₃ = (p...,)
    x1, x2, x3, y1, y2, y3 = (u...,)

    dx1 = -ω₁*(x2 + x3) + c_yx*(y1 - x1)
    dx2 = ω₁*x1 + a₁*x2
    dx3 = a₂ + x3*(x1 - a₃)

    dy1 = -ω₂*(y2 + y3) + c_xy*(x1 - y1)
    dy2 = ω₂*y1 + b₁*y2
    dy3 = b₂ + y3*(y1 - b₃)

    return SVector{6}(dx1, dx2, dx3, dy1, dy2, dy3)
end

"""
    rossler_rossler_bidir(; u0 = rand(6),
        ω₁ = 1.015, ω₂ = 0.985,
        c_xy = 0.1, c_yx = 0.1,
        a₁ = 0.15, a₂ = 0.2, a₃ = 10,
        b₁ = 0.15, b₂ = 0.2, b₃ = 10)

Initialise a system of two bidirectionally coupled 3D Rössler systems.
This system has been modified from [1] to allow other parameterisations,
but default parameters are as in [1].

The ``X`` and ``Y`` subsystems are mostly synchronized for
`c_xy > 0.1` or `c_yx > 0.1`.

## Equations of motion

The dynamics is generated by the following vector field

```math
\\begin{aligned}
\\dot{x_1} &= -\\omega_1(x_2 + x_3) + c_{yx}(y_1 - x_1) \\\\
\\dot{x_2} &= \\omega_1 x_1 + a_1 x_2 \\\\
\\dot{x_3} &= a_2 + x_3 (x_1 - a_3) \\\\
\\dot{y_1} &= -\\omega_2 (y_2 + y_3) + c_{xy}(x_1 - y_1) \\\\
\\dot{y_2} &= \\omega_2 y_1 + b_1 y_2 \\\\
\\dot{y_3} &= b_2 + y_3 (y_1 - b_3)
\\end{aligned}
```

## References

1. Amigó, José M., and Yoshito Hirata. "Detecting directional couplings from
    multivariate flows by the joint distance distribution." Chaos: An
    Interdisciplinary Journal of Nonlinear Science 28.7 (2018): 075302.
"""
function rossler_rossler_bidir(; u0 = rand(6),
        ω₁ = 1.015, ω₂ = 0.985,
        c_xy = 0.1, c_yx = 0.1,
        a₁ = 0.15, a₂ = 0.2, a₃ = 10,
        b₁ = 0.15, b₂ = 0.2, b₃ = 10)

    p = @LArray [ω₁, ω₂, c_xy, c_yx, a₁, a₂, a₃, b₁, b₂, b₃] (:ω₁, :ω₂, :c_xy, :c_yx, :a₁, :a₂, :a₃, :b₁, :b₂, :b₃)
    ContinuousDynamicalSystem(eom_rossler_rossler_bidir, u0, p)
end

function rossler_rossler_bidir_trajectory(npts, sample_dt; n_transient = 2000, dt = 0.2,
        u0 = rand(6), ω₁ = 1.015, ω₂ = 0.985,
        c_xy = 0.2, c_yx = 0.2,
        a₁ = 0.15, a₂ = 0.2, a₃ = 10,
        b₁ = 0.15, b₂ = 0.2, b₃ = 10)

    s = rossler_rossler_bidir(u0 = u0,
        c_xy = c_xy, c_yx = c_yx,
        ω₁ = ω₁, ω₂ = ω₂,
        a₁ = a₁, a₂ = a₂, a₃ = a₃,
        b₁ = b₁, b₂ = b₂, b₃ = b₃)

    # the system is recorded at times t0:dt:T
    T = npts*dt*sample_dt
    o = trajectory(s, T, dt = dt, Ttr = n_transient*dt, alg = SimpleDiffEq.SimpleATsit5())[1:sample_dt:end-1, :] #alg = SimpleDiffEq.SimpleATsit5()
end

function good_rossler_rossler_bidir_trajectory(npts; sample_dt = 1,
        Da₁ = Uniform(0.12, 0.17),
        Da₂ = Uniform(0.18, 0.22),
        Da₃ = Uniform(9.0, 11.0),
        Db₁ = Uniform(0.10, 0.20),
        Db₂ = Uniform(0.18, 0.22),
        Db₃ = Uniform(9.0, 11.0),
        Dω₁ = Uniform(0.95, 0.999), #Uniform(0.97, 1.03)
        Dω₂ = Uniform(1.001, 1.05), #Uniform(0.97, 1.03)
        a₁ = nothing,
        a₂ = nothing,
        a₃ = nothing,
        b₁ = nothing,
        b₂ = nothing,
        b₃ = nothing,
        ω₁ = nothing,
        ω₂ = nothing,
        c_xy = 0.2,  c_yx = 0.2,
        Ttr = 5000, dt = 0.2,
        u0 = rand(6),
        n_maxtries = 300)

    n_tries = 0

    while n_tries <= n_maxtries
        ω₁ == nothing ? ω₁ = rand(Dω₁) : ω₁ = ω₁
        ω₂ == nothing ? ω₂ = rand(Dω₂) : ω₂ = ω₂
        a₁ == nothing ? a₁ = rand(Da₁) : a₁ = a₁
        a₂ == nothing ? a₂ = rand(Da₂) : a₂ = a₂
        a₃ == nothing ? a₃ = rand(Da₃) : a₃ = a₃
        b₁ == nothing ? b₁ = rand(Db₁) : b₁ = b₁
        b₂ == nothing ? b₂ = rand(Db₂) : b₂ = b₂
        b₃ == nothing ? b₃ = rand(Db₃) : b₃ = b₃

        pts = rossler_rossler_bidir_trajectory(npts, sample_dt, dt = dt, n_transient = Ttr,
            c_xy = c_xy,  c_yx = c_yx, ω₁ = ω₁, ω₂ = ω₂, a₁ = a₁, a₂ = a₂,a₃ = a₃,b₁ = b₁,  b₂ = b₂, b₃ = b₃)

        if all(Matrix(pts) .< 1e10) && length(unique(pts)) > npts/2
            return pts
        end
        println("no attractor found. trying with new initial condition and parameters")
        n_tries += 1
    end
end


using Test
using FastGaussQuadrature
using IntervalArithmetic
# using Revise
using ImplicitDomainQuadrature

IDQ = ImplicitDomainQuadrature

function allequal(v1,v2)
    return all(v1 .≈ v2)
end

function allapprox(v1, v2; tol = 1e-14)
    @assert length(v1) == length(v2)
    flags = [isapprox(v1[i], v2[i], atol = tol) for i = 1:length(v1)]
    return all(flags)
end

p = [+1.0 0.0 1.0]
w = [1/3,2/3,1/3]
@test_throws AssertionError IDQ.ReferenceQuadratureRule(p,w,-1.0,1.0)

p = [-1.0 0.0 1.0]
@test_throws AssertionError IDQ.ReferenceQuadratureRule(p,w,1.,1.)
@test_throws AssertionError IDQ.ReferenceQuadratureRule(p,w,-1.,1.)

float_type = typeof(1.0)
w = [0.5,1.0,0.5]
quad = IDQ.ReferenceQuadratureRule(p,w,-1.,1.)
@test typeof(quad) == IDQ.ReferenceQuadratureRule{float_type}
@test typeof(quad) <: IDQ.AbstractQuadratureRule{1}
@test allequal(quad.points,p)
@test allequal(quad.weights,w)
@test quad.lo ≈ -1.0
@test quad.hi ≈ +1.0

p = [1.0 0.0 1.0
     1.0 0.0 1.0]
w = [0.5,1.0,0.5]
@test_throws AssertionError IDQ.ReferenceQuadratureRule(p,w,-1.0,1.0)
p = [-1.0 0.0 1.0 2.0]
@test_throws AssertionError IDQ.ReferenceQuadratureRule(p,w,-1.0,1.0)
p = [-1.0 0.0 1.0]
quad = IDQ.ReferenceQuadratureRule(p,w,-1.0,1.0)
@test allequal(quad.points,p)
@test allequal(quad.weights,w)
@test quad.lo ≈ -1.0
@test quad.hi ≈ +1.0

p = [0.0,0.5,1.0]
w = [1/3,1/3,1/3]
quad = IDQ.ReferenceQuadratureRule(p,w,0.0,1.0)
@test allequal(quad.points,p')
@test allequal(quad.weights,w)
@test quad.lo ≈ 0.0
@test quad.hi ≈ 1.0

quad = IDQ.ReferenceQuadratureRule(4)
p,w = gausslegendre(4)
@test allequal(quad.points,p')
@test allequal(quad.weights,w)
@test quad.lo ≈ -1.0
@test quad.hi ≈ +1.0
@test typeof(quad) == IDQ.ReferenceQuadratureRule{float_type}
@test typeof(quad) <: IDQ.AbstractQuadratureRule{1}

function test_iteration(quad,points,weights)
    count = 1
    flag = true
    for (p,w) in quad
        flag = flag && allequal(p,points[count])
        flag = flag && w == weights[count]
        count += 1
    end
    return flag
end

@test test_iteration(quad,p,w)
qp,qw = quad[3]
@test allequal(qp,p[3])
@test qw == w[3]

@test IDQ.affine_map(0.0,-1.,1.,0.,1.) ≈ 0.5
@test IDQ.affine_map(0.75,0.,1.,5.,10.) ≈ 8.75
points = [-0.5 0.0 0.5]
@test allequal(IDQ.affine_map.(points,-1.,1.,0.,1.),[0.25 0.5 0.75])

@test IDQ.affine_map_derivative(1.0,-1.,1.,0.,1.) ≈ 0.5
@test IDQ.affine_map_derivative(1.0,-1.,1.,-5.,5.) ≈ 5.0
weights = [1.,2.,3.]
@test allequal(IDQ.affine_map_derivative.(weights,2.,3.,5.,10.),5weights)

points = [2. 3. 4.]
weights = [5.,7.,9.]
p,w = IDQ.transform(points,weights,1.,5.,-1.,1.)
@test allequal(p,[-0.5 0.0 0.5])
@test allequal(w,0.5weights)

quad = IDQ.ReferenceQuadratureRule(3)
p,w = IDQ.transform(quad,5.,10.)
@test allequal(p,5.0 .+ 2.5*(quad.points .+ 1))
@test allequal(w,2.5*quad.weights)

p,w = IDQ.transform(quad,3..4)
@test allequal(p,3 .+ 0.5*(quad.points .+ 1))
@test allequal(w,0.5*quad.weights)

p = zeros(4,3)
w = [0.5,1.0,0.5]
@test_throws AssertionError IDQ.QuadratureRule(p,w)
p = [-1. -1 +1 +1
     -1  +1 -1 +1]
w = [1.,1.,1.,1.,1.]
@test_throws AssertionError IDQ.QuadratureRule(p,w)

w = [1.,1.,1.,1.]
quad = IDQ.QuadratureRule(p,w)
@test allequal(quad.points,p)
@test allequal(quad.weights,w)

p = [-1. -1 +1 +1
     -1  +1 -1 +1
     -1  +1 -1 +1
     -1  +1 -1 +1]
w = [1.,1.,1.,1.]
@test_throws AssertionError IDQ.QuadratureRule(p,w)

p = [-1. -1 +1 +1
     -1  +1 -1 +1]
w = [1.,1.,1.,1.,1.]
@test_throws AssertionError IDQ.QuadratureRule(p,w)

w = [1.,1.,1.,1.]
quad = IDQ.QuadratureRule(p,w)
@test typeof(quad) == IDQ.QuadratureRule{2,float_type}
@test supertype(typeof(quad)) == IDQ.AbstractQuadratureRule{2}
@test allequal(quad.points,p)
@test allequal(quad.weights,w)

p1 = [1. 2.]
p2 = [3. 4.]
p = IDQ.tensor_product_points(p1,p2)
testp = [1. 1.  2.  2.
         3. 4.  3.  4.]
@test allequal(p,testp)

box = IntervalBox(-1..1,2)
quad = IDQ.ReferenceQuadratureRule(2)
tq = IDQ.tensor_product(quad,box)
p = quad.points
testp = [p[1]  p[1]  p[2]  p[2]
         p[1]  p[2]  p[1]  p[2]]
testw = [w[1]*w[1],w[1]*w[2],w[2]*w[1],w[2]*w[2]]
@test allequal(tq.points,testp)
@test allequal(tq.weights,testw)

box = IntervalBox(0..1,2..3)
tq = IDQ.tensor_product(quad,box)
p1 = 0.5*(quad.points .+ 1.)
p2 = 2. .+ 0.5*(quad.points .+ 1.)
w = 0.5*quad.weights
testp = [p1[1]  p1[1]  p1[2]  p1[2]
         p2[1]  p2[2]  p2[1]  p2[2]]
testw = kron(w,w)
@test allequal(tq.points,testp)
@test allequal(tq.weights,testw)

@test_throws AssertionError IDQ.tensor_product_quadrature(2,0)
@test_throws AssertionError IDQ.tensor_product_quadrature(2,-1)

tq = IDQ.tensor_product_quadrature(2,2)
quad = IDQ.ReferenceQuadratureRule(2)
p = quad.points
testp = [p[1]  p[1]  p[2]  p[2]
         p[1]  p[2]  p[1]  p[2]]
testw = kron(quad.weights,quad.weights)
@test allequal(tq.points,testp)
@test allequal(tq.weights,testw)

p = [1.0 2.0 3.0]
w = [1.0, 2.0]
@test_throws AssertionError IDQ.TemporaryQuadrature(p, w)

float_type = typeof(1.0)
p = [-1.0 1.0]
w = [0.5, 0.5]
quad = IDQ.TemporaryQuadrature(p, w)
@test typeof(quad) == IDQ.TemporaryQuadrature{float_type}
@test allapprox(quad.points, p)
@test allapprox(quad.weights, w)

@test_throws AssertionError IDQ.TemporaryQuadrature(float_type, 0)
@test_throws AssertionError IDQ.TemporaryQuadrature(float_type, 4)

quad = IDQ.TemporaryQuadrature(float_type, 2)
@test size(quad.points) == (2, 0)
@test length(quad.weights) == 0

p = [0.0, 1.0]
w = [1.0, 2.0]
@test_throws AssertionError IDQ.update!(quad, p, w)
p = [0.0]
w = [1.0]
@test_throws AssertionError IDQ.update!(quad, p, w)

p = [0.0, 1.0]
w = [1.0]
IDQ.update!(quad, p, w)
@test allapprox(quad.points, p)
@test allapprox(quad.weights, w)


### A Pluto.jl notebook ###
# v0.14.5

using Markdown
using InteractiveUtils

# ╔═╡ d941d6c2-55bf-11eb-0002-35c7474e4050
using NiLang, Test

# ╔═╡ 2061b434-0ad1-46eb-a0c7-1a5f432bfa62
begin
	twocol(left, right; llabel="forward", rlabel="backward") = HTML("

	$llabel
	$rlabel

$(html(left))

$(html(right))

")
	example(str) = HTML("""
$str
""")
	title1(str) = HTML("""

'_'))>$str

""")
	title2(str) = HTML("""
'_'))>$str
""")
	titleref(str) = HTML("""'_'))">$str""")
	using PlutoUI: TableOfContents
	using Pkg
	pkgversion(m::Module) = Pkg.TOML.parsefile(NiLang.project_relative_path("Project.toml"))["version"]
	hightlight(str) = HTML("$str")
end;

# ╔═╡ 8c2c4fa6-172f-4dde-a279-5d0aecfdbe46
module M
using NiLang

# define two functions
function new_forward(x)
	if x > 0
		return x * 2
	elseif x < 0
		return x / 2
	end
end

function new_backward(x)
	if x > 0
		return x / 2
	elseif x < 0
		return x * 2
	end
end

# declare them as reversible to each other
@dual new_forward new_backward

# The following is need only when your function is differentiable
using NiLang.AD: GVar
function new_backward(x::GVar)
	if x.x > 0
		GVar(new_backward(x.x), x.g * 2)
	elseif x.x < 0
		GVar(new_backward(x.x), x.g / 2)
	end
end

function new_forward(x::GVar)
	if x.x > 0
		GVar(new_forward(x.x), x.g / 2)
	elseif x.x < 0
		GVar(new_forward(x.x), x.g * 2)
	end
end
end

# ╔═╡ 3199a048-7b39-40f8-8183-6a54cccd91b6
using BenchmarkTools

# ╔═╡ 0e1ba158-a6bc-401c-9ba7-ed78020ad068
using Base.Threads

# ╔═╡ a4e76427-f051-4b29-915a-fdfce3a299bb
html"""

  <svg class="octicon octicon-mark-github v-align-middle" height="32" viewBox="0 0 16 16" version="1.1" width="32" aria-hidden="true"><path fill-rule="evenodd" d="M8 0C3.58 0 0 3.58 0 8c0 3.54 2.29 6.53 5.47 7.59.4.07.55-.17.55-.38 0-.19-.01-.82-.01-1.49-2.01.37-2.53-.49-2.69-.94-.09-.23-.48-.94-.82-1.13-.28-.15-.68-.52-.01-.53.63-.01 1.08.58 1.23.82.72 1.21 1.87.87 2.33.66.07-.52.28-.87.51-1.07-1.78-.2-3.64-.89-3.64-3.95 0-.87.31-1.59.82-2.15-.08-.2-.36-1.02.08-2.12 0 0 .67-.21 2.2.82.64-.18 1.32-.27 2-.27.68 0 1.36.09 2 .27 1.53-1.04 2.2-.82 2.2-.82.44 1.1.16 1.92.08 2.12.51.56.82 1.27.82 2.15 0 3.07-1.87 3.75-3.65 3.95.29.25.54.73.54 1.48 0 1.07-.01 1.93-.01 2.2 0 .21.15.46.55.38A8.013 8.013 0 0016 8c0-4.42-3.58-8-8-8z"></path></svg>

Download this notebook

"""

# ╔═╡ c2c7b4d4-f8c9-4ebf-8da2-0103f03136e7
md"# NiLang's (v$(pkgversion(NiLang))) Documentation

NiLang is a embeded domain specific language (eDSL) in Julia, so one need to install [Julia](https://julialang.org/) first. Before reading this documentation, you need to know basic Julia grammar, and how to install and use packages.
Also, it might be good to read the [README](https://github.com/GiggleLiu/NiLang.jl) first. In this tutorial, we focus on

* NiLang grammar and design patterns,
* Automatic differentiation based on reversible programming.
"

# ╔═╡ 12f07cc7-979c-43c3-9dc9-36ea1463c1f6
md"The symbols used in this notebook"

# ╔═╡ 611b577f-4722-42bf-8f8e-aeb2fb30be71
md"""
|  symbol  |  meaning  | how to type |
| -------  | --------- | ----------- |
| ← | allocate |  \leftarrow + TAB |
| → | deallocate | \rightarrow + TAB |
| ↔ | exchange | \leftrightarrow + TAB |
| ∅ | empty variable | \emptyset + TAB |
| ~ | inverse | ~ |
"""

# ╔═╡ 605872cf-f3fd-462e-a2b1-7d1c5ae45efd
title1("Getting started")

# ╔═╡ fb3dee44-5fa9-4773-8b7f-a83c44358545
md"
After installing NiLang in a Julia REPL by typing `]add NiLang`, one can use NiLang and use the macro `@i` to define a reversible function ``f_1: (x, y) → (x+5y, y)``."

# ╔═╡ 70088425-6779-4a2d-ba6d-b0a34c8e93a6
@i @inline function f1(x, y; constant)
	x += y * constant
end

# ╔═╡ af738f89-3214-429c-9c7d-18a6ea0d9401
f1(1.0, 2.0; constant=5.0)  # call

# ╔═╡ 48d7ebc1-5def-4a57-9ec1-3fc370a4543f
 (~f1)(11.0, 2.0; constant = 5.0)  # uncall

# ╔═╡ f0e94247-f615-472b-8218-3fa287b38aa1
md"A NiLang function defines a ``\mathbb{R}^n\rightarrow\mathbb{R}^n`` (notice inputs and outputs have the same shapes) mapping, it can take both keyword and positional arguments, where positional arguments can only be used as constants. There is no `return` statement because this function returns input non-keyword variables automatically, it is forbiden to write `return` statement inside a NiLang function.
One can aslo put macros like `@inline` after NiLang's `@i` macro. NiLang's macro will render the body of the function first and pass it to other macros.
"

# ╔═╡ 2581aa33-1dc5-40b1-aa9f-6a11cc750c93
md"`x += y * constant` is an instruction that defines a bijective mapping ``\mathbb{R}^3\rightarrow\mathbb{R}^3``. All NiLang instructions change the variable inplace. Here, we accumulate the result to `x` rather than using `y *= constant` to modify the variable directly. This is because in a regular number system, one can easily use the zero element as the eraser to erase all information, which will cause irreversibility."

# ╔═╡ 60575978-081a-4bca-a3ed-2b51cd6abc92
md"One can differentiate a NiLang function with the `NiLang.gradient(f, args; iloss, kwargs...)` where `args` and `kwargs` are positional and keyword arguments for `f`, and `iloss` is the index of the loss variable."

# ╔═╡ f98305cb-4ba2-404a-a5c3-65510e059504
NiLang.AD.gradient(f1, (1.0, 2.0); iloss=1, constant=5.0)

# ╔═╡ e8cd6667-597f-458b-8465-1822e09a7891
md"Here, we specify the first variable as the one that stores the loss. We get
```math
\begin{cases}
\frac{\partial x+5y}{\partial x}=1\\
\frac{\partial x+5y}{\partial y}=5\\
\end{cases}
```
"

# ╔═╡ 20145d75-004a-4c2f-b7ff-c400ca846d42
let
	content = md"""The above macro generates two functions, one is `f` and another is `~f` (or `Inv(f)`). The `x += y * constant` is translated to function `(x, y, constant) = PlusEq(*)(x, y, constant)`, where the function `PlusEq(*)` is bijective.

```julia
julia> using MacroTools

julia> MacroTools.prettify(@macroexpand @i function f(x, y; constant)
           x += y * constant
       end)
quote
    $(Expr(:meta, :doc))
    function $(Expr(:where, :(f(x, y; constant))))
        hare = wrap_tuple(((PlusEq)(*))(x, y, constant))
        x = hare[1]
        y = hare[2]
        constant = hare[3]
        (x, y)
    end
    if (NiLangCore)._typeof(f) != _typeof(~f)
        function $(Expr(:where, :((newt::_typeof(~f))(x, y; constant))))
            boar = wrap_tuple(((MinusEq)(*))(x, y, constant))
            x = boar[1]
            y = boar[2]
            constant = boar[3]
            (x, y)
        end
    end
end
```
"""
	HTML("How does the compiler work?$(html(content))")
end

# ╔═╡ c682a17f-600f-4034-bfe3-a851ab645c10
title1("Instructions and operands")

# ╔═╡ 5239dfe2-ea6d-4e07-a1b1-90954fe8ddc9
md"""
The basic form of an instruction is `y ⊙= f(args...)`, where `⊙` can be `+`, `-`, `*`, `/` or `⊻`, where `*=` and `/=` are only reversible in the logarithmic number systems. See section $(titleref("Integers, floating-point numbers, fixed-point numbers and logarithmic numbers")) for details.

A function/instruction can be used in NiLang only if its reverse is defined. A NiLang function is differentiable if the function body is composed of differentiable instructions, differentiable NiLang functions and NiLang control flows. A list of differentiable NiLang instructions are

| instruction | output   |
| ----------- | ---------- |
| ``{\rm FLIP}(y)`` | ``\sim y`` |
| ``{\rm NEG}(y)`` | ``-y`` |
| ``{\rm INC}(y)`` | ``y+1`` |
| ``{\rm DEC}(y)`` | ``y-1`` |
| ``{\rm INV}(y)`` | ``y ^ {-1}`` |
| ``{\rm HADAMARD}(x, y)`` | ``\frac{1}{\sqrt{2}}(x+y), \frac{1}{\sqrt{2}}(x-y)``
| ``{\rm SWAP}(a, b)`` | ``b, a`` |
| ``{\rm ROT}(a, b, \theta)`` | ``a \cos\theta - b\sin\theta, b \cos\theta + a\sin\theta, \theta`` |
| ``{\rm IROT}(a, b, \theta)`` | ``a \cos\theta + b\sin\theta, b \cos\theta - a\sin\theta, \theta`` |
| ``y \mathrel{\{+,-\}}= f_{+-}(args...)`` | ``y\{+, -\}f_{+-}(args...), args...`` |
| ``y \mathrel{\{*, /\}}= f_{*/}(args...)`` | ``y\{*, /\}f_{*/}(args...), args...`` |

Functions ``f_{+-} ∈ \rm \{identity, +, -, *, /, ^\wedge, abs, abs2, sqrt, exp, log, sin, sinh, asin, cos, cosh,`` ``\rm acos, tan, tanh, atan, sincos, convert\}`` and ``f_{*/}∈\rm \{identity, +, -, *, /, ^\wedge, convert\}.``
Functions `FLIP`, `NEG`, `INV`, `HADAMARD`, `SWAP` and `y ⊻= f_{⊻}(args...)` are self-reversible (or reflexive). {`ROT`, `IROT`} and {`INC`, `DEC`}, {`y += f_{+-}(args...)`, `y -= f_{+-}(args...)`} and {`y *= f_{*/}(args...)`, `y /= f_{*/}(args...)`} are pair-wise reversible.

For Jacobians and Hessians defined on these instructions, please check this [blog post](https://giggleliu.github.io/2020/01/18/jacobians.html).
This set of instructions is extensible, see section $(titleref("Extending the instruction set")) for an example.
"""

# ╔═╡ e4d86a5a-e820-4a70-8a87-08bac416291b
md"The operands of an instruction can be a composition of the following expressions"

# ╔═╡ c307b6a4-906d-4be7-9fd7-57c942aded51
md"""

| expression | meaning |
| ---------- | ------- |
| x | change a variable |
| x.field | change a field |
| x.:1    | change a tuple element |
| x'      | change the adjoint |
| -x      | change the inverse |
| x[i]    | change an array/dict element |
| (x, y, z)  | change multiple variables |
| @fields x | change the fields of an object |
| A{Float64}(x, y) | wrap, update fields and unwrap |
| x \|> subarray(:,i) | change the view of an array |
| x \|> f | change a field map or a bijective mapping |
| @const 3 | target is a constant that can not be changed |
| @skip! f(x) | an expression that can not be assigned back |
"""

# ╔═╡ 47b502d4-e8af-4d58-9067-9700784ea435
md"In the following example, one modifies the negated real part of a complex number in a vector inside a tuple directly."

# ╔═╡ faecc0a7-55d7-42a1-8e9f-7e30143eef9c
@i function dataview_func1(x::Vector, y::Tuple, θ)
	(-x[2].re, y.:1) += sincos(θ)
end

# ╔═╡ 4cbe69d6-b68f-4bda-a0dd-209f9ee54f18
dataview_func1([1+2.0im, 3+4im], (1.0,2.0), π/6)

# ╔═╡ 7dc82f28-77bf-40da-b520-800ed1bc80c9
md"""One can check the real part of the second element of `x` is decreased by `sin(π/6)`, while the first element of the tuple is increased by `cos(π/6)`. A NiLang instruction is different to a regular Julia instruction in that $(hightlight("A NiLang instruction changes variables inplace")), eventhough `ComplexF64`, `-x` and `tuple` are all considered as immutable in Julia."""

# ╔═╡ 9f5f9de3-9558-4c18-9d98-b77d19b570ec
example("Complex valued log")

# ╔═╡ 6dfcfa19-f78f-4dac-89f7-d3c5dbe17987
@i function complex_log(y!::Complex{T}, x::Complex{T}) where T
    n ← zero(T)
    n += abs(x)

    y!.re += log(n)
    y!.im += angle(x)

    n -= abs(x)
    n → zero(T)
end; @test complex_log(0.0im, 3.0+2.0im)[1] ≈ log(3.0+2.0im)

# ╔═╡ dad1c6c0-d61b-4f9f-a71e-e683fe143aaa
title2("Broadcasting")

# ╔═╡ 3e4a3916-8fe4-4262-bcd3-3014822717a3
md"The interfaces is similar to julia's native broadcast, but expanded to native NiLang loops."

# ╔═╡ 34906208-e6f1-4a67-860a-a7b056a86dde
@i function complex_log_broadcast!(ys!, xs)
	complex_log.(ys!, xs)
end; @test (x=[1.0+2im, 3.0+4im, 4.0+5im]; complex_log_broadcast!(zeros(ComplexF64, 3), x)[1] ≈ log.(x))

# ╔═╡ 4601df35-679f-465d-9191-c18748b2fd83
title2("Avoid shared read-write")

# ╔═╡ 4fc72b9d-19a2-40f1-a4a8-5e97d3d5e529
md"Shared read-write is not allowed"

# ╔═╡ a0fde16f-8454-4f5c-a29c-a9e415c0c311
# `y -= y` effectively clears the content in `y`, this is why shared read-write is so dangerous.
@test_throws LoadError macroexpand(NiLang, :(@i function shared_readwrite_error(y)
	y -= y
end))

# ╔═╡ 633ff8f3-8d93-4f73-bec2-c42070e6ece9
md"NiLang is more restrictive, it forbids shared read too. This is in purpose, shared read will become shared write to the gradient field in the back-propagation of gradients - the main goal of NiLang."

# ╔═╡ 57f2d890-b5a0-47b7-9e3d-af4d03b10605
# `shared read is also forbidden`
@test_throws LoadError macroexpand(NiLang, :(@i function shared_read_error(x, y)
	x -= y * y  # should be written as `x -= y^2`
end))

# ╔═╡ 34063cd0-171e-46ce-80dd-52a341fa50a1
md"The correct way of avoiding shared read is renaming one of the variable."

# ╔═╡ 5d5d01db-8ff9-434c-8771-1fec6393e1fb
@i function avoid_shared_read(x, y)
	tmp ← zero(y)
	tmp += y
	x -= y * tmp
	tmp -= y
	tmp → zero(y)
end

# ╔═╡ 10d85a50-f2f9-403e-8f6c-baef61cf702a
avoid_shared_read(0.0, 3.0)

# ╔═╡ b52648bf-a28a-48af-8912-31729d943ce0
md"Shared read-write issue is more tricky when one uses NiLang to write kernel function in a parallel program (multi-threading, MPI and CUDA). See $(titleref(\"Multi-threading and CUDA\")) for details."

# ╔═╡ f45db10f-a836-40f3-9d8d-054ea6540e87
title2("Protect constant variables")

# ╔═╡ 1903563e-ccc2-44d9-9dbe-e5dede275b3c
md"""
To achieve the goal of "everything is mutable", NiLang assigns the output value back to the input variable after a call. Sometime it can cause issues for special variables like function, type parameters and constants (including the results generated from a function call). One can use `@const` (assert a variable is a constant) or `@skip!` (skip assigning back) to avoid such complications.
"""

# ╔═╡ 583f2585-15a3-47c6-a70e-e2f002754028
md"When using a function (e.g. `exp` as shown bellow) as an variable, one should be careful about the scope issue."

# ╔═╡ 60e6ff80-3593-4ae4-a273-914847f692db
@i function func_arg(y, f, x)
	y += f(x)
end

# ╔═╡ 9e5cfd68-b58d-4d83-aae2-447e5f805c97
@i function use_func_arg(y, x)
	func_arg(y, exp, x)
end

# ╔═╡ dc85a942-cf52-4405-ad03-32a768e1b6e7
@test_throws UndefVarError use_func_arg(0.0, 3.0)

# ╔═╡ 5cdd346b-10a5-485c-ba78-4c0b3cb0e02f
md"We see an error, but why calling `f2g` causes an error? If one check the generated code with `macroexpand`, one will see the `exp` is assigned in the local scope. The compiler takes it as a local variable and compaints that `exp` is not defined."

# ╔═╡ af9287b7-6131-46f6-beb8-6885e55e1975
macroexpand(NiLang, :(@i function use_func_arg(y, x)
	f2g(y, exp, x)
end)) |> NiLangCore.rmlines

# ╔═╡ e20eeabf-1c80-431e-8cfc-4d1b79c52b5a
@i function avoid_assignback(y, x)
	func_arg(y, (@const exp), x)
end; @test avoid_assignback(0.0, 3.0)[1] == exp(3)

# ╔═╡ 90d30eea-53de-48a0-9700-ff35681fdf38
md"Type parameter can not be assigned back too."

# ╔═╡ 390f58a5-6f5f-4d3a-bb16-ba04e43a07e7
@test_throws ErrorException Core.eval(NiLang, :(@i function type_arg(t::Type{T}, x) where T
	x += one(T)
end))

# ╔═╡ 2b57443e-a516-434b-be86-80616a98e2f5
@i function avoid_type_assignback(t::Type{T}, x) where T
	x += one(@skip! T)
end; @test avoid_type_assignback(Float64, 0.0)[2] == 1.0

# ╔═╡ fc2e27f9-b7ba-44cc-a953-6745548ad733
md"A function call that returning a constant should also be decorated with the `@const` or `@skip!` macro."

# ╔═╡ fc744931-360b-4478-9f77-c50f048de243
@test_throws LoadError macroexpand(NiLang, :(@i function funccall_arg(y, x::Matrix) where T
	y += size(x, 1) * size(x, 2)
end))

# ╔═╡ 9a152b36-f377-44da-9700-ca9e05e365ff
@i function avoid_funccall_assignback(y, x::Matrix) where T
	y += (@const size(x, 1)) * (@const size(x, 2))
end; @test avoid_funccall_assignback(0, randn(3,4))[1] == 12

# ╔═╡ b6dcd18c-606f-4340-b2ec-163e8bad03f5
title1("Variable manipulation")

# ╔═╡ a1a29f34-f8a9-4e9f-9afe-7d0096771440
title2("Allocate and deallocate a variable")

# ╔═╡ 90bd6ad4-3dd8-4e7c-b445-aed1d248a2ec
md"""
One can allocate a new variable `x` like `x ← constant` and deallocate a variable with known value with `x → known value`. They are reversible to each other with the following relation.

$(
twocol(md"
```julia
x ← constant
```", md"
```julia
x → constant
```
")
)

For example
"""

# ╔═╡ c0259e48-1973-486c-a828-1fcd3e4331c6
@i function alloc_func1()
	tmp ← 1
	# some code that uses `tmp` for computing and restores it to `1`
	tmp → 1
end

# ╔═╡ 8bbffa31-04a6-49ca-b36f-4d4140d75992
md"allocate multiple variable of the same type at one time"

# ╔═╡ a6f18c34-80ee-4b52-9ff8-f3c1b1d80f90
@i function power12(y, x::T) where T
	@zeros T a b c  # three variable of type `T`
	a += x^2
	b += a^2
	c += b*a
	y += c^2
	c -= b*a
	b -= a^2
	a -= x^2
	@safe @show a b c x y
	~@zeros T a b c
end

# ╔═╡ a694132b-4f52-467f-8bc4-dc32fe2812db
@test power12(0, 2)[1] == 4096

# ╔═╡ 8c2c82f2-1240-4f2f-830e-ee8021c1a41a
md"One can copy and push a value into a stack and use it later. It inverse operation will pop out a variable and assert its value."

# ╔═╡ 6203cf10-f8cc-4fb9-b814-7552b68c01dc
twocol(md"
```julia
stack[end+1] ← variable
```", md"
```julia
stack[end] → variable
```
")

# ╔═╡ f97a6bab-b9f9-4b95-98a9-381c51397526
@i function stack_push_and_pop!(stack, x, y, z)
	z += y
	stack[end+1] ← x  # copy a variable into a stack
	stack[end+1] ← y
	stack[end] → z  # pop a variable from a stack, `z` must have the same value as the variable.
end

# ╔═╡ 2a2970f4-ab01-486b-89a2-6ff96f734018
md"A less recommended approach is using the global stacks in NiLang, since NiLang is an eDSL, it can not guarantee the access order. Available global stacks are `FLOAT64_STACK`, `COMPLEXF64_STACK`, `INT64_STACK` and their 32 bit counter parts, as well as a `BOOL_STACK`."

# ╔═╡ 0b80d9be-53d7-4bf3-a558-659607af4709
@i function stack_push_and_pop!(x, y)
	GLOBAL_STACK[end+1] ← x  # copy a variable into a stack
	FLOAT64_STACK[end+1] ← y
end

# ╔═╡ 4ca48a2e-43da-457a-8e9f-6476097e4d7b
let
	stack = FastStack{Float64}(1000) # a preallocated stack of size 1000
	stack_push_and_pop!(stack, 5.0, 1.0, 0.0)  # you will get a stack of size 1
	@test length(stack) == 1
end

# ╔═╡ 92362fda-bae2-4e35-bfe4-dcaea853d50b
let
	NiLang.empty_global_stacks!()  # empty stacks
	stack_push_and_pop!(5.0, 1.0)
	@test length(GLOBAL_STACK) == 1
	@test length(FLOAT64_STACK) == 1
end

# ╔═╡ db9e7940-39f1-4ccf-ac70-146a521daa6e
md"one can also allocate and deallocate on dicts"

# ╔═╡ 93936612-1447-4114-b864-aba43adef4bd
md"""
The forward pass of dictionary allocation adds an entry to the dictionary (cause key error if the key already exists), the backward pass. The backward pass checks the variable in the dictionary is consistent with the asserted variable, and delete the key.

$(
twocol(md"
```julia
dict[key] ← variable
```", md"
```julia
dict[key] → variable
```
")
)
"""

# ╔═╡ b2add70c-c5d6-4e0f-a153-43e21a197181
@i function dict_assign(dict::Dict)
	var1 ← 3.14

	# copy a new variable to a dict
	dict["data"] ← var1

	# deallocate the original variables
	var1 → dict["data"]
end

# ╔═╡ 5c03d5a5-99f0-4efd-9a32-ce6d7c2b266c
dict_assign(Dict())

# ╔═╡ 2349e3ea-3053-42a4-b9d9-f97a76e4abd7
title2("Exchange two variables")

# ╔═╡ 269f18ee-3cd8-466a-a522-7c624503e31b
let
	expr = twocol(md"
```julia
var1 ↔ var2
```", md"
```julia
var1 ↔ var2
```
")
md"""One can exchange two variables is using `↔`.
$expr
"""
end

# ╔═╡ 89139719-c478-4066-9452-f9893f36d561
@i function exchange_func1(x, y)
	x ↔ y
end

# ╔═╡ d620c5ee-7d9c-4d3f-9e87-0c828dfab9ca
exchange_func1(3, 5)

# ╔═╡ 255d01b9-a873-4e63-9298-9d8f073348b0
md"One can also make a \"link\" by exchanging a variable with an empty variable such as `var::∅` and `stack[end+1]`. The forward pass push a variable to the stack and deallocate variable. The backward pass pops a variable and asserts its value."

# ╔═╡ f6cf1729-766c-4ed7-b004-c8c8ec6c7e07
let
	expr = twocol(md"
```julia
stack[end+1] ↔ var2
var1::∅ ↔ var2
```", md"
```julia
stack[end] ↔ var2::∅
var1 ↔ var2::∅
```
")
end

# ╔═╡ 3645d672-423f-4ac8-805f-0452793fee5a
@i function exchange_func2(x, y)
	anc ← 0.0
	anc += x * y
	anc ↔ z::∅  # declare `z` as an empty variable
	# after exchange, `anc` is empty and deallocated automatically.
	z -= x * y
	z → 0.0
end

# ╔═╡ c2a0024e-11dd-4ef7-8346-4374d98cafc0
exchange_func2(3, 4)

# ╔═╡ b20004e9-3c73-4dfb-8fd5-f377786fd53b
md"When exchanging with a stack top + 1, it means push and deallocate."

# ╔═╡ 5c1952b1-5016-4c87-b23c-8e6a235bf8cd
@i function stack_exchange(stack, y, x)
	stack[end+1] ↔ y  # push a variable into a stack and deallocate `y`
	y ← 1.2  # since `y` is deallocated, you can assign any value to it
	stack[end] ↔ anc::∅  # pop a variable to `anc`
	anc ↔ x    # exchange `anc` and x
	stack[end+1] ↔ anc  # push `anc` back to stack
end

# ╔═╡ 8e4470ee-01da-4547-b091-c4f65cd729b0
let
	stack = FastStack{Float64}(1000) # a preallocated stack
	stack_exchange(stack, 2.0, 3.0)  # you will get a stack of size 2
	@test length(stack) == 1 && stack.data[1] == 3.0
end

# ╔═╡ 0863bd06-cc70-4dde-b3b2-0a466805a356
md"""$(title2("Integers, floating-point numbers, fixed-point numbers and logarithmic numbers"))

A fixed-point zero with 43 fraction bits can be declared as `x ← Fixed43(0.0)` or `x ← zero(Fixed43)`, a logarithmic one can be declared as `x ← ULogarithmic(1.0)`, `x ← one(ULogarithmic{Float64})` or `x ← ULogarithmic(Fixed43(1.0))`. A complex number zero can be defined as `x ← Complex(0.0, 0.0)`.

| Number Type | +=  | *= | ⊻= | Source |
| ---- | ---- | ---- | ---- | ---- |
| boolean | - |  - | ✓ | JuliaLang |
| integer | ✓ |  × | ✓ | JuliaLang |
| floating-point number | ✓ (rounding error) |  × | - | JuliaLang |
| fixed-point number | ✓ |  × | - | [FixedPointNumbers.jl](https://github.com/JuliaMath/FixedPointNumbers.jl)
| logarithmic number | × |  ✓ | - |  [LogarithmicNumbers.jl](https://github.com/cjdoris/LogarithmicNumbers.jl)

* `✓`: an operation has its reverse when operating on a number type.
* `×`: an operation does not have a reverse when operating on a number type.
* `-`: an operation does not apply on a number type.

The `+=` operation is not regoriously reversible on floating point numbers, but we ignore the rounding errors in NiLang and use the reversibility check to detect the potential too-large rounding errors. Whether logarithmic number has rounding errors depends on its content type. If it uses floating point numbers as storage, then yes, otherwise if it uses fixed point number as the content type, then no.

One can use the `y ⊙= convert(x)` statement to convert `x` to the target type `typeof(y)` and accumulate it to `y`. Here `⊙=` can be one of `+=`, `*=` and `⊻=` that has its reverse on type `typeof(y)`.
"""

# ╔═╡ a0bae195-04e1-4642-9e14-fe4691e0906b
md"Fixed point numbers and Floating point numbers are reversible under the `+=` operation"

# ╔═╡ 20d6e8a0-2cf5-48ad-9549-60506b42b970
@i function fixed_pluseq(y1::T, x::T) where T<:Union{Fixed43, AbstractFloat}
	y1 += x
end; @test fixed_pluseq(Fixed43(0.5), Fixed43(0.6))[1] === Fixed43(1.1) && fixed_pluseq(0.5, 0.6)[1] === 1.1

# ╔═╡ 7dee5748-ed73-4e13-aa80-7a50efbc8449
@i function fixed_muleq(y1::T, x::T) where T<:Union{Fixed43, AbstractFloat}
	y1 *= x
end; @test_throws MethodError fixed_muleq(Fixed43(1.0), Fixed43(2.0))

# ╔═╡ 4c719e9b-641e-404e-9ab7-59e89135f3ba
md"Logarithmic numbers are reversible under the `*=` operation"

# ╔═╡ 77947e00-42c3-4c9e-b62a-b4b29489db43
@i function ulog_pluseq(y1::ULogarithmic{T}, x::ULogarithmic{T}) where T
	y1 += x
end; @test_throws MethodError ulog_pluseq(ULogarithmic(1.0), ULogarithmic(3.0))

# ╔═╡ 2614127d-34fb-4c3d-b678-42693f3c9341
@i function ulog_muleq(y1::ULogarithmic{T}, x::ULogarithmic{T}) where T
	y1 *= x
end; @test ulog_muleq(ULogarithmic(1.0), ULogarithmic(3.0))[1] === ULogarithmic(3.0)

# ╔═╡ 8f169235-3bd1-4cc4-a083-79736d306ad5
example("computing x ^ 3 with logarithmic numbers")

# ╔═╡ dfc9d305-5bce-4555-bfa3-d8d61fe4ca09
@i function power3(y::ULogarithmic{T}, x::T) where T
	for i=1:3
		y *= convert(x)
	end
end; @test power3(ULogarithmic(1.0), 3.0)[1] ≈ 27

# ╔═╡ f6bfa015-c101-45e8-995c-2bb6a3b7dc7d
title2("Types and arrays")

# ╔═╡ 8651d7ec-6bcd-4dbe-a062-c4bde32e5e91
md"
A Julia type can be accessed in NiLang if its default constructor is not overloaded, because NiLang requires the default constructor to \"modify\" a field of a immutable type.
"

# ╔═╡ edaa9fdb-3af8-4554-a701-0e3bff2107a5
md"It is also possbile to extract the fields directly."

# ╔═╡ 7551a880-340e-4e3f-815b-188e73f7eb9a
@i function complex_add(y::Complex{T}, x::Complex{T}) where T
    ((a, b), (c, d))::∅ ↔ ((@fields x), (@fields y))
	a += c
	b += d
    ((a, b), (c, d)) ↔ ((@fields x), (@fields y))::∅
end

# ╔═╡ 0489e51b-781f-4441-bb7f-ff3bd2e848ad
@test complex_add(1+2im, 3+4im) == (1+2im, 4+6im)

# ╔═╡ 7b0d30d6-39ff-4f6e-b13c-0ddbfcb576e5
md"Type cast is also possible"

# ╔═╡ 042297d8-6ab3-4ae6-b6e7-3b1ab2d5553b
@i function add4(a, b, c, d)
	complex_add(Complex{}(a, b), Complex{}(c, d))  # do not omit `{}`
end

# ╔═╡ 57d65a36-bfa8-4dc2-8e11-d87fa1324122
@test add4(1, 2, 3, 4) == (1, 2, 4, 6)

# ╔═╡ c21d81c3-981f-4472-ad61-d1661bfe5c4e
example("Implementing \"axpy\" function")

# ╔═╡ 99d6fe7b-d704-48f3-b115-2b3159a78068
md"`axpy` function is defined as ``\vec y += a \vec x``. One can modify the `Array` directly"

# ╔═╡ 1950ff70-54eb-4ece-a26d-a23fd0e90f5a
@i function arrayaxpy!(y!::Vector{T}, a::T, x::Vector{T}) where T
    for i=1:length(x)
		y![i] += a * x[i]
	end
end; @test arrayaxpy!(zeros(10), 2.0, collect(1.0:10.0))[1] ≈ collect(2.0:2.0:20.0)

# ╔═╡ 21458f81-9007-46f8-92e0-7a17c60beb36
md"To modify an element of a `Tuple`, we need to use a different style to avoid confusion with array"

# ╔═╡ 7813f4ce-6e98-45f3-94a8-7f5981129f2b
@i function tupleaxpy!(y!::NTuple{N,T}, a::T, x::NTuple{N,T}) where {N, T}
    for i=1:length(x)
		(y! |> tget(i)) += a * (x |> tget(i))
	end
end; @test tupleaxpy!((0,0,0), 2, (1,2,3))[1] == (2, 4, 6)

# ╔═╡ 59ec7cb7-6011-456d-9f57-a55bb8ea51a0
md"Here `data |> tget(i)` represents the `i`th field of the tuple (note it is not allowed to write `data.:i`)."

# ╔═╡ aacf63a2-9708-40db-8928-049621a7bbc4
md"## Control flows"

# ╔═╡ ad0097e7-c8ad-457a-82a9-18b998a9e9fb
md"""
#### If statement

The condition expression in `if` statement contains two parts, one is precondition and another is postcondition.

$(
twocol(md"
```julia
if (precondition[, postcondition])
	...
end
```
", md"
```julia
if (postcondition[, precondition])
	~(...)
end
```
")
)

where `...` are statements and `~(...)` are the backward execution of them.
"""

# ╔═╡ 94cd1345-3132-4882-86fe-d2429f610d1d
md"""If no postcondition is provided, it means the precondition is the same as the postcondition. It is translated to `if (cond, cond) ... end`. `elseif` is also supported."""

# ╔═╡ 4a558bd3-6e42-4c61-bd23-888b7f33ae25
@i function f7a(x)
	if x > 1
		x -= 1
	end
end; @test_throws InvertibilityError f7a(1.2)

# ╔═╡ 004f727a-e0c8-49cb-8858-dfdf4d3ac57a
@i function f7b(x, branch_keeper)
	branch_keeper ⊻= x > 1
	if (x > 1, branch_keeper)
		x += 1
	end
end; @test f7b(1.2, false)[1] == 2.2

# ╔═╡ 4c03cde9-b643-40ff-b275-f1795f88949e
title2("For statement")

# ╔═╡ fae7c74e-d25e-4c1e-ac97-199e6dae3365
md"""
The reversible `for` statement is similar to its irreversible counterpart.

$(
twocol(md"
```julia
for iter = start:step:stop
	...
end
```
", md"
```julia
for iter = stop:-step:start
	~(...)
end
```
")
)
"""

# ╔═╡ 2f2b24ea-66d0-4b3b-a460-53b6b3f28ef0
md"The iterator length should not be changed during the iteration."

# ╔═╡ e9dbb64a-27b9-443a-b917-69d55c290235
@i function f7c(x, y)
	for i=1:length(x)
		POP!(x, y[i])
	end
end; @test_throws InvertibilityError f7c([1,2,3], [0,0,0])

# ╔═╡ 0d56ce96-81a5-4102-acbc-7d88f80adcb3
md"There is an `InvertibilityError` because the length of the `x` has been changed. The inverse execution will give incorrect result."

# ╔═╡ 95c41bd1-e50a-42e8-93c3-3b754a458c13
title2("While statement")

# ╔═╡ b8629aeb-6c9a-44ed-87a1-9ab22d9485ed
md"""
The reversible `while` statement starts with the `@from` macro.

$(
twocol(md"
```julia
@from condition1 while condition2
	...
end
```
", md"
```julia
@from !(condition2) while !(condition1)
    ~(...)
end
```
")
)
"""

# ╔═╡ 3b211406-041f-4b41-acae-3958e4a37224
md"where the `condition1` in the forward pass is a condition that holds before entering the loop body, but broken at the first iteration, while `condition2` is just a normal while condition. In the backward pass, `!(condition1)` becomes the criteria to break the loop."

# ╔═╡ 62522772-cb59-4d13-acdd-d5067b223910
@i function f7d(x, i)
	@from i==0 while i<10
		i += 1
		x += 1
	end
end

# ╔═╡ 2ba68a0f-6e36-4ea2-a91d-6af43741bad1
f7d(1, 0)

# ╔═╡ 75cebaf1-38de-475f-892e-346fd2b46f6f
@test (~f7d)(11, 10) == (1, 0)

# ╔═╡ 72dcf2fe-eb48-4dee-8121-efafc87637e3
title2("Compute-copy-uncompute statement")

# ╔═╡ 84321198-93d4-4d22-8c0f-a5a10b884e1f
md"The *compute-copy-uncompute* statement is a widely used design pattern in reversible programming.
We compute the forward pass for the result, then we copy the result to the output variable, and run the backward pass to erase intermediate results.

For example, to compute `y = x * exp(k)`, we might write the following code"

# ╔═╡ 32244789-afbf-4215-97cd-15483f438eee
@i function f7e(y, x, k)
	expk ← zero(k)
	expk += exp(k)
	y += x * expk
	# uncompute the ancilla and deallocate it
	expk -= exp(k)
	expk → zero(k)
end; @test f7e(0.0, 2.0, 3.0)[1] ≈ 2.0 * exp(3.0)

# ╔═╡ 4b7f0baf-0316-4da7-9ded-50c064ddbaa3
md"It is equivalent to the following statement that generates the backward pass automatically for you."

# ╔═╡ 0e02952c-7589-4606-b006-16a9f3e52ae1
@i function f7f(y, x, k)
	# record the forward pass
	@routine begin
		expk ← zero(k)
		expk += exp(k)
	end
	y += x * expk
	# reverse execute the recorded the program
	~@routine
end; @test f7f(0.0, 2.0, 3.0)[1] ≈ 2.0 * exp(3.0)

# ╔═╡ f0904d3f-1bf1-459c-9959-b53c0f774e3f
example("Computing Fibonacci numbers")

# ╔═╡ 19bb2af5-2a67-453d-82b0-7d3059b1fa47
md"The sequence of Fibonacci numbers are: 1, 1, 2, 3, 5, 8"

# ╔═╡ 5b5858bf-63ac-4e31-a516-055a9cd18ffe
@i function rfib(out!, n::T) where T
    @routine begin
		n1 ← zero(T)
		n2 ← zero(T)
        n1 += n - 1
        n2 += n - 2
    end
    if (value(n) <= 2, ~)
        out! += 1
    else
        rfib(out!, n1)
        rfib(out!, n2)
    end
    ~@routine
end

# ╔═╡ 95060588-f24b-4eeb-9b0b-ed7159962a3c
@test rfib(0, 6)[1] == 8

# ╔═╡ c4cd9f88-9cd6-4364-b016-78f90aba6a66
title1("Extending the instruction set")
# They are translated to `y += f(args...)` is translated to `PlusEq(f)(y, args...)`, `y -= f(args...)` is translated to `MinusEq(f)(y, args...)`, `y *= f(args...)` is translated to `MulEq(f)(y, args...)`, `y /= f(args...)` is translated to `DivEq(f)(y, args...)` and `y ⊻= f(args...)` is translated to `XorEq(f)(y, args...)`.

# ╔═╡ f6049c78-7468-47ce-a4a5-84fab34d115a
title2("How to create a new elementary reversible function")

# ╔═╡ b0f73825-bbb1-448c-b491-bf634fdd398a
md"To define a pair of elementary functions that **reverse to each other**,
1. declare two functions `f` and `g` that each of them defines a mapping ``\mathbb{R}^n \rightarrow \mathbb{R}^n``
2. use `@dual f g` to tell NiLang they are reversible to each other.
3. if you want to make `f` and `g` differentiable, you can specify backward rules on these two function by defining two mappings on ``\mathbb{G}^n\rightarrow \mathbb{G}^n``, where ``\mathbb{G}`` is a 2-tuple of ``\mathbb{R}`` (or `NiLang.AD.GVar`) in NiLang. It is similar to `ForwardDiff.Dual` (check [ForwardDiff](https://github.com/JuliaDiff/ForwardDiff.jl)) but defined for the backward pass.

To define a **self-reversible** elementary function
1. declare a functions `f` that defines a mapping ``\mathbb{R}^n \rightarrow \mathbb{R}^n``
2. use `@selfdual f`
3. define the backward rule on `f` to make it differentiable.
"

# ╔═╡ 648cdcd6-f4f5-461f-a525-4b350cae9eb0
example("defining a new elementary function")

# ╔═╡ d6b1abd6-749d-4591-99e8-64aaa9199ab5
md"""
One can use the invertibility checker to check if the function is really reversible (under a certain tolerance `NiLangCore.GLOBAL_ATOL[]` = $(NiLangCore.GLOBAL_ATOL[])).
"""

# ╔═╡ f502b8c1-9b80-4e67-80e8-a64ddb88fb0f
@test NiLang.check_inv(M.new_forward, (3.0,))

# ╔═╡ 0bce342e-9a8e-4005-8b88-82da2d2c7163
md"""
To check of the gradients are properly defined, one can use `NiLang.AD.check_grad`
"""

# ╔═╡ fd9cf757-2698-4886-9f0a-c6c23ff0d331
@test NiLang.AD.check_grad(M.new_forward, (3.0,); iloss=1)

# ╔═╡ dda6652a-d063-4511-8041-e869bb88ca26
@test NiLang.AD.check_grad(M.new_backward, (3.0,); iloss=1)

# ╔═╡ a7d47e83-7f44-49d0-a43d-e01316fc6eba
title1("Performance Tips")

# ╔═╡ eca3efef-f35b-4623-8af8-0b830a55566d
md"The following trick still work in NiLang
* Removing boundary check with `@inbounds` (it works on FastStack),
* Add `@inline` before short functions,
* Add `@simd` before a for loop,

Other tricks like type stability are introduced in the [Julia documentation](https://docs.julialang.org/en/v1/manual/performance-tips/).
"

# ╔═╡ 45985244-adbf-4d6d-9732-a963cca62212
title2("Remove reversibility check")

# ╔═╡ 83d7e75f-7273-4c6a-bec1-a2180ebc3fb9
md"This can be done by putting an `@invcheckoff` before a code block."

# ╔═╡ 7fb05c65-f47c-430a-b588-c2f9bade40a9
example("computing the exp function by Taylor expansion")

# ╔═╡ 14c0caa1-51ea-448c-a7dc-d06e34dd0895
md"Note: this is not a clever implementation. There is an approach of defining it without allocation."

# ╔═╡ 457d07bb-e999-413e-8f29-58714670296f
@i function exp_with_reversibility_check(y::T, x::T) where T
	@routine begin
		N ← 1_000
		anc ← zeros(T, N)
		anc[1] += 1
		anc_y ← T(2.0)
		for i=2:N
			@routine begin
				temp ← zero(T)
				temp += x * anc[i-1]
			end
			anc[i] += temp / i
			anc_y += anc[i]
			~@routine
		end
	end
	y += anc_y
	~@routine
end

# ╔═╡ ac53eac0-1a59-4407-8bf6-3d8b966a9bff
@benchmark exp_with_reversibility_check(0.0, 1.0) seconds=0.3

# ╔═╡ 85c8ac7b-54f5-47dc-bd50-e78ffd6cf1cf
@i function exp_without_reversibility_check(y::T, x::T) where T
	@routine @invcheckoff begin
		N ← 1_000
		anc ← zeros(T, N)
		anc[1] += 1
		anc_y ← T(2.0)
		for i=2:N
			@routine begin
				temp ← zero(T)
				temp += x * anc[i-1]
			end
			anc[i] += temp / i
			anc_y += anc[i]
			~@routine
		end
	end
	y += anc_y
	~@routine
end

# ╔═╡ 95c55847-0591-4f7f-b9a1-aa974ccfef69
@benchmark exp_without_reversibility_check(0.0, 1.0) seconds=0.3

# ╔═╡ 91f8cfc6-e261-4945-8506-eed8caa607c2
title1("Multi-threading and CUDA")

# ╔═╡ c82b3b5c-c4e2-4bf6-b4ec-0d05ba9a669b
@i function multi_thread_exp(y::Vector, x::Vector)
	# check the size of `x` and `y`. `@assert` is not a valid statement in NiLang, so one should decorate it with `@safe` to tell the compiler, doing this is safe, do not check this statement.
	@safe @assert length(x) == length(y)
	@threads for i=1:length(y)
		y[i] += exp(x[i])
	end
end; @test multi_thread_exp(zeros(3), [1.0, 2.0, 3.0])[1] ≈ [exp(1.0), exp(2.0), exp(3.0)]

# ╔═╡ 32d75270-60d7-4326-a4ff-8674d0fbd491
md"With [CUDA](https://github.com/JuliaGPU/CUDA.jl), one can also define parallel reversile and differentiable GPU device functions. Use the broadcasting version `y += x` as an example."

# ╔═╡ 4b8834c1-8bb3-49f2-ae9e-1dbb8832d7f0
@i function addkernel(target, source)
    @invcheckoff begin
		@routine b ← (blockIdx().x-1) * blockDim().x + threadIdx().x
		if (b <= length(target), ~)
			@inbounds target[b] += source[b]
		end
		~@routine
	end
end

# ╔═╡ d8f0ae56-e643-48a1-86ee-1cd907ecb662
md"One can launch the kernel function in NiLang with `@cuda`"

# ╔═╡ e72b9dc2-dfac-4631-b114-01ec14297427
md"""
```julia
using CUDA

@i function :(+=)(identity)(target::CuArray, source::CuArray)
    @safe @assert length(target) == length(source)
    @cuda threads=256 blocks=ceil(Int,length(target)/256) addkernel(target, source)
end
```
"""

# ╔═╡ 8c93a773-edc0-4ec2-88ef-1b58b7deddc5
title2("Shared read-write in parallel computing and autodiff")

# ╔═╡ 16d08950-0575-4a4b-afc8-11ddca3198c7
md"The parallel code may suffer from the shared read issue when computing gradients.
Let's take a look at a parallel code that computes the loss ``\mathcal{L} = \sum x \vec z``."

# ╔═╡ 7c594d19-59fc-433a-bffa-c63bad46869e
@i function shared_read(loss::Real, y::Vector, x::Real, z::Vector)
	@safe @assert length(z) == length(y)
	@threads for i=1:length(y)
		y[i] += x * z[i]
	end
	for i=1:length(y)
		loss += y[i]
	end
end; @test shared_read(0.0, zeros(3), 2.0, [1.0, 2.0, 3.0])[1] ≈ 12

# ╔═╡ 345b344d-afda-4ce1-a0e7-ce6063a69206
md"However, when computing the gradients, the gradient on `x` will not be computed correctly."

# ╔═╡ c55eb045-daca-42f9-a357-095edef24644
let
	z = randn(100)
	_, gy, gx, gz = NiLang.AD.gradient(shared_read, (0.0, zeros(100), 2.0, z); iloss=1)
	@test gx ≈ sum(z)
end

# ╔═╡ c6903b65-a8b8-4aef-8c43-c822077b9d0e
md"The error is expected. Because the variable `x` is shared by multiple threads, when updating the gradient field of `x` in the backward pass, all threads will try to update the same gradient field, this is famous [race condition](https://en.wikipedia.org/wiki/Race_condition) in parallel computing."

# ╔═╡ f80353d6-0dfe-4b0a-a1af-655d344473bf
title1("Resources")

# ╔═╡ 4ca276fb-859d-4c5a-81c3-8e4b28922fa4
title2("Help and Discussion")

# ╔═╡ 3d020209-b8dd-4605-9329-78a985f0a6a3
md"""
`reversible-computing` channel of [Julia slack](https://julialang.org/slack/) and [Julia Zulip](https://julialang.zulipchat.com/register/).
"""

# ╔═╡ c7ec3496-79ea-4956-976c-b88dd22207c7
title2("Learning")

# ╔═╡ 8530a9d1-5a27-4a1d-883f-4b033a6f8fe4
md"""
1. Reversible computing book: Perumalla, Kalyan S. Introduction to reversible computing. CRC Press, 2013.
2. Our paper: Liu, Jin-Guo, and Taine Zhao. "Differentiate Everything with a Reversible Programming Language." arXiv:2003.04617 (2020).
"""

# ╔═╡ 7ce31932-0447-4445-99aa-7ebced7d0bad
TableOfContents()

# ╔═╡ Cell order:
# ╟─2061b434-0ad1-46eb-a0c7-1a5f432bfa62
# ╟─a4e76427-f051-4b29-915a-fdfce3a299bb
# ╟─c2c7b4d4-f8c9-4ebf-8da2-0103f03136e7
# ╟─12f07cc7-979c-43c3-9dc9-36ea1463c1f6
# ╟─611b577f-4722-42bf-8f8e-aeb2fb30be71
# ╟─605872cf-f3fd-462e-a2b1-7d1c5ae45efd
# ╟─fb3dee44-5fa9-4773-8b7f-a83c44358545
# ╠═d941d6c2-55bf-11eb-0002-35c7474e4050
# ╠═70088425-6779-4a2d-ba6d-b0a34c8e93a6
# ╠═af738f89-3214-429c-9c7d-18a6ea0d9401
# ╠═48d7ebc1-5def-4a57-9ec1-3fc370a4543f
# ╟─f0e94247-f615-472b-8218-3fa287b38aa1
# ╟─2581aa33-1dc5-40b1-aa9f-6a11cc750c93
# ╟─60575978-081a-4bca-a3ed-2b51cd6abc92
# ╠═f98305cb-4ba2-404a-a5c3-65510e059504
# ╟─e8cd6667-597f-458b-8465-1822e09a7891
# ╟─20145d75-004a-4c2f-b7ff-c400ca846d42
# ╟─c682a17f-600f-4034-bfe3-a851ab645c10
# ╟─5239dfe2-ea6d-4e07-a1b1-90954fe8ddc9
# ╟─e4d86a5a-e820-4a70-8a87-08bac416291b
# ╟─c307b6a4-906d-4be7-9fd7-57c942aded51
# ╟─47b502d4-e8af-4d58-9067-9700784ea435
# ╠═faecc0a7-55d7-42a1-8e9f-7e30143eef9c
# ╠═4cbe69d6-b68f-4bda-a0dd-209f9ee54f18
# ╟─7dc82f28-77bf-40da-b520-800ed1bc80c9
# ╟─9f5f9de3-9558-4c18-9d98-b77d19b570ec
# ╠═6dfcfa19-f78f-4dac-89f7-d3c5dbe17987
# ╟─dad1c6c0-d61b-4f9f-a71e-e683fe143aaa
# ╟─3e4a3916-8fe4-4262-bcd3-3014822717a3
# ╠═34906208-e6f1-4a67-860a-a7b056a86dde
# ╟─4601df35-679f-465d-9191-c18748b2fd83
# ╟─4fc72b9d-19a2-40f1-a4a8-5e97d3d5e529
# ╠═a0fde16f-8454-4f5c-a29c-a9e415c0c311
# ╟─633ff8f3-8d93-4f73-bec2-c42070e6ece9
# ╠═57f2d890-b5a0-47b7-9e3d-af4d03b10605
# ╟─34063cd0-171e-46ce-80dd-52a341fa50a1
# ╠═5d5d01db-8ff9-434c-8771-1fec6393e1fb
# ╠═10d85a50-f2f9-403e-8f6c-baef61cf702a
# ╟─b52648bf-a28a-48af-8912-31729d943ce0
# ╟─f45db10f-a836-40f3-9d8d-054ea6540e87
# ╟─1903563e-ccc2-44d9-9dbe-e5dede275b3c
# ╟─583f2585-15a3-47c6-a70e-e2f002754028
# ╠═60e6ff80-3593-4ae4-a273-914847f692db
# ╠═9e5cfd68-b58d-4d83-aae2-447e5f805c97
# ╠═dc85a942-cf52-4405-ad03-32a768e1b6e7
# ╟─5cdd346b-10a5-485c-ba78-4c0b3cb0e02f
# ╠═af9287b7-6131-46f6-beb8-6885e55e1975
# ╠═e20eeabf-1c80-431e-8cfc-4d1b79c52b5a
# ╟─90d30eea-53de-48a0-9700-ff35681fdf38
# ╠═390f58a5-6f5f-4d3a-bb16-ba04e43a07e7
# ╠═2b57443e-a516-434b-be86-80616a98e2f5
# ╟─fc2e27f9-b7ba-44cc-a953-6745548ad733
# ╠═fc744931-360b-4478-9f77-c50f048de243
# ╠═9a152b36-f377-44da-9700-ca9e05e365ff
# ╟─b6dcd18c-606f-4340-b2ec-163e8bad03f5
# ╟─a1a29f34-f8a9-4e9f-9afe-7d0096771440
# ╟─90bd6ad4-3dd8-4e7c-b445-aed1d248a2ec
# ╠═c0259e48-1973-486c-a828-1fcd3e4331c6
# ╟─8bbffa31-04a6-49ca-b36f-4d4140d75992
# ╠═a6f18c34-80ee-4b52-9ff8-f3c1b1d80f90
# ╠═a694132b-4f52-467f-8bc4-dc32fe2812db
# ╟─8c2c82f2-1240-4f2f-830e-ee8021c1a41a
# ╟─6203cf10-f8cc-4fb9-b814-7552b68c01dc
# ╠═f97a6bab-b9f9-4b95-98a9-381c51397526
# ╠═4ca48a2e-43da-457a-8e9f-6476097e4d7b
# ╟─2a2970f4-ab01-486b-89a2-6ff96f734018
# ╠═0b80d9be-53d7-4bf3-a558-659607af4709
# ╠═92362fda-bae2-4e35-bfe4-dcaea853d50b
# ╟─db9e7940-39f1-4ccf-ac70-146a521daa6e
# ╟─93936612-1447-4114-b864-aba43adef4bd
# ╠═b2add70c-c5d6-4e0f-a153-43e21a197181
# ╠═5c03d5a5-99f0-4efd-9a32-ce6d7c2b266c
# ╟─2349e3ea-3053-42a4-b9d9-f97a76e4abd7
# ╟─269f18ee-3cd8-466a-a522-7c624503e31b
# ╠═89139719-c478-4066-9452-f9893f36d561
# ╠═d620c5ee-7d9c-4d3f-9e87-0c828dfab9ca
# ╟─255d01b9-a873-4e63-9298-9d8f073348b0
# ╟─f6cf1729-766c-4ed7-b004-c8c8ec6c7e07
# ╠═3645d672-423f-4ac8-805f-0452793fee5a
# ╠═c2a0024e-11dd-4ef7-8346-4374d98cafc0
# ╟─b20004e9-3c73-4dfb-8fd5-f377786fd53b
# ╠═5c1952b1-5016-4c87-b23c-8e6a235bf8cd
# ╠═8e4470ee-01da-4547-b091-c4f65cd729b0
# ╟─0863bd06-cc70-4dde-b3b2-0a466805a356
# ╟─a0bae195-04e1-4642-9e14-fe4691e0906b
# ╠═20d6e8a0-2cf5-48ad-9549-60506b42b970
# ╠═7dee5748-ed73-4e13-aa80-7a50efbc8449
# ╟─4c719e9b-641e-404e-9ab7-59e89135f3ba
# ╠═77947e00-42c3-4c9e-b62a-b4b29489db43
# ╠═2614127d-34fb-4c3d-b678-42693f3c9341
# ╟─8f169235-3bd1-4cc4-a083-79736d306ad5
# ╠═dfc9d305-5bce-4555-bfa3-d8d61fe4ca09
# ╟─f6bfa015-c101-45e8-995c-2bb6a3b7dc7d
# ╟─8651d7ec-6bcd-4dbe-a062-c4bde32e5e91
# ╟─edaa9fdb-3af8-4554-a701-0e3bff2107a5
# ╠═7551a880-340e-4e3f-815b-188e73f7eb9a
# ╠═0489e51b-781f-4441-bb7f-ff3bd2e848ad
# ╟─7b0d30d6-39ff-4f6e-b13c-0ddbfcb576e5
# ╠═042297d8-6ab3-4ae6-b6e7-3b1ab2d5553b
# ╠═57d65a36-bfa8-4dc2-8e11-d87fa1324122
# ╟─c21d81c3-981f-4472-ad61-d1661bfe5c4e
# ╟─99d6fe7b-d704-48f3-b115-2b3159a78068
# ╠═1950ff70-54eb-4ece-a26d-a23fd0e90f5a
# ╟─21458f81-9007-46f8-92e0-7a17c60beb36
# ╠═7813f4ce-6e98-45f3-94a8-7f5981129f2b
# ╟─59ec7cb7-6011-456d-9f57-a55bb8ea51a0
# ╟─aacf63a2-9708-40db-8928-049621a7bbc4
# ╟─ad0097e7-c8ad-457a-82a9-18b998a9e9fb
# ╟─94cd1345-3132-4882-86fe-d2429f610d1d
# ╠═4a558bd3-6e42-4c61-bd23-888b7f33ae25
# ╠═004f727a-e0c8-49cb-8858-dfdf4d3ac57a
# ╟─4c03cde9-b643-40ff-b275-f1795f88949e
# ╟─fae7c74e-d25e-4c1e-ac97-199e6dae3365
# ╟─2f2b24ea-66d0-4b3b-a460-53b6b3f28ef0
# ╠═e9dbb64a-27b9-443a-b917-69d55c290235
# ╟─0d56ce96-81a5-4102-acbc-7d88f80adcb3
# ╟─95c41bd1-e50a-42e8-93c3-3b754a458c13
# ╟─b8629aeb-6c9a-44ed-87a1-9ab22d9485ed
# ╟─3b211406-041f-4b41-acae-3958e4a37224
# ╠═62522772-cb59-4d13-acdd-d5067b223910
# ╠═2ba68a0f-6e36-4ea2-a91d-6af43741bad1
# ╠═75cebaf1-38de-475f-892e-346fd2b46f6f
# ╟─72dcf2fe-eb48-4dee-8121-efafc87637e3
# ╟─84321198-93d4-4d22-8c0f-a5a10b884e1f
# ╠═32244789-afbf-4215-97cd-15483f438eee
# ╟─4b7f0baf-0316-4da7-9ded-50c064ddbaa3
# ╠═0e02952c-7589-4606-b006-16a9f3e52ae1
# ╟─f0904d3f-1bf1-459c-9959-b53c0f774e3f
# ╟─19bb2af5-2a67-453d-82b0-7d3059b1fa47
# ╠═5b5858bf-63ac-4e31-a516-055a9cd18ffe
# ╠═95060588-f24b-4eeb-9b0b-ed7159962a3c
# ╟─c4cd9f88-9cd6-4364-b016-78f90aba6a66
# ╟─f6049c78-7468-47ce-a4a5-84fab34d115a
# ╟─b0f73825-bbb1-448c-b491-bf634fdd398a
# ╟─648cdcd6-f4f5-461f-a525-4b350cae9eb0
# ╠═8c2c4fa6-172f-4dde-a279-5d0aecfdbe46
# ╟─d6b1abd6-749d-4591-99e8-64aaa9199ab5
# ╠═f502b8c1-9b80-4e67-80e8-a64ddb88fb0f
# ╟─0bce342e-9a8e-4005-8b88-82da2d2c7163
# ╠═fd9cf757-2698-4886-9f0a-c6c23ff0d331
# ╠═dda6652a-d063-4511-8041-e869bb88ca26
# ╟─a7d47e83-7f44-49d0-a43d-e01316fc6eba
# ╟─eca3efef-f35b-4623-8af8-0b830a55566d
# ╟─45985244-adbf-4d6d-9732-a963cca62212
# ╟─83d7e75f-7273-4c6a-bec1-a2180ebc3fb9
# ╟─7fb05c65-f47c-430a-b588-c2f9bade40a9
# ╟─14c0caa1-51ea-448c-a7dc-d06e34dd0895
# ╠═457d07bb-e999-413e-8f29-58714670296f
# ╠═3199a048-7b39-40f8-8183-6a54cccd91b6
# ╠═ac53eac0-1a59-4407-8bf6-3d8b966a9bff
# ╠═85c8ac7b-54f5-47dc-bd50-e78ffd6cf1cf
# ╠═95c55847-0591-4f7f-b9a1-aa974ccfef69
# ╟─91f8cfc6-e261-4945-8506-eed8caa607c2
# ╠═0e1ba158-a6bc-401c-9ba7-ed78020ad068
# ╠═c82b3b5c-c4e2-4bf6-b4ec-0d05ba9a669b
# ╟─32d75270-60d7-4326-a4ff-8674d0fbd491
# ╠═4b8834c1-8bb3-49f2-ae9e-1dbb8832d7f0
# ╟─d8f0ae56-e643-48a1-86ee-1cd907ecb662
# ╟─e72b9dc2-dfac-4631-b114-01ec14297427
# ╟─8c93a773-edc0-4ec2-88ef-1b58b7deddc5
# ╟─16d08950-0575-4a4b-afc8-11ddca3198c7
# ╠═7c594d19-59fc-433a-bffa-c63bad46869e
# ╟─345b344d-afda-4ce1-a0e7-ce6063a69206
# ╠═c55eb045-daca-42f9-a357-095edef24644
# ╟─c6903b65-a8b8-4aef-8c43-c822077b9d0e
# ╟─f80353d6-0dfe-4b0a-a1af-655d344473bf
# ╟─4ca276fb-859d-4c5a-81c3-8e4b28922fa4
# ╟─3d020209-b8dd-4605-9329-78a985f0a6a3
# ╟─c7ec3496-79ea-4956-976c-b88dd22207c7
# ╟─8530a9d1-5a27-4a1d-883f-4b033a6f8fe4
# ╟─7ce31932-0447-4445-99aa-7ebced7d0bad


using ModelingToolkit, StaticArrays, LinearAlgebra, SparseArrays
using DiffEqBase
using Test

canonequal(a, b) = isequal(simplify(a), simplify(b))

# Calculus
@parameters t σ ρ β
@variables x y z

eqs = [σ*(y-x),
       x*(ρ-z)-y,
       x*y - β*z]

simpexpr = [
   :(σ * (y - x))
   :(x * (ρ - z) - y)
   :(x * y - β * z)
   ]

for i in 1:3
   @test ModelingToolkit.simplified_expr.(eqs)[i] == simpexpr[i]
   @test ModelingToolkit.simplified_expr.(eqs)[i] == simpexpr[i]
end

∂ = ModelingToolkit.jacobian(eqs,[x,y,z])
for i in 1:3
    ∇ = ModelingToolkit.gradient(eqs[i],[x,y,z])
    @test canonequal(∂[i,:],∇)
end

@test all(canonequal.(ModelingToolkit.gradient(eqs[1],[x,y,z]),[σ * -1,σ,0]))
@test all(canonequal.(ModelingToolkit.hessian(eqs[1],[x,y,z]),0))

Joop,Jiip = eval.(ModelingToolkit.build_function(∂,[x,y,z],[σ,ρ,β],t))
J = Joop([1.0,2.0,3.0],[1.0,2.0,3.0],1.0)
@test J isa Matrix
J2 = copy(J)
Jiip(J2,[1.0,2.0,3.0],[1.0,2.0,3.0],1.0)
@test J2 == J

Joop,Jiip = eval.(ModelingToolkit.build_function(vcat(∂,∂),[x,y,z],[σ,ρ,β],t))
J = Joop([1.0,2.0,3.0],[1.0,2.0,3.0],1.0)
@test J isa Matrix
J2 = copy(J)
Jiip(J2,[1.0,2.0,3.0],[1.0,2.0,3.0],1.0)
@test J2 == J

Joop,Jiip = eval.(ModelingToolkit.build_function(hcat(∂,∂),[x,y,z],[σ,ρ,β],t))
J = Joop([1.0,2.0,3.0],[1.0,2.0,3.0],1.0)
@test J isa Matrix
J2 = copy(J)
Jiip(J2,[1.0,2.0,3.0],[1.0,2.0,3.0],1.0)
@test J2 == J

∂3 = cat(∂,∂,dims=3)
Joop,Jiip = eval.(ModelingToolkit.build_function(∂3,[x,y,z],[σ,ρ,β],t))
J = Joop([1.0,2.0,3.0],[1.0,2.0,3.0],1.0)
@test size(J) == (3,3,2)
J2 = copy(J)
Jiip(J2,[1.0,2.0,3.0],[1.0,2.0,3.0],1.0)
@test J2 == J

s∂ = sparse(∂)
@test nnz(s∂) == 8
Joop,Jiip = eval.(ModelingToolkit.build_function(s∂,[x,y,z],[σ,ρ,β],t,linenumbers=true))
J = Joop([1.0,2.0,3.0],[1.0,2.0,3.0],1.0)
@test length(nonzeros(s∂)) == 8
J2 = copy(J)
Jiip(J2,[1.0,2.0,3.0],[1.0,2.0,3.0],1.0)
@test J2 == J

# Function building

@parameters σ() ρ() β()
@variables x y z
eqs = [σ*(y-x),
       x*(ρ-z)-y,
       x*y - β*z]
f1,f2 = ModelingToolkit.build_function(eqs,[x,y,z],[σ,ρ,β])
f = eval(f1)
out = [1.0,2,3]
o1 = f([1.0,2,3],[1.0,2,3])
f = eval(f2)
f(out,[1.0,2,3],[1.0,2,3])
@test all(o1 .== out)

function test_worldage()
   @parameters σ() ρ() β()
   @variables x y z
   eqs = [σ*(y-x),
          x*(ρ-z)-y,
          x*y - β*z]
   f, f_iip = ModelingToolkit.build_function(eqs,[x,y,z],[σ,ρ,β];expression=Val{false})
   out = [1.0,2,3]
   o1 = f([1.0,2,3],[1.0,2,3])
   f_iip(out,[1.0,2,3],[1.0,2,3])
end
test_worldage()

## No parameters
@variables x y z
eqs = [(y-x)^2,
       x*(x-z)-y,
       x*y - y*z]
f1,f2 = ModelingToolkit.build_function(eqs,[x,y,z])
f = eval(f1)
out = zeros(3)
o1 = f([1.0,2,3])
f = eval(f2)
f(out,[1.0,2,3])
@test all(out .== o1)

function test_worldage()
   @variables x y z
   eqs = [(y-x)^2,
          x*(x-z)-y,
          x*y - y*z]
   f, f_iip = ModelingToolkit.build_function(eqs,[x,y,z];expression=Val{false})
   out = zeros(3)
   o1 = f([1.0,2,3])
   f_iip(out,[1.0,2,3])
end
test_worldage()

@test_nowarn muladd(x, y, ModelingToolkit.Constant(0))
@test promote(x, ModelingToolkit.Constant(0)) == (x, identity(0))
@test_nowarn [x, y, z]'