# this test is designed to test the ability of the package to use a state space that is not an AbstractArray.
# using Revise
using Test, Random, Setfield
using PseudoArcLengthContinuation
const PALC = PseudoArcLengthContinuation
####################################################################################################
# We start with a simple Fold problem
using LinearAlgebra
function F0(x::Vector, r)
	out = r .+  x .- x.^3
end

opt_newton0 = PALC.NewtonPar(tol = 1e-11, verbose = true)
	out0, hist, flag = @time PALC.newton(
		x -> F0(x, 1.),
		[0.8],
		opt_newton0)

opts_br0 = PALC.ContinuationPar(dsmin = 0.001, dsmax = 0.021, ds= -0.01, pMax = 4.1, pMin = -1., newtonOptions = setproperties(opt_newton0; maxIter = 70, tol = 1e-8), detectBifurcation = 0, maxSteps = 150, detectFold = true)

	br0, u1 = @time PALC.continuation(
		F0,
		out0, 1.0,
		opts_br0, printSolution = (x,p) -> x[1])

# plotBranch(br0);title!("a")

outfold, hist, flag = @time PALC.newtonFold(
	F0,
	(x, r) -> diagm(0 => 1 .- 3 .* x.^2),
	(x, r) -> diagm(0 => 1 .- 3 .* x.^2),
	(x, r, v1, v2) -> -6 .* x .* v1 .* v2,
	br0, 2, #index of the fold point
	opts_br0.newtonOptions)
		flag && printstyled(color=:red, "--> We found a Fold Point at α = ",outfold.p, ", from ", br0.foldpoint[1][3], "\n")

####################################################################################################
# Here is a more involved example

function Fb(x::BorderedArray, r, s = 1.0)
	BorderedArray(r .+  s .* x.u .- (x.u).^3, x.p - 0.0)
end

# there is no finite differences defined, so we need to provide a linearsolve
# we could also have used GMRES. We define a custom Jacobian which will be used for evaluation and jacobian inverse
struct Jacobian
	x
	r
	s
end

# We express the jacobian operator
function (J::Jacobian)(dx)
	BorderedArray( (J.s .- 3 .* ((J.x).u).^2) .* dx.u, dx.p)
end

struct linsolveBd <: PALC.AbstractBorderedLinearSolver end

function (l::linsolveBd)(J, dx)
	x = J.x
	r = J.r
	out = BorderedArray(dx.u ./ (J.s .- 3 .* (x.u).^2), dx.p)
	out, true, 1
end

sol = BorderedArray([0.8], 0.0)

opt_newton = PALC.NewtonPar(tol = 1e-11, verbose = true, linsolver = linsolveBd())
out, hist, flag = @time PALC.newton(
		x -> Fb(x, 1., 1.),
		x -> Jacobian(x, 1., 1.),
		sol,
		opt_newton)

opts_br = PALC.ContinuationPar(dsmin = 0.001, dsmax = 0.05, ds= -0.01, pMax = 4.1, pMin = -1., newtonOptions = setproperties(opt_newton; maxIter = 70, tol = 1e-8), detectBifurcation = 0, maxSteps = 150)

	br, u1 = @time PALC.continuation(
		(x, r) -> Fb(x, r, 1.),
		(x, r) -> Jacobian(x, r, 1.),
		out, 1.,
		opts_br; printSolution = (x,p) -> x.u[1])

# plotBranch(br);title!("")

outfold, hist, flag = @time PALC.newtonFold(
	(x, r) -> Fb(x, r),
	(x, r) -> Jacobian(x, r, 1.),
	(x, r) -> Jacobian(x, r, 1.),
	(x, r, v1, v2) -> BorderedArray(-6 .* x.u .* v1.u .* v2.u, 0.),
	br, 1, #index of the fold point
	opts_br.newtonOptions)
		flag && printstyled(color=:red, "--> We found a Fold Point at α = ", outfold.p, ", from ", br.foldpoint[1][3],"\n")

outfoldco, hist, flag = @time PALC.continuationFold(
	(x, r, s) -> Fb(x, r, s),
	(x, r, s) -> Jacobian(x, r, s),
	(x, r, s) -> Jacobian(x, r, s),
	s -> ((x, r, v1, v2) -> BorderedArray(-6 .* x.u .* v1.u .* v2.u, 0.)),
	br, 1,
	1.0, plot = false, opts_br)

# try with newtonDeflation
printstyled(color=:green, "--> test with Newton deflation 1")
deflationOp = DeflationOperator(2.0, (x, y) -> dot(x, y), 1.0, [zero(sol)])
soldef0 = BorderedArray([0.1], 0.0)
soldef1, _, _ = @time PALC.newton(
	x -> Fb(x, 0., 1.),
	x -> Jacobian(x, 0., 1.),
	soldef0,
	opt_newton, deflationOp)

push!(deflationOp, soldef1)

Random.seed!(1231)
printstyled(color=:green, "--> test with Newton deflation 2")
soldef2, _, _ = @time PALC.newton(
	x -> Fb(x, 0., 1.),
	x -> Jacobian(x, 0., 1.),
	rmul!(soldef0,rand()),
	opt_newton, deflationOp)

# test indexing
deflationOp[1]

# test length
length(deflationOp)

# test pop!
pop!(deflationOp)

####################################################################################################
using KrylovKit

function Fr(x::RecursiveVec, r, s = 1.0)
	out = similar(x)
	for ii=1:length(x)
		out[ii] .= r .+  s .* x[ii] .- x[ii].^3
	end
	out
end

# there is no finite differences defined, so we need to provide a linearsolve
# we could also have used GMRES. We define a custom Jacobian which will be used for evaluation and jacobian inverse
struct JacobianR
	x
	s
end

# We express the jacobian operator
function (J::JacobianR)(dx)
	out = similar(dx)
	for ii=1:length(out)
		out[ii] .= (J.s .- 3 .* (J.x[ii]).^2) .* dx[ii]
	end
	return out
end

struct linsolveBd_r <: PALC.AbstractBorderedLinearSolver end

function (l::linsolveBd_r)(J, dx)
	x = J.x
	out = similar(dx)
	for ii=1:length(out)
		out[ii] .= dx[ii] ./ (J.s .- 3 .* (x[ii]).^2)
	end
	out, true, 1
end

opt_newton0 = PALC.NewtonPar(tol = 1e-11, verbose = false, linsolver = linsolveBd_r())
	out0, hist, flag = @time PALC.newton(
		x -> Fr(x, 1.),
		x -> JacobianR(x, 1),
		RecursiveVec([1 .+ 0.1*rand(10) for _=1:2]),
		opt_newton0)

opts_br0 = PALC.ContinuationPar(dsmin = 0.001, dsmax = 0.2, ds= -0.01, pMin = -1.1, pMax = 1.1, newtonOptions = opt_newton0)

br0, u1 = @time PALC.continuation(
	(x, r) -> Fr(x, r),
	(x, r) -> JacobianR(x, 1),
	out0, 0.9, opts_br0;
	plot = false,
	printSolution = (x,p) -> x[1][1])

br0, u1 = @time PALC.continuation(
	(x, r) -> Fr(x, r),
	(x, r) -> JacobianR(x, 1),
	out0, 0.9, opts_br0;
	tangentAlgo = BorderedPred(),
	plot = false,
	printSolution = (x,p) -> x[1][1])

outfold, hist, flag = @time PALC.newtonFold(
	(x, r) -> Fr(x, r),
	(x, r) -> JacobianR(x, 1),
	(x, r) -> JacobianR(x, 1),
	(x, r, v1, v2) -> RecursiveVec([-6 .* x[ii] .* v1[ii] .* v2[ii] for ii=1:length(x)]),
	br0, 1, #index of the fold point
	opt_newton0)
		flag && printstyled(color=:red, "--> We found a Fold Point at α = ", outfold.p, ", from ", br0.foldpoint[1][3],"\n")

outfoldco, hist, flag = @time PALC.continuationFold(
	(x, r, s) -> Fr(x, r, s),
	(x, r, s) -> JacobianR(x, s),
	(x, r, s) -> JacobianR(x, s),
	s -> ((x, r, v1, v2) -> RecursiveVec([-6 .* x[ii] .* v1[ii] .* v2[ii] for ii=1:length(x)])),
	br0, 1,	1.0, opts_br0;
	tangentAlgo = SecantPred(), plot = false)

outfoldco, hist, flag = @time PALC.continuationFold(
	(x, r, s) -> Fr(x, r, s),
	(x, r, s) -> JacobianR(x, s),
	(x, r, s) -> JacobianR(x, s),
	s -> ((x, r, v1, v2) -> RecursiveVec([-6 .* x[ii] .* v1[ii] .* v2[ii] for ii=1:length(x)])),
	br0, 1,	1.0, opts_br0;
	tangentAlgo = BorderedPred(), plot = false)

# try with newtonDeflation
printstyled(color=:green, "--> test with Newton deflation 1")
deflationOp = DeflationOperator(2.0, (x, y) -> dot(x, y), 1.0, [(out0)])
soldef1, _, _ = @time PALC.newton(
	x -> Fr(x, 0.),
	x -> JacobianR(x, 0.),
	out0,
	opt_newton0, deflationOp)

push!(deflationOp, soldef1)


# A unit fraction contains 1 in the numerator. The decimal representation of
# the unit fractions with denominators 2 to 10 are given:

# 1/2 =   0.5
# 1/3 =   0.(3)
# 1/4 =   0.25
# 1/5 =   0.2
# 1/6 =   0.1(6)
# 1/7 =   0.(142857)
# 1/8 =   0.125
# 1/9 =   0.(1)
# 1/10    =   0.1
#
# Where 0.1(6) means 0.166666..., and has a 1-digit recurring cycle. It can be
# seen that 1/7 has a 6-digit recurring cycle.
#
# Find the value of d < 1000 for which 1/d contains the longest recurring cycle
# in its decimal fraction part.

using ProjectEulerSolutions

# Solve using greedy regex to pull out the longest repeated sequence of digits.
# Could be made slightly faster to consider only prime values between 2 and n.
function p026solution_regex(n::Integer=3)::Integer
    numdigits = 2500
    startvalue = BigInt(10)^numdigits

    regexstring = r"(\d+?)\1"
    lengths = zeros(Integer, n)

    @simd for i = 3:2:n
        stringnum = string(fld(startvalue, i))
        @inbounds lengths[i] = mapreduce(m->length(m.captures[1]), max, eachmatch(regexstring, stringnum))
    end

    return findmax(lengths)[2]
end

# Cycles through a sequence of digits building up and storing larger numbers
# until it sees one it has seen before.  Could be made slightly faster to
# consider only prime values between 2 and n.
function p026solution_repeats(n::Integer=3)::Integer
    totalmax = 0
    maxind = 1
    for i = 3:2:n
        seennumbers = Set{Integer}()
        r = 10
        j = 0
        while ! in(r, seennumbers)
            push!(seennumbers, r)
            r = 10 * (r % i)
            j += 1
        end

        if j > totalmax
            totalmax = j
            maxind = i
        end
    end

    return maxind
end

p026 = Problems.Problem(Dict("Regex" => p026solution_regex,
                             "Repeats" => p026solution_repeats))

Problems.benchmark(p026, 1000)


include("../src/LarSurf.jl")
# include("arr_fcn.jl")

import SparseArrays.spzeros
using Plasm
using LinearAlgebraicRepresentation
Lar = LinearAlgebraicRepresentation

# using arr_fcn

block_size = [2, 2, 2]
# function get_boundary3(block_size)
V, CVill = Lar.cuboidGrid([block_size[1], block_size[2], block_size[3]])

# A lot of work can be done by this:
# V, (VV, EV, FV, CV) = Lar.cuboidGrid([2,2,2], true)

# CVill to FVi
nc = size(CVill)[1]
nfaces_per_C = 6
nfaces = nfaces_per_C * nc
FVi = Array{Int64}(undef, nfaces, 4)
#     println(sz)
# produce faces
for k in 1:nc
    cube_points = CVill[k]
    FVi[(k - 1) * nfaces_per_C + 1, :] = cube_points[[1, 2, 4, 3]]
    FVi[(k - 1) * nfaces_per_C + 2, :] = cube_points[[1, 5, 6, 2]]
    FVi[(k - 1) * nfaces_per_C + 3, :] = cube_points[[1, 3, 7, 5]]
    FVi[(k - 1) * nfaces_per_C + 4, :] = cube_points[[3, 4, 8, 7]]
    FVi[(k - 1) * nfaces_per_C + 5, :] = cube_points[[2, 6, 8, 4]]
    FVi[(k - 1) * nfaces_per_C + 6, :] = cube_points[[6, 5, 8, 7]]
end

