# using Revise, Plots
using OrdinaryDiffEq, ForwardDiff, Test
	using BifurcationKit, LinearAlgebra, Parameters, Setfield
	const BK = BifurcationKit
	const FD = ForwardDiff

norminf = x -> norm(x, Inf)

function Fsl!(f, u, p, t = 0)
	@unpack r, μ, ν, c3, c5 = p
	u1 = u[1]
	u2 = u[2]

	ua = u1^2 + u2^2

	f[1] = r * u1 - ν * u2 - ua * (c3 * u1 - μ * u2) - c5 * ua^2 * u1
	f[2] = r * u2 + ν * u1 - ua * (c3 * u2 + μ * u1) - c5 * ua^2 * u2
	return f
end

Fsl(x, p) = Fsl!(similar(x), x, p)
dFsl(x, dx, p) = FD.derivative(t -> Fsl(x .+ t .* dx, p), 0.)
jet = BK.getJet(Fsl; matrixfree = false)

par_sl = (r = 0.5, μ = 0., ν = 1.0, c3 = 1.0, c5 = 0.0,)
par_hopf = (@set par_sl.r = 0.1)
u0 = [.001, .001]

function FslMono!(f, x, p, t)
	u = x[1:2]
	du = x[3:4]
	Fsl!(f[1:2], u, p, t)
	f[3:4] .= dFsl(u, du, p)
end
####################################################################################################
# continuation
optconteq = ContinuationPar(ds = -0.01, detectBifurcation = 3, pMin = -0.5, nInversion = 4)
br, = continuation(Fsl, u0, par_hopf, (@lens _.r), optconteq)
####################################################################################################
prob = ODEProblem(Fsl!, u0, (0., 100.), par_hopf)
probMono = ODEProblem(FslMono!, vcat(u0, u0), (0., 100.), par_hopf)
####################################################################################################
sol = solve(probMono, KenCarp4(), abstol =1e-9, reltol=1e-6)
sol = solve(prob, KenCarp4(), abstol =1e-9, reltol=1e-6)
# plot(sol[1,:], sol[2,:])
####################################################################################################
section(x, T) = x[1] #* x[end]
section(x, T, dx, dT) = dx[1] #* x[end]
# standard simple shooting
M = 1
dM = 1
_pb = ShootingProblem(prob, KenCarp4(), 1, section; abstol =1e-10, reltol=1e-9)

initpo = [0.13, 0., 6.]
res = _pb(initpo, par_hopf)

# test of the differential of the shooting method

_dx = rand(3)
resAD = FD.derivative(z -> _pb(initpo .+ z .* _dx, par_hopf), 0.)
resFD = (_pb(initpo .+ 1e-8 .* _dx, par_hopf) - _pb(initpo, par_hopf)) .* 1e8
resAN = _pb(initpo, par_hopf, _dx; δ = 1e-8)
@test norm(resAN - resFD, Inf) < 4e-6
@test norm(resAN - resAD, Inf) < 4e-6
####################################################################################################
# test shooting interface M = 1
_pb = ShootingProblem(prob, Rodas4(), [initpo[1:end-1]]; abstol =1e-10, reltol=1e-9)
res = _pb(initpo, par_hopf)
res = _pb(initpo, par_hopf, initpo)

# test the jacobian of the functional in the case M=1
_Jad = FD.jacobian( x -> _pb(x, par_hopf), initpo)
_Jana = _pb(Val(:JacobianMatrix), initpo, par_hopf)
@test norm(_Jad - _Jana, Inf) < 1e-7

_pb2 = ShootingProblem(prob, Rodas4(), probMono, Rodas4(autodiff=false), [initpo[1:end-1]]; abstol = 1e-10, reltol = 1e-9)
res = _pb2(initpo, par_hopf)
res = _pb2(initpo, par_hopf, initpo)
BK.isSimple(_pb2)

# we test this using Newton - Continuation
optn = NewtonPar(verbose = false, tol = 1e-9,  maxIter = 20)
# deflationOp = BK.DeflationOperator(2, (x,y) -> dot(x[1:end-1], y[1:end-1]), 1.0, [zeros(3)])
outpo, _, flag,_ = newton(_pb,
	initpo, par_hopf,
	optn;
	jacobianPO = :autodiffDense,
	normN = norminf)
	@test flag

BK.getPeriod(_pb, outpo, par_hopf)
BK.getAmplitude(_pb, outpo, par_hopf)
BK.getMaximum(_pb, outpo, par_hopf)
BK.getPeriodicOrbit(_pb, outpo, par_hopf)

opts_po_cont = ContinuationPar(dsmin = 0.001, dsmax = 0.01, ds= -0.01, pMax = 4.0, maxSteps = 30, detectBifurcation = 2, nev = 2, newtonOptions = (@set optn.tol = 1e-7), precisionStability = 1e-5)
	br_pok2, = continuation(_pb,
		outpo, par_hopf, (@lens _.r),
		opts_po_cont;
		tangentAlgo = BorderedPred(),
		verbosity = 0, plot = false,
		jacobianPO = :autodiffDense,
		recordFromSolution = (u, p) -> norm(u[1:2]),
		normC = norminf)
# plot(br_pok2)
####################################################################################################
# test automatic branch switching
br_pok2, = continuation(jet..., br, 1, opts_po_cont, ShootingProblem(1, prob, KenCarp4();  abstol = 1e-10, reltol = 1e-9); normC = norminf, jacobianPO = :autodiffDense)

# test matrix-free computation of floquet coefficients
eil = EigKrylovKit(dim = 2, x₀=rand(2))
opts_po_contMF = @set opts_po_cont.newtonOptions.eigsolver = eil
opts_po_contMF = @set opts_po_cont.detectBifurcation = 0
br_pok2, = continuation(jet...,br,1, opts_po_contMF, ShootingProblem(1, prob, Rodas4(); abstol = 1e-10, reltol = 1e-9); jacobianPO = :autodiffDense, normC = norminf, plot=false)

# case with 2 sections
br_pok2_s2, = continuation(jet..., br, 1, opts_po_cont, ShootingProblem(2, prob, KenCarp4();  abstol = 1e-10, reltol = 1e-9); normC = norminf, jacobianPO = :autodiffDense)

####################################################################################################
# test shooting interface M > 1
_pb = ShootingProblem(prob, KenCarp4(), [initpo[1:end-1],initpo[1:end-1],initpo[1:end-1]]; abstol =1e-10, reltol=1e-9)
initpo = [0.13, 0, 0, 0.13, 0, 0.13 , 6.3]
res = _pb(initpo, par_hopf)
res = _pb(initpo, par_hopf, initpo)
# test the jacobian of the functional in the case M=1
_Jad = FD.jacobian( x -> _pb(x, par_hopf), initpo)
_Jana = _pb(Val(:JacobianMatrix), initpo, par_hopf)
@test norm(_Jad - _Jana, Inf) < 1e-7
####################################################################################################
# Single Poincaré Shooting with hyperplane parametrization
normals = [[-1., 0.]]
centers = [zeros(2)]

probPsh = PoincareShootingProblem(2, prob, Rodas4(), probMono, Rodas4(); abstol=1e-10, reltol=1e-9)
probPsh = PoincareShootingProblem(2, prob, Rodas4(); rtol = abstol=1e-10, reltol=1e-9)

probPsh = PoincareShootingProblem(prob, Rodas4(),
		probMono, Rodas4(),
		normals, centers; abstol = 1e-10, reltol = 1e-9)

initpo_bar = BK.R(probPsh, [0, 0.4], 1)

BK.E(probPsh, [1.0,], 1)
initpo_bar = [0.4]

probPsh(initpo_bar, par_hopf)

# test of the analytical formula for jacobian of the functional
_Jad = BifurcationKit.finiteDifferences( x-> probPsh(x, par_hopf), initpo_bar)
_Jana = probPsh(Val(:JacobianMatrix), initpo_bar, par_hopf)
@test norm(_Jad - _Jana, Inf) < 1e-4

ls = DefaultLS()
	eil = EigKrylovKit(dim = 1, x₀ = rand(1))
	optn = NewtonPar(verbose = false, tol = 1e-8,  maxIter = 140, linsolver = ls, eigsolver = eil)
	deflationOp = BK.DeflationOperator(2, dot, 1.0, [zero(initpo_bar)])
	outpo, _, flag, = newton(probPsh,
			initpo_bar, par_hopf,
			optn; normN = norminf, jacobianPO = :autodiffDenseAnalytical)
	@test flag

BK.getPeriod(probPsh, outpo, par_hopf)
BK.getAmplitude(probPsh, outpo, par_hopf)
BK.getMaximum(probPsh, outpo, par_hopf)
BK.getPeriodicOrbit(probPsh, outpo, par_hopf)

probPsh = PoincareShootingProblem(prob, Rodas4(),
		# probMono, Rodas4(autodiff=false),
		normals, centers; abstol = 1e-10, reltol = 1e-9)

probPsh(outpo, par_hopf)
probPsh(outpo, par_hopf, outpo)
probPsh([0.30429879744900434], par_hopf)
probPsh([0.30429879744900434], (r = 0.09243096156871472, μ = 0.0, ν = 1.0, c3 = 1.0, c5 = 0.0))
BK.evolve(probPsh.flow,[0.0, 0.30429879744900434], (r = 0.094243096156871472, μ = 0.0, ν = 1.0, c3 = 1.0, c5 = 0.0), Inf64) # this gives an error in DiffEqBase

opts_po_cont = ContinuationPar(dsmin = 0.001, dsmax = 0.015, ds= 0.01, pMax = 4.0, maxSteps = 30, newtonOptions = setproperties(optn; tol = 1e-7, eigsolver = eil), detectBifurcation = 0)
	# br_pok2, = continuation(
	# 	probPsh, outpo, par_hopf, (@lens _.r),
	# 	opts_po_cont; verbosity = 2,
	# 	tangentAlgo = BorderedPred(),
	# 	jacobianPO = :autodiffDenseAnalytical,
	# 	plot = false, normC = norminf)
# plot(br_pok2)
####################################################################################################
# normals = [[-1., 0.], [1, -1]]
# centers = [zeros(2), zeros(2)]
# initpo_bar = [1.04, -1.04/√2]

normals = [[-1., 0.], [1, 0]]
centers = [zeros(2), zeros(2)]
initpo_bar = [0.2, -0.2]

probPsh = PoincareShootingProblem(prob, KenCarp4(), normals, centers; abstol=1e-10, reltol=1e-9)
# version with analytical jacobian
probPsh2 = PoincareShootingProblem(prob, KenCarp4(), normals, centers; abstol=1e-10, reltol=1e-9, δ = 0)

# test of the analytical formula for jacobian of the functional
_Jad = BifurcationKit.finiteDifferences( x-> probPsh(x, par_hopf), initpo_bar)
# _Jphifd = BifurcationKit.finiteDifferences(x->probPsh.flow(x, par_hopf, Inf64), [0,0.4]; δ=1e-8)
_Jana = probPsh(Val(:JacobianMatrix), initpo_bar, par_hopf)
@test norm(_Jad - _Jana, Inf) < 1e-5

ls = DefaultLS()
	eil = EigKrylovKit(dim = 1, x₀=rand(1))
	optn = NewtonPar(verbose = false, tol = 1e-9,  maxIter = 140, linsolver = ls, eigsolver = eil)
	deflationOp = DeflationOperator(2.0, dot, 1.0, [zero(initpo_bar)])
	outpo, = newton(probPsh2, initpo_bar, par_hopf, optn; normN = norminf, jacobianPO = :autodiffDenseAnalytical)
	outpo, = newton(probPsh, initpo_bar, par_hopf, optn; normN = norminf, jacobianPO = :autodiffDenseAnalytical)

getPeriod(probPsh, outpo, par_hopf)

opts_po_cont = ContinuationPar(dsmin = 0.0001, dsmax = 0.025, ds= -0.01, pMax = 4.0, maxSteps = 50, newtonOptions = (@set optn.tol = 1e-9), detectBifurcation = 3, nev = 2)
	br_pok2, = continuation(probPsh, outpo, par_hopf, (@lens _.r),
		opts_po_cont; verbosity = 0,
		tangentAlgo = BorderedPred(),
		jacobianPO = :autodiffDenseAnalytical,
		plot = false, normC = norminf)

####################################################################################################
normals = [[-1., 0.], [1, 0], [0, 1]]
centers = [zeros(2), zeros(2), zeros(2)]
initpo = [[0., 0.4], [0, -.3], [0.3, 0]]

probPsh = PoincareShootingProblem(prob, KenCarp4(), normals, centers; abstol=1e-10, reltol=1e-9)

initpo_bar = reduce(vcat, [BK.R(probPsh, initpo[ii], ii) for ii in eachindex(centers)])
# same with projection function
initpo_bar = reduce(vcat, BK.projection(probPsh, initpo))

