# ====================================================================
# Component structure
#
mutable struct Gaussian_1D <: AbstractComponent
    norm::Parameter
    center::Parameter
    sigma::Parameter

    function Gaussian_1D(norm::Number, center::Number, sigma::Number)
        @assert norm  > 0
        @assert sigma > 0

        out = new(Parameter(norm), Parameter(center), Parameter(sigma))
        out.norm.low = 0
        out.sigma.low = 0
        return out
    end
end

mutable struct Gaussian_2D <: AbstractComponent
    norm::Parameter
    centerX::Parameter
    centerY::Parameter
    sigmaX::Parameter
    sigmaY::Parameter
    angle::Parameter

    function Gaussian_2D(norm::Number, centerX::Number, centerY::Number, sigmaX::Number, sigmaY::Number, angle::Number)
        @assert norm   > 0
        @assert sigmaX > 0
        @assert sigmaY > 0

        out = new(Parameter(norm), Parameter(centerX), Parameter(centerY), Parameter(sigmaX), Parameter(sigmaY), Parameter(angle))
        out.norm.low = 0
        out.sigmaX.low = 0
        out.sigmaY.low = 0
        return out
    end
end

Gaussian(norm, center, sigma) = Gaussian_1D(norm, center, sigma)
Gaussian(norm, centerX, centerY, sigmaX, sigmaY, angle) = Gaussian_2D(norm, centerX, centerY, sigmaX, sigmaY, angle)
function Gaussian(norm, centerX, centerY, sigma)
    out = Gaussian_2D(norm, centerX, centerY, sigma, sigma, 0.)
    out.sigmaX.free = false
    out.sigmaX.expr = "this_sigmaY"
    out.angle.free = false
    return out
end

# ====================================================================
# Prepare component `cdata`
mutable struct Gaussian_cdata
    ix::Vector{Int}
    iy::Vector{Int}
    Gaussian_cdata() = new(Vector{Int}(), Vector{Int}())
end

ceval_data(domain::Domain_1D, comp::Gaussian_1D) = (Gaussian_cdata(), length(domain))
ceval_data(domain::Domain_2D, comp::Gaussian_2D) = (Gaussian_cdata(), length(domain))

# ====================================================================
# Evaluate component
function evaluate(c::CompEval{Domain_1D, Gaussian_1D},
                  norm, center, sigma)
    # TODO: optimize using cdata
    x = c.domain[1]
    @. (c.eval = exp( ((x - center) / sigma)^2. / (-2.)) /
        2.5066282746310002 / sigma * norm) # sqrt(2pi) = 2.5066282746310002
end

function evaluate(c::CompEval{Domain_2D, Gaussian_2D},
                   norm, centerX, centerY, sigmaX, sigmaY, angle)
    angle *= -pi / 180.
    a =  (cos(angle) / sigmaX)^2 / 2  +  (sin(angle) / sigmaY)^2 / 2
    b = -sin(2angle) / sigmaX^2  / 2  +  sin(2angle) / sigmaY^2  / 2
    c =  (sin(angle) / sigmaX)^2 / 2  +  (cos(angle) / sigmaY)^2 / 2

    # TODO: optimize using cdata
    x = c.domain[1]
    y = c.domain[2]

    @. (c.eval = norm *
        exp(
            -(
                a * (x - centerX)^2. +
                b * (x - centerX) * (y - centerY) +
                c *                 (y - centerY)^2.
            )
        ) / 6.283185307179586 / sigmaX / sigmaY) # 2pi = 6.283185307179586
end


# The Computer Language Benchmarks Game
# https://salsa.debian.org/benchmarksgame-team/benchmarksgame/

# based on Oleg Mazurov's Java Implementation and Jeremy Zerfas' C implementation
# transliterated by Hamza Yusuf Çakır

global const preferred_num_blocks = 12

struct Fannkuch
    n::Int64
    blocksz::Int64
    maxflips::Vector{Int32}
    chksums::Vector{Int32}

    function Fannkuch(n, nthreads)
        nfact = factorial(n)

        blocksz = nfact ÷ (nfact < preferred_num_blocks ? 1 : preferred_num_blocks)
        maxflips = zeros(Int32,nthreads)
        chksums = zeros(Int32, nthreads)

        new(n, blocksz, maxflips, chksums)
    end
end

struct Perm
    p::Vector{Int8}
    pp::Vector{Int8}
    count::Vector{Int32}

    function Perm(n)
        p = zeros(Int8, n)
        pp = zeros(Int8, n)
        count = zeros(Int32, n)

        new(p, pp, count)
    end
end

Base.@propagate_inbounds function first_permutation(perm::Perm, idx)
    p = perm.p
    pp = perm.pp

    for i = 2:length(p)
        p[i] = i - 1
    end

    for i = length(p):-1:2
        ifact = factorial(i-1)
        d = idx ÷ ifact
        perm.count[i] = d
        idx = idx % ifact

        for j = 1:i
            pp[j] = p[j]
        end

        for j = 1:i
            p[j] = j+d <= i ? pp[j+d] : pp[j+d-i]
        end
    end
end

Base.@propagate_inbounds function next_permutation(perm::Perm)
    p = perm.p
    count = perm.count

    first = p[2]
    p[2]  = p[1]
    p[1]  = first

    i = 2
    while (count[i] + 1) >= i
        count[i] = 0
        i += 1

        next = p[1] = p[2]

        for j = 1:i-1
            p[j] = p[j+1]
        end

        p[i] = first
        first = next
    end
    count[i] += 1
end

Base.@propagate_inbounds @inline function count_flips(perm::Perm)
    p = perm.p
    pp = perm.pp

    flips = 1
    first = p[1] + 1
    if p[first] != 0

        copyto!(pp, 2, p, 2, length(p) - 1)

        while true
            flips += 1

            new_first = pp[first]
            pp[first] = first - 1

            if first > 3
                lo = 2; hi = first-1

                # see the note in Jeremy Zerfas' C implementation
                for k = 1:15
                    t = pp[lo]
                    pp[lo] = pp[hi]
                    pp[hi] = t

                    !((lo + 3) <= hi) && break
                    lo += 1
                    hi -= 1
                end
            end

            first = new_first + 1
            pp[first] == 0 && break
        end
    end

    return flips
end

Base.@propagate_inbounds function run_task(f::Fannkuch, perm::Perm, idxmin, idxmax)

    maxflips = 1
    chksum = 0

    let
        i = idxmin
        while true
            if perm.p[1] != 0
                flips = count_flips(perm)
                maxflips = max(maxflips, flips)
                chksum += iseven(i) ? flips : -flips
            end
            if i == idxmax
                break
            end
            i += 1
            next_permutation(perm)
        end
    end
    id = Threads.threadid()
    f.maxflips[id] = max(f.maxflips[id], maxflips)
    f.chksums[id] += chksum
end

function runf(f::Fannkuch)
    factn = factorial(f.n)

    Threads.@threads for idxmin=0:f.blocksz:factn-1
        perm = Perm(f.n)
        first_permutation(perm, idxmin)
        idxmax = idxmin + f.blocksz - 1
        @inbounds run_task(f, perm, idxmin, idxmax)
    end
end

function fannkuchredux(n)
    f = Fannkuch(n,Threads.nthreads())

    runf(f)

    # reduce results
    chk = sum(f.chksums)
    res = maximum(f.maxflips)

    println(chk, "\nPfannkuchen(", n, ") = ", res)
end

n = parse(Int, ARGS[1])
fannkuchredux(n)


import LinearAlgebra: qr, lu, eigen, diag

function feast!(X::AbstractMatrix, A::AbstractMatrix;
				nodes::Integer=8, iter::Integer=10, c=complex(0.0,0.0), r=1.0, ϵ=1e-12,
				debug=false, store=false, mixed_prec=false, factorizer=lu, left_divider=ldiv!)
	contour = circular_contour_trapezoidal(c, r, nodes)
	feast!(X, A, contour; iter=iter, debug=debug, ϵ=ϵ, store=store, mixed_prec=mixed_prec, factorizer=factorizer, left_divider=left_divider)
end

finalize!(x::Any) = nothing