VFi = convert(Array{Int64, 2}, FVi')

# convert from list of list to 2D array
# CVi = Array{Int64}(undef, size(CVill)[1], size(CVill[1])[1])
CVi = LarSurf.ll2array(CVill)

nvertices = size(V)[2]

# CVill to sparse
print("CVill to sparse, sz: ", nc, " nvertices: ", nvertices)
VF01 = LarSurf.ind_to_sparse(VFi, nvertices, 1)
CV01 = LarSurf.ind_to_sparse(CVi, nvertices, 2)

boundary_numbers = CV01 * VF01

# take just o

boundary = boundary_numbers .== 4

boundary = convert(Array{Int8, 2}, boundary)

# list of lists

FVill = [FVi[k, :] for k=1:size(FVi,1)]

# Plasm.view(V, CVill)
# Plasm.viewexploded(V, CVill)(2,2,2)
# Plasm.viewexploded(V, FVill)(2,2,2)
print("boundary stats (#true, # false): ", sum(boundary .== true), ", ", sum(boundary .== false) )
# return boundary, V, FVi
# end


using MathOptInterface
const MOI = MathOptInterface

# Simple evaluator for testing

struct ExprEvaluator <: MOI.AbstractNLPEvaluator
    constraints::Vector{Expr}
end

function MOI.initialize(d::ExprEvaluator, requested_features::Vector{Symbol})
    if requested_features != [:ExprGraph]
        error("Only supports expression graph!")
    end
end

MOI.constraint_expr(evaluator::ExprEvaluator, i) = evaluator.constraints[i]

@testset "linear expressions" begin
    optimizer = SCIP.Optimizer(display_verblevel=0)

    @test MOI.supports(optimizer, MOI.NLPBlock()) == true

    # min. -x
    # s.t. 2x + y <= 1.5
    #      x,y >= 0

    x, y = MOI.add_variables(optimizer, 2)

    # bounds
    MOI.add_constraint(optimizer, x, MOI.GreaterThan(0.0))
    MOI.add_constraint(optimizer, y, MOI.GreaterThan(0.0))

    # "nonlinear" constraints (x => x[1], y => x[2])
    data = MOI.NLPBlockData(
        [MOI.NLPBoundsPair(-Inf, 1.5)],
        ExprEvaluator([:(2 * x[$x] + x[$y] <= 1.5)]),
        false # no objective
    )
    MOI.set(optimizer, MOI.NLPBlock(), data)

    # objective
    MOI.set(optimizer, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
            MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([-1.0], [x]), 0.0))
    MOI.set(optimizer, MOI.ObjectiveSense(), MOI.MIN_SENSE)

    # solve
    MOI.optimize!(optimizer)
    @test MOI.get(optimizer, MOI.TerminationStatus()) == MOI.OPTIMAL
    @test MOI.get(optimizer, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT

    atol, rtol = 1e-6, 1e-6
    @test MOI.get(optimizer, MOI.ObjectiveValue()) ≈ -0.75 atol=atol rtol=rtol
    @test MOI.get(optimizer, MOI.VariablePrimal(), x) ≈ 0.75 atol=atol rtol=rtol
    @test MOI.get(optimizer, MOI.VariablePrimal(), y) ≈ 0.0 atol=atol rtol=rtol
end

@testset "pot pourri" begin
    optimizer = SCIP.Optimizer(display_verblevel=0)

    @test MOI.supports(optimizer, MOI.NLPBlock()) == true

    num_vars = 20
    x = MOI.add_variables(optimizer, num_vars)
    for i in 1:num_vars
        MOI.add_constraint(optimizer, x[i], MOI.Interval(0.1, 10.0))
    end

    # Nonlinear equations.
    rhs = 2.0
    expressions = Expr[
        :(x[$(x[1])] + 0.5                      == rhs), # VARIDX / CONSTANT
        :(x[$(x[2])] - x[$(x[3])]               == rhs), # MINUS (binary)
        :(-x[$(x[4])] + 4.0                     == rhs), # MINUS (unary)
        :(x[$(x[5])] + x[$(x[6])] + x[$(x[7])]  == rhs), # SUM
        :(x[$(x[8])] * x[$(x[9])] * x[$(x[10])] == rhs), # PRODUCT
        :((x[$(x[11])])^3                       == rhs), # INTPOWER
        :((x[$(x[12])])^0.8                     == rhs), # REALPOWER
        :(x[$(x[13])] / x[$(x[14])]             == rhs), # DIV
        :(sqrt(x[$(x[15])])                     == rhs), # SQRT
        :(exp(x[$(x[16])])                      == rhs), # EXP
        :(log(x[$(x[17])])                      == rhs), # LOG
        :(abs(x[$(x[18])] - 11)                 == rhs), # ABS
        :(min(x[$(x[19])], x[$(x[20])]) + 1     == rhs), # MIN
        :(max(x[$(x[19])], x[$(x[20])]) - 1     == rhs), # MAX
    ]

    data = MOI.NLPBlockData(
        [MOI.NLPBoundsPair(rhs, rhs) for i in 1:length(expressions)],
        ExprEvaluator(expressions),
        false # no objective
    )
    MOI.set(optimizer, MOI.NLPBlock(), data)

    # Solve and get results.
    MOI.optimize!(optimizer)
    @test MOI.get(optimizer, MOI.TerminationStatus()) == MOI.OPTIMAL
    @test MOI.get(optimizer, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
    sol = MOI.get(optimizer, MOI.VariablePrimal(), x)

    # Evaluate lhs of all the expressions with solution values.
    atol, rtol = 1e-6, 1e-6
    @test sol[1] + 0.5              ≈ rhs  atol=atol rtol=rtol
    @test sol[2] - sol[3]           ≈ rhs  atol=atol rtol=rtol
    @test -sol[4] + 4.0             ≈ rhs  atol=atol rtol=rtol
    @test sol[5] + sol[6] + sol[7]  ≈ rhs  atol=atol rtol=rtol
    @test sol[8] * sol[9] * sol[10] ≈ rhs  atol=atol rtol=rtol
    @test (sol[11])^3               ≈ rhs  atol=atol rtol=rtol
    @test (sol[12])^0.8             ≈ rhs  atol=atol rtol=rtol
    @test sol[13] / sol[14]         ≈ rhs  atol=atol rtol=rtol
    @test sqrt(sol[15])             ≈ rhs  atol=atol rtol=rtol
    @test exp(sol[16])              ≈ rhs  atol=atol rtol=rtol
    @test log(sol[17])              ≈ rhs  atol=atol rtol=rtol
    @test abs(sol[18] - 11)         ≈ rhs  atol=atol rtol=rtol
    @test min(sol[19], sol[20])     ≈ 1.0  atol=atol rtol=rtol
    @test max(sol[19], sol[20])     ≈ 3.0  atol=atol rtol=rtol
end

@testset "add nonlinear constraint after solve" begin
    optimizer = SCIP.Optimizer()
    MOI.set(optimizer, MOI.RawOptimizerAttribute("display/verblevel"), 0)

    x, y = MOI.add_variables(optimizer, 2)

    data1 = MOI.NLPBlockData(
        [MOI.NLPBoundsPair(rhs, rhs) for rhs in [1.0, 2.0]],
        ExprEvaluator([:(x[$x] == 1.0), :(x[$y] == 2.0)]),
        false
    )
    MOI.set(optimizer, MOI.NLPBlock(), data1)

    MOI.optimize!(optimizer)

    data2 = MOI.NLPBlockData(
        [MOI.NLPBoundsPair(rhs, rhs) for rhs in [1.0, 2.0, 3.0]],
        ExprEvaluator([:(x[$x] == 1.0), :(x[$y] == 2.0),
                       :(x[$x] + x[$y] == 3.0)]),
        false
    )
    MOI.set(optimizer, MOI.NLPBlock(), data2)
end


using Statistics: mean

# Test demean function
@testset "Demean Test" begin
    arr = [1, 2, 3, 4, 5]
    arr_copy = [1, 2, 3, 4, 5]
    demean!(arr)
    @test arr == (arr_copy .-= mean(arr_copy))
end


mutable struct ConstantMean{T<:Real} <: PriorMean{T}
    C::T
    opt::Optimizer
end

"""
**ConstantMean**
```julia
ConstantMean(c::T=1.0;opt::Optimizer=Adam(α=0.01))
```

Construct a prior mean with constant `c`
Optionally set an optimizer `opt` (`Adam(α=0.01)` by default)
"""
function ConstantMean(c::T=1.0;opt::Optimizer=Adam(α=0.01)) where {T<:Real}
    ConstantMean{T}(c,opt)
end

function update!(μ::ConstantMean{T},grad::AbstractVector{T}) where {T<:Real}
    μ.C += update(μ.opt,sum(grad))
end

(μ::ConstantMean{T})(x::Real) where {T<:Real} = μ.C
(μ::ConstantMean{T})(x::AbstractArray) where {T<:Real} = μ.C*ones(T,size(X))


using DifferentialEquations
using ApproxFun, Sundials
using PyPlot

#using Pkg
#Pkg.add("DiffEqOperators")
using DiffEqOperators, LinearAlgebra
begin
	S = Fourier()
	#S = Chebyshev()
	n = 2^8
	x = points(S, n)
	T = ApproxFun.plan_transform(S, n)
	Ti = ApproxFun.plan_itransform(S, n)

	d = 0.2
	nu = 2e-2
	dt = 1e-3
	Nmax =  (3000+1)
	t0 = Nmax * dt

	u₀ = T*(cos.(x).*(sin.(x).+1));
end

function d_dx(S, n, N)
    D = Derivative(S,n)
    return D[1:N,1:N]
end

Dx(S,m) = d_dx(S,m,n)
L1 = Dx(S,1)
L2 = Dx(S,2)
L4 = Dx(S,4)

A = DiffEqArrayOperator(-d*Diagonal(L2)-nu*Diagonal(L4))

function f(dû,û,tmp,t)
  # Transform u back to point-space
  mul!(tmp,Ti,û)
  tmp= 0.5.*tmp.*tmp
  mul!(tmp, T, tmp)
  mul!(dû,-L1,tmp)
  #mul!(dû,T,tmp) # Transform back to Fourier space
end

prob = SplitODEProblem(A, f, u₀, (0.0,t0), similar(u₀));

sol = solve(prob, ETDRK4(), dt=dt)

plot(Ti*sol(0.001))
plot(Ti*sol(3))
#plot(Ti*sol(3))

#
u = zeros(eltype(sol),(n, Nmax))
for it in 1:Nmax
   print(dt*it)
   u[:,it] = Ti*sol(dt*it)
   #println(u[:,it])
end
size(u)

IM = PyPlot.imshow(abs.(u))
cb = PyPlot.colorbar(IM, orientation="horizontal")
#gcf()


using Plots
using MLDatasets
using LinearAlgebra

function DenseLayer(batch_size, input_size, output_size, act_f, d_act_f)
    w = [rand(output_size, input_size) for _ in 1:batch_size]
    b = [rand(output_size) for _ in 1:batch_size]

    dlocaly_dpreact = [Array{Float32}(undef, output_size, output_size) for _ in 1:batch_size]
    dpreact_dlocalX = [Array{Float32}(undef, output_size, input_size) for _ in 1:batch_size]
    dpreact_dlocalw = [Array{Float32}(undef, input_size) for _ in 1:batch_size]

    function forward!(X) # expects X :: Array{Array{Number, 2}}
        y = [similar(X[i]) for i in 1:batch_size]

        for i in eachindex(X, y)
            y_ = w[i] * X[i] + b[i]

            dlocaly_dpreact[i] = d_act_f(y_) # not point wise because of softmax activation (relu here returns a diagonal matrix)
            y[i] = act_f(y_)

            dpreact_dlocalX[i] = w[i]
            dpreact_dlocalw[i] = X[i]
        end

        return y
    end

    function backward!(dprediction_dlocaly, lr)
        dprediction_dlocalX = [zeros(input_size) for _ in 1:BATCH_SIZE]

        for i in 1:batch_size
            dprediction_dpreact = dlocaly_dpreact[i] * dprediction_dlocaly[i]
            dprediction_dlocalX[i] = Array(dpreact_dlocalX[i]' * dprediction_dpreact)

            w[i] .-= lr .* Array(dprediction_dpreact * dpreact_dlocalw[i]')
            b[i] .-= lr .* dprediction_dpreact
        end

        return dprediction_dlocalX
    end

    return forward!, backward!
end

function BatchNormLayer(batch_size, in_out_size)
    α = rand()
    β = rand()
    ϵ = 0.00001 # to prevent devision by zero

    dlocaly_dα = [zeros(in_out_size) for _ in 1:batch_size]
    dlocaly_dlocalX = [zeros(in_out_size, in_out_size) for _ in 1:batch_size]

    function forward!(X) # ϵ to avoid dividing by zero
        y = [similar(X[i]) for i in 1:batch_size]

        N = sum([length(X[i]) for i in 1:batch_size])
        μ = sum([sum(X[i]) for i in 1:batch_size]) / N
        σ = sum([sum((X[i] .- μ) .^ 2) for i in 1:batch_size]) / N

        for i in 1:batch_size
            bn = (X[i] .- μ) ./ sqrt(σ + ϵ)
            y[i] = α .* bn .+ β

            d_BN = (sqrt(σ + ϵ) .- (X[i] .- μ) .* 0.5 .* (1/(σ + ϵ))) ./ (sqrt(σ + ϵ) ^ 2)
            dlocaly_dlocalX[i] .= α .* d_BN # sqrt(σ + ϵ)
            dlocaly_dα[i] .= bn
        end

        return y
    end

    function backward!(dprediction_dlocaly, lr)
        dprediction_dlocalX = [similar(dprediction_dlocaly[i]) for i in 1:batch_size]

        for i in 1:batch_size
            dprediction_dlocalX[i] .= dlocaly_dlocalX[i]' * dprediction_dlocaly[i]

            α -= lr .* (dprediction_dlocaly[i]' * dlocaly_dα[i])
            β -= lr .* sum(dprediction_dlocaly[i])
        end

        return dprediction_dlocalX
    end

    return forward!, backward!
end

function feed_forward!(X::Array, forward_layer_functions::Array)
    for ff_layer in forward_layer_functions
        X = ff_layer(X)
    end
    return X
end

function feed_backwards!(d_pred::Array, backward_layer_functions::Array; lr=0.006)
    for fb_layer in reverse(backward_layer_functions)
        d_pred = fb_layer(d_pred, lr)
    end
    return nothing
end

# ::: ACTIVATION FUNCTIONS
relu(A) = map(a->a * (a>0), A)
d_relu(A) =  Matrix{Float32}(I, length(A), length(A)) .* map(a->(a>0), A)
sigmoid(x; base=ℯ) = 1 ./ (1 .+ ℯ .^ .-x)
d_sigmoid(x; base=ℯ) = Matrix(I, length(x), length(x)) .* (sigmoid(x, base=base) .* (1 .- sigmoid(x, base=base)))
softmax(A::AbstractArray) = [exp(a)/sum(exp.(A)) for a in A]
softmax(A::AbstractArray, ind::Integer) = exp(A[ind])/sum(exp.(A))
function d_softmax(A::AbstractArray)
    out = zeros(length(A), length(A))
    for i in eachindex(A)
        for j in eachindex(A)
            out[i, j] = softmax(A, i) * ((i == j) - softmax(A, j))
        end
    end
    out
end
function d_softmax(A::AbstractArray, ind::Integer)
    out = zeros(length(A))
    for i in eachindex(A)
        out[i] = softmax(A, i) * ((i == ind) - softmax(A, ind))
    end
end

# ::: LOSS FUNCTIONS
crossentropy(p::AbstractArray, q::AbstractArray; base=ℯ) = -sum(q .* log.(base, p))
d_crossentropy(p::AbstractArray, q::AbstractArray; base=ℯ) = -(q .* d_log.(p, base=base) .* log.(base, p))
mse(y_, y) = sum((y .- y_) .^ 2)
d_mse(y_, y) = 2 .* (y .- y_)

# ::: HELPER FUNCTIONS
d_log(x; base=ℯ) = 1/(x*log(base))
function onehot(i, max)
    r = zeros(max)
    r[i] = 1
    return r
end
function get_MNIST_minibatch(inputs::Array, labels::Array, batch_size)
    x_out = Array{Array}(undef, batch_size)
    y_out = Array{Array}(undef, batch_size)
    for i in 1:batch_size
        rn = rand(1:length(inputs[1, :]))
        x_out[i] = inputs[:, rn]
        y_out[i] = labels[:, rn]
    end
    return x_out, y_out
end

# TESTS

# MNIST data
train_x, train_y = MNIST.traindata()
train_x = reshape(Array{Float32}(train_x), (size(train_x)[1] * size(train_x)[2], size(train_x)[3]))
train_y = reshape(collect(Iterators.flatten([onehot(i+1, 10) for i in train_y])), (10, 60000))
test_x, test_y = MNIST.testdata()
test_x = reshape(Array{Float32}(test_x), (size(test_x)[1] * size(test_x)[2], size(test_x)[3]))
test_y = reshape(collect(Iterators.flatten([onehot(i+1, 10) for i in test_y])), (10, 10000))

# params
TRAIN_EPISODES = 500
TEST_EPISODES = 1000
BATCH_SIZE = 35
LEARNING_RATE = 0.005
INPUT_SIZE = length(train_x[:, 1])
OUTPUT_SIZE = length(train_y[:, 1])

# network
l1, dl1 = DenseLayer(BATCH_SIZE, INPUT_SIZE, (INPUT_SIZE / 2) |> Int, relu, d_relu)
l2, dl2 = BatchNormLayer(BATCH_SIZE, (INPUT_SIZE / 2) |> Int)
l3, dl3 = DenseLayer(BATCH_SIZE, (INPUT_SIZE / 2) |> Int, (INPUT_SIZE / 4) |> Int, relu, d_relu)
l4, dl4 = BatchNormLayer(BATCH_SIZE, (INPUT_SIZE / 4) |> Int)
l5, dl5 = DenseLayer(BATCH_SIZE, (INPUT_SIZE / 4) |> Int, (INPUT_SIZE / 8) |> Int, relu, d_relu)
l6, dl6 = BatchNormLayer(BATCH_SIZE, (INPUT_SIZE / 8) |> Int)
l7, dl7 = DenseLayer(BATCH_SIZE, (INPUT_SIZE / 8) |> Int, OUTPUT_SIZE, softmax, d_softmax)

ff_net = [l1,l2,l3,l4,l5,l6,l7]
fb_net = [dl1,dl2,dl3,dl4,dl5,dl6,dl7]

x1, y1 = get_MNIST_minibatch(train_x, train_y, BATCH_SIZE)

y_ = x1 |> l1 |> l2 |> l3 |> l4 |> l4 |> l5 |> l6 |> l7
y_ = map(x->clamp.(x, 0.000001, 0.999999), y_)

[crossentropy(y_[i], y1[i]) for i in 1:BATCH_SIZE]
dloss = [d_crossentropy(y_[i], y1[i]) for i in 1:BATCH_SIZE]

d1 = dl7(dloss, LEARNING_RATE)
d2 = dl6(d1, LEARNING_RATE)

# training
train_metric = [NaN for _ in 1:TRAIN_EPISODES]; for e in 1:TRAIN_EPISODES

    x, y = get_MNIST_minibatch(train_x, train_y, BATCH_SIZE)

    y_ = feed_forward!(x, ff_net)
    y_ = map(x->clamp.(x, 0.00001, 0.99999), y_)

    losses = [crossentropy(y_[i], y[i]) for i in 1:BATCH_SIZE]
    d_losses = [.-d_crossentropy(y_[i], y[i]) for i in 1:BATCH_SIZE]

    for derivs in d_losses
        for di in eachindex(derivs)
            if isnan(derivs[di])
                derivs[di] = 0
            end
        end
    end

    feed_backwards!(d_losses, fb_net, lr = LEARNING_RATE)

    train_metric[e] = sum(losses) / length(losses)
    println("episode: $e/$TRAIN_EPISODES\t---\tmean loss: $(sum(losses) / length(losses))")
end

# testing
test_metric = []; for i in 1:TEST_EPISODES

    x, y = get_MNIST_minibatch(train_x, train_y, BATCH_SIZE)

    y_ = feed_forward!(x, ff_net)
    y_ = map(x->clamp.(x, 0.00001, 0.99999), y_)

    # losses = [crossentropy(y_[i], y[i]) for i in 1:BATCH_SIZE]

    avg_pred = [y[i][argmax(y_[i])] for i in 1:BATCH_SIZE]
    append!(test_metric, [sum(avg_pred)/length(avg_pred)])

    println("testing: $i/$TEST_EPISODES")
end

plot(train_metric, label="mean squared error [scaled by a factor of $LOSS_SCALING] (batch size: $BATCH_SIZE, lr: $LEARNING_RATE)", xlabel="train episode")

println("\ntest prediction accuracy = $((sum(test_metric))/TEST_EPISODES)\n")


using FirstOrderOptimization

lr = LowRankOperator(rand(5), rand(4)')
lr*rand(4)
lr'*rand(5)

svds(lr, nsv=1)

pobs = .5
obs = sprand(5, 4, pobs).>0
iobs = find(obs)
io = IndexingOperator(5, 4, iobs)
io*rand(5,4)
io'*rand(length(iobs))

# svds(io, nsv=1)

ilr = IndexedLowRankOperator(io,rand(length(iobs)))
size(ilr)==(5,4)
ilr*rand(4)
ilr'*rand(5)

svds(ilr, nsv=1)

# test thin update

u,s,v = FirstOrderOptimization.svd(rand(4,5), 2)
A = LowRankOperator(u,spdiagm(s),v, transpose=(:N,:N,:T))
Delta = LowRankOperator(rand(4,1), rand(5,1), transpose=(:N,:T))
Aup = Array(A) + Array(Delta)
thin_update!(A,Delta,1.)
@assert norm(Array(A) - Aup) <= 1e-10

# test dot between sparse and low rank (svd) operator

m,n = size(A)
G = sprand(m,n,.3)
@assert norm(dot(G,A) - dot(full(G),Array(A))) <= 1e-10


export ∥

@doc raw"""
# Function call

`∥(z...)`

Type `\parallel` and hit the `tab` key to autocomplete the parallel symbol ∥

# Description

This calculates the parallel connection of impedances. Optionally, the
calculation can be performed using units from the module Unitful.

**Important note:** Use the ∥ operator in parentheses, since its precedence
is equal to the (low) precedence of relation operators or use the function call
with arguments to avoid miscalculations

# Variables

`z` Vector of impedances

# Examples

```julia
julia> using Unitful,Unitful.DefaultSymbols,ElectricalEngineering
julia> ∥(4,6)
2.4
julia> ∥(4Ω,6Ω)
2.4 Ω
julia> 2Ω+(4Ω∥6Ω) # Parentheses are ALWAYS required since
4.4 Ω
julia> 2Ω+4Ω∥6Ω # Is equal to (2Ω+4Ω)∥6Ω which is usually NOT intended
3.0 Ω
julia> (2.4Ω∥(4Ω∥6Ω))
1.2 Ω
julia> ∥(4Ω,6Ω,2.4Ω)
1.2 Ω

```

Type `\Omega` and hit the `tab` key to autocomplete the parallel symbol Ω
"""
function ∥(z...)
    ypar = 0.0/unit(z[1])
    for k=1:length(z)
        ypar = ypar + 1/z[k]
    end
    return 1/ypar
end


@testset "Polar Decomposition" begin

Random.seed!(1103)

# For each variant, decomposing A to U, H, verify:
# U is unitary
# H is hermitian and positive semi-definite
# backward error is small

for T in [Float64, Complex{Float64}]
    @testset "Newton $T" begin
        for n in [1,3,10]
            A = rand(T,n,n)
            U,H = polar(A, alg = :newton)
            @test U'*U ≈ Matrix(I,n,n) atol=1e-7
            @test ishermitian(H)
            for i in eigvals(H)
                @test i >= 0.
            end
            @test A ≈ U*H

            # also check the properties here (once should suffice)
            r = polar(A, alg = :newton)
            @test U == r.U
            @test H == r.H

            r = polar(A, alg = :newton, verbose = true)
            @test U == r.U
            @test H == r.H

            m = n + 2
            A = rand(T,m,n)
            U,H = polar(A, alg = :newton)
            @test U'*U ≈ Matrix(I,n,n) atol=1e-7
            @test ishermitian(H)
            for i in eigvals(H)
                @test i >= 0.
            end
            @test A ≈ U*H
        end
        A = rand(T,5,7)
        @test_throws ArgumentError U,H = polar(A, alg = :newton)
    end
end

for T in [Float64, Complex{Float64}]
    @testset "Halley $T" begin
        for n in [1,3,10]
            A = rand(T, n, n)
            r = polar(A, alg = :halley)
            U = r.U
            H = r.H
            @test U'*U ≈ Matrix(I,n,n) atol=1e-7
            @test ishermitian(H)
            for i in eigvals(H)
                @test i >= 0.
            end
            @test A ≈ U*H
        end
    end

    @testset "QDWH $T" begin
      for m in [1,3,10]
        B = rand(T, m, m)
        U2,H2 = polar(B, alg = :qdwh)
        @test U2'*U2 ≈ Matrix(I,m,m) atol=1e-7
        @test ishermitian(H2)
        for i in eigvals(H2)
            @test i >= 0.
        end
        @test B ≈ U2*H2
      end
    end
end

for T in [Float64, Complex{Float64}]
    @testset "SVD $T" begin
        for n in [1,3,10]
            m = n + rand(0:5)
            A = rand(T, m, n)
            U,H = polar(A, alg = :svd)
            @test U'*U ≈ Matrix(I,n,n) atol=1e-7
            @test ishermitian(H)
            for i in eigvals(H)
                @test i >= 0.
            end
            @test A ≈ U*H
        end
    end
end

for T in [Float64, Complex{Float64}]
    @testset "Newton-Schulz $T" begin
        for n in [1,3,10]
            m = n
            A = rand(T, m, n)
            nA = opnorm(A)
            A = (1.5 / nA) * A
            U,H = polar(A, alg = :schulz)
            @test U'*U ≈ Matrix(I,n,n) atol=1e-7
            @test ishermitian(H)
            for i in eigvals(H)
                @test i >= 0.
            end
            @test A ≈ U*H
        end
    end
end

for T in [Float64, Complex{Float64}]
    @testset "Hybrid $T" begin
        for n in [1,3,10]
            A = rand(T,n,n)
            U,H = polar(A, alg =:hybrid)
            @test U'*U ≈ Matrix(I,n,n) atol=1e-7
            @test ishermitian(H)
            for i in eigvals(H)
                @test i >= 0.
            end
            @test A ≈ U*H atol=1e-7
        end
    end
end

@testset "Float32" begin
    for n in [1,3,10]
        A = Array{Float32}(undef, n, n)
        copyto!(A, rand(n,n))

        r1 = polar(A, alg =:newton)
        r2 = polar(A, alg =:qdwh)
        r3 = polar(A, alg =:halley)
        r4 = polar(A, alg =:svd)
        r5 = polar(A, alg =:hybrid)

        @test r1.U ≈ r2.U
        @test r2.U ≈ r3.U
        @test r3.U ≈ r4.U
        @test r4.U ≈ r5.U
    end
end

end


"""
    gaussian_puff_factory(scenario)

Generates a gaussian dispersion model on the given scenario and returns a
function giving the concentration of the form
c(x, y, z, t)

"""

function gaussian_puff_factory(scenario)

    required_params = [:mass_emission_rate, :release_duration, :release_height,
                       :windspeed, :pasquill_gifford]
    if all(key -> !(ismissing(getproperty(scenario,key))), required_params)

        # parameters of the jet
        m = scenario.mass_emission_rate*scenario.release_duration
        h = scenario.release_height

        # parameters of the environment
        u = scenario.windspeed
        stability = scenario.pasquill_gifford

    else
        missing_params = [ String(i) for i in filter(key -> ismissing(getproperty(scenario,key)), required_params)]
        error_string = "These parameters cannot be missing: " * join(missing_params, ", ")
        e = MissingException(error_string)
        throw(e)
    end

    # Pasquill-Gifford dispersion
    σx = crosswind_dispersion(stability; plume=false)
    σy = σx
    σz = vertical_dispersion(stability; plume=false)

    function gaussian_puff(x,y,z,t)
        sx = σx(x)
        sy = σy(x)
        sz = σz(x)

        C1 = m / ( (2*π)^(1.5) * sx * sy * sz )
        C2 = exp(-0.5*((x-u*t)/sx)^2)
        C3 = exp(-0.5*(y/sy)^2)
        C4 = ( exp(-0.5*((z-h)/sz)^2) + exp(-0.5*((z+h)/sz)^2) )

        return C1*C2*C3*C4
    end

    return gaussian_puff

end


# copyright Lénaïc Chizat
# the code shows the training dynamics in parameter and predictor space for a 2-layer relu NN on a two class classification task
# we optimize the exponential loss (the behavior is sensibly the same with the logistic loss)
# we use a specific step-size schedule motivated by the theory in this ref https://arxiv.org/abs/2002.04486 (sections 4 & 5)

# HOW TO RUN THE CODE
# in a prompt first run:
# include("exp_output.jl")
# Then for fast prototyping, run for instance the following code (takes 1 min)
# illustration(4, 60, 200, 4, 100000, 10, 0.05);
# These parameters are explained in the header of the function "illustration"
# A good illustration is obtained as follows (takes 20 min):
# Random.seed!(9); # 9 generates a nice random problem
# illustration(4, 60, 400, 1, 3000000, 400, 0.005);

using LinearAlgebra, Random
using PyPlot, ProgressMeter # these packages need to be installed (via "] add NamePackage" )
using3D()

"""
Gradient descent to train a 2-layers ReLU neural network for the exponential loss
We use the step-size schedule from in https://arxiv.org/abs/2002.04486 (sections 4 & 5)
INPUT: X (training input), Y (training output), m (nb neurons), both: training
        both layers or just the output layer
OUTPUT: Ws (training trajectory)
"""
function twonet(X, Y, m, stepsize, niter)

    (n,d) = size(X) # n samples in R^d
    # initialize
    W = randn(m, d+1)
    # input weights initialized uniformly on the unit sphere
    W[:,1:d] = W[:,1:d] ./ sqrt.(sum(W[:,1:d].^2, dims=2))
    # output weights initialized by -1 or 1 with equal proba
    W[1:div(m,2),end] .= 0.0
    W[div(m,2)+1:end,end] .= 0.0

    Ws    = zeros(m, d+1, niter) # store optimization path
    loss  = zeros(niter) # loss is -log of the empirical risk
    margins = zeros(niter)
    betas = zeros(niter)

    @showprogress 1 "Training neural network..." for iter = 1:niter
        Ws[:,:,iter] = W
        act  =  max.( W[:,1:end-1] * X', 0.0) # activations
        out  =  (1/m) * sum( W[:,end] .* act , dims=1) # predictions of the network
        perf = Y .* out[:]
        margin = minimum(perf)
        temp = exp.(margin .- perf) # numerical stabilization of exp
        gradR = temp .* Y ./ sum(temp)' # gradient of the loss
        grad_w1 = zeros(m,d) #(W[:,end] .* float.(act .> 0) * ( X .* gradR  ))  # gradient for input weights
        grad_w2 = act * gradR  # gradient for output weights

        grad = cat(grad_w1, grad_w2, dims=2) # size (m × d+1)
        betas[iter] = max(1, sqrt(sum(W[:,end].^2)/m))
        loss[iter] = margin - log(sum(exp.(margin .- perf))/n)
        margins[iter] = margin/betas[iter] # margin in F_1 norm
        W = W + stepsize * grad * betas[iter]/(sqrt(iter+1))
    end

    return Ws, loss, margins, betas
end

##
"Coordinates of the 2d cluster centers, p is k^2 the number of clusters"
function cluster_center(p,k)
    p1 = mod.(p .- 1,k) .+ 1
    p2 = div.(p .- 1,k) .+ 1
    Delta = 1/(3k-1)
    x1 =  Delta*(1 .+ 3*(p1 .- 1)) .- 1/2
    x2 =  Delta*(1 .+ 3*(p2 .- 1)) .- 1/2
    return x1, x2
end

"""
Plot the classifier for a test case, comparing training both  layer
k: number of clusters per dimension in the data set (choose 3, 4 or 5)
n: number of training samples
m: number of neurons
stepsize, niter: parameters for training
nframes, resolution: parameters for the plots
"""
function illustration(k, n, m, stepsize, niter, nframes, resolution)
    # data distribution (itself random)
    sd = 0 # number of spurious dimensions with random noise
    Delta = 1/(3k-1) # interclass distance
    A = ones(k^2) # cluster affectation
    A[randperm(k^2)[1:div(k^2,2)]] .= -1

    # sample from this data distribution
    P = rand(1:k^2,n) # cluster label
    T = 2π*rand(n)  # shift angle
    R = Delta*rand(n) # shift magnitude
    X = cat(ones(n), cluster_center(P,k)[1] .+ R .* cos.(T),cluster_center(P,k)[2] + R .* sin.(T), (rand(n,sd) .- 1/2), dims=2)
    Y = A[P]

    # plot training set
    X1 = X[(Y .== 1),:]
    X2 = X[(Y .== -1),:]
    fig = figure(figsize=[3,3])
    plot(X1[:,2],X1[:,3],"+k")
    plot(X2[:,2],X2[:,3],"_k")
    axis("equal");axis("off");
    display(fig)

    # train the neural network
    Ws, loss, margins, betas = twonet(X, Y, m, stepsize, niter)

    # define the sequence of time steps ts to be plotted
    a = (niter-1)/(nframes-1)^4
    ts = setdiff(Int.(floor.(a*(0:nframes-1).^4)) .+ 1)
    Ws = Ws[:,:,ts]

    # these are the positions we plot for the particles
    Wproj = Ws[:,1:end-1,:] .* abs.(Ws[:,end:end,:])
    WN   = sqrt.(sum(Wproj.^2, dims = 2))
    Wdir = Wproj ./ WN
    Wlog = tanh.(0.5*WN) .* Wdir

    @showprogress 1 "Plotting images..." for k = 1:length(ts)
        ioff() # turns off interactive plotting
        fig = figure(figsize=[7,4])
        ax1 = subplot(121, projection="3d")
        ax1.set_position([0,0.1,0.5,0.8])

        if k<11
            indt = 1:k
        else
            indt = (k-10):k
        end

        for i = 1:size(Wlog,1)
            plot3D(Wlog[i,2,indt],Wlog[i,3,indt],Wlog[i,1,indt], color="k", linewidth=0.2) # tail
        end
        ind_p = (Ws[:,end,k] .>= 0.0)[:]
        ind_m = (Ws[:,end,k] .< 0.0)[:]
        plot3D(Wlog[ind_p,2,k],Wlog[ind_p,3,k],Wlog[ind_p,1,k],"o",color="C3", markersize=1)
        plot3D(Wlog[ind_m,2,k],Wlog[ind_m,3,k],Wlog[ind_m,1,k],"o",color="C0", markersize=1)

        ax1.set_xlim3d(-1, 1)
        ax1.set_ylim3d(-1, 1)
        ax1.set_zlim3d(-1, 1)
        ax1.set_xticks([-1/2, 0, 1/2])
        ax1.set_yticks([-1/2, 0, 1/2])
        ax1.set_zticks([-1/2, 0, 1/2])
        ax1.view_init(25-20*sinpi(k/length(ts)),45-45*cospi(k/length(ts)))

        ax2 = subplot(122)
        ax2.set_position([0.45,0.25,0.5,0.5])

        f(x1,x2,k) = (1/m) * sum( Ws[:,end,k] .* max.( Ws[:,1:3,k] * [1;x1;x2], 0.0)) # prediction function

        xs = -0.8:resolution:0.8
        tab = [f(xs[i],xs[j],k) for i=1:length(xs), j=1:length(xs)]
        pcolormesh(xs', xs, tanh.(tab'), cmap="coolwarm", shading="gouraud", vmin=-1.0, vmax=1.0, edgecolor="face")

        xs = -0.8:resolution:0.8
        tab = [f(xs[i],xs[j],k) for i=1:length(xs), j=1:length(xs)]
        contour(xs', xs, tanh.(tab'), levels =0, colors="k", antialiased = true, linewidths=2)

        # plot training set
        X1 = X[(Y .== 1),:]
        X2 = X[(Y .== -1),:]
        plot(X1[:,2],X1[:,3],"+k")
        plot(X2[:,2],X2[:,3],"_k")
        axis("equal");axis("off");
        ax2.set_xticks([-1/2, 0, 1/2])
        ax2.set_yticks([-1/2, 0, 1/2])

        savefig("dynamics_output_$(k).png",bbox_inches="tight", dpi=300)
        close(fig)
    end

end

##


# This file is part of IntegerSequences.
# Copyright Peter Luschny. License is MIT.

(@__DIR__) ∉ LOAD_PATH && push!(LOAD_PATH, (@__DIR__))

module DedekindEta
using Nemo

export ModuleDedekindEta
export DedekindEtaPowers, RamanujanTau, RamanujanTauList, PartitionNumberList
export L010815, L002107, L010816, L000727, L000728, L000729, L000730, L000731
export L010817, L010819, L000735, L010820, L010821, L010822, L000739, L010823
export L010824, L010825, L010826, L010827, L010828, L010829, L000594, L010830
export L010831, L010832, L010833, L010834, L010835, L010836, L010837, L010840
export L010838, L010839, L010841, L000041, L000712, L000716, L023003, L023004
export L023005, L023006, L023007, L023008, L023009, L023010, L005758, L023011
export L023012, L023013, L023014, L023015, L023016, L023017, L023018, L023019
export L023020, L023021, L006922, L082556, L082557, L082558, L082559

"""

* DedekindEtaPowers, RamanujanTau, RamanujanTauList, PartitionNumberList
"""
const ModuleDedekindEta = ""

"""

Compute the ``q``-expansion to length len of the Dedekind ``η`` function (without
 the leading factor ``q^{1/24}``) raised to the power ``r``, i.e.
``{(q^{-1/24} η(q))^r = ∏_{k ≥ 1} (1 - q^k)^r.}`` In particular, ``r = -1`` returns the generating function of the Partition function ``p(k)`` and ``r = 24`` gives the Ramanujan tau function ``τ(k)``.
"""
function DedekindEtaPowers(len::Int, r::Int)
    len ≤ 0 && return fmpz[]
    R, x = PolynomialRing(ZZ, "x")
    e = eta_qexp(r, len, x)
    [coeff(e, j) for j ∈ 0:len-1]
end

"""

Compute the expansion of ``∏ _{m≥1} (1 - q^m)``.
"""
L010815(len::Int) = DedekindEtaPowers(len, 1)

"""

Compute the expansion of ``∏ _{m≥1} (1 - q^m)^2``.
"""
L002107(len::Int) = DedekindEtaPowers(len, 2)

"""

Compute the expansion of ``∏ _{m≥1} (1 - q^m)^3``.
"""
L010816(len::Int) = DedekindEtaPowers(len, 3)

"""

Compute the expansion of ``∏ _{m≥1} (1 - q^m)^4``.
"""
L000727(len::Int) = DedekindEtaPowers(len, 4)

"""

Compute the expansion of ``∏ _{m≥1} (1 - q^m)^5``.
"""
L000728(len::Int) = DedekindEtaPowers(len, 5)

"""

Compute the expansion of ``∏ _{m≥1} (1 - q^m)^6``.
"""
L000729(len::Int) = DedekindEtaPowers(len, 6)

"""

Compute the expansion of ``∏ _{m≥1} (1 - q^m)^7``.
"""
L000730(len::Int) = DedekindEtaPowers(len, 7)

"""

Compute the expansion of ``∏ _{m≥1} (1 - q^m)^8``.
"""
L000731(len::Int) = DedekindEtaPowers(len, 8)

"""

Compute the expansion of ``∏ _{m≥1} (1 - q^m)^9``.
"""
L010817(len::Int) = DedekindEtaPowers(len, 9)

"""

Compute the expansion of ``∏_{m≥1} (1 - q^m)^{10}``.
"""
L010818(len::Int) = DedekindEtaPowers(len, 10)

"""

Compute the expansion of ``∏_{m≥1} (1 - q^m)^{11}``.
"""
L010819(len::Int) = DedekindEtaPowers(len, 11)

"""

Compute the expansion of ``∏_{m≥1} (1 - q^m)^{12}``.
"""
L000735(len::Int) = DedekindEtaPowers(len, 12)

"""

Compute the expansion of ``∏_{m≥1} (1 - q^m)^{13}``.
"""
L010820(len::Int) = DedekindEtaPowers(len, 13)

"""

Compute the expansion of ``∏_{m≥1} (1 - q^m)^{14}``.
"""
L010821(len::Int) = DedekindEtaPowers(len, 14)

"""

Compute the expansion of ``∏_{m≥1} (1 - q^m)^{15}``.
"""
L010822(len::Int) = DedekindEtaPowers(len, 15)

"""

Compute the expansion of ``∏_{m≥1} (1 - q^m)^{16}``.
"""
L000739(len::Int) = DedekindEtaPowers(len, 16)

"""

Compute the expansion of ``∏_{m≥1} (1 - q^m)^{17}``.
"""
L010823(len::Int) = DedekindEtaPowers(len, 17)

"""

Compute the expansion of ``∏_{m≥1} (1 - q^m)^{18}``.
"""
L010824(len::Int) = DedekindEtaPowers(len, 18)

"""

Compute the expansion of ``∏_{m≥1} (1 - q^m)^{19}``.
"""
L010825(len::Int) = DedekindEtaPowers(len, 19)

"""

Compute the expansion of ``∏_{m≥1} (1 - q^m)^{20}``.
"""
L010826(len::Int) = DedekindEtaPowers(len, 20)

"""

Compute the expansion of ``∏_{m≥1} (1 - q^m)^{21}``.
"""
L010827(len::Int) = DedekindEtaPowers(len, 21)

"""

Compute the expansion of ``∏_{m≥1} (1 - q^m)^{22}``.
"""
L010828(len::Int) = DedekindEtaPowers(len, 22)

"""

Compute the expansion of ``∏_{m≥1} (1 - q^m)^{23}``.
"""
L010829(len::Int) = DedekindEtaPowers(len, 23)

"""

Compute the expansion of ``∏_{m≥1} (1 - q^m)^{24}``.
"""
L000594(len::Int) = DedekindEtaPowers(len, 24)

"""

Compute the expansion of ``∏_{m≥1} (1 - q^m)^{25}``.
"""
L010830(len::Int) = DedekindEtaPowers(len, 25)

"""

Compute the expansion of ``∏_{m≥1} (1 - q^m)^{26}``.
"""
L010831(len::Int) = DedekindEtaPowers(len, 26)

"""

Compute the expansion of ``∏_{m≥1} (1 - q^m)^{27}``.
"""
L010832(len::Int) = DedekindEtaPowers(len, 27)

"""

Compute the expansion of ``∏_{m≥1} (1 - q^m)^{28}``.
"""
L010833(len::Int) = DedekindEtaPowers(len, 28)

"""

Compute the expansion of ``∏_{m≥1} (1 - q^m)^{29}``.
"""
L010834(len::Int) = DedekindEtaPowers(len, 29)

"""

Compute the expansion of ``∏_{m≥1} (1 - q^m)^{30}``.
"""
L010835(len::Int) = DedekindEtaPowers(len, 30)

"""

Compute the expansion of ``∏_{m≥1} (1 - q^m)^{31}``.
"""
L010836(len::Int) = DedekindEtaPowers(len, 31)

"""

Compute the expansion of ``∏_{m≥1} (1 - q^m)^{32}``.
"""
L010837(len::Int) = DedekindEtaPowers(len, 32)

"""

Compute the expansion of ``∏_{m≥1} (1 - q^m)^{40}``.
"""
L010840(len::Int) = DedekindEtaPowers(len, 40)

"""

Compute the expansion of ``∏_{m≥1} (1 - q^m)^{44}``.
"""
L010838(len::Int) = DedekindEtaPowers(len, 44)

"""

Compute the expansion of ``∏_{m≥1} (1 - q^m)^{48}``.
"""
L010839(len::Int) = DedekindEtaPowers(len, 48)

"""

Compute the expansion of ``∏_{m≥1} (1 - q^m)^{64}``.
"""
L010841(len::Int) = DedekindEtaPowers(len, 64)

"""

List of the first values of the Ramanujan tau function, the Fourier coefficients of the Weierstrass Delta-function.
"""
RamanujanTauList(len::Int) = DedekindEtaPowers(len, 24)

"""

Return Ramanujan's tau(n).
"""
RamanujanTau(n::Int) = DedekindEtaPowers(n, 24)[end]

"""

Return the first n numbers of integer partitions.
"""
PartitionNumberList(len::Int) = DedekindEtaPowers(len, -1)

"""

Return the first n numbers of integer partitions.
"""
L000041(len::Int) = DedekindEtaPowers(len, -1)

"""

Return the number of partitions of n into parts of 2 kinds.
"""
L000712(len::Int) = DedekindEtaPowers(len, -2)

"""

Return the number of partitions of n into parts of 3 kinds.
"""
L000716(len::Int) = DedekindEtaPowers(len, -3)

"""

Return the number of partitions of n into parts of 4 kinds.
"""
L023003(len::Int) = DedekindEtaPowers(len, -4)

"""

Return the number of partitions of n into parts of 5 kinds.
"""
L023004(len::Int) = DedekindEtaPowers(len, -5)

"""

Return the number of partitions of n into parts of 6 kinds.
"""
L023005(len::Int) = DedekindEtaPowers(len, -6)

"""

Return the number of partitions of n into parts of 7 kinds.
"""
L023006(len::Int) = DedekindEtaPowers(len, -7)

"""

Return the number of partitions of n into parts of 8 kinds.
"""
L023007(len::Int) = DedekindEtaPowers(len, -8)

"""

Return the number of partitions of n into parts of 9 kinds.
"""
L023008(len::Int) = DedekindEtaPowers(len, -9)

"""

Return the number of partitions of n into parts of 10 kinds.
"""
L023009(len::Int) = DedekindEtaPowers(len, -10)

"""

Return the number of partitions of n into parts of 11 kinds.
"""
L023010(len::Int) = DedekindEtaPowers(len, -11)

"""

Return the number of partitions of n into parts of 12 kinds.
"""
L005758(len::Int) = DedekindEtaPowers(len, -12)

"""

Return the number of partitions of n into parts of 13 kinds.
"""
L023011(len::Int) = DedekindEtaPowers(len, -13)

"""

Return the number of partitions of n into parts of 14 kinds.
"""
L023012(len::Int) = DedekindEtaPowers(len, -14)

"""

Return the number of partitions of n into parts of 15 kinds.
"""
L023013(len::Int) = DedekindEtaPowers(len, -15)

"""

Return the number of partitions of n into parts of 16 kinds.
"""
L023014(len::Int) = DedekindEtaPowers(len, -16)

"""

Return the number of partitions of n into parts of 17 kinds.
"""
L023015(len::Int) = DedekindEtaPowers(len, -17)

"""

Return the number of partitions of n into parts of 18 kinds.
"""
L023016(len::Int) = DedekindEtaPowers(len, -18)

"""

Return the number of partitions of n into parts of 19 kinds.
"""
L023017(len::Int) = DedekindEtaPowers(len, -19)

"""

Return the number of partitions of n into parts of 20 kinds.
"""
L023018(len::Int) = DedekindEtaPowers(len, -20)

"""

Return the number of partitions of n into parts of 21 kinds.
"""
L023019(len::Int) = DedekindEtaPowers(len, -21)

"""

Return the number of partitions of n into parts of 22 kinds.
"""
L023020(len::Int) = DedekindEtaPowers(len, -22)

"""

Return the number of partitions of n into parts of 23 kinds.
"""
L023021(len::Int) = DedekindEtaPowers(len, -23)

"""

Return the number of partitions of n into parts of 24 kinds.
"""
L006922(len::Int) = DedekindEtaPowers(len, -24)

"""

Return the number of partitions of n into parts of 30 kinds.
"""
L082556(len::Int) = DedekindEtaPowers(len, -30)

"""

Return the number of partitions of n into parts of 32 kinds.
"""
L082557(len::Int) = DedekindEtaPowers(len, -32)

"""

Return the number of partitions of n into parts of 48 kinds.
"""
L082558(len::Int) = DedekindEtaPowers(len, -48)

"""

Return the number of partitions of n into parts of 64 kinds.
"""
L082559(len::Int) = DedekindEtaPowers(len, -64)

#START-TEST-########################################################

using Test, SeqTests

function test()
    @testset "DedekindEta" begin

        @test DedekindEtaPowers(0, 1) == fmpz[]
        @test isa(DedekindEtaPowers(30, 1)[10], fmpz)
        @test isa(DedekindEtaPowers(30, -1)[10], fmpz)

        @test RamanujanTau(20) == -7109760
        P0 = [1, 1, 2, 3, 5, 7, 11, 15, 22, 30, 42, 56, 77, 101]
        P1 = PartitionNumberList(14)
        @test all(P0 .== P1)

        if data_installed()

            L = [
                L010815,
                L002107,
                L010816,
                L000727,
                L000728,
                L000729,
                L000730,
                L000731,
                L010817,
                L010819,
                L000735,
                L010820,
                L010821,
                L010822,
                L000739,
                L010823,
                L010824,
                L010825,
                L010826,
                L010827,
                L010828,
                L010829,
                L000594,
                L010830,
                L010831,
                L010832,
                L010833,
                L010834,
                L010835,
                L010836,
                L010837,
                L010840,
                L010838,
                L010839,
                L010841,
                L000041,
                L000712,
                L000716,
                L023003,
                L023004,
                L023005,
                L023006,
                L023007,
                L023008,
                L023009,
                L023010,
                L005758,
                L023011,
                L023012,
                L023013,
                L023014,
                L023015,
                L023016,
                L023017,
                L023018,
                L023019,
                L023020,
                L023021,
                L006922,
                L082556,
                L082557,
                L082558,
                L082559
            ]

            SeqTest(L, 'L')
        end
    end
end

function demo()
    for n ∈ 0:6
        println(n, ": ", DedekindEtaPowers(8, n))
    end

    for n ∈ 0:6
        println(-n, ": ", DedekindEtaPowers(8, -n))
    end
end

"""

PartitionNumberList(10000)
    0.071317 seconds (10.01 k allocations: 234.813 KiB)
L000731(10000)
    0.002731 seconds (10.01 k allocations: 234.813 KiB)
RamanujanTauList(10000)
    0.012097 seconds (10.01 k allocations: 234.813 KiB)
"""
function perf()
    @time PartitionNumberList(10000)
    @time L000731(10000)
    @time RamanujanTauList(10000)
end

function main()
    test()
    demo()
    perf()
end

main()

end # module


# Please run from the `basic_gppp` directory.
using Pkg
Pkg.activate(@__DIR__)
Pkg.instantiate()

using AbstractGPs, Plots, Random, Stheno

###########################  Define our model  ###########################

# Define a distribution over f₁, f₂, and f₃, where f₃(x) = f₁(x) + f₂(x).
f = @gppp let
    f1 = GP(randn(), SEKernel())
    f2 = GP(SEKernel())
    f3 = f1 + f2
end;

# Randomly sample `N₁` and `N₃` locations at which to observe `f₁` and `f₃` respectively.
rng, N1, N3 = MersenneTwister(123546), 10, 11;
x1 = GPPPInput(:f1, sort(rand(rng, N1) * 10));
x3 = GPPPInput(:f3, sort(rand(rng, N3) * 10));
x = BlockData(x1, x3);

# Generate some toy of `f1` and `f3`, `y1` and `y3` respectively.
fx = f(x);
y = rand(rng, f(x));
y1, y3 = split(x, y);

# Compute the posterior processes.
f_post = posterior(fx, y);

# Define some plotting stuff.
Np, S = 500, 25;
xp_ = range(-2.5, stop=12.5, length=Np);
xp = BlockData(GPPPInput(:f1, xp_), GPPPInput(:f2, xp_), GPPPInput(:f3, xp_));

# Sample jointly from the posterior over each process.
f_samples = rand(rng, f_post(xp, 1e-9), S);
f′1_xp, f′2_xp, f′3_xp = split(xp, f_samples);

# Compute posterior marginals.
ms = marginals(f_post(xp, 1e-9));
f′1_m, f′2_m, f′3_m = split(xp, mean.(ms));
f′1_s, f′2_s, f′3_s = split(xp, std.(ms));

###########################  Plot results  ###########################

gr();
posterior_plot = plot();

# Plot posterior over f1.
plot!(posterior_plot, xp_, f′1_m; ribbon=3f′1_s, color=:red, label="f1", fillalpha=0.3);
plot!(posterior_plot, xp_, f′1_xp; color=:red, label="", alpha=0.2, linewidth=1);
scatter!(posterior_plot, x1.x, y1;
    markercolor=:red,
    markershape=:circle,
    markerstrokewidth=0.0,
    markersize=4,
    markeralpha=0.7,
    label="",
);

# Plot posterior over f2.
plot!(posterior_plot, xp_, f′2_m; ribbon=3f′2_s, color=:green, label="f2", fillalpha=0.3);
plot!(posterior_plot, xp_, f′2_xp; color=:green, label="", alpha=0.2, linewidth=1);

# Plot posterior over f3
plot!(posterior_plot, xp_, f′3_m; ribbon=3f′3_s, color=:blue, label="f3", fillalpha=0.3);
plot!(posterior_plot, xp_, f′3_xp; color=:blue, label="", alpha=0.2, linewidth=1);
scatter!(posterior_plot, x3.x, y3;
    markercolor=:blue,
    markershape=:circle,
    markerstrokewidth=0.0,
    markersize=4,
    markeralpha=0.7,
    label="",
);

display(posterior_plot);


using Inflation
using SymPy

using Random
using Roots
using ForwardDiff
using LinearAlgebra

d = 2

# sample masses from the Marchenko-Pastur distribution
m = 1e-5
β = 1/2
a = m^2 * (1-sqrt(β))^2
b = m^2 * (1+sqrt(β))^2
function MarchenkoPasturPDF(x)
    return sqrt( (b-x)*(x-a)) / (2*pi*x*β*m^2)
end
support = (a,b)

function rejection_sample(pdf,num,support=(-Inf,Inf))
    # for now assume pdf has exactly one critical point
    # will probably fail without finite support
    a = support[1]
    b = support[2]
    dpdf(x) = ForwardDiff.derivative(pdf,x)
    xmax = find_zero(dpdf,(a,b))
    pdf_max = pdf(xmax)
    samples = Array{Float64}(undef,num)
    i = 1
    while i <= num
        x = a+(b-a)*rand()
        s = rand()*pdf_max
        if s < pdf(x)
            samples[i] = x
            i += 1
        end
    end
    return samples
end

masses = rejection_sample(MarchenkoPasturPDF,d,support)

params = [symbols("m$i") for i in 1:d]
pvalues = masses

function G(Phi,p)
    d = length(Phi)
    g = Matrix{Float64}(I,d,d)
    return g
end

function V(Phi,p)
    m = p[1:end]
    v = sum( m .* Phi.^2 ) / 2
    return v
end

# pick an initial radius to give N0 e-folds
N0 = 100
# get a vector on the sphere, use gaussian trick
Phi0 = randn(d)
Phi0 ./= norm(Phi0)
Phi0 .*= sqrt(4*N0)

Pi0 = zeros(d)

funcs = inflation_setup(d,V,G,params)

sol = background_evolve(Phi0,Pi0,pvalues,funcs,verbose=true)

tsol = transport_perturbations(sol,pvalues,funcs,verbose=true)

output_data("Nflation_test.jld2",funcs,pvalues,sol,tsol)


using MortalityTables
using Test

@testset "basic `Ultimate` MortalityTable" begin
    tbl = MortalityTable([i/100 for i in 1:100],0)
    @test q(tbl,0,1) ≈ 0.01
    @test q(tbl,0,2) ≈ 1.0 - (1.0 - 0.01) * (1.0 - 0.02)
    @test q(tbl,0) ≈ 0.01

    @test p(tbl,0) ≈ 0.99
    @test p(tbl,0,1) ≈ 0.99
    @test p(tbl,0,2) ≈ 1.0 - q(tbl,0,2)
    @test q(tbl,0,1,1) ≈ 0.02
    @test p(tbl,0,1,1) ≈ 0.98
    @test_throws BoundsError q(tbl,200)
end

@testset "basic `Ultimate` MortalityTable with non-zero start age" begin
    tbl = MortalityTable([i/100 for i in 5:100],5)
    @test ismissing(q(tbl,0))
    @test q(tbl,5) ≈ 0.05
    @test q(tbl,5,1) ≈ 0.05
    @test q(tbl,5,2) ≈ 1.0 - (1.0 - 0.05) * (1.0 - 0.06)

    @test p(tbl,5) ≈ 0.95
    @test p(tbl,5,1) ≈ 0.95
    @test p(tbl,5,2) ≈ 1.0 - q(tbl,5,2)

    @test q(tbl,5,1,1) ≈ 0.06
    @test p(tbl,5,1,1) ≈ 0.94
end

@testset "test table names" begin
    tbl = MortalityTable([i/100 for i in 1:100],0,TableMetaData(name="test"))
    @test q(tbl,0,1) ≈ 0.01
    @test tbl.d.name == "test"
end

@testset "SOA tables" begin

    tables = MortalityTables.tables()
    cso2001 = tables["2001 CSO Super Preferred Select and Ultimate - Male Nonsmoker, ANB"]
    vbt2001 = tables["2001 VBT Residual Standard Select and Ultimate - Male Nonsmoker, ANB"]
    cso1980 = tables["1980 CSO Basic Table – Male, ANB"]

    @test q(cso1980,35,0,1) ≈ .00118
    @test qx(cso1980,35,1) ≈ .00118
    @test qx(cso1980,35,61) ≈ .27302
    @test qx(cso1980,95) ≈ .27302
    @test ismissing(qx(cso1980,35,95))
    @test ismissing(qx(cso1980,101))

    @test qx(cso2001,35,1) ≈ .00037
    @test qx(cso2001,35,61) ≈ .26719
    @test qx(cso2001,16) ≈ .00041
    @test qx(cso2001,95) ≈ .26719
    @test ismissing(qx(cso2001,15))
    @test_throws BoundsError qx(cso2001,150)
    @test ismissing(qx(cso2001,35,95))

    @test qx(vbt2001,35,1) ≈ .00036
    @test qx(vbt2001,35,61) ≈ .24298
    @test qx(vbt2001,95) ≈ .24298
    @test ismissing(qx(vbt2001,35,95))
    @test_throws BoundsError qx(vbt2001,150)

    @test cso2001[29,1:3] == [0.00029, 0.00035, 0.0004]

end


function Log_LikeliHood(m::AbstractModel, om, leaves)
    if leaves
        #= if the animal has left we want the probability of this to happen
        at that omission, given that it has not left before.
        Since we assume that the events are independent the product of each likelihood
        should give me the total probability. Using the log we can simply sum them =#
        return log(cdf(m,om) - cdf(m,om-1))
    else
        #= if the animal hasn't left we want the probability of not leaving until om =#
        # this is equivalent log(1-cdf(dist,om)) but it's safe when cdf is near 1
        return log(1-cdf(m, om))
    end
end

function  nll_data(m::AbstractModel,cm::Dict)
    -sum(v*Log_LikeliHood(m, k...) for (k, v) in cm)
end

function Distributions.fit(::Type{T}, b::DataFrames.AbstractDataFrame) where T<:AbstractModel
    cm = countmap(StructArray((Omissions_plus_one = b.Omissions_plus_one, Leave = b.Leave)))
    p = init(T, cm)
    res = optimize(params(p)) do param
        all(param .> 1e-10) || return 1e10
        params(p) .= param
        return nll_data(p, cm)
    end
    params(p) .= Optim.minimizer(res)
    return p
end

###
function simulate(m::AbstractModel,n_samples=1)
    w = [cdf(m,x)-cdf(m,x-1) for x in 1:maximum(m)]
    weights = StatsBase.weights(w)
    [sample(1:maximum(m),weights) for _ in 1:n_samples]
end


#############################################################################
# JuMP
# An algebraic modelling langauge for Julia
# See http://github.com/JuliaOpt/JuMP.jl
#############################################################################
# speed2.jl
#
# Runs some JuMP benchmarks to test for speed-related regressions.
# Test time to build model in memory.
#############################################################################

using JuMP

function pMedian(numFacility::Int,numCustomer::Int,numLocation::Int)
    srand(10)
    customerLocations = [rand(1:numLocation) for a = 1:numCustomer ]

    tic()
    m = Model()

    # Facility locations
    @variable(m, 0 <= s[1:numLocation] <= 1)

    # Aux. variable: x_a,i = 1 iff the closest facility to a is at i
    @variable(m, 0 <= x[1:numLocation,1:numCustomer] <= 1)

    # Objective: min distance
    @objective(m, Max, sum{abs(customerLocations[a]-i)*x[i,a], a = 1:numCustomer, i = 1:numLocation} )

    # Constraints
    for a in 1:numCustomer
        # Subject to linking x with s
        for i in 1:numLocation
            @constraint(m, x[i,a] - s[i] <= 0)
        end
        # Subject to one of x must be 1
        @constraint(m, sum{x[i,a],i=1:numLocation} == 1 )
    end

    # Subject to must allocate all facilities
    @constraint(m, sum{s[i],i=1:numLocation} == numFacility )
    buildTime = toq()

    tic()
    JuMP.build(m)
    writeTime = toq()

    return buildTime, writeTime
end

function cont5(n)
    m = n
    n1 = n-1
    m1 = m-1
    dx = 1/n
    T = 1.58
    dt = T/m
    h2 = dx^2
    a = 0.001
    yt = [0.5*(1 - (j*dx)^2) for j=0:n]

    tic()
    mod = Model()
    @variable(mod,  0 <= y[0:m,0:n] <= 1)
    @variable(mod, -1 <= u[1:m] <= 1)
    @objective(mod, Min, 0.25*dx*( (y[m,0] - yt[1])^2 +
       2*sum{ (y[m,j]-yt[j+1])^2, j=1:n1} + (y[m,n]-yt[n+1])^2) +
       0.25*a*dt*(2*sum{u[i]^2,i=1:m1} + u[m]^2))

    # PDE
    for i = 0:m1
        for j = 1:n1
            @constraint(mod, h2*(y[i+1,j] - y[i,j]) == 0.5*dt*(y[i,j-1] - 2*y[i,j] + y[i,j+1] + y[i+1,j-1] - 2*y[i+1,j] + y[i+1,j+1]) )
        end
    end

    # IC
    for j = 0:n
        @constraint(mod, y[0,j] == 0)
    end

    # BC
    for i = 1:m
        @constraint(mod, y[i,2]   - 4*y[i,1]  + 3*y[i,0] == 0)
        @constraint(mod, y[i,n-2] - 4*y[i,n1] + 3*y[i,n] == (2*dx)*(u[i] - y[i,n]))
    end
    buildTime = toq()

    tic()
    JuMP.build(mod)
    writeTime = toq()

    return buildTime, writeTime
end

function RunTests()
    # Pmedian
    pmedian_build = Float64[]
    pmedian_write = Float64[]
    for runs = 1:9
        bt, wt = pMedian(100,100,5000)
        push!(pmedian_build, bt)
        push!(pmedian_write, wt)
    end
    sort!(pmedian_build)
    sort!(pmedian_write)
    print("PMEDIAN BUILD MIN=",minimum(pmedian_build),"  MED=",pmedian_build[5],"\n")
    print("PMEDIAN INTRN MIN=",minimum(pmedian_write),"  MED=",pmedian_write[5],"\n")

    # Cont5
    cont5_build = Float64[]
    cont5_write = Float64[]
    for runs = 1:9
        bt, wt = cont5(500)
        push!(cont5_build, bt)
        push!(cont5_write, wt)
    end
    sort!(cont5_build)
    sort!(cont5_write)
    print("CONT5 BUILD   MIN=",minimum(cont5_build),"  MED=",cont5_build[5],"\n")
    print("CONT5 INTRN   MIN=",minimum(cont5_write),"  MED=",cont5_write[5],"\n")

end

RunTests()


using Distributed, StatsBase, Distributions, PyPlot
using TransitionalMCMC

lb  = -15
ub  = 15

fT(x) = logpdf(Uniform(lb,ub), x[1]) .+ logpdf(Uniform(lb,ub), x[2])
sample_fT(Nsamples) = rand(Uniform(lb,ub), Nsamples, 2)

log_fD_T(x) = log.(pdf(MvNormal([0,0],[1 -0.5; -0.5 1]), x) + pdf(MvNormal([5,5],[1 0.5; 0.5 1]), x) + pdf(MvNormal([-5,5],[1 0.9; 0.9 1]), x))

samps, acc =tmcmc(log_fD_T, fT, sample_fT, 2000)

plt.scatter(samps[:,1], samps[:,2])


# MIP01 from CPLEX.jl
function int1test(model::MOI.ModelLike, config::TestConfig)
    atol = config.atol
    rtol = config.rtol
    # an example on mixed integer programming
    #
    #   maximize 1.1x + 2 y + 5 z
    #
    #   s.t.  x + y + z <= 10
    #         x + 2 y + z <= 15
    #
    #         x is continuous: 0 <= x <= 5
    #         y is integer: 0 <= y <= 10
    #         z is binary

    @test MOIU.supports_default_copy_to(model, #=copy_names=# false)
    @test MOI.supports(model, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}())
    @test MOI.supports(model, MOI.ObjectiveSense())
    @test MOI.supports_constraint(model, MOI.ScalarAffineFunction{Float64}, MOI.LessThan{Float64})
    @test MOI.supports_constraint(model, MOI.SingleVariable, MOI.GreaterThan{Float64})
    @test MOI.supports_constraint(model, MOI.SingleVariable, MOI.ZeroOne)
    @test MOI.supports_constraint(model, MOI.SingleVariable, MOI.Integer)

    MOI.empty!(model)
    @test MOI.is_empty(model)

    v = MOI.add_variables(model, 3)
    @test MOI.get(model, MOI.NumberOfVariables()) == 3

    cf = MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1.0,1.0,1.0], v), 0.0)
    c = MOI.add_constraint(model, cf, MOI.LessThan(10.0))
    @test MOI.get(model, MOI.NumberOfConstraints{MOI.ScalarAffineFunction{Float64},MOI.LessThan{Float64}}()) == 1

    cf2 = MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1.0,2.0,1.0], v), 0.0)
    c2 = MOI.add_constraint(model, cf2, MOI.LessThan(15.0))
    @test MOI.get(model, MOI.NumberOfConstraints{MOI.ScalarAffineFunction{Float64},MOI.LessThan{Float64}}()) == 2

    vc1 = MOI.add_constraint(model, MOI.SingleVariable(v[1]), MOI.Interval(0.0, 5.0))
    # We test this after the creation of every `SingleVariable` constraint
    # to ensure a good coverage of corner cases.
    @test vc1.value == v[1].value
    @test MOI.get(model, MOI.NumberOfConstraints{MOI.SingleVariable,MOI.Interval{Float64}}()) == 1

    vc2 = MOI.add_constraint(model, MOI.SingleVariable(v[2]), MOI.Interval(0.0, 10.0))
    @test vc2.value == v[2].value
    @test MOI.get(model, MOI.NumberOfConstraints{MOI.SingleVariable,MOI.Interval{Float64}}()) == 2
    vc3 = MOI.add_constraint(model, MOI.SingleVariable(v[2]), MOI.Integer())
    @test vc3.value == v[2].value
    @test MOI.get(model, MOI.NumberOfConstraints{MOI.SingleVariable,MOI.Integer}()) == 1

    vc4 = MOI.add_constraint(model, MOI.SingleVariable(v[3]), MOI.ZeroOne())
    @test vc4.value == v[3].value
    @test MOI.get(model, MOI.NumberOfConstraints{MOI.SingleVariable,MOI.ZeroOne}()) == 1

    objf = MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1.1, 2.0, 5.0], v), 0.0)
    MOI.set(model, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(), objf)
    MOI.set(model, MOI.ObjectiveSense(), MOI.MAX_SENSE)

    @test MOI.get(model, MOI.ObjectiveSense()) == MOI.MAX_SENSE

    if config.solve
        @test MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMIZE_NOT_CALLED

        MOI.optimize!(model)

        @test MOI.get(model, MOI.TerminationStatus()) == config.optimal_status

        @test MOI.get(model, MOI.ResultCount()) >= 1

        @test MOI.get(model, MOI.PrimalStatus()) in [ MOI.FEASIBLE_POINT, MOI.NEARLY_FEASIBLE_POINT ]

        @test MOI.get(model, MOI.ObjectiveValue()) ≈ 19.4 atol=atol rtol=rtol

        @test MOI.get(model, MOI.VariablePrimal(), v) ≈ [4,5,1] atol=atol rtol=rtol

        @test MOI.get(model, MOI.ConstraintPrimal(), c) ≈ 10 atol=atol rtol=rtol

        @test MOI.get(model, MOI.ConstraintPrimal(), c2) ≈ 15 atol=atol rtol=rtol

        @test MOI.get(model, MOI.ObjectiveBound()) >= 19.4 - atol
        # FIXME the following are currently not implemented in MockOptimizer