# test of the other projection function
BK.projection(probPsh, reduce(hcat, initpo)')

probPsh(initpo_bar, par_hopf; verbose = true)

# test of the analytical formula for jacobian of the functional
_Jad = BifurcationKit.finiteDifferences( x -> probPsh(x, par_hopf), initpo_bar)
_Jana = probPsh(Val(:JacobianMatrix), initpo_bar, par_hopf)
@test norm(_Jad - _Jana, Inf) < 1e-5

outpo, = newton(probPsh, initpo_bar, par_hopf, optn; normN = norminf, jacobianPO = :autodiffDenseAnalytical)

for ii=1:length(normals)
	@show BK.E(probPsh	, [outpo[ii]], ii)
end

getPeriod(probPsh, outpo, par_hopf)

opts_po_cont = ContinuationPar(dsmin = 0.0001, dsmax = 0.025, ds= -0.005, pMax = 4.0, maxSteps = 10, newtonOptions = setproperties(optn; tol = 1e-8), detectBifurcation = 3)
	br_hpsh, upo , _= continuation(probPsh, outpo, par_hopf, (@lens _.r),
		opts_po_cont; normC = norminf, jacobianPO = :autodiffDenseAnalytical)
# plot(br_hpsh)
####################################################################################################
# test automatic branch switching with most possible options
# calls with analytical jacobians
br_psh, = continuation(jet..., br, 1, (@set opts_po_cont.ds = 0.005), PoincareShootingProblem(2, prob, KenCarp4(); abstol=1e-10, reltol=1e-9); normC = norminf, jacobianPO = :autodiffDenseAnalytical)

# test Iterative Floquet eigen solver
@set! opts_po_cont.newtonOptions.eigsolver.dim = 20
@set! opts_po_cont.newtonOptions.eigsolver.x₀ = rand(2)
br_sh, = continuation(jet..., br, 1, ContinuationPar(opts_po_cont; ds = 0.005, saveSolEveryStep = 1), ShootingProblem(2, prob, KenCarp4(); abstol=1e-10, reltol=1e-9); normC = norminf, jacobianPO = :autodiffDenseAnalytical)

# test MonodromyQaD
# BK.MonodromyQaD(br_sh.functional, br.sol)

ls = GMRESIterativeSolvers(reltol = 1e-7, N = length(initpo_bar), maxiter = 500, verbose = false)
@set! opts_po_cont.detectBifurcation = 0
@set! opts_po_cont.newtonOptions.linsolver = ls

for M in (1,2), jacobianPO in (:autodiffMF, :MatrixFree, :autodiffDenseAnalytical, :FiniteDifferencesDense)
	@info M, jacobianPO, "PS"

	# specific to Poincaré Shooting
	jacPO = jacobianPO == :autodiffMF ? :FiniteDifferencesDense : jacobianPO

	br_psh, = continuation(jet..., br, 1, (@set opts_po_cont.ds = 0.005), PoincareShootingProblem(M, prob, Rodas4P(); abstol=1e-10, reltol=1e-9, parallel = true); normC = norminf, updateSectionEveryStep = 2, linearAlgo = BorderingBLS(solver = (@set ls.N = M), checkPrecision = false), jacobianPO = jacPO, verbosity = 0)

	br_ssh, = continuation(jet..., br, 1, (@set opts_po_cont.ds = 0.005),
	ShootingProblem(M, prob, Rodas4P(); abstol=1e-10, reltol=1e-9, parallel = true); normC = norminf, updateSectionEveryStep = 2, jacobianPO = jacobianPO,
	linearAlgo = BorderingBLS(solver = (@set ls.N = 2M + 1), checkPrecision = false), verbosity = 0)
end


# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:light
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.4'
#       jupytext_version: 1.1.3
#   kernelspec:
#     display_name: Julia 1.1.1
#     language: julia
#     name: julia-1.1
# ---

# +
using Loess
using Random
using Statistics

Random.seed!(100)
xs = 10 .* rand(100)
ys = sin.(xs) .+ 0.5 * rand(100)

model = loess(xs, ys)

us = collect(minimum(xs):0.1:maximum(xs))
vs = predict(model, us);

# +
using Plots

plot( us, vs)
scatter!( xs, ys, markersize=2)
# -


# ---
# title: 273. Integer to English Words
# id: problem273
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Math, String
# link: <https://leetcode.com/problems/integer-to-english-words/description/>
# hidden: true
# ---
#
# Convert a non-negative integer `num` to its English words representation.
#
#
#
# **Example 1:**
#
#
#
#     Input: num = 123
#     Output: "One Hundred Twenty Three"
#
#
# **Example 2:**
#
#
#
#     Input: num = 12345
#     Output: "Twelve Thousand Three Hundred Forty Five"
#
#
# **Example 3:**
#
#
#
#     Input: num = 1234567
#     Output: "One Million Two Hundred Thirty Four Thousand Five Hundred Sixty Seven"
#
#
# **Example 4:**
#
#
#
#     Input: num = 1234567891
#     Output: "One Billion Two Hundred Thirty Four Million Five Hundred Sixty Seven Thousand Eight Hundred Ninety One"
#
#
#
#
# **Constraints:**
#
#   * `0 <= num <= 231 - 1`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


export lsqr, lsqr!

# Adapted from the BSD-licensed Matlab implementation at
#    http://www.stanford.edu/group/SOL/software/lsqr.html
#
# LSQR solves Ax = b or min ||b - Ax||_2 if damp = 0,
# or   min ||(b) - (  A   )x||   otherwise.
#          ||(0)   (damp*I) ||_2

#-----------------------------------------------------------------------
# LSQR uses an iterative (conjugate-gradient-like) method.
# For further information, see
# 1. C. C. Paige and M. A. Saunders (1982a).
#    LSQR: An algorithm for sparse linear equations and sparse least squares,
#    ACM TOMS 8(1), 43-71.
# 2. C. C. Paige and M. A. Saunders (1982b).
#    Algorithm 583.  LSQR: Sparse linear equations and least squares problems,
#    ACM TOMS 8(2), 195-209.
# 3. M. A. Saunders (1995).  Solution of sparse rectangular systems using
#    LSQR and CRAIG, BIT 35, 588-604.
#
# Input parameters:
# atol, btol  are stopping tolerances.  If both are 1.0e-9 (say),
#             the final residual norm should be accurate to about 9 digits.
#             (The final x will usually have fewer correct digits,
#             depending on cond(A) and the size of damp.)
# conlim      is also a stopping tolerance.  lsqr terminates if an estimate
#             of cond(A) exceeds conlim.  For compatible systems Ax = b,
#             conlim could be as large as 1.0e+12 (say).  For least-squares
#             problems, conlim should be less than 1.0e+8.
#             Maximum precision can be obtained by setting
#             atol = btol = conlim = zero, but the number of iterations
#             may then be excessive.
# maxiter      is an explicit limit on iterations (for safety).
#
#              Michael Saunders, Systems Optimization Laboratory,
#              Dept of MS&E, Stanford University.
#
# Adapted for Julia by Timothy E. Holy with the following changes:
#    - Allow an initial guess for x
#    - Eliminate printing
#-----------------------------------------------------------------------
function lsqr_method!(log::ConvergenceHistory, x, A, b; damp=0, atol=sqrt(eps(Adivtype(A,b))), btol=sqrt(eps(Adivtype(A,b))), conlim=real(one(Adivtype(A,b)))/sqrt(eps(Adivtype(A,b))), maxiter::Int=max(size(A,1), size(A,2)))
    # Sanity-checking
    m = size(A,1)
    n = size(A,2)
    length(x) == n || error("x should be of length ", n)
    length(b) == m || error("b should be of length ", m)
    for i = 1:n
        isfinite(x[i]) || error("Initial guess for x must be finite")
    end

    # Initialize
    T = Adivtype(A, b)
    Tr = real(T)
    itn = istop = 0
    ctol = conlim > 0 ? convert(Tr, 1/conlim) : zero(Tr)
    Anorm = Acond = ddnorm = res2 = xnorm = xxnorm = z = sn2 = zero(Tr)
    cs2 = -one(Tr)
    dampsq = abs2(damp)
    tmpm = similar(b, T, m)
    tmpn = similar(x, T, n)

    log[:atol] = atol
    log[:btol] = btol
    log[:ctol] = ctol

    # Set up the first vectors u and v for the bidiagonalization.
    # These satisfy  beta*u = b-A*x,  alpha*v = A'u.
    u = b - A*x
    v = copy(x)
    beta = norm(u)
    alpha = zero(Tr)
    if beta > 0
        log.mtvps=1
        scale!(u, inv(beta))
        Ac_mul_B!(v,A,u)
        alpha = norm(v)
    end
    if alpha > 0
        scale!(v, inv(alpha))
    end
    w = copy(v)
    wrho = similar(w)

    Arnorm = alpha*beta
    if Arnorm == 0
        return
    end

    rhobar = alpha
    phibar = bnorm = rnorm = r1norm = r2norm = beta

    #------------------------------------------------------------------
    #     Main iteration loop.
    #------------------------------------------------------------------
    while itn < maxiter
        nextiter!(log,mvps=1)
        itn += 1

        # Perform the next step of the bidiagonalization to obtain the
        # next beta, u, alpha, v.  These satisfy the relations
        #      beta*u  =  A*v  - alpha*u,
        #      alpha*v  =  A'*u - beta*v.

        # Note that the following three lines are a band aid for a GEMM: X: C := αAB + βC.
        # This is already supported in A_mul_B! for sparse and distributed matrices, but not yet dense
        A_mul_B!(tmpm, A, v)
        scale!(u, -alpha)
        LinAlg.axpy!(one(eltype(tmpm)), tmpm, u)
        beta = norm(u)
        if beta > 0
            log.mtvps+=1
            scale!(u, inv(beta))
            Anorm = sqrt(abs2(Anorm) + abs2(alpha) + abs2(beta) + dampsq)
            # Note that the following three lines are a band aid for a GEMM: X: C := αA'B + βC.
            # This is already supported in Ac_mul_B! for sparse and distributed matrices, but not yet dense
            Ac_mul_B!(tmpn, A, u)
            scale!(v, -beta)
            LinAlg.axpy!(one(eltype(tmpn)), tmpn, v)
            alpha  = norm(v)
            if alpha > 0
                scale!(v, inv(alpha))
            end
        end

        # Use a plane rotation to eliminate the damping parameter.
        # This alters the diagonal (rhobar) of the lower-bidiagonal matrix.
        rhobar1 = sqrt(abs2(rhobar) + dampsq)
        cs1     = rhobar/rhobar1
        sn1     = damp  /rhobar1
        psi     = sn1*phibar
        phibar  = cs1*phibar

        # Use a plane rotation to eliminate the subdiagonal element (beta)
        # of the lower-bidiagonal matrix, giving an upper-bidiagonal matrix.
        rho     =   sqrt(abs2(rhobar1) + abs2(beta))
        cs      =   rhobar1/rho
        sn      =   beta   /rho
        theta   =   sn*alpha
        rhobar  = - cs*alpha
        phi     =   cs*phibar
        phibar  =   sn*phibar
        tau     =   sn*phi

        # Update x and w
        t1      =   phi  /rho
        t2      = - theta/rho

        LinAlg.axpy!(t1, w, x)
        scale!(w, t2)
        LinAlg.axpy!(one(t2), v, w)
        copy!(wrho, w)
        scale!(wrho, inv(rho))
        ddnorm += norm(wrho)

        # Use a plane rotation on the right to eliminate the
        # super-diagonal element (theta) of the upper-bidiagonal matrix.
        # Then use the result to estimate  norm(x).
        delta   =   sn2*rho
        gambar  =  -cs2*rho
        rhs     =   phi - delta*z
        zbar    =   rhs/gambar
        xnorm   =   sqrt(xxnorm + abs2(zbar))
        gamma   =   sqrt(abs2(gambar) + abs2(theta))
        cs2     =   gambar/gamma
        sn2     =   theta /gamma
        z       =   rhs   /gamma
        xxnorm +=   abs2(z)

        # Test for convergence.
        # First, estimate the condition of the matrix  Abar,
        # and the norms of  rbar  and  Abar'rbar.
        Acond   =   Anorm*sqrt(ddnorm)
        res1    =   abs2(phibar)
        res2    =   res2 + abs2(psi)
        rnorm   =   sqrt(res1 + res2)
        Arnorm  =   alpha*abs(tau)

        # 07 Aug 2002:
        # Distinguish between
        #    r1norm = ||b - Ax|| and
        #    r2norm = rnorm in current code
        #           = sqrt(r1norm^2 + damp^2*||x||^2).
        #    Estimate r1norm from
        #    r1norm = sqrt(r2norm^2 - damp^2*||x||^2).
        # Although there is cancellation, it might be accurate enough.
        r1sq    =   abs2(rnorm) - dampsq*xxnorm
        r1norm  =   sqrt(abs(r1sq));   if r1sq < 0 r1norm = - r1norm; end
        r2norm  =   rnorm
        push!(log, :resnorm, r1norm)

        # Now use these norms to estimate certain other quantities,
        # some of which will be small near a solution.
        test1   =   rnorm /bnorm
        test2   =   Arnorm/(Anorm*rnorm)
        test3   =   inv(Acond)
        t1      =   test1/(1 + Anorm*xnorm/bnorm)
        rtol    =   btol + atol*Anorm*xnorm/bnorm
        push!(log, :cnorm, test3)
        push!(log, :anorm, test2)
        push!(log, :rnorm, test1)

        # The following tests guard against extremely small values of
        # atol, btol  or  ctol.  (The user may have set any or all of
        # the parameters  atol, btol, conlim  to 0.)
        # The effect is equivalent to the normal tests using
        # atol = eps,  btol = eps,  conlim = 1/eps.
        if itn >= maxiter  istop = 7; end
        if 1 + test3  <= 1 istop = 6; end
        if 1 + test2  <= 1 istop = 5; end
        if 1 + t1     <= 1 istop = 4; end

        # Allow for tolerances set by the user
        if  test3 <= ctol  istop = 3; end
        if  test2 <= atol  istop = 2; end
        if  test1 <= rtol  istop = 1; end
    end
    shrink!(log)
    setconv(log, istop > 0)
    x
end

function lsqr!(x, A, b; maxiter::Int=max(size(A,1), size(A,2)), kwargs...)
    T = Adivtype(A, b)
    z = zero(T)
    history = ConvergenceHistory()
    reserve!(history,[:resnorm,:anorm,:rnorm,:cnorm],maxiter)
    lsqr_method!(history, x, A, b; maxiter=maxiter, kwargs...)
    x, history
end

lsqr(A, b; kwargs...) = lsqr!(zerox(A, b), A, b; kwargs...)


using AxisAlgorithms
using Base.Test

include("tridiag.jl")
include("matmul.jl")
include("woodbury.jl")


#############################################################
# Making sure embedding is invariant under forward linear map
# of its vertices
#############################################################
import Distributions
import Statistics
import StaticArrays: SVector, MVector
import DelayEmbeddings: Dataset

"""
    invariantize(pts; noise_factor = 0.01, step = 5)

Checks whether the last point in `pts` lies inside the convex hull of the
previus points. If not, the last point is iteratively moved towards the
center of the convex hull until it reaches the interior of the convex hull (`step` percent
at a time).

This functions adds a small amount of noise to the data, equivalent to `noise_factor` times
the standard deviation for each dynamical variable.

Valid inputs are `AbstractArray{T, 2}` (the largest dimension will be treated as the index),
`Vector{Vector{T}}`, `Vector{SVector{D, T}}`, `Vector{MVector{D, T}}` (each vector is a
point) or a `Dataset` instance.
"""
function invariantize end

"""
    forwardlinearmap_invariant(pts)

Return true if the last point in `pts` lies inside the convex hull of the preceding points.

Valid inputs are `AbstractArray{T, 2}` (the largest dimension will be treated as the index),
`Vector{Vector{T}}`, `Vector{SVector{D, T}}`, `Vector{MVector{D, T}}` (each vector is a
point) or a `Dataset` instance.
"""
function forwardlinearmap_invariant end

function forwardlinearmap_invariant(pts::Vector{SVector{D, T}}) where {D, T}
    forwardlinearmap_invariant(Array(hcat(pts...,)))
end

function forwardlinearmap_invariant(pts::Vector{MVector{D, T}}) where {D, T}
    forwardlinearmap_invariant(Array(hcat(pts...,)))
end

function forwardlinearmap_invariant(pts::Dataset) where {D, T}
    forwardlinearmap_invariant(Array(hcat(pts...,)))
end

function forwardlinearmap_invariant(pts::Vector{Vector{T}}) where T
    forwardlinearmap_invariant(hcat(pts...,))
end

function forwardlinearmap_invariant(pts::AbstractArray{T, 2}) where T

    # Transpose if rows represent points
    if size(pts, 1) > size(pts, 2)
       pts = transpose(pts)
    end

    lastpoint = pts[:, end]
    dim = size(pts, 1)
    # Triangulate the embedding using all points but the last
    # Returns a vector of index vectors, one for each simplex
    t = delaunay(transpose(pts[:, 1:end-1]))
    n_simplices = length(t)
    points = pts[:, 1:end-1]

    # Loop over valid simplices and check whether the corresponding simplex
    # contains the last point.
    i = 1
    lastpoint_contained = false

    #=
    Continue checking if simplices contain the last point until some simplex
    contains it. Then the point must necessarily be inside the convex hull of
    the triangulation formed by those simplices.
    =#
    while i <= n_simplices && !lastpoint_contained
        sᵢ = transpose(points[:, t[i]])
        orientation_sᵢ = Statistics.det([ones(dim + 1, 1) sᵢ])

        # Insert the last point of the embedding in place
        # of the j-th vertex of sᵢ and check whether the
        # convex expansion coefficient β𝚥 stays non-negative.
        # If β𝚥 becomes negative, the point is not contained.
        j = 1
        β𝚥 = 1
        while j <= dim + 1 && β𝚥 >= 0
            tmp = copy(sᵢ)
            tmp[j, :] = lastpoint
            β𝚥 = Statistics.det(hcat(ones(dim + 1, 1), tmp)) * sign(orientation_sᵢ)
            j = j + 1
        end

        #=
        If the last convex expansion coefficient is positive, the last point is
        contained in the triangulation (because all the previous coefficients
        must have been nonnegative)
        =#
        if β𝚥 >= 0
            lastpoint_contained = true
        end

        i = i + 1
    end
    return lastpoint_contained
end

forwardlinearmap_invariant(E::AbstractEmbedding) = forwardlinearmap_invariant(E.points)

function invariantize(E::AbstractEmbedding{D, T};
                        noise_factor = 0.01, step = 3) where {D, T}
    pts = E.points
    dim = size(E.points, 1)

    if size(unique(pts, dims = 2), 2) < size(pts, 2)

       @warn """Embedding points not unique. Adding a little noise ($noise_factor times the maximum of the the standard deviations along each axis)"""
       # Find standard deviation along each axis
       dim = size(pts, 1)
       σ = Statistics.std(E.points, dims = 2)

       for i = 1:dim
          pts[i, :] .+= rand(Distributions.Uniform(-σ[i], σ[i])) .* noise_factor
       end
    end

    #=
    # Keep track of the embedding's center point and the original position of the
    # last point in the embedding, so we can move the last point along a line
    # from its original position towards the embedding's center, until the point
    # lies inside the convex hull of the preceding points.
    =#
    ce = sum(E.points, dims = 2)/size(E.points, 2) # embedding center
    lp = E.points[:, end] # last point of the embedding
    # What direction should we move?
    dir = ce - lp

    dir = dropdims(ce .- lp, dims = 2)

    # Points along line toward the center of the embedding.
    steps = 1:step:100
    ptsonline = [lp .+ dir .* (pct_moved/100) for pct_moved in 1:step:100]

    for i = 1:length(ptsonline)
       pt = ptsonline[i]
       P = hcat(pts[:, 1:(end - 1)], pt)
       if forwardlinearmap_invariant(P)

          embeddingdata = EmbeddingData{dim, T}(
               float.(E.embeddingdata.dataseries),  # the dataseries
               E.embeddingdata.labels, # empty labels by default
               E.embeddingdata.descriptions, # empty descriptions by default
               E.embeddingdata.in_which_pos, # the positions in which the dataseries appear
               E.embeddingdata.at_what_lags # the embedding lag for each column
               )

          return LinearlyInvariantEmbedding(
              hcat(pts[:, 1:(end-1)], pt), embeddingdata)
       end
    end
    @warn """Could not make embedding invariant. Returning unmodified $E."""
    return E
end

function invariantize(pts::AbstractArray{Float64, 2};
                        noise_factor = 0.01, step = 3)
   dim = size(pts, 1)

   if size(pts, 1) > size(pts, 2)
      error("`pts` must be an array of size dim-by-n_points, where each column is a point")
   end

   if size(unique(pts, dims = 2), 2) < size(pts, 2)

      @warn """Points not unique. Adding a little noise ($noise_factor times the maximum of the the standard deviations along each axis)"""
      # Find standard deviation along each axis
      dim = size(pts, 1)
      σ = Statistics.std(pts, dims = 2)

      for i = 1:dim
         pts[i, :] .+= rand(Distributions.Uniform(-σ[i], σ[i])) .* noise_factor
      end
   end

   #=
   # Keep track of the embedding's center point and the original position of the
   # last point in the embedding, so we can move the last point along a line
   # from its original position towards the embedding's center, until the point
   # lies inside the convex hull of the preceding points.
   =#
   ce = sum(pts, dims = 2)/size(pts, 2) # embedding center
   lp = pts[:, end] # last point of the embedding
   # What direction should we move?
   dir = ce - lp

   dir = dropdims(ce .- lp, dims = 2)

   # Points along line toward the center of the embedding.
   ptsonline = [lp .+ dir .* (pct_moved/100) for pct_moved in 1:step:100]
   for pt in ptsonline
      P = hcat(pts[:, 1:(end - 1)], pt)
      if forwardlinearmap_invariant(P)
         return P
      end
   end

   @warn """Could not make point set invariant under forward linear map. Returning unmodified points."""
   return pts
end

function invariantize(pts::Vector{VT}; step = 3) where {VT <: Union{SVector{D, T}, MVector{D, T}, Vector{T}} where {D, T}}
   pts_arr = Array(hcat(pts...,))
   inv_pts = invariantize(pts_arr, step = step)
   return [VT(inv_pts[:, i]) for i = 1:length(pts)]
end

function invariantize(pts::Vector{Vector{T}}; step = 3) where T
   pts_arr = hcat(pts...,)
   inv_pts = invariantize(pts_arr, step = step)
   return [inv_pts[:, i] for i = 1:length(pts)]
end

function invariantize(pts::Dataset; step = 3)
   inv_pts = invariantize(hcat([Vector(x) for x in pts]...,), step = step)
   Dataset([inv_pts[:, i] for i = 1:length(pts)])
end

export
forwardlinearmap_invariant,
invariantize


@testset "d_separated" begin
    # Examples from <https://doi.org/10.1073/PNAS.1510507113>.
    nodes = [:Z1, :W1, :X, :Z3, :W3, :Y, :Z2, :W2]
    mapping = Dict(zip(nodes, 1:length(nodes)))
    m = mapping
    E = [
        :Z1 => :W1,
        :W1 => :X,
        :Z1 => :Z3,
        :Z3 => :X,
        :Z2 => :Z3,
        :Z2 => :W2,
        :Z3 => :Y,
        :W2 => :Y,
        :W3 => :Y,
        :X => :W3
    ]
    edges_num = [m[pair.first] => m[pair.second] for pair in E]
    G = SimpleDiGraph(Edge.(edges_num))
    X = m[:Z1]
    Y = m[:Y]
    Z = Set([m[:X], m[:Z3], m[:W2]])
    @test M.d_separated(G, X, Y, Z)
    Z = Set([m[:X], m[:Z3], m[:W3]])
    @test !(M.d_separated(G, X, Y, Z))
end

@testset "without_in_out" begin
    # From https://youtu.be/pZkCecwE-xE.
    # See also https://youtu.be/9X4pR4jvKmM?t=1025 for the same example with more
    # explanation.
    s = 1  # smoking
    t = 2  # tar
    c = 3  # cancer
    g = 4  # genotype
    E = [
        s => t,
        t => c,
        g => s,
        g => c
    ]
    E = Edge.(E)
    G = SimpleDiGraphFromIterator(E)
    @test Set(M.nodes(G)) == Set([1, 2, 3, 4])

    G_without_in = M.without_incoming(G, Set([s, t]))
    @test G_without_in == SimpleDiGraph(Edge.([t => c, g => c]))

    G_without_out = M.without_outgoing(G, Set([s, t]))
    @test G_without_out == SimpleDiGraph(Edge.([g => s, g => c]))
end


function sum_vector(x::Array{Float64, 1})
s = 0.0
for i = 1:length(x)
s = s + x[i]
end
return s
end
function sum_cols_matrix(x::Array{Float64, 2})
num_cols = size(x, 2)
s = zeros(num_cols)
for i = 1:num_cols
s[i] = sum_vector(x[:, i])
end
return s
end


struct InnerProductVec{F,T}
    vec::T
    dotf::F
end
Base.eltype(v::InnerProductVec) = eltype(typeof(v))
Base.eltype(::Type{InnerProductVec{F,T}}) where {F,T} = eltype(T)

Base.:-(v::InnerProductVec) = InnerProductVec(-v.vec, v.dotf)
Base.:+(v::InnerProductVec{F}, w::InnerProductVec{F}) where F = InnerProductVec(v.vec+w.vec, v.dotf)
Base.:-(v::InnerProductVec{F}, w::InnerProductVec{F}) where F = InnerProductVec(v.vec-w.vec, v.dotf)
Base.:*(v::InnerProductVec, a) = InnerProductVec(v.vec*a, v.dotf)
Base.:*(a, v::InnerProductVec) = InnerProductVec(a*v.vec, v.dotf)
Base.:/(v::InnerProductVec, a) = InnerProductVec(v.vec/a, v.dotf)
Base.:\(a, v::InnerProductVec) = InnerProductVec(a\v.vec, v.dotf)

Base.similar(v::InnerProductVec, ::Type{T} = eltype(v)) where {T} = InnerProductVec(similar(v.vec), v.dotf)

Base.getindex(v::InnerProductVec) = v.vec

function Base.fill!(v::InnerProductVec, a)
    fill!(v.vec, a)
    return v
end

function Base.copyto!(w::InnerProductVec{F}, v::InnerProductVec{F}) where F
    copyto!(w.vec, v.vec)
    return w
end

function LinearAlgebra.mul!(w::InnerProductVec{F}, a, v::InnerProductVec{F}) where F
    mul!(w.vec, a, v.vec)
    return w
end

function LinearAlgebra.mul!(w::InnerProductVec{F}, v::InnerProductVec{F}, a) where F
    mul!(w.vec, v.vec, a)
    return w
end

function LinearAlgebra.rmul!(v::InnerProductVec, a)
    rmul!(v.vec, a)
    return v
end

function LinearAlgebra.axpy!(a, v::InnerProductVec{F}, w::InnerProductVec{F}) where F
    axpy!(a, v.vec, w.vec)
    return w
end
function LinearAlgebra.axpby!(a, v::InnerProductVec{F}, b, w::InnerProductVec{F}) where F
    axpby!(a, v.vec, b, w.vec)
    return w
end

LinearAlgebra.dot(v::InnerProductVec{F}, w::InnerProductVec{F}) where {F} = v.dotf(v.vec, w.vec)
LinearAlgebra.norm(v::InnerProductVec) = sqrt(real(dot(v,v)))


using BenchmarkTools, PyPlot
using Base.Math.@horner, ParPoly

bench(f::Function, X) = mapreduce(x -> reinterpret(Unsigned,x), |, f(x) for x in X)
X = randn(1_000_000)

nn = 5:25
time_h = Float64[]
time_hs = Float64[]
time_hss = Float64[]
for n in nn
    println("running bench for order $n")
    c = randn(Float64,n+1)

    @gensym f1
    @gensym f2
    @gensym f3
    @eval begin
        $f1(x) = @horner x $(c...)
        $f2(x) = @horner_split x $(c...)
        $f3(x) = @horner_split_simd x $(c...)
    end
    @eval begin
        $f1(2.); $f2($2.); $f3($2.) # precompile them
    end
    t_h    = @benchmark bench($f1, $X)
    t_hs   = @benchmark bench($f2, $X)
    t_hss  = @benchmark bench($f3, $X)

    push!(time_h,   mean(t_h).time)
    push!(time_hs,  mean(t_hs).time)
    push!(time_hss, mean(t_hss).time)
end

plot(nn, time_h*0.000001,   label="@horner")
plot(nn, time_hs*0.000001,  label="@horner_split")
plot(nn, time_hss*0.000001, label="@horner_split_simd")

legend(loc="upper left")
xlabel("polynomial order")
ylabel("time (ms)")
savefig("out.jpg")


#=

Code to implement piecewise linear interpolation for an arbitrary number
of dimensions

=#

function piecewise_linear_nodes(n::S,domain = [1.0,-1.0]) where {S <: Integer}

  if n <= 0
    error("The number of nodes must be positive.")
  end

  nodes = zeros(n)

  if isodd(n)
    nodes[Int((n-1)/2)+1] = (domain[1]+domain[2])/2.0
  end

  if n == 1

    return nodes

  elseif n == 2

    return [domain[2], domain[1]]

  else

    nodes[1] = domain[2]
    nodes[n] = domain[1]

    inc = (domain[1]-domain[2])/(n-1)
    for i = 2:div(n,2)
      nodes[i]     = domain[2] + (i-1)*inc
      nodes[n-i+1] = domain[1] - (i-1)*inc
    end

    return nodes

  end

end

const linear_nodes = piecewise_linear_nodes

function bracket_nodes(x::Array{T,1},point::R) where {T <: AbstractFloat,R<:Number}

  if real(point) <= x[1] # Real is used because complex numbers are occasionally used in NLboxsolve.jl
    return (1,2)
  elseif real(point) >= x[end]
    return (length(x) - 1, length(x))
  else
    y = 0
    for i in x
      if i < real(point)
        y += 1
      else
        break
      end
    end
    return (y, y+1)
  end

end

function bracket_nodes(x::Union{NTuple{N,Array{T,1}},Array{Array{T,1},1}},point::Array{R,1}) where {T <: AbstractFloat, N,R<:Number}

  bracketing_nodes = Array{Int64,2}(undef,2,length(point))
  for i = 1:length(point)
    bracketing_nodes[:,i] .= bracket_nodes(x[i],point[i])
  end

  return bracketing_nodes

end

function piecewise_linear_weight(x::Array{T,1},point::R) where {T<:AbstractFloat,R<:Number}

    bracketing_nodes = bracket_nodes(x,point)
    weight = (point-x[bracketing_nodes[1]])/(x[bracketing_nodes[2]]-x[bracketing_nodes[1]])

    return weight

end

function piecewise_linear_weights(x::Union{NTuple{N,Array{T,1}},Array{Array{T,1},1}},point::Array{R,1}) where {T<:AbstractFloat,N,R<:Number}

    weights = Array{R,1}(undef,length(point))
    for i = 1:length(point)
        weights[i] = piecewise_linear_weight(x[i],point[i])
    end

    return weights

end

function select_bracketing_nodes(bounds::Array{S,2} where {S <: Integer})

  d = size(bounds,2)

  number_relevant_grid_points = 2^d
  bracketing_grid_points      = Array{Int64}(undef,number_relevant_grid_points,d)

  for i = 1:d
    bracketing_grid_points[:,i] .= repeat(repeat(bounds[:,i],inner=[2^(d-i),1]),2^(i-1))[:]
  end

  return bracketing_grid_points

end

function select_relevant_data(y::AbstractArray{T,N},bracketing_grid_points::Array{S,2}) where {T <: AbstractFloat, S <: Integer, N}

  data = zeros(size(bracketing_grid_points,1))
  for i = 1:length(data)
    data[i] = y[CartesianIndex(tuple(bracketing_grid_points[i,:]...))]
  end

  return data

end

function piecewise_linear_evaluate(y::AbstractArray{T,N},x::Union{NTuple{N,Array{T,1}},Array{Array{T,1},1}},point::Union{R,Array{R,1}}) where {T<:AbstractFloat,N,R<:Number}

  b = bracket_nodes(x,point)
  w = piecewise_linear_weights(x,point)

  d = size(b,2)

  relevant_points = select_bracketing_nodes(b)
  data = select_relevant_data(y,relevant_points)

  for j = d:-1:1

    new_data = zeros(R,div(length(data),2))
    for i = 1:length(new_data)
      new_data[i] = data[2*(i-1)+1] + w[j]*(data[2*i]-data[2*(i-1)+1])
    end

    data = copy(new_data)

  end

  return data[1]

end

function piecewise_linear_evaluate(y::AbstractArray{T,N},x::Union{NTuple{N,Array{T,1}},Array{Array{T,1},1}}) where {T <: AbstractFloat, N}

  function approximating_function(point::Union{R,Array{R,1}}) where {R<:Number}

    return piecewise_linear_evaluate(y,x,point)

  end

  return approximating_function

end

function piecewise_linear_evaluate(y::AbstractArray{T,N},x::Union{NTuple{N,Array{T,1}},Array{Array{T,1},1}},point::Union{R,Array{R,1}},integrals::Union{T,Array{T,1}}) where {T<:AbstractFloat,N,R<:Number}

  # This function is only needed to facilitate compatibility with SolveDSGE

  b = bracket_nodes(x,point)
  w = piecewise_linear_weights(x,point)
  w = w.*integrals

  d = size(b,2)

  relevant_points = select_bracketing_nodes(b)
  data = select_relevant_data(y,relevant_points)

  for j = d:-1:1

    new_data = zeros(R,Int(length(data)/2))
    for i = 1:length(new_data)
      new_data[i] = data[2*(i-1)+1] + w[j]*(data[2*i]-data[2*(i-1)+1])
    end

    data = copy(new_data)

  end

  return data[1]

end

# Functions to handle the 1-D case

function piecewise_linear_evaluate(y::Array{T,1},x::Array{T,1},point::R) where {T<:AbstractFloat,R<:Number}

    b = bracket_nodes(x,point)
    w = piecewise_linear_weight(x,point)

    y_estimate = y[b[1]] + w*(y[b[2]]-y[b[1]])

    return y_estimate

end

function piecewise_linear_evaluate(y::Array{T,1},nodes::Array{T,1}) where {T <: AbstractFloat}

  function approximating_function(point::R) where {R<:Number}

      return piecewise_linear_evaluate(y,nodes,point)

  end

  return approximating_function

end

# Function to transform from one uniform grid to another

function grid_reshape(f::AbstractArray{T,N},grid::NTuple{N,Array{T,1}}) where {T<:AbstractFloat,N}

  #1. Construct the old grid

  old_grid = Array{Array{T,1},1}(undef,N)
  for i = 1:N
    piecewise_linear_nodes
    old_grid[i] = piecewise_linear_nodes(size(f,i),[grid[i][end],grid[i][1]])
  end

  #2. Initialize the new y

  y = Array{T,N}(undef,tuple(length.(grid)...))

  #3. Interpolate to fill y

  point = Array{T,1}(undef,N)
  for i in CartesianIndices(y)
    for j = 1:N
      point[j] = grid[j][i[j]]
    end
    y[i] = piecewise_linear_evaluate(f,old_grid,point)
  end

  return y

end


cd(@__DIR__)
using Pkg; Pkg.activate("."); Pkg.instantiate()

using DiffEqFlux, Flux
using DifferentialEquations
using Plots, DelimitedFiles, Statistics
using Sundials

function FENEP!(out,du,u,p,t,γd)
  # DAE definition for FENE-p
  θ₁₁,θ₂₂,θ₁₂, τ₁₁,τ₂₂,τ₁₂ = u
  λ,η,L = p
  a = L^2 /(L^2 -3)
  fλ = (L^2 + (λ/η/a)*(τ₂₂+τ₁₁))/(L^2 - 3)
  out[1] =  τ₁₁ + du[1] - 2*λ*γd(t)*τ₁₂/fλ
  out[2] =  τ₂₂ + du[2]
  out[3] =  τ₁₂ + du[3] - λ*γd(t)*τ₂₂/fλ - η/fλ * γd(t)

  out[4] = θ₁₁ - λ*τ₁₁/fλ
  out[5] = θ₂₂ - λ*τ₂₂/fλ
  out[6] = θ₁₂ - λ*τ₁₂/fλ
end

function find_σ_exact(tsave,γd)
  # finds the exact solution of the FENE-p equations given a strain rate fn and
  # time points to save at
  λ = 2.0
  L = 2.0
  η = 4.0
  p = [λ,η,L]
  u₀ = zeros(6)
  du₀ = [0.0, 0.0, η*γd(0.0)*(L^2-3)/L^2, 0.0,0.0,0.0]
  tspan = (Float64(tsave[1]),Float64(tsave[end]))
  differential_vars = [true,true,true,false,false,false]
  prob = DAEProblem((out,du,u,p,t) -> FENEP!(out,du,u,p,t,γd),du₀,u₀,tspan,
                    p=p,differential_vars=differential_vars)
  sol = solve(prob,IDA(),saveat=Float64.(tsave))
  return [Float32(σ[6]) for σ in sol.u]
end

function mode_loss(f0,f1,γd)
    # Given the NN f0 and f1 and strain rate γd, this returns a function
    # which when evaluated gives the loss.
    tspan = (0.0f0,6.2831f0)
    tsave = range(tspan[1],tspan[2],length=100)
    σ_exact = find_σ_exact(tsave,γd)
    p = []
    σ0 =0.0f0
    u0 = Tracker.param([σ0])

    dudt_(u::TrackedArray,p,t) = f1(Flux.Tracker.collect(vcat(u, [γd(t)])))
    dudt_(u::AbstractArray,p,t) = Flux.data(f1(vcat(u,[γd(t)])))

    prob = ODEProblem(dudt_,u0,tspan,p)
    loss_rd() =   begin
        P_RD = vcat(Flux.Tracker.collect(diffeq_rd(p,prob,Tsit5(),u0=u0,saveat=tsave))
                      ,γd.(tsave)')
        σ_out = [f0(P_RD[:,i])[1] for i = 1:size(P_RD,2)]
        return sum( (σ_out .- σ_exact).^2 )
    end
    return loss_rd
end

function test_NN(γd,f0,f1)
    # Test the trained NN against a new strain rate
    tspan = (0.0f0,10.0f0)
    tsave = range(tspan[1],tspan[2],length=100)
    σ_exact = find_σ_exact(tsave,γd)
    dudt_(u::AbstractArray,p,t) = Flux.data(f1(vcat(u,[γd(t)])))
    prob = ODEProblem(dudt_,[0.0f0],tspan,p)
    ode_solve = solve(prob,u0=[0.0f0],Tsit5())
    f0_ = Flux.data(f0)
    σ_approx = [f0_([ode_solve(t)[1],γd(t)]).data[1] for t in tsave]

    return tsave, σ_approx, σ_exact
end
function test_err(γd,f0,f1)
    tsave, σ_approx, σ_exact = test_NN(γd,f0,f1)
    return sum( (σ_approx .- σ_exact).^2 )
end

## --- Define NN and linear model ---
f0_n = Chain(Dense(2,4,tanh), Dense(4,1))
f1_n = Chain(Dense(2,4,tanh), Dense(4,1))
f0_l = Chain(Dense(2,1))
f1_l = Chain(Dense(2,1))
# --- Define the total loss across various strain rates ----
t_loss(f0,f1) =  sum(mode_loss(f0,f1,t -> 12.0f0*cos.(ω.*t))() for ω in 1.0f0:0.2f0:2.0f0)
data = Iterators.repeated((), 5000)
opt = ADAM(0.015)
p = []
# --- Define test strain rate here, so call back can see how all errors improve
#     with time ----
γd_test = t -> 12f0*cos.(1.5f0.*t) #+ 7f0*cos.(2.5f0.*t)
err_l = []
err_n = []
cb(err,f0,f1) = push!(err,[t_loss(f0,f1).data,test_err(γd_test,f0,f1)])
Flux.train!(() -> t_loss(f0_n,f1_n), params(f0_n,f1_n,p), data, opt,cb= () -> cb(err_n,f0_n,f1_n))
Flux.train!(() -> t_loss(f0_l,f1_l), params(f0_l,f1_l,p), data, opt,cb= () -> cb(err_l,f0_l,f1_l))

# --- Make plot of how error goes with time ------
Er = (i,er) -> [e[i] for e in er]
p1 = plot(Er(1,err_n),xscale=:log10,yscale=:log10,ylabel="Error",xlabel="Training steps",
        label="Training error, Neural net",legend=:bottomleft)
plot!(p1,Er(2,err_n),label="Testing error, Neural net")
plot!(p1,Er(1,err_l), label="Training error, linear model")
plot!(p1,Er(2,err_l),label="Testing error, linear model")

# --- Plot approximations vs true sol ----
tsave, σ_approx_n, σ_exact = test_NN(γd_test,f0_n,f1_n)
tsave, σ_approx_l, σ_exact = test_NN(γd_test,f0_l,f1_l)
p2 = plot(tsave,σ_approx_n,m=:circle,label="NN solution",ylabel="stress",xlabel="time")
plot!(p2,tsave,σ_approx_l,m=:hexagon,label="Linear model")
plot!(p2,tsave,σ_exact,label="True solution",leg=:bottomright,size=(700,400))

# program takes time to run, so save data for tweaking figures
open("er_data.txt", "w") do io # or other path
           writedlm(io, [Er(1,err_n) Er(2,err_n) Er(1,err_l) Er(2,err_l)], ',')
end;

open("plt_data.txt", "w") do io # or other path
           writedlm(io, [[t for t in tsave]  σ_approx_n σ_approx_l σ_exact], ',')
end;

# Timing
using StaticArrays, BenchmarkTools
W1 = convert(SMatrix{4,2},Tracker.data(f1_n[1].W))
b1 = convert(SVector{4},Tracker.data(f1_n[1].b))
W2 = convert(SMatrix{1,4},Tracker.data(f1_n[2].W))
b2 = convert(SVector{1},Tracker.data(f1_n[2].b))

tspan = (0.0f0,10.0f0)
tsave = range(tspan[1],tspan[2],length=100)
σ_exact = find_σ_exact(tsave,γd_test)
dudt_(u,p,t) = first(W2*broadcast(tanh,W1*@SVector([u[1],γd_test(t)])+b1)+b2)
prob = ODEProblem(dudt_,0.0f0,tspan,p)
@btime ode_solve = solve(prob,Tsit5(),saveat=tsave)
@btime find_σ_exact(tsave,γd_test)


function adversarial_errors(network, X, Y; coeffs = [-0.74, -0.44],
                            δ = 0.02, solver = GLPK.Optimizer, use_gurobi = false,
                            PGD = false, step_factor = 0.1, iterations = 600)
    """ Computes absolute errors with no adversarial perturbations
        Args:
            - network: neural network (NNV model)
            - X: training examples (each column is an example)
            - Y: labels (each column is a label)
            ---------------
            - coeffs: coefficients to get output (linear combination of network outputs)
            - δ: radius for hyperrectangle to check error within
            - solver: solver for mixed integer programs (ignored if PGD true)
            - use_gurobi: whether we are using gurobi (still need to put in solver)
            - PGD: whether or not to use PGD (if not, solve exactly with MIP Verify)
            - step_factor: step factor for PGD (ignored if PGD false)
            - iterations: number of steps for PGD (ignored if PGD false)
    """

    N_samples = size(Y, 2)

    truevals = vec(coeffs' * Y)
    predvals = [coeffs' * compute_output(network, X[:, i]) for i = 1:N_samples]

    minvals = zeros(N_samples)
    maxvals = zeros(N_samples)

    if PGD
        @showprogress "Getting adversial results..." for i = 1:N_samples
            lbs = X[:, i] .- Float32.(δ * ones(size(X, 1)))
            ubs = X[:, i] .+ Float32.(δ * ones(size(X, 1)))
            maxvals[i] = pgd(network, X[:, i], lbs, ubs, coeffs = coeffs,
                        step_size = step_factor, iterations = iterations)
            minvals[i] = -pgd(network, X[:, i], lbs, ubs, coeffs = -coeffs,
                        step_size = step_factor, iterations = iterations)
        end
    else
        @showprogress "Getting adversial results..." for i = 1:N_samples
            input_set = Hyperrectangle(X[:, i], Float32.(δ * ones(size(X, 1))))
            maxvals[i] = optimize_ouput(network, input_set, coeffs = coeffs,
            solver = solver, use_gurobi = use_gurobi)
            minvals[i] = optimize_ouput(network, input_set, coeffs = coeffs,
            solver = solver, use_gurobi = use_gurobi, maximize = false)
        end
    end

    return max.(abs.(truevals .- maxvals), abs.(truevals .- minvals)),
           max.(abs.(predvals .- maxvals), abs.(predvals .- minvals))
end

function optimize_ouput(network, input_set;
    coeffs = [-0.74, -0.44], solver = GLPK.solver, use_gurobi = false, maximize=true,
    obj_threshold=(maximize ? -Inf : Inf))
    """ Function from Chris to get max output
    """
    # Get your bounds
    bounds = NeuralVerification.get_bounds(Ai2z(), network, input_set; before_act=true)

    # Create your model
    if use_gurobi
        model = Model(with_optimizer(solver, OutputFlag = 0))
    else
        model = Model(with_optimizer(solver))
    end
    z = init_vars(model, network, :z, with_input=true)
    δ = init_vars(model, network, :δ, binary=true)
    # get the pre-activation bounds:
    model[:bounds] = bounds
    model[:before_act] = true

    # Add the input constraint
    NeuralVerification.add_set_constraint!(model, input_set, first(z))

    # Encode the network as an MIP
    encode_network!(model, network, BoundedMixedIntegerLP())
    obj = coeffs'*last(z)
    if maximize
        @objective(model, Max, obj)
        if obj_threshold != -Inf
            @constraint(model, obj >= obj_threshold)
        end
    else
        @objective(model, Min, obj)
        if obj_threshold != Inf
            @constraint(model, obj >= obj_threshold)
        end
    end

    # Set lower and upper bounds
    #for the first layer it's special because it has no ẑ
    set_lower_bound.(z[1], low(bounds[1]))
    set_upper_bound.(z[1], high(bounds[1]))
    for i = 2:length(z)-1
        # Set lower and upper bounds for the intermediate layers
        ẑ_i =  _ẑᵢ₊₁(model, i-1)
        z_i = z[i]
        # @constraint(model, ẑ_i .>= low(bounds[i])) These empirically seem to slow it down?
        # @constraint(model, ẑ_i .<= high(bounds[i]))
        z_low = max.(low(bounds[i]), 0.0)
        z_high = max.(high(bounds[i]), 0.0)
        set_lower_bound.(z_i, z_low)
        set_upper_bound.(model[:z][i], z_high)
    end

    # Set lower and upper bounds for the last layer special because
    # it has no ReLU
    set_lower_bound.(z[end], low(bounds[end]))
    set_upper_bound.(z[end], high(bounds[end]))

    optimize!(model)
    if termination_status(model) == OPTIMAL
        return objective_value(model)
    elseif termination_status(model) == INFEASIBLE
        @warn "Infeasible result, did you have an output threshold? If not, then it should never return infeasible"
        return maximize ? -Inf : Inf
    else
        @assert false "Non optimal result"
    end
end

function fgsm(network, x, lbs, ubs; coeffs = [-0.74, -0.44], step_size = 0.1)
    grad_full = NeuralVerification.get_gradient(network, x)
    grad = grad_full' * coeffs
    return clamp.(x + grad * step_size, lbs, ubs)
end

function pgd(network, x, lbs, ubs; coeffs = [-0.74, -0.44], step_size = 0.01,
    iterations = 600)
    cur_x = x
    for i = 1:iterations
        cur_x = fgsm(network, cur_x, lbs, ubs; coeffs = coeffs, step_size = step_size)
    end
    return coeffs' * compute_output(network, cur_x)
end


"""
    xopt, fopt, converged = fminunc(obj, startval)

Minimize the function obj, starting at startval. This is a convenience function, provided for
students. I recommend learning to use Optim.jl rather than using this.

fminunc() with no arguments will run an example, execute edit(fminunc,()) to see the code.
"""
function fminunc(obj, x; xtol = 1e-08, gtol=1e-6, ftol=1e-12)
    results = try
        Optim.optimize(obj, x, LBFGS(),
                            Optim.Options(
                            g_tol = gtol,
                            x_tol = xtol,
                            f_tol=ftol); autodiff=:forward)
    catch
        Optim.optimize(obj, x, LBFGS(),
                            Optim.Options(
                            g_tol = gtol,
                            x_tol = xtol,
                            f_tol = ftol))
    end
    return results.minimizer, results.minimum, Optim.converged(results)
end

function fminunc()
    println("with no arguments, fminunc() runs a simple example")
    println("type edit(fminunc, ()) to see the example code")
    # return of objective should be real valued, thus the [1] to pull value out of 1-dim array
    obj = x -> (x.^2)[1]
    # obj = x -> x'x # this would also work, as returns a Float64
    x = [2.0] # argument to objective function should be a vector, thus the brackets
    results = fminunc(obj, x)
end


using Base.Test
using QuantumOptics

Nmin = 2
Nmax = 4
N = Nmax - Nmin + 1
basis = FockBasis(Nmin, Nmax)
bra = Bra(basis)
ket = Ket(basis)

@test_throws DimensionMismatch Bra(basis, [1, 2])
@test_approx_eq 0. norm(bra-Bra(basis, zeros(Int, N)))
@test_approx_eq 0. norm(ket-Ket(basis, zeros(Int, N)))
@test_approx_eq 0. bra*ket
@test_throws bases.IncompatibleBases bra*Ket(FockBasis(Nmin, Nmax+1))

bra = Bra(basis, [1im, 0, 1])
ket = Ket(basis, [0, -1im, 1])

@test_approx_eq 0. norm(5*bra - Bra(basis, [5im, 0, 5]))
@test_approx_eq 0. norm(5*ket - Ket(basis, [0, -5im, 5]))
@test_approx_eq 0. norm(5*ket - ket/0.2)
@test_throws bases.IncompatibleBases bra + Bra(FockBasis(Nmin, Nmax+1))
@test_throws bases.IncompatibleBases ket + Ket(FockBasis(Nmin, Nmax+1))

# Norm
basis = FockBasis(0, 1)
bra = Bra(basis, [3im, -4])
ket = Ket(basis, [-4im, 3])
@test_approx_eq 5. norm(bra)
@test_approx_eq 5. norm(ket)

ubra = normalize(bra)
uket = normalize(ket)
@test_approx_eq 5. norm(bra)
@test_approx_eq 5. norm(ket)
@test_approx_eq 1. norm(ubra)
@test_approx_eq 1. norm(uket)

bra_ = normalize!(bra)
ket_ = normalize!(ket)
@test bra_ === bra
@test ket_ === ket
@test_approx_eq 1. norm(bra)
@test_approx_eq 1. norm(ket)

# Partial Trace
basis = FockBasis(0, 1)
a = normalize(Bra(basis, [1im, 0]))
b = normalize(Bra(basis, [1, 2]))
c = normalize(Bra(CompositeBasis(basis, basis), [1im, 2im, 0, 0]))
@test_approx_eq 0. norm(tensor(a, b) - c)
@test_approx_eq_eps 0. tracedistance(operators.ptrace(c, 1), tensor(dagger(b), b)) 1e-15
@test_approx_eq_eps 0. tracedistance(operators.ptrace(c, 2), tensor(dagger(a), a)) 1e-15


# ------------------------------------------------------------------------------------------
# # Learning to recognize handwritten digits using a neural network
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# We have now reached the point where we can tackle a very interesting task: applying the
# knowledge we have gained with machine learning in general, and `Flux.jl` in particular, to
# create a neural network that can recognize handwritten digits! The data are from a data
# set called MNIST, which has become a classic in the machine learning world.
#
# [We could also try to apply the techniques to the original images of fruit instead.
# However, the fruit images are much larger than the MNIST images, which makes the learning
# a suitable neural network too slow.]
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ## Data munging
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# As we know, the first difficulty with any new data set is locating it, understanding what
# format it is stored in, reading it in and decoding it into a useful data structure in
# Julia.
#
# The original MNIST data is available [here](http://yann.lecun.com/exdb/mnist); see also
# the [Wikipedia page](https://en.wikipedia.org/wiki/MNIST_database). However, the format
# that the data is stored in is rather obscure.
#
# Fortunately, various packages in Julia provide nicer interfaces to access it. We will use
# the one provided by `Flux.jl`.
#
# The data are images of handwritten digits, and the corresponding labels that were
# determined by hand (i.e. by humans). Our job will be to get the computer to **learn** to
# recognize digits by learning, as usual, the function that relates the input and output
# data.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ### Loading and examining the data
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# First we load the required packages:
# ------------------------------------------------------------------------------------------

using Flux, Flux.Data.MNIST

# ------------------------------------------------------------------------------------------
# Now we read in the data:
# ------------------------------------------------------------------------------------------

labels = MNIST.labels();
images = MNIST.images();  # the semi-colon (`;`) here is important: it prevents Julia from displaying the object

# ------------------------------------------------------------------------------------------
# #### Exercise 1
#
# Examine the `labels` data. Then examine the first few images. *Do not try to view the
# whole of the `images` object!* Try to drill down to discover how the data is laid out.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# #### Exercise 2
#
# Convert the first image to a matrix of `Float64`.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ### Munging the data
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# In the previous notebooks, we arranged the input data for Flux as a `Vector` of `Vector`s.
# Now we will use an alternative arrangement, as a matrix, since that allows `Flux` to use
# matrix operations, which are more efficient.
#
# The column $i$ of the matrix is a vector consisting of the $i$th data point
# $\mathbf{x}^{(i)}$.  Similarly, the desired outputs are given as a matrix, with the $i$th
# column being the desired output $\mathbf{y}^{(i)}$.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# #### Exercise 3
#
# An image is a matrix of colours, but now we need a vector instead. To do so, we just
# arrange all of the elements of the matrix in a certain way into a single list;
# fortunately, Julia already provides the function `vec` to do so!
#
# 1. Which order does `vec` use? [This reflects the underlying way in which the matrix is
# stored in memory.]
#
# 2. How can you convert an image into a `Vector` of `Float64`?
#
# 3. Define a variable $n$ that is the length of these vectors.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# #### Exercise 4
# Make a function `rewrite` that accepts a range and converts that range of images to
# floating-point vectors and stacks them horizontally using `hcat` and the "splat" operator
# `...`.
#
# We also want a matrix of one-hot vectors. `Flux` provides a function `onehotbatch` to do
# this (you will need to import it). It works like `onehot`, but takes in a vector of labels
# and outputs a matrix `Y`.
#
# Return the pair `(X, Y)`.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ## Setting up the neural network
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# Now we must set up a neural network. Since the data is complicated, we may expect to need
# several layers.
# But we can start with a single layer.
#
# - The network will take as inputs the vectors $\mathbf{x}^{(i)}$, so the input layer has
# $n$ nodes.
#
# - The output will be a one-hot vector encoding the digit from 1 to 9 or 0 that is desired.
# There are 10 possible categories, so we need an output layer of size 10.
#
# It is then our task as neural network designers to insert layers between these input and
# output layers, whose weights will be tuned during the learning process. *This is an art,
# not a science*! But major advances have come from finding a good structure for the
# network.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ### Softmax
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# We will make a network with a single layer; let's choose each neuron in the layer to use
# the `relu` activation function.
# The output `relu` can be arbitrarily large, but in the end we will wish to compare the
# network's output with one-hot vectors, i.e. values between $0$ and $1$.
#
# In order to make this work, we will thus use an extra function at the end that takes in a
# vector of arbitrary real numbers and maps it ("squashes it down") to a vector of numbers
# between $0$ and $1$.
#
# The most used function with this property is $\mathrm{softmax}$. Firstly we take the
# exponential of each input variable to make them positive. Then we divide by the sum to
# make sure they lie between $0$ and $1$.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# $$\mathrm{softmax}(\mathbf{x})_i := \frac{\exp (x_i)}{\sum_j \exp(x_j)}$$
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# Note that here we have written the result for the $i$th component of the function
# $\mathbf{R}^n \to \mathbf{R}^n$. Note also that the function returns a vector of numbers
# that are positive, and whose components sum to $1$. Thus, in fact, they can be thought of
# as probabilities.
#
# In the neural network context, using a `softmax` after the final layer thus allows us to
# interpret the outputs as probabilities, in our case the probability that the network
# assigns that a given image represents each possible output value ($0$-$9$)!
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# #### Exercise 5
#
# Make a neural network with one single layer, using the function $\sigma$, and a final
# `softmax`.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ## Training
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# As we know, **training** consists of iteratively adjusting the model's parameters to
# decrease the `loss` function. Which parameters need to be adjusted? All of them!
#
# Since the `loss` function contains a call to the `model` function, calling `back!` on the
# result of the loss function updates the information about the gradient of the loss
# function with respect to *every node in the network!*:
# ------------------------------------------------------------------------------------------

l = loss(X, Y)

Flux.Tracker.back!(l)

# ------------------------------------------------------------------------------------------
# This is what is going on inside the `train!` function.
# In fact, `train!(loss, data, opt)` iterates over each object in `data` and runs this
# function.
# For this reason, `data` must consist of an iterable object that returns pairs `(X, Y)` at
# each step.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# The simplest possibility is
# ------------------------------------------------------------------------------------------

data = ((X, Y), )  # one-element tuple

# ------------------------------------------------------------------------------------------
# Alternatively, we can make one call to the `train!` function iterate over several copies
# of `data`, using `repeated`. This is an **iterator**; it does not copy the data 100 times,
# which would be very wasteful; it just gives an object that repeatedly loops over the same
# data:
# ------------------------------------------------------------------------------------------

dataset = Base.Iterators.repeated((X, Y), 100)

# ------------------------------------------------------------------------------------------
# #### Exercise 6
#
# Train the model on a subset of $N$ images with $N = 5000$.
# ------------------------------------------------------------------------------------------

N = 5_000
X, Y = rewrite(1:N)

# ------------------------------------------------------------------------------------------
# The function `loss` evaluated on the matrices gives the overall loss:
# ------------------------------------------------------------------------------------------

loss(X, Y)

@time Flux.train!(loss, data, opt)

@time Flux.train!(loss, dataset, opt)

# ------------------------------------------------------------------------------------------
# This is (approximately) equivalent to just doing a `for` loop to run the previous `train!`
# command 100 times.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ### Using callbacks
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# The `train!` function can take an optional keyword argument, `cb` (short for
# "*c*all*b*ack"). A callback function is a function that you provide as an argument to a
# function `f`, which "calls back" your function every so often.
#
# This provides the possibility to provide a function that is called at each step or every
# so often during the training process.
# A common use case is to provide a visual trace of the training process by printing out the
# current value of the `loss` function:
# ------------------------------------------------------------------------------------------

callback() = @show(loss(X, Y))

Flux.train!(loss, data, opt; cb = callback)

Flux.train!(loss, dataset, opt; cb = callback)

# ------------------------------------------------------------------------------------------
# However, it is expensive to calculate the complete `loss` function and it is not necessary
# to output it every step. So `Flux` also provides a function `throttle`, that provides a
# mechanism to call a given function at most once every certain number of seconds:
# ------------------------------------------------------------------------------------------

Flux.train!(loss, dataset, opt; cb = Flux.throttle(callback, 1))

for i in 1:100
    Flux.train!(loss, dataset, opt; cb = Flux.throttle(callback, 1))
end

# ------------------------------------------------------------------------------------------
# ## Testing phase
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# We now have trained a model, i.e. we have found the parameters `W` and `b` for the network
# layer(s). In order to **test** if the learning procedure was really successful, we check
# how well the resulting trained network performs when we test it with images that the
# network has not yet seen!
#
# Often, a dataset is split up into "training data" and "test (or validation) data" for this
# purpose, and indeed the MNIST data set has a separate pool of training data. We can
# instead use the images that we have not included in our reduced training process.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# #### Exercise 7
#
# Take the next 100 images after those that were used for training. How well does it do?
# ------------------------------------------------------------------------------------------

X_test, Y_test = rewrite(N+1:N+100)

loss(X_test, Y_test)

display(images[N+1])
labels[N+1]

[model(X_test[:,1]) Y_test[:,1]]

loss(X_test[:,1], Y_test[:,1])

loss(X_test, Y_test)

# ------------------------------------------------------------------------------------------
# #### Exercise 8
#
# Use the `indmax` function to write a function `prediction` that reports which digit
# `model` predicts, as the index with the maximum weight.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# #### Exercise 9
#
# Count the number of correct predictions over the whole data set, and hence the percentage
# of images that are correctly predicted. [This percentage is what is used to compare
# different machine learning techniques.]
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ## Improving the prediction
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# So far we have used a single layer. In order to improve the prediction, we probably need
# to use more layers.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# #### Exercise 10
#
# Introduce an intermediate, hidden layer. Does it give a better prediction?
# ------------------------------------------------------------------------------------------

.nbexports/introductory-tutorials/broader-topics-and-ecosystem/intro-to-ml/19. Recognizing handwriting using a neural network.jl


# Utility operators for computing prox iteratively, e.g. using CG
#
# The idea is that whatever type has the A_mul_B!, Ac_mul_B!, size and eltype
# methods implemented is a linear operator.

import Base: *, size, eltype
import LinearAlgebra: mul!

abstract type LinOp end

function (*)(Op::LinOp, x)
  # Is this the right thing to do?
  # Or maybe just: y = zeros(eltype(x), size(Op, 1))
  y = zeros(promote_type(eltype(Op), eltype(x)), size(Op, 1))
  mul!(y, Op, x)
end

size(Op::LinOp, i::Integer) = i <= 2 ? size(Op)[i] : 1

# AAc (Gram matrix)

struct AAc{M, T} <: LinOp
  A::M
  buf::AbstractArray{T}
  function AAc{M, T}(A::M) where {M, T}
    new(A, zeros(T, size(A, 2)))
  end
end

AAc(A::M) where {M} = AAc{M, eltype(A)}(A)

function mul!(y, Op::AAc, x)
  mul!(Op.buf, adjoint(Op.A), x)
  mul!(y, Op.A, Op.buf)
end

mul!(y, Op::Adjoint{AAc}, x) = mul!(y, adjoint(Op), x)

size(Op::AAc) = size(Op.A, 1), size(Op.A, 1)
eltype(Op::AAc) = eltype(Op.A)

# AcA (Covariance matrix)

struct AcA{M, T} <: LinOp
  A::M
  buf::AbstractArray{T}
  function AcA{M, T}(A::M) where {M, T}
    new(A, zeros(T, size(A, 1)))
  end
end

AcA(A::M) where {M} = AcA{M, eltype(A)}(A)

function mul!(y, Op::AcA, x)
  mul!(Op.buf, Op.A, x)
  mul!(y, adjoint(Op.A), Op.buf)
end

mul!(y, Op::Adjoint{AcA}, x) = mul!(y, adjoint(Op), x)

size(Op::AcA) = size(Op.A, 2), size(Op.A, 2)
eltype(Op::AcA) = eltype(Op.A)

# Shifted symmetric linear operator

struct Shift{M, T} <: LinOp
  A::M
  rho::T
  function Shift{M, T}(A::M, rho::T) where {M, T}
    if eltype(A) != T
      error("type of rho is incompatible with A")
    end
    new(A, rho)
  end
end

Shift(A::M, rho::T) where {M, T} = Shift{M, T}(A, rho)

function mul!(y, Op::Shift, x)
  mul!(y, Op.A, x)
  y .+= Op.rho .* x
end

function mul!(y, Op::Adjoint{Shift}, x)
  mul!(y, adjoint(Op.A), x)
  y .+= Op.rho .* x
end

size(Op::Shift) = size(Op.A, 2), size(Op.A, 2)
eltype(Op::Shift) = eltype(Op.A)


# Parts of this code were taken / derived from NetworkX. See LICENSE for
# licensing details.

"""Creates a complete graph with `n` vertices. A complete graph has edges
connecting each pair of vertices.
"""
function CompleteGraph(n::Integer)
    g = Graph(n)
    for i = 1:n, j=1:n
        if i < j
            add_edge!(g, Edge(i,j))
        end
    end
    return g
end

"""Creates a complete bipartite graph with `n1+n2` vertices. It has edges
connecting each pair of vertices in the two sets.
"""
function CompleteBipartiteGraph(n1::Integer, n2::Integer)
    g = Graph(n1+n2)
    for i = 1:n1, j=n1+1:n1+n2
        add_edge!(g, Edge(i,j))
    end
    return g
end

"""Creates a complete digraph with `n` vertices. A complete digraph has edges
connecting each pair of vertices (both an ingoing and outgoing edge).
"""
function CompleteDiGraph(n::Integer)
    g = DiGraph(n)
    for i = 1:n, j=1:n
        if i != j
            add_edge!(g, Edge(i,j))
        end
    end
    return g
end

"""Creates a star graph with `n` vertices. A star graph has a central vertex
with edges to each other vertex.
"""
function StarGraph(n::Integer)
    g = Graph(n)
    for i = 2:n
        add_edge!(g, Edge(1,i))
    end
    return g
end

"""Creates a star digraph with `n` vertices. A star digraph has a central
vertex with directed edges to every other vertex.
"""
function StarDiGraph(n::Integer)
    g = DiGraph(n)
    for i = 2:n
        add_edge!(g, Edge(1,i))
    end
    return g
end

"""Creates a path graph with `n` vertices. A path graph connects each
successive vertex by a single edge."""
function PathGraph(n::Integer)
    g = Graph(n)
    for i = 2:n
        add_edge!(g, Edge(i-1, i))
    end
    return g
end

"""Creates a path digraph with `n` vertices. A path graph connects each
successive vertex by a single directed edge."""
function PathDiGraph(n::Integer)
    g = DiGraph(n)
    for i = 2:n
        add_edge!(g, Edge(i-1, i))
    end
    return g
end

"""Creates a cycle graph with `n` vertices. A cycle graph is a closed path graph.
"""
function CycleGraph(n::Integer)
    g = Graph(n)
    for i = 1:n-1
        add_edge!(g, Edge(i, i+1))
    end
    add_edge!(g, Edge(n, 1))
    return g
end

"""Creates a cycle digraph with `n` vertices. A cycle digraph is a closed path digraph.
"""
function CycleDiGraph(n::Integer)
    g = DiGraph(n)
    for i = 1:n-1
        add_edge!(g, Edge(i, i+1))
    end
    add_edge!(g, Edge(n, 1))
    return g
end

"""Creates a wheel graph with `n` vertices. A wheel graph is a star graph with
the outer vertices connected via a closed path graph.
"""
function WheelGraph(n::Integer)
    g = StarGraph(n)
    for i = 3:n
        add_edge!(g, Edge(i-1, i))
    end
    if n != 2
        add_edge!(g, Edge(n, 2))
    end
    return g
end

"""Creates a wheel digraph with `n` vertices. A wheel graph is a star digraph
with the outer vertices connected via a closed path graph.
"""
function WheelDiGraph(n::Integer)
    g = StarDiGraph(n)
    for i = 3:n
        add_edge!(g, Edge(i-1, i))
    end
    if n != 2
        add_edge!(g, Edge(n, 2))
    end
    return g
end

"""
    Grid{T<:Integer}(dims::AbstractVector{T}; periodic=false)

Creates a `d`-dimensional cubic lattice, with `d=length(dims)` and length  `dims[i]` in dimension `i`.
If `periodic=true` the resulting lattice will have periodic boundary condition in each dimension.
"""
function Grid{T<:Integer}(dims::AbstractVector{T}; periodic=false)
    func = periodic ? CycleGraph : PathGraph
    g = func(dims[1])
    for d in dims[2:end]
        g = cartesian_product(func(d), g)
    end
    return g
end

"""create a binary tree with k-levels vertices are numbered 1:2^levels-1"""
function BinaryTree(levels::Int)
    g = Graph(2^levels-1)
    for i in 0:levels-2
        for j in 2^i:2^(i+1)-1
            add_edge!(g, j, 2j)
            add_edge!(g, j, 2j+1)
        end
    end
    return g
end

"""create a double complete binary tree with k-levels
used as an example for spectral clustering by Guattery and Miller 1998."""
function DoubleBinaryTree(levels::Int)
    gl = BinaryTree(levels)
    gr = BinaryTree(levels)
    g = blkdiag(gl, gr)
    add_edge!(g,1, nv(gl)+1)
    return g
end

"""The Roach Graph from Guattery and Miller 1998"""
function RoachGraph(k::Int)
    dipole = CompleteGraph(2)
    nopole = Graph(2)
    antannae = crosspath(k, nopole)
    body = crosspath(k,dipole)
    roach = blkdiag(antannae, body)
    add_edge!(roach, nv(antannae)-1, nv(antannae)+1)
    add_edge!(roach, nv(antannae), nv(antannae)+2)
    return roach
end

"""This function generates `n` connected k-cliques """
function CliqueGraph(k::Integer, n::Integer)
    g = Graph(k*n)
    for c=1:n
        for i=(c-1)*k+1:c*k-1, j=i+1:c*k
            add_edge!(g, i, j)
        end
    end
    for i=1:n-1
        add_edge!(g, (i-1)*k+1, i*k+1)
    end
    add_edge!(g, 1, (n-1)*k+1)
    return g
end


using ParameterizedFunctions, Test
using OrdinaryDiffEq, Calculus, ForwardDiff

f = @ode_def begin
    dx = -a
    dy = b
end a b

cb = ContinuousCallback((u,t,i) -> u[1], (integrator)->(println("Stopped.");integrator.p[2]=zero(integrator.p[2])))
function test_f(p)
  prob = ODEProblem(f,eltype(p).([1.0,0.0]),eltype(p).((0.0,1.0)),copy(p))
  integrator = init(prob,Tsit5(),abstol=1e-14,reltol=1e-14,callback=cb)
  step!(integrator)
  solve!(integrator).u[end]
end
p = [2.0, 1.0]
findiff = Calculus.finite_difference_jacobian(test_f,p)
fordiff = ForwardDiff.jacobian(test_f,p)
@test findiff ≈ fordiff

f2 = @ode_def begin
    dx = -x
    dy = b
end a b

function test_f2(p)
  prob = ODEProblem(f2,eltype(p).([1.0,0.0]),eltype(p).((0.0,1.0)),copy(p))
  integrator = init(prob,Tsit5(),abstol=1e-14,reltol=1e-14,callback=cb)
  step!(integrator)
  solve!(integrator).u[end]
end
p = [2.0, 1.0]
findiff = Calculus.finite_difference_jacobian(test_f2,p)
fordiff = ForwardDiff.jacobian(test_f2,p)
@test findiff ≈ fordiff


struct Solver{T}
    p::Problem{T}
    nlp_bounds::Vector{MOI.NLPBoundsPair}
    block_data::MOI.NLPBlockData
    solver::Ipopt.Optimizer
    z::Vector{MOI.VariableIndex}
end

function Solver(trajopt::TrajectoryOptimizationProblem; eval_hess=false, options=nothing)
    p = Problem(trajopt, eval_hess=eval_hess)

    nlp_bounds = MOI.NLPBoundsPair.(p.con_bnds...)
    block_data = MOI.NLPBlockData(nlp_bounds, p, true)

    solver = Ipopt.Optimizer()

    # TODO: options
    solver.options["max_iter"] = 1000
    solver.options["tol"] = 1.0e-3
    solver.options["constr_viol_tol"] = 1.0e-3

    z = MOI.add_variables(solver, p.num_var)

    for i = 1:p.num_var
        MOI.add_constraint(solver, z[i], MOI.LessThan(p.var_bnds[2][i]))
        MOI.add_constraint(solver, z[i], MOI.GreaterThan(p.var_bnds[1][i]))
    end

    MOI.set(solver, MOI.NLPBlock(), block_data)
    MOI.set(solver, MOI.ObjectiveSense(), MOI.MIN_SENSE)

    return Solver(p, nlp_bounds, block_data, solver, z)
end

function initialize!(s::Solver, z)
    for i = 1:s.p.num_var
        MOI.set(s.solver, MOI.VariablePrimalStart(), s.z[i], z[i])
    end
end

function solve!(s::Solver)
    MOI.optimize!(s.solver)
end


module TestWoodburyFactorizations
using WoodburyFactorizations
using LazyInverses
using LinearAlgebra
using Test

function getW(n, m; diagonal = true, symmetric = false)
    A = diagonal ? Diagonal(exp.(.1randn(n))) + I(n) : randn(n, n)
    A = symmetric && !diagonal ? A'A : A
    U = randn(n, m)
    C = diagonal ? Diagonal(exp.(.1randn(m))) + I(m) : randn(m, m)
    C = symmetric && !diagonal ? C'C : C
    V = symmetric ? U' : randn(m, n)
    W = Woodbury(A, U, C, V)
end

Base.AbstractMatrix(F::BunchKaufman) = Symmetric(Matrix(F))

@testset "Woodbury" begin
    n = 5
    m = 2
    W = getW(n, m, symmetric = true)
    MW = Matrix(W)
    @test size(W) == (n, n)
    x = randn(size(W, 2))

    # multiplication
    @test W*x ≈ MW*x
    @test x'W ≈ x'MW
    @test dot(x, W, x) ≈ dot(x, MW*x)
    X = randn(size(W, 2), 3)
    @test W*X ≈ MW*X
    @test X'*W ≈ X'*MW
    @test W == copy(W) # testing isequal and copy
    @test W == deepcopy(W) # testing isequal and deepcopy

    # in-place multiplication
    y = randn(size(W, 1)) # with vector
    α, β = randn(2)
    b = α*(W*x) + β*y
    mul!(y, W, x, α, β)
    @test y ≈ b

    Y = randn(size(W, 1), size(X, 2)) # with matrix
    α, β = randn(2)
    B = α*(W*X) + β*Y
    mul!(Y, W, X, α, β)
    @test Y ≈ B

    @test α*W isa Woodbury
    @test Matrix(α * W) ≈ α * Matrix(W) # with scalar
    @test Matrix(W * α) ≈ Matrix(W) * α # with scalar
    @test Matrix(W / α) ≈ Matrix(W) / α # with scalar
    @test Matrix(α \ W) ≈ α \ Matrix(W) # with scalar
    @test logdet(abs(α)*W) ≈ logdet(abs(α)*Matrix(W)) # checking that logdet works correctly

    @test eltype(W) == Float64
    @test issymmetric(W)
    @test ishermitian(W)
    # @test ishermitian(MW)
    # @test ishermitian(Woodbury(W.A, W.U, W.C, copy(W.U)')) # testing edge case
    @test !issymmetric(getW(n, m))
    @test !ishermitian(getW(n, m))

    # test solves
    @test W\(W*x) ≈ x
    @test (x'*W)/W ≈ x'

    # factorization
    n = 1024
    m = 3
    W = getW(n, m, symmetric = true)
    MW = Matrix(W)
    F = factorize(W)
    @test F isa Inverse
    x = randn(n)
    @test W \ x ≈ F \ x
    @test MW \ x ≈ F \ x

    n = 4
    m = 3
    W = getW(n, m, symmetric = true)
    MW = Matrix(W)

    ## determinant
    @test det(W) ≈ det(MW)
    @test logdet(W) ≈ logdet(MW)
    @test logabsdet(W)[1] ≈ logabsdet(MW)[1]
    @test logabsdet(W)[2] ≈ logabsdet(MW)[2]
    @test det(inverse(W)) ≈ det(Inverse(W))
    @test det(inverse(W)) ≈ det(inv(MW))
    @test logdet(inverse(W)) ≈ logdet(Inverse(W))
    @test logdet(inverse(W)) ≈ logdet(inv(MW))
    @test all(logabsdet(inverse(W)) .≈ logabsdet(Inverse(W)))
    @test all(logabsdet(inverse(W)) .≈ logabsdet(inv(MW)))

    # indexing
    for i in 1:n, j in 1:n
        @test W[i, j] ≈ MW[i, j]
    end
    i = 1:2
    j = 3:3
    Wij = W[i, j]
    @test Wij isa Woodbury
    @test Wij[1, 1] == W[1, 3]
    Wij = @view W[i, j]
    @test Wij isa Woodbury
    @test Wij.U isa SubArray

    # trace
    @test tr(W) ≈ tr(MW)

    # factorize
    F = factorize(W)
    @test Matrix(inverse(F)) ≈ inv(MW)
    @test inv(W) ≈ inv(MW)
    @test logdet(F) ≈ logdet(MW)
    @test isposdef(F)

    # factorize and logdet after factorization
    F = factorize(W)
    @test logabsdet(F)[1] ≈ logabsdet(MW)[1]
    @test logabsdet(F)[2] ≈ logabsdet(MW)[2]
    @test logdet(F) ≈ logdet(MW)
    @test Matrix(inverse(F)) ≈ inv(MW)

    # test recursive stacking of woodbury objects
    b = randn(n, 1)
    W2 = Woodbury(W, b, ones(1, 1), b')
    F2 = factorize(W2)
    M2 = Matrix(W2)
    @test logdet(W2) ≈ logdet(M2)
    @test logdet(F2) ≈ logdet(M2)

    # test factorize_D behavior on non-hermitian 1x1 matrix
    using WoodburyFactorizations: factorize_D
    FD = factorize_D(W2, fill(-rand(), (1, 1)))
    @test FD isa Matrix
    @test size(FD) == (1, 1)

    # test rank 1 constructor
    x = randn(n)
    A = Diagonal(2 .+ exp.(randn(n)))
    u = randn(n)
    W = Woodbury(A, u)
    @test Matrix(W) ≈ A + u*u'
    b = W*x
    @test W \ b ≈ x

    v = randn(n)
    W = Woodbury(A, u, v')
    @test Matrix(W) ≈ A + u*v'
    b = W*x
    @test W \ b ≈ x

    # LowRank constructor
    U = randn(n, 1)
    W = Woodbury(I(n), U, U')

    # CholeskyPivoted constructor
    A = randn(1, n)
    A = A'A
    C = cholesky(A, RowMaximum(), check = false)
    W = Woodbury(A, C)
    @test W isa Woodbury
    @test Matrix(W) ≈ 2A
    # tests with α = -
    W = Woodbury(I(n), C, -1)
    # display(W)
    @test Matrix(W) ≈ I(n) - A
    @test inv(W) ≈ inv(I(n) - A)

    FW = factorize(W)
    @test FW isa Inverse
    @test Matrix(FW) ≈ I(n) - A
    @test Matrix(inverse(FW)) ≈ inv(I(n) - A)
    MW = I(n) - A
    # indexing
    for i in 1:n, j in 1:n
        @test W[i, j] ≈ MW[i, j]
    end
    @test tr(W) ≈ tr(MW)
    @test diag(W) ≈ diag(MW)

    # Woodbury structure with 1d outer dimension
    a, b, c, d = randn(4)
    W = Woodbury(a, b, c, d)
    MW = Matrix(W)
    @test size(MW) == (1, 1)
    @test MW[1, 1] ≈ a + *(b, c, d)

    a = randn()
    B = randn(1, n)
    C = randn(n, n)
    D = randn(n, 1)
    W = Woodbury(a, B, C, D)
    MW = Matrix(W)
    @test size(MW) == (1, 1)
    @test MW[1, 1] ≈ a + dot(B, C, D)

    @test inv(W) isa Real
    @test factorize(W) isa Real
end

end # TestWoodburyFactorizations


module SquarenessTest

using Squareness
using Test

@testset "Bool" begin
    @test issquare(false)
    @test issquare(true)
end

@testset "Int8" begin
    @testset "is a square" begin
        dataset_sq_i8 = Int8[i^2 for i=0:11]
        @test all(issquare, dataset_sq_i8)
    end

    @testset "is not a square" begin
        dataset_nosq_i8 = setdiff(typemin(Int8):typemax(Int8), Int8[i^2 for i=0:11])
        @test all(!issquare, dataset_nosq_i8)
    end
end

@testset "UInt8" begin
    @testset "is a square" begin
        dataset_sq_ui8 = UInt8[i^2 for i=0:15]
        @test all(issquare, dataset_sq_ui8)
    end

    @testset "is not a square" begin
        dataset_nosq_ui8 = setdiff(typemin(UInt8):typemax(UInt8), UInt8[i^2 for i=0:15])
        @test all(!issquare, dataset_nosq_ui8)
    end
end

@testset "Int16" begin
    @testset "is a square" begin
        dataset_sq_i16 = Int16[i^2 for i=0:181]
        @test all(issquare, dataset_sq_i16)
    end

    @testset "is not a square" begin
        dataset_nosq_i16 = setdiff(typemin(Int16):typemax(Int16), Int16[i^2 for i=0:181])
        @test all(!issquare, dataset_nosq_i16)
    end
end

@testset "UInt16" begin
    @testset "is a square" begin
        dataset_sq_ui16 = UInt16[i^2 for i=0:255]
        @test all(issquare, dataset_sq_ui16)
    end

    @testset "is not a square" begin
        dataset_nosq_ui16 = setdiff(typemin(UInt16):typemax(UInt16), UInt16[i^2 for i=0:255])
        @test all(!issquare, dataset_nosq_ui16)
    end
end

@testset "Int32" begin
    @testset "is a square" begin
        dataset_sq_i32 = Int32[i^2 for i=0:46340]
        @test all(issquare, dataset_sq_i32)
    end

    @testset "is not a square" begin
        # dataset_nosq_i32 = setdiff(typemin(Int32):typemax(Int32), Int32[i^2 for i=0:46340])
        dataset_nosq_i32 = setdiff(rand(Int32, 1000000), Int32[i^2 for i=0:46340])
        @test all(!issquare, dataset_nosq_i32)
    end
end

@testset "UInt32" begin
    @testset "is a square" begin
        dataset_sq_ui32 = UInt32[i^2 for i=0:65535]
        @test all(issquare, dataset_sq_ui32)
    end

    @testset "is not a square" begin
        # dataset_nosq_ui32 = setdiff(typemin(UInt32):typemax(UInt32), UInt32[i^2 for i=0:65535])
        dataset_nosq_ui32 = setdiff(rand(UInt32, 1000000), UInt32[i^2 for i=0:65535])
        @test all(!issquare, dataset_nosq_ui32)
    end
end

@testset "Int64" begin
    @testset "is a square" begin
        # dataset_sq_i64 = Int64[i^2 for i=0:3037000499]
        dataset_sq_i64 = Int64[i^2 for i in rand(0:3037000499, 1000000)]
        @test all(issquare, dataset_sq_i64)
    end

    @testset "random (may not be a square)" begin
        dataset_random_i64 = rand(Int64, 1000000)
        @test all(@.(issquare(dataset_random_i64) == Squareness.issquare_with_isqrt(dataset_random_i64)))
    end
end

@testset "UInt64" begin
    @testset "is a square" begin
        # dataset_sq_ui64 = UInt64[i^2 for i=0:4294967295]
        dataset_sq_ui64 = UInt64[i^2 for i in rand(UInt64(0):UInt64(4294967295), 1000000)]
        @test all(issquare, dataset_sq_ui64)
    end

    @testset "random (may not be a square)" begin
        dataset_random_ui64 = rand(UInt64, 1000000)
        @test all(@.(issquare(dataset_random_ui64) == Squareness.issquare_with_isqrt(dataset_random_ui64)))
    end
end

@testset "Int128" begin
    @testset "is a square" begin
        # dataset_sq_i128 = Int128[i^2 for i=0:13043817825332782212]
        dataset_sq_i128 = Int128[i^2 for i in rand(int128"0":int128"13043817825332782212", 1000000)]
        @test all(issquare, dataset_sq_i128)
    end

    @testset "random (may not be a square)" begin
        dataset_random_i128 = rand(Int128, 1000000)
        @test all(@.(issquare(dataset_random_i128) == Squareness.issquare_with_isqrt(dataset_random_i128)))
    end
end

@testset "UInt128" begin
    @testset "is a square" begin
        # dataset_sq_ui128 = UInt128[i^2 for i=0:18446744073709551615]
        dataset_sq_ui128 = UInt128[i^2 for i in rand(uint128"0":uint128"18446744073709551615", 1000000)]
        @test all(issquare, dataset_sq_ui128)
    end

    @testset "random (may not be a square)" begin
        dataset_random_ui128 = rand(UInt128, 1000000)
        @test all(@.(issquare(dataset_random_ui128) == Squareness.issquare_with_isqrt(dataset_random_ui128)))
    end
end

@testset "BigInt" begin
    @testset "is a square" begin
        dataset_sq_mpz = BigInt[i^2 for i in (big.(rand(UInt128, 10000)) .<< 128 .+ rand(UInt128, 10000))]
        @test all(issquare, dataset_sq_mpz)
    end

    @testset "random (may not be a square)" begin
        dataset_random_mpz = big.(rand(Int128, 10000)) .<< 128 .+ rand(Int128, 10000)
        @test all(@.(issquare(dataset_random_mpz) == Squareness.issquare_with_isqrt(dataset_random_mpz)))
    end
end

end


# n = mult(3, 4) #> ErrorException("mult not defined"))
function mult(x, y)
    println("x is $x and y is $y")
    x * y
end

n = mult(3, 4) #> 12

function mult(x, y)
    println("x is $x and y is $y")
    if x == 1
        return y
    end
    x * y
end

function multi(n, m)
    n * m, div(n, m), n % m
end
x, y, z = multi(8, 2)  #> (16,4,0)

# variable number of arguments:
function varargs(n, m, args...)
    println("arguments: $n $m $args")
end
varargs(1, 2, 3, 4) #> arguments: 1 2 (3, 4)

function varargs2(args...)
    println("arguments2: $args")
end

x = (3, 4)
varargs2(1, 2, x...) # args is (1,2,3,4)
x = [10, 11, 12]
varargs2(1, 2, x...) # args is (1,2,10,11,12)

function insert_elem(arr)
    push!(arr, -10)
end

arr = [ 2, 3, 4]
insert_elem(arr)
println("arr is now $arr") #> arr is now [ 2, 3, 4, -10 ]

# typing the arguments:
function mult(x::Float64, y::Float64)
    println("x is $x and y is $y")
    x * y
end
# following gives error:
# mult(5, 6)
# ERROR: `mult` has no method matching mult(::Int64, ::Int64)

# one-line assignment syntax:
mult(x, y) = x * y
f(x, y) = x^3 - y + x*y
println(f(3,2) ) # 31

∑(x,y) = x + y
∑(3, 4) #> 7


using Flux
println("forward_tests")
using DiffEqFlux
println("Starting Soon!")
using ModelingToolkit
using DiffEqBase
using Test, NeuralPDE
println("Starting Soon!")
using SciMLBase
using DomainSets
import ModelingToolkit: Interval

@testset "ODE" begin
    @parameters x
    @variables u(..)

    Dx = Differential(x)
    eq = Dx(u(x)) ~ 0.
    bcs = [u(0.) ~ u(0.)]
    domains = [x ∈ Interval(0.0,1.0)]
    chain = FastChain((x,p) -> x.^2)

    chain([1],Float64[])
    strategy_ = NeuralPDE.GridTraining(0.1)
    discretization = NeuralPDE.PhysicsInformedNN(chain,strategy_;init_params = Float64[])
    @named pde_system = PDESystem(eq,bcs,domains,[x],[u(x)])
    prob = NeuralPDE.discretize(pde_system,discretization)
    discretized_functions = NeuralPDE.discretize_inner_functions(pde_system,discretization)

    eqs = pde_system.eqs
    bcs = pde_system.bcs
    domains = pde_system.domain
    dx = strategy_.dx
    eltypeθ = eltype(discretized_functions.flat_initθ)
    depvars,indvars,dict_indvars,dict_depvars, dict_depvar_input = NeuralPDE.get_vars(pde_system.ivs,
                                                        pde_system.dvs)

    train_sets = generate_training_sets(domains,dx,eqs,bcs,eltypeθ,
                                        dict_indvars,dict_depvars)

    pde_train_sets, bcs_train_sets = train_sets
    pde_train_sets = NeuralPDE.adapt(eltypeθ,pde_train_sets)[1]

    train_data =pde_train_sets
    pde_loss_function = discretized_functions.inner_pde_loss_functions[1]

    dudx(x) = @. 2*x
    @test pde_loss_function(train_data, Float64[]) ≈ dudx(train_data) rtol = 1e-8
end

@testset "derivatives" begin
    chain = FastChain(FastDense(2,16,Flux.σ),FastDense(16,16,Flux.σ),FastDense(16,1))
    initθ = Float64.(DiffEqFlux.initial_params(chain))

    eltypeθ = eltype(initθ)
    parameterless_type_θ = DiffEqBase.parameterless_type(initθ)
    phi = NeuralPDE.get_phi(chain,parameterless_type_θ)
    derivative = NeuralPDE.get_numeric_derivative()

    u_ = (cord, θ, phi)->sum(phi(cord, θ))

    phi([1,2], initθ)

    phi_ = (p) -> phi(p, initθ)[1]
    dphi = Zygote.gradient(phi_,[1.,2.])

    function get_ε(dim, der_num,eltypeθ)
        epsilon = cbrt(eps(eltypeθ))
        ε = zeros(eltypeθ, dim)
        ε[der_num] = epsilon
        ε
    end

    eps_x = get_ε(2, 1,Float64)
    eps_y = get_ε(2, 2,Float64)

    dphi_x = derivative(phi,u_,[1.,2.],[eps_x],1,initθ)
    dphi_y = derivative(phi,u_,[1.,2.],[eps_y],1,initθ)

    #first order derivatives
    @test isapprox(dphi[1][1], dphi_x, atol=1e-8)
    @test isapprox(dphi[1][2], dphi_y, atol=1e-8)

    dphi_x = derivative(phi,u_,[1.,2.],[[ 0.0049215667, 0.0]],1,initθ)
    dphi_y = derivative(phi,u_,[1.,2.],[[0.0,  0.0049215667]],1,initθ)

    hess_phi = Zygote.hessian(phi_,[1,2])

    dphi_xx = derivative(phi,u_,[1.,2.],[eps_x,eps_x],2,initθ)
    dphi_xy = derivative(phi,u_,[1.,2.],[eps_x,eps_y],2,initθ)
    dphi_yy = derivative(phi,u_,[1.,2.],[eps_y,eps_y],2,initθ)

    #second order derivatives
    @test isapprox(hess_phi[1], dphi_xx, atol=1e-5)
    @test isapprox(hess_phi[2], dphi_xy, atol=1e-5)
    @test isapprox(hess_phi[4], dphi_yy, atol=1e-5)
end

@testset "Integral" begin
    #semi-infinite intervals
    @parameters x
    @variables u(..)
    I = Integral(x in ClosedInterval(0, Inf))
    eq = I(u(x)) ~ 0
    bcs = [u(1.) ~ exp(1)/(exp(2) + 3)]
    domains = [x ∈ Interval(1.0, 2.0)]
    chain = FastChain((x,p) -> exp.(x) ./ (exp.(2 .*x) .+ 3))
    chain([1],Float64[])
    strategy_ = NeuralPDE.GridTraining(0.1)
    discretization = NeuralPDE.PhysicsInformedNN(chain,strategy_;init_params = Float64[])
    @named pde_system = PDESystem(eq,bcs,domains,[x],[u(x)])
    sym_prob = SciMLBase.symbolic_discretize(pde_system, discretization)
    prob = NeuralPDE.discretize(pde_system,discretization)
    discretized_functions = NeuralPDE.discretize_inner_functions(pde_system,discretization)
    inner_loss =discretized_functions.inner_pde_loss_functions[1]
    exact_u = π/(3*sqrt(3))
    @test  inner_loss(ones(1,1), Float64[])[1] ≈ exact_u  rtol = 1e-5

    #infinite intervals
    @parameters x
    @variables u(..)
    I = Integral(x in ClosedInterval(-Inf, Inf))
    eqs = I(u(x)) ~ 0
    domains = [x ∈ Interval(1.0, 2.0)]
    bcs = [u(1) ~ u(1)]
    chain = FastChain((x,p) -> x .* exp.(-x .^2))
    chain([1],Float64[])

    discretization = NeuralPDE.PhysicsInformedNN(chain,strategy_;init_params = Float64[])
    @named pde_system = PDESystem(eqs, bcs, domains, [x], [u(x)])
    sym_prob = SciMLBase.symbolic_discretize(pde_system, discretization)
    prob = SciMLBase.discretize(pde_system, discretization)
    discretized_functions = NeuralPDE.discretize_inner_functions(pde_system,discretization)
    inner_loss =discretized_functions.inner_pde_loss_functions[1]
    exact_u = 0
    @test  inner_loss(ones(1,1), Float64[])[1] ≈ exact_u  rtol = 1e-9
end


struct NNPDEHan{C1,C2,O} <: NeuralNetDiffEqAlgorithm
    u0::C1
    σᵀ∇u::C2
    opt::O
end
NNPDEHan(u0,σᵀ∇u;opt=Flux.ADAM(0.1)) = NNPDEHan(u0,σᵀ∇u,opt)

function DiffEqBase.solve(
    prob::TerminalPDEProblem,
    alg::NNPDEHan;
    abstol = 1f-6,
    verbose = false,
    maxiters = 300,
    save_everystep = false,
    dt,
    trajectories)

    X0 = prob.X0
    ts = prob.tspan[1]:dt:prob.tspan[2]
    d  = length(X0)
    g,f,μ,σ,p = prob.g,prob.f,prob.μ,prob.σ,prob.p

    data = Iterators.repeated((), maxiters)

    #hidden layer
    opt = alg.opt
    u0 = alg.u0
    σᵀ∇u = alg.σᵀ∇u
    ps = Flux.params(u0, σᵀ∇u...)

    function sol()
        map(1:trajectories) do j
            u = u0(X0)[1]
            X = X0
            for i in 1:length(ts)-1
                t = ts[i]
                _σᵀ∇u = σᵀ∇u[i](X)
                dW = sqrt(dt)*randn(d)
                u = u - f(X, u, _σᵀ∇u, p, t)*dt + _σᵀ∇u'*dW
                X  = X .+ μ(X,p,t)*dt .+ σ(X,p,t)*dW
            end
            X,u
        end
    end

    function loss()
        mean(sum(abs2,g(X) - u) for (X,u) in sol())
    end

    iters = eltype(X0)[]

    cb = function ()
        save_everystep && push!(iters, u0(X0)[1].data)
        l = loss()
        verbose && println("Current loss is: $l")
        l < abstol && Flux.stop()
    end

    Flux.train!(loss, ps, data, opt; cb = cb)

    save_everystep ? iters : u0(X0)[1].data
end #pde_solve