text:
# R. Crandall
# 2013 October
# Julia implementation of BM3D denoising
#
# Reference
# K. Dabov, A. Foi, V. Katkovnik, and K. Egiazarian, “Image denoising by
# sparse 3-d transform-domain collaborative filtering,” Image Processing,
# IEEE Transactions on, vol. 16, no. 8, pp. 2080–2095, 2007.
#
# An open source C++ implementation and description:
# http://www.ipol.im/pub/art/2012/l-bm3d/
#
# TODO:
# -----------------------------------------------------------------------------
# - Only works on matrix of Float64 for now, add methods for Image classes
# - This algorithm is parallelizable.  See e.g. Marc Lebrun's C++ implementation
# - Using DCT instead of wavelet for 1st step; implement bior1.5 (is there a package for this?)
# - This implementation is rather slow, since it computes the entire 3D group spectrum at
#   once.  A better implementation uses a rolling buffer to only compute the 3D groups in
#   the current search window
# - Remove magic numbers / hard-coded parameters
# - Add Kaiser window to reduce border effects between patches

module BM3D
export bm3d_thr, bm3d_wie

using Hadamard
using NumericExtensions

# 1st step of BM3D denoising: hard thresholding
function bm3d_thr(img::Matrix{Float64}, sigma::Float64)

	# Hard code algorithm parameters for now...
	patchSize = [8;8]
	stepSize = [3;3]
	nBorder = [0;0]
	searchWin = [19;19]
	nMatch = 31
	thresh3D = 2.7

	# Block matching
	(Ilist,Jlist) = get_reference_pixels([size(img,1);size(img,2)],patchSize,stepSize,nBorder)
	matchTable = match_patches(img,Ilist,Jlist,patchSize,searchWin,nMatch)

	G3D = form_groups(img,matchTable,Ilist,Jlist,patchSize)

	# Filter 3D groups by hard thresholding
	G3D[abs(G3D) .< sigma*thresh3D] = 0

	W = zeros(size(G3D))
	for j = 1:length(Jlist)
		for i = 1:length(Ilist)
			T = nnz(G3D[:,:,:,i,j])
			W[:,:,:,i,j] = T > 0 ? 1.0/T : 1.0
		end
	end
	G3D .*= W

	imgOut = invert_groups([size(img,1);size(img,2)],G3D,matchTable,Ilist,Jlist,patchSize)

	Wout = zeros(Float64,size(img))
	groups_to_image!(Wout,W,matchTable,Ilist,Jlist,patchSize)

	return imgOut./Wout

end

# 2nd step of BM3D
# img: input noisy image
# imgBasic: denoised image from first step of BM3D (hard thresholding)
# sigma: known or assumed standard deviation of noise
function bm3d_wie(img::Matrix{Float64}, imgBasic::Matrix{Float64}, sigma::Float64)

	# Hard code algorithm parameters for now...
	patchSize = [8;8]
	stepSize = [3;3]
	nBorder = [0;0]
	searchWin = [11;11]
	nMatch = 15
	thresh3D = 2.7

	# block matching step
	(Ilist,Jlist) = get_reference_pixels([size(img,1);size(img,2)],patchSize,stepSize,nBorder)
	matchTable = match_patches(imgBasic,Ilist,Jlist,patchSize,searchWin,nMatch)

	# Compute 3D group spectrum
	G3D = form_groups(img,matchTable,Ilist,Jlist,patchSize)
	G3Dbasic = form_groups(imgBasic,matchTable,Ilist,Jlist,patchSize)

	# Wiener filtering of 3D groups, using basic estimate as target spectrum
	WC = G3Dbasic.^2./(G3Dbasic.^2 + sigma^2) # devec?
	G3D .*= WC

	# Weight groups
	W = zeros(Float64,size(G3D))
	for j = 1:length(Jlist)
		for i = 1:length(Ilist)
			T = sumsq(WC[:,:,:,i,j])
			W[:,:,:,i,j] = T > 0 ? 1.0/T : 1.0
		end
	end

	G3D .*= W

	Wout = zeros(Float64,size(img))
	groups_to_image!(Wout,W,matchTable,Ilist,Jlist,patchSize)

	imgOut = invert_groups([size(img,1);size(img,2)],G3D,matchTable,Ilist,Jlist,patchSize)

	return imgOut./Wout

end

# Forward BM3D groupings (full transform... inefficient!)
function form_groups(img::Matrix{Float64},
					 matchTable::Array{Float64,4},
					 Ilist::Vector{Int64},
					 Jlist::Vector{Int64},
					 patchSize::Vector{Int64})

	(t,Nmatch,N1,N2) = size(matchTable)

	G3D = zeros(Float64,Nmatch+1,patchSize[1],patchSize[2],N1,N2)

	# Form table of 3D groups
	image_to_groups!(img,G3D,matchTable,Ilist,Jlist,patchSize)

	# Apply 3D DCT on groups
	dct!(G3D,1:3)

	# Apply normalized Walsh-Hadamard transform on 3rd dimension of groups
	#G3D = fwht(G3D,1)*sqrt(float(Nmatch+1))

	return G3D

end

# Inverse BM3D
function invert_groups(imgSize::Vector{Int64},
					   G3D::Array{Float64,5},
					   matchTable::Array{Float64,4},
					   Ilist::Vector{Int64},
					   Jlist::Vector{Int64},
					   patchSize::Vector{Int64})

	(t,Nmatch,N1,N2) = size(matchTable)

	# Allocate image and weight table
	img = zeros(Float64,imgSize[1],imgSize[2])

	# Normalized inverse Walsh-Hadamard transform on 3rd dimension of groups
	#G3D = ifwht(G3D,1)/sqrt(float(Nmatch+1))

	# Apply inverse 3D DCT on groups
	idct!(G3D,1:3)

	groups_to_image!(img,G3D,matchTable,Ilist,Jlist,patchSize)

	return img

end

# Return filtered patches to their place in the image
function groups_to_image!(img::Matrix{Float64},
						  G3D::Array{Float64,5},
						  matchTable::Array{Float64,4},
						  Ilist::Vector{Int64},
						  Jlist::Vector{Int64},
						  patchSize::Vector{Int64})

	Nmatch = size(matchTable,2)

	for j1 = 1:length(Jlist)
		for i1 = 1:length(Ilist)

			for jj = 1:patchSize[2]
				for ii = 1:patchSize[1]
					img[Ilist[i1]+ii-1,Jlist[j1]+jj-1] += G3D[1,ii,jj,i1,j1]
				end
			end

			for k = 1:Nmatch

				i2 = i1 + matchTable[1,k,i1,j1]
				j2 = j1 + matchTable[2,k,i1,j1]

				for jj = 1:patchSize[2]
					for ii = 1:patchSize[1]
						img[Ilist[i2]+ii-1,Jlist[j2]+jj-1] += G3D[k+1,ii,jj,i1,j1]
					end
				end
			end

		end
	end
end

function image_to_groups!(img::Matrix{Float64},
	                      G3D::Array{Float64,5},
	                      matchTable::Array{Float64,4},
	                      Ilist::Vector{Int64},
	                      Jlist::Vector{Int64},
	                      patchSize::Vector{Int64})

	Nmatch = size(matchTable,2)

	for j1 = 1:length(Jlist)
		for i1 = 1:length(Ilist)
			for jj = 1:patchSize[2]
				for ii = 1:patchSize[1]
					G3D[1,ii,jj,i1,j1] = img[Ilist[i1]+ii-1,Jlist[j1]+jj-1]
				end
			end

			for k = 1:Nmatch

				i2 = i1 + matchTable[1,k,i1,j1]
				j2 = j1 + matchTable[2,k,i1,j1]

				for jj = 1:patchSize[2]
					for ii = 1:patchSize[1]
						G3D[k+1,ii,jj,i1,j1] = img[Ilist[i2]+ii-1,Jlist[j2]+jj-1]
					end
				end
			end
		end
	end
end

# Get locations of reference pixels, the upper-left corner of each patch.

# imgSize - size of input image, including border
# patchSize - size of patch in pixels
# stepSize - step in x and y between reference pixels
# nBorder - number of border pixels on each side in x and y

function get_reference_pixels(imgSize::Vector{Int64},
							  patchSize::Vector{Int64},
							  stepSize::Vector{Int64},
							  nBorder::Vector{Int64})

    ph = imgSize[1] - 2*nBorder[1] - patchSize[1] + 1
    pw = imgSize[2] - 2*nBorder[2] - patchSize[2] + 1

    I = [nBorder[1] + 1:stepSize[1]:ph]
    J = [nBorder[2] + 1:stepSize[2]:pw]

    # Make sure there is a patch touching the lower and right borders
    if(maximum(I) < nBorder[1] + ph)
        I = [I; nBorder[1] + ph]
    end
    if(maximum(J) < nBorder[2] + pw)
        I = [J; nBorder[2] + pw]
    end

    return (I,J)

end

# Full-search block matching algorithm for BM3D
function match_patches(img::Matrix{Float64},
					   Ilist::Vector{Int64},
					   Jlist::Vector{Int64},
					   patchSize::Vector{Int64},
					   searchWin::Vector{Int64},
					   nMatch::Int64)

	# Dimensions of patch table
	N1 = length(Ilist)
	N2 = length(Jlist)

	# Allocate
	matchTable = zeros(Float64,(3,nMatch,N1,N2))
    matchMaxTable = ones(Float64,(2,N1,N2))
    matchTable[3,:,:,:] = realmax(Float64)
    matchMaxTable[2,:,:] = realmax(Float64)

    for j1 = 1:N2
    	for i1 = 1:N1

    		for j2 = j1:minimum([N2;j1 + searchWin[2]])
    		# Lower bound on columns
    		LB = (j1 == j2) ? (i1+1) : (i1 - searchWin[1])
    		for i2 = maximum([1,LB]):minimum([i1+searchWin[1];N1])

    				d2 = 0.0
    				for jj = 1:patchSize[2]
    					for ii = 1:patchSize[1]
    						d2 += (img[Ilist[i1] + ii - 1,Jlist[j1] + jj - 1] - img[Ilist[i2] + ii - 1,Jlist[j2] + jj - 1])^2
    					end
    				end
    				d2 /= prod(patchSize)

    				 # Check current maximum for patch (i1,j1)
	                 if (d2 < matchMaxTable[2,i1,j1])
	                 	kmatch = matchMaxTable[1,i1,j1]
	                 	matchTable[1,kmatch,i1,j1] = i2-i1
	                 	matchTable[2,kmatch,i1,j1] = j2-j1
	                 	matchTable[3,kmatch,i1,j1] = d2

	                 	(tmp2,tmp1) = findmax(matchTable[3,:,i1,j1])
	                    matchMaxTable[1,i1,j1] = tmp1
	                    matchMaxTable[2,i1,j1] = tmp2
	                 end

	                 # Check current maximum for patch (i2,j2)
	                 if (d2 < matchMaxTable[2,i2,j2])
	                 	kmatch = matchMaxTable[1,i2,j2]
	                 	matchTable[1,kmatch,i2,j2] = i1-i2
	                 	matchTable[2,kmatch,i2,j2] = j1-j2
	                 	matchTable[3,kmatch,i2,j2] = d2

	                    (tmp2,tmp1) = findmax(matchTable[3,:,i2,j2])
	                    matchMaxTable[1,i2,j2] = tmp1
	                    matchMaxTable[2,i2,j2] = tmp2
	                 end

    			end
    		end

    	end
    end

	return matchTable

end

end #module


text:
using NLSolvers, StaticArrays, Test
@testset "mixed optimization problems" begin
function theta(x)
   if x[1] > 0
       return atan(x[2] / x[1]) / (2.0 * pi)
   else
       return (pi + atan(x[2] / x[1])) / (2.0 * pi)
   end
end
f(x) = 100.0 * ((x[3] - 10.0 * theta(x))^2 + (sqrt(x[1]^2 + x[2]^2) - 1.0)^2) + x[3]^2

function f∇f!(x, ∇f)
    if !(∇f==nothing)
        if ( x[1]^2 + x[2]^2 == 0 )
            dtdx1 = 0;
            dtdx2 = 0;
        else
            dtdx1 = - x[2] / ( 2 * pi * ( x[1]^2 + x[2]^2 ) );
            dtdx2 =   x[1] / ( 2 * pi * ( x[1]^2 + x[2]^2 ) );
        end
        ∇f[1] = -2000.0*(x[3]-10.0*theta(x))*dtdx1 +
            200.0*(sqrt(x[1]^2+x[2]^2)-1)*x[1]/sqrt( x[1]^2+x[2]^2 );
        ∇f[2] = -2000.0*(x[3]-10.0*theta(x))*dtdx2 +
            200.0*(sqrt(x[1]^2+x[2]^2)-1)*x[2]/sqrt( x[1]^2+x[2]^2 );
        ∇f[3] =  200.0*(x[3]-10.0*theta(x)) + 2.0*x[3];
    end

    fx = f(x)
    objective_return(fx, ∇f)
end

function f∇f(x, ∇f)
    if !(∇f == nothing)
        ∇f = similar(x)
    end
    fx, ∇f = f∇f!(gx, x)
    objective_return(fx, ∇f)
end
function f∇fs(x, ∇f)
    if !(∇f == nothing)
        if ( x[1]^2 + x[2]^2 == 0 )
            dtdx1 = 0;
            dtdx2 = 0;
        else
            dtdx1 = - x[2] / ( 2 * pi * ( x[1]^2 + x[2]^2 ) )
            dtdx2 =   x[1] / ( 2 * pi * ( x[1]^2 + x[2]^2 ) )
        end

        s1 = -2000.0*(x[3]-10.0*theta(x))*dtdx1 +
            200.0*(sqrt(x[1]^2+x[2]^2)-1)*x[1]/sqrt( x[1]^2+x[2]^2 )
        s2 = -2000.0*(x[3]-10.0*theta(x))*dtdx2 +
            200.0*(sqrt(x[1]^2+x[2]^2)-1)*x[2]/sqrt( x[1]^2+x[2]^2 )
        s3 = 200.0*(x[3]-10.0*theta(x)) + 2.0*x[3]
        ∇f = @SVector [s1, s2, s3]
    end
    objective_return(f(x), ∇f)
end
obj_inplace = OnceDiffed(f∇f!)
obj_outofplace = OnceDiffed(f∇f)
obj_static = OnceDiffed(f∇fs)

x0 = [-1.0, 0.0, 0.0]
xopt = [1.0, 0.0, 0.0]
x0s = @SVector [-1.0, 0.0, 0.0]

println("Starting  from: ", x0)
println("Targeting     : ", xopt)