#        @test MOI.get(model, MOI.RelativeGap()) >= 0.0
#        @test MOI.get(model, MOI.SolveTime()) >= 0.0
#        @test MOI.get(model, MOI.SimplexIterations()) >= 0
#        @test MOI.get(model, MOI.BarrierIterations()) >= 0
#        @test MOI.get(model, MOI.NodeCount()) >= 0
    end
end

# sos from CPLEX.jl" begin
function int2test(model::MOI.ModelLike, config::TestConfig)
    atol = config.atol
    rtol = config.rtol
    @testset "SOSI" begin
        @test MOIU.supports_default_copy_to(model, #=copy_names=# false)
        @test MOI.supports(model, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}())
        @test MOI.supports(model, MOI.ObjectiveSense())
        @test MOI.supports_constraint(model, MOI.VectorOfVariables, MOI.SOS1{Float64})
        @test MOI.supports_constraint(model, MOI.SingleVariable, MOI.LessThan{Float64})

        MOI.empty!(model)
        @test MOI.is_empty(model)

        v = MOI.add_variables(model, 3)
        @test MOI.get(model, MOI.NumberOfVariables()) == 3
        vc1 = MOI.add_constraint(model, MOI.SingleVariable(v[1]), MOI.LessThan(1.0))
        @test vc1.value == v[1].value
        vc2 = MOI.add_constraint(model, MOI.SingleVariable(v[2]), MOI.LessThan(1.0))
        @test vc2.value == v[2].value
        vc3 = MOI.add_constraint(model, MOI.SingleVariable(v[3]), MOI.LessThan(2.0))
        @test vc3.value == v[3].value

        c1 = MOI.add_constraint(model, MOI.VectorOfVariables([v[1], v[2]]), MOI.SOS1([1.0, 2.0]))
        c2 = MOI.add_constraint(model, MOI.VectorOfVariables([v[1], v[3]]), MOI.SOS1([1.0, 2.0]))
        @test MOI.get(model, MOI.NumberOfConstraints{MOI.VectorOfVariables,MOI.SOS1{Float64}}()) == 2

        #=
            To allow for permutations in the sets and variable vectors
            we're going to sort according to the weights
        =#
        cs_sos = MOI.get(model, MOI.ConstraintSet(), c2)
        cf_sos = MOI.get(model, MOI.ConstraintFunction(), c2)
        p = sortperm(cs_sos.weights)
        @test cs_sos.weights[p] ≈ [1.0, 2.0] atol=atol rtol=rtol
        @test cf_sos.variables[p] == v[[1,3]]

        objf = MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([2.0, 1.0, 1.0], v), 0.0)
        MOI.set(model, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(), objf)
        MOI.set(model, MOI.ObjectiveSense(), MOI.MAX_SENSE)
        @test MOI.get(model, MOI.ObjectiveSense()) == MOI.MAX_SENSE

        if config.solve
            @test MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMIZE_NOT_CALLED

            MOI.optimize!(model)

            @test MOI.get(model, MOI.TerminationStatus()) == config.optimal_status

            @test MOI.get(model, MOI.ResultCount()) >= 1

            @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT

            @test MOI.get(model, MOI.ObjectiveValue()) ≈ 3 atol=atol rtol=rtol

            @test MOI.get(model, MOI.VariablePrimal(), v) ≈ [0,1,2] atol=atol rtol=rtol
        end

        MOI.delete(model, c1)
        MOI.delete(model, c2)

        if config.solve
            MOI.optimize!(model)

            @test MOI.get(model, MOI.TerminationStatus()) == config.optimal_status

            @test MOI.get(model, MOI.ResultCount()) >= 1

            @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT

            @test MOI.get(model, MOI.ObjectiveValue()) ≈ 5 atol=atol rtol=rtol

            @test MOI.get(model, MOI.VariablePrimal(), v) ≈ [1,1,2] atol=atol rtol=rtol
        end
    end
    @testset "SOSII" begin
        @test MOIU.supports_default_copy_to(model, #=copy_names=# false)
        @test MOI.supports(model, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}())
        @test MOI.supports(model, MOI.ObjectiveSense())
        @test MOI.supports_constraint(model, MOI.VectorOfVariables, MOI.SOS1{Float64})
        @test MOI.supports_constraint(model, MOI.VectorOfVariables, MOI.SOS2{Float64})
        @test MOI.supports_constraint(model, MOI.SingleVariable, MOI.ZeroOne)
        @test MOI.supports_constraint(model, MOI.ScalarAffineFunction{Float64}, MOI.EqualTo{Float64})

        MOI.empty!(model)
        @test MOI.is_empty(model)

        v = MOI.add_variables(model, 10)
        @test MOI.get(model, MOI.NumberOfVariables()) == 10

        bin_constraints = []
        for i in 1:8
            vc = MOI.add_constraint(model, MOI.SingleVariable(v[i]), MOI.Interval(0.0, 2.0))
            @test vc.value == v[i].value
            push!(bin_constraints, MOI.add_constraint(model, MOI.SingleVariable(v[i]), MOI.ZeroOne()))
            @test bin_constraints[i].value == v[i].value
        end

        MOI.add_constraint(model,
            MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1.0,2.0,3.0,-1.0], v[[1,2,3,9]]), 0.0),
            MOI.EqualTo(0.0)
        )

        MOI.add_constraint(model,
            MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([5.0,4.0,7.0,2.0,1.0,-1.0], v[[4,5,6,7,8,10]]), 0.0),
            MOI.EqualTo(0.0)
        )

        MOI.add_constraint(model,
            MOI.VectorOfVariables(v[[1, 2, 3]]),
            MOI.SOS1([1.0, 2.0, 3.0])
        )

        vv   = MOI.VectorOfVariables(v[[4,5,6,7,8]])
        sos2 = MOI.SOS2([5.0, 4.0, 7.0, 2.0, 1.0])
        c = MOI.add_constraint(model, vv, sos2)

        #=
            To allow for permutations in the sets and variable vectors
            we're going to sort according to the weights
        =#
        cs_sos = MOI.get(model, MOI.ConstraintSet(), c)
        cf_sos = MOI.get(model, MOI.ConstraintFunction(), c)
        p = sortperm(cs_sos.weights)
        @test cs_sos.weights[p] ≈ [1.0, 2.0, 4.0, 5.0, 7.0] atol=atol rtol=rtol
        @test cf_sos.variables[p] == v[[8,7,5,4,6]]

        objf = MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1.0, 1.0], [v[9], v[10]]), 0.0)
        MOI.set(model, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(), objf)
        MOI.set(model, MOI.ObjectiveSense(), MOI.MAX_SENSE)
        @test MOI.get(model, MOI.ObjectiveSense()) == MOI.MAX_SENSE

        if config.solve
            @test MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMIZE_NOT_CALLED

            MOI.optimize!(model)

            @test MOI.get(model, MOI.TerminationStatus()) == config.optimal_status

            @test MOI.get(model, MOI.ResultCount()) >= 1

            @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT

            @test MOI.get(model, MOI.ObjectiveValue()) ≈ 15.0 atol=atol rtol=rtol

            @test MOI.get(model, MOI.VariablePrimal(), v) ≈ [0.0, 0.0, 1.0, 1.0, 0.0, 1.0, 0.0, 0.0, 3.0, 12.0] atol=atol rtol=rtol
        end

        for cref in bin_constraints
            MOI.delete(model, cref)
        end

        if config.solve
            MOI.optimize!(model)

            @test MOI.get(model, MOI.TerminationStatus()) == config.optimal_status

            @test MOI.get(model, MOI.ResultCount()) >= 1

            @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT

            @test MOI.get(model, MOI.ObjectiveValue()) ≈ 30.0 atol=atol rtol=rtol

            @test MOI.get(model, MOI.VariablePrimal(), v) ≈ [0.0, 0.0, 2.0, 2.0, 0.0, 2.0, 0.0, 0.0, 6.0, 24.0] atol=atol rtol=rtol
        end
    end