function feast!(X::AbstractMatrix, A::AbstractMatrix, contour::Contour;
					 iter::Integer=10, ϵ=1e-12, debug=false, store=false, mixed_prec=false, factorizer=lu, left_divider=ldiv!)
	N, m₀ = size(X)
	if size(A, 1) != size(A, 2)
		 error("Incorrect dimensions of A, must be square")
	elseif size(A,1) != N
		 error("Incorrect dimensions of X, must match A")
	end

	Ctype = if mixed_prec ComplexF32 else ComplexF64 end

	Λ, resolvent, res = zeros(ComplexF64, m₀), zeros(ComplexF64, m₀), zeros(m₀)
	temp, R, Q = zeros(Ctype, N, m₀), similar(X, ComplexF64), deepcopy(X)
	Aq, Xq = zeros(ComplexF64, m₀, m₀), zeros(ComplexF64, m₀, m₀)

	ZmA = similar(A, Ctype)
	nodes = size(contour.nodes, 1)

	if store
		ZmA .= (A - I*contour.nodes[1])
		facts1 = factorizer(ZmA)
		facts = Array{typeof(facts1)}(undef, nodes)
		facts[1] = facts1

		Threads.@threads for i=2:nodes
			  ZmA .= (A - I*contour.nodes[i])
			  facts[i] = factorizer(ZmA)
		end
	end

	for nit=0:iter
		Q .= Matrix(qr(Q).Q)
		mul!(R, A, Q) ## why does this one allocate?
		mul!(Aq, Q', R) ### Aq = Q' * A * Q
		# mul!(Bq, Q', Q) ### Bq = Q' * Q = I
		F = eigen!(Aq)
		Λ .= F.values
		Xq .= F.vectors
		mul!(X, Q, Xq) ### Recover eigenvectors from Ritz vectors ( X = Q * Xq )
		update_R!(X, R, Λ, A) ### compute residual vectors R for RII update
		residuals!(res, R, Λ, A) ### compute actual residuals
		contour_nonempty = reduce(|, in_contour(Λ, contour))
		if debug iter_debug_print(nit, Λ, res, contour, 1e-5) end
		if contour_nonempty && maximum(res[in_contour(Λ, contour)]) < ϵ
			if debug println("converged in $nit iteration") end
			break
		end
		if nit < iter ### Do not solve linear systems / form Q on last iteration
			Q .= 0.00
			for i=1:nodes
				resolvent .= 1.0 ./(contour.nodes[i] .- Λ)
				if store
					left_divider(temp, facts[i], R)
				else
					ZmA .= A - I*contour.nodes[i]
					linsolve!(temp, ZmA, R, factorizer, left_divider)
				end

				temp .= X - temp
				rmul!(temp, Diagonal(resolvent .* contour.weights[i]))
				Q .+= temp
			end
		end
	end
	if store
		foreach(finalize!, facts)
	end
	contour_nonempty = reduce(|, in_contour(Λ, contour))
	if !contour_nonempty println("no eigenvalues found in contour!") end
	Λ[in_contour(Λ, contour)], X[:,in_contour(Λ, contour)], res[in_contour(Λ, contour)]
end

function gen_feast!(X::AbstractMatrix, A::AbstractMatrix, B::AbstractMatrix;
					nodes::Integer=8, iter::Integer=10, c=complex(0.0,0.0), r=1.0,
					debug=false, store=false, ϵ=1e-12, factorizer=lu, left_divider=ldiv!)
	contour = circular_contour_trapezoidal(c, r, nodes)
	gen_feast!(X, A, B, contour; iter=iter, debug=debug, ϵ=ϵ, factorizer=factorizer, left_divider=ldiv!)
end

function gen_feast!(X::AbstractMatrix, A::AbstractMatrix, B::AbstractMatrix, contour::Contour;
					iter::Integer=10, debug=false, store=false, ϵ=1e-12, factorizer=lu, left_divider=ldiv!)
	N, m₀ = size(X)
	if size(A, 1) != size(A, 2)
		error("Incorrect dimensions of A, must be square")
	elseif size(A,1) != N
		error("Incorrect dimensions of X, must match A")
	end

	Λ, resolvent, res = zeros(ComplexF64, m₀), zeros(ComplexF64, m₀), zeros(m₀)
	temp, R, Q = zeros(ComplexF64, N, m₀), similar(X, ComplexF64), copy(X)
	Aq, Bq, Xq = zeros(ComplexF64, m₀, m₀), zeros(ComplexF64, m₀, m₀), zeros(ComplexF64, m₀, m₀)
	ZmA = similar(A, ComplexF64)
	nodes = size(contour.nodes, 1)

	if store
		 ZmA .= (A - B*contour.nodes[1])
		 facts1 = factorizer(ZmA)
		 facts = Array{typeof(facts1)}(undef, nodes)
		 facts[1] = facts1

		 Threads.@threads for i=2:nodes
				ZmA .= (A - B*contour.nodes[i])
				facts[i] = factorizer(ZmA)
		 end
	end

	for nit=0:iter
		Q .= Matrix(qr(Q).Q)
		mul!(R, A, Q) ## why does this one allocate?
		mul!(Aq, Q', R) ### Aq = Q' * A * Q
		mul!(R, B, Q)
		mul!(Bq, Q', R) ### Bq = Q' * Q = I
		F = eigen!(Aq, Bq)
		Λ .= F.values
		Xq .= F.vectors
		mul!(X, Q, Xq) ### Recover eigenvectors from Ritz vectors ( X = Q * Xq )
		update_R!(X, R, Λ, A, B) ### compute residual vectors R for RII update
		residuals!(res, R, Λ, A) ### compute actual residuals
		contour_nonempty = reduce(|, in_contour(Λ, contour))
		if debug iter_debug_print(nit, Λ, res, contour, 1e-5) end
		if contour_nonempty && maximum(res[in_contour(Λ, contour)]) < ϵ
			  if debug println("converged in $nit iteration") end
			  break
		end
		if nit < iter ### Do not solve linear systems / form Q on last iteration
			Q .= 0.00
			for i=1:nodes
				resolvent .= (1.0 ./(contour.nodes[i] .- Λ))
				if store
					left_divider(temp, facts[i], R)
				else
					ZmA .= (A - B*contour.nodes[i])
					linsolve!(temp, ZmA, R, factorizer, left_divider)
				end
				temp .= X - temp
				rmul!(temp, Diagonal(resolvent .* contour.weights[i]))
				Q .+= temp
			end
		end
	end
	if store
		foreach(finalize!, facts)
	end
	contour_nonempty = reduce(|, in_contour(Λ, contour))
	if !contour_nonempty println("no eigenvalues found in contour!") end
	Λ[in_contour(Λ, contour)], X[:,in_contour(Λ, contour)], res[in_contour(Λ, contour)]
end

function dual_gen_feast!(Xr::AbstractMatrix, Xl::AbstractMatrix, A::AbstractMatrix, B::AbstractMatrix;
					nodes::Integer=8, iter::Integer=10, c=complex(0.0,0.0), r=1.0,
					debug=false, store=false, ϵ=1e-12, factorizer=lu, left_divider=ldiv!)
	contour = circular_contour_trapezoidal(c, r, nodes)
	dual_gen_feast!(Xr, Xl, A, B, contour; iter=iter, debug=debug, ϵ=ϵ, factorizer=factorizer, left_divider=ldiv!)
end

function dual_gen_feast!(Xr::AbstractMatrix, Xl::AbstractMatrix, A::AbstractMatrix, B, contour::Contour;
					iter::Integer=10, debug=false, store=false, ϵ=1e-12, factorizer=lu, left_divider=ldiv!, resize_subspace=true)
	N, m₀ = size(Xl)
	if size(A, 1) != size(A, 2)
		error("Incorrect dimensions of A, must be square")
	elseif size(A,1) != N
		error("Incorrect dimensions of X, must match A")
	end

	Λ, resolvent, resr, resl, temp, Rr, Rl, Ql, Qr, Aq, Bq, Xq = init(Xr, Xl, N, m₀)
	ZmA = similar(A, ComplexF64)
	nodes = size(contour.nodes, 1)

	if store
		if debug print("Precomputing factorizations: ") end
			ZmA .= (A - B*contour.nodes[1])
			rfacts1 = factorizer(ZmA)
			rfacts = Array{typeof(rfacts1)}(undef, nodes)
			rfacts[1] = rfacts1
			if debug print("1R ") end
			ZmA .= (A - B*contour.nodes[1])'
			lfacts1 = factorizer(ZmA)
			lfacts = Array{typeof(lfacts1)}(undef, nodes)
			lfacts[1] = lfacts1
			if debug print("1L ") end
			for i=2:nodes
				ZmA .= (A - B*contour.nodes[i])
				rfacts[i] = factorizer(ZmA)
				if debug print("$(i)R ") end
				ZmA .= (A - B*contour.nodes[i])'
				lfacts[i] = factorizer(ZmA)
				if debug print("$(i)L ") end
			end
		if debug println() end
	end

	try
	nin = 0
	for nit=0:iter
		if resize_subspace
			new_m₀ = new_subspace_size(nin, m₀)
			if new_m₀ > 0 && new_m₀ != m₀
				if debug println("Resizing subspace from $(m₀) to $(new_m₀)") end
				m₀ = new_m₀
				Λ, Xr, Xl, resolvent, resr, resl, temp, Rr, Rl, Ql, Qr, Aq, Bq, Xq = reinit(Λ, Xr, Xl, resolvent, resr, resl, temp, Rr, Rl, Ql, Qr, Aq, Bq, Xq, new_m₀)
			end
		end
		S = svd!(Ql'*B*Qr)
		Qr .= Qr*S.V*Diagonal(1.0/S.S)
		Ql .= Ql*S.U*Diagonal(1.0/S.S)
		mul!(Rr, A, Qr) ## why does this one allocate?
		mul!(Aq, Ql', Rr) ### Aq = Q' * A * Q
		mul!(Rr, B, Qr)
		mul!(Bq, Ql', Rr) ### Bq = Q' * Q = I
		F = eigen(Aq, Bq)
		Λ .= F.values
		Xq .= F.vectors
		mul!(Xr, Qr, Xq) ### Recover eigenvectors from Ritz vectors ( X = Q * Xq )
		F = eigen(Aq', Bq')
		Xq .= F.vectors
		mul!(Xl, Ql, Xq) ### Recover eigenvectors from Ritz vectors ( X = Q * Xq )
		update_R!(Xr, Rr, Λ, A, B) ### compute residual vectors R for RII update
		update_R!(Xl, Rl, Λ, A', B') ### compute residual vectors R for RII update
		residuals!(resr, Rr, Λ, A) ### compute actual residuals
		incntr = in_contour(Λ, contour)
		nin = sum(resr[incntr] .< 1e-5)
		contour_nonempty = reduce(|, incntr)
		if debug iter_debug_print(nit, Λ, resr, contour, 1e-5) end
		if contour_nonempty && maximum(resr[in_contour(Λ, contour)]) < ϵ
			if debug println("converged in $nit iteration") end
			break
		end
		if nit < iter ### Do not solve linear systems / form Q on last iteration
			Qr .= 0.00
			Ql .= 0.00
			for i=1:nodes
				resolvent .= (1.0 ./(contour.nodes[i] .- Λ))
				if store
					left_divider(temp, rfacts[i], Rr)
				else
					ZmA .= (A - B*contour.nodes[i])
					linsolve!(temp, ZmA, Rr, factorizer, left_divider)
				end
				temp .= Xr - temp
				rmul!(temp, Diagonal(resolvent .* contour.weights[i]))
				Qr .+= temp

				resolvent .= (1.0 ./(conj(contour.nodes[i]) .- conj.(Λ)))
				if store
					left_divider(temp, lfacts[i], Rl)
				else
					ZmA .= (A - B*contour.nodes[i])'
					linsolve!(temp, ZmA, Rl, factorizer, left_divider)
				end
				temp .= Xl - temp
				rmul!(temp, Diagonal(resolvent .* conj(contour.weights[i])))
				Ql .+= temp
			end
		end
	end

	contour_nonempty = reduce(|, in_contour(Λ, contour))
	if !contour_nonempty
		println("no eigenvalues found in contour!")
	end
	output = Λ[in_contour(Λ, contour)], Xr[:,in_contour(Λ, contour)], Xl[:,in_contour(Λ, contour)], resr[in_contour(Λ, contour)]
	return output
	finally
		if store
			foreach(finalize!, rfacts)
			foreach(finalize!, lfacts)
		end
	end
end

function linsolve!(Y, C, X, factorizer, left_divider)
	F = factorizer(C)
	left_divider(Y, F, X)
	finalize!(F)
end

function init(Xr, Xl, N::Int, m₀::Int)
	Λ, resolvent, resr, resl = zeros(ComplexF64, m₀), zeros(ComplexF64, m₀), zeros(m₀), zeros(m₀)
	temp, Rr, Rl, Ql, Qr = zeros(ComplexF64, N, m₀), similar(Xr, ComplexF64), similar(Xr, ComplexF64), copy(Xl), copy(Xr)
	Aq, Bq, Xq = zeros(ComplexF64, m₀, m₀), zeros(ComplexF64, m₀, m₀), zeros(ComplexF64, m₀, m₀)
	return Λ, resolvent, resr, resl, temp, Rr, Rl, Ql, Qr, Aq, Bq, Xq
end

function reinit(Λ, Xr, Xl, resolvent, resr, resl, temp, Rr, Rl, Ql, Qr, Aq, Bq, Xq, new_m₀, fillf::F=x->randn(typeof(x))) where F
	N, m₀ = size(temp,1), length(Λ)
	Xrn, Xln = zeros(ComplexF64, N, new_m₀), zeros(ComplexF64, N, new_m₀)
	map!(fillf, Xrn, Xrn)
	map!(fillf, Xln, Xln)
	copyto!(Xrn, Xr)
	copyto!(Xln, Xl)
	Λn, resolventn, resrn, resln, tempn, Rrn, Rln, Qln, Qrn, Aqn, Bqn, Xqn = init(Xrn, Xln, N, new_m₀)
	copyto!(Λn, Λ)
	copyto!(resolventn, resolvent)
	copyto!(resrn, resr)
	copyto!(resln, resl)
	copyto!(tempn, temp)
	copyto!(Rrn, Rr)
	copyto!(Rln, Rl)
	copyto!(Qln, Ql)
	copyto!(Qrn, Qr)
	copyto!(Aqn, Aq)
	copyto!(Bqn, Bq)
	copyto!(Xqn, Xqn)
	return Λn, Xrn, Xln, resolventn, resrn, resln, tempn, Rrn, Rln, Qln, Qrn, Aqn, Bqn, Xqn
end

function new_subspace_size(nin, m₀, fac_test=0.75, fac_new=2.0)
	if nin > ceil(Int, fac_test*m₀)
		return ceil(Int, fac_new*nin)
	else
		return 0
	end
end


using KernelDensityEstimate, Gadfly

toggleYTicks()

p = resample(kde!([0.0],[10.0]),300);
q = resample(kde!([-8.0; 13.0],[1.5]),300);

# easy product operator example, use prodAppxMSGibbsS(...) for more options and multiple terms
# dummy = kde!(rand(1,300),[1.0]);
# pGM, = prodAppxMSGibbsS(dummy, [p;q], Union{}, Union{}, 5)
pq = p*q

pl = plotKDE([p;q;pq],c=["blue";"red";"deepskyblue"],legend=["p(Θ | Y)";"p(Θ | Z)";"p(Θ | Y, Z)"],xlbl="p(Θ | Y, Z) ∝ p(Θ | Y) p(Θ | Z)")

r = resample(kde!([-35;-11.0; 26.0],[2.5]),300);

dummy = kde!(rand(1,300),[1.0]);
pGM, = prodAppxMSGibbsS(dummy, [p;q;r], Union{}, Union{}, 5)
pqr = kde!(pGM);

pl2 = plotKDE([p;q;r;pqr],c=["blue";"red";"green";"deepskyblue"],legend=["p(Θ | Y)";"p(Θ | Z1)";"p(Θ | Z2)";"p(Θ | Y, Z:)"],xlbl="p(Θ | Y, Z:) ∝ p(Θ | Y) p(Θ | Z1) p(Θ | Z2)")

draw(PDF("consensus.pdf",12cm,12cm),vstack(pl,pl2))


"""
    Pendulum{R}

The most simple canonical nonlinear system, consisting of a simple pendulum controlled by
torque at the base. Has 2 state and 1 control.

# Constructor
    Pendulum(; kwargs...)

with keyword arguments
* `mass` - mass of the pendulum, in kg (default = 1.0)
* `length` - length of the pendulum, in m (default = 0.5)
* `b` - friction coefficient (default=0.1)
* `I` - inertia, in kg⋅m² (default=0.25)
* `g` - gravity, in kg/m² (default=9.81)
"""

@autodiff mutable struct Pendulum{T} <: ContinuousDynamics
    mass::T
    len::T
    b::T
    lc::T
    I::T
    g::T
end
function Pendulum(;mass=1., len=0.5, b=0.1, lc=0.5, I=0.25, g=9.81)
    T = eltype(promote(mass, len, b, lc, I, g))
    Pendulum{T}(mass, len, b, lc, I, g)
end

RobotDynamics.state_dim(::Pendulum) = 2
RobotDynamics.control_dim(::Pendulum) = 1

function dynamics(p::Pendulum, x, u)
    m = p.mass * p.lc * p.lc
    @SVector [x[2],
              u[1]/m - p.g*sin(x[1])/p.lc - p.b*x[2]/m]
end
function dynamics!(p::Pendulum, ẋ, x, u)
    ẋ .= dynamics(p, SA[x[1], x[2]], SA[u[1]])
end
@inline Base.position(::Pendulum, x) = @SVector zeros(3)
orientation(::Pendulum, x) = expm((pi-x[1])*@SVector [1,0,0.])


J = j[1].RB2.J;
β2 = j[1].RB2.x[4:7]
# β2 = []
function M_β(β::Vector{T}) where T<:Real
    E = genE(β)
    M = 4*transpose(E)*J*E
    if typeof(β[1]) == Float64
        Ms = real(sqrt(M))
        # C = cholesky(M)
        # println("L = ", norm(Ms - C.L))
        # println(norm(Ms - C.U))
        # Ms = (sqrt.(diag(M))).*Matrix{T}(I,size(M))
    else # ForwardDiff.Dual variable
        # C = cholesky(M)
        # Ms = C.U

        # Ms = (sqrt.(diag(M))).*Matrix{T}(I,size(M))

        # diag_Mfinal = diag(Mfinal);
        # vals_sq = zeros(T,size(Mfinal,1))
        # for i=1:size(Mfinal,1)
        #     vals_sq[i] = sqrt(diag_Mfinal[i])
        # end
        # Ms = diagm(vals_sq);

        # Ms = (Mfinal^(0.5))

        vals, vecs = eigen(M);
        vals_sq = zeros(T,size(vals))
        for i=1:length(vals)
            vals_sq[i] = sqrt(vals[i])
        end
        Ms = vecs*diagm(vals_sq)*(vecs\I(size(vecs,1)))
    end
    return M, Ms
end
# M_β(β2)
m = central_fdm(10,1)
temp1 = FiniteDifferences.jacobian(m, z->M_β(z)[2], β2)[1]
temp2 = ForwardDiff.jacobian(z->M_β(z)[1], β2)
temp1_nz = findall(x -> abs(x)>1e-6, temp1)
temp2_nz = findall(x -> abs(x)>1e-6, temp2)
norm(temp1 - temp2)

# function kronSum(A::Array{T,2},B::Array{T,2}) where T<:Real
#     out = kron(A, Matrix{T}(I,(size(B)))) + kron(Matrix{T}(I,(size(A))),B)
#     return out
# end

M_check, Mhalf_check = M_β(β2)
dM_half = inv(kronSum(permutedims(Mhalf_check), Mhalf_check))
norm(dM_half*temp2 - temp1)

## trying svd for square root
# clearconsole()
# using  FiniteDifferences
# using LinearAlgebra
# using GenericLinearAlgebra
# using ForwardDiff
# using Revise
# function sqrtM(x::Array{T}) where T<:Real
#     # M = assembleM(x,j)
#     M = genMatM(x,j[1].RB2)
#     M1 = M[4:7,4:7]#[5:7,5:7]
#     # E = genE(q)
#     # M = 4*transpose(E)*J*E
#     if typeof(M[1]) == Float64
#         Ms = real(sqrt(M1))
#         # println("Ms = ", Ms)
#         # F = GenericLinearAlgebra.svd!(M)
#         # Ms2 = permutedims(F.Vt)*diagm(sqrt.(F.S))*F.Vt
#         # println("Mserr = ", norm(Ms - Ms2))
#     else
#         F = GenericLinearAlgebra.svd!(M1)
#         Ms = permutedims(F.Vt)*diagm(sqrt.(F.S))*F.Vt
#     end
#     return M, Ms
# end
# function sqrtMq(q::Array{T}) where T<:Real
#     E = genE(q)
#     M = 4*transpose(E)*j[1].RB2.J*E
#     if typeof(M[1]) == Float64
#         out = real(sqrt(M))
#         # F = GenericLinearAlgebra.svd!(M)
#         # Ms = permutedims(F.Vt)*diagm(sqrt.(F.S))*F.Vt
#         # println("Mserr = ", norm(out - Ms))
#     else
#         F = GenericLinearAlgebra.svd!(M)
#         out = permutedims(F.Vt)*diagm(sqrt.(F.S))*F.Vt
#     end
#     return M, out
# end
# # x0 = [0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.24999999999999997, 0.0, -0.43301270189221935, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.294735414289422];
# # q2 = [sqrt(2)/2; sqrt(2)/2;zeros(2)]
# x0 = [0.0; 0.0; 0.0; 1.0; 0.0; 0.0; 0.0; 0.0; 0.0; 0.0; 0.0; 0.0; 0.0; 0.0; 0.24999999999999997; 0.0; -0.43301270189221935; [1;zeros(3)]; zeros(7)];
# # x0 = [0.0; 0.0; 0.0; 1.0; 0.0; 0.0; 0.0; 0.0; 0.0; 0.0; 0.0; 0.0; 0.0; 0.0; 0.24999999999999997; 0.0; -0.43301270189221935; randomQuat(); zeros(7)];
# updateRigidBody!(j[1].RB2,x0)
# # x0 = randomQuat();
# _,tempMs = sqrtM(x0)
# # temp = rand(10)
# m = central_fdm(10,1)
# finJ = FiniteDifferences.jacobian(m, z->sqrtM(z)[2], x0)[1]
# fdJ = ForwardDiff.jacobian(z->sqrtM(z)[2],x0)
# finJ_nz = findall(x -> abs(x)>1e-6, finJ)
# fdJ_nz = findall( x-> abs(x)>0, fdJ)
# println()
# println(norm(fdJ - finJ))
#
# # q2 = randomQuat();
# q2 = (x0[18:21])
# sqrtMq(q2)
# finJ2 = FiniteDifferences.jacobian(m, z->sqrtMq(z)[2], q2)[1]
# fdJ2 = ForwardDiff.jacobian(z->sqrtMq(z)[2],q2)
# println(norm(fdJ2 - finJ2))
##
# J = j[1].RB2.J;
# function pinvM(q::Array{T}) where T<:Real
#     E = genE(q)
#     M = 4*transpose(E)*J*E
#     Ms = sqrt(diagm(M))
#     F = GenericLinearAlgebra.svd!(diagm(M))
#     out = pinv(diagm(M))
#     return out
# end
# ForwardDiff.jacobian(pinvM, temp)
## symbolic
# using SymPy
#
# u0, u1, u2, u3 = symbols("u0 u1 u2 u3")
#
# E = [u0 u1 u2 u3;
#     -u1 u0 u3 -u2;
#     -u2 -u3 u0 u1;
#     -u3 u2 -u1 u0];
#
# J = j[1].RB2.J;
# M = 4*permutedims(E)*J*E
# F = GenericLinearAlgebra.svd!(M)
# # Ms = (M)^(2)
# diffM = [M.diff((u0,1))[:] M.diff((u1,1))[:] M.diff((u2,1))[:] M.diff((u3,1))[:]]
#
# diffM_eval = zeros(size(diffM))
# for i=1:length(diffM)
#     diffM_eval[i] = diffM[i](1.0, 0.0, 0.0, 0.0)
# end
# diffM_eval
# M[1,1](1.0, 0.0, 0.0, 0.0)
##
# function invM(x)
#     M = diagm(x)
#     return inv(M)
# end
# r=rand(3)
# finJ_M = FiniteDifferences.jacobian(m,z->invM(z), r)[1]
# fdJ_M = ForwardDiff.jacobian(z->invM(z), r)
# println("err= ", norm(finJ_M - fdJ_M))
## Mdiff function
clearconsole()
function kronSum(A::Array{T,2},B::Array{T,2}) where T<:Real
    out = kron(A, Matrix{T}(I,(size(B)))) + kron(Matrix{T}(I,(size(A))),B)
    return out
end
# x0 = [0.0; 0.0; 0.0; 1.0; 0.0; 0.0; 0.0; 0.0; 0.0; 0.0; 0.0; 0.0; 0.0; 0.0; 0.24999999999999997; 0.0; -0.43301270189221935; [1;zeros(3)]; zeros(7)];
# x0 = [0.0; 0.0; 0.0; 1.0; 0.0; 0.0; 0.0; 0.0; 0.0; 0.0; 0.0; 0.0; 0.0; 0.0; 0.24999999999999997; 0.0; -0.43301270189221935; randomQuat(); zeros(7)];

function Mdiff(x::Vector{T}) where T<:Real
    M = assembleM(x,j)
    Ms = real(sqrt(M))
    M_inv = inv(M)
    Ms_inv = real(sqrt(M_inv))

    dM = ForwardDiff.jacobian(z -> assembleM(z,j),x)
    dM_inv = -(kron(permutedims(inv(M)),inv(M)))*ForwardDiff.jacobian(z->assembleM(z,j), x)
    # dM_inv = ForwardDiff.jacobian(z -> inv(assembleM(z,j)),x)
    dMs = inv(kronSum(permutedims(Ms), Ms))*dM
    dMs_inv = -kron(permutedims(Ms_inv),Ms_inv)*(dMs);
    return M, Ms, Ms_inv, dM, dMs, dMs_inv, M_inv, dM_inv
end
m = central_fdm(5,1);
# finJ_dMs = FiniteDifferences.jacobian(m, z->Mdiff(z)[2], x0)[1]
# finJ_dMs_inv = FiniteDifferences.jacobian(m, z->Mdiff(z)[3], x0)[1]
# println()
# println(norm(finJ_dMs - Mdiff(x0)[5]))
# println(norm(finJ_dMs_inv - Mdiff(x0)[6]))
#
# _,_,Ms_inv,_,dMs, fdJ_dMs_inv = Mdiff(x0);
# # fdJ_dMs_inv = -kron(permutedims(Ms_inv),Ms_inv)*(dMs);
# println(norm(finJ_dMs_inv - fdJ_dMs_inv))
##
clearconsole()
function t1Fn(x)
    A, b = Ab_VecOfMat(x,j)
    M = assembleM(x,j)
    M, _, _, _, _, _, M_inv, dM_inv = Mdiff(x);
    # M = genMatM(x,j[1].RB2)
    # M1 = M#[4:7,4:7]
    # F = svd(M);
    # M_inv = F.V*diagm(1./(F.S))*(permutedims(F.U))
    # println("err = ", norm(inv(M) - pinv(M)))
    # sleep(1000)
    # M, _, Ms_inv, _, _, dMs_inv = Mdiff(x)
    sz_A1 = size(A,1)
    # t1 = permutedims(A)*((A*(M\(permutedims(A))))\I(sz_A1))
    t1 = permutedims(A)*inv(A*inv(M)*permutedims(A))

    dA = ForwardDiff.jacobian(z->Ab_VecOfMat(z,j)[1],x);
    dA_tr = ForwardDiff.jacobian(z->permutedims(Ab_VecOfMat(z,j)[1]),x);

    Y = A*M_inv*permutedims(A); invY = inv(Y);
    dY1 = kron(A*M_inv,1.0I(sz_A1))*dA;
    dY2 = kron(A,A)*dM_inv;
    dY3 = kron(1.0I(sz_A1), A*M_inv)*dA_tr;
    dY = dY1 + dY2 + dY3;
    dY_inv = -kron(permutedims(invY),invY)*dY;

    dt1_1 = kron(permutedims(invY),1.0I(size(A,2)))*dA_tr;
    dt1_2 = kron(1.0I(sz_A1),permutedims(A))*dY_inv;
    dt1 = dt1_1 + dt1_2;
    return t1, dt1
end
function Gdiff(x::Vector{T}) where T<:Real
    A, b = Ab_VecOfMat(x,j)
    M, _, Ms_inv, _, _, dMs_inv, _, _ = Mdiff(x)

    G = A*Ms_inv;
    Gp = permutedims(G)*((G*permutedims(G))\I(size(G,1)))

    t1, dt1 = t1Fn(x)
    dGp1 = kron(permutedims(t1),Matrix{T}(I,size(M)))*dMs_inv

    sz_t2 =size(t1,2)
    t2 = kron(Matrix{T}(I,(sz_t2,sz_t2)), Ms_inv)
    dGp2 = t2*dt1;

    dGp = dGp1 + dGp2
    return Gp, dGp
end
# Gdiff(x0)
finJ_G = FiniteDifferences.jacobian(m,z->Gdiff(z)[1],x0)[1]
fdJ_G = Gdiff(x0)[2]#ForwardDiff.jacobian(z->Gdiff(z)[2],x0)
println("\nerr_Gp", norm(finJ_G - fdJ_G))
#
# finJ_t1 = FiniteDifferences.jacobian(m,z->t1Fn(z)[1], x0)[1]
# fdJ_t1 = t1Fn(x0)[2]
# println("\nerr_t1= ", norm(finJ_t1 - fdJ_t1))

##
x0, u = getXU_0(j);

function hdiff(x)
    Fu = assembleF(x,u,j,GravityInInertial)
    A, b = Ab_VecOfMat(x,j)
    M = assembleM(x,j)

    a = inv(M)*Fu
    h = b-A*a
    return h
end
function Fcdiff(x::Vector{T}) where T<:Real
    Fu = assembleF(x,u,j,GravityInInertial)
    A, b = Ab_VecOfMat(x,j)
    M, Ms, _, _, dMs, _, _, _ = Mdiff(x)

    a = M\Fu
    h = b-A*a

    Gp, dGp = Gdiff(x)

    Fc = Ms*Gp*h

    dFc1 = kron(permutedims(Gp*h),I(size(M,1)))*dMs
    dFc2 = kron(permutedims(h),Ms)*dGp
    dFc3 = Ms*Gp*ForwardDiff.jacobian(z->hdiff(z),x)

    dFc = dFc1 + dFc2 + dFc3;

    dFu = ForwardDiff.jacobian(z->assembleF(z,u,j,GravityInInertial),x)
    return Fc, dFc, Fu, dFu
end
finJ_Fc = FiniteDifferences.jacobian(m,z->Fcdiff(z)[1],x0)[1]
fdJ_Fc = Fcdiff(x0)[2]
println("\nerr_Fc = ", norm(finJ_Fc - fdJ_Fc))

finJ_Fu = FiniteDifferences.jacobian(m,z->Fcdiff(z)[3],x0)[1]
fdJ_Fu = Fcdiff(x0)[4]
println("\nerr_Fu = ", norm(finJ_Fu - fdJ_Fu))

## fxdotdiff
invJ = j[1].RB2.invJ;
nJ = length(j); # also equal to number of bodies (excluding inertial frame)
function accdiff(x::Vector{T}) where T<:Real
    len_x = length(x)-14
    # dacc = zeros(T,(len_x,len_x))

    M = assembleM(x,j)
    Fc, dFc, Fu, dFu = Fcdiff(x)

    # dM_inv = ForwardDiff.jacobian(z ->inv(assembleM(z,j)), x)

    # dacc1 = kron(permutedims(Fc+Fu),1.0I(size(M,1)))*dM_inv
    # dacc2 = inv(M)*(dFc + dFu)
    #
    # dacc = dacc1 + dacc2
    dacc = zeros(T,(7*nJ,length(x)))
    for k=1:nJ
        ct1 = 7*(k-1); ct2 = 14*k;
        dacc[ct1+1:ct1+3,:] = 1/j[k].RB2.m*(dFc[ct1+1:ct1+3,:] +   dFu[ct1+1:ct1+3,:])

        E = genE(x[ct2+4:ct2+7])
        fdJ_E = ForwardDiff.jacobian(z->genE(z[ct2+4:ct2+7]),x)
        fdJ_trE = ForwardDiff.jacobian(z->transpose(genE(z[ct2+4:ct2+7])),x)

        dacc2_1 = kron(permutedims(invJ*E*((Fc+Fu)[ct1+4:ct1+7])),I(size(E,1)))*fdJ_trE
        dacc2_2 = kron(permutedims((Fc+Fu)[ct1+4:ct1+7]),permutedims(E)*invJ)*fdJ_E
        dacc2_3 = permutedims(E)*invJ*E*(dFc[ct1+4:ct1+7,:] + dFu[ct1+4:ct1+7,:])

        # println("\ndacc2_1 = ", dacc2_1)
        # println("\ndacc2_2 = ", dacc2_2)
        # println("\ndacc2_3 = ", dacc2_3)
        # sleep(1000);
        dacc[ct1+4:ct1+7,:] = 1/4*(dacc2_1 + dacc2_2 + dacc2_3)
    end
    return dacc
end
# accdiff(x0)
##
function fxdotdiff(x::Vector{T}) where T<:Real
    len_x = length(x)
    dFx = zeros(T,(len_x,len_x))
    # dFx = Matrix{T}(undef,(len_x,len_x))

    dacc = accdiff(x);
    # println("dacc = ", dacc)

    for k=1:length(j)
        dFx[14*k+1:14*k+7,14*k+8:14*k+14] = 1.0*I(7)
        dFx[14*k+8:14*k+14,:] = dacc[7*(k-1)+1:7*k,:]
    end
    return dFx
end

finJ_fx = FiniteDifferences.jacobian(m,z->fxdot(z,u,j,GravityInInertial),x0)[1];

fdJ_fx = fxdotdiff(x0);
# println("\nerr_fx = ", norm(finJ_fx[25:end,:] - fdJ_fx[25:end,:]))
println("\nerr_fx = ", norm(finJ_fx- fdJ_fx))
## M for β = [1.0;zeros(3)]
q = deepcopy(x0); q[18:21] = [1.0;zeros(3)]
function Mq(x::AbstractVector{T}) where T<:Real
    M = assembleM(x,j)
    M1 = M[4:7,4:7]
    out = (M1)#\I(size(M1,1))
    return out
end
finJ_M = FiniteDifferences.jacobian(m,z->inv(Mq(z)), q)[1]
M1 = Mq(q);
fdJ_M = -(kron(permutedims(inv(M1)),inv(M1)))*ForwardDiff.jacobian(Mq, q)
println("\nerr_invM = ", norm(finJ_M - fdJ_M))
## symbolic trying in login node (hprc)
# using SymPy
# using LinearAlgebra
# u0, u1, u2, u3 = symbols("u0 u1 u2 u3")
# uSym = [u0;u1;u2;u3];
# E = [u0 u1 u2 u3;
#     -u1 u0 u3 -u2;
#     -u2 -u3 u0 u1;
#     -u3 u2 -u1 u0];
#
# J = j[1].RB2.J;
# M = 4*permutedims(E)*J*E
# invM = (M)^(-1)s
## forwarddiff benchmarks
# function tempFn(z)
#     out = z.^2;
#     return out
# end
# y = rand(100);
# fdJ_fn = x->ForwardDiff.jacobian(tempFn,x)
# @btime fdJ_fn(y);
# @btime ForwardDiff.jacobian(tempFn,y);

jb1_nz = findall(x -> abs(x) > 1e-6, jb1)
jb2_nz = findall(x -> abs(x) > 0, jb2)


ds = Systems.standardmap()
p_index = 1

ds, pmin, pmax = Systems.standardmap(), 0.6, 1.2

integ = integrator(ds)

for i in 1:2
od, xmin, xmax = InteractiveChaos.minimal_normalized_od(integ, i, p_index, pmin, pmax,
                 100, 100, 100, get_state(ds))

@test xmin ≥ 0
@test xmax ≤ 2π
end

od, xmin, xmax = InteractiveChaos.minimal_normalized_od(integ, 1, p_index, pmin, pmax,
                 100, 100, 100, get_state(ds))

# Now to simulate selecting a rectangle:
rect = Observable((origin = [0., 0], widths = [1., 1.]))
OD = Observable(od)
od0 = copy(od)

on(rect) do r
    pmin, xmin = r.origin
    pmax, xmax = r.origin + r.widths

    OD[] = InteractiveChaos.minimal_normalized_od(
        integ, 1,  p_index, pmin, pmax,
        100, 100, 100, get_state(ds), xmin, xmax
    )
end

# change rectangle:
rect[] = (origin = [0.8, 0.2], widths = [1., 5.0])

@test od0 != OD[]


module Staggered
using SparseArrays, LinearAlgebra
import sbp
using sbp.Operator: Operators2D
using sbp.Metrics: CovariantBasis
using sbp.Sparse: block_matrix, block_matrix_insert

function boundary_matrix_p(n, is_sparse=true)
        if is_sparse
                B = spzeros(n, n + 1)
        else
                B = zeros(n, n + 1)
        end

        B[1,1] = -1
        B[end,end] = 1

        return B
end

function boundary_matrix_m(n, is_sparse=true)
        if is_sparse
                B = spzeros(n, n - 1)
        else
                B = zeros(n, n - 1)
        end

        B[1,1] = -1
        B[end,end] = 1

        return B
end

function build_operators_2d(Dx::AbstractArray,
                            Dy::AbstractArray,
                            Px::AbstractArray,
                            Py::AbstractArray,
                            Hx::AbstractArray,
                            Hy::AbstractArray,
                            Bx::AbstractArray,
                            By::AbstractArray
                            )
        nx = size(Dx, 1)
        ny = size(Dy, 1)

        Ix = sparse(1.0I, nx, nx)
        Iy = sparse(1.0I, ny, ny)

        Dx2 = kron(Dx, Iy)
        Dy2 = kron(Ix, Dy)

        Px2 = kron(Px, Iy)
        Py2 = kron(Ix, Py)

        Hx2 = kron(Hx, Iy)
        Hy2 = kron(Ix, Hy)

        Bx2 = kron(Bx, Iy)
        By2 = kron(Ix, By)

        return Operators2D(nx, ny, Dx2, Dy2, Px2, Py2, Hx2, Hy2, Bx2, By2)

end

function build_all_operators_2d(builder::Function, nx::Int64, ny::Int64)
        xp, xm, h, Dxp, Dxm, Hxp, Hxm, Pxp, Pxm = builder(nx)
        yp, ym, h, Dyp, Dym, Hyp, Hym, Pyp, Pym = builder(ny)

        Bxp = boundary_matrix_p(nx)
        Bxm = boundary_matrix_m(nx+1)
        Byp = boundary_matrix_p(ny)
        Bym = boundary_matrix_m(ny+1)

        pp = build_operators_2d(Dxp, Dyp, Pxp, Pyp, Hxp, Hyp, Bxp, Byp)
        mm = build_operators_2d(Dxm, Dym, Pxm, Pym, Hxm, Hym, Bxm, Bym)
        pm = build_operators_2d(Dxp, Dym, Pxp, Pym, Hxp, Hym, Bxp, Bym)
        mp = build_operators_2d(Dxm, Dyp, Pxm, Pyp, Hxm, Hyp, Bxm, Byp)

        return pp, mm, pm, mp
end

"""
Construct the covariant basis at the cell centers (mm) using SBP operators.
The mapping function f(x,y) = (fx, fy) must be defined at the nodes (pp).
"""
function build_covariant_basis_mm(fx::AbstractArray, fy::AbstractArray,
                               mm::Operators2D, pm::Operators2D)
        x_r1 = mm.Dx * pm.Py * fx
        x_r2 = mm.Px * pm.Dy * fx
        y_r1 = mm.Dx * pm.Py * fy
        y_r2 = mm.Px * pm.Dy * fy

        return CovariantBasis(x_r1, y_r1, x_r2, y_r2)
end

"""
Construct the covariant basis at the v^1 components (pm) using SBP operators.
The mapping function f(x,y) = (fx, fy) must be defined at the nodes (pp).
"""
function build_covariant_basis_pm(fx::AbstractArray, fy::AbstractArray,
                               mm::Operators2D, pm::Operators2D,
                               mp::Operators2D)
        x_r1 = pm.Px * mm.Py * mp.Dx * fx
        x_r2 = pm.Dy * fx
        y_r1 = pm.Px * mm.Py * mp.Dx * fy
        y_r2 = pm.Dy * fy

        return CovariantBasis(x_r1, y_r1, x_r2, y_r2)
end

"""
Construct the covariant basis at the v^2 components (mp) using SBP operators.
The mapping function f(x,y) = (fx, fy) must be defined at the nodes (pp).
"""
function build_covariant_basis_mp(fx::AbstractArray, fy::AbstractArray,
                               mm::Operators2D, pm::Operators2D,
                               mp::Operators2D)
        x_r1 = mp.Dx * fx
        x_r2 = mp.Py * mm.Px * pm.Dy * fx
        y_r1 = mp.Dx * fy
        y_r2 = mp.Py * mm.Px * pm.Dy * fy

        return CovariantBasis(x_r1, y_r1, x_r2, y_r2)
end

end


export integral, integrate

const RTOL_DEFAULT = 1e-4

integral(args...; kwargs...) = integrate(args...; kwargs...).val

function integrate(vc::VoxelwiseConstant{K,AN}) where {K,AN}
    CI = CartesianIndices(size(vc))
    val = zero(eltype(AN))
    for ci = CI
        ∆lvxl_ci = map((∆lvxlₖ,iₖ) -> ∆lvxlₖ[iₖ], vc.∆lvxl, ci.I)
        val += vc.val[ci] * prod(∆lvxl_ci)
    end

    return (val=val, err=0.0)
end

function integrate(vc::VoxelwiseConstant,
                   dom;  # dom(x) returns true if x is inside domain; false otherwise
                   kwargs...)
    val, err = hcubature(x -> vc(x) * dom(x), bounds(vc)...; rtol=RTOL_DEFAULT, kwargs...)
    return (val=val, err=err)
end


# load

using Flux
using Flux: onehotbatch, argmax, crossentropy, throttle, @epochs
using BSON: @save, @load
using Base.Iterators: partition

m = Chain(
  Conv((2,2), 1=>16, relu),
  x -> maxpool(x, (2,2)),
  Conv((2,2), 16=>8, relu),
  x -> maxpool(x, (2,2)),
  x -> reshape(x, :, size(x, 4)),
  Dense(288, 10), softmax) |> gpu

loss(x, y) = crossentropy(m(x), y)

accuracy(x, y) = mean(argmax(m(x)) .== argmax(y))

evalcb = throttle(() -> @show(accuracy(tX, tY)), 1)
opt = ADAM(params(m))

@load "fashion_mnist_model.bson" m
@load "fashion_mnist_weights.bson" weights
Flux.loadparams!(m, weights)

using MLDatasets # FashionMNIST
using ColorTypes: N0f8, Gray

const Img = Matrix{Gray{N0f8}}

function prepare_test()
    # load full test set
    test_x,  test_y  = FashionMNIST.testdata() # 10_000
    test_imgs = Img.([test_x[:,:,i] for i in 1:10_000])

    # Prepare test set
    tX = cat(4, float.(test_imgs[1:10_000])...) |> gpu
    tY = onehotbatch(test_y[1:10_000], 0:9) |> gpu
    tX, tY
end

tX, tY = prepare_test()
@show(accuracy(tX, tY))
#  accuracy(tX, tY) = 0.3329


function maximum_hourglass()
    # Defined 2d Matrix
    B = Matrix{Int64}(undef,6,6)

    for i in 1:1:6
        A = map((x) -> parse(Int64,x) , split(readline(), " "))
        for j in 1:1:6
            B[i,j] = A[j]
        end
    end
    max_val = -100

    for i in 2:5
        for j in 3:5
            max_val = max(max_val, B[i,j] + B[i-1,j] + B[i-1,j-1] + B[i-1,j+1] + B[i+1, j] + B[i+1, j-1] + B[i+1,j+1])
        end
    end
    println(max_val)

end

maximum_hourglass()

Hackerrank/30 Days of Code/Julia/day 11.jl


function hc4(X::IntervalBox{N,T}, constraints::Vector{Constraint{T}}, tol=1e-5) where{N, T}
    n = length(constraints)
    while true
        X_temp = X
        for i in 1:n
            X = invokelatest(constraints[i].C,  constraints[i].bound, X)
        end
        if isempty(X) || sum(dist.(X, X_temp)) < tol
            break
        end
    end

    #new_constraints = Constraint{T}[]

    #for i in 1:n
    #    if !(invokelatest(constraints[i].C, X) ⊆ constraints[i].bound)
    #        push!(new_constraints, constraints[i])
    #    end
    #end
    #return new_constraints, X
    return constraints, X
end

function contraction(f::Function, C, global_min::Float64, X::IntervalBox{N,T}, constraints::Vector{Constraint{T}}, tol=1e-5) where {N, T}
    lb = -Inf
    while true
        X_temp = X
        X = invokelatest(C, -Inf..global_min, X)
        lb = inf(f(X))
        constraints, X = hc4(X, constraints)

        if isempty(X) || sum(dist.(X, X_temp)) < tol
            break
        end
    end

    return lb, X, constraints
end

function generate_random_feasible_point(X::IntervalBox{N, T}, constraints::Vector{Constraint{T}}) where{N, T}
    for i in 1:30
        point = [X[j].lo + (1-rand())*(X[j].hi - X[j].lo) for j in 1:length(X)]      # discover a random point in interval box X
        for j in 1:length(constraints)
            if !(invokelatest(constraints[j].C, point) ⊆ constraints[j].bound)
                break
            end
            if j == length(constraints)
                return (true, point)    # return a feasible point
            end
        end
        if i == 30
            return (false, point)    # returns a infeasible point
        end
    end

end

function convex_hull(vec::Vector{IntervalBox{N,T}}) where{N, T}

    x_big = Interval{T}[]
    num_variables = N
    num_boxes = length(vec)

    for i in 1:num_variables
        lower_bound = Inf
        upper_bound = -Inf
        for j in 1:num_boxes
            if vec[j][i].lo < lower_bound
                lower_bound = vec[j][i].lo
            end
            if upper_bound < vec[j][i].hi
                upper_bound = vec[j][i].hi
            end
        end
        append!(x_big, Interval(lower_bound, upper_bound))
    end
    return IntervalBox(x_big...)
end

function check_feasiblity(point::SVector{N,T}, constraints::Vector{Constraint{T}}) where{N, T}
    for constraint in constraints
        if !(invokelatest(constraint.C, point) ∈ constraint.bound)
            return false
        end
    end
    return true
end

function ibc_minimise(f::Function , X::IntervalBox{N,T}, constraints::Vector{Constraint{T}}; ibc_chnl = RemoteChannel(()->Channel{Tuple{IntervalBox{N,T}, T}}(0)),
               diffevol_chnl = RemoteChannel(()->Channel{Tuple{SVector{N, T}, T, Union{Nothing, IntervalBox{N, T}}}}(0)), structure = SortedVector, debug = false, tol=1e-6, ibc_ind = true) where{N, T}

    vars = [Variable(Symbol(:(x[$i])))() for i in 1:N]
    g(x...) = f(x)
    C = BasicContractor(vars, g)

    working = structure([(X, inf(f(X)))], x->x[2]) # list of boxes with corresponding lower bound, arranged according to selected structure :
    minimizers = IntervalBox{N,T}[]
    global_min = ∞  # upper bound

    info = Information(0, 0, 0)
    num_bisections = 0

    while !isempty(working)

        info.iterations= info.iterations + 1

        if !ibc_ind
            if isready(ibc_chnl)
                from_diff = take!(ibc_chnl)     # Receiving best individual from ibc_minimise
                if debug
                    println("Box recevied from DifferentialEvolution: ", from_diff)
                end
                global_min = min(from_diff[2], global_min)
                info.de_to_ibc = info.de_to_ibc + 1
            end
        end

        (X, X_min) = popfirst!(working)

        gra = gradient(f, X.v)
        (lbb, fcb, cb) = bauman_form(X, f, gra)      # ----- second order form

        if global_min > lbb
            X_min = max(X_min, lbb)
            if fcb < global_min && check_feasiblity(cb, constraints)
                global_min = fcb
                x_best = cb
            end
        else
            continue
        end

        if debug
            println("New search-space : ", X)
        end

        X_min, X, constraints = contraction(f, C, global_min, X, constraints)

        if debug
            println("Contracted search_space: ", X)
        end

        if X_min > global_min    # X_min == inf(f(X))
            continue
        end

        status, output = generate_random_feasible_point(X, constraints)   # finding a feasible point in the interval box if there present any

        if status                           # true if we got any feasible point
            feas_point = output
            m = sup(f(Interval.(feas_point)))  # find candidate for upper bound of global minimum by evaluating f on a feasible point
            if m < global_min
                global_min = m
                x_best = SVector(mid(X))
                if !ibc_ind
                    if debug
                        println("Box send to DifferentialEvolution: ", x_best )
                    end
                    if info.iterations % 200 == 0
                       x_big = convex_hull(working.data)
                        put!(diffevol_chnl, (x_best, global_min, x_big))  # sending best individual to diffevol
                    else
                        put!(diffevol_chnl, (x_best, global_min, nothing))
                    end
                    info.ibc_to_de = info.ibc_to_de + 1
                end
            end

            filter_elements!(working , (X, global_min) )   # Remove all boxes whose lower bound is greater than the current one:

            if diam(X) < tol
                push!(minimizers, X)
            else
                X1, X2 = bisect(X)
                push!( working, (X1, inf(f(X1))) )
                push!( working, (X2, inf(f(X2))) )
                num_bisections += 1
            end
        else
            if diam(X) < tol
                push!(minimizers, X)
            else
                X1, X2 = bisect(X)
                push!( working, (X1, inf(f(X1))) )
                push!( working, (X2, inf(f(X2))) )
                num_bisections += 1
            end
        end

    end

    if debug
        println("IBC search is ended")
    end

    if !ibc_ind
        if isready(diffevol_chnl)
            take!(diffevol_chnl)
            put!(diffevol_chnl,(SVector(mid(X)), Inf, nothing) )
        else
            put!(diffevol_chnl,(SVector(mid(X)), Inf, nothing) )
        end
        if debug
            println("DifferentialEvolution is also terminated")
        end
        if isready(ibc_chnl)
            take!(ibc_chnl)
        end
    end

    lower_bound = minimum(inf.(f.(minimizers)))

    return Interval(lower_bound, global_min), minimizers, info

end

function ibc_maximise(f::Function , X::IntervalBox{N,T}, constraints::Vector{Constraint{T}}; ibc_chnl = RemoteChannel(()->Channel{Tuple{IntervalBox{N,T}, T}}(0)),
               diffevol_chnl = RemoteChannel(()->Channel{Tuple{SVector{N, T}, T, Union{Nothing, IntervalBox{N, T}}}}(0)), structure = SortedVector, debug = false, tol=1e-6, ibc_ind = true) where{N, T}
    bound, minimizer, info = ibc_minimise(x -> -f(x), X, constraints, ibc_chnl = ibc_chnl, diffevol_chnl = diffevol_chnl, debug = debug, tol = tol, ibc_ind = ibc_ind)
    return -bound, minimizer, info
end


function CNDEV(U)

A = [2.50662823884, -18.61500062529, 41.39119773534, -25.44106049637]
b = [-8.4735109309, 23.083336743743, -21.06224101826, 3.13082909833]
c = [0.337475482272615, 0.976169019091719, 0.160797971491821, 0.0276438810333863, 0.0038405729373609, 0.0003951896511919, 3.21767881767818e-5, 2.888167364e-7, 3.960315187e-7]

X = U - 0.5

if abs(X) < 0.92

     r = X^2
     r = X * (((A[4] * r + A[3]) * r * A[2]) * r + A[1]) / ((((b[4] * r + b[3]) * r + b[2]) * r + b[1]) * r + 1)
     cndev = r
     return cndev
end

     r = U

if X >= 0
     r = 1 - U
end

     r = log(-log(r))
     r = c[1] + r * (c[2] + r * (c[3] + r * (c[4] + r + (c[5] + r * (c[6] + r * (c[7] + r * (c[8] + r * c[9])))))))

if X < 0
     r = -r
end

     cndev = r

end


# Adaptation of some examples in MosekTools
# Refer to:
# https://github.com/JuliaOpt/MosekTools.jl/blob/master/test/jump_soc.jl
#
# MIT License
# Copyright (c) 2017, Ulf Worsøe, Mosek Aps

using JuMP

@testset "Second-order Cone Programming" begin
    solver = KNITRO.Optimizer

    @testset "SOC1" begin
        m = Model(optimizer_with_attributes(solver, "outlev" => 0))
        @variable(m, x)
        @variable(m, y)
        @variable(m, t >= 0.0)
        @objective(m, Min, t)
        @constraint(m, x + y >= 1.0)
        @constraint(m, [t,x,y] in MOI.SecondOrderCone(3))

        JuMP.optimize!(m)

        @test JuMP.has_values(m)
        @test JuMP.termination_status(m) == MOI.LOCALLY_SOLVED
        @test JuMP.primal_status(m) == MOI.FEASIBLE_POINT

        @test JuMP.objective_value(m) ≈ sqrt(1/2) atol=1e-6

        @test JuMP.value.([x, y, t]) ≈ [0.5, 0.5, sqrt(1/2)] atol=1e-3

    end

    @testset "RotatedSOC1" begin
        m = Model(optimizer_with_attributes(solver, "outlev" => 0))
        @variable(m, x[1:5] >= 0.0)
        @variable(m, 0.0 <= u <= 5.0)
        @variable(m, v)
        @variable(m, t1 == 1.0)
        @variable(m, t2 == 1.0)

        @objective(m, Max, v)

        @constraint(m, [t1/sqrt(2),t2/sqrt(2),x...] in MOI.RotatedSecondOrderCone(7))
        @constraint(m, [x[1]/sqrt(2), u/sqrt(2), v] in MOI.RotatedSecondOrderCone(3))

        JuMP.optimize!(m)

        @test JuMP.has_values(m)
        @test JuMP.termination_status(m) == MOI.LOCALLY_SOLVED
        @test JuMP.primal_status(m) == MOI.FEASIBLE_POINT

        @test JuMP.value.(x) ≈ [1,0,0,0,0] atol=1e-2
        @test JuMP.value(u) ≈ 5.0 atol=1e-4
        @test JuMP.value(v) ≈ sqrt(5.0) atol=1e-4
    end
end


"""
    GAE(enc[, σ])

Graph autoencoder.

# Arguments
- `enc`: encoder. It can be any graph convolutional layer.

Encoder is specified by user and decoder will be `InnerProductDecoder` layer.
"""
struct GAE{T,S}
    encoder::T
    decoder::S
end

GAE(enc, σ::Function=identity) = GAE(enc, InnerProductDecoder(σ))

@functor GAE

function (g::GAE)(X::AbstractMatrix)
    Z = g.encoder(X)
    A = g.decoder(Z)
    A
end

"""
    VGAE(enc[, σ])

Variational graph autoencoder.

# Arguments
- `enc`: encoder. It can be any graph convolutional layer.

Encoder is specified by user and decoder will be `InnerProductDecoder` layer.
"""
struct VGAE{T,S}
    encoder::T
    decoder::S
end

function VGAE(enc, h_dim::Integer, z_dim::Integer, σ::Function=identity)
    VGAE(VariationalEncoder(enc, h_dim, z_dim), InnerProductDecoder(σ))
end

@functor VGAE

function (g::VGAE)(fg::FeaturedGraph)
    fg_ = g.encoder(fg)
    fg_ = g.decoder(fg_)
    fg_
end

"""
    InnerProductDecoder(σ)

Inner-product decoder layer.

# Arguments
- `σ`: activation function.
"""
struct InnerProductDecoder
    σ
end

@functor InnerProductDecoder

(i::InnerProductDecoder)(Z::AbstractMatrix)::AbstractMatrix = i.σ.(Z'*Z)

function (i::InnerProductDecoder)(fg::FeaturedGraph)::FeaturedGraph
    Z = node_feature(fg)
    A = i(Z)
    FeaturedGraph(graph(fg), nf=A)
end

"""
    VariationalEncoder(nn, h_dim, z_dim)

Variational encoder layer.

# Arguments
- `nn`: neural network. It can be any graph convolutional layer.
- `h_dim`: dimension of hidden layer. This should fit the output dimension of `nn`.
- `z_dim`: dimension of latent variable layer. This will be parametrized into `μ` and `logσ`.

Encoder can be any graph convolutional layer.
"""
struct VariationalEncoder
    nn
    μ
    logσ
    z_dim::Integer
end

function VariationalEncoder(nn, h_dim::Integer, z_dim::Integer)
    VariationalEncoder(nn,
                       GCNConv(h_dim=>z_dim),
                       GCNConv(h_dim=>z_dim),
                       z_dim)
end

@functor VariationalEncoder

function (ve::VariationalEncoder)(fg::FeaturedGraph)::FeaturedGraph
    μ, logσ = summarize(ve, fg)
    Z = sample(μ, logσ)
    FeaturedGraph(graph(fg), nf=Z)
end

function summarize(ve::VariationalEncoder, fg::FeaturedGraph)
    fg_ = ve.nn(fg)
    fg_μ, fg_logσ = ve.μ(fg_), ve.logσ(fg_)
    node_feature(fg_μ), node_feature(fg_logσ)
end

sample(μ::AbstractArray{T}, logσ::AbstractArray{T}) where {T<:Real} =
    μ + exp.(logσ) .* randn(T, size(logσ))


using Flux
println("NNPDE_tests")
using DiffEqFlux
println("Starting Soon!")
using ModelingToolkit
using DiffEqBase
using Test, NeuralPDE
println("Starting Soon!")
using GalacticOptim
using Optim
using Quadrature,Cubature, Cuba
using QuasiMonteCarlo
using SciMLBase
using OrdinaryDiffEq
import ModelingToolkit: Interval, infimum, supremum

using Random
Random.seed!(100)

cb = function (p,l)
    println("Current loss is: $l")
    return false
end

## Example 1, 1D ode
function test_ode(strategy_)
    println("Example 1, 1D ode: strategy: $(nameof(typeof(strategy_)))")
    @parameters θ
    @variables u(..)
    Dθ = Differential(θ)

    # 1D ODE
    eq = Dθ(u(θ)) ~ θ^3 + 2*θ + (θ^2)*((1+3*(θ^2))/(1+θ+(θ^3))) - u(θ)*(θ + ((1+3*(θ^2))/(1+θ+θ^3)))

    # Initial and boundary conditions
    bcs = [u(0.) ~ 1.0]

    # Space and time domains
    domains = [θ ∈ Interval(0.0,1.0)]

    # Neural network
    chain = FastChain(FastDense(1,12,Flux.σ),FastDense(12,1))
    initθ = Float64.(DiffEqFlux.initial_params(chain))

    discretization = NeuralPDE.PhysicsInformedNN(chain,
                                                 strategy_;
                                                 init_params = nothing,
                                                 phi = nothing,
                                                 derivative = nothing,
                                                 )

    pde_system = PDESystem(eq,bcs,domains,[θ],[u])
    prob = NeuralPDE.discretize(pde_system,discretization)
    sym_prob = NeuralPDE.symbolic_discretize(pde_system,discretization)

    res = GalacticOptim.solve(prob, ADAM(0.1); cb = cb, maxiters=1000)
    prob = remake(prob,u0=res.minimizer)
    res = GalacticOptim.solve(prob, ADAM(0.01); cb = cb, maxiters=1000)
    prob = remake(prob,u0=res.minimizer)
    res = GalacticOptim.solve(prob, ADAM(0.001); cb = cb, maxiters=1000)
    phi = discretization.phi

    analytic_sol_func(t) = exp(-(t^2)/2)/(1+t+t^3) + t^2
    ts = [infimum(d.domain):0.01:supremum(d.domain) for d in domains][1]
    u_real  = [analytic_sol_func(t) for t in ts]
    u_predict  = [first(phi(t,res.minimizer)) for t in ts]

    @test u_predict ≈ u_real atol = 0.1
    # using Plots
    # t_plot = collect(ts)
    # plot(t_plot ,u_real)
    # plot!(t_plot ,u_predict)
end

grid_strategy = NeuralPDE.GridTraining(0.1)
quadrature_strategy = NeuralPDE.QuadratureTraining(quadrature_alg=CubatureJLh(),
                                                    reltol=1e-3,abstol=1e-3,
                                                    maxiters =50, batch=100)
stochastic_strategy = NeuralPDE.StochasticTraining(400; bcs_points= 50) #points
quasirandom_strategy = NeuralPDE.QuasiRandomTraining(400; #points
                                                     sampling_alg = LatinHypercubeSample(),
                                                     resampling =false,
                                                     minibatch = 100
                                                    )
quasirandom_strategy_resampling = NeuralPDE.QuasiRandomTraining(400; #points
                                                     bcs_points= 50,
                                                     sampling_alg = LatinHypercubeSample(),
                                                     resampling = true,
                                                     minibatch = 0)

strategies = [grid_strategy,stochastic_strategy, quadrature_strategy,quasirandom_strategy,quasirandom_strategy_resampling]

for strategy_ in strategies
    test_ode(strategy_)
end

## Example 2, 2D Poisson equation
function test_2d_poisson_equation(chain_, strategy_)
    println("Example 2, 2D Poisson equation, chain: $(nameof(typeof(chain_))), strategy: $(nameof(typeof(strategy_)))")
    @parameters x y
    @variables u(..)
    Dxx = Differential(x)^2
    Dyy = Differential(y)^2

    # 2D PDE
    eq  = Dxx(u(x,y)) + Dyy(u(x,y)) ~ -sin(pi*x)*sin(pi*y)

    # Initial and boundary conditions
    bcs = [u(0,y) ~ 0.f0, u(1,y) ~ -sin(pi*1)*sin(pi*y),
           u(x,0) ~ 0.f0, u(x,1) ~ -sin(pi*x)*sin(pi*1)]
    # Space and time domains
    domains = [x ∈ Interval(0.0,1.0),
               y ∈ Interval(0.0,1.0)]

    # chain_ = FastChain(FastDense(2,12,Flux.σ),FastDense(12,12,Flux.σ),FastDense(12,1))
    discretization = NeuralPDE.PhysicsInformedNN(chain_,
                                                 strategy_)

    pde_system = PDESystem(eq,bcs,domains,[x,y],[u])
    prob = NeuralPDE.discretize(pde_system,discretization)
    sym_prob = NeuralPDE.symbolic_discretize(pde_system,discretization)
    res = GalacticOptim.solve(prob, ADAM(0.1); cb = cb, maxiters=2000)
    phi = discretization.phi

    dx = 0.1
    xs,ys = [infimum(d.domain):dx/10:supremum(d.domain) for d in domains]
    analytic_sol_func(x,y) = (sin(pi*x)*sin(pi*y))/(2pi^2)

    u_predict = reshape([first(phi([x,y],res.minimizer)) for x in xs for y in ys],(length(xs),length(ys)))
    u_real = reshape([analytic_sol_func(x,y) for x in xs for y in ys], (length(xs),length(ys)))
    diff_u = abs.(u_predict .- u_real)

    @test u_predict ≈ u_real atol = 2.0

    # p1 = plot(xs, ys, u_real, linetype=:contourf,title = "analytic");
    # p2 = plot(xs, ys, u_predict, linetype=:contourf,title = "predict");
    # p3 = plot(xs, ys, diff_u,linetype=:contourf,title = "error");
    # plot(p1,p2,p3)
end

fastchain = FastChain(FastDense(2,12,Flux.σ),FastDense(12,12,Flux.σ),FastDense(12,1))
fluxchain = Chain(Dense(2,12,Flux.σ),Dense(12,12,Flux.σ),Dense(12,1))
chains = [fluxchain, fastchain]
for chain in chains
    test_2d_poisson_equation(chain, grid_strategy)
end

for strategy_ in strategies
    chain_ = FastChain(FastDense(2,12,Flux.σ),FastDense(12,12,Flux.σ),FastDense(12,1))
    test_2d_poisson_equation(chain_, strategy_)
end

algs = [CubatureJLp()] #CubatureJLh(),
for alg in algs
    chain_ = FastChain(FastDense(2,12,Flux.σ),FastDense(12,12,Flux.σ),FastDense(12,1))
    strategy_ =  NeuralPDE.QuadratureTraining(quadrature_alg = alg,reltol=1e-4,abstol=1e-3,maxiters=30, batch=10)
    test_2d_poisson_equation(chain_, strategy_)
end

## Example 3, high-order ode
println("Example 3, high-order ode")
@parameters x
@variables u(..) ,Dxu(..) ,Dxxu(..)
Dxxx = Differential(x)^3
Dx = Differential(x)

# ODE
eq = Dx(Dxxu(x)) ~ cos(pi*x)

# Initial and boundary conditions
bcs_ = [u(0.) ~ 0.0,
        u(1.) ~ cos(pi),
        Dxu(1.) ~ 1.0]

der = [Dx(u(x)) ~ Dxu(x),
       Dx(Dxu(x)) ~ Dxxu(x)]

bcs = [bcs_;der]
# Space and time domains
domains = [x ∈ Interval(0.0,1.0)]

# Neural network
chain = [FastChain(FastDense(1,8,Flux.σ),FastDense(8,1)) for _ in 1:3]
quasirandom_strategy = NeuralPDE.QuasiRandomTraining(100; #points
                                                     sampling_alg = LatinHypercubeSample())
initθ = map(c -> Float64.(c), DiffEqFlux.initial_params.(chain))

discretization = NeuralPDE.PhysicsInformedNN(chain,quasirandom_strategy;init_params = initθ)
pde_system = PDESystem(eq,bcs,domains,[x],[u,Dxu,Dxxu])
prob = NeuralPDE.discretize(pde_system,discretization)
sym_prob = NeuralPDE.symbolic_discretize(pde_system,discretization)

pde_inner_loss_functions = prob.f.f.loss_function.pde_loss_function.pde_loss_functions.contents
bcs_inner_loss_functions = prob.f.f.loss_function.bcs_loss_function.bc_loss_functions.contents

cb_ = function (p,l)
    println("loss: ", l )
    println("pde_losses: ", map(l_ -> l_(p), pde_inner_loss_functions))
    println("bcs_losses: ", map(l_ -> l_(p), bcs_inner_loss_functions))
    return false
end

res = GalacticOptim.solve(prob, ADAM(0.1); cb = cb_, maxiters=1000)
prob = remake(prob,u0=res.minimizer)
res = GalacticOptim.solve(prob, ADAM(0.01); cb = cb_, maxiters=1000)
prob = remake(prob,u0=res.minimizer)
res = GalacticOptim.solve(prob, ADAM(0.001); cb = cb_, maxiters=1000)
phi = discretization.phi[1]

analytic_sol_func(x) = (π*x*(-x+(π^2)*(2*x-3)+1)-sin(π*x))/(π^3)

xs = [infimum(d.domain):0.01:supremum(d.domain) for d in domains][1]
u_real  = [analytic_sol_func(x) for x in xs]
u_predict  = [first(phi(x,res.minimizer)) for x in xs]

@test u_predict ≈ u_real atol = 0.1

# x_plot = collect(xs)
# plot(x_plot ,u_real)
# plot!(x_plot ,u_predict)

## Example 4, system of pde
println("Example 4, system of pde")
@parameters x, y
@variables u1(..), u2(..)
Dx = Differential(x)
Dy = Differential(y)

# System of pde
eqs = [Dx(u1(x,y)) + 4*Dy(u2(x,y)) ~ 0,
      Dx(u2(x,y)) + 9*Dy(u1(x,y)) ~ 0]
      # 3*u1(x,0) ~ 2*u2(x,0)]

# Initial and boundary conditions
bcs = [u1(x,0) ~ 2*x, u2(x,0) ~ 3*x]

# Space and time domains
domains = [x ∈ Interval(0.0,1.0), y ∈ Interval(0.0,1.0)]

# Neural network
chain1 = FastChain(FastDense(2,15,Flux.σ),FastDense(15,1))
chain2 = FastChain(FastDense(2,15,Flux.σ),FastDense(15,1))

strategy = NeuralPDE.QuadratureTraining()
chain = [chain1,chain2]
initθ = map(c -> Float64.(c), DiffEqFlux.initial_params.(chain))
discretization = NeuralPDE.PhysicsInformedNN(chain,strategy; init_params = initθ)
pde_system = PDESystem(eqs,bcs,domains,[x,y],[u1,u2])
prob = NeuralPDE.discretize(pde_system,discretization)
sym_prob = NeuralPDE.symbolic_discretize(pde_system,discretization)

res = GalacticOptim.solve(prob,BFGS(); cb = cb, maxiters=1000)
phi = discretization.phi

analytic_sol_func(x,y) =[1/3*(6x - y), 1/2*(6x - y)]
xs,ys = [infimum(d.domain):0.01:supremum(d.domain) for d in domains]
u_real  = [[analytic_sol_func(x,y)[i] for x in xs  for y in ys] for i in 1:2]

initθ = discretization.init_params
acum =  [0;accumulate(+, length.(initθ))]
sep = [acum[i]+1 : acum[i+1] for i in 1:length(acum)-1]
minimizers = [res.minimizer[s] for s in sep]
u_predict  = [[phi[i]([x,y],minimizers[i])[1] for x in xs  for y in ys] for i in 1:2]

@test u_predict[1] ≈ u_real[1] atol = 0.1
@test u_predict[2] ≈ u_real[2] atol = 0.1

# p1 =plot(xs, ys, u_predict, st=:surface);
# p2 = plot(xs, ys, u_real, st=:surface);
# plot(p1,p2)

## Example 5, 2d wave equation, neumann boundary condition
println("Example 5, 2d wave equation, neumann boundary condition")
#here we use low level api for build solution
@parameters x, t
@variables u(..)
Dxx = Differential(x)^2
Dtt = Differential(t)^2
Dt = Differential(t)

#2D PDE
C=1
eq  = Dtt(u(x,t)) ~ C^2*Dxx(u(x,t))

# Initial and boundary conditions
bcs = [u(0,t) ~ 0.,# for all t > 0
       u(1,t) ~ 0.,# for all t > 0
       u(x,0) ~ x*(1. - x), #for all 0 < x < 1
       Dt(u(x,0)) ~ 0. ] #for all  0 < x < 1]

# Space and time domains
domains = [x ∈ Interval(0.0,1.0),
           t ∈ Interval(0.0,1.0)]

# Neural network
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()

indvars = [x,t]
depvars = [u]
dim = length(domains)
quadrature_strategy = NeuralPDE.QuadratureTraining()

_pde_loss_function = NeuralPDE.build_loss_function(eq,indvars,depvars,phi,derivative,
                                                   chain,initθ,quadrature_strategy)
_pde_loss_function(rand(2,10), initθ)

bc_indvars = NeuralPDE.get_argument(bcs,indvars,depvars)
_bc_loss_functions = [NeuralPDE.build_loss_function(bc,indvars,depvars, phi, derivative,
                                                    chain,initθ,quadrature_strategy,
                                                    bc_indvars = bc_indvar) for (bc,bc_indvar) in zip(bcs,bc_indvars)]
map(loss_f -> loss_f(rand(1,10), initθ),_bc_loss_functions)

dx = 0.1
train_sets = NeuralPDE.generate_training_sets(domains,dx,[eq],bcs,eltypeθ,indvars,depvars)
pde_train_set,bcs_train_set = train_sets
pde_bounds, bcs_bounds = NeuralPDE.get_bounds(domains,[eq],bcs,eltypeθ,indvars,depvars,quadrature_strategy)

lbs,ubs = pde_bounds
pde_loss_functions = [NeuralPDE.get_loss_function(_pde_loss_function,
                                                lbs[1],ubs[1],
                                                eltypeθ, parameterless_type_θ,
                                                quadrature_strategy)]

pde_loss_functions[1](initθ)

lbs,ubs = bcs_bounds
bc_loss_functions = [NeuralPDE.get_loss_function(_loss,lb,ub,
                                                 eltypeθ, parameterless_type_θ,
                                                 quadrature_strategy)
                                                 for (_loss,lb,ub) in zip(_bc_loss_functions, lbs,ubs)]

map(l->l(initθ) ,bc_loss_functions)

loss_functions =  [pde_loss_functions;bc_loss_functions]

function loss_function(θ,p)
    sum(map(l->l(θ) ,loss_functions))
end

f_ = OptimizationFunction(loss_function, GalacticOptim.AutoZygote())
prob = GalacticOptim.OptimizationProblem(f_, initθ)

cb_ = function (p,l)
    println("loss: ", l )
    println("losses: ", map(l -> l(p), loss_functions))
    return false
end

res = GalacticOptim.solve(prob,Optim.BFGS(); cb = cb_, maxiters=1000)

xs,ts = [infimum(d.domain):dx:supremum(d.domain) for d in domains]
analytic_sol_func(x,t) =  sum([(8/(k^3*pi^3)) * sin(k*pi*x)*cos(C*k*pi*t) for k in 1:2:50000])

u_predict = reshape([first(phi([x,t],res.minimizer)) for x in xs for t in ts],(length(xs),length(ts)))
u_real = reshape([analytic_sol_func(x,t) for x in xs for t in ts], (length(xs),length(ts)))

@test u_predict ≈ u_real atol = 0.1

# diff_u = abs.(u_predict .- u_real)
# p1 = plot(xs, ts, u_real, linetype=:contourf,title = "analytic");
# p2 =plot(xs, ts, u_predict, linetype=:contourf,title = "predict");
# p3 = plot(xs, ts, diff_u,linetype=:contourf,title = "error");
# plot(p1,p2,p3)

## Example 6, pde with mixed derivative
println("Example 6, pde with mixed derivative")
@parameters x y
@variables u(..)
Dxx = Differential(x)^2
Dyy = Differential(y)^2
Dx = Differential(x)
Dy = Differential(y)

eq = Dxx(u(x,y)) + Dx(Dy(u(x,y))) - 2*Dyy(u(x,y)) ~  -1.

# Initial and boundary conditions
bcs = [u(x,0) ~ x,
       Dy(u(x,0)) ~ x,
       u(x,0) ~ Dy(u(x,0))]

# Space and time domains
domains = [x ∈ Interval(0.0,1.0), y ∈ Interval(0.0,1.0)]

strategy = NeuralPDE.QuadratureTraining()
# Neural network
chain = FastChain(FastDense(2,12,Flux.tanh),FastDense(12,12,Flux.tanh),FastDense(12,1))
initθ = Float64.(DiffEqFlux.initial_params(chain))
discretization = NeuralPDE.PhysicsInformedNN(chain, strategy; init_params = initθ)
pde_system = PDESystem(eq,bcs,domains,[x,y],[u])
prob = NeuralPDE.discretize(pde_system,discretization)

res = GalacticOptim.solve(prob,BFGS(); cb = cb, maxiters=1000)
phi = discretization.phi

analytic_sol_func(x,y) = x + x*y +y^2/2
xs,ys = [infimum(d.domain):0.01:supremum(d.domain) for d in domains]

u_predict = reshape([first(phi([x,y],res.minimizer)) for x in xs for y in ys],(length(xs),length(ys)))
u_real = reshape([analytic_sol_func(x,y) for x in xs for y in ys], (length(xs),length(ys)))
diff_u = abs.(u_predict .- u_real)

@test u_predict ≈ u_real atol = 1.0

# p1 = plot(xs, ys, u_real, linetype=:contourf,title = "analytic");
# p2 = plot(xs, ys, u_predict, linetype=:contourf,title = "predict");
# p3 = plot(xs, ys, diff_u,linetype=:contourf,title = "error");
# plot(p1,p2,p3)

## Example 7, Fokker-Planck equation
println("Example 7, Fokker-Planck equation")
# the example took from this article https://arxiv.org/abs/1910.10503
@parameters x
@variables p(..)
Dx = Differential(x)
Dxx = Differential(x)^2

#2D PDE
α = 0.3
β = 0.5
_σ = 0.5
# Discretization
dx = 0.05
# here we use normalization condition: dx*p(x) ~ 1, in order to get non-zero solution.
#(α - 3*β*x^2)*p(x) + (α*x - β*x^3)*Dx(p(x)) ~ (_σ^2/2)*Dxx(p(x))
eq  = [Dx((α*x - β*x^3)*p(x)) ~ (_σ^2/2)*Dxx(p(x)) + dx*p(x) - 1.]

# Initial and boundary conditions
bcs = [p(-2.2) ~ 0. ,p(2.2) ~ 0., p(-2.2) ~ p(2.2)]

# Space and time domains
domains = [x ∈ Interval(-2.2,2.2)]

# Neural network
inn = 18
chain = FastChain(FastDense(1,inn,Flux.σ),
                  FastDense(inn,inn,Flux.σ),
                  FastDense(inn,1))
initθ = Float64.(DiffEqFlux.initial_params(chain))

discretization = NeuralPDE.PhysicsInformedNN(chain,
                                             NeuralPDE.GridTraining(dx);
                                             init_params = initθ)

pde_system = PDESystem(eq,bcs,domains,[x],[p])
prob = NeuralPDE.discretize(pde_system,discretization)

pde_inner_loss_functions = prob.f.f.loss_function.pde_loss_function.pde_loss_functions.contents
bcs_inner_loss_functions = prob.f.f.loss_function.bcs_loss_function.bc_loss_functions.contents

cb_ = function (p,l)
    println("loss: ", l )
    println("pde_losses: ", map(l_ -> l_(p), pde_inner_loss_functions))
    println("bcs_losses: ", map(l_ -> l_(p), bcs_inner_loss_functions))
    return false
end

res = GalacticOptim.solve(prob,Optim.BFGS(),cb = cb_,maxiters=2000)

phi = discretization.phi
analytic_sol_func(x) = 28*exp((1/(2*_σ^2))*(2*α*x^2 - β*x^4))
xs = [infimum(d.domain):dx:supremum(d.domain) for d in domains][1]
u_real  = [analytic_sol_func(x) for x in xs]
u_predict  = [first(phi(x,res.minimizer)) for x in xs]

@test u_predict ≈ u_real atol = 20.0

# plot(xs ,u_real, label = "analytic")
# plot!(xs ,u_predict, label = "predict")

## Example 8, Lorenz System (Parameter Estimation)
println("Example 8, Lorenz System")

Random.seed!(1234)
@parameters t ,σ_ ,β, ρ
@variables x(..), y(..), z(..)
Dt = Differential(t)
eqs = [Dt(x(t)) ~ σ_*(y(t) - x(t)),
       Dt(y(t)) ~ x(t)*(ρ - z(t)) - y(t),
       Dt(z(t)) ~ x(t)*y(t) - β*z(t)]

bcs = [x(0) ~ 1.0, y(0) ~ 0.0, z(0) ~ 0.0]
domains = [t ∈ Interval(0.0,1.0)]
dt = 0.05

input_ = length(domains)
n = 8
chain = [FastChain(FastDense(input_,n,Flux.σ),FastDense(n,n,Flux.σ),FastDense(n,1)) for _ in 1:3]
#Generate Data
function lorenz!(du,u,p,t)
 du[1] = 10.0*(u[2]-u[1])
 du[2] = u[1]*(28.0-u[3]) - u[2]
 du[3] = u[1]*u[2] - (8/3)*u[3]
end

u0 = [1.0;0.0;0.0]
tspan = (0.0,1.0)
prob = ODEProblem(lorenz!,u0,tspan)
sol = solve(prob, Tsit5(), dt=0.1)
ts = [infimum(d.domain):dt/5:supremum(d.domain) for d in domains][1]

function getData(sol)
    data = []
    us = hcat(sol(ts).u...)
    ts_ = hcat(sol(ts).t...)
    return [us,ts_]
end

data = getData(sol)

#Additional Loss Function
initθs = map(c -> Float64.(c), DiffEqFlux.initial_params.(chain))
acum =  [0;accumulate(+, length.(initθs))]
sep = [acum[i]+1 : acum[i+1] for i in 1:length(acum)-1]
(u_ , t_) = data
len = length(data[2])

function additional_loss(phi, θ , p)
    return sum(sum(abs2, phi[i](t_ , θ[sep[i]]) .- u_[[i], :])/len for i in 1:1:3)
end

discretization = NeuralPDE.PhysicsInformedNN(chain,
                                             NeuralPDE.GridTraining(dt);
                                             init_params =initθs,
                                             param_estim=true,
                                             additional_loss=additional_loss)
pde_system = PDESystem(eqs,bcs,domains,
                      [t],[x, y, z],[σ_, ρ, β],
                      defaults=Dict([p .=> 1.0 for p in [σ_, ρ, β]]))
prob = NeuralPDE.discretize(pde_system,discretization)
sym_prob = NeuralPDE.symbolic_discretize(pde_system,discretization)

res = GalacticOptim.solve(prob, Optim.BFGS(); cb = cb, maxiters=6000)
p_ = res.minimizer[end-2:end]
@test sum(abs2, p_[1] - 10.00) < 0.1
@test sum(abs2, p_[2] - 28.00) < 0.1
@test sum(abs2, p_[3] - (8/3)) < 0.1
#Plotting the system
# initθ = discretization.init_params
# acum =  [0;accumulate(+, length.(initθ))]
# sep = [acum[i]+1 : acum[i+1] for i in 1:length(acum)-1]
# minimizers = [res.minimizer[s] for s in sep]
# ts = [infimum(d.domain):dt/10:supremum(d.domain) for d in domains][1]
# u_predict  = [[discretization.phi[i]([t],minimizers[i])[1] for t in ts] for i in 1:3]
# plot(sol)
# plot!(ts, u_predict, label = ["x(t)" "y(t)" "z(t)"])


using Arithmetic
using Test

# 4
@testset "Evaluate constant" begin
  @test evaluate(Constant(4)) == 4
end

# 1 * 3
@testset "Simple multiplication" begin
  @test evaluate(BinaryOp(Constant(1), MUL, Constant(3))) == 3
end

# 2 + 4 / 4
@testset "Nested operations" begin
  @test evaluate(BinaryOp(Constant(2), ADD, BinaryOp(Constant(4), DIV, Constant(4)))) == 3
end


# ==============================================================================
# Max-Plus Productive (flowshops).
# Note: we use the Julia SimpleWeightedGraphs package for converting the flowshop
# to a grap. We use the Julia TikzPictures package instead of GraphPlot.jl for
# displaying them.
# ==============================================================================

using SimpleWeightedGraphs, TikzPictures

# ==============================================================================
# Save a max-plus matrice to a directed graph.

mpgraph(A::ArrMP) = SimpleWeightedDiGraph(A)
mpgraph(A::SpaMP) = SimpleWeightedDiGraph(A)

# ==============================================================================
# shift

# ==============================================================================
# From the matrix of time of tasks (machines, coins) build a max+ linear system

function flowshop(E::ArrMP{T}) where T
    nmach, npiece = size(E)

    A = mpzeros(T, nmach * npiece, nmach * npiece)
    B = mpzeros(T, nmach * npiece, nmach + npiece)
    C = mpzeros(T, nmach + npiece, nmach * npiece)
    D = mpzeros(T, nmach * npiece, nmach * npiece)
    X0 = mpzeros(T, nmach * npiece, 1)

    l = zeros(Int64, 1, nmach)
    d = zeros(Int64, 1, npiece)
    ij = 0
    for i in 1:npiece
        for j in 1:nmach
            ij = i + (j-1) * npiece
            if iszero(E[j, i])
                !iszero(l[j]) && (l[j] += 1)
                !iszero(d[i]) && (d[i] += 1)
            else
                !iszero(l[j]) && (D[ij, ij - l[j]] = E[j, i - l[j]])
                !iszero(d[i]) && (D[ij, ij - d[i] * npiece] = E[j - d[i],i])
                iszero(d[i])  && (B[ij, i] = mp1);  d[i] = 1
                iszero(l[j])  && (B[ij, j + npiece] = mp1); l[j] = 1
            end
        end
        C[i, ij - (d[i] - 1) * npiece] = E[nmach - d[i] + 1, i]
    end
    for j in 1:nmach
        C[j + npiece, j * npiece - l[j] + 1] = E[j, npiece - l[j] + 1]
    end

    mpsyslin(A, B, C, D, X0)
end

# ==============================================================================
#

#function flowshop_graph(E::arrMP, m::arrMP, p::arrMP)
#    (g, T, N)
#end

# ==============================================================================
#

#function flowshop_simu(E::arrMP, m::arrMP, p::arrMP, u::arrMP)
#    mpsimul()
#end

# ==============================================================================
#

#function critical_graph

# ==============================================================================

# SimpleWeightedDiGraph
function LaTeXG(PT::ArrMP, O::Float64)
    # Define colors and shapes for arrows.
    settings="\\tikzset{
              darrow/.style={->,draw=black,line width=1pt},
              barrow/.style={->,draw=blue,line width=1pt},
              yarrow/.style={->,draw=yellow,line width=1pt}}"

    # Define colors for places
    opts="entree/.style={circle,draw=blue,line width=1pt},
          sortie/.style={circle,draw=green,fill=green,line width=1pt},
          vide/.style={circle,draw=black,line width=1pt},
          pleine/.style={circle,draw=green,line width=1pt}"

    nm, np = size(PT)

    code = "%% Piece inputs -- Piece outputs\n"
    for i in 1:np
        code *= "\\node[entree] (pi$(i)) at ($(i),0) {}; \\node[sortie] (po$(i)) at ($(i),$(nm+1+O)) {};\n"
    end

    code *= "%% Machine inputs -- Machine outputs\n"
    for i in 1:nm
        code *= "\\node[entree] (mi$(i)) at (0,$(nm-i+1)) {}; \\node[sortie] (mo$(i)) at ($(np+1+O),$(nm-i+1)) {};\n"
    end

    code *= "%% Places\n"
    for j in 1:np
        code *= "% Row $(j)\n"
        for i in 1:nm
            if !iszero(PT[i,j])
                code *= "\\node[vide] (p$(i)$(j)) at ($(j+O),$(i+O)) {$(i)$(j)};\n"
            end
        end
        for i in 1:nm
            code *= "\\draw[yarrow] (mo$(i)) -- (mi$(i));\n"
        end
        code *= "\\draw[yarrow] (po$(j)) -- (pi$(j));\n"
    end

    code *= "%% Transitions\n"

    TikzPicture(code, options=opts, preamble=settings)
end


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

using ValueShapes
using Test

using LinearAlgebra, Random, Statistics
using StatsBase, Distributions, ArraysOfArrays

@testset "reshaped_dist" begin
    scshape = ScalarShape{Real}()

    a1shape = ArrayShape{Real}(3)
    a2shape = ArrayShape{Real}(2, 3)

    ntshape = NamedTupleShape(
        a = ArrayShape{Real}(2, 3),
        b = ScalarShape{Real}(),
        c = ConstValueShape(4.2),
        x = ConstValueShape([11 21; 12 22]),
        y = ArrayShape{Real}(4)
    )

    names = (:a, :b, :c, :x, :y)

    _mvndist(n::Integer) = MvNormal(fill(2.0, n), Diagonal(fill(3.0, n)))
    _mvnormalrd(shape::AbstractValueShape) = ReshapedDist(_mvndist(totalndof(shape)), shape)
    _mvnormalrd2(shape::AbstractValueShape) = shape(_mvndist(totalndof(shape)))

    @testset "ctors" begin
        @test @inferred(_mvnormalrd(scshape)) isa Distribution{Univariate, Continuous}
        @test @inferred(_mvnormalrd(a1shape)) isa Distribution{Multivariate, Continuous}
        @test @inferred(_mvnormalrd(a2shape)) isa Distribution{Matrixvariate, Continuous}
        @test @inferred(_mvnormalrd(ntshape)) isa Distribution{ValueShapes.NamedTupleVariate{names}, Continuous}

        @test_throws ArgumentError ReshapedDist(MvNormal(Diagonal(fill(3.0, 2))), a2shape)

        @test @inferred(_mvnormalrd2(scshape)) isa ReshapedDist{Univariate, Continuous}
        @test @inferred(_mvnormalrd2(a1shape)) isa MvNormal
        @test @inferred(_mvnormalrd2(a2shape)) isa MatrixReshaped
        @test @inferred(_mvnormalrd2(ntshape)) isa ReshapedDist{ValueShapes.NamedTupleVariate{names}, Continuous}

        @inferred(varshape(_mvnormalrd2(ntshape))) == ntshape
        @inferred(unshaped(_mvnormalrd2(ntshape))) == MvNormal(fill(2.0, totalndof(ntshape)), Diagonal(fill(3.0, totalndof(ntshape))))
    end

    @testset "rand" begin
        @test @inferred(rand(_mvnormalrd(scshape))) isa Real
        @test @inferred(rand(_mvnormalrd(scshape), ())) isa AbstractArray{<:Real,0}
        @test @inferred(rand(_mvnormalrd(scshape), 7)) isa AbstractArray{<:Real,1}
        @test @inferred(rand(_mvnormalrd(scshape), (7,))) isa AbstractArray{<:Real,1}
        @test size(rand(_mvnormalrd(scshape), 7)) == (7,)

        let d = _mvndist(totalndof(a1shape))
            @test @inferred(a1shape(d)) === d
        end
        @test @inferred(rand(_mvnormalrd(a1shape))) isa AbstractVector{<:Real}
        @test @inferred(rand(_mvnormalrd(a1shape), ())) isa AbstractArray{<:AbstractVector{<:Real},0}
        @test @inferred(rand(_mvnormalrd(a1shape), 7)) isa AbstractArray{<:Real,2}
        @test @inferred(rand(_mvnormalrd(a1shape), (7,))) isa AbstractArray{<:AbstractVector{<:Real},1}
        @test size(rand(_mvnormalrd(a1shape), 7)) == (3, 7)
        @test size(rand(_mvnormalrd(a1shape), (7,))) == (7,)

        let d = _mvndist(totalndof(a2shape))
            @test @inferred(a2shape(d)) isa MatrixReshaped
            @test unshaped(a2shape(d)) === d
        end
        @test @inferred(rand(_mvnormalrd(a2shape))) isa AbstractMatrix{<:Real}
        @test @inferred(rand(_mvnormalrd(a2shape), 7)) isa AbstractArray{<:AbstractMatrix{<:Real},1}
        @test size(rand(_mvnormalrd(a2shape), 7)) == (7,)

        @test @inferred(rand(_mvnormalrd(ntshape))) isa NamedTuple{names}
        @test @inferred(rand(_mvnormalrd(ntshape), ())) isa ShapedAsNT{<:NamedTuple{names}}
        @test @inferred(rand(_mvnormalrd(ntshape), 7)) isa ShapedAsNTArray{<:NamedTuple{names},1}
        @test size(rand(_mvnormalrd(ntshape), 7)) == (7,)
        let X = rand(_mvnormalrd(ntshape), 7)
            @test @inferred(rand!(_mvnormalrd(ntshape), view(X, 1))) == view(X, 1)
            @test @inferred(rand!(_mvnormalrd(ntshape), X)) === X
        end
    end

    @testset "stats functions" begin
        @test @inferred(mean(_mvnormalrd(scshape))) ≈ 2
        @test @inferred(mode(_mvnormalrd(scshape))) ≈ 2
        @test @inferred(var(_mvnormalrd(scshape))) ≈ 3
        let rd = _mvnormalrd(scshape), ux = rand(unshaped(rd)), vs = varshape(rd)
            @test @inferred(pdf(rd, vs(ux)[])) == pdf(unshaped(rd), ux)
            @test @inferred(pdf(rd, vs(ux))) == pdf(unshaped(rd), ux)
            @test @inferred(logpdf(rd, vs(ux)[])) == logpdf(unshaped(rd), ux)
            @test @inferred(logpdf(rd, vs(ux))) == logpdf(unshaped(rd), ux)
            @test @inferred(insupport(rd, vs(ux)[])) == true
        end

        @test @inferred(mean(_mvnormalrd(a1shape))) ≈ fill(2, 3)
        @test @inferred(mode(_mvnormalrd(a1shape))) ≈ fill(2, 3)
        @test @inferred(var(_mvnormalrd(a1shape))) ≈ fill(3, 3)
        @test @inferred(cov(_mvnormalrd(a1shape))) ≈ Diagonal(fill(3.0, 3))
        let rd = _mvnormalrd(a1shape), ux = rand(unshaped(rd)), vs = varshape(rd)
            @test @inferred(pdf(rd, vs(ux))) == pdf(unshaped(rd), ux)
            @test @inferred(logpdf(rd, vs(ux))) == logpdf(unshaped(rd), ux)
            @test @inferred(insupport(rd, vs(ux))) == true
        end

        @test @inferred(mean(_mvnormalrd(a2shape))) ≈ fill(2, 2, 3)
        @test @inferred(mode(_mvnormalrd(a2shape))) ≈ fill(2, 2, 3)
        @test @inferred(var(_mvnormalrd(a2shape))) ≈ fill(3, 2, 3)
        let rd = _mvnormalrd(a2shape), ux = rand(unshaped(rd)), vs = varshape(rd)
            @test @inferred(pdf(rd, vs(ux))) == pdf(unshaped(rd), ux)
            @test @inferred(logpdf(rd, vs(ux))) == logpdf(unshaped(rd), ux)
            @test @inferred(insupport(rd, vs(ux))) == true
        end

        @test @inferred(mean(_mvnormalrd(ntshape))) == (a = [2.0 2.0 2.0; 2.0 2.0 2.0], b = 2.0, c = 4.2, x = [11 21; 12 22], y = [2.0, 2.0, 2.0, 2.0])
        @test @inferred(mode(_mvnormalrd(ntshape))) == (a = [2.0 2.0 2.0; 2.0 2.0 2.0], b = 2.0, c = 4.2, x = [11 21; 12 22], y = [2.0, 2.0, 2.0, 2.0])
        @test @inferred(var(_mvnormalrd(ntshape))) == (a = [3.0 3.0 3.0; 3.0 3.0 3.0], b = 3.0, c = 0.0, x = [0 0; 0 0], y = [3.0, 3.0, 3.0, 3.0])
        let rd = _mvnormalrd(ntshape), ux = rand(unshaped(rd)), vs = varshape(rd)
            @test @inferred(pdf(rd, vs(ux)[])) == pdf(unshaped(rd), ux)
            @test @inferred(pdf(rd, vs(ux))) == pdf(unshaped(rd), ux)
            @test @inferred(logpdf(rd, vs(ux)[])) == logpdf(unshaped(rd), ux)
            @test @inferred(logpdf(rd, vs(ux))) == logpdf(unshaped(rd), ux)
            @test @inferred(insupport(rd, vs(ux)[])) == true
            @test @inferred(insupport(rd, vs(ux))) == true
        end

        let rd = ReshapedDist(_mvndist(5), ArrayShape{Real}(5)), X = rand(rd, 10), Xn = nestedview(X)
            @test @inferred(Distributions._pdf(rd, Xn[1])) == pdf(rd, Xn[1])
            @test @inferred(pdf(rd, X)) == pdf.(Ref(rd), Xn)

            @test @inferred(Distributions._logpdf(rd, Xn[1])) == logpdf(rd, Xn[1])
            @test @inferred(logpdf(rd, X)) == logpdf.(Ref(rd), Xn)
        end
    end
end


#############################################################
# Running Simulations
#############################################################

#Include KG_Module to run simulation
include("../src/KG_Module.jl")

p0 = Param(

    #Temporal variables
    tmax           = 1000.0     ,
    out_every_t    = 0.1        ,
    deltat         = 0.05       ,

    #Spacial variables
    xmin           =  -30.0     ,
    xmax           =  30.0      ,
    xnodes         =  128       ,

    #Spacial variables
    ymin           =  -30.0     ,
    ymax           =  30.0      ,
    ynodes         =   128      ,

    #Initial Conditions
    A0              =   3.5     ,
    σ               =   1.5     ,
    r0              =   17.0    ,
    ω               =   0.5     ,
    m               =   2       ,

    #Klein Gordon Mass
    μ               =   0.0     ,

    #Dissipation terms
    dissipation     =  :true    ,

    R_orbit         = 0.0       ,
    R1              = 1.5       ,
    R2              = 0.0       ,

    Ω               = 0.45      , #(3)^(-3/2) ,      ,
    alpha1          = 10.0      ,
    alpha2          = 0.0       ,

    #Periodic BC?
    Boundaries      = :radial   , # :square :periodic

    #Data variables
    folder          = "/test_runs/"  ,       #must have a bar before the folder
    fname           = "run_004"

)

solve_wave_equation_2D(p0)

#############################################################
# Analysis
#############################################################

#Include plot utils
include("../src/KG_Plot_Module.jl")

#name of file to analyse
#filename = "./test_runs/run_003.h5"

#Slices Observer (also works with filename)
slices_observer(p0)

#Generate animation
#generate_animation_heatmap_makie2( filename , "test_animation.mp4")


function complete_lower_bound_grad(; y=y, idx=idx, S=S, μα=μα, Σα=Σα, v²ₐ=v²ₐ, μβ=μβ, Σβ=Σβ, v²ᵦ=v²ᵦ, μf=μf, Σf=Σf, K=K)

    #
    # ∫ q(α) q(β) q(f) log𝓝(y| Af + b, S) + log𝓝(f|0,K) + log𝓝(a|0,v²ₐI) + log𝓝(b|0,v²ᵦ) dα dβ df
    #                           (1)               (2)             (3)             (4)

    gradα, gradΣα, gradβ, gradΣβ = exact_expectation_inside_log_of_lower_bound_grad(; y=y, idx=idx, S=S, μα=μα, Σα=Σα, v²ₐ=v²ₐ, μβ=μβ, Σβ=Σβ, v²ᵦ=v²ᵦ, μf=μf, Σf=Σf, K=K)

    #
    # H[q(α)]  +  H[q(β)]  +  H[q(f)]
    #

    # remember we parametrise the n-th entry of the diagonal covariance
    # Σa with elements exp(pₙ), i.e. gradient calculated below is dependent
    # on this parametrisation (see "unpack" function)
    gradΣα += 0.5 * ones(size(Σα, 1))

    # same comment as above applies here
    gradΣβ += 0.5 * ones(size(Σβ, 1))

    return gradα, gradΣα, gradβ, gradΣβ

end


"""
    FM90(;c1=0.10, c2=0.70, c3=3.23, c4=0.41, x0=4.60, gamma=0.9)
    FM90(coeffs, x0=4.60, gamma=0.9)

Fitzpatrick & Massa (1990) generative model for ultraviolet dust extinction. The default values are the published values for the Milky Way average.

## Parameters
* `c1` - y-intercept of linear term
* `c2` - slope of liner term
* `c3` - amplitude of 2175 Å bump
* `c4` - amplitude of FUV rise
* `x0` - centroid of 2175 Å bump
* `gamma` - width of 2175 Å bump

If `λ` is a `Unitful.Quantity` it will be automatically converted to Å and the
returned value will be `UnitfulAstro.mag`.

## Examples
```jldoctest
julia> model = FM90(c1=0.2, c2=0.7, c3=3.23, c4=0.41, x0=4.6, gamma=0.99);

julia> model(1500)
5.2521258452800135

julia> FM90()(1500)
5.152125845280013

julia> FM90(c1=0.2, c2=0.7, c3=3.23, c4=0.41, x0=4.6, gamma=0.99).([1000, 1200, 1800])
3-element Array{Float64,1}:
 12.562237969522851
  7.769215017329513
  4.890128210972148

```
# Extended Help

The model has form ``c_1 + c_2x + c_3D(x; \\gamma, x_0) + c_4 F(x)`` where ``x`` is the wavenumber in inverse microns, ``D(x)``
is a Drude profile (modified Lorentzian) used to model the 2175 Å bump with the scale-free parameters ``x_0`` (central wavenumber)
and ``\\gamma`` (damping coefficient), and ``F(x)``, a piecewise function for the far-UV. Note that the coefficients will change
the overall normalization, possibly changing the expected behavior of reddening via the parameter ``A_V``.

## References
[Fitzpatrick & Massa (1990)](https://ui.adsabs.harvard.edu/abs/1990ApJS...72..163F)
"""
@with_kw struct FM90{T<:Number}  <: ExtinctionLaw @deftype T
    c1 = 0.10
    c2 = 0.70
    c3 = 3.23
    c4 = 0.41
    x0 = 4.60
    gamma = 0.99
    @assert x0 ≥ 0 "`x0` must be ≥ 0, got $x0"
    @assert gamma ≥ 0 "`gamma` must be ≥ 0, got $gamma"
end

FM90(c1, c2, c3, c4, x0, gamma) = FM90(promote(c1, c2, c3, c4, x0, gamma)...)
FM90(coeffs, x0=4.60, gamma=0.99) = FM90(coeffs..., x0, gamma)

bounds(::Type{<:FM90}) = (912, 3200)

function (law::FM90)(wave::T) where T
    checkbounds(law, wave) || return zero(float(T))

    x = aa_to_invum(wave)
    exvebv = law.c1 + law.c2 * x

    x2 = x^2
    exvebv += law.c3 * (x2 / ((x2 - law.x0^2)^2 + x2 * (law.gamma^2)))

    if x >= 5.9
        y = x - 5.9
        exvebv += law.c4 * @evalpoly y 0 0 0.5392 0.05644
    end

    return exvebv
end


include("SortSparse.jl")
include("KoLesky.jl")
using LinearAlgebra
using Distributions
using Distances
using Random
using Profile
using JLD
using MAT
Random.seed!(123)

NTrain = 8000
NTest = 100
N = NTest + NTrain
d =  2
y = randn(N)
x = randn(d, N)

setSplit = randperm(N)
trainInds = setSplit[1:NTrain]
testInds = setSplit[(NTrain + 1) : end]

xTrain = x[ :, trainInds ]
yTrain = y[ trainInds ]
xTest = x[:, testInds ]
yTest = y[ testInds ]

ρ = 3.5
λ = 1.0

l = 0.1
cov = r -> exp(-r/l)

@time colptr, rowval, P, revP, distances = sortSparseRev(xTrain,ρ, 1)
xTrain = xTrain[:, P]
yTrain = yTrain[P]
x = hcat(xTrain, xTest)

skeletons = construct_skeletons(colptr, rowval, distances, λ)
colptr = 0
length(skeletons)
rowval = 0
GC.gc()

#Comparison to full kodel
LinearAlgebra.BLAS.set_num_threads(8)
@time begin
  Θ = cov.(pairwise(Euclidean(), x, dims=2))
  μFull = Θ[(NTrain + 1) : end, 1 : NTrain] * ( Θ[1 : NTrain, 1 : NTrain] \ yTrain)
end

Θ = 0
GC.gc()
@show mean(abs.(yTest - μFull))
@show mean(abs.(yTest - μFull)) / std(yTest)

@time μ, σ = predict(skeletons, (NTrain+1): N, yTrain, x, cov)

mean(abs.(yTest - μ)) / std(yTest)
mean(abs.(yTest - μ))


using LinearAlgebra
using Test

using CompScienceMeshes
using SauterSchwabQuadrature
using StaticArrays

pI = point(1,5,3)
pII = point(2,5,3)
pIII = point(7,1,0)
pIV = point(5,1,-3)

Sourcechart = simplex(pI,pIII,pII)
Testchart = simplex(pI,pIV,pII)

Accuracy = 12
ce = CommonEdge(SauterSchwabQuadrature._legendre(Accuracy,0.0,1.0))

function integrand(x,y)
			return(((x-pI)'*(y-pII))*exp(-im*1*norm(x-y))/(4pi*norm(x-y)))
end

function INTEGRAND(û,v̂)
	n1 = neighborhood(Testchart, û)
	n2 = neighborhood(Sourcechart, v̂)
	x = cartesian(n1)
	y = cartesian(n2)
	output = integrand(x,y)*jacobian(n1)*jacobian(n2)

return(output)
end

result = sauterschwab_parameterized(INTEGRAND, ce)-
		   verifintegral2(Sourcechart, Testchart, integrand, Accuracy)

@test norm(result) < 1.e-3

kernel(x,y) = 1/norm(cartesian(x)-cartesian(y))

t1 = simplex(
	@SVector[0.180878, -0.941848, -0.283207],
	@SVector[0.0, -0.92388, -0.382683],
 	@SVector[0.0, -0.980785, -0.19509])
t2 = simplex(
	@SVector[0.180878, -0.941848, -0.283207],
	@SVector[0.0, -0.92388, -0.382683],
	@SVector[0.158174, -0.881178, -0.44554])

@test indexin(t1.vertices, t2.vertices) == [1, 2, nothing]

rt = RTRefSpace{Float64}()
igd = generate_integrand_uv(kernel, rt, rt, t1, t2)

i5 = sauterschwab_parameterized(igd, CommonEdge(SauterSchwabQuadrature._legendre(5,0.0,1.0)))
i10 = sauterschwab_parameterized(igd, CommonEdge(SauterSchwabQuadrature._legendre(10,0.0,1.0)))
i15 = sauterschwab_parameterized(igd, CommonEdge(SauterSchwabQuadrature._legendre(15,0.0,1.0)))

# brute numerical approach
q1 = quadpoints(t1, 10)
q2 = quadpoints(t2, 10)

M = N = numfunctions(rt)
iref = zero(i5)
for (x,w1) in q1
    f = rt(x)
    for (y,w2) in q2
        g = rt(y)
        G = kernel(x,y)
        ds = w1*w2
        global iref += SMatrix{M,N}([dot(f[i][1], G*g[j][1])*ds for i=1:M, j=1:N])
    end
end

include(joinpath(dirname(@__FILE__,),"numquad.jl"))
ibf = numquad(kernel, rt, rt, t1, t2, zero(i5))

@test i5  ≈ iref atol=1e-3
@test i10 ≈ iref atol=1e-3
@test i10 ≈ ibf  atol=1e-3
@test i10 ≈ i15  atol=1e-5

# Test the more (or less) singular case of the second kind kernel
function kernel2nd(x,y)

	r = cartesian(x) - cartesian(y)
	R = norm(r)

	gradgreen = - r / R^3

	@SMatrix [
		0             -gradgreen[3]  gradgreen[2]
		 gradgreen[3]             0 -gradgreen[1]
		-gradgreen[2]  gradgreen[1]             0 ]
end

igd = generate_integrand_uv(kernel2nd, rt, rt, t1, t2)
i10 = sauterschwab_parameterized(igd, CommonEdge(SauterSchwabQuadrature._legendre(10,0.0,1.0)))
i15 = sauterschwab_parameterized(igd, CommonEdge(SauterSchwabQuadrature._legendre(15,0.0,1.0)))
i20 = sauterschwab_parameterized(igd, CommonEdge(SauterSchwabQuadrature._legendre(20,0.0,1.0)))

iref = numquad(kernel2nd, rt, rt, t1, t2, zero(i15))

# # Compare to BEAST:
# tqd = CompScienceMeshes.quadpoints(rt, [t1], (12,))
# bqd = CompScienceMeshes.quadpoints(rt, [t2], (13,))
#
# SE_strategy = BEAST.WiltonSEStrategy(
#   tqd[1,1],
#   BEAST.DoubleQuadStrategy(
# 	tqd[1,1],
# 	bqd[1,1]))
#
# op = BEAST.MWDoubleLayer3D(0.0)
# z2 = zeros(3,3)
# BEAST.momintegrals!(op, rt, rt, t1, t2, z2, SE_strategy)


struct Context{Ns, Nb}
    β :: Float64

    system_space :: StateSpace{Ns}
    system_state :: Vector{Float64}

    bath_space :: StateSpace{Nb}
    bath_state :: Vector{Float64}

    degeneracies :: Vector{NTuple{2, NTuple{2, Int}}}
    P :: Vector{BigFloat}
    Pid :: BigFloat
    Pcum :: Vector{BigFloat}

    function Context{Ns, Nb}(β::Float64, system_space::StateSpace{Ns}, system_state::AbstractVector,
                             bath_space::StateSpace{Nb}) where{Ns, Nb}
        bath_state = boltzmann(β, bath_space)
        deg, P, Pid = degeneracies(system_space, bath_space)
        Pcum = cumsum(P) + Pid
        new(β, system_space, system_state, bath_space, bath_state, deg, P, Pid, Pcum)
    end
end

function Context(β::Float64, system_space::StateSpace{N}, system_state::AbstractVector) where {N}
    Context{N,N}(β, system_space, system_state, system_space)
end

function timestep(rng::AbstractRNG, ctx::Context)
    p, i = BigFloat(rand(rng)), 0
    if p > ctx.Pid
        i = 1
        while i < length(ctx.Pcum) && p > ctx.Pcum[i]
            i += 1
        end
        if i <= length(ctx.P)
            ((a,b), (u,v)) = ctx.degeneracies[i]
            Pa, Pb = ctx.system_state[a], ctx.bath_state[b]
            Pu, Pv = ctx.system_state[u], ctx.bath_state[v]
            degab = ctx.system_space.deg[a] * ctx.bath_space.deg[b]
            deguv = ctx.system_space.deg[u] * ctx.bath_space.deg[v]
            dP = ((Pu*Pv)/deguv) - ((Pa*Pb)/degab)
            ctx.system_state[a] += dP
            ctx.system_state[u] -= dP
        end
    end
    ctx.system_state, i
end

timestep(ctx::Context) = timestep(Base.GLOBAL_RNG, ctx)


using XConv, LinearAlgebra, Flux, CUDA, PyPlot, BenchmarkTools
BLAS.set_num_threads(2)

nx = 128
ny = 128
batchsize=10
n_in = 2
n_out = 2
stride = 1
n_bench = 5
nw   = 3;

X = randn(Float32, nx, ny, n_in, batchsize) |> gpu
Y = randn(Float32, nx, ny, n_out, batchsize) |> gpu
g20 = grad_ev(X, Y, 5, nw, stride);
# Flux network
C = Conv((nw, nw), n_in=>n_out, identity;pad=1, stride=stride) |> gpu

batches = [2^k for k=0:8]

tf = zeros(length(batches))
t1 = similar(tf)
t10 = similar(tf)
t50 = similar(tf)
t100 = similar(tf)
angles = zeros(length(batches), 4)

close("all")

# Init plot
fig, axsl = subplots(3, 3, figsize=(10, 5))
fig.suptitle("Conv layer gradient chi-$(n_in), cho-$(n_out)")

for (i, b)=enumerate(batches)
    println("Gradient for batchsize=$b")

    local X = rand([-1f0, 1f0], nx, ny, n_in, b) |> gpu
    local Y = C(X) - gpu(randn(Float32, nx, ny, n_out, b))

    cdims = DenseConvDims(X, C.weight; stride=C.stride, padding=C.pad, dilation=C.dilation)

    # Run once to get gradient
    @time local g1 = ∇conv_filter(X, Y, cdims)
    @time local g20 = grad_ev(X, Y, 5, nw, stride);
    @time local g21 = grad_ev(X, Y, 10, nw, stride);
    @time local g22 = grad_ev(X, Y, 50, nw, stride);
    @time local g23 = grad_ev(X, Y, 100, nw, stride);

    # Compute similarities
    angles[i, 1] = dot(g20, g1)/(norm(g20)*norm(g1))
    angles[i, 2] = dot(g21, g1)/(norm(g21)*norm(g1))
    angles[i, 3] = dot(g22, g1)/(norm(g22)*norm(g1))
    angles[i, 4] = dot(g23, g1)/(norm(g23)*norm(g1))

    # Benchmark runtime
    tf[i] = @belapsed ∇conv_filter($X, $Y, $cdims) samples=n_bench
    t1[i] = @belapsed grad_ev($X, $Y, 5, $nw, $stride) samples=n_bench
    t10[i] = @belapsed grad_ev($X, $Y, 10, $nw, $stride) samples=n_bench
    t50[i] = @belapsed grad_ev($X, $Y, 50, $nw, $stride) samples=n_bench
    t100[i] = @belapsed grad_ev($X, $Y, 100, $nw, $stride) samples=n_bench

    # Plot result
    axsl[i].plot(vec(g20)/norm(g20, Inf), label="LR(s=5)", "-r")
    axsl[i].plot(vec(g21)/norm(g21, Inf), label="LR(s=10)", "-b")
    axsl[i].plot(vec(g22)/norm(g22, Inf), label="LR(s=50)", "-g")
    axsl[i].plot(vec(g23)/norm(g23, Inf), label="LR(s=100))", "-c")
    axsl[i].plot(vec(g1)/norm(g1, Inf), label="Flux", "-k")
    axsl[i].set_title("batchsize=$b")
end
lines, labels = fig.axes[end].get_legend_handles_labels()
fig.legend(lines, labels, loc = "upper left")
tight_layout()

figure()
title("10 Gradient computation speedup")
plot(batches, tf, label="Flux")
plot(batches,t1, label="LR(s=5)")
plot(batches,t10, label="LR(s=10)")
plot(batches,t50, label="LR(s=50)")
plot(batches,t100, label="LR(s=100)")
xlabel("Batch size")
ylabel("Runtime(s)")
legend()

figure()
title("Similarty x' x2 / (||x|| ||x2||)")
plot(batches,angles[:, 1], label="LR(s=5)")
plot(batches,angles[:, 2], label="LR(s=10)")
plot(batches,angles[:, 3], label="LR(s=50)")
plot(batches,angles[:, 4], label="LR(s=100)")
legend()

@show angles


import CollectiveSpins

# Band structure (Omega) and collective decay rate (Gamma) for 1D chain for x-axis (along chain)
# and in the yz-plane (plane perpendicular to the atom chain) polarized atoms.
# See Asenjo-Garcia et al 10.1103/PhysRevX.7.031024 Figs. 1b and 1c.

# Define chain parameters
a = 0.2   # spin-spin distance
polarization_par = [1, 0, 0]
polarization_perp = [0, 1, 1]

k_max = 10000
k_list = [iii*pi/(a*k_max) for iii=-k_max:k_max]

band_structure_par = []
band_structure_perp = []
decay_par = []
decay_perp = []
for k in k_list
    append!(band_structure_par, CollectiveSpins.Omega_k_chain(k, a, polarization_par))
    append!(band_structure_perp, CollectiveSpins.Omega_k_chain(k, a, polarization_perp))
    append!(decay_par, CollectiveSpins.Gamma_k_chain(k, a, polarization_par))
    append!(decay_perp, CollectiveSpins.Gamma_k_chain(k, a, polarization_perp))
end

using PyPlot

subplot(211)
plot(k_list*a/pi, band_structure_par)
plot(k_list*a/pi, band_structure_perp, color="red")
xlabel("k d/pi")
ylabel("Omega(k)/gamma0")

subplot(212)
plot(k_list*a/pi, decay_par)
plot(k_list*a/pi, decay_perp, color="red")
xlabel("k d/pi")
ylabel("Gamma(k)/gamma0")

tight_layout()
savefig("band_structure_decay.svg")