res = minimize(obj_inplace, copy(x0), NelderMead(), MinOptions())
print("NN  $(summary(NelderMead()))         ")
@printf("%2.2e  %2.2e %2.2e %d\n",  norm(res.info.minimizer-xopt,Inf), norm(res.info.nm_obj, Inf), res.info.minimum, res.info.iter)
res = minimize(obj_inplace, copy(x0), SimulatedAnnealing(), MinOptions())
print("NN  $(summary(SimulatedAnnealing()))         ")
@printf("%2.2e  %2.2e %d\n",  norm(res.info.minimizer-xopt,Inf),  res.info.minimum, res.info.iter)
res = minimize(obj_inplace, copy(x0), ConjugateGradient(update=HZ()), MinOptions())
print("NN  $(summary(ConjugateGradient(update=HZ())))         ")
@printf("%2.2e  %2.2e %d\n",  norm(res.info.minimizer-xopt,Inf),  res.info.minimum, res.info.iter)
res = minimize(obj_inplace, copy(x0), ConjugateGradient(update=CD()), MinOptions())
print("NN  $(summary(ConjugateGradient(update=CD())))         ")
@printf("%2.2e  %2.2e %d\n",  norm(res.info.minimizer-xopt,Inf),  res.info.minimum, res.info.iter)
res = minimize(obj_inplace, copy(x0), ConjugateGradient(update=HS()), MinOptions())
print("NN  $(summary(ConjugateGradient(update=HS())))         ")
@printf("%2.2e  %2.2e %d\n",  norm(res.info.minimizer-xopt,Inf),  res.info.minimum, res.info.iter)
res = minimize(obj_inplace, copy(x0), ConjugateGradient(update=FR()), MinOptions())
print("NN  $(summary(ConjugateGradient(update=FR())))         ")
@printf("%2.2e  %2.2e %d\n",  norm(res.info.minimizer-xopt,Inf),  res.info.minimum, res.info.iter)
res = minimize(obj_inplace, copy(x0), ConjugateGradient(update=PRP()), MinOptions())
print("NN  $(summary(ConjugateGradient(update=PRP(;plus=false))))         ")
@printf("%2.2e  %2.2e %d\n",  norm(res.info.minimizer-xopt,Inf),  res.info.minimum, res.info.iter)
res = minimize(obj_inplace, copy(x0), ConjugateGradient(update=PRP(plus=true)), MinOptions())
print("NN  $(summary(ConjugateGradient(update=PRP(;plus=true))))         ")
@printf("%2.2e  %2.2e %d\n",  norm(res.info.minimizer-xopt,Inf),  res.info.minimum, res.info.iter)
res = minimize(obj_inplace, copy(x0), ConjugateGradient(update=VPRP()), MinOptions())
print("NN  $(summary(ConjugateGradient(update=VPRP())))         ")
@printf("%2.2e  %2.2e %d\n",  norm(res.info.minimizer-xopt,Inf),  res.info.minimum, res.info.iter)
res = minimize(obj_inplace, copy(x0), ConjugateGradient(update=LS()), MinOptions())
print("NN  $(summary(ConjugateGradient(update=LS())))         ")
@printf("%2.2e  %2.2e %d\n",  norm(res.info.minimizer-xopt,Inf),  res.info.minimum, res.info.iter)
res = minimize(obj_inplace, copy(x0), ConjugateGradient(update=DY()), MinOptions())
print("NN  $(summary(ConjugateGradient(update=DY())))         ")
@printf("%2.2e  %2.2e %d\n",  norm(res.info.minimizer-xopt,Inf),  res.info.minimum, res.info.iter)

for _method in (GradientDescent, LBFGS, BFGS, DBFGS, DFP, SR1)
# for _method in (GradientDescent, BFGS, DBFGS, DFP, SR1)
    methodtxt = summary(_method())
    for m in (Inverse(), Direct())
        mtxt = m isa Inverse ? "(inverse): " : "(direct):  "
        if _method == LBFGS && m isa Inverse
            res = minimize!(obj_inplace, copy(x0), LineSearch(_method(m)), MinOptions())
            print("NN! $_method    $mtxt")
            @printf("%2.2e  %2.2e %2.2e  %d\n", norm(res.info.minimizer-xopt,Inf), norm(res.info.∇fz, Inf), res.info.minimum, res.info.iter)
        elseif _method !== LBFGS
            res = minimize(obj_inplace, copy(x0), LineSearch(_method(m)), MinOptions())
            print("NN  $_method    $mtxt")
            @printf("%2.2e  %2.2e %2.2e  %d\n", norm(res.info.minimizer-xopt,Inf), norm(res.info.∇fz, Inf), res.info.minimum, res.info.iter)
            res = minimize!(obj_inplace, copy(x0), LineSearch(_method(m)), MinOptions())
            print("NN! $_method    $mtxt")
            @printf("%2.2e  %2.2e %2.2e  %d\n", norm(res.info.minimizer-xopt,Inf), norm(res.info.∇fz, Inf), res.info.minimum, res.info.iter)
            res = minimize(obj_static, x0s, LineSearch(_method(m)), MinOptions())
            print("NN  $_method(S) $mtxt")
            @printf("%2.2e  %2.2e %2.2e  %d\n", norm(res.info.minimizer-xopt,Inf), norm(res.info.∇fz, Inf), res.info.minimum, res.info.iter)
            if m isa Direct && !(_method == GradientDescent)
                println("Trust region: NWI")
                res = minimize!(obj_inplace, copy(x0), TrustRegion(_method(Direct()), NWI()), MinOptions())
                print("NN! $_method    $mtxt")
                @printf("%2.2e  %2.2e %2.2e  %d\n", norm(res.info.minimizer-xopt,Inf), norm(res.info.∇fz, Inf), res.info.minimum, res.info.iter)
                println("Trust region: NTR")
                res = minimize!(obj_inplace, copy(x0), TrustRegion(_method(Direct()), NTR()), MinOptions())
                print("NN! $_method    $mtxt")
                @printf("%2.2e  %2.2e %2.2e  %d\n", norm(res.info.minimizer-xopt,Inf), norm(res.info.∇fz, Inf), res.info.minimum, res.info.iter)
            end
        end
    end
end
println()

res = minimize(obj_inplace, x0, LineSearch(BFGS(Inverse())), MinOptions())
@test res.info.iter == 30
@printf("NN  BFGS    (inverse): %2.2e  %2.2e %2.2e %d\n", norm(res.info.minimizer-xopt,Inf), norm(res.info.∇fz, Inf), res.info.minimum, res.info.iter)
res = minimize!(obj_inplace, copy(x0), LineSearch(BFGS(Inverse()), Backtracking()), MinOptions())
@test res.info.iter == 30
@printf("NN! BFGS    (inverse): %2.2e  %2.2e %2.2e %d\n",  norm(res.info.minimizer-xopt,Inf), norm(res.info.∇fz, Inf), res.info.minimum, res.info.iter)
res = minimize!(obj_inplace, copy(x0), LineSearch(BFGS(Inverse()), Backtracking(interp=FFQuadInterp())), MinOptions())
@test res.info.iter == 30
@printf("NN! BFGS    (inverse, quad): %2.2e %2.2e %2.2e %d\n",  norm(res.info.minimizer-xopt,Inf), norm(res.info.∇fz, Inf), res.info.minimum, res.info.iter)
res = minimize(obj_static, x0s, LineSearch(BFGS(Inverse())), MinOptions())
@test res.info.iter == 30
@printf("NN  BFGS(S) (inverse): %2.2e %2.2e %2.2e %d\n",  norm(res.info.minimizer-xopt,Inf), norm(res.info.∇fz, Inf), res.info.minimum, res.info.iter)
res = minimize(obj_static, x0s, LineSearch(BFGS(Inverse()), Backtracking(interp=FFQuadInterp())), MinOptions())
@test res.info.iter == 30
@printf("NN  BFGS(S) (inverse, quad): %2.2e %2.2e %2.2e %d\n",  norm(res.info.minimizer-xopt,Inf), norm(res.info.∇fz, Inf), res.info.minimum, res.info.iter)

x0 = rand(3)
x0s = SVector{3}(x0)
println("\nFrom a random point: ", x0)
res = minimize(obj_inplace, copy(x0), NelderMead(), MinOptions())
print("NN  $(summary(NelderMead()))         ")
@printf("%2.2e  %2.2e %2.2e %d\n",  norm(res.info.minimizer-xopt,Inf), norm(res.info.nm_obj, Inf), res.info.minimum, res.info.iter)
res = minimize(obj_inplace, copy(x0), SimulatedAnnealing(), MinOptions())
print("NN  $(summary(SimulatedAnnealing()))         ")
@printf("%2.2e  %2.2e %d\n",  norm(res.info.minimizer-xopt,Inf),  res.info.minimum, res.info.iter)
res = minimize(obj_inplace, copy(x0), ConjugateGradient(), MinOptions())
print("NN  $(summary(ConjugateGradient()))         ")
@printf("%2.2e  %2.2e %d\n",  norm(res.info.minimizer-xopt,Inf),  res.info.minimum, res.info.iter)

for _method in (GradientDescent, LBFGS, BFGS, DBFGS, DFP, SR1)
    methodtxt = summary(_method())
    for m in (Inverse(), Direct())
        mtxt = m isa Inverse ? "(inverse): " : "(direct):  "
        if _method == LBFGS && m isa Inverse
            res = minimize!(obj_inplace, copy(x0), LineSearch(_method(m)), MinOptions())
            print("NN! $_method    $mtxt")
            @printf("%2.2e  %2.2e %2.2e  %d\n", norm(res.info.minimizer-xopt,Inf), norm(res.info.∇fz, Inf), res.info.minimum, res.info.iter)
        elseif _method !== LBFGS
            res = minimize(obj_inplace, copy(x0), LineSearch(_method(m)), MinOptions())
            print("NN  $_method    $mtxt")
            @printf("%2.2e  %2.2e %2.2e %d\n",  norm(res.info.minimizer-xopt,Inf), norm(res.info.∇fz, Inf), res.info.minimum, res.info.iter)
            res = minimize!(obj_inplace, copy(x0), LineSearch(_method(m)), MinOptions())
            print("NN! $_method    $mtxt")
            @printf("%2.2e  %2.2e %2.2e %d\n",  norm(res.info.minimizer-xopt,Inf), norm(res.info.∇fz, Inf), res.info.minimum, res.info.iter)
            res = minimize(obj_static, x0s, LineSearch(_method(m)), MinOptions())
            print("NN  $_method(S) $mtxt")
            @printf("%2.2e  %2.2e %2.2e %d\n",  norm(res.info.minimizer-xopt,Inf), norm(res.info.∇fz, Inf), res.info.minimum, res.info.iter)
            if m isa Direct && !(_method == GradientDescent)
                println("Trust region: NWI")
                res = minimize!(obj_inplace, copy(x0), TrustRegion(_method(Direct()), NWI()), MinOptions())
                print("NN! $_method    $mtxt")
                @printf("%2.2e  %2.2e %2.2e %d\n",  norm(res.info.minimizer-xopt,Inf), norm(res.info.∇fz, Inf), res.info.minimum, res.info.iter)
                println("Trust region: NTR")
                res = minimize!(obj_inplace, copy(x0), TrustRegion(_method(Direct()), NTR()), MinOptions())
                print("NN! $_method    $mtxt")
                @printf("%2.2e  %2.2e %2.2e %d\n",  norm(res.info.minimizer-xopt,Inf), norm(res.info.∇fz, Inf), res.info.minimum, res.info.iter)
            end
        end
    end
end
println()

function himmelblau!(x, ∇f)
    if !(∇f == nothing)
        ∇f[1] = 4.0 * x[1]^3 + 4.0 * x[1] * x[2] -
            44.0 * x[1] + 2.0 * x[1] + 2.0 * x[2]^2 - 14.0
        ∇f[2] = 2.0 * x[1]^2 + 2.0 * x[2] - 22.0 +
            4.0 * x[1] * x[2] + 4.0 * x[2]^3 - 28.0 * x[2]
    end

    fx = (x[1]^2 + x[2] - 11)^2 + (x[1] + x[2]^2 - 7)^2
    objective_return(fx, ∇f)
end

function himmelblaus(x, ∇f)
    fx = (x[1]^2 + x[2] - 11)^2 + (x[1] + x[2]^2 - 7)^2
    if !(∇f == nothing)
        ∇f1 = 4.0 * x[1]^3 + 4.0 * x[1] * x[2] -
            44.0 * x[1] + 2.0 * x[1] + 2.0 * x[2]^2 - 14.0
        ∇f2 = 2.0 * x[1]^2 + 2.0 * x[2] - 22.0 +
            4.0 * x[1] * x[2] + 4.0 * x[2]^3 - 28.0 * x[2]
        ∇f = @SVector([∇f1, ∇f2])
    end
    objective_return(fx, ∇f)
end

function himmelblau(x, ∇f)
    g = ∇f == nothing ? ∇f : similar(x)

    return himmelblau!(x, g)
end

him_inplace = OnceDiffed(himmelblau!)
him_static = OnceDiffed(himmelblaus)
him_outofplace = OnceDiffed(himmelblau)