end

# CPLEX #76
function int3test(model::MOI.ModelLike, config::TestConfig)
    atol = config.atol
    rtol = config.rtol
    # integer knapsack problem
    # max   z - 0.5 ( b1 + b2 + b3) / 40
    # s.t.  0 <= z - 0.5 eᵀ b / 40 <= 0.999
    #       b1, b2, ... b10 ∈ {0, 1}
    #       z in {0, 1, 2, ..., 100}

    MOI.empty!(model)
    @test MOI.is_empty(model)

    @test MOIU.supports_default_copy_to(model, #=copy_names=# false)
    @test MOI.supports(model, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}())
    @test MOI.supports(model, MOI.ObjectiveSense())
    @test MOI.supports_constraint(model, MOI.SingleVariable, MOI.ZeroOne)
    @test MOI.supports_constraint(model, MOI.SingleVariable, MOI.Integer)
    @test MOI.supports_constraint(model, MOI.SingleVariable, MOI.Interval{Float64})
    @test MOI.supports_constraint(model, MOI.ScalarAffineFunction{Float64}, MOI.Interval{Float64})

    z = MOI.add_variable(model)
    vc1 = MOI.add_constraint(model, MOI.SingleVariable(z), MOI.Integer())
    @test vc1.value == z.value
    vc2 = MOI.add_constraint(model, MOI.SingleVariable(z), MOI.Interval(0.0, 100.0))
    @test vc2.value == z.value

    b = MOI.add_variables(model, 10)

    for bi in b
        vc = MOI.add_constraint(model, MOI.SingleVariable(bi), MOI.ZeroOne())
        @test vc.value == bi.value
    end

    c = MOI.add_constraint(model, MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(vcat(1.0, fill(-0.5 / 40, 10)), vcat(z, b)), 0.0), MOI.Interval(0.0, 0.999))

    MOI.set(model, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(), MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(vcat(1.0, fill(-0.5 / 40, 3)), vcat(z, b[1:3])), 0.0))
    MOI.set(model, MOI.ObjectiveSense(), MOI.MAX_SENSE)

    if config.solve
        @test MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMIZE_NOT_CALLED

        MOI.optimize!(model)

        @test MOI.get(model, MOI.TerminationStatus()) == config.optimal_status

        @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT

        @test MOI.get(model, MOI.ObjectiveValue()) ≈ 1 atol=atol rtol=rtol

        # test for CPLEX.jl #76
        MOI.optimize!(model)

        @test MOI.get(model, MOI.TerminationStatus()) == config.optimal_status

        @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT

        @test MOI.get(model, MOI.ObjectiveValue()) ≈ 1 atol=atol rtol=rtol
    end