println("\nHimmelblau function")
x0 = [3.0, 1.0]
x0s = SVector{2}(x0)
minimizers = [[3.0,2.0],[-2.805118,3.131312],[-3.779310,-3.283186],[3.584428,-1.848126]]
for _method in (GradientDescent, LBFGS, BFGS, DBFGS, DFP, SR1)
# for _method in (GradientDescent, BFGS, DBFGS, DFP, SR1)
    methodtxt = summary(_method())
    for m in (Inverse(), Direct())
        mtxt = m isa Inverse ? "(inverse): " : "(direct):  "
        if _method == LBFGS && m isa Inverse
            res = minimize!(him_inplace, copy(x0), LineSearch(_method(m)), MinOptions())
            print("NN! $_method    $mtxt")
            @printf("%2.2e  %2.2e %2.2e  %d\n", minimum([norm(res.info.minimizer-xopt,Inf) for xopt in minimizers]), norm(res.info.∇fz, Inf), res.info.minimum, res.info.iter)
        elseif _method !== LBFGS
            res = minimize(him_outofplace, copy(x0), LineSearch(_method(m)), MinOptions())
            print("NN  $_method    $mtxt")
            @printf("%2.2e  %2.2e %2.2e  %d\n", minimum([norm(res.info.minimizer-xopt,Inf) for xopt in minimizers]), norm(res.info.∇fz, Inf), res.info.minimum, res.info.iter)
            res = minimize!(him_inplace, copy(x0), LineSearch(_method(m)), MinOptions())
            print("NN! $_method    $mtxt")
            @printf("%2.2e  %2.2e %2.2e  %d\n", minimum([norm(res.info.minimizer-xopt,Inf) for xopt in minimizers]), norm(res.info.∇fz, Inf), res.info.minimum, res.info.iter)
            res = minimize(him_static, x0s, LineSearch(_method(m)), MinOptions())
            print("NN  $_method(S) $mtxt")
            @printf("%2.2e  %2.2e %2.2e  %d\n", minimum([norm(res.info.minimizer-xopt,Inf) for xopt in minimizers]), norm(res.info.∇fz, Inf), res.info.minimum, res.info.iter)
            if m isa Direct && !(_method == GradientDescent)
                println("Trust region: NWI")
                res = minimize!(him_inplace, copy(x0), TrustRegion(_method(Direct()), NWI()), MinOptions())
                print("NN! $_method    $mtxt")
                @printf("%2.2e  %2.2e %2.2e  %d\n", minimum([norm(res.info.minimizer-xopt,Inf) for xopt in minimizers]), norm(res.info.∇fz, Inf), res.info.minimum, res.info.iter)
                println("Trust region: NTR")
                res = minimize!(him_inplace, copy(x0), TrustRegion(_method(Direct()), NTR()), MinOptions())
                print("NN! $_method    $mtxt")
                @printf("%2.2e  %2.2e %2.2e  %d\n", minimum([norm(res.info.minimizer-xopt,Inf) for xopt in minimizers]), norm(res.info.∇fz, Inf), res.info.minimum, res.info.iter)
            end
        end
    end
end
println()
println()
xrand = rand(2)
xrands = SVector{2}(xrand)
println("\nFrom a random point: ", xrand)

for _method in (GradientDescent, LBFGS, BFGS, DBFGS, SR1)
# for _method in (GradientDescent, BFGS, DBFGS, DFP, SR1)
    methodtxt = summary(_method())
    for m in (Inverse(), Direct())
        mtxt = m isa Inverse ? "(inverse): " : "(direct):  "
        if _method == LBFGS && m isa Inverse
            res = minimize!(him_inplace, copy(x0), LineSearch(_method(m)), MinOptions())
            print("NN! $_method    $mtxt")
            @printf("%2.2e  %2.2e %2.2e  %d\n", minimum([norm(res.info.minimizer-xopt,Inf) for xopt in minimizers]), norm(res.info.∇fz, Inf), res.info.minimum, res.info.iter)
        elseif _method !== LBFGS
            res = minimize(him_outofplace, copy(xrand), LineSearch(_method(m)), MinOptions())
            print("NN  $_method    $mtxt")
            @printf("%2.2e  %2.2e %2.2e  %d\n", minimum([norm(res.info.minimizer-xopt,Inf) for xopt in minimizers]), norm(res.info.∇fz, Inf), res.info.minimum, res.info.iter)
            res = minimize!(him_inplace, copy(xrand), LineSearch(_method(m)), MinOptions())
            print("NN! $_method    $mtxt")
            @printf("%2.2e  %2.2e %2.2e  %d\n", minimum([norm(res.info.minimizer-xopt,Inf) for xopt in minimizers]), norm(res.info.∇fz, Inf), res.info.minimum, res.info.iter)
            res = minimize(him_static, xrands, LineSearch(_method(m)), MinOptions())
            print("NN  $_method(S) $mtxt")
            @printf("%2.2e  %2.2e %2.2e  %d\n", minimum([norm(res.info.minimizer-xopt,Inf) for xopt in minimizers]), norm(res.info.∇fz, Inf), res.info.minimum, res.info.iter)
            if m isa Direct && !(_method == GradientDescent)
                println("Trust region: NWI")
                res = minimize!(him_inplace, copy(xrand), TrustRegion(_method(Direct()), NWI()), MinOptions())
                print("NN! $_method    $mtxt")
                @printf("%2.2e  %2.2e %2.2e  %d\n", minimum([norm(res.info.minimizer-xopt,Inf) for xopt in minimizers]), norm(res.info.∇fz, Inf), res.info.minimum, res.info.iter)
                println("Trust region: NTR")
                res = minimize!(him_inplace, copy(xrand), TrustRegion(_method(Direct()), NTR()), MinOptions())
                print("NN! $_method    $mtxt")
                @printf("%2.2e  %2.2e %2.2e  %d\n", minimum([norm(res.info.minimizer-xopt,Inf) for xopt in minimizers]), norm(res.info.∇fz, Inf), res.info.minimum, res.info.iter)
            end
        end
    end
end
println()

end

ext: jl


text:
using WAVI, Plots
"""
    MISMIP Experiments EXP1_1: advance to a steady state on a linear slope

"""
function MISMIP_EXP1_1(A)
    #grid and bc
    nx = 300
    ny = 2
    nσ = 4
    x0 = 0.
    y0 = 0.
    dx = round(18.e5 / nx)
    dy = dx
    h_mask = trues(nx,ny)
    grid = Grid(nx = nx,
            ny = ny,
            nσ = nσ,
            x0 = x0,
            y0 = y0,
            dx = dx,
            dy = dy,
            h_mask = h_mask)

    #bed
    function b(x,y)
        B = 720 - 778.5 * x ./ (750e3)
        return B
    end

    #solver parameters
    maxiter_picard = 1
    solver_params = SolverParams(maxiter_picard = maxiter_picard)

    #Physical parameters
    default_thickness = 50.0 #set the initial condition this way
    accumulation_rate = 0.3
    sec_per_year = 365.25*24*60*60
    glen_a_ref = A * sec_per_year

    params = Params(default_thickness = default_thickness,
                    accumulation_rate = accumulation_rate,
                    glen_a_ref = glen_a_ref)

    #make the model
    model = Model(grid = grid,
    bed_elevation = b,
    params = params,
    solver_params = solver_params)

    #timestepping parameters
    niter0 = 0
    dt = 0.5
    end_time = 250.
    timestepping_params = TimesteppingParams(niter0 = niter0,
                                            dt = dt,
                                            end_time = end_time)

    folder = joinpath(@__DIR__, "./outputs_mismip_exp1")
    isdir(folder) && rm(folder, force = true, recursive = true)
    mkdir(folder) #make a clean folder for outputs
    outputs = (h = model.fields.gh.h,
                u = model.fields.gh.u,
                v = model.fields.gh.v,
                b = model.fields.gh.b,
                s = model.fields.gh.s,
                grfrac = model.fields.gh.grounded_fraction) #output velocities and thickness

    output_freq = 10.
    output_params = OutputParams(outputs = outputs,
                            output_path = folder,
                            output_freq = output_freq,
                            output_format = "jld2",
                            zip_format = "none")

    simulation = Simulation(model = model,
    timestepping_params = timestepping_params,
    output_params = output_params)

    #perform the simulation
    run_simulation!(simulation)

    return simulation
end

A = 4.6416e-24
simulation = MISMIP_EXP1_1(A)

################################################
function plot_evolution()
outfolder = joinpath(@__DIR__, "outputs_mismip_exp1")
files = [joinpath(outfolder, file) for file in readdir(outfolder) if endswith( joinpath(outfolder, file), ".jld2") ]
nout = length(files)

#get the bed
dd = load(files[1])
bed  = dd["b"][:,1]
x    = dd["x"][:,1]
nx = length(bed[:,1])
h_out = zeros(nx, nout)
s_out = zeros(nx, nout)
ib_out = zeros(nx, nout)
t_out = zeros(1,nout)

#load solution files into matrix
for i = 1:nout
    fname = files[i]
    d = load(fname)

    h_out[:,i] = d["h"][:,1]
    ib_out[:,i] = d["s"][:,1] .- d["h"][:,1]
    s_out[:,i] = d["s"][:,1]
    t_out[i] = d["t"]
end

#plot things
plot1 = Plots.plot(x, bed, legend = false)
Plots.plot!(x, s_out)
Plots.plot!(x, ib_out)

display(plot1)
return nothing
end

plot_evolution()


text:
using Dhall
using Test
using FileIO

function load(path::AbstractString)
    # TODO delete this function after registration
    f = FileIO.File(FileIO.format"dhall", path)
    Dhall.load(f)
end

function testdatapath(args...)
    joinpath(@__DIR__, "assets", args...)
end

@testset "load" begin
    @test load(testdatapath("hello_world.dhall")) == ["hello", "world"]
    @test load(testdatapath("hello_world2.dhall")) == ["hello", "world"]
    @test load(testdatapath("hello_world3.dhall")) == ["hello", "world"]
    @test load(testdatapath("hello_struct.dhall")) == Dict("x"=>1, "y"=>2)
end

@testset "smoketest executables $name" for (name, path) in pairs(Dhall.PATHS)
    @test ispath(path)
    @test isfile(path)
    version = read(`$path --version`, String)
    @show version
end


text:
export Point

struct Point{T<:Float64}
    x::T
    y::T
end

function Point(x::Int, y::Int)
    return Point(convert(Float64, x), convert(Float64, y))
end

Base.:(==)(p::Point, q::Point) = (p.x == q.x) & (p.y == q.y)
Base.isless(p::Point, q::Point) = (p.y > q.y) | ((p.y == q.y) & (p.x < q.x))
Base.isapprox(p::Point, q::Point; rtol = 1e-6) =
    isapprox(p.x, q.x, rtol = rtol) && isapprox(p.y, q.y, rtol = rtol)


text:
# This file is a part of SolidStateDetectors.jl, licensed under the MIT License (MIT).

include("Units.jl")

include("Points.jl")
include("Vectors.jl")

abstract type AbstractGeometry{T, N} end
abstract type AbstractVolumePrimitive{T, N} <: AbstractGeometry{T, N} end
abstract type AbstractSurfacePrimitive{T, N} <: AbstractGeometry{T, N} end
abstract type AbstractSet{T, N} <: AbstractGeometry{T, N} end

function Geometry(T, geom_dict::Union{Dict{String,Any},Dict{Any,Any}}, inputunit_dict::Dict{String, Unitful.Units} )
    Geometry(T, Val{Symbol(geom_dict["type"])}(), Dict{Any,Any}(geom_dict), inputunit_dict)
end
# function Geometry(T, geom_sub_dict::Dict, geom_unit)
#     haskey(geom_sub_dict, "positive") ? geometry_positive = Geometries(T, geom_sub_dict["positive"], geom_unit) : geometry_positive = AbstractGeometry[]
#     haskey(geom_sub_dict, "negative") ? geometry_negative = Geometries(T, geom_sub_dict["negative"], geom_unit) : geometry_negative = AbstractGeometry[]
#     geometry = geometry_positive[1]
#     for g in sort!(vcat(geometry_positive,geometry_negative))
#         g in geometry_negative ? geometry -= g : nothing
#         g in geometry_positive ? geometry += g : nothing
#     end
#     return geometry, geometry_positive, geometry_negative
# end

function Geometry(T, geom_sub_dict::Vector, geom_unit)
    geometry_positive = Geometries(T, geom_sub_dict, geom_unit)
    geometry_negative = AbstractGeometry[]
    geometry = geometry_positive[1]
    for g in sort!(vcat(geometry_positive,geometry_negative))
        g in geometry_negative ? geometry -= g : nothing
        g in geometry_positive ? geometry += g : nothing
    end
    return geometry, geometry_positive, geometry_negative
end

function in(pt::AbstractCoordinatePoint{T}, vg::Vector{AbstractGeometry{T}})::Bool where {T <: SSDFloat}
    inside::Bool = false
    for g in vg
        if in(pt, g)
            inside = true
            break
        end
    end
    return inside
end

include("VolumePrimitives/VolumePrimitives.jl")
include("SurfacePrimitives/SurfacePrimitives.jl")
include("LinePrimitives/LinePrimitives.jl")

include("Sets.jl")

include("IO.jl")
function get_decomposed_volumes(vol::AbstractGeometry{T})::Tuple{Vector{<:AbstractGeometry},Vector{<:AbstractGeometry}} where T
    positive_volumes = AbstractVolumePrimitive[]
    negative_volumes = AbstractVolumePrimitive[]
    translate = CartesianVector{T}(0.0,0.0,0.0)
    decompose_volume(vol, positive_volumes, negative_volumes, translate,  1)
    positive_volumes, negative_volumes
end

function decompose_volume(vol::Difference, pos, neg, translate, flag = 1)::Nothing
    ismissing(vol.translate) ? nothing : translate += vol.translate
    decompose_volume(vol.a, pos, neg, translate, flag)
    decompose_volume(vol.b, pos, neg, translate, flag *= -1)
    nothing
end

function decompose_volume(vol::GeometryUnion, pos, neg, translate, flag = 1)::Nothing
    ismissing(vol.translate) ? nothing : translate += vol.translate
    decompose_volume(vol.a ,pos, neg, translate, flag)
    decompose_volume(vol.b, pos, neg, translate, flag)
    nothing
end

function decompose_volume(vol::Intersection, pos, neg, translate, flag = 1)::Nothing
    ismissing(vol.translate) ? nothing : translate += vol.translate
    decompose_volume(vol.a ,pos, neg, translate, flag)
    decompose_volume(vol.b, pos, neg, translate, flag)
    nothing
end

function decompose_volume(vol::AbstractVolumePrimitive, pos, neg, translate, flag = 1)::Nothing
    # ismissing(vol.translate) ? vol.translate = translate : vol.translate += translate
    flag == 1 ? push!(pos, vol+translate) : push!(neg, vol+translate)
    nothing
end

function sample(geometry::AbstractGeometry{T}, stepsize::Vector{T}) where T
    return filter(x-> x in geometry, unique!(vcat([sample(g, stepsize) for g in get_decomposed_volumes(geometry)[1] ]...)))
end


text:
module FinMetrics

using Distributions
using MathProgBase
using ForwardDiff
using StatsBase
using StatsFuns
using KNITRO
using Ipopt
import Distributions: rand, length, insupport, _logpdf
import Base: size

FinMetricsDefaultSolver() = IpoptSolver()

include("univariate/garch.jl")

end # module


text:
@with_kw struct Grid{T<:Real,Tprog<:Real}

    # Parameters taken from Parameter struct
    nx::Int                             # number of grid cells in x-direction
    Lx::Real                            # length of the domain in x-direction [m]
    L_ratio::Real                       # Domain aspect ratio of Lx/Ly
    bc::String                          # boundary condition, "periodic" or "nonperiodic"
    g::Real                             # gravitational acceleration [m/s]
    H::Real                             # layer thickness at rest [m]
    cfl::Real                           # CFL number (1.0 recommended for RK4, 0.6 for RK3)
    Ndays::Real                         # number of days to integrate for
    nstep_diff::Int                     # diffusive terms every nstep_diff time steps.
    nstep_advcor::Int                   # nonlinear terms every nstep_advcor time steps.
    Uadv::Real                          # Velocity scale [m/s] for tracer advection
    output_dt::Real                     # output time step [hours]
    ω::Real                             # Earth's angular frequency [s^-1]
    ϕ::Real                             # central latitue of the domain (for coriolis) [°]
    R::Real                             # Earth's radius [m]
    scale::Real                         # multiplicative scale for momentum equations [1]

    # DOMAIN SIZES
    Δ::Real=Lx / nx                         # grid spacing
    ny::Int=Int(round(Lx / L_ratio / Δ))    # number of grid cells in y-direction
    Ly::Real=ny * Δ                         # length of domain in y-direction

    # NUMBER OF GRID POINTS
    nux::Int = if (bc == "periodic") nx else nx-1 end   # u-grid in x-direction
    nuy::Int = ny                                       # u-grid in y-direction
    nvx::Int = nx                                       # v-grid in x-direction
    nvy::Int = ny-1                                     # v-grid in y-direction
    nqx::Int = if (bc == "periodic") nx else nx+1 end   # q-grid in x-direction
    nqy::Int = ny+1                                     # q-grid in y-direction

    # TOTAL NUMBER OF GRID POINTS
    nT::Int = nx*ny                     # T-grid
    nu::Int = nux*nuy                   # u-grid
    nv::Int = nvx*nvy                   # v-grid
    nq::Int = nqx*nqy                   # q-grid

    # GRID VECTORS
    x_T::AbstractVector = Δ*Array(1:nx) .- Δ/2
    y_T::AbstractVector = Δ*Array(1:ny) .- Δ/2
    x_u::AbstractVector = if (bc == "periodic") Δ*Array(0:nx-1) else Δ*Array(1:nx-1) end
    y_u::AbstractVector = y_T
    x_v::AbstractVector = x_T
    y_v::AbstractVector = Δ*Array(1:ny-1)
    x_q::AbstractVector = if bc == "periodic" x_u else Δ*Array(1:nx+1) .- Δ end
    y_q::AbstractVector = Δ*Array(1:ny+1) .- Δ

    # HALO SIZES
    halo::Int=2                         # halo size for u,v (Biharmonic stencil requires 2)
    haloη::Int=1                        # halo size for η
    halosstx::Int=1                     # halo size for tracer sst in x
    halossty::Int=0                     # halo size for tracer sst in y

    # EDGE POINT (1 = yes, 0 = no)
    ep::Int = if bc == "periodic" 1 else 0 end  # is there a u-point on the left edge?

    # GRID VECTORS WITH HALO
    x_T_halo::AbstractVector = Δ*Array(0:nx+1) .- Δ/2
    y_T_halo::AbstractVector = Δ*Array(0:ny+1) .- Δ/2
    x_u_halo::AbstractVector = if (bc == "periodic") Δ*Array(-2:nx+1) else Δ*Array(-1:nx+1) end
    y_u_halo::AbstractVector = Δ*Array(-1:ny+2) .- Δ/2
    x_v_halo::AbstractVector = Δ*Array(-1:nx+2) .- Δ/2
    y_v_halo::AbstractVector = Δ*Array(-1:ny+1)
    x_q_halo::AbstractVector = if bc == "periodic" x_u_halo else Δ*Array(-1:nx+3) .- Δ end
    y_q_halo::AbstractVector = Δ*Array(-1:ny+3) .- Δ

    # TIME STEPS
    c::Real = √(g*H)                            # shallow water gravity wave speed
    dtint::Int = Int(floor(cfl*Δ/c))            # dt converted to Int
    nt::Int = Int(ceil(Ndays*3600*24/dtint))    # number of time steps to integrate
    dt::T = T(dtint)                            # time step [s]
    Δt::T = T(dtint/Δ)                          # time step divided by grid spacing [s/m]
    Δt_diff::Tprog = Tprog(nstep_diff*dtint/Δ)  # time step for diffusive terms

    # TIME STEPS FOR ADVECTION
    nadvstep::Int = max(1,Int(floor(Δ/Uadv/dtint)))         # advection each n time steps
    nadvstep_half::Int = nadvstep ÷ 2                       # nadvstep ÷ 2
    dtadvint::Int = nadvstep*dtint                          # advection time step [s]
    # divide by scale here to undo the scaling in u,v for tracer advection
    dtadvu::T = T(dtadvint*nx/Lx/scale)                     # Rescaled advection time step for u [s/m]
    dtadvv::T = T(dtadvint*ny/Ly/scale)                     # Rescaled advection time step for v [s/m]
    half_dtadvu::T = T(dtadvint*nx/Lx/2/scale)              # dtadvu/2
    half_dtadvv::T = T(dtadvint*ny/Ly/2/scale)              # dtadvv/2

    # N TIME STEPS FOR OUTPUT
    nout::Int = max(1,Int(floor(output_dt*3600/dtint)))     # output every n time steps
    nout_total::Int = (nt ÷ nout)+1                         # total number of output time steps
    t_vec::AbstractVector = Array(0:nout_total-1)*dtint     # time vector

    # CORIOLIS
    f₀::Float64 = coriolis_at_lat(ω,ϕ)                      # Coriolis parameter
    β::Float64 = β_at_lat(ω,R,ϕ)                            # Derivate of Coriolis parameter wrt latitude
    # scale only f_q as it's used for non-linear advection
    f_q::Array{T,2} = T.(scale*Δ*(f₀ .+ β*(yy_q(bc,x_q_halo,y_q_halo) .- Ly/2)))  # same on the q-grid
    # f_u, f_v are only used for linear dynamics (scaling implicit)
    f_u::Array{T,2} = T.(Δ*(f₀ .+ β*(meshgrid(x_u,y_u)[2] .- Ly/2)))        # f = f₀ + βy on the u-grid
    f_v::Array{T,2} = T.(Δ*(f₀ .+ β*(meshgrid(x_v,y_v)[2] .- Ly/2)))        # same on the v-grid
end

"""Helper function to create yy_q based on the boundary condition bc."""
function yy_q(bc::String,x_q_halo::AbstractVector,y_q_halo::AbstractVector)
    if bc == "periodic"
        # points on the right edge needed too
        _,yy_q = meshgrid(x_q_halo[3:end-1],y_q_halo[3:end-2])
    else
        _,yy_q = meshgrid(x_q_halo[3:end-2],y_q_halo[3:end-2])
    end

    return yy_q
end

"""Generator function for the Grid struct."""
function Grid{T,Tprog}(P::Parameter) where {T<:Real,Tprog<:Real}
    @unpack nx,Lx,L_ratio = P
    @unpack bc,g,H,cfl = P
    @unpack Ndays,nstep_diff,nstep_advcor = P
    @unpack Uadv,output_dt = P
    @unpack ϕ,ω,R,scale = P

    return Grid{T,Tprog}(nx=nx,Lx=Lx,L_ratio=L_ratio,bc=bc,g=g,H=H,cfl=cfl,Ndays=Ndays,
                nstep_diff=nstep_diff,nstep_advcor=nstep_advcor,Uadv=Uadv,output_dt=output_dt,
                ϕ=ϕ,ω=ω,R=R,scale=scale)
end

"""Meter per 1 degree of latitude (or longitude at the equator)."""
function m_per_lat(R::Real)
    return 2π*R/360.
end

"""Coriolis parameter f [1/s] at latitude ϕ [°] given Earth's rotation ω [1/s]."""
function coriolis_at_lat(ω::Real,ϕ::Real)
    return 2*ω*sind(ϕ)
end

"""Coriolis parameter's derivative β wrt latitude [(ms)^-1] at latitude ϕ, given
Earth's rotation ω [1/s] and radius R [m]."""
function β_at_lat(ω::Real,R::Real,ϕ::Real)
    return 2*ω/R*cosd(ϕ)
end

"""Similar to the numpy meshgrid function:
repeats x length(y)-times and vice versa. Returns two matrices xx,yy of same shape so that
each row of xx is x and each column of yy is y."""
function meshgrid(x::AbstractVector,y::AbstractVector)
    m,n = length(x),length(y)

    # preallocate preserving the data type of x,y
    xx = zeros(eltype(x),m,n)
    yy = zeros(eltype(y),m,n)

    for i in 1:m
        xx[i,:] .= x[i]
    end

    for i in 1:n
        yy[:,i] .= y[i]
    end

    return xx,yy
end


text:
#=
     The most significand bit in the exponent bit field
         is freed from floating point valuation, that
         it may elucidate state or elaborate context.

In all IEEE-754 2008 conformant binary floating point representations, the exponent bit field
is adjacent to the sign bit at the most significant region of the floating point representation.
The most significant bit of the exponent is the second most significant bit in the representation.
This is the bit that is to be freed and made available as annotation, demarcation or connection.
Releasing this bit necessarily halves number of distinct exponents that may be stored and that
reduces the quantity of representable finite values by half.  The resulting, adjusted domain
for these stateful projections of Float64|32 values onto Float64|32 is given with the mapping.

    Most computing uses numerical values within 1e-153 .. 1e+153:  all remain available.

  nomenclature

    An e-bit carrying floating point value is an `augmented` float.
    An augmented float that has its e-bit clear (0b0) is 'situated'.
    An augmented float that has its e-bit set   (0b1) is 'enhanced'.

=#

import Base: sign_mask, exponent_mask, exponent_bias, significand_bits

#=
    The 'e-bit' is the most significant bit of the exponent in an IEEE754-2008 Std Binary float.
    These inline ops are used for direct access and immediate manipulation of the 'e-bit'.
=#

@inline ebit_mask{T<:AbstractFloat}(::Type{T}) =  (sign_mask(T) >> 1)
@inline ebit_filt{T<:AbstractFloat}(::Type{T}) = ~(ebit_mask(T))
@inline ebit_mask{T<:Unsigned}(x::T) = ebit_mask(reinterpret(AbstractFloat,T))
@inline ebit_filt{T<:Unsigned}(x::T) = ebit_mask(reinterpret(AbstractFloat,T))

@inline set_ebit{T<:AbstractFloat}(x::T) = reinterpret(T,(reinterpret(Unsigned,x) | ebit_mask(T)))
@inline clr_ebit{T<:AbstractFloat}(x::T) = reinterpret(T,(reinterpret(Unsigned,x) & ebit_filt(T)))
@inline tst_ebit{T<:AbstractFloat}(x::T) = zero(reinterpret(Unsigned,T)) != (reinterpret(Unsigned,x) & ebit_mask(T))

#=
   this is the exponent governance
=#
@inline exponent_filt{T<:AbstractFloat}(::Type{T})  = ~exponent_mask(T)

for m in (:sign_mask, :exponent_mask, :exponent_filt, :exponent_bias, :significand_bits)
          #  obtain masks within the context of an Unsigned type
    @eval ($m){T<:Unsigned}(::Type{T}) = ($m)(reinterpret(AbstractFloat,T))
end

# these serve utilizing raw, biased exponent values

@inline clear_exponent{T<:Unsigned}(x::T)     = (x & exponent_filt(T))
@inline isolate_exponent{T<:Unsigned}(x::T)   = (x & exponent_mask(T))
@inline get_exponent{T<:Unsigned}(x::T)       = (isolate_exponent(x) >> significand_bits(T))
@inline put_exponent{T<:Unsigned}(x::T, e::T) = (isolate_exponent(e << significand_bits(T)) | clear_exponent(x))

@inline get_exponent(x::AbstractFloat)  = get_exponent(reinterpret(Unsigned,x))
@inline put_exponent(x::AbstractFloat, e::Unsigned)= put_exponent(reinterpret(Unsigned,x),e)

#=
  nomenclature
    An e-bit enfolding floating point value is an 'augmented' float.
    An augmented float that has its e-bit clear (0b0) is 'situated'.
    An augmented float that has its e-bit set   (0b1) is 'enhanced'.
=#

# query state
enhanced{T<:AbstractFloat}(x::T) =  tst_ebit(x)  #      is this enhanced?
situated{T<:AbstractFloat}(x::T) = !tst_ebit(x)  #      is this situated?

# force state
 enhance!{T<:AbstractFloat}(x::T) =  set_ebit(x)  # this is now  enhanced!
 situate!{T<:AbstractFloat}(x::T) =  clr_ebit(x)  # this is now  situated!


text:
import AlphaZero
#include("game.jl")
#g = GI.read_state(Game)
#g = GI.init(GameSpec())
#update_status!(g)
# GI.render(g)
#print(" Move: ")
#a = GI.parse_action(g, readline())
# a = GI.parse_action(g, "8c*n")
# GI.play!(g, GI.parse_action(g, "00/s"))
# GI.render(g)
# GI.play!(g, GI.parse_action(g, "44/n"))
# GI.render(g)
# GI.play!(g, GI.parse_action(g, "80/n"))
# GI.render(g)
# GI.play!(g, GI.parse_action(GI.spec(g), "40/e"))
# GI.render(g)
# GI.render(GI.play!(g, GI.parse_action(g, "0c*e")))
AlphaZero.Scripts.test_game("bonbon-rectangle")


text:
export
    _complete,
    depth,
    expander_complete,
    expansions,
    expanded,
    expansion,
    subruntime,
    clear_analysis!,
    ensure_expansion_state!,
    expand!,
    expander_probs

function expander_complete(runtime :: Runtime, v :: Variable, parent_ranges)
    for p in parent_ranges[1]
        if !expanded(runtime, v, p) return false end
        subnet = expansion(runtime, v, p)
        subrun = subruntime(runtime, v, subnet)
        if !_complete(subrun) return false end
    end
    return true
end

# This method _only_ works if the model being operated upon contains an
# operated called 'support_quality'
function _complete(runtime :: Runtime)
    net = runtime.network
    order = topsort(get_initial_graph(net))
    for var in order
        if !has_instance(runtime, var)
            return false
        end
        inst = current_instance(runtime, var)
        sf = get_sfunc(inst)

        if !has_range(runtime, inst)
            return false
        end
        parranges = parent_ranges(runtime, var)
        imp = get_imp(MultiInterface.get_policy(), Support, sf, parranges, 0, output_type(sf)[])
        # if is_fixed(var.model) && isa(sf, Expander)
        if isa(sf, Expander)
            if !expander_complete(runtime, var, Tuple(parranges))
                return false
            end
        elseif !(support_quality(imp, sf, parranges) == :CompleteSupport)
            return false
        end
    end
    return true
end