end

# Mixed-integer linear problems

function knapsacktest(model::MOI.ModelLike, config::TestConfig)
    atol = config.atol
    rtol = config.rtol
    # integer knapsack problem
    # max 5a + 3b + 2c + 7d + 4e
    # st  2a + 8b + 4c + 2d + 5e <= 10
    #                  a,b,c,d,e ∈ binary

    MOI.empty!(model)
    @test MOI.is_empty(model)

    @test MOIU.supports_default_copy_to(model, #=copy_names=# false)
    @test MOI.supports(model, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}())
    @test MOI.supports(model, MOI.ObjectiveSense())
    @test MOI.supports_constraint(model, MOI.SingleVariable, MOI.ZeroOne)
    @test MOI.supports_constraint(model, MOI.ScalarAffineFunction{Float64}, MOI.LessThan{Float64})

    v = MOI.add_variables(model, 5)
    @test MOI.get(model, MOI.NumberOfVariables()) == 5

    for vi in v
        vc = MOI.add_constraint(model, MOI.SingleVariable(vi), MOI.ZeroOne())
        @test vc.value == vi.value
    end
    @test MOI.get(model, MOI.NumberOfConstraints{MOI.SingleVariable,MOI.ZeroOne}()) == 5
    c = MOI.add_constraint(model, MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([2.0, 8.0, 4.0, 2.0, 5.0], v), 0.0), MOI.LessThan(10.0))
    @test MOI.get(model, MOI.NumberOfConstraints{MOI.ScalarAffineFunction{Float64},MOI.LessThan{Float64}}()) == 1

    MOI.set(model, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(), MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([5.0, 3.0, 2.0, 7.0, 4.0], v), 0.0))
    MOI.set(model, MOI.ObjectiveSense(), MOI.MAX_SENSE)

    if MOI.supports(model, MOI.VariablePrimalStart(), MOI.VariableIndex)
        MOI.set(model, MOI.VariablePrimalStart(), v, [0.0, 0.0, 0.0, 0.0, 0.0])
    end

    if config.solve
        @test MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMIZE_NOT_CALLED

        MOI.optimize!(model)

        @test MOI.get(model, MOI.TerminationStatus()) == config.optimal_status

        @test MOI.get(model, MOI.PrimalStatus()) in [ MOI.FEASIBLE_POINT, MOI.NEARLY_FEASIBLE_POINT ]

        @test MOI.get(model, MOI.ObjectiveValue()) ≈ 16 atol=atol rtol=rtol

        @test MOI.get(model, MOI.VariablePrimal(), v) ≈ [1, 0, 0, 1, 1] atol=atol rtol=rtol
    end