function expansions(runtime::InstantRuntime, var::Variable) :: Dict
    exp :: Expander = make_initial(var.model)
    mod = var.model
    has_state(runtime, :subnets) || return Dict()
    exps = get_state(runtime, :subnets)
    mod in keys(exps) || return Dict()
    return exps[mod]
end

function expanded(runtime::Runtime, var::Variable, arg)
    return arg in keys(expansions(runtime, var))
end

function expansion(runtime :: Runtime, var :: Variable, arg) :: Network
    return expansions(runtime, var)[arg]
end

function subruntime(runtime::Runtime, var::Variable, net)
    subruns = get_state(runtime, :subruntimes)
    return subruns[net]
end

function depth(runtime::Runtime) :: Int
    if has_state(runtime, :depth)
        get_state(runtime, :depth)
    else
        0
    end
end

function ensure_expansion_state!(runtime::Runtime)
    has_state(runtime, :subnets) || set_state!(runtime, :subnets, Dict())
    has_state(runtime, :subruntimes) || set_state!(runtime, :subruntimes, Dict())
end

# Returns the expanded subnet and a flag indicating whether it was newly expanded
function expand!(runtime::InstantRuntime, var::Variable, arg)
    ensure_expansion_state!(runtime)
    exp :: Expander = make_initial(var.model)
    mod = var.model
    exps = get_state(runtime, :subnets)
    subruns = get_state(runtime, :subruntimes)
    if !(mod in keys(exps))
        exps[mod] = Dict()
    elseif arg in keys(exps[mod])
        return (exps[mod][arg], false)
    end

    subnet = apply(exp, arg)
    subrun = Runtime(subnet)
    # push all parent variables to subruntime
    set_state!(subrun, :parent_env, get_env(runtime))
    # we associate a subrun with the network, in case the same network is
    # produced for multiple arguments, so we avoid repeating computation for all
    #  the arguments that produce the same network
    subruns[subnet] = subrun
    ensure_all!(subrun)
    exps[mod][arg] = subnet
    return (subnet, true)
end

# analysis is 'per network', so we have to make sure we track per-network
network_analysis = Dict{Network, Dict{Symbol, Any}}()

clear_analysis!() = empty!(network_analysis)

function has_analysis(net,sym)
    haskey(network_analysis, net) && haskey(network_analysis[net], sym)
end

get_analysis(net,sym) = network_analysis[net][sym]

function add_analysis!(net,sym,val)
    get!(network_analysis, net, Dict{Symbol, Any}())[sym] = val
end

function expander_probs(runtime::Runtime, fn::Function, v::Variable, depth::Int) where {I,O}
    net = get_network(runtime)
    par = get_parents(net, v)[1]
    parinst = current_instance(runtime, par)
    inst = current_instance(runtime, v)
    if has_range(runtime, parinst)
        parrange = get_range(runtime, parinst, depth)
    else
        parrange = []
    end
    if has_range(runtime, inst)
        range = get_range(runtime, inst, depth)
    else
        range = []
    end
    lowers = Float64[]
    uppers = Float64[]

    function make_unexpanded_row()
        for i in 1:length(range)
            push!(lowers, 0)
            push!(uppers, 1)
        end
    end

    for p in parrange
        if depth < 2 || !expanded(runtime, v, p)
            make_unexpanded_row()
        else
            subnet = expansion(runtime, v, p)
            subrun = subruntime(runtime, v, subnet)
            output = get_node(subnet, :out)
            need_to_compute = true
            if has_analysis(subnet, :depthsolution)
                (saved_depth, saved_range, solution) =
                    get_analysis(subnet, :depthsolution)
                # To be able to reuse the solution to define this node's CPD,
                # it must have the required depth and range
                if saved_depth >= depth && saved_range == range
                    need_to_compute = false
                    (lfact, ufact) = solution
                end
            end
            if need_to_compute
                order = topsort(get_initial_graph(subnet))
                ((lfact, ufact), _) = fn(subrun, order, [output], depth - 1)
                add_analysis!(subnet, :depthsolution,
                              (depth, range, (lfact, ufact)))
            end
            lsum = sum(lfact.entries)
            output_range = get_range(subrun, output, depth - 1)
            inds = indexin(range, output_range)
            for i in 1:length(range)
                if isempty(inds) || isnothing(inds[i])
                    # If this value is not in the output range, then it
                    # might not be in the range of the output, so has
                    # probability 0, or it might be in the range but not
                    # expanded yet. But it's probability cannot be more than
                    # 1 - lsum
                    push!(lowers, 0)
                    push!(uppers, 1 - lsum)
                else
                    if isempty(lfact.entries)
                        push!(lowers, 0)
                    else
                        push!(lowers, lfact.entries[inds[i]])
                    end
                    if isempty(ufact.entries)
                        push!(uppers, 0)
                    else
                        push!(uppers, ufact.entries[inds[i]])
                    end
                end
            end
        end
    end
    return (lowers, uppers)
end


text:
"""
A type storing the status at the end of a call to `ssa`:

- **termination_status** : whether the simulation stops at the final time (`finaltime`) or early due to zero propensity function (`zeroprop`)
- **nsteps** : the number of steps taken during the simulation.
"""

struct SSAStats
    termination_status :: String
    nsteps :: Int64
end

"""
A type storing the call to `ssa`:

- **initial_state** : a `Vector` of `Int64`, representing the initial states of the system.
- **F** : a `Function` or a callable type, which itself takes two arguments; x, a `Vector` of `Int64` representing the states, and parameters, a `Vector` of `Float64` representing the parameters of the system.
- **rates** : a `Matrix` of `Int64`, representing the transitions of the system, organised by row.
- **parameters** : a `Vector` of `Float64` representing the parameters of the system.
- **tf** : the final simulation time (`Float64`).
- **alg** : the algorithm used (`Symbol`, either `:gillespie`, `jensen`, or `tjc`).
- **tvc** : whether rates are time varying.
"""
struct SSAArgs{X, Ftype, N, P}
    initial_state :: X
                F :: Ftype
            rates :: N
       parameters :: P
               tf :: Float64
              alg :: Symbol
              tvc :: Bool
end

"
This type stores the output of `ssa`, and comprises of:

- **time** : a `Vector` of `Float64`, containing the times of simulated events.
- **data** : a `Matrix` of `Int64`, containing the simulated states.
- **stats** : an instance of `SSAStats`.
- **args** : arguments passed to `ssa`.

"
struct SSAResult
     time :: Vector{Float64}
     data :: Matrix{Int64}
    stats :: SSAStats
     args :: SSAArgs
end

"
This function is a substitute for `StatsBase.sample(wv::WeightVec)`,
which avoids recomputing the sum and size of the weight vector,
as well as a type conversion of the propensity vector.
It takes the following arguments:

- weights : an `Array{Float64, 1}`, representing propensity function weights.
- total : the sum of `w`.
- n : the length of `w`.

"
function propensity_sample(weights, s, n)
    t = rand() * s
    i = 1
    cw = weights[1]
    while cw < t && i < n
        i += 1
        @inbounds cw += weights[i]
    end
    return i
end

const pfsample = propensity_sample

"""
This function performs Gillespie's stochastic simulation algorithm. It takes the following arguments:

- initial_state: a `Vector` of `Int64`, representing the initial states of the system.

- F: a `Function` or a callable type, which itself takes two arguments;
     x, a `Vector` of `Int64` representing the states, and parameters,
     a `Vector` of `Float64` representing the parameters of the system.

- rates: a `Matrix` of `Int64`, representing the transitions of the system, organised by row.

- parameters: a `Vector` of `Float64` representing the parameters of the system.

- stop_time: the final simulation time (`Float64`).
"""
function gillespie(initial_state, propensity_function, rates, parameters, stop_time)

    # Args
    args = SSAArgs(initial_state, propensity_function, rates, parameters, stop_time, :gillespie, false)

    # Set up time array
    times = Float64[]
    time = 0.0
    push!(times, time)

    # Set up initial state
    nstates = length(initial_state)
    state = copy(initial_state')
    data = copy(Array(initial_state))

    # Number of propensity functions
    n_propensities = size(rates, 1)

    # Main loop
    termination_status = "finaltime"
    nsteps = 0

    while time <= stop_time
        propensity = propensity_function(state, parameters)

        # Update time
        total_propensity = sum(propensity)

        if total_propensity == 0.0
            termination_status = "zeroprop"
            break
        end

        # Time step
        time_step = rand(Exponential(1 / total_propensity))
        time += time_step
        push!(times, time)

        # Update event
        event = propensity_sample(propensity, total_propensity, n_propensities)

        if state isa SVector
            @inbounds state[1] += rates[event, :]
        else
            Δx = view(rates, event, :)

            for i in 1:nstates
                @inbounds state[1, i] += Δx[i]
            end
        end

        push!(data, state...)

        # Update nsteps
        nsteps += 1
    end

    statistics = SSAStats(termination_status, nsteps)
    data = transpose(reshape(data, length(state), nsteps+1))

    return SSAResult(times, data, statistics, args)
end

include("truejump.jl")
include("jensen.jl")


text:
@testset "runge_kutta" begin
    function f(t, y, θ)
        return y
    end

    x0 = 1.0
    out = rk4(f, 1.0, 100, x0)
    @test abs(run(sess, out)[end]-exp(1))<1e-5
    out = ode45(f, 1.0, 100, x0)
    @test abs(run(sess, out)[end]-exp(1))<1e-5

    function f(t, y, θ)
        return y
    end

    x0 = [1.0;1.0]
    out = rk4(f, 1.0, 100, x0)
    @test norm(run(sess, out)[end,:] .- exp(1))<1e-5
    out = ode45(f, 1.0, 100, x0)
    @test norm(run(sess, out)[end,:] .- exp(1))<1e-5

    x0 = [1.0 1.0
          1.0 1.0]
    out = rk4(f, 1.0, 100, x0)
    @test norm(run(sess, out)[end,:,:] .- exp(1))<1e-5
    out = ode45(f, 1.0, 100, x0)
    @test norm(run(sess, out)[end,:,:] .- exp(1))<1e-5

end

@testset "alpha scheme" begin
    d0 = zeros(2)
    v0 = [2.;3.]
    a0 = zeros(2)
    Δt = 1e-3
    n = 9001
    Δt = Δt * ones(n)
    M = sparse([1.0 0.0;0.0 1.0])
    C = sparse(zeros(2,2))
    K = sparse([4.0 0.0;0.0 9.0])
    F = zeros(n, 2)

    d, v, a = αscheme(M, C, K, F, d0, v0, a0, Δt)
    d_, v_, a_ = run(sess, [d, v, a])

    tspan = 0:1e-3:(n-1)*1e-3
    @test norm(d_[1:end-1,1] - sin.(2tspan))<1e-3
    @test norm(d_[1:end-1,2] - sin.(3tspan))<1e-3
end


text:
module Domains1D

using ..Boundaries
using ..DielectricFunctions
using ..PumpFunctions
using ..Dispersions
using ..Lattices
using ..Points
using ..Shapes
using RecipesBase

import ..Symmetric, ..Unsymmetric
import LinearAlgebra.norm

import ..LatticeDomain

function LatticeDomain(
            boundary::Boundary{1},
            lattice::Lattice{1},
            n::Number = 1;
            type::Symbol = :generic,
            name::Symbol = :anonymous,
            fit::Bool = true
            )

    bnd = boundary
    lat = lattice
    if fit
        nx = ceil(Int,bnd.shape.a/lat.dx)
        dx = bnd.shape.a/nx
        lat = Lattices.Cartesian(dx; origin=bnd.shape.start+dx/2)
    end

    i1 = latticeindex(lat,bnd.shape.start)
    i2 = latticeindex(lat,bnd.shape.stop)
    imin, imax = floor(Int,min(i1,i2)), floor(Int,max(i1,i2))
    x = lat[imin:imax]
    inds_keep = bnd.shape.(x)
    x = x[inds_keep]
    indices = map(CartesianIndex,(imin:imax)[inds_keep])
    interior = trues(size(x))
    surface = falses(size(x))
    surface[1] = surface[end] = true
    bulk = interior .& .!surface
    corner = falses(size(x))
    nnm = (vcat(0,1:length(x)-1),)
    nnp = (vcat(2:length(x),0),)

    return LatticeDomain{Symmetric}(boundary, lat, n, x, indices, interior,
                            bulk, surface, corner, nnm, nnp; type=type, name=name)
end

function Base.propertynames(::LatticeDomain{1}, private=false)
    if private
        return fieldnames(LatticeDomain)
    else
        return (:boundary, :lattice, :shape, :n, :ε, :type, :name, :x)
    end
end

end

using .Domains1D


text:
"""
    Descent(η = 1f-1)

Classic gradient descent optimiser with learning rate `η`.
For each parameter `p` and its gradient `dp`, this runs `p -= η*dp`.

# Parameters
- Learning rate (`η`): Amount by which gradients are discounted before updating
                       the weights.
"""
struct Descent{T}
  eta::T
end
Descent() = Descent(1f-1)

init(o::Descent, x::AbstractArray) = nothing

function apply!(o::Descent, state, x, dx)
  η = convert(float(eltype(x)), o.eta)

  return state, @lazy dx * η  # @lazy creates a Broadcasted, will later fuse with x .= x .- dx
end

"""
    Momentum(η = 1f-2, ρ = 9f-1)

Gradient descent optimizer with learning rate `η` and momentum `ρ`.

# Parameters
- Learning rate (`η`): Amount by which gradients are discounted before updating
                       the weights.
- Momentum (`ρ`): Controls the acceleration of gradient descent in the
                  prominent direction, in effect dampening oscillations.
"""
struct Momentum{T}
  eta::T
  rho::T
end
Momentum(η = 1f-2, ρ = 9f-1) = Momentum{typeof(η)}(η, ρ)

init(o::Momentum, x::AbstractArray) = zero(x)

function apply!(o::Momentum, state, x, dx)
  η, ρ, mvel = o.eta, o.rho, state
  @.. mvel = ρ * mvel + η * dx  # Macro @.. broadcasts into mvel if it can, else @. of rhs.

  return mvel, mvel
end

"""
    Nesterov(η = 1f-3, ρ = 9f-1)

Gradient descent optimizer with learning rate `η` and Nesterov momentum `ρ`.

# Parameters
- Learning rate (`η`): Amount by which gradients are discounted before updating
                       the weights.
- Nesterov momentum (`ρ`): Controls the acceleration of gradient descent in the
                           prominent direction, in effect dampening oscillations.
"""
struct Nesterov{T}
  eta::T
  rho::T
end
Nesterov(η = 1f-3, ρ = 9f-1) = Nesterov{typeof(η)}(η, ρ)

init(o::Nesterov, x::AbstractArray) = zero(x)

function apply!(o::Nesterov, state, x, dx)
  η, ρ, vel = o.eta, o.rho, state

  newdx = @. - ρ^2 * vel + (1+ρ) * η * dx  # Cannot be lazy as this needs the old velocity
  @.. vel = ρ * vel - η * dx

  return vel, newdx
end

"""
    RMSProp(η = 1f-3, ρ = 9f-1, ϵ = eps(typeof(η)))

Optimizer using the
[RMSProp](https://www.cs.toronto.edu/~tijmen/csc321/slides/lecture_slides_lec6.pdf)
algorithm. Often a good choice for recurrent networks. Parameters other than learning rate
generally don't need tuning.

# Parameters
- Learning rate (`η`): Amount by which gradients are discounted before updating
                       the weights.
- Momentum (`ρ`): Controls the acceleration of gradient descent in the
                  prominent direction, in effect dampening oscillations.
- Machine epsilon (`ϵ`): Constant to prevent division by zero
                         (no need to change default)
"""
struct RMSProp{T}
  eta::T
  rho::T
  epsilon::T
end
RMSProp(η = 1f-3, ρ = 9f-1, ϵ = eps(typeof(η))) = RMSProp{typeof(η)}(η, ρ, ϵ)

init(o::RMSProp, x::AbstractArray) = zero(x)

function apply!(o::RMSProp, state, x, dx)
  η, ρ, ϵ, acc = o.eta, o.rho, o.epsilon, state

  @.. acc = ρ * acc + (1 - ρ) * abs2(dx)
  dx′ = @lazy dx * (η / (sqrt(acc) + ϵ))

  return acc, dx′
end

"""
    ADAM(η = 1f-3, β = (9f-1, 9.99f-1), ϵ = eps(typeof(η)))

[ADAM](https://arxiv.org/abs/1412.6980) optimiser.

# Parameters
- Learning rate (`η`): Amount by which gradients are discounted before updating
                       the weights.
- Decay of momentums (`β::Tuple`): Exponential decay for the first (β1) and the
                                   second (β2) momentum estimate.
- Machine epsilon (`ϵ`): Constant to prevent division by zero
                         (no need to change default)
"""
struct ADAM{T}
  eta::T
  beta::Tuple{T, T}
  epsilon::T
end
ADAM(η = 1f-3, β = (9f-1, 9.99f-1), ϵ = eps(typeof(η))) = ADAM{typeof(η)}(η, β, ϵ)

init(o::ADAM, x::AbstractArray) = (zero(x), zero(x), o.beta)

function apply!(o::ADAM, state, x, dx)
  η, β, ϵ = o.eta, o.beta, o.epsilon
  mt, vt, βt = state

  @.. mt = β[1] * mt + (1 - β[1]) * dx
  @.. vt = β[2] * vt + (1 - β[2]) * abs2(dx)
  dx′ = @lazy mt / (1 - βt[1]) / (sqrt(vt / (1 - βt[2])) + ϵ) * η

  return (mt, vt, βt .* β), dx′
end

"""
    RADAM(η = 1f-3, β = (9f-1, 9.99f-1), ϵ = eps(typeof(η)))

[Rectified ADAM](https://arxiv.org/abs/1908.03265) optimizer.

# Parameters
- Learning rate (`η`): Amount by which gradients are discounted before updating
                       the weights.
- Decay of momentums (`β::Tuple`): Exponential decay for the first (β1) and the
                                   second (β2) momentum estimate.
- Machine epsilon (`ϵ`): Constant to prevent division by zero
                         (no need to change default)
"""
struct RADAM{T}
  eta::T
  beta::Tuple{T, T}
  epsilon::T
end
RADAM(η = 1f-3, β = (9f-1, 9.99f-1), ϵ = eps(typeof(η))) = RADAM{typeof(η)}(η, β, ϵ)

init(o::RADAM, x::AbstractArray) = (zero(x), zero(x), o.beta, 1)

function apply!(o::RADAM, state, x, dx)
  η, β, ϵ = o.eta, o.beta, o.epsilon
  ρ∞ = 2/(1-β[2])-1

  mt, vt, βt, t = state

  @.. mt = β[1] * mt + (1 - β[1]) * dx
  @.. vt = β[2] * vt + (1 - β[2]) * abs2(dx)
  ρ = ρ∞ - 2*t * βt[2] / (1 - βt[2])
  if ρ > 4
    r = sqrt((ρ - 4) * (ρ - 2) * ρ∞/((ρ∞ - 4) * (ρ∞ - 2) * ρ))
    dx′ = @lazy mt / (1 - βt[1]) / (sqrt(vt / (1 - βt[2])) + ϵ) * η * r
  else
    dx′ = @lazy mt / (1 - βt[1]) * η
  end

  return (mt, vt, βt .* β, t + 1), dx′
end

"""
    AdaMax(η = 1f-3, β = (9f-1, 9.99f-1), ϵ = eps(typeof(η)))

[AdaMax](https://arxiv.org/abs/1412.6980) is a variant of ADAM based on the ∞-norm.

# Parameters
- Learning rate (`η`): Amount by which gradients are discounted before updating
                       the weights.
- Decay of momentums (`β::Tuple`): Exponential decay for the first (β1) and the
                                   second (β2) momentum estimate.
- Machine epsilon (`ϵ`): Constant to prevent division by zero
                         (no need to change default)
"""
struct AdaMax{T}
  eta::T
  beta::Tuple{T, T}
  epsilon::T
end
AdaMax(η = 1f-3, β = (9f-1, 9.99f-1), ϵ = eps(typeof(η))) = AdaMax{typeof(η)}(η, β, ϵ)

init(o::AdaMax, x::AbstractArray) = (zero(x), zero(x), o.beta)

function apply!(o::AdaMax, state, x, dx)
  η, β, ϵ = o.eta, o.beta, o.epsilon
  mt, ut, βt = state

  @.. mt = β[1] * mt + (1 - β[1]) * dx
  @.. ut = max(β[2] * ut, abs(dx))
  dx′ = @lazy (η/(1 - βt[1])) * mt/(ut + ϵ)

  return (mt, ut, βt .* β), dx′
end

"""
    OADAM(η = 1f-3, β = (5f-1, 9f-1), ϵ = eps(typeof(η)))

[OADAM](https://arxiv.org/abs/1711.00141) (Optimistic ADAM)
is a variant of ADAM adding an "optimistic" term suitable for adversarial training.

# Parameters
- Learning rate (`η`): Amount by which gradients are discounted before updating
                       the weights.
- Decay of momentums (`β::Tuple`): Exponential decay for the first (β1) and the
                                   second (β2) momentum estimate.
- Machine epsilon (`ϵ`): Constant to prevent division by zero
                         (no need to change default)
"""
struct OADAM{T}
  eta::T
  beta::Tuple{T, T}
  epsilon::T
end
OADAM(η = 1f-3, β = (5f-1, 9f-1), ϵ = eps(typeof(η))) = OADAM{typeof(η)}(η, β, ϵ)

init(o::OADAM, x::AbstractArray) = (zero(x), zero(x), o.beta, zero(x))

function apply!(o::OADAM, state, x, dx)
  η, β, ϵ = o.eta, o.beta, o.epsilon
  mt, vt, βt, term = state

  @.. mt = β[1] * mt + (1 - β[1]) * dx
  @.. vt = β[2] * vt + (1 - β[2]) * abs2(dx)
  prev = copy(term)
  @.. term = η * mt / (1 - βt[1]) / (sqrt(vt / (1 - βt[2])) + ϵ)
  dx′ = @lazy 2 * term - prev

  return (mt, vt, βt .* β, term), dx′
end

"""
    ADAGrad(η = 1f-1, ϵ = eps(typeof(η)))

[ADAGrad](http://www.jmlr.org/papers/volume12/duchi11a/duchi11a.pdf) optimizer. It has
parameter specific learning rates based on how frequently it is updated.
Parameters don't need tuning.

# Parameters
- Learning rate (`η`): Amount by which gradients are discounted before updating
                       the weights.
- Machine epsilon (`ϵ`): Constant to prevent division by zero
                         (no need to change default)
"""
struct ADAGrad{T}
  eta::T
  epsilon::T
end
ADAGrad(η = 1f-1, ϵ = eps(typeof(η))) = ADAGrad{typeof(η)}(η, ϵ)

init(o::ADAGrad, x::AbstractArray) = onevalue(o.epsilon, x)

function apply!(o::ADAGrad, state, x, dx)
  η, ϵ = o.eta, o.epsilon
  acc = state

  @.. acc = acc + abs2(dx)
  dx′ = @lazy dx * η / (sqrt(acc) + ϵ)

  return acc, dx′
end

"""
    ADADelta(ρ = 9f-1, ϵ = eps(typeof(ρ)))

[ADADelta](https://arxiv.org/abs/1212.5701) is a version of ADAGrad adapting its learning
rate based on a window of past gradient updates.
Parameters don't need tuning.

# Parameters
- Rho (`ρ`): Factor by which the gradient is decayed at each time step.
- Machine epsilon (`ϵ`): Constant to prevent division by zero
                         (no need to change default)
"""
struct ADADelta{T}
  rho::T
  epsilon::T
end
ADADelta(ρ = 9f-1, ϵ = eps(typeof(ρ))) = ADADelta{typeof(ρ)}(ρ, ϵ)

init(o::ADADelta, x::AbstractArray) = (zero(x), zero(x))

function apply!(o::ADADelta, state, x, dx)
  ρ, ϵ = o.rho, o.epsilon
  acc, Δacc = state

  @.. acc = ρ * acc + (1 - ρ) * abs2(dx)
  # DON'T remove epsilon from numerator or even out of the square roots!
  dx′ = @. dx * sqrt(Δacc + ϵ) / sqrt(acc + ϵ)  # Cannot be lazy as this needs the old Δacc
  @.. Δacc = ρ * Δacc + (1 - ρ) * abs2(dx′)

  return (acc, Δacc), dx′
end

"""
    AMSGrad(η = 1f-3, β = (9f-1, 9.99f-1), ϵ = eps(typeof(η)))

The [AMSGrad](https://openreview.net/forum?id=ryQu7f-RZ) version of the ADAM
optimiser. Parameters don't need tuning.

# Parameters
- Learning rate (`η`): Amount by which gradients are discounted before updating
                       the weights.
- Decay of momentums (`β::Tuple`): Exponential decay for the first (β1) and the
                                   second (β2) momentum estimate.
- Machine epsilon (`ϵ`): Constant to prevent division by zero
                         (no need to change default)
"""
struct AMSGrad{T}
  eta::T
  beta::Tuple{T, T}
  epsilon::T
end
AMSGrad(η = 1f-3, β = (9f-1, 9.99f-1), ϵ = eps(typeof(η))) = AMSGrad{typeof(η)}(η, β, ϵ)

init(o::AMSGrad, x::AbstractArray) =
  (onevalue(o.epsilon, x), onevalue(o.epsilon, x), onevalue(o.epsilon, x))

function apply!(o::AMSGrad, state, x, dx)
  η, β, ϵ = o.eta, o.beta, o.epsilon
  mt, vt, v̂t = state

  @.. mt = β[1] * mt + (1 - β[1]) * dx
  @.. vt = β[2] * vt + (1 - β[2]) * abs2(dx)
  @.. v̂t = max(v̂t, vt)
  dx′ = @lazy η * mt / (sqrt(v̂t) + ϵ)

  return (mt, vt, v̂t), dx′
end

"""
    NADAM(η = 1f-3, β = (9f-1, 9.99f-1), ϵ = eps(typeof(η)))

[NADAM](https://openreview.net/forum?id=OM0jvwB8jIp57ZJjtNEZ) is a Nesterov variant of ADAM.
Parameters don't need tuning.

# Parameters
- Learning rate (`η`): Amount by which gradients are discounted before updating
                       the weights.
- Decay of momentums (`β::Tuple`): Exponential decay for the first (β1) and the
                                   second (β2) momentum estimate.
- Machine epsilon (`ϵ`): Constant to prevent division by zero
                         (no need to change default)
"""
struct NADAM{T}
  eta::T
  beta::Tuple{T, T}
  epsilon::T
end
NADAM(η = 1f-3, β = (9f-1, 9.99f-1), ϵ = eps(typeof(η))) = NADAM{typeof(η)}(η, β, ϵ)

init(o::NADAM, x::AbstractArray) = (zero(x), zero(x), o.beta)

function apply!(o::NADAM, state, x, dx)
  η, β, ϵ = o.eta, o.beta, o.epsilon

  mt, vt, βt = state

  @.. mt = β[1] * mt + (1 - β[1]) * dx
  @.. vt = β[2] * vt + (1 - β[2]) * abs2(dx)
  dx′ = @lazy (β[1] * mt / (1 - β[1] * βt[1]) + (1 - β[1]) * dx / (1 - βt[1])) /
          (sqrt(vt * β[2] / (1 - βt[2])) + ϵ) * η

  return (mt, vt, βt .* β), dx′
end

"""
    ADAMW(η = 1f-3, β = (9f-1, 9.99f-1), γ = 0, ϵ = eps(typeof(η)))

[ADAMW](https://arxiv.org/abs/1711.05101) is a variant of ADAM fixing (as in repairing) its
weight decay regularization.

# Parameters
- Learning rate (`η`): Amount by which gradients are discounted before updating
                       the weights.
- Decay of momentums (`β::Tuple`): Exponential decay for the first (β1) and the
                                   second (β2) momentum estimate.
- Weight decay (`γ`): Decay applied to weights during optimisation.
- Machine epsilon (`ϵ`): Constant to prevent division by zero
                         (no need to change default)
"""
ADAMW(η = 1f-3, β = (9f-1, 9.99f-1), γ = 0, ϵ = eps(typeof(η))) =
  OptimiserChain(ADAM{typeof(η)}(η, β, ϵ), WeightDecay{typeof(η)}(γ))

"""
    AdaBelief(η = 1f-3, β = (9f-1, 9.99f-1), ϵ = eps(typeof(η)))

The [AdaBelief](https://arxiv.org/abs/2010.07468) optimiser is a variant of the well-known
ADAM optimiser.

# Parameters
- Learning rate (`η`): Amount by which gradients are discounted before updating
                       the weights.
- Decay of momentums (`β::Tuple`): Exponential decay for the first (β1) and the
                                   second (β2) momentum estimate.
- Machine epsilon (`ϵ::Float32`): Constant to prevent division by zero
                                  (no need to change default)
"""
struct AdaBelief{T}
  eta::T
  beta::Tuple{T, T}
  epsilon::T
end
AdaBelief(η = 1f-3, β = (9f-1, 9.99f-1), ϵ = eps(typeof(η))) = AdaBelief{typeof(η)}(η, β, ϵ)

init(o::AdaBelief, x::AbstractArray) = (zero(x), zero(x))

function apply!(o::AdaBelief, state, x, dx)
  η, β, ϵ = o.eta, o.beta, o.epsilon
  mt, st = state

  @.. mt = β[1] * mt + (1 - β[1]) * dx
  @.. st = β[2] * st + (1 - β[2]) * abs2(dx - mt)
  dx′ = @lazy η * mt / (sqrt(st) + ϵ)

  return (mt, st), dx′
end

"""
    WeightDecay(γ = 5f-4)

Decay weights by ``γ``, that is, add `γ .* x` to the gradient `x̄` which will be
subtracted from `x`.

Typically composed  with other optimisers as the first transformation in an [`OptimiserChain`](@ref).
This is equivalent to adding ``L_2`` regularization with coefficient ``γ`` to the loss.

# Parameters
- Weight decay (`γ`): Decay applied to weights during optimisation.
"""
struct WeightDecay{T}
  gamma::T
end
WeightDecay() = WeightDecay(5f-4)

init(o::WeightDecay, x::AbstractArray) = nothing

function apply!(o::WeightDecay, state, x, dx)
  dx′ = @lazy dx + o.gamma * x

  return state, dx′
end

"""
    ClipGrad(δ = 10f0)

Restricts every gradient component to obey `-δ ≤ dx[i] ≤ δ`.

See also [`ClipNorm`](@ref).
"""
struct ClipGrad{T<:Real}
  delta::T
end
ClipGrad() = ClipGrad(10f0)

init(o::ClipGrad, x::AbstractArray) = nothing

function apply!(o::ClipGrad, state, x, dx)
  δ = convert(float(eltype(x)), o.delta)
  dx′ = @lazy clamp(dx, -δ, δ)

  return state, dx′
end

"""
    ClipNorm(ω = 10f0, p = 2; throw = true)

Scales any gradient array for which `norm(dx, p) > ω`
to stay at this threshold (unless `p==0`).

Throws an error if the norm is infinite or `NaN`,
which you can turn off with `throw = false`.

See also [`ClipGrad`](@ref).
"""
struct ClipNorm{T<:Real}
  omega::T
  p::T
  throw::Bool
end
ClipNorm(ω = 10f0, p = 2; throw::Bool = true) = ClipNorm{typeof(ω)}(ω, p, throw)

init(o::ClipNorm, x::AbstractArray) = nothing

function apply!(o::ClipNorm, state, x, dx)
  nrm = norm(dx, o.p)
  if o.throw && !isfinite(nrm)
    throw(DomainError("gradient has $(o.p)-norm $nrm, for array $(summary(x))"))
  end
  λ = min(o.omega / nrm, 1)

  return state, @lazy dx * λ
end

"""
    OptimiserChain(opts...)

Compose a sequence of optimisers so that each `opt` in `opts`
updates the gradient, in the order specified.

With an empty sequence, `OptimiserChain()` is the identity,
so `update!` will subtract the full gradient from the parameters.
This is equivalent to `Descent(1)`.

# Example
```jldoctest
julia> o = OptimiserChain(ClipGrad(1), Descent(0.1));

julia> m = (zeros(3),);

julia> s = Optimisers.setup(o, m)
(Leaf(OptimiserChain(ClipGrad{Int64}(1), Descent{Float64}(0.1)), [nothing, nothing]),)

julia> Optimisers.update(s, m, ([0.3, 1, 7],))[2]  # clips before discounting
([-0.03, -0.1, -0.1],)
```
"""
struct OptimiserChain{O<:Tuple}
  opts::O
end
OptimiserChain(opts...) = OptimiserChain(opts)

init(o::OptimiserChain, x::AbstractArray) = [init(opt, x) for opt in o.opts]

function apply!(o::OptimiserChain, states, x, dx, dxs...)
  new_states = similar(states)
  for (i, (opt, state)) in enumerate(zip(o.opts, states))
    new_states[i], dx = apply!(opt, state, x, dx, dxs...)
  end

  return new_states, dx
end

function Base.show(io::IO, c::OptimiserChain)
  print(io, "OptimiserChain(")
  join(io, c.opts, ", ")
  print(io, ")")
end


text:
const allow_variational_hidden_dropout_mask_reset = Ref(false)

@inline is_variational_hidden_dropout_mask_reset_allowed() = allow_variational_hidden_dropout_mask_reset[]

@inline function _update_is_variational_hidden_dropout_mask_reset_allowed(val::Bool)
    return allow_variational_hidden_dropout_mask_reset[] = val
end

"""
    update_is_variational_hidden_dropout_mask_reset_allowed()
    update_is_variational_hidden_dropout_mask_reset_allowed(val::Bool)

Allow/Disallows variational_hidden_dropout mask reset.
"""
@inline function update_is_variational_hidden_dropout_mask_reset_allowed(val::Bool=!is_variational_hidden_dropout_mask_reset_allowed())
    return Zygote.hook(Δ -> begin
                           _update_is_variational_hidden_dropout_mask_reset_allowed(!val)
                           return Δ
                       end, _update_is_variational_hidden_dropout_mask_reset_allowed(val))
end

mutable struct VariationalHiddenDropout{F,M}
    p::F
    mask::M
    active::Union{Bool,Nothing}
end

Flux.trainable(::VariationalHiddenDropout) = ()

Flux.@functor VariationalHiddenDropout

"""
    VariationalHiddenDropout(p, s)

Like [`Flux.Dropout`](@ref) but with a mask which is not reset after each function call.
This is important for models like DeepEquilibriumModels where mask is reset
only after the solve call is complete.

# Arguments

* `p`: Probability of dropout.
* `s`: Size of the Input. Pass `1` as the last value of the tuple

See also [`update_is_variational_hidden_dropout_mask_reset_allowed`](@ref)
"""
function VariationalHiddenDropout(p, s)
    @assert 0 ≤ p ≤ 1
    mask = zeros(Float32, s)
    vd = VariationalHiddenDropout(p, mask, nothing)
    reset_mask!(vd)
    return vd
end

function reset_mask!(a::VariationalHiddenDropout)
    Flux.rand!(a.mask)
    return a.mask .= Flux._dropout_kernel.(a.mask, a.p, 1 - a.p)
end

Zygote.@nograd reset_mask!

function (a::VariationalHiddenDropout)(x)
    Flux._isactive(a) || return x
    is_variational_hidden_dropout_mask_reset_allowed() && reset_mask!(a)
    return variational_hidden_dropout(x, a.mask; active=true)
end

function Flux.testmode!(m::VariationalHiddenDropout, mode=true)
    return (m.active = (isnothing(mode) || mode == :auto) ? nothing : !mode; m)
end

function Base.show(io::IO, d::VariationalHiddenDropout)
    print(io, "VariationalDropout(", d.p)
    print(io, ", size = $(repr(size(d.mask)))")
    return print(io, ")")
end

function variational_hidden_dropout(x, mask; active::Bool=true)
    active || return x
    return x .* mask
end

Zygote.@adjoint function variational_hidden_dropout(x, mask; active::Bool=true)
    active || return x, Δ -> (Δ, nothing)
    return x .* mask, Δ -> (Δ .* mask, nothing)
end


text:
# One step entry decision

make_test_entry_one_step(J, nc) =
    EntryOneStepSwitches{Float64}(nTypes = J, nColleges = nc);

function init_entry_decision(objId :: ObjectId,
    switches :: EntryOneStepSwitches{F1},
    st :: SymbolTable) where F1

    pEntryPref = init_entry_prefscale(switches, st);
    pvec = ParamVector(objId, [pEntryPref]);
    return EntryOneStep(objId, pvec, ModelParams.value(pEntryPref), switches)
end

# ------------


text:
function toeplitz(c,r)
    @assert c[1] == r[1]
    nc = length(c)
    nr = length(r)
    A  = zeros(eltype(c), nc, nr)
    A[:,1] = c
    A[1,:] = r
    for i in 2:nr
        A[2:end,i] = A[1:end-1,i-1]
    end
    A
end

obslength(y::AbstractMatrix) = size(y,1)
obslength(y::AbstractVector) = length(y[1])

Base.oftype(x::Vector{<:Vector}, y::Vector{<:Vector}) = y
Base.oftype(x::AbstractVector, y::Vector{<:Vector}) = getindex.(y, 1)
Base.oftype(x::Matrix, y::Vector{<:Vector}) = reduce(hcat,y)
Base.oftype(x::Matrix, y::Matrix) = y
Base.oftype(x::Vector{<:Vector}, y::Matrix) = [y[:,i] for i in 1:size(y,2)]
function Base.oftype(x::Number, y::AbstractArray)
    length(y) == 1 || throw(ArgumentError("Can only autoconvert one-element arrays"))
    y[]
end

rms(x::AbstractVector) = sqrt(mean(abs2,x))
sse(x::AbstractVector) = x⋅x
mse(x::AbstractVector) = sse(x)/length(x)

rms(x::AbstractMatrix) = sqrt.(mean(abs2.(x),dims=2))[:]
sse(x::AbstractMatrix) = sum(abs2,x,dims=2)[:]
modelfit(y,yh) = 100 * (1 .-rms(y.-yh)./rms(y.-mean(y,dims=2)))
modelfit(y::T,yh::T) where T <: AbstractVector = 100 * (1 .-rms(y.-yh)./rms(y.-mean(y)))
aic(x::AbstractVector,d) = log(sse(x)) .+ 2d/size(x,2)
const nrmse = modelfit


text:
# This file was generated, do not modify it. # hide
W()             # transforms all data
W(rows=test, )  # transforms only test data
W(X[3:4, :])    # transforms specific data


text:
export geogebra

"""
`geogebra(G,file_name)` takes a `SimpleGraph` and writes out a
script to produce a drawing of this graph in GeoGebra.

Here is the secret sauce to make this work.

* Run `geogebra(G,file_name)` to save the script to `file_name`. By
  default, the file name is `geogebra.txt`.
* Copy the contents of `file_name` to the clipboard.
* Create a new GeoGebra document.
* In the **Input** zone at the bottom, enter `Button[]` to create a new
  button.
* Right click the button and select the **Object Properties...** menu option.
* Go to the **Scripting** tab and paste in the copied commands.
* Save and close the properties window.
* Press the newly created button.

Some properties of the vertices can be specified in this function with
named parameters as follows:

* `vertex_labels`: If set to `true`, the vertices in the drawing have
  labels. Default is `false` (no labels drawn). Note that the color of
  the labels matches the color of the vertices, so if you set the
  color to `white` the labels will be invisible.
* `vertex_color`: Use this to specify the fill color of the
  vertices. The default is `"black"`.
* `vertex_colors`: This is a dictionary mapping vertices to color
  names. Vertices are given the color specified. If, however, a vertex
  is missing from this dictionary, then `vertex_color` is used
  instead.
* `vertex_size`: Use this to specify the radius of the vertices. The
  default is `3`.
"""
function geogebra(G::SimpleGraph,
                  file_name::String="geogebra.txt";
                  vertex_labels::Bool=false,
                  vertex_color::String="black",
                  vertex_colors::Dict=Dict(),
                  vertex_size::Int=3
                  )
    X = get_embedding_direct(G)
    VV = vlist(X.G)
    n = NV(X.G)
    F = open(file_name,"w")

    for i=1:n
        v = VV[i]
        vname = string(v)
        (x,y) = X.xy[v]
        x = round(x,3)
        y = round(y,3)
        println(F,"v_{$i} = CopyFreeObject[Point[{$x,$y}]]")
        println(F,"SetPointSize[v_{$i}, $(string(vertex_size))]")

        color = ""
        try
            color = vertex_colors[v]
        catch
            color = vertex_color
        end
        println(F,"SetColor[v_{$i}, \"$color\"]")

        println(F,"SetCaption[v_{$i}, \"$vname\"]")
        println(F,"ShowLabel[v_{$i}, $(string(vertex_labels))]")
    end

    for i=1:n-1
        u = VV[i]
        for j=i+1:n
            v = VV[j]
            if has(X.G,u,v)
                println(F,"e_{$i,$j} = Segment[v_{$i},v_{$j}]")
                println(F,"ShowLabel[e_{$i,$j}, false]")
            end
        end
    end
    println(F, "ShowAxes[false]")
    println(F, "ShowGrid[false]")
    println(F, "Execute[{\"Delete[button1]\"}]")
    close(F)
end


text:
# Implements a simple interpreter for julia's lowered AST

getlhs(pc) = SSAValue(pc.next_stmt)

isassign(fr) = isassign(fr, fr.pc)
isassign(fr, pc) = (pc.next_stmt in fr.used)

lookup_var(frame, val::SSAValue) = frame.ssavalues[val.id+1]
lookup_var(frame, ref::GlobalRef) = getfield(ref.mod, ref.name)
lookup_var(frame, slot::SlotNumber) = something(frame.locals[slot.id])
function lookup_var(frame, e::Expr)
    isexpr(e, :the_exception) && return frame.last_exception[]
    isexpr(e, :boundscheck) && return true
    isexpr(e, :static_parameter) || error()
    frame.sparams[e.args[1]]
end

function finish!(frame)
    pc = frame.pc
    while true
        new_pc = _step_expr(frame, pc)
        new_pc == nothing && return pc
        pc = new_pc
    end
end

instantiate_type_in_env(arg, spsig, spvals) =
    ccall(:jl_instantiate_type_in_env, Any, (Any, Any, Ptr{Any}), arg, spsig, spvals)

function evaluate_call(frame, call_expr)
    args = Array{Any}(undef, length(call_expr.args))
    for i = 1:length(args)
        arg = call_expr.args[i]
        if isa(arg, QuoteNode)
            args[i] = arg.value
        elseif isa(arg, Union{SSAValue, GlobalRef, Slot})
            args[i] = lookup_var(frame, arg)
        elseif isexpr(arg, :&)
            args[i] = Expr(:&, lookup_var(frame, arg.args[1]))
        elseif isa(arg, Expr)
            args[i] = eval_rhs(frame, arg)
        else
            args[i] = arg
        end
    end
    # Don't go through eval since this may have unqouted, symbols and
    # exprs
    if isexpr(call_expr, :foreigncall)
        args = map(args) do arg
            isa(arg, Symbol) ? QuoteNode(arg) : arg
        end
        if !isempty(frame.sparams)
            args[2] = instantiate_type_in_env(args[2], frame.meth.sig, frame.sparams)
            args[3] = Core.svec(map(args[3]) do arg
                instantiate_type_in_env(arg, frame.meth.sig, frame.sparams)
            end...)
        end
        ret = Core.eval(frame.meth.module, Expr(:foreigncall, args...))
    else
        f = to_function(args[1])
        if isa(f, CodeInfo)
            ret = finish!(enter_call_expr(frame, call_expr))
        else
            # Don't go through eval since this may have unqouted, symbols and
            # exprs
            ret = f(args[2:end]...)
        end
    end
    return ret
end

function do_assignment!(frame, lhs, rhs)
    if isa(lhs, SSAValue)
        frame.ssavalues[lhs.id+1] = rhs
    elseif isa(lhs, Slot)
        frame.locals[lhs.id] = Some(rhs)
        frame.last_reference[frame.code.slotnames[lhs.id]] =
            lhs.id
    elseif isa(lhs, GlobalRef)
        Base.eval(lhs.mod,:($(lhs.name) = $(QuoteNode(rhs))))
    end
end

eval_rhs(frame, node) = eval(node)
function eval_rhs(frame, node::Expr)
    if isexpr(node, :new)
        new_expr = Expr(:new, map(x->QuoteNode(lookup_var_if_var(frame, x)),
            node.args)...)
        rhs = Core.eval(frame.meth.module, new_expr)
    elseif isexpr(node, :isdefined)
        rhs = check_isdefined(frame, node.args[1])
    else
        rhs = (isexpr(node, :call) || isexpr(node, :foreigncall)) ?
            evaluate_call(frame, node) :
            lookup_var_if_var(frame, node)
    end
    if isa(rhs, QuoteNode)
        rhs = rhs.value
    end
    return rhs
end

eval_rhs(frame, node::Union{SSAValue, GlobalRef, SlotNumber}) = lookup_var(frame, node)
eval_rhs(frame, node::QuoteNode) = node.value
check_isdefined(frame, node::Slot) = isdefined(frame.locals, slot.id)
function check_isdefined(frame, node::Expr)
    node.head == :static_parameter && return isdefined(frame.sparams, node.args[1])
end

function _step_expr(frame, pc)
    node = pc_expr(frame, pc)
    try
        if isassign(frame, pc)
            lhs = getlhs(pc)
            rhs = eval_rhs(frame, node)
            do_assignment!(frame, lhs, rhs)
        elseif isa(node, Expr)
            if node.head == :(=)
                lhs = node.args[1]
                rhs = eval_rhs(frame, node.args[2])
                do_assignment!(frame, lhs, rhs)
                # Special case hack for readability.
                # ret = rhs
            elseif node.head == :&
            elseif node.head == :gotoifnot
                arg = eval_rhs(frame, node.args[1])
                if !isa(arg, Bool)
                    throw(TypeError(frame.meth.name, "if", Bool, node.args[1]))
                end
                if !arg
                    return JuliaProgramCounter(node.args[2])
                end
            elseif node.head == :call || node.head == :foreigncall
                evaluate_call(frame, node)
            elseif node.head == :static_typeof
            elseif node.head == :type_goto || node.head == :inbounds
            elseif node.head == :enter
                push!(frame.exception_frames, node.args[1])
            elseif node.head == :leave
                for _ = 1:node.args[1]
                    pop!(frame.exception_frames)
                end
            elseif node.head == :static_parameter
            elseif node.head == :gc_preserve_end || node.head == :gc_preserve_begin
            elseif node.head == :return
                return nothing
            else
                ret = eval(node)
            end
        elseif isa(node, GotoNode)
            return JuliaProgramCounter(node.label)
        elseif isa(node, QuoteNode)
            ret = node.value
        else
            ret = eval_rhs(frame, node)
        end
    catch err
        isempty(frame.exception_frames) && rethrow(err)
        frame.last_exception[] = err
        return JuliaProgramCounter(frame.exception_frames[end])
    end
    return JuliaProgramCounter(pc.next_stmt + 1)
end
step_expr(frame) = _step_expr(frame, frame.pc)

function is_call(node)
    isexpr(node, :call) ||
    (isexpr(node, :(=)) && isexpr(node.args[2], :call))
end

function next_until!(f, frame, pc=frame.pc)
    while (pc = _step_expr(frame, pc)) != nothing
        f(pc_expr(frame, pc)) && return pc
    end
    return nothing
end
next_call!(frame, pc=frame.pc) = next_until!(node->is_call(node)||isexpr(node,:return), frame, pc)

function changed_line!(expr, line, fls)
    if length(fls) == 1 && isa(expr, LineNumberNode)
        return expr.line != line
    elseif length(fls) == 1 && isa(expr, Expr) && isexpr(expr, :line)
        return expr.args[1] != line
    else
        if is_loc_meta(expr, :pop_loc)
            pop!(fls)
        elseif is_loc_meta(expr, :push_loc)
            push!(fls,(expr.args[2],0))
        end
        return false
    end
end

isgotonode(node) = isa(node, GotoNode) || isexpr(node, :gotoifnot)

"""
Determine whether we are calling a function for which the current function
is a wrapper (either because of optional arguments or becaue of keyword arguments).
"""
function iswrappercall(expr)
    isexpr(expr, :(=)) && (expr = expr.args[2])
    isexpr(expr, :call) && any(x->x==SlotNumber(1), expr.args)
end

pc_expr(frame, pc) = frame.code.code[pc.next_stmt]
pc_expr(frame) = pc_expr(frame, frame.pc)

function find_used(code::CodeInfo)
    used = BitSet()
    stmts = code.code
    for stmt in stmts
        Core.Compiler.scan_ssa_use!(push!, used, stmt)
    end
    return used
end

function maybe_next_call!(frame, pc)
    call_or_return(node) = is_call(node) || isexpr(node, :return)
    call_or_return(pc_expr(frame, pc)) ||
        (pc = next_until!(call_or_return, frame, pc))
    pc
end
maybe_next_call!(frame) = maybe_next_call!(frame, frame.pc)

location(frame) = location(frame, frame.pc)
location(frame, pc) = frame.code.codelocs[pc.next_stmt] + frame.meth.line - 1
function next_line!(frame, stack = nothing)
    initial = location(frame)
    first = true
    pc = frame.pc
    while location(frame, pc) == initial
        # If this is a return node, interrupt execution. This is the same
        # special case as in `s`.
        (!first && isexpr(pc_expr(frame, pc), :return)) && return pc
        first = false
        # If this is a goto node, step it and reevaluate
        if isgotonode(pc_expr(frame, pc))
            pc = _step_expr(frame, pc)
            pc == nothing && return nothing
        elseif stack !== nothing && iswrappercall(pc_expr(frame, pc))
            # With splatting it can happen that we do something like ssa = tuple(#self#), _apply(ssa), which
            # confuses the logic here, just step into the first call that's not a builtin
            while true
                stack[1] = JuliaStackFrame(frame, pc; wrapper = true)
                call_expr = pc_expr(frame, pc)
                isexpr(call_expr, :(=)) && (call_expr = call_expr.args[2])
                call_expr = Expr(:call, map(x->lookup_var_if_var(frame, x), call_expr.args)...)
                new_frame = enter_call_expr(call_expr)
                if new_frame !== nothing
                    pushfirst!(stack, new_frame)
                    frame = new_frame
                    pc = frame.pc
                    break
                else
                    pc = _step_expr(frame, pc)
                    pc == nothing && return nothing
                end
            end
        elseif isa(pc_expr(frame, pc), LineNumberNode)
            line != pc_expr(frame, pc).line && break
            pc = _step_expr(frame, pc)
        else
            pc = _step_expr(frame, pc)
            pc == nothing && return nothing
        end
    end
    maybe_next_call!(frame, pc)
end


text:
test = """
    """

function parse_input(input::AbstractString)
    lines = split(input, r"\s+"; keepempty = false)
    return
end

q1(v) = 0
q2(v) = 0

let v = parse_input(test)
    @assert q1(v) == 0
    @assert q2(v) == 0
end

v = parse_input(read("day00.in", String))
println("Q1: ", q1(v))
println("Q2: ", q2(v))

max_stars_repo_name: dhanak/competitive-coding

max_issues_repo_name: dhanak/competitive-coding

max_forks_repo_name: dhanak/competitive-coding


text:
# additional method to render the intersection environment with obstacles
function AutoViz.render!(rendermodel::RenderModel, env::IntersectionEnv)
    render!(rendermodel, env.roadway)

    for obs in env.obstacles
        pts = Array{Float64}(2, obs.npts)
        for (i, pt) in enumerate(obs.pts)
            pts[1,i] = pt.x
            pts[2,i] = pt.y
        end

        add_instruction!(rendermodel, render_fill_region, (pts, colorant"gray"))
    end

    stop_line = posg(Frenet(env.roadway[LaneTag(6,1)], env.params.stop_line), env.roadway)
    x_pos, y_pos = stop_line.x, stop_line.y
    stop_pts = zeros(2,2)
    stop_pts[1,:] =  [(x_pos - env.params.lane_width/2) , (x_pos + env.params.lane_width/2)]
    stop_pts[2,:] =  [y_pos, y_pos]
    add_instruction!(rendermodel, render_line, (stop_pts, colorant"white", 1.0, Cairo.CAIRO_LINE_CAP_BUTT))
    return rendermodel
end

function animate_record(rec::SceneRecord, env, overlays::Vector{SceneOverlay} = SceneOverlay[]; cam=FitToContentCamera(0.0), sim_dt::Float64=0.1)
    duration =rec.nframes*sim_dt
    fps = Int(1/sim_dt)
    function render_rec(t, dt)
        frame_index = Int(floor(t/dt)) + 1
        return AutoViz.render(rec[frame_index-nframes(rec)], overlays, env, cam=cam, car_colors=get_colors(scene))
    end
    return duration, fps, render_rec
end

function animate_hist(pomdp::OIPOMDP, hist; overlays = SceneOverlay[], cam = FitToContentCamera(0.))
    duration = n_steps(hist)*pomdp.ΔT
    fps = Int(1/pomdp.ΔT)
    function render_hist(t, dt)
        stateindex = Int(floor(t/dt)) + 1
        scene = state_hist(hist)[stateindex]
        return AutoViz.render(scene, pomdp.env, overlays, cam = cam,  car_colors=get_colors(scene))
    end
    return duration, fps, render_hist
end

max_stars_repo_name: MaximeBouton/AutomotivePOMDPs

max_issues_repo_name: MaximeBouton/AutomotivePOMDPs

max_forks_repo_name: MaximeBouton/AutomotivePOMDPs


text:
@inline function _add_edge!(graph::PharmGraph,
                            pair::Pair{PharmClass,PharmClass})
    g = graph.g
    from_class = pair[1]
    to_class = pair[2]
    from_vertex_integer = _class_to_vertex_integer!(graph, from_class)
    to_vertex_integer = _class_to_vertex_integer!(graph, to_class)
    MetaGraphs.add_edge!(g, from_vertex_integer, to_vertex_integer)
    return nothing
end


text:
using PyCall

function series_expansion(x)
    N=100
    cx=x
    sx=(pi/2)*(1/3)*x^3
    for k in 1:N
        cx*=-(4k+1)/((2k+1)*(4k+5))*(pi/2)^2*x^4
        sx*=-(4k+3)/(2k+3)*(pi/2)*x^4
    end
    return cx,sx
end

function continued_fraction(x)
    N=100
    cx=x
    sx=(pi/2)*(1/3)*x^3
    for k in 1:N
        cx*=-(4k+1)/((2k+1)*(4k+5))*(pi/2)^2*x^4
        sx*=-(4k+3)/(2k+3)*(pi/2)*x^4
    end
    return cx,sx
end

function fresnel(x)
    cx,sx=0.0,0.0

    if x==0
        return 0,0
    end

    xa=abs(x)
    if xa<2
        cx,sx=series_expansion(xa)
    elseif xa>=2
        cx,sx=continued_fraction(xa)
    end

    if x<0
        cx,sx=-cx,-sx
    end
    return cx,sx
end

function main_()
    n=1000
    h=0.001
    for i in -n:n
        x=i*h
        cx,sx=fresnel(x)
        print(x,cx,sx)
    end
end

function main()
    xs=linspace(-pi,pi,100)
    ys=sin.(xs)
    @pyimport matplotlib.pyplot as plt
    plt.plot(xs,ys)
    plt.show()
end
main()

max_stars_repo_name: terasakisatoshi/Fortran

max_issues_repo_name: terasakisatoshi/Fortran

max_forks_repo_name: terasakisatoshi/Fortran


text:
using Test, Pkg
using MPI
using Random

if get(ENV,"JULIA_MPI_TEST_ARRAYTYPE","") == "CuArray"
    using CuArrays
    ArrayType = CuArray
else
    ArrayType = Array
end

MPI.Init()
comm = MPI.COMM_WORLD
rank = MPI.Comm_rank(comm)
sz = MPI.Comm_size(comm)
filename = MPI.bcast(tempname(), 0, comm)

# Write
MPI.Barrier(comm)

f = MPI.File.open(comm, filename, write=true)
MPI.File.set_view!(f, 0, MPI.Datatype(Int64), MPI.Datatype(Int64))
MPI.File.write_at(f, rank*2, ArrayType([Int64(rank+1) for i = 1:2]))
close(f)

MPI.Barrier(comm)

if rank == 0
    @test read!(filename, zeros(Int64, (2,sz))) == [j for i = 1:2, j=1:sz]
end

MPI.Barrier(comm)

f = MPI.File.open(comm, filename, write=true)
MPI.File.set_view!(f, 0, MPI.Datatype(Int64), MPI.Datatype(Int64))
MPI.File.write_at_all(f, rank*2, ArrayType([Int64(rank+1) for i = 1:2]))
close(f)

MPI.Barrier(comm)

if rank == 0
    @test read!(filename, zeros(Int64, (2,sz))) == [j for i = 1:2, j=1:sz]
end

MPI.Barrier(comm)

# Read
if rank == 0
    write(filename, [Float64(j) for i = 1:3, j = 1:sz])
end

MPI.Barrier(comm)

f = MPI.File.open(comm, filename, read=true)
MPI.File.set_view!(f, 0, MPI.Datatype(Float64), MPI.Datatype(Float64))

data = ArrayType(zeros(Float64, 3))
MPI.File.read_at!(f, rank*3, data)
@test data == Float64[rank+1 for i = 1:3]

MPI.Barrier(comm)

data = ArrayType(zeros(Float64, 3))
MPI.File.read_at_all!(f, rank*3, data)
@test data == Float64[rank+1 for i = 1:3]