end

function indicator1_test(model::MOI.ModelLike, config::TestConfig)
    atol = config.atol
    rtol = config.rtol
    # linear problem with indicator constraint
    # max  2x1 + 3x2
    # s.t. x1 + x2 <= 10
    #      z1 ==> x2 <= 8
    #      z2 ==> x2 + x1/5 <= 9
    #      z1 + z2 >= 1

    MOI.empty!(model)
    @test MOI.is_empty(model)

    @test MOI.supports(model, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}())
    @test MOI.supports(model, MOI.ObjectiveSense())
    @test MOI.supports_constraint(model, MOI.SingleVariable, MOI.ZeroOne)
    @test MOI.supports_constraint(model, MOI.SingleVariable, MOI.Interval{Float64})
    @test MOI.supports_constraint(model, MOI.ScalarAffineFunction{Float64}, MOI.Interval{Float64})
    @test MOI.supports_constraint(model, MOI.VectorAffineFunction{Float64}, MOI.IndicatorSet{MOI.ACTIVATE_ON_ONE, MOI.LessThan{Float64}})
    x1 = MOI.add_variable(model)
    x2 = MOI.add_variable(model)
    z1 = MOI.add_variable(model)
    z2 = MOI.add_variable(model)
    MOI.add_constraint(model, z1, MOI.ZeroOne())
    MOI.add_constraint(model, z2, MOI.ZeroOne())
    f1 = MOI.VectorAffineFunction(
        [MOI.VectorAffineTerm(1, MOI.ScalarAffineTerm(1.0, z1)),
         MOI.VectorAffineTerm(2, MOI.ScalarAffineTerm(1.0, x2)),
        ],
        [0.0, 0.0]
    )
    iset1 = MOI.IndicatorSet{MOI.ACTIVATE_ON_ONE}(MOI.LessThan(8.0))
    MOI.add_constraint(model, f1, iset1)

    f2 = MOI.VectorAffineFunction(
        [MOI.VectorAffineTerm(1, MOI.ScalarAffineTerm(1.0, z2)),
         MOI.VectorAffineTerm(2, MOI.ScalarAffineTerm(0.2, x1)),
         MOI.VectorAffineTerm(2, MOI.ScalarAffineTerm(1.0, x2)),
        ],
        [0.0, 0.0],
    )
    iset2 = MOI.IndicatorSet{MOI.ACTIVATE_ON_ONE}(MOI.LessThan(9.0))

    MOI.add_constraint(model, f2, iset2)

    # Additional regular constraint.
    MOI.add_constraint(model,
        MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, x1), MOI.ScalarAffineTerm(1.0, x2)], 0.0),
        MOI.LessThan(10.0),
    )

    # Disjunction z1 ⋁ z2
    MOI.add_constraint(model,
        MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, z1), MOI.ScalarAffineTerm(1.0, z2)], 0.0),
        MOI.GreaterThan(1.0),
    )

    MOI.set(model, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([2.0, 3.0], [x1, x2]), 0.0)
    )
    MOI.set(model, MOI.ObjectiveSense(), MOI.MAX_SENSE)

    if config.solve
        @test MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMIZE_NOT_CALLED

        MOI.optimize!(model)

        @test MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMAL
        @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
        @test MOI.get(model, MOI.ObjectiveValue()) ≈ 28.75 atol=atol rtol=rtol
        @test MOI.get(model, MOI.VariablePrimal(), x1) ≈ 1.25 atol=atol rtol=rtol
        @test MOI.get(model, MOI.VariablePrimal(), x2) ≈ 8.75 atol=atol rtol=rtol
        @test MOI.get(model, MOI.VariablePrimal(), z1) ≈ 0.0 atol=atol rtol=rtol
        @test MOI.get(model, MOI.VariablePrimal(), z2) ≈ 1.0 atol=atol rtol=rtol
    end
end

function indicator2_test(model::MOI.ModelLike, config::TestConfig)
    atol = config.atol
    rtol = config.rtol
    # linear problem with indicator constraint
    # max  2x1 + 3x2 - 30 z2
    # s.t. x1 + x2 <= 10
    #      z1 ==> x2 <= 8
    #      z2 ==> x2 + x1/5 <= 9
    #      z1 + z2 >= 1

    MOI.empty!(model)
    @test MOI.is_empty(model)

    # This is the same model as indicator_test1, except that the penalty on z2 forces z1 to be 1.

    x1 = MOI.add_variable(model)
    x2 = MOI.add_variable(model)
    z1  = MOI.add_variable(model)
    z2  = MOI.add_variable(model)
    MOI.add_constraint(model, z1, MOI.ZeroOne())
    MOI.add_constraint(model, z2, MOI.ZeroOne())
    f1 = MOI.VectorAffineFunction(
        [MOI.VectorAffineTerm(1, MOI.ScalarAffineTerm(1.0, z1)),
         MOI.VectorAffineTerm(2, MOI.ScalarAffineTerm(1.0, x2)),
        ],
        [0.0, 0.0]
    )
    iset1 = MOI.IndicatorSet{MOI.ACTIVATE_ON_ONE}(MOI.LessThan(8.0))
    MOI.add_constraint(model, f1, iset1)

    f2 = MOI.VectorAffineFunction(
        [MOI.VectorAffineTerm(1, MOI.ScalarAffineTerm(1.0, z2)),
         MOI.VectorAffineTerm(2, MOI.ScalarAffineTerm(0.2, x1)),
         MOI.VectorAffineTerm(2, MOI.ScalarAffineTerm(1.0, x2)),
        ],
        [0.0, 0.0],
    )
    iset2 = MOI.IndicatorSet{MOI.ACTIVATE_ON_ONE}(MOI.LessThan(9.0))

    MOI.add_constraint(model, f2, iset2)

    # additional regular constraint
    MOI.add_constraint(model,
        MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, x1), MOI.ScalarAffineTerm(1.0, x2)], 0.0),
        MOI.LessThan(10.0),
    )

    # disjunction z1 ⋁ z2
    MOI.add_constraint(model,
        MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, z1), MOI.ScalarAffineTerm(1.0, z2)], 0.0),
        MOI.GreaterThan(1.0),
    )

    # objective penalized on z2
    MOI.set(model, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([2.0, 3.0, -30.0], [x1, x2, z2]), 0.0)
    )
    MOI.set(model, MOI.ObjectiveSense(), MOI.MAX_SENSE)

    if config.solve
        @test MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMIZE_NOT_CALLED

        MOI.optimize!(model)

        @test MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMAL
        @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
        @test MOI.get(model, MOI.ObjectiveValue()) ≈ 28.0 atol=atol rtol=rtol
        @test MOI.get(model, MOI.VariablePrimal(), x1) ≈ 2.0  atol=atol rtol=rtol
        @test MOI.get(model, MOI.VariablePrimal(), x2) ≈ 8.0  atol=atol rtol=rtol
        @test MOI.get(model, MOI.VariablePrimal(), z1) ≈ 1.0 atol=atol rtol=rtol
        @test MOI.get(model, MOI.VariablePrimal(), z2) ≈ 0.0 atol=atol rtol=rtol
    end
end

function indicator3_test(model::MOI.ModelLike, config::TestConfig)
    atol = config.atol
    rtol = config.rtol
    # linear problem with indicator constraint
    # similar to indicator1_test with reversed z1
    # max  2x1 + 3x2
    # s.t. x1 + x2 <= 10
    #      z1 == 0 ==> x2 <= 8
    #      z2 == 1 ==> x2 + x1/5 <= 9
    #      (1-z1) + z2 >= 1 <=> z2 - z1 >= 0

    MOI.empty!(model)
    @test MOI.is_empty(model)

    @test MOI.supports(model, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}())
    @test MOI.supports(model, MOI.ObjectiveSense())
    @test MOI.supports_constraint(model, MOI.SingleVariable, MOI.ZeroOne)
    @test MOI.supports_constraint(model, MOI.SingleVariable, MOI.Interval{Float64})
    @test MOI.supports_constraint(model, MOI.ScalarAffineFunction{Float64}, MOI.Interval{Float64})
    @test MOI.supports_constraint(model, MOI.VectorAffineFunction{Float64}, MOI.IndicatorSet{MOI.ACTIVATE_ON_ONE, MOI.LessThan{Float64}})
    x1 = MOI.add_variable(model)
    x2 = MOI.add_variable(model)
    z1 = MOI.add_variable(model)
    z2 = MOI.add_variable(model)
    vc1 = MOI.add_constraint(model, z1, MOI.ZeroOne())
    @test vc1.value == z1.value
    vc2 = MOI.add_constraint(model, z2, MOI.ZeroOne())
    @test vc2.value == z2.value
    f1 = MOI.VectorAffineFunction(
        [MOI.VectorAffineTerm(1, MOI.ScalarAffineTerm(1.0, z1)),
         MOI.VectorAffineTerm(2, MOI.ScalarAffineTerm(1.0, x2)),
        ],
        [0.0, 0.0]
    )
    iset1 = MOI.IndicatorSet{MOI.ACTIVATE_ON_ZERO}(MOI.LessThan(8.0))
    MOI.add_constraint(model, f1, iset1)

    f2 = MOI.VectorAffineFunction(
        [MOI.VectorAffineTerm(1, MOI.ScalarAffineTerm(1.0, z2)),
         MOI.VectorAffineTerm(2, MOI.ScalarAffineTerm(0.2, x1)),
         MOI.VectorAffineTerm(2, MOI.ScalarAffineTerm(1.0, x2)),
        ],
        [0.0, 0.0],
    )
    iset2 = MOI.IndicatorSet{MOI.ACTIVATE_ON_ONE}(MOI.LessThan(9.0))

    MOI.add_constraint(model, f2, iset2)

    # Additional regular constraint.
    MOI.add_constraint(model,
        MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, x1), MOI.ScalarAffineTerm(1.0, x2)], 0.0),
        MOI.LessThan(10.0),
    )

    # Disjunction (1-z1) ⋁ z2
    MOI.add_constraint(model,
        MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(-1.0, z1), MOI.ScalarAffineTerm(1.0, z2)], 0.0),
        MOI.GreaterThan(0.0),
    )

    MOI.set(model, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([2.0, 3.0], [x1, x2]), 0.0)
    )
    MOI.set(model, MOI.ObjectiveSense(), MOI.MAX_SENSE)

    if config.solve
        @test MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMIZE_NOT_CALLED

        MOI.optimize!(model)

        @test MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMAL
        @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
        @test MOI.get(model, MOI.ObjectiveValue()) ≈ 28.75 atol=atol rtol=rtol
        @test MOI.get(model, MOI.VariablePrimal(), x1) ≈ 1.25 atol=atol rtol=rtol
        @test MOI.get(model, MOI.VariablePrimal(), x2) ≈ 8.75 atol=atol rtol=rtol
        @test MOI.get(model, MOI.VariablePrimal(), z1) ≈ 1.0 atol=atol rtol=rtol
        @test MOI.get(model, MOI.VariablePrimal(), z2) ≈ 1.0 atol=atol rtol=rtol
    end
end

function indicator4_test(model::MOI.ModelLike, config::TestConfig)
    atol = config.atol
    rtol = config.rtol
    # equivalent to indicator1_test with left-hand-side partially in LHS constant
    # linear problem with indicator constraint and
    # max  2x1 + 3x2
    # s.t. x1 + x2 <= 10
    #      z1 ==> x2 - 1 <= 7
    #      z2 ==> x2 + x1/5 + 1 <= 10
    #      z1 + z2 >= 1

    MOI.empty!(model)
    @test MOI.is_empty(model)

    @test MOI.supports(model, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}())
    @test MOI.supports(model, MOI.ObjectiveSense())
    @test MOI.supports_constraint(model, MOI.SingleVariable, MOI.ZeroOne)
    @test MOI.supports_constraint(model, MOI.SingleVariable, MOI.Interval{Float64})
    @test MOI.supports_constraint(model, MOI.ScalarAffineFunction{Float64}, MOI.Interval{Float64})
    @test MOI.supports_constraint(model, MOI.VectorAffineFunction{Float64}, MOI.IndicatorSet{MOI.ACTIVATE_ON_ONE, MOI.LessThan{Float64}})
    x1 = MOI.add_variable(model)
    x2 = MOI.add_variable(model)
    z1 = MOI.add_variable(model)
    z2 = MOI.add_variable(model)
    MOI.add_constraint(model, z1, MOI.ZeroOne())
    MOI.add_constraint(model, z2, MOI.ZeroOne())
    f1 = MOI.VectorAffineFunction(
        [MOI.VectorAffineTerm(1, MOI.ScalarAffineTerm(1.0, z1)),
         MOI.VectorAffineTerm(2, MOI.ScalarAffineTerm(1.0, x2)),
        ],
        [0.0, -1.0]
    )
    iset1 = MOI.IndicatorSet{MOI.ACTIVATE_ON_ONE}(MOI.LessThan(7.0))
    MOI.add_constraint(model, f1, iset1)

    f2 = MOI.VectorAffineFunction(
        [MOI.VectorAffineTerm(1, MOI.ScalarAffineTerm(1.0, z2)),
         MOI.VectorAffineTerm(2, MOI.ScalarAffineTerm(0.2, x1)),
         MOI.VectorAffineTerm(2, MOI.ScalarAffineTerm(1.0, x2)),
        ],
        [0.0, 1.0],
    )
    iset2 = MOI.IndicatorSet{MOI.ACTIVATE_ON_ONE}(MOI.LessThan(10.0))

    MOI.add_constraint(model, f2, iset2)

    # Additional regular constraint.
    MOI.add_constraint(model,
        MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, x1), MOI.ScalarAffineTerm(1.0, x2)], 0.0),
        MOI.LessThan(10.0),
    )

    # Disjunction z1 ⋁ z2
    MOI.add_constraint(model,
        MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, z1), MOI.ScalarAffineTerm(1.0, z2)], 0.0),
        MOI.GreaterThan(1.0),
    )

    MOI.set(model, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([2.0, 3.0], [x1, x2]), 0.0)
    )
    MOI.set(model, MOI.ObjectiveSense(), MOI.MAX_SENSE)

    if config.solve
        @test MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMIZE_NOT_CALLED

        MOI.optimize!(model)

        @test MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMAL
        @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
        @test MOI.get(model, MOI.ObjectiveValue()) ≈ 28.75 atol=atol rtol=rtol
        @test MOI.get(model, MOI.VariablePrimal(), x1) ≈ 1.25 atol=atol rtol=rtol
        @test MOI.get(model, MOI.VariablePrimal(), x2) ≈ 8.75 atol=atol rtol=rtol
        @test MOI.get(model, MOI.VariablePrimal(), z1) ≈ 0.0 atol=atol rtol=rtol
        @test MOI.get(model, MOI.VariablePrimal(), z2) ≈ 1.0 atol=atol rtol=rtol
    end
end

function _semitest(model::MOI.ModelLike, config::TestConfig{T}, int::Bool) where T
    atol = config.atol
    rtol = config.rtol

    @test MOIU.supports_default_copy_to(model, #=copy_names=# false)
    @test MOI.supports(model, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{T}}())
    @test MOI.supports(model, MOI.ObjectiveSense())
    if !int
        @test MOI.supports_constraint(model, MOI.SingleVariable, MOI.Semicontinuous{T})
    else
        @test MOI.supports_constraint(model, MOI.SingleVariable, MOI.Semiinteger{T})
    end

    # 2 variables
    #
    # min  x
    # st   x >= y
    #      if !int
    #           x ∈ {0.0} U [2.0,3.0]
    #      if int
    #           x ∈ {0.0} U {2.0} U {3.0}
    #      y = 0.0

    MOI.empty!(model)
    @test MOI.is_empty(model)

    v = MOI.add_variables(model, 2)
    @test MOI.get(model, MOI.NumberOfVariables()) == 2

    if !int
        vc1 = MOI.add_constraint(model, MOI.SingleVariable(v[1]), MOI.Semicontinuous(T(2), T(3)))
        @test MOI.get(model, MOI.NumberOfConstraints{MOI.SingleVariable,MOI.Semicontinuous{T}}()) == 1
    else
        vc1 = MOI.add_constraint(model, MOI.SingleVariable(v[1]), MOI.Semiinteger(T(2), T(3)))
        @test MOI.get(model, MOI.NumberOfConstraints{MOI.SingleVariable,MOI.Semiinteger{T}}()) == 1
    end

    vc2 = MOI.add_constraint(model, MOI.SingleVariable(v[2]), MOI.EqualTo(zero(T)))
    @test MOI.get(model, MOI.NumberOfConstraints{MOI.SingleVariable,MOI.EqualTo{T}}()) == 1

    cf = MOI.ScalarAffineFunction{T}(MOI.ScalarAffineTerm{T}.([one(T), -one(T)], v), zero(T))
    c = MOI.add_constraint(model, cf, MOI.GreaterThan(zero(T)))
    @test MOI.get(model, MOI.NumberOfConstraints{MOI.ScalarAffineFunction{T},MOI.GreaterThan{T}}()) == 1

    objf = MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1.0, 0.0], v), 0.0)
    MOI.set(model, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(), objf)
    MOI.set(model, MOI.ObjectiveSense(), MOI.MIN_SENSE)

    @test MOI.get(model, MOI.ObjectiveSense()) == MOI.MIN_SENSE

    if config.solve
        @test MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMIZE_NOT_CALLED

        MOI.optimize!(model)

        @test MOI.get(model, MOI.TerminationStatus()) == config.optimal_status
        @test MOI.get(model, MOI.ResultCount()) >= 1
        @test MOI.get(model, MOI.PrimalStatus()) in [ MOI.FEASIBLE_POINT, MOI.NEARLY_FEASIBLE_POINT ]
        @test MOI.get(model, MOI.ObjectiveValue()) ≈ 0.0 atol=atol rtol=rtol
        @test MOI.get(model, MOI.VariablePrimal(), v) ≈ [0, 0] atol=atol rtol=rtol
        @test MOI.get(model, MOI.ConstraintPrimal(), c) ≈ 0.0 atol=atol rtol=rtol
        @test MOI.get(model, MOI.ObjectiveBound()) <= 0.0 + atol
    end

    # Change y fixed value

    MOI.set(model, MOI.ConstraintSet(), vc2, MOI.EqualTo(one(T)))

    if config.solve
        MOI.optimize!(model)

        @test MOI.get(model, MOI.TerminationStatus()) == config.optimal_status
        @test MOI.get(model, MOI.ResultCount()) >= 1
        @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
        @test MOI.get(model, MOI.ObjectiveValue()) ≈ 2.0 atol=atol rtol=rtol
        @test MOI.get(model, MOI.VariablePrimal(), v) ≈ [2.0, 1.0] atol=atol rtol=rtol
        @test MOI.get(model, MOI.ConstraintPrimal(), c) ≈ 1.0 atol=atol rtol=rtol
        @test MOI.get(model, MOI.ObjectiveBound()) <= 2.0 + atol
    end

    MOI.set(model, MOI.ConstraintSet(), vc2, MOI.EqualTo(T(2)))

    if config.solve
        MOI.optimize!(model)

        @test MOI.get(model, MOI.TerminationStatus()) == config.optimal_status
        @test MOI.get(model, MOI.ResultCount()) >= 1
        @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
        @test MOI.get(model, MOI.ObjectiveValue()) ≈ 2.0 atol=atol rtol=rtol
        @test MOI.get(model, MOI.VariablePrimal(), v) ≈ [2.0, 2.0] atol=atol rtol=rtol
        @test MOI.get(model, MOI.ConstraintPrimal(), c) ≈ 0.0 atol=atol rtol=rtol
        @test MOI.get(model, MOI.ObjectiveBound()) <= 2.0 + atol
    end

    MOI.set(model, MOI.ConstraintSet(), vc2, MOI.EqualTo(T(5//2)))

    if config.solve
        MOI.optimize!(model)

        @test MOI.get(model, MOI.TerminationStatus()) == config.optimal_status
        @test MOI.get(model, MOI.ResultCount()) >= 1
        @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
        if !int
            @test MOI.get(model, MOI.ObjectiveValue()) ≈ 2.5 atol=atol rtol=rtol
            @test MOI.get(model, MOI.VariablePrimal(), v) ≈ [2.5, 2.5] atol=atol rtol=rtol
            @test MOI.get(model, MOI.ConstraintPrimal(), c) ≈ 0.0 atol=atol rtol=rtol
            @test MOI.get(model, MOI.ObjectiveBound()) <= 2.5 + atol
        else
            @test MOI.get(model, MOI.ObjectiveValue()) ≈ 3.0 atol=atol rtol=rtol
            @test MOI.get(model, MOI.VariablePrimal(), v) ≈ [3.0, 2.5] atol=atol rtol=rtol
            @test MOI.get(model, MOI.ConstraintPrimal(), c) ≈ 0.5 atol=atol rtol=rtol
            @test MOI.get(model, MOI.ObjectiveBound()) <= 3.0 + atol
        end
    end

    MOI.set(model, MOI.ConstraintSet(), vc2, MOI.EqualTo(T(3)))

    if config.solve
        MOI.optimize!(model)

        @test MOI.get(model, MOI.TerminationStatus()) == config.optimal_status
        @test MOI.get(model, MOI.ResultCount()) >= 1
        @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
        @test MOI.get(model, MOI.ObjectiveValue()) ≈ 3.0 atol=atol rtol=rtol
        @test MOI.get(model, MOI.VariablePrimal(), v) ≈ [3.0, 3.0] atol=atol rtol=rtol
        @test MOI.get(model, MOI.ConstraintPrimal(), c) ≈ 0.0 atol=atol rtol=rtol
        @test MOI.get(model, MOI.ObjectiveBound()) <= 3.0 + atol
    end

    MOI.set(model, MOI.ConstraintSet(), vc2, MOI.EqualTo(T(4)))

     if config.solve
        MOI.optimize!(model)

        @test MOI.get(model, MOI.TerminationStatus()) == MOI.INFEASIBLE ||
            MOI.get(model, MOI.TerminationStatus()) == MOI.INFEASIBLE_OR_UNBOUNDED
    end
end

semiconttest(model::MOI.ModelLike, config::TestConfig) = _semitest(model, config, false)
semiinttest(model::MOI.ModelLike, config::TestConfig) = _semitest(model, config, true)

const intlineartests = Dict("knapsack" => knapsacktest,
                            "int1"     => int1test,
                            "int2"     => int2test,
                            "int3"     => int3test,
                            "indicator1"  => indicator1_test,
                            "indicator2"  => indicator2_test,
                            "indicator3"  => indicator3_test,
                            "indicator4"  => indicator4_test,
                            "semiconttest" => semiconttest,
                            "semiinttest" => semiinttest
                           )

@moitestset intlinear