using LinearAlgebra
dt = 0.002
function flow(u::Array{Float64,2},s::Array{Float64,1})
    sigma, rho, beta = s
    x, y, z = view(u,1,:), view(u,2,:), view(u,3,:)
    return [sigma.*(y - x)  x.*(rho .- z) - y  x.*y - beta.*z]'
end
function flow(u::Array{Float64,1},s::Array{Float64,1})
    sigma, rho, beta = s
    x, y, z = u[1], u[2], u[3]
    return [sigma*(y - x), x*(rho - z) - y, x*y - beta*z]
end
function mag_flow(u::Array{Float64,2},s::Array{Float64,1})
    d, n = size(u)
    mags = zeros(n)
    for i = 1:n
    	mags[i] = norm(flow(u[:,i],s))
    end
    return mags
end
function mag_flow(u::Array{Float64,1},s::Array{Float64,1})
    return norm(flow(u,s))
end
function dmag_flow(u::Array{Float64,2},s::Array{Float64,1})
    sigma, rho, beta = s
    d, n = size(u)
    db = zeros(d, n)
    for i = 1:n
    	x, y, z = u[1,i],u[2,i],u[3,i]
    	term = [sigma*(y - x), x*(rho - z) - y,
    		x*y - beta*z]
    	b = norm(term)
    	db[1,i] = dot([-sigma, (rho - z), y], term./b)
    	db[2,i] = dot([sigma, -1, x], term./b)
    	db[3,i] = dot([0., -x, -beta], term./b)
    end
    return db
end
function dmag_flow(u::Array{Float64,1},s::Array{Float64,1})
    sigma, rho, beta = s
    x, y, z = u[1], u[2], u[3]

    term = [sigma*(y - x), x*(rho - z) - y,
    		x*y - beta*z]
    b = norm(term)
    dxb2 = dot([-sigma, (rho - z), y], term./b)
    dyb2 = dot([sigma, -1, x], term./b)
    dzb2 = dot([0., -x, -beta], term./b)
    return [dxb2, dyb2, dzb2]
end
function dunit_flow(u::Array{Float64,1},s::Array{Float64,1})
    sigma, rho, beta = s
    x, y, z = u[1], u[2], u[3]
    F = [sigma*(y - x), x*(rho - z) - y, x*y - beta*z]
    b2 = dot(F,F)
    db = dmag_flow(u, s)
    dF = -(db'/b2).*[F F F]
    dF = dF .+ (1.0/sqrt(b2))*dflow(u,s)
    return dF
end

function step(u0, s, n)
    u_trj = zeros(3, n+1)
    u_trj[:,1] = u0
    n = n+1
    for i = 2:n
        u = u_trj[:,i-1]
        k1 = dt*flow(u, s)
        k2 = dt*flow(u .+ k1, s)
        u_trj[:,i] = @. u + (k1 + k2)/2
    end
    return u_trj
end
function next(u, s)
    k1 = dt*flow(u, s)
    k2 = dt*flow(u .+ k1, s)
    u_next = @. u + (k1 + k2)/2
    return u_next
end
function dflow(u::Array{Float64,2}, s::Array{Float64,1})
    sigma, rho, beta = s
    d, n = size(u)
    x = view(u,1,:)
    y = view(u,2,:)
    z = view(u,3,:)
    du = zeros(n, d, d)
    @. du[:,1,1] = -sigma
    @. du[:,1,2] = sigma
    @. du[:,2,1] = (rho - z)
    @. du[:,2,2] = -1.0
    @. du[:,2,3] = -x
    @. du[:,3,1] = y
    @. du[:,3,2] = x
    @. du[:,3,3] = -beta
    return reshape([du[:,:,1]'; du[:,:,2]';
                	du[:,:,3]'], d, d, n)
end
function dflow(u::Array{Float64,1},s::Array{Float64,1})
    du = zeros(3,3)
    x, y, z = u[1], u[2], u[3]
    sigma, rho, beta = s
    du[1,1] = -sigma
    du[1,2] = sigma
    du[2,1] = (rho - z)
    du[2,2] = -1.0
    du[2,3] = -x
    du[3,1] = y
    du[3,2] = x
    du[3,3] = -beta
    return du
end
function dstep(u::Array{Float64,2}, s::Array{Float64,1})
    d, n = size(u)
    du = zeros(d, d, n)
    dui = zeros(d,d)
    for i = 1:n
        ui = u[:,i]
        k1 = dt*flow(ui, s)
        dk1 = dt*dflow(ui, s)
        dk2 = dt*dflow(ui .+ k1, s)
        dui .= 1.0*I(d) .+ 0.5*(dk1 .+ dk2*(1.0*I(d) .+ dk1))
        du[:,:,i] .= dui
    end
    return du
end
function dstep(u::Array{Float64,1}, s::Array{Float64,1})
    d = 3
    du = zeros(d, d)
    k1 = dt*flow(u, s)
    dk1 = dt*dflow(u, s)
    dk2 = dt*dflow(u .+ k1, s)
    du = 1.0*I(d) .+ 0.5*(dk1 .+ dk2*(1.0*I(d) .+ dk1))
    return du
end

function d2flow(u::Array{Float64,2}, s::Array{Float64,1})
    n = size(u)[2]
    ddu = zeros(3,3,3,n)
    for i = 1:n
        ddu[2,3,1,i] = -1.0
        ddu[3,2,1,i] = 1.0
        ddu[3,1,2,i] = 1.0
        ddu[2,1,3,i] = -1.0
    end
    return ddu
end
function d2flow(u::Array{Float64,1}, s::Array{Float64,1})
    ddu = zeros(3,3,3)
    ddu[2,3,1] = -1.0
    ddu[3,2,1] = 1.0
    ddu[3,1,2] = 1.0
    ddu[2,1,3] = -1.0
    return ddu
end
function d2step(u::Array{Float64,2}, s::Array{Float64,1})
    d, n = size(u)
    ddu = zeros(d, d, d, n)
    k1 = dt*flow(u, s)
    dk1 = dt*dflow(u,s)
    dk2 = dt*dflow(u .+ k1, s)
    d2 = dt*d2flow(u, s)
    d2p = permutedims(d2,[1, 3, 2, 4])
    A = zeros(d,d,d)
    for k = 1:n
        for j = 1:d
        	A[:,:,j] .= d2p[:,:,j,k]*(1.0*I(d) .+ dk1[:,:,k])
        end
        A = permutedims(A, [1,3,2])
        for i = 1:d
        		ddu[:,:,i,k] = 0.5*(d2[:,:,i,k] .+
        				dk2[:,:,k]*d2[:,:,i,k] .+ A[:,:,i]*
        				(1.0*I(d) .+ dk1[:,:,k]))
        end
    end
    return ddu
end
function d2step(u::Array{Float64,1}, s::Array{Float64,1})
    ddu = zeros(3, 3, 3)
    k1 = dt*flow(u, s)
    dk1 = dt*dflow(u,s)
    dk2 = dt*dflow(u .+ k1, s)
    d2 = dt*d2flow(u, s)
    d2p = permutedims(d2,[1, 3, 2])
    A = zeros(3,3,3)
    for j = 1:3
        A[:,:,j] .= d2p[:,:,j]*(1.0*I(3) .+ dk1)
    end
    A = permutedims(A, [1,3,2])
    for i = 1:3
        ddu[:,:,i] = 0.5*(d2[:,:,i] .+ dk2*d2[:,:,i] .+ A[:,:,i]*
        					  (1.0*I(3) .+ dk1))
    end

    return ddu
end
function pert(u::Array{Float64,2},s::Array{Float64,1})
    n = size(u)[2]
    dsflow(x) = [0., x, 0.]
    pp = zeros(3, n)
    for i = 1:n
    	ui = u[:,i]
    	k1 = dt*flow(ui, s)
    	dk2 = dt*dflow(ui .+ k1,s)
    	dsk1 = dt*dsflow(ui[1])
    	dsk2 = dt*dsflow(ui[1] .+ k1[1]) .+
    		   dk2*dsk1
    	pp[:,i] = 0.5*(dsk1 + dsk2)
    end
    return pp
end
function pert(u::Array{Float64,1},s::Array{Float64,1})
    dsflow(x) = [0., x, 0.]
    k1 = dt*flow(u, s)

    dk2 = dt*dflow(u .+ k1,s)
    dsk1 = dt*dsflow(u[1])
    dsk2 = dt*dsflow(u[1] .+ k1[1]) .+
    		   dk2*dsk1
    pp = 0.5*(dsk1 + dsk2)
    return pp
end
function dpert_next(u_arr::Array{Float64,2},s::Array{Float64,1})
    d, n = size(u_arr)
    dpp = zeros(3,3,n)
    ddsflow = zeros(3,3)
    ddsflow[2,1] = 1.0
    dsflow(x) = [0., x, 0.]
    for j = 1:n
    	u = u_arr[:,j]
    	k1 = dt*flow(u, s)
    	dk1 = dt*dflow(u,s)
    	dk2 = dt*dflow(u .+ k1,s)
    	ddk2 = dt*d2flow(u .+ k1, s)
    	dsk1 = dt*dsflow(u[1])
    	ddsk1 = dt*ddsflow
    	d2 = dt*d2flow(u, s)
        d2p = permutedims(d2,[1, 3, 2])
        A = zeros(3,3,3)
        for j = 1:3
        	A[:,:,j] .= d2p[:,:,j]*(1.0*I(3) .+ dk1)
        end
        A = permutedims(A, [1,3,2])
    	ddsk2 = dt*ddsflow*(1.0*I(3) .+ dk1) .+
    		dk2*ddsk1
    	for i = 1:3
    		ddsk2[:,i] .+= A[:,:,i]*dsk1
    	end
    	dpp[:,:,j] .= 0.5*(ddsk1 + ddsk2)

      end
    return dpp
end
function dpert_next(u::Array{Float64,1},s::Array{Float64,1})
    dpp = zeros(3,3)
    ddsflow = zeros(3,3)
    ddsflow[2,1] = 1.0

    k1 = dt*flow(u, s)
    dk1 = dt*dflow(u,s)

    dk2 = dt*dflow(u .+ k1,s)
    ddk2 = dt*d2flow(u .+ k1, s)

    dsflow(x) = [0., x, 0.]
    dsk1 = dt*dsflow(u[1])
    ddsk1 = dt*ddsflow

    d2 = dt*d2flow(u, s)
    d2p = permutedims(d2,[1, 3, 2])
    A = zeros(3,3,3)
    for j = 1:3
        A[:,:,j] .= d2p[:,:,j]*(1.0*I(3) .+ dk1)
    end
    A = permutedims(A, [1,3,2])

    ddsk2 = dt*ddsflow*(1.0*I(3) .+ dk1) .+
    		dk2*ddsk1
    for j = 1:3
    	ddsk2[:,j] .+= A[:,:,j]*dsk1
    end
    dpp .= 0.5*(ddsk1 + ddsk2)
    return dpp
end
function dpert(u::Array{Float64, 1}, s::Array{Float64,1})
    du = dstep(u,s)
    dpp_next = dpert_next(u,s)
    dpp = du'\dpp_next'
    return dpp'
end
function dpert(u::Array{Float64, 2}, s::Array{Float64,1})
    du = dstep(u,s)
    dpp_next = dpert_next(u,s)
    d, n = size(u)
    dpp = zeros(3,3,n)
    for i = 1:n
    	dpp[:,:,i] .= du[:,:,i]'\dpp_next[:,:,i]'
    	dpp[:,:,i] .= dpp[:,:,i]'
    end
    return dpp
end

### AD functions
function lorenz63_rhs_ad(du, u, s, t)
    du[1] = s[1]*(u[2] - u[1])
    du[2] = u[1]*(s[2] - u[3]) - u[2]
    du[3] = u[1]*u[2] - s[3]*u[3]
end
function lorenz63_ad(u0, s, n)
    t = n*dt
    prob = ODEProblem(lorenz63_rhs_ad, u0, (0.,t), s)
    sol = Array(solve(prob, Tsit5(), saveat=dt))
    return sol[:,end]
end
function obj_fun(u0, s)
    prob = ODEProblem(lorenz63_rhs_ad, u0, (0.,1.1), s)
    #_prob = remake(prob,u0=u0,p=s)
    sol = solve(prob, Tsit5(), saveat=0.005)
    sum(sol[3,:])/size(sol,2)
end


function GasInflow(pr,J,GasPVT;n=10)

    #Calculate PseudoPressure
    GasPVT.Ps=map((x,y,z)->2*x/(y*z),GasPVT.P,GasPVT.Mug,GasPVT.Z)

    #Integral Matrix
    Integral = DataFrame()

    #Trasnsfer Pressure to Integrate Matrix
    Integral.P= GasPVT.P[2:end]

    #Pressure Differences
    Integral.DiffP=diff(GasPVT.P)

    #PseudoPressure Average
    Integral.AvgPs=0.0
    for i=1:size(GasPVT,1)-1
        Integral.AvgPs[i]=(GasPVT.Ps[i+1]+GasPVT.Ps[i])/2
    end
    Integral.Cum=accumulate(+, Integral.DiffP.*Integral.AvgPs)

    #Interpolation Integral
    mp=LinearInterpolation(Integral.P.*1, Integral.Cum.*1)

    #IPR Estimation

    pres=zeros(n,size(pr,1))
    Qg=zeros(n,size(pr,1)*size(J,1))
    leg=Array{String,1}(undef,size(pr,1)*size(J,1))
    count=0
    for p=1:size(pr,1)

        pres[:,p]=range(14.7,stop=pr[p]-50,length=n)

        for j=1:size(J,1)
    count=count+1
    leg[count]="Res. Pressure=$(pr[p]) psi; J=$(J[j]) Mscf/d*psi*cp"
      Qg[:,count]=map(pwf->J[j]*(mp(pr[p])-mp(pwf)),pres[:,p])

        end
    end

    return pres, Qg, leg

end #End Function


import Base: -, trace, inv, det, logdet, transpose, ctranspose, vecnorm, kron

############################# Unary sensitivities #############################
push!(ops, DiffOp(:(Base.:-), :(Tuple{DLA.AA}), [true]))
∇(::typeof(-), ::Arg1, p, Y::AA, Ȳ::AA, X::AA) = map(-, Ȳ)

push!(ops, DiffOp(:(Base.trace), :(Tuple{DLA.AM}), [true]))
∇(::typeof(trace), ::Arg1, p, Y::Real, Ȳ::Real, X::AM) = Diagonal(fill!(similar(X), Ȳ))

push!(ops, DiffOp(:(Base.inv), :(Tuple{DLA.AM}), [true]))
∇(::typeof(inv), ::Arg1, p, Y::AM, Ȳ::AM, X::AM) = -transpose(Y) * Ȳ * transpose(Y)

push!(ops, DiffOp(:(Base.det), :(Tuple{DLA.AM}), [true]))
∇(::typeof(det), ::Arg1, p, Y::Real, Ȳ::Real, X::AM) = Y * Ȳ * transpose(inv(X))

push!(ops, DiffOp(:(Base.logdet), :(Tuple{DLA.AM}), [true]))
∇(::typeof(logdet), ::Arg1, p, Y::Real, Ȳ::Real, X::AM) = Ȳ * transpose(inv(X))

push!(ops, DiffOp(:(Base.transpose), :(Tuple{DLA.AVM}), [true]))
∇(::typeof(transpose), ::Arg1, p, Y::AVM, Ȳ::AVM, X::AVM) = transpose(Ȳ)

push!(ops, DiffOp(:(Base.ctranspose), :(Tuple{DLA.AVM}), [true]))
∇(::typeof(ctranspose), ::Arg1, p, Y::AVM, Ȳ::AVM, X::AVM) = ctranspose(Ȳ)

push!(ops, DiffOp(:(Base.vecnorm), :(Tuple{DLA.AA}), [true]))
∇(::typeof(vecnorm), ::Arg1, p, Y::Real, Ȳ::Real, X::AA) = Ȳ ./ Y .* abs2.(X) ./ X

push!(ops, DiffOp(:(Base.vecnorm), :(Tuple{Real}), [true]))
∇(::typeof(vecnorm), ::Arg1, p, Y::Real, Ȳ::Real, X::Real) = Ȳ * sign(X)

############################# Binary sensitivities #############################
push!(ops, DiffOp(:(Base.:*), :(Tuple{DLA.AVM, DLA.AVM}), [true, true]))
∇(::typeof(*), ::Arg1, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = Ȳ * B'
∇(::typeof(*), ::Arg2, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = A' * Ȳ

push!(ops, DiffOp(:(Base.At_mul_B), :(Tuple{DLA.AVM, DLA.AVM}), [true, true]))
∇(::typeof(At_mul_B), ::Arg1, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = B * Ȳ'
∇(::typeof(At_mul_B), ::Arg2, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = A * Ȳ

push!(ops, DiffOp(:(Base.A_mul_Bt), :(Tuple{DLA.AVM, DLA.AVM}), [true, true]))
∇(::typeof(A_mul_Bt), ::Arg1, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = Ȳ * B
∇(::typeof(A_mul_Bt), ::Arg2, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = Ȳ' * A

push!(ops, DiffOp(:(Base.At_mul_Bt), :(Tuple{DLA.AVM, DLA.AVM}), [true, true]))
∇(::typeof(At_mul_Bt), ::Arg1, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = B' * Ȳ'
∇(::typeof(At_mul_Bt), ::Arg2, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = Ȳ' * A'

push!(ops, DiffOp(:(Base.Ac_mul_B), :(Tuple{DLA.AVM, DLA.AVM}), [true, true]))
∇(::typeof(Ac_mul_B), ::Arg1, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = B * Ȳ'
∇(::typeof(Ac_mul_B), ::Arg2, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = A * Ȳ

push!(ops, DiffOp(:(Base.A_mul_Bc), :(Tuple{DLA.AVM, DLA.AVM}), [true, true]))
∇(::typeof(A_mul_Bc), ::Arg1, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = Ȳ * B
∇(::typeof(A_mul_Bc), ::Arg2, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = Ȳ' * A

push!(ops, DiffOp(:(Base.Ac_mul_Bc), :(Tuple{DLA.AVM, DLA.AVM}), [true, true]))
∇(::typeof(Ac_mul_Bc), ::Arg1, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = B' * Ȳ'
∇(::typeof(Ac_mul_Bc), ::Arg2, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = Ȳ' * A'

push!(ops, DiffOp(:(Base.:/), :(Tuple{DLA.AVM, DLA.AVM}), [true, true]))
∇(::typeof(/), ::Arg1, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = Ȳ / B'
∇(::typeof(/), ::Arg2, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = -Y' * (Ȳ / B')

push!(ops, DiffOp(:(Base.At_rdiv_B), :(Tuple{DLA.AVM, DLA.AVM}), [true, true]))
∇(::typeof(At_rdiv_B), ::Arg1, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = B \ Ȳ'
∇(::typeof(At_rdiv_B), ::Arg2, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = -Y' * (B \ Ȳ')'

push!(ops, DiffOp(:(Base.A_rdiv_Bt), :(Tuple{DLA.AVM, DLA.AVM}), [true, true]))
∇(::typeof(A_rdiv_Bt), ::Arg1, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = Ȳ / B
∇(::typeof(A_rdiv_Bt), ::Arg2, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = -(B' \ Ȳ') * Y

push!(ops, DiffOp(:(Base.At_rdiv_Bt), :(Tuple{DLA.AVM, DLA.AVM}), [true, true]))
∇(::typeof(At_rdiv_Bt), ::Arg1, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = B' \ Ȳ'
∇(::typeof(At_rdiv_Bt), ::Arg2, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = -(B' \ Ȳ') * Y

push!(ops, DiffOp(:(Base.Ac_rdiv_B), :(Tuple{DLA.AVM, DLA.AVM}), [true, true]))
∇(::typeof(Ac_rdiv_B), ::Arg1, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = B \ Ȳ'
∇(::typeof(Ac_rdiv_B), ::Arg2, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = -Y' * (B \ Ȳ')'

push!(ops, DiffOp(:(Base.A_rdiv_Bc), :(Tuple{DLA.AVM, DLA.AVM}), [true, true]))
∇(::typeof(A_rdiv_Bc), ::Arg1, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = Ȳ / B
∇(::typeof(A_rdiv_Bc), ::Arg2, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = -(B' \ Ȳ') * Y

push!(ops, DiffOp(:(Base.Ac_rdiv_Bc), :(Tuple{DLA.AVM, DLA.AVM}), [true, true]))
∇(::typeof(Ac_rdiv_Bc), ::Arg1, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = B' \ Ȳ'
∇(::typeof(Ac_rdiv_Bc), ::Arg2, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = -(B' \ Ȳ') * Y

push!(ops, DiffOp(:(Base.:\), :(Tuple{DLA.AVM, DLA.AVM}), [true, true]))
∇(::typeof(\), ::Arg1, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = -(A' \ Ȳ) * Y'
∇(::typeof(\), ::Arg2, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = A' \ Ȳ

push!(ops, DiffOp(:(Base.At_ldiv_B), :(Tuple{DLA.AVM, DLA.AVM}), [true, true]))
∇(::typeof(At_ldiv_B), ::Arg1, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = -Y *(A \ Ȳ)'
∇(::typeof(At_ldiv_B), ::Arg2, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = A \ Ȳ

push!(ops, DiffOp(:(Base.A_ldiv_Bt), :(Tuple{DLA.AVM, DLA.AVM}), [true, true]))
∇(::typeof(A_ldiv_Bt), ::Arg1, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = -(Ȳ' / A)' * Y'
∇(::typeof(A_ldiv_Bt), ::Arg2, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = Ȳ' / A

push!(ops, DiffOp(:(Base.At_ldiv_Bt), :(Tuple{DLA.AVM, DLA.AVM}), [true, true]))
∇(::typeof(At_ldiv_Bt), ::Arg1, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = -Y * (Ȳ' / A')
∇(::typeof(At_ldiv_Bt), ::Arg2, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = Ȳ' / A'

push!(ops, DiffOp(:(Base.Ac_ldiv_B), :(Tuple{DLA.AVM, DLA.AVM}), [true, true]))
∇(::typeof(Ac_ldiv_B), ::Arg1, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = -Y * (A \ Ȳ)'
∇(::typeof(Ac_ldiv_B), ::Arg2, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = A \ Ȳ

push!(ops, DiffOp(:(Base.A_ldiv_Bc), :(Tuple{DLA.AVM, DLA.AVM}), [true, true]))
∇(::typeof(A_ldiv_Bc), ::Arg1, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = -(Ȳ' / A)' * Y'
∇(::typeof(A_ldiv_Bc), ::Arg2, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = Ȳ' / A

push!(ops, DiffOp(:(Base.Ac_ldiv_Bc), :(Tuple{DLA.AVM, DLA.AVM}), [true, true]))
∇(::typeof(Ac_ldiv_Bc), ::Arg1, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = -Y * (Ȳ' / A')
∇(::typeof(Ac_ldiv_Bc), ::Arg2, p, Y::ASVM, Ȳ::ASVM, A::AVM, B::AVM) = Ȳ' / A'

push!(ops, DiffOp(:(Base.vecnorm), :(Tuple{DLA.AA, Real}), [true, true]))
∇(::typeof(vecnorm), ::Arg1, p, Y::Real, Ȳ::Real, A::AA, B::Real) =
    Ȳ .* Y^(1 - B) .* abs.(A).^B ./ A
∇(::typeof(vecnorm), ::Arg2, p, Y::Real, Ȳ::Real, A::AA, B::Real) =
    Ȳ * (Y^(1 - B) * sum(abs.(A).^B .* log.(abs.(A))) - Y * log(Y)) / B

push!(ops, DiffOp(:(Base.vecnorm), :(Tuple{Real, Real}), [true, true]))
∇(::typeof(vecnorm), ::Arg1, p, Y::Real, Ȳ::Real, A::Real, B::Real) = Ȳ * sign(A)
∇(::typeof(vecnorm), ::Arg2, p, Y::Real, Ȳ::Real, A::Real, B::Real) = 0

push!(ops, DiffOp(:(Base.kron), :(Tuple{AM, AM}), [true, true]))
∇(::typeof(kron), ::Type{Val{1}}, p, Y::AM, Ȳ::AM, A::AM, B::AM) =
    ∇(zero(A), kron, Val{1}, p, Y, Ȳ, A, B)
∇(::typeof(kron), ::Type{Val{2}}, p, Y::AM, Ȳ::AM, A::AM, B::AM) =
    ∇(zero(B), kron, Val{2}, p, Y, Ȳ, A, B)
function ∇(Ā::AM, ::typeof(kron), ::Type{Val{1}}, p, Y::AM, Ȳ::AM, A::AM, B::AM)
    @assert size(Ā) == size(A)
    (I, J), (K, L), m = size(A), size(B), length(Y)
    @inbounds for j = reverse(1:J), l = reverse(1:L), i = reverse(1:I)
        āij = Ā[i, j]
        for k = reverse(1:K)
            āij += Ȳ[m] * B[k, l]
            m -= 1
        end
        Ā[i, j] = āij
    end
    return Ā
end
function ∇(B̄::AM, ::typeof(kron), ::Type{Val{2}}, p, Y::AM, Ȳ::AM, A::AM, B::AM)
    @assert size(B̄) == size(B)
    (I, J), (K, L), m = size(A), size(B), length(Y)
    @inbounds for j = reverse(1:J), l = reverse(1:L), i = reverse(1:I)
        aij = A[i, j]
        for k = reverse(1:K)
            B̄[k, l] += aij * Ȳ[m]
            m -= 1
        end
    end
    return B̄
end


struct VirtualGate{T<:AbstractFloat} <: Wall{T}
    sp::SVector{2,T}
    ep::SVector{2,T}
    normal::SVector{2,T}
    width::T
    center::SVector{2,T}
    color::Symbol
    name::String
end
function VirtualGate(sp::AbstractVector, ep::AbstractVector,
    n::AbstractVector, color::Symbol, name::String = "Virtual Gate")
    T = eltype(sp) <: Integer ? Float64 : eltype(sp)
    n = normalize(n)
    d = dot(n, ep-sp)
    if abs(d) > 10eps(T)
        error("Normal vector is not actually normal to the wall: dot = $d")
    end
    w = norm(ep - sp)
    center = @. (ep+sp)/2
    return VirtualGate{T}(SVector{2,T}(sp), SVector{2,T}(ep), SVector{2,T}(n), w, SVector{2,T}(center), color, name)
end

import DynamicalBilliards.specular!
@inline function specular!(p::AbstractParticle{T}, r::VirtualGate{T})::Void where {T}
    n = normalvec(r, p.pos)
    φ = atan2(n[2], n[1]) + 0.95(π*rand() - π/2) #this cannot be exactly π/2
    p.vel = SVector{2,T}(cos(φ), sin(φ))
    return nothing
end

import DynamicalBilliards.translation
translation(w::VirtualGate, vec) = VirtualGate(w.sp + vec, w.ep + vec, w.normal, w.color)

import DynamicalBilliards.collisiontime
function collisiontime(p::Particle{T}, w::VirtualGate{T})::T where {T}
    n = normalvec(w, p.pos)
    denom = dot(p.vel, n)
    # case of velocity pointing away of wall:
    denom ≥ 0 && return Inf
    posdot = dot(w.sp-p.pos, n)
    # Case of particle starting behind finite wall:
    posdot ≥ 0 && return Inf
    colt = posdot/denom
    intersection = p.pos .+ colt .* p.vel
    dfc = norm(intersection - w.center)
    if dfc > w.width/2
        return Inf
    else
        return colt
    end
end

@inline function project_to_line(point, c, n)
    posdot = dot(c - point, n)
    intersection = point + posdot*n
end

import DynamicalBilliards.distance_init
function distance_init(pos::SVector{2,T}, w::VirtualGate{T})::T where {T}
    n = normalvec(w, pos)
    posdot = dot(w.sp - pos, n)
    if posdot ≥ 0 # I am behind wall
        intersection = project_to_line(pos, w.center, n)
        dfc = norm(intersection - w.center)
        if dfc > w.width/2
            return +1.0 # but not directly behind
        else
            return -1.0
        end
    end
    v1 = pos - w.sp
    dot(v1, n)
end


module ssa_varimax

export varimax_rotate!

using LinearAlgebra

"""
Varimax rotation as per
   Groth, A. and M. Ghil, 2011: Multivariate singular spectrum analysis and the
      road to phase synchronization, Physical Review E, 84, 036206.
   Groth, A. and M. Ghil, 2015: Monte Carlo Singular Spectrum Analysis (SSA)
      revisited: Detecting oscillator clusters in multivariate datasets, Journal
      of Climate, 28, 7873-7893.

Code based on Matlab implementation by Andreas Groth at
https://www.mathworks.com/matlabcentral/fileexchange/65939-multichannel-singular-spectrum-analysis-varimax-tutorial

Copyright (c) 2018, Andreas Groth
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

* Redistributions of source code must retain the above copyright notice, this
  list of conditions and the following disclaimer.

* Redistributions in binary form must reproduce the above copyright notice,
  this list of conditions and the following disclaimer in the documentation
  and/or other materials provided with the distribution

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
"""
function varimax(A::Array{float_type, 3}, reltol=sqrt(eps(float_type)),
                 maxit=1000, normalize=true, G=[]) where {float_type<:AbstractFloat}
   M, D, S = size(A)
   B = A

   B = reshape(B, D*M, S, 1)[:, :, 1]
   D = D*M

   T = diagm(0=>ones(S))

   for iter=1:maxit
      maxTheta = 0
      for i = 1:(S-1)
         for j = (i+1):S
            u = B[:, i].^2 - B[:,j].^2
            v = 2*B[:,i].*B[:,j]
            if M>1
               u = sum(reshape(u, M, :), dims=1)'
               v = sum(reshape(v, M, :), dims=1)'
            end

            usum = sum(u, dims=1);
            vsum = sum(v, dims=1);
            numer = 2*u'*v - 2*usum*vsum/D;
            denom = u'*u - v'*v - (usum^2 - vsum^2)/D;
            theta = atan(numer[1, 1], denom[1, 1])/4;
            maxTheta = max(maxTheta, abs(theta));
            Tij = [cos(theta) -sin(theta); sin(theta) cos(theta)];
            B[:, [i, j]] = B[:, [i, j]] * Tij;
            T[:, [i, j]] = T[:, [i, j]] * Tij;
         end
      end
      if (maxTheta < reltol)
         break
      end
   end

   return B, T
end

function varimax_rotate!(EW, EV, M, D, S)
   EVscaled=EV[:,1:S]*diagm(0=>sqrt.(EW[1:S]))
   EVreshape=reshape(EVscaled,M,D,S)

   _, T = varimax(EVreshape)

   EV[:,1:S]=EV[:,1:S]*T
   EW[1:S]=diag(T'*diagm(0=>EW[1:S])*T)

   EV = EV[:, sortperm(EW, rev=true)]
   EW = sort(EW, rev=true)

   return EW, EV
end

end


function logsumexp{I<:Number}(x::AbstractArray{I})
    m = maximum(x)
    r = 0.0
    for i = 1:endof(x)
        r += exp(x[i] - m)
    end
    return log(r) + m
end


# Laplace distribution

export Laplace

@parameterized Laplace() ≪ (1/2) * Lebesgue(ℝ)

for N in AFFINEPARS
    @eval begin
        proxy(d::Laplace{$N}) = affine(params(d), Laplace())
        logdensity(d::Laplace{$N}, x) = logdensity(proxy(d), x)
        basemeasure(d::Laplace{$N}) = basemeasure(proxy(d))
    end
end

# @affinepars Laplace

function logdensity(d::Laplace{()} , x)
    return -abs(x)
end

Base.rand(rng::AbstractRNG, μ::Laplace{()}) = rand(rng, Dists.Laplace())

≪(::Laplace, ::Lebesgue{X}) where X <: Real = true

TV.as(::Laplace) = asℝ

distproxy(::Laplace{()}) = Dists.Laplace()
distproxy(d::Laplace{(:μ,)}) = Dists.Laplace(d.μ, 1.0)
distproxy(d::Laplace{(:σ,)}) = Dists.Laplace(0.0, d.σ)
distproxy(d::Laplace{(:μ,:σ)}) = Dists.Laplace(d.μ, d.σ)
distproxy(d::Laplace{(:ω,)}) = Dists.Laplace(0.0, inv(d.ω))
distproxy(d::Laplace{(:μ,:ω)}) = Dists.Laplace(d.μ, inv(d.ω))


# Make a macro that will create matrix R and vector r for restricting the VAR
# The user will write beta[1,2,1] - beta[2,1,3] = 10 (beta[from, to, lag]) and through
# regular expressions this will get parsed into a vector of the matrix
# macro restrict(ex) gamaa restricts the coefficients

function β(from, to, lag)
	col = fill(0.0, (vars^2)*lags + vars*ntyp)
	println(1 + vars*ntyp + (lag-1)*(vars^2) + (from-1)*vars + (to-1))
	col[1 + vars*ntyp + (lag-1)*(vars^2) + (from-1)*vars + (to-1)] = 1.0
	return col
end

function betaa(from, to, lag)
	return β(from, to, lag)
end

function γ(eq)
	col = fill(0.0, (vars^2)*lags + vars*ntyp)
	col[eq] = 1.0
	return col
end

function gamaa(eq)
	return γ(eq)
end

macro restrictions(ex)
	ind = [2:2:length(ex.args)]
	R = apply(Expr, prepend!([:($(ex.args[i].args[1])) for i=ind], [:call, :hcat]));
	r = [:($(ex.args[i].args[2].args[2])) for i=ind];
	return :(tuple($R, $r))
end

# vars, lags, and ntyp have to be defined
# vars = e.vars
# ntyp = e.ntyp
# lags = e.lags

# @restrictions begin
#     gamaa(3)=1
# end

# @restrictions begin
#     gamaa(4) + betaa(1,2,3) + betaa(3,2,1) = 10
# 	betaa(3,4,2) - betaa(3,1,3) = 5
# end

# Heh, hard
# function makeR(R::Matrix{Float64})
# 	base = eye(size(R)[1])

# 	if numberInColumn==1
# 		base

# 	end
# end

# function testRestrictions(R::Matrix{Float64})
# 	@assume rank(R)==apply(max, size(mat)) "The constraints are not independent. One can be expressed as linear combination of other."
# 	# Alright, this is making me nuts :-D
# 	@assume all(mapslices((x) -> sum(map((y) -> y==0.0, x)), getR(mat), [2]) .>= apply(min, size(R))) "Your constants yield the model overidentified. For example, beta(1,2,3) + beta(2,3,1) = 0 and beta(1,2,3) + beta(2,2,1) = 0 imply only equality of beta(2,2,1) - beta(2,3,1) = 0 and does not depend on beta(1,2,3)."
# end


#=
disk-phantom.jl
generate random disk phantoms
2019-07-03, Jeff Fessler, University of Michigan
=#

export disk_phantom_params

# using MIRT: jim
using Random: seed!

"""
`params = disk_phantom_params( ; ...)`

Generate ellipse phantom parameters for a head-sized disk
plus many disks within it,
designed so that the disks have some minimum separation `minsep`
to avoid overlap and to simplify patch-based model fitting.
# Arguments
- `fov::Real = 240` image field of view in mm
- `rhead::Real = 100` background radius for "head" [mm]
- `muhead::Real = 1000` "mu" (intensity) value for background head disk
- `rmin::Real = 10` min radius for random disks
- `rmax::Real = 20` max radius for random disks
- `mumin::Real = 100` range of "mu" values for disks
- `mumax::Real = 300`
- `ndisk::Integer = 10` # of random disks
- `minsep::Real = 8` minimum disk separation in mm
- `maxtry::Integer = 500` give up on adding more disks if this is reached
- `warn::Bool = false` warn if maxtry reached?
- `seed::Integer = 0`if nonzero then use this seed
"""
function disk_phantom_params( ;
		fov::Real = 240,
		rhead::Real = 100, # background radius for "head" [mm]
		muhead::Real = 1000, # "mu" value for background head
		rmin::Real = 10, # min radius for disks
		rmax::Real = 20, # max radius for disks
		mumin::Real = 100, # range of "mu" values for disks
		mumax::Real = 300,
		ndisk::Integer = 10,
		minsep::Real = 8, # minimum separation in mm
		maxtry::Integer = 500, # give up on adding more disks if this is reached
		warn::Bool = false, # warn if maxtry reached?
		seed::Integer = 0,
	)

	params = zeros(Float32, ndisk+1, 6)
	params[end,:] = [0, 0, rhead, rhead, 0, muhead]

	randu = (a, b; f = x->x) -> a + f(rand()) * (b - a)
	seper = (a,b,r1,x,y,r2) -> sqrt.((x .- a).^2 + (y .- b).^2) .- r1 .- r2

	(seed != 0) && seed!(seed)

	for id=1:ndisk
		trial = 0
		for ii=1:maxtry
			rc = randu(0, rhead - rmin - minsep, f = x->x^0.5)
			phi = rand() * 2*pi
			rad = randu(rmin, rmax)
			val = randu(mumin, mumax)
			trial = [rc*cos(phi), rc*sin(phi), rad, rad, 0, val]

			id == 1 && break # first one is always fine

			# see if trial is too close to another one
			c = 1:(id-1) # check these
			sep = seper(params[c,1], params[c,2], params[c,3],
				trial[1], trial[2], trial[3])

			(minimum(sep) > minsep) && (rhead - (rc + rad) > minsep) && break
			ii == maxtry && warn && @warn("need more tries")
			ii == maxtry && return params[[1:id; end],:]
		end

#=
		if id > 1
			c = 1:(id-1) # check these
			sep = seper(params[c,1], params[c,2], params[c,3],
				trial[1], trial[2], trial[3])
			@show minimum(sep)
		end
		@show id, trial
=#

		params[id,:] = trial
	end

	return params
end

"""
`disk_phantom_params(:test)`
self test

`disk_phantom_params(:show)`
show example
"""
function disk_phantom_params(test::Symbol)
	ig = image_geom(nx = 128, fov = 240)
	params = disk_phantom_params( ; minsep = ig.dx*8)

	if test == :show
		tmp = ellipse_im(ig, params, oversample=3)
		return jim(ig.x, ig.y, tmp)
	end

	test != :test && throw(ArgumentError("test $test"))
	disk_phantom_params(:show)
	true
end

# disk_phantom_params(:test)


using ToyFHE
using OffsetArrays
using Test

# For now use the BFV parameter estimator and just copy it over.
# We need to figure out the right way to do this in CKKS
bfv_params = BFVParams(
    7, # plaintext modulus
    ; eval_mult_count = 3
)

params = CKKSParams(bfv_params.ℛ, bfv_params.relin_window,
    bfv_params.σ)

scale = 2^40
Tscale = FixedRational{scale}
Tscalesq = FixedRational{Int128(scale)^2}

plain = CKKSEncoding{Tscale}(zero(plaintext_space(params)))
plain .= OffsetArray(LinRange(0.0, 1.0, 2048), 0:2047)

re = convert(ToyFHE.NTT.RingElement, plain)

# TODO: Can we do better here by being more careful with the FFT?
@test real.(collect(CKKSEncoding{Tscalesq}(re*re))) ≈ LinRange(0.0, 1.0, 2048).^2 atol=10^-4

# First test the encoder in isolation without the extra noise from encryption

kp = keygen(params)
c = encrypt(kp, plain)

@test real.(collect(decrypt(kp, c))) ≈ LinRange(0.0, 1.0, 2048) atol=10^-4
@test real.(collect(decrypt(kp, c*c))) ≈ LinRange(0.0, 1.0, 2048).^2 atol=10^-4


module RunTests

using UnitJuMP
using Test

import MutableArithmetics
const _MA = MutableArithmetics

function runtests()
    for name in names(@__MODULE__; all = true)
        if startswith("$(name)", "test_")
            @testset "$(name)" begin
                getfield(@__MODULE__, name)()
            end
        end
    end
    return
end

function test_examples()
    examples = joinpath(@__DIR__, "..", "examples")
    for filename in filter(f -> endswith(f, ".jl"), readdir(examples))
        include(joinpath(examples, filename))
    end
    return
end

function test_jump_variables()
    m = Model()

    @variable(m, x ≥ 0, u"m/s")
    @test x == UnitJuMP.UnitVariableRef(x.variable, u"m/s")
    @test unit(x) == u"m/s"
    @test owner_model(x) === m

    @variable(m, y[1:4], u"km/hr")
    @test y[1] == UnitJuMP.UnitVariableRef(y[1].variable, u"km/hr")
    return
end

function test_jump_constraints()
    m = Model()

    @variable(m, u[1:2] ≥ 0, u"m")
    @variable(m, v[1:4], u"km/hr")
    @variable(m, w, Bin)
    @variable(m, y)

    maxspeed = 4u"ft/s"
    speed = 2.3u"m/s"

    # Various combination of coefficient and variables with and without units
    @constraint(m, 2v[1] ≤ maxspeed)
    @test num_constraints(m, AffExpr, MOI.LessThan{Float64}) == 1

    @constraint(m, 2.3v[1] ≤ maxspeed)
    @test num_constraints(m, AffExpr, MOI.LessThan{Float64}) == 2

    @constraint(m, c1a, speed * w ≤ maxspeed)
    @test num_constraints(m, AffExpr, MOI.LessThan{Float64}) == 3
    @test unit(c1a) == u"m/s"

    @constraint(m, c1b, 40u"km/hr" * w ≤ 15u"m/s")
    @test num_constraints(m, AffExpr, MOI.LessThan{Float64}) == 4
    @test unit(c1b) == u"km/hr"

    @constraint(m, c1c, w * speed ≤ maxspeed)
    @test num_constraints(m, AffExpr, MOI.LessThan{Float64}) == 5
    @test unit(c1c) == u"m/s"

    @constraint(m, c1d, sum(v[i] for i in 1:4) ≤ maxspeed)
    @test num_constraints(m, AffExpr, MOI.LessThan{Float64}) == 6
    @test unit(c1d) == u"km/hr"

    @constraint(m, c1, 2v[1] + 4v[2] ≤ maxspeed)
    @test typeof(c1) <: UnitJuMP.UnitConstraintRef
    @test unit(c1) == u"km/hr"

    @constraint(m, c2, 2v[1] + 4v[2] ≤ maxspeed, u"m/s")
    @test unit(c2) == u"m/s"
    @test normalized_rhs(c2.constraint) ==
          convert(Float64, Unitful.uconvert(u"m/s", maxspeed).val)

    @variable(m, z, Bin)

    maxlength = 1000u"yd"
    period = 1.5u"hr"
    @constraint(m, c3, u[2] + period * v[2] ≤ maxlength * z, u"cm")
    @test unit(c3) == u"cm"

    @constraint(m, c3b, u[2] + 1.5u"hr" * v[2] ≤ 1000u"yd" * z, u"cm")
    @test unit(c3b) == u"cm"
    return
end

function test_mutable_arithmetics()
    m = Model()
    @variable(m, x ≥ 0)
    @variable(m, y)

    xu = UnitJuMP.UnitVariableRef(x, u"m/s")
    yu = UnitJuMP.UnitVariableRef(y, u"km/hr")

    @test _MA.@rewrite(xu) == UnitJuMP.UnitExpression(1x, u"m/s")
    @test _MA.@rewrite(-xu) == UnitJuMP.UnitExpression(-x, u"m/s")
    @test _MA.@rewrite(5xu) == UnitJuMP.UnitExpression(5x, u"m/s")
    @test _MA.@rewrite(xu * 5) == UnitJuMP.UnitExpression(5x, u"m/s")
    @test _MA.@rewrite(xu / 5) == UnitJuMP.UnitExpression(0.2x, u"m/s")

    @test _MA.@rewrite(xu + yu) == UnitJuMP.UnitExpression(x + y / 3.6, u"m/s")
    @test _MA.@rewrite(xu - yu) == UnitJuMP.UnitExpression(x - y / 3.6, u"m/s")
    @test _MA.@rewrite(-xu + yu) ==
          UnitJuMP.UnitExpression(-x + y / 3.6, u"m/s")

    speed = 10u"m/s"
    length = 10u"m"
    freq = 1u"1/s"
    @test _MA.@rewrite(speed * x) == UnitJuMP.UnitExpression(10x, u"m/s")
    @test _MA.@rewrite(x * speed) == UnitJuMP.UnitExpression(10x, u"m/s")
    @test _MA.@rewrite(xu * 10) == UnitJuMP.UnitExpression(10x, u"m/s")
    @test _MA.@rewrite(speed * xu) == UnitJuMP.UnitExpression(10x, u"m^2/s^2")
    @test _MA.@rewrite(xu * speed) == UnitJuMP.UnitExpression(10x, u"m^2/s^2")
    @test _MA.@rewrite(x / speed) == UnitJuMP.UnitExpression(0.1x, u"s/m")
    @test _MA.@rewrite(speed + 2xu) == UnitJuMP.UnitExpression(10 + 2x, u"m/s")
    @test _MA.@rewrite(-speed + xu) == UnitJuMP.UnitExpression(-10 + x, u"m/s")
    @test _MA.@rewrite(2xu - speed) == UnitJuMP.UnitExpression(2x - 10, u"m/s")

    @test _MA.@rewrite(xu + 0.5 * speed) ==
          UnitJuMP.UnitExpression(x + 5, u"m/s")
    @test _MA.@rewrite(xu - 0.5 * speed) ==
          UnitJuMP.UnitExpression(x - 5, u"m/s")
    @test _MA.@rewrite(xu + speed * 0.5) ==
          UnitJuMP.UnitExpression(x + 5, u"m/s")
    @test _MA.@rewrite(xu - speed * 0.5) ==
          UnitJuMP.UnitExpression(x - 5, u"m/s")
    @test _MA.@rewrite(xu + length * freq) ==
          UnitJuMP.UnitExpression(x + 10, u"m/s")
    @test _MA.@rewrite(xu - length * freq) ==
          UnitJuMP.UnitExpression(x - 10, u"m/s")
    @test _MA.@rewrite(0.5 * speed + xu) ==
          UnitJuMP.UnitExpression(x + 5, u"m/s")

    @test _MA.@rewrite(4xu + 2yu - speed) ==
          UnitJuMP.UnitExpression(4x + 2y / 3.6 - 10, u"m/s")
    @test _MA.@rewrite(4xu - 2yu + speed) ==
          UnitJuMP.UnitExpression(4x - 2y / 3.6 + 10, u"m/s")

    ex = xu + yu
    @test _MA.@rewrite(ex + yu + speed) ==
          UnitJuMP.UnitExpression(x + 2y / 3.6 + 10, u"m/s")
    ex = xu + yu
    @test _MA.@rewrite(yu + ex) == UnitJuMP.UnitExpression(3.6x + 2y, u"km/hr")
    ex = xu + yu
    @test _MA.@rewrite(yu - 2ex) == UnitJuMP.UnitExpression(-7.2x - y, u"km/hr")

    @test_throws ErrorException _MA.@rewrite(xu + y)
    @test_throws ErrorException _MA.@rewrite(xu - y)
    @test_throws ErrorException _MA.@rewrite(xu + 5y)
    @test_throws ErrorException _MA.@rewrite(xu - 3y)
    @test_throws ErrorException _MA.@rewrite(y + xu)
    @test_throws ErrorException _MA.@rewrite(y - xu)
    @test_throws ErrorException _MA.@rewrite(y + 5xu)
    @test_throws ErrorException _MA.@rewrite(y - 4xu)
    @test_throws ErrorException _MA.@rewrite(xu + 4)
    @test_throws ErrorException _MA.@rewrite(xu - 4)

    @test_throws ErrorException @constraint(m, x <= xu)
    @test_throws ErrorException @constraint(m, xu <= y)

    return
end

function test_operators()
    m = Model()

    @variable(m, x)
    @variable(m, y)
    @variable(m, z)

    xu = UnitJuMP.UnitVariableRef(x, u"km")
    yu = UnitJuMP.UnitVariableRef(y, u"hr")
    zu = UnitJuMP.UnitVariableRef(z, u"s")

    @test -xu == UnitJuMP.UnitExpression(-x, u"km")
    @test xu + 200u"m" == UnitJuMP.UnitExpression(x + 0.2, u"km")
    @test 200u"m" + xu == UnitJuMP.UnitExpression(1000x + 200, u"m")
    @test xu - 400u"m" == UnitJuMP.UnitExpression(x - 0.4, u"km")
    @test 400u"m" - xu == UnitJuMP.UnitExpression(-1000x + 400, u"m")

    @test 200u"km" * x == UnitJuMP.UnitExpression(200 * x, u"km")
    @test x * 200u"km" == UnitJuMP.UnitExpression(200 * x, u"km")
    @test 200 * xu == UnitJuMP.UnitExpression(200 * x, u"km")
    @test xu * 200 == UnitJuMP.UnitExpression(200 * x, u"km")
    @test 1.5u"s" * xu == UnitJuMP.UnitExpression(1.5 * x, u"s*km")
    @test xu * 1.5u"s" == UnitJuMP.UnitExpression(1.5 * x, u"s*km")

    @test xu / 0.5 == UnitJuMP.UnitExpression(2 * x, u"km")
    @test x / 0.5u"s" == UnitJuMP.UnitExpression(2 * x, u"s^-1")
    @test xu / 0.5u"s" == UnitJuMP.UnitExpression(2 * x, u"km/s")

    @test yu + zu == UnitJuMP.UnitExpression(y + z / 3600, u"hr")
    @test yu - zu == UnitJuMP.UnitExpression(y - z / 3600, u"hr")

    expr = yu + 1800 * zu

    @test expr + 3600u"s" == UnitJuMP.UnitExpression(y + 0.5z + 1, u"hr")
    @test expr - 1800u"s" == UnitJuMP.UnitExpression(y + 0.5z - 0.5, u"hr")
    @test 2u"s" + expr == UnitJuMP.UnitExpression(3600y + 1800z + 2, u"s")
    @test 3600u"s" - expr ==
          UnitJuMP.UnitExpression(-3600y - 1800z + 3600, u"s")

    @test 2 * expr == UnitJuMP.UnitExpression(2y + z, u"hr")
    @test expr * 2 == UnitJuMP.UnitExpression(2y + z, u"hr")
    @test expr / 0.5 == UnitJuMP.UnitExpression(2y + z, u"hr")
    @test 2u"kW" * expr == UnitJuMP.UnitExpression(2y + z, u"kW*hr")
    @test expr * 2u"kW" == UnitJuMP.UnitExpression(2y + z, u"kW*hr")
    @test expr / 0.5u"km" == UnitJuMP.UnitExpression(2y + z, u"hr/km")

    @variable(m, w)
    wu = UnitJuMP.UnitVariableRef(w, u"minute")

    @test expr + wu == UnitJuMP.UnitExpression(y + 0.5z + w / 60, u"hr")
    @test expr - wu == UnitJuMP.UnitExpression(y + 0.5z - w / 60, u"hr")
    @test wu + expr == UnitJuMP.UnitExpression(60y + 30z + w, u"minute")
    @test wu - expr == UnitJuMP.UnitExpression(-60y - 30z + w, u"minute")
    @test expr + yu == UnitJuMP.UnitExpression(2y + 0.5z, u"hr")

    expr2 = 30wu + yu

    @test expr + expr2 == UnitJuMP.UnitExpression(2y + 0.5z + 0.5w, u"hr")
    @test expr2 + expr == UnitJuMP.UnitExpression(120y + 30z + 30w, u"minute")
    @test 2 * expr - expr2 == UnitJuMP.UnitExpression(y + z - 0.5w, u"hr")

    # Combination of unit expression with unitless variables/expressions
    # should throw error
    @test_throws ErrorException w + wu
    @test_throws ErrorException w - wu
    @test_throws ErrorException 4w - wu
    @test_throws ErrorException wu - 4w
    @test_throws ErrorException 4 - wu
    @test_throws ErrorException wu + 4
    @test_throws ErrorException wu - 4
    @test_throws ErrorException x + 2y - wu
    @test_throws ErrorException xu + 2y - wu

    return
end

function test_number_times_quantity_times_variable()
    a = 9.81u"m/s^2"
    model = Model()
    @variable(model, y, u"s")
    expr = @expression(model, 0.5 * a * y)
    @test Unitful.unit(expr) == u"m/s"
    return
end

function test_number_times_quantity_times_variable_to_existing()
    a = 9.81u"m/s^2"
    model = Model()
    @variable(model, x, u"m/s")
    @variable(model, y, u"s")
    expr = @expression(model, x + 0.5 * a * y)
    @test Unitful.unit(expr) == u"m/s"
    return
end

function test_unit_variable()
    model = Model()
    @variable(model, x, u"kW*hr")
    @variable(model, y, u"s")

    # Start values
    set_start_value(x, 1u"MW*hr")
    @test start_value(x) == 1000u"kW*hr"
    @test_throws ErrorException set_start_value(x, 13.2)

    # Lower bounds
    set_lower_bound(x, 10u"MJ")
    @test has_lower_bound(x)
    @test lower_bound(x) == (10 / 3.6)u"kW*hr"
    lref = LowerBoundRef(x)
    @test typeof(lref) <: ConstraintRef
    delete_lower_bound(x)
    @test !has_lower_bound(x)
    @test_throws ErrorException set_lower_bound(x, -4.0)

    # Upper bounds
    set_upper_bound(y, 10u"hr")
    @test has_upper_bound(y)
    @test upper_bound(y) == 36000u"s"
    uref = UpperBoundRef(y)
    @test typeof(uref) <: ConstraintRef
    delete_upper_bound(y)
    @test !has_lower_bound(y)
    @test_throws ErrorException set_upper_bound(y, 42)

    # Fix values
    fix(y, 1u"hr")
    @test is_fixed(y)
    @test fix_value(y) == 3600u"s"
    fref = FixRef(y)
    @test typeof(fref) <: ConstraintRef
    unfix(y)
    @test !is_fixed(y)
    @test_throws ErrorException fix(y, 10)

    # Integer and binary
    set_integer(x)
    @test is_integer(x)
    iref = IntegerRef(x)
    @test typeof(iref) <: ConstraintRef
    unset_integer(x)
    @test !is_integer(x)
    set_binary(y)
    @test is_binary(y)
    bref = BinaryRef(y)
    @test typeof(bref) <: ConstraintRef
    unset_binary(y)
    @test !is_binary(y)

    # Naming
    set_name(x, "xvar")
    @test name(x) == "xvar"

    return
end

end

RunTests.runtests()


using BlockDiagonalMatrices
using SparseArrays
using LinearAlgebra
using TimerOutputs

abstract type Model end

struct ModelParameters
    Nj::Integer
    kind::Real
    kcoll::Real
    Gamma::Real
    omega::Real
    eta::Real
    dt::Real
    Tfinal::Real
    Ntime::Integer
    outpoints::Integer
    mu::Real
    _outsteps::Integer

    function ModelParameters(; Nj::Integer=1,
                               kind::Real=1.0,
                               kcoll::Real=1.0,
                               Gamma::Real=1.0,
                               omega::Real=1.0,
                               eta::Real=1.0,
                               dt::Real=0.0001,
                               Tfinal::Real=1.0,
                               outpoints::Integer=0,
                               mu::Real=0.0)

        Ntime = Int(floor(Tfinal / dt)) # Number of timesteps

        outsteps = 1
        if outpoints > 0

            try
                outsteps = Int(round(Tfinal / dt / outpoints, digits=3))
            catch InexactError
                @warn "The requested $outpoints output points does not divide
                the total time steps. Using the full time output."
            end
        end
        # outpoints = Ntime

        new(Nj, kind, kcoll, Gamma, omega, eta, dt, Tfinal, Ntime, outpoints, mu, outsteps)
    end
end

function get_time(mp::ModelParameters)
    t = (1:mp.Ntime) * mp.dt
    t[mp._outsteps:mp._outsteps:end]
end

struct KrausOperator
    M0::BlockDiagonal
    M::BlockDiagonal
end


@testset "Pre Whitening filter" begin
    test_signal = complex.(randn(1000, 2), randn(1000, 2)) / sqrt(2)
    Rxx = calc_variance_covariance(test_signal)
    pre_whitening_filter = calc_prewhitening_filter(Rxx)
    @test pre_whitening_filter ≈ I atol = 0.1

    test_signal = complex.(randn(1000, 2), randn(1000, 2)) / sqrt(2) * sqrt(100) # Power: 100
    Rxx = calc_variance_covariance(test_signal)
    pre_whitening_filter = calc_prewhitening_filter(Rxx)
    @test pre_whitening_filter ≈ I atol = 0.1
end

@testset "Amplitude filter" begin
    test_signal = complex.(randn(1000, 2), randn(1000, 2)) / sqrt(2)
    Rxx = calc_variance_covariance(test_signal)
    amplitude_filter = calc_amplitude_filter(Rxx)
    @test amplitude_filter ≈ I atol = 0.1

    test_signal = complex.(randn(1000, 2), randn(1000, 2)) / sqrt(2) * sqrt(100) # Power: 100
    Rxx = calc_variance_covariance(test_signal)
    pre_whitening_filter = calc_amplitude_filter(Rxx)
    @test pre_whitening_filter ≈ I atol = 0.1
end

@testset "Eigen beamformer" begin
    test_signal = complex.(randn(1000, 4), randn(1000, 4)) / sqrt(2) / 5 .+ ones(1000, 4)
    Rxx = calc_variance_covariance(test_signal)
    beamformer = calc_eigen_beamformer(Rxx)
    @test beamformer ≈ -1 * [0.5, 0.5, 0.5, 0.5] atol = 0.03
end


using Plots
Plots.GRBackend()

function flood(f::Array;shift=(0.,0.),cfill=:RdBu_11,clims=(),kv...)
    if length(clims)==2
        @assert clims[1]<clims[2]
        @. f=min(clims[2],max(clims[1],f))
    else
        clims = (minimum(f),maximum(f))
    end
    Plots.contour(axes(f,1).+shift[1],axes(f,2).+shift[2],f',
        linewidth=0,linecolor=:black,
        fill = (true,palette(cfill)),
        clims = clims, aspect_ratio=:equal; kv...)
end

addbody(x,y;c=:black) = Plots.plot!(Shape(x,y), c=c, legend=false)
function μbody_plot!(flow::Flow{m};levels=[0.5],lines=:black) where {m}
    @inside flow.σ[I] = sum(WaterLily.ϕ(i,CartesianIndex(I,i),flow.μ₀) for i ∈ 1:m)/m
    BC!(flow.σ)
    contour!(flow.σ';levels,lines)
end

function sim_gif!(sim;duration=1,step=0.1,verbose=true,
                  remeasure=false,μbody=false,size=(700,600),kv...)
    t₀ = round(sim_time(sim))
    t = range(t₀,t₀+duration;step)
    gr(show=false;size)
    @time @gif for tᵢ in t
        sim_step!(sim,tᵢ;remeasure)
        @inside sim.flow.σ[I] = WaterLily.curl(3,I,sim.flow.u)*sim.L/sim.U
        flood(sim.flow.σ;shift=(-0.5,-0.5),clims=(-5,5),kv...)
        μbody && μbody_plot!(sim.flow)
        verbose && println("tU/L=",round(tᵢ,digits=4),
            ", Δt=",round(sim.flow.Δt[end],digits=3))
    end
    return
end

function vectors(f::Array)
    x = axes(f,1)' .* ones(size(f,2))
    y = ones(size(f,2))' .* axes(f,2)
    u,v = f[:,:,1]',f[:,:,2]'
    Plots.quiver(x[:],y[:],quiver=(u[:],v[:]),
                aspect_ratio=:equal)
end


#   Copyright (c) 2020-2022 Velexi Corporation
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

"""
multiply.jl defines methods for the *(x, y) function
"""

# --- Exports

import Base.:(*)

# --- Method definitions

# ------ Specializations involving an AbstractMultivector instance

# M::AbstractMultivector, B::AbstractBlade
# B::AbstractBlade, M::AbstractMultivector
*(M::AbstractMultivector, B::AbstractBlade) =
    Multivector(map(C -> C * B, blades(M)))

*(B::AbstractBlade, M::AbstractMultivector) =
    Multivector(map(C -> B * C, blades(M)))

# M::AbstractMultivector, B::AbstractScalar,
# B::AbstractScalar, M::AbstractMultivector
*(M::AbstractMultivector, B::AbstractScalar) = B * M
*(B::AbstractScalar, M::AbstractMultivector) = value(B) * M

# M::AbstractMultivector, B::One
# B::One, M::AbstractMultivector
*(M::AbstractMultivector, B::One) = M
*(B::One, M::AbstractMultivector) = M

# M::AbstractMultivector, B::Zero
# B::Zero, M::AbstractMultivector
*(M::AbstractMultivector, B::Zero) = B
*(B::Zero, M::AbstractMultivector) = B

# M::AbstractMultivector, x::Real
# x::Real, M::AbstractMultivector
*(M::AbstractMultivector, x::Real) = x * M
*(x::Real, M::AbstractMultivector) = Multivector(map(B -> x * B, blades(M)))

# M::AbstractMultivector, v::Vector
# v::Vector, M::AbstractMultivector
*(M::AbstractMultivector, v::Vector) = Multivector(map(B -> B * v, blades(M)))
*(v::Vector, M::AbstractMultivector) = Multivector(map(B -> v * B, blades(M)))

# ------ Specializations involving an AbstractBlade instance

# B::AbstractBlade, C::AbstractScalar
# B::AbstractScalar, B::AbstractBlade
*(B::AbstractBlade, C::AbstractScalar) =
    Blade(B, volume=volume(B) * value(C), copy_basis=false)
*(C::AbstractScalar, B::AbstractBlade) = B * C

# B::AbstractBlade, C::One
# B::One, C::AbstractBlade
*(B::AbstractBlade, C::One) = B
*(B::One, C::AbstractBlade) = C

# B::AbstractBlade, C::Zero
# B::Zero, C::AbstractBlade
*(B::AbstractBlade, C::Zero) = C
*(B::Zero, C::AbstractBlade) = B

# B::AbstractBlade, x::Real
# x::Real, B::AbstractBlade
*(B::AbstractBlade, x::Real) = Blade(B, volume=volume(B) * x)
*(x::Real, B::AbstractBlade) = B * x

# ------ Specializations involving a Blade instance

# B::Blade, C::Blade
function *(B::Blade, C::Blade)
    # Check arguments
    assert_dim_equal(B, C)

    # Compute geometric product
    M = C
    for i in grade(B):-1:1
        M = basis(B)[:, i] * M
    end
    M = volume(B) * M
end

# B::Blade, C::Pseudoscalar
# B::Pseudoscalar, C::Blade
*(B::Blade, C::Pseudoscalar) = B ⋅ C
*(B::Pseudoscalar, C::Blade) = zero(B)

# B::Blade, v::Vector
# v::Vector, B::Blade
function *(B::Blade, v::Vector{<:Real})
    # --- Check arguments

    # Check that B and v are from the same real vector space
    assert_dim_equal(v, B)

    # Check that grade(B) == 1. B * v is not defined when grade(B) > 1
    if grade(B) > 1
        error("Geometric product B * v is not defined when grade(B) > 1")
    end

    # Compute geometric product
    B_dot_v = B ⋅ v
    B_wedge_v = B ∧ v

    if B_dot_v == zero(B_dot_v)
        return B_wedge_v
    elseif B_wedge_v == zero(B_wedge_v)
        return B_dot_v
    end

    Multivector([B_dot_v, B_wedge_v])
end

function *(v::Vector{<:Real}, B::Blade)
    # Check arguments
    assert_dim_equal(v, B)

    # Compute geometric product
    v_dot_B = v ⋅ B
    v_wedge_B = v ∧ B

    if v_dot_B == zero(v_dot_B)
        return v_wedge_B
    elseif v_wedge_B == zero(v_wedge_B)
        return v_dot_B
    end

    Multivector([v_dot_B, v_wedge_B])
end

# ------ Specializations involving a Pseudoscalar instance

# B::Pseudoscalar, C::Pseudoscalar
*(B::Pseudoscalar, C::Pseudoscalar) = contract_left(B, C)

# B::Pseudoscalar, C::AbstractScalar
# B::AbstractScalar, C::Pseudoscalar
*(B::Pseudoscalar, C::AbstractScalar) =
    Pseudoscalar(B, value=value(B) * value(C))

*(B::AbstractScalar, C::Pseudoscalar) = C * B

# B::Pseudoscalar, C::One
# B::One, C::Pseudoscalar
*(B::Pseudoscalar, C::One) = B
*(B::One, C::Pseudoscalar) = C

# B::Pseudoscalar, C::Zero
# B::Zero, C::Pseudoscalar
*(B::Pseudoscalar, C::Zero) = C
*(B::Zero, C::Pseudoscalar) = B

# B::Pseudoscalar, x::Real
# x::Real, B::Pseudoscalar
*(B::Pseudoscalar, x::Real) = Pseudoscalar(B, value=x * value(B))
*(x::Real, C::Pseudoscalar) = C * x

# ------ Specializations involving an AbstractScalar instance

# B::AbstractScalar, C::One
# B::One, C::AbstractScalar
*(B::AbstractScalar, C::One) = B
*(B::One, C::AbstractScalar) = C

# B::AbstractScalar, C::Zero
# B::Zero, C::AbstractScalar
*(B::AbstractScalar, C::Zero) = C
*(B::Zero, C::AbstractScalar) = B

# B::AbstractScalar, x::Real
# x::Real, B::AbstractScalar
*(B::AbstractScalar, x::Real) = Scalar{typeof(value(B))}(value(B) * x)
*(x::Real, B::AbstractScalar) = B * x

# B::AbstractScalar, v::Vector
# v::Vector, B::AbstractScalar
*(v::Vector{<:Real}, B::AbstractScalar) = Blade(value(B) * v)
*(B::AbstractScalar, v::Vector{<:Real}) = v * B

# ------ Specializations involving a Scalar instance

# B::Scalar, C::Scalar
*(B::Scalar, C::Scalar) = Scalar{typeof(value(B))}(value(B) * value(C))

# ------ Specializations involving a One instance

# B::One, C::One
*(B::One, C::One) = B

# B::One, C::Zero
# B::Zero, C::One
*(B::One, C::Zero) = C
*(B::Zero, C::One) = B

# ------ Specializations involving a Zero instance

# B::Zero, C::Zero
*(B::Zero, C::Zero) = B


module _GridEn
export GridEn
using Plots
using Plots.PlotMeasures
using StatsBase: Histogram, fit
    """
        GDE, GDR, _ = GridEn(Sig)

    Returns the gridded distribution entropy (`GDE`) and the gridded
    distribution rate (`GDR`) estimated from the data sequence (`Sig`) using
    the default  parameters:
    grid coarse-grain = 3, time delay = 1, logarithm = base 2

        GDE, GDR, PIx, GIx, SIx, AIx = GridEn(Sig)

    In addition to GDE and GDR, GridEn returns the following indices
    estimated for the data sequence (`Sig`) using the default  parameters:
    [`PIx`]   - Percentage of points below the line of identity (LI)
    [`GIx`]   - Proportion of point distances above the LI
    [`SIx`]   - Ratio of phase angles (w.r.t. LI) of the points above the LI
    [`AIx`]   - Ratio of the cumulative area of sectors of points above the LI

        GDE, GDR, ..., = GridEn(Sig::AbstractArray{T,1} where T<:Real; m::Int=3, tau::Int=1, Logx::Real=exp(1), Plotx::Bool=false)

    Returns the gridded distribution entropy (`GDE`) estimated from the data
    sequence (`Sig`) using the specified 'keyword' arguments:

    # Arguments:
    `m`     - Grid coarse-grain (m x m sectors), an integer > 1 \n
    `tau`   - Time Delay, a positive integer    \n
    `Logx`  - Logarithm base, a positive scalar \n
    `Plotx` - When Plotx == true, returns gridded Poicaré plot and a bivariate
              histogram of the grid point distribution (default: false)  \n

    # See also `PhasEn`, `CoSiEn`, `SlopEn`, `BubbEn`, `MSEn`

    # References:
        [1] Chang Yan, et al.,
                "Novel gridded descriptors of poincaré plot for analyzing
                heartbeat interval time-series."
                Computers in biology and medicine
                109 (2019): 280-289.

        [2] Chang Yan, et al.
                "Area asymmetry of heart rate variability signal."
                Biomedical engineering online
                16.1 (2017): 1-14.

        [3] Alberto Porta, et al.,
                "Temporal asymmetries of short-term heart period variability
                are linked to autonomic regulation."
                American Journal of Physiology-Regulatory, Integrative and
                Comparative Physiology
                295.2 (2008): R550-R557.

        [4] C.K. Karmakar, A.H. Khandoker and M. Palaniswami,
                "Phase asymmetry of heart rate variability signal."
                Physiological measurement
                36.2 (2015): 303.

    """
    function GridEn(Sig::AbstractArray{T,1} where T<:Real; m::Int=3, tau::Int=1,
        Logx::Real=exp(1), Plotx::Bool=false)

    N = size(Sig,1)
    (N > 10) ? nothing : error("Sig:   must be a numeric vector")
    (m > 1) ? nothing :  error("m:     must be an integer > 1")
    (tau >0) ? nothing : error("tau:   must be an integer > 0")
    (Logx>0) ? nothing : error("Logx:  must be a positive number > 0")

    Sig_n = (Sig.-minimum(Sig))/(maximum(Sig)-minimum(Sig))
    Temp = hcat(Sig_n[1:end-tau], Sig_n[1+tau:end])
    Edges = collect(0:1/m:1); Edges[end] = 1.1
    N = fit(Histogram,(Temp[:,1],Temp[:,2]),(Edges,Edges)).weights
    Pj = reverse(transpose(N),dims=1)/size(Temp,1); Ppi = Pj[Pj.>0];
    if round(sum(Ppi),digits=4) != 1
        @warn("Potential error of estimated probabilities: P = $(sum(Ppi))")
    end
    GDE = -sum(Ppi.*(log.(Ppi)/log(Logx)));
    GDR = sum(N.!=0)/(m*m);

    T2   = atand.(Temp[:,2]./Temp[:,1])
    Dup  = sum(abs.(diff(Temp[T2.>45,:],dims=2)))
    Dtot = sum(abs.(diff(Temp[T2.!=45,:],dims=2)))
    Sup  = sum((T2[T2.>45].-45))
    Stot = sum(abs.(T2[T2.!=45].-45))
    Aup  = sum(abs.((T2[T2.>45].-45).*sqrt.(sum(Temp[T2.>45,:].^2,dims=2))))
    Atot = sum(abs.((T2[T2.!=45].-45).*sqrt.(sum(Temp[T2.!=45,:].^2,dims=2))))

    PIx = 100*sum(T2.<45)/sum(T2.!=45)
    GIx = 100*Dup/Dtot
    SIx = 100*Sup/Stot
    AIx = 100*Aup/Atot

    if Plotx
        ntrvl = range(0,1,length=m+1)
        p1 = plot(Sig_n[1:end-tau],Sig_n[tau+1:end],seriestype =:scatter,
            c=:magenta, markerstrokecolor=:magenta, markersize = 2)
        plot!(hcat(ntrvl,ntrvl)',hcat(zeros(m+1),ones(m+1))',c=:cyan,lw=3)
        plot!(hcat(zeros(m+1),ones(m+1))',hcat(ntrvl,ntrvl)',c=:cyan,lw=3)
        plot!([0, 1],[0, 1],c=:black, legend=false, aspect_ratio=:equal,
            xlim = (0, 1), ylim = (0, 1),
            xticks=[0, 1], yticks=[0, 1],  xlabel= "X ₙ" , ylabel="X ₙ₊ₜ")

        p2 = histogram2d(Temp[:,1],Temp[:,2], nbins=ntrvl,
            xticks=[0, 1], yticks=[0, 1], xlabel= "X ₙ" , ylabel="X ₙ₊ₜ",
            xlim = (0, 1), ylim = (0, 1), aspect_ratio=:equal,
            colorbar_entry=false, seriescolor=:cool,show_empty_bins=true)

        xx = plot(p1,p2,layout=(1,2),legend=false, framestyle=:grid, margin=7mm)
        display(xx)
    end

    return GDE, GDR, PIx, GIx, SIx, AIx
    end

end

"""
Copyright 2021 Matthew W. Flood, EntropyHub

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

For Terms of Use see https://github.com/MattWillFlood/EntropyHub

"""


using Plots, NipponColors

D = dict_color()

x = 0:0.02:1 |> collect
plot(x, sin.(2π .* x), color = D["mizu"])
plot!(x, cos.(2π .* x), color = D["梔子"])


using LinearAlgebra
function vech(x)
    a = triu(x)
    a = a[findall(x->x!=0,a)]
end


abstract type KLIEPSolver end

struct CD_KLIEP <: KLIEPSolver end

### interface

KLIEP(Ψx, Ψy, ::CD_KLIEP) = KLIEP!(SparseIterate(size(Ψx, 1)), Ψx, Ψy)

function KLIEP!(x::SparseIterate, Ψx, Ψy)
    f = CDKLIEPLoss(Ψx, Ψy)
    g = ProxZero()
    coordinateDescent!(x, f, g)
end

function spKLIEP!(x::SparseIterate, Ψx, Ψy, λ)
    f = CDKLIEPLoss(Ψx, Ψy)
    g = ProxL1(λ)
    coordinateDescent!(x, f, g)
end

spKLIEP!(x::SparseIterate, f::CDKLIEPLoss, g::ProxL1) = coordinateDescent!(x, f, g)

function _compute_loadings(x, Ψx, Ψy)
    p, nx = size(Ψx)
    ny = size(Ψy, 2)
    Ψyr = transpose( rhat(x, Ψy) .* transpose(Ψy) )
    s = zeros(p)
    for k = 1:p
        s[k] = sqrt( (var(Ψx[k, :]) / nx) + (var(Ψyr[k, :]) / ny) )
    end
    s
end

function spKLIEP_scaled!(x::SparseIterate, Ψx, Ψy, λ)
    f = CDKLIEPLoss(Ψx, Ψy)

    γ = _compute_loadings(x, Ψx, Ψy)
    g = ProxL1(λ, γ)
    coordinateDescent!(x, f, g)

    γ = _compute_loadings(x, Ψx, Ψy)
    g = ProxL1(λ, γ)
    coordinateDescent!(x, f, g, CDOptions(; warmStart=true))
end

spKLIEP(Ψx, Ψy, λ, ::CD_KLIEP; loadings=true) =
    (loadings ? spKLIEP_scaled!(SparseIterate(size(Ψx, 1)), Ψx, Ψy, λ) : spKLIEP!(SparseIterate(size(Ψx, 1)), Ψx, Ψy, λ))

# Invert row of a Hessian using the scaled lasso
function Hinv_row(H, row, λ0)
    p = size(H, 1)
    e = zeros(p)
    e[row] = -1.

    x = SparseIterate(p)
    x[row] = 1.
    f = CDQuadraticLoss(H, e)

    # init sigma
    σ = sqrt( dot(x, H * x) )
    λ = e .+ 1.

    for iter = 1:10
        coordinateDescent!(x, f, ProxL1(λ0 * σ, λ))
        σnew = sqrt( dot(x, H * x) )
        if abs(σnew - σ) / σ < 1e-3
            break
        end
        σ = σnew
    end

    x
end


export result_beam_force,result_beam_displacement

"""
    result_beam_force(assembly,lc_id,beam_id)
读取已求解的assembly实例的单个beam单元内力
# Arguments
- `assembly`: Assembly类型实例
- `lc_id`: 工况id
- `beam_id`: 单元id
# 返回
- 包含梁单元坐标系下的12个节点力向量
"""
function result_beam_force(assembly,lc_id,beam_id)
    path=assembly.working_path
    beam_id=string(beam_id)
    if !(id in keys(assembly.structure.beams))
        throw("beam id "*string(beam_id)*" doesn't exists!")
    end
    lc_id=string(lc_id)
    if !(lc_id in keys(assembly.lcset.cases))
        throw("loadcase with id "*lc_id*" doesn't exist!")
    end
    elm=assembly.structure.beams[beam_id]
    i=elm.node1.hid
    j=elm.node2.hid
    iDOFs=6*i-5:6*i
    jDOFs=6*j-5:6*j
    u=read_vector(path*"/.analysis",lc_id*"_u.v")
    T=elm.T
    uᵉ=T*[d[iDOFs];u[jDOFs]]

    rDOF=findall(x->x==true,elm.release)
    Kᵉ=integrateK(elm)
    K̃ᵉ,P̃ᵉ=static_condensation(K̃ᵉ,zeros(12),rDOF)
    return K̃ᵉ*uᵉ
end

"""
    result_beam_displacement(assembly,lc_id,beam_id)
读取已求解的assembly实例的单个beam单元两端节点位移
# 参数
- `assembly`: Assembly类型实例
- `lc_id`: 工况id
- `beam_id`: 单元id
# 返回
- 包含梁单元坐标系下的12个节点位移向量
"""
function result_beam_displacement(assembly,lc_id,beam_id)
    path=assembly.working_path
    beam_id=string(beam_id)
    if !(id in keys(assembl.structure.beams))
        throw("beam id "*string(beam_id)*" doesn't exists!")
    end
    lc_id=string(lc_id)
    if !(lc_id in keys(assembly.lcset.cases))
        throw("loadcase with id "*lc_id*" doesn't exist!")
    end
    loadcase=assembly.lcset.cases[lc_id]
    u=read_vector(path*"/.analysis",lc_id*"_u.v")
    elm=assembly.structure.beams[beam_id]
    i=elm.node1.hid
    j=elm.node2.hid
    iDOFs=6*i-5:6*i
    jDOFs=6*j-5:6*j
    if beam_id in keys(loadcase.beam_forces)
        Pᵉ=zeros(12)
    else
        force=loadcase.beam_forces[beam_id]
        Pᵉ=FEBeam.integrateP!(elm,force)
    end
    T=elm.T
    uᵉ=T*[u[iDOFs];u[jDOFs]]
    K̄ᵉ=integrateK(elm)
    i=.!elm.release
    j=elm.release
    Kᵢᵢ=K̄ᵉ[i,i]
    Kᵢⱼ=K̄ᵉ[i,j]
    Kⱼᵢ=K̄ᵉ[j,i]
    Kⱼⱼ=K̄ᵉ[j,j]
    Pⱼ=Pᵉ[j]
    ūᵉ=zeros(12)
    ūᵉ[i]=uᵉ[i]
    ūᵉ[j]=inv(Kⱼⱼ)*(Pⱼ-Kⱼᵢ*uᵉ[i])
    return ūᵉ
end

"""
    result_beam_displacement(assembly,lc_id)
读取已求解的assembly实例的所有beam单元两端节点位移
# 参数
- `assembly`: Assembly类型实例
- `lc_id`: 工况id
# 返回
- 所有梁在单元坐标系下的12个节点位移向量
"""
function result_beam_displacement(assembly,lc_id)
    path=assembly.working_path
    lc_id=string(lc_id)
    if !(lc_id in keys(assembly.lcset.cases))
        throw("loadcase with id "*lc_id*" doesn't exist!")
    end
    loadcase=assembly.lcset.cases[lc_id]
    u=read_vector(path*"/.analysis",lc_id*"_u.v")
    res=Dict{String,Array{Float64}}()
    for elm in values(assembly.structure.beams)
        i=elm.node1.hid
        j=elm.node2.hid
        iDOFs=6*i-5:6*i
        jDOFs=6*j-5:6*j
        if !(elm.id in keys(loadcase.beam_forces))
            Pᵉ=zeros(12)
        else
            force=loadcase.beam_forces[elm.id]
            Pᵉ=BeamModule.integrateP!(elm,force)
        end
        T=elm.T
        uᵉ=T*[u[iDOFs];u[jDOFs]]
        K̄ᵉ=integrateK(elm)
        i=.!elm.release
        j=elm.release
        Kᵢᵢ=K̄ᵉ[i,i]
        Kᵢⱼ=K̄ᵉ[i,j]
        Kⱼᵢ=K̄ᵉ[j,i]
        Kⱼⱼ=K̄ᵉ[j,j]
        Pⱼ=Pᵉ[j]
        ūᵉ=zeros(12)
        ūᵉ[i]=uᵉ[i]
        ūᵉ[j]=inv(Kⱼⱼ)*(Pⱼ-Kⱼᵢ*uᵉ[i])
        res[elm.id]=ūᵉ
    end
    return res
end

function result_element()
    Δdᵉ=(elm.T)'*[Δd[6*hid1-5:6*hid1];Δd[6*hid2-5:6*hid2]]
    Δϵᵉ=Bᵉ*Δdᵉ
    Δσᵉ=∫D*Δϵᵉ
    σᵉ+=Δσᵉ
end


# # Single stack with HEVI solver tutorial based on the 3D Burgers + tracer equations

# This tutorial implements the Burgers equations with a tracer field
# in a single element stack. The flow is initialized with a horizontally
# uniform profile of horizontal velocity and uniform initial temperature. The fluid
# is heated from the bottom surface. Gaussian noise is imposed to the horizontal
# velocity field at each node at the start of the simulation. The tutorial demonstrates how to
#
#   * Initialize a [`BalanceLaw`](@ref ClimateMachine.BalanceLaws.BalanceLaw) in a single stack configuration;
#   * Return the horizontal velocity field to a given profile (e.g., large-scale advection);
#   * Remove any horizontal inhomogeneities or noise from the flow.
#   * Use horizontal explicit vertical implicit (HEVI) solver in the Single stack setup
#
# The second and third bullet points are demonstrated imposing Rayleigh friction, horizontal
# diffusion and 2D divergence damping to the horizontal momentum prognostic equation.

# Equations solved in balance law form:

# ```math
# \begin{align}
# \frac{∂ ρ}{∂ t} =& - ∇ ⋅ (ρ\mathbf{u}) \\
# \frac{∂ ρ\mathbf{u}}{∂ t} =& - ∇ ⋅ (-μ ∇\mathbf{u}) - ∇ ⋅ (ρ\mathbf{u} \mathbf{u}') - γ[ (ρ\mathbf{u}-ρ̄\mathbf{ū}) - (ρ\mathbf{u}-ρ̄\mathbf{ū})⋅ẑ ẑ] - ν_d ∇_h (∇_h ⋅ ρ\mathbf{u}) \\
# \frac{∂ ρcT}{∂ t} =& - ∇ ⋅ (-α ∇ρcT) - ∇ ⋅ (\mathbf{u} ρcT)
# \end{align}
# ```

# Boundary conditions:
# ```math
# \begin{align}
# z_{\mathrm{min}}: & ρ = 1 \\
# z_{\mathrm{min}}: & ρ\mathbf{u} = \mathbf{0} \\
# z_{\mathrm{min}}: & ρcT = ρc T_{\mathrm{fixed}} \\
# z_{\mathrm{max}}: & ρ = 1 \\
# z_{\mathrm{max}}: & ρ\mathbf{u} = \mathbf{0} \\
# z_{\mathrm{max}}: & -α∇ρcT = 0
# \end{align}
# ```

# where
#  - ``t`` is time
#  - ``ρ`` is the density
#  - ``\mathbf{u}`` is the velocity (vector)
#  - ``\mathbf{ū}`` is the horizontally averaged velocity (vector)
#  - ``μ`` is the dynamic viscosity tensor
#  - ``γ`` is the Rayleigh friction frequency
#  - ``ν_d`` is the horizontal divergence damping coefficient
#  - ``T`` is the temperature
#  - ``α`` is the thermal diffusivity tensor
#  - ``c`` is the heat capacity
#  - ``ρcT`` is the thermal energy

# Solving these equations is broken down into the following steps:
# 1) Preliminary configuration
# 2) PDEs
# 3) Space discretization
# 4) Time discretization
# 5) Solver hooks / callbacks
# 6) Solve
# 7) Post-processing

# # Preliminary configuration

# ## [Loading code](@id Loading-code-burgers-bjfnk)

# First, we'll load our pre-requisites
#  - load external packages:
using MPI
using Distributions
using OrderedCollections
using Plots
using StaticArrays
using LinearAlgebra: Diagonal, tr

#  - load CLIMAParameters and set up to use it:
using CLIMAParameters
struct EarthParameterSet <: AbstractEarthParameterSet end
const param_set = EarthParameterSet()

#  - load necessary ClimateMachine modules:
using ClimateMachine
using ClimateMachine.Mesh.Topologies
using ClimateMachine.Mesh.Grids
using ClimateMachine.Writers
using ClimateMachine.DGMethods
using ClimateMachine.DGMethods.NumericalFluxes
using ClimateMachine.BalanceLaws:
    BalanceLaw, Prognostic, Auxiliary, Gradient, GradientFlux, parameter_set

using ClimateMachine.Mesh.Geometry: LocalGeometry
using ClimateMachine.MPIStateArrays
using ClimateMachine.GenericCallbacks
using ClimateMachine.SystemSolvers
using ClimateMachine.ODESolvers
using ClimateMachine.VariableTemplates
using ClimateMachine.SingleStackUtils

#  - import necessary ClimateMachine modules: (`import`ing enables us to
#  provide implementations of these structs/methods)
using ClimateMachine.Orientations:
    Orientation,
    FlatOrientation,
    init_aux!,
    vertical_unit_vector,
    projection_tangential

import ClimateMachine.BalanceLaws:
    vars_state,
    source!,
    flux_second_order!,
    flux_first_order!,
    compute_gradient_argument!,
    compute_gradient_flux!,
    init_state_auxiliary!,
    init_state_prognostic!,
    BoundaryCondition,
    boundary_conditions,
    boundary_state!

# ## Initialization

# Define the float type (`Float64` or `Float32`)
const FT = Float64;
# Initialize ClimateMachine for CPU.
ClimateMachine.init(; disable_gpu = true);

const clima_dir = dirname(dirname(pathof(ClimateMachine)));

# Load some helper functions for plotting
include(joinpath(clima_dir, "docs", "plothelpers.jl"));

# # Define the set of Partial Differential Equations (PDEs)

# ## Define the model

# Model parameters can be stored in the particular [`BalanceLaw`](@ref
# ClimateMachine.BalanceLaws.BalanceLaw), in this case, the `BurgersEquation`:

Base.@kwdef struct BurgersEquation{FT, APS, O} <: BalanceLaw
    "Parameters"
    param_set::APS
    "Orientation model"
    orientation::O
    "Heat capacity"
    c::FT = 1
    "Vertical dynamic viscosity"
    μv::FT = 1e-4
    "Horizontal dynamic viscosity"
    μh::FT = 1
    "Vertical thermal diffusivity"
    αv::FT = 1e-2
    "Horizontal thermal diffusivity"
    αh::FT = 1
    "IC Gaussian noise standard deviation"
    σ::FT = 5e-2
    "Rayleigh damping"
    γ::FT = 5
    "Domain height"
    zmax::FT = 1
    "Initial conditions for temperature"
    initialT::FT = 295.15
    "Bottom boundary value for temperature (Dirichlet boundary conditions)"
    T_bottom::FT = 300.0
    "Top flux (α∇ρcT) at top boundary (Neumann boundary conditions)"
    flux_top::FT = 0.0
    "Divergence damping coefficient (horizontal)"
    νd::FT = 1
end

# Create an instance of the `BurgersEquation`:
orientation = FlatOrientation()

m = BurgersEquation{FT, typeof(param_set), typeof(orientation)}(
    param_set = param_set,
    orientation = orientation,
);

# This model dictates the flow control, using [Dynamic Multiple
# Dispatch](https://en.wikipedia.org/wiki/Multiple_dispatch), for which
# kernels are executed.

# ## Define the variables

# All of the methods defined in this section were `import`ed in
# [Loading code](@ref Loading-code-burgers) to let us provide
# implementations for our `BurgersEquation` as they will be used
# by the solver.

# Specify auxiliary variables for `BurgersEquation`
function vars_state(m::BurgersEquation, st::Auxiliary, FT)
    @vars begin
        coord::SVector{3, FT}
        orientation::vars_state(m.orientation, st, FT)
    end
end

# Specify prognostic variables, the variables solved for in the PDEs, for
# `BurgersEquation`
vars_state(::BurgersEquation, ::Prognostic, FT) =
    @vars(ρ::FT, ρu::SVector{3, FT}, ρcT::FT);

# Specify state variables whose gradients are needed for `BurgersEquation`
vars_state(::BurgersEquation, ::Gradient, FT) =
    @vars(u::SVector{3, FT}, ρcT::FT, ρu::SVector{3, FT});

# Specify gradient variables for `BurgersEquation`
vars_state(::BurgersEquation, ::GradientFlux, FT) = @vars(
    μ∇u::SMatrix{3, 3, FT, 9},
    α∇ρcT::SVector{3, FT},
    νd∇D::SMatrix{3, 3, FT, 9}
);

# ## Define the compute kernels

# Specify the initial values in `aux::Vars`, which are available in
# `init_state_prognostic!`. Note that
# - this method is only called at `t=0`.
# - `aux.coord` is available here because we've specified `coord` in `vars_state(m, aux, FT)`.
function nodal_init_state_auxiliary!(
    m::BurgersEquation,
    aux::Vars,
    tmp::Vars,
    geom::LocalGeometry,
)
    aux.coord = geom.coord
end;

# `init_aux!` initializes the auxiliary gravitational potential field needed for vertical projections
function init_state_auxiliary!(
    m::BurgersEquation,
    state_auxiliary::MPIStateArray,
    grid,
    direction,
)
    init_aux!(m, m.orientation, state_auxiliary, grid, direction)

    init_state_auxiliary!(
        m,
        nodal_init_state_auxiliary!,
        state_auxiliary,
        grid,
        direction,
    )
end;

# Specify the initial values in `state::Vars`. Note that
# - this method is only called at `t=0`.
# - `state.ρ`, `state.ρu` and`state.ρcT` are available here because we've specified `ρ`, `ρu` and `ρcT` in `vars_state(m, state, FT)`.
function init_state_prognostic!(
    m::BurgersEquation,
    state::Vars,
    aux::Vars,
    localgeo,
    t::Real,
)
    z = aux.coord[3]
    ε1 = rand(Normal(0, m.σ))
    ε2 = rand(Normal(0, m.σ))
    state.ρ = 1
    ρu = 1 - 4 * (z - m.zmax / 2)^2 + ε1
    ρv = 1 - 4 * (z - m.zmax / 2)^2 + ε2
    ρw = 0
    state.ρu = SVector(ρu, ρv, ρw)

    state.ρcT = state.ρ * m.c * m.initialT
end;

# The remaining methods, defined in this section, are called at every
# time-step in the solver by the [`BalanceLaw`](@ref
# ClimateMachine.BalanceLaws.BalanceLaw) framework.

# Since we have second-order fluxes, we must tell `ClimateMachine` to compute
# the gradient of `ρcT`, `u` and `ρu`. Here, we specify how `ρcT`, `u` and `ρu` are computed. Note that
# e.g. `transform.ρcT` is available here because we've specified `ρcT` in `vars_state(m, ::Gradient, FT)`.
function compute_gradient_argument!(
    m::BurgersEquation,
    transform::Vars,
    state::Vars,
    aux::Vars,
    t::Real,
)
    transform.ρcT = state.ρcT
    transform.u = state.ρu / state.ρ
    transform.ρu = state.ρu
end;

# Specify where in `diffusive::Vars` to store the computed gradient from
# `compute_gradient_argument!`. Note that:
#  - `diffusive.μ∇u` is available here because we've specified `μ∇u` in `vars_state(m, ::GradientFlux, FT)`.
#  - `∇transform.u` is available here because we've specified `u` in `vars_state(m, ::Gradient, FT)`.
#  - `diffusive.μ∇u` is built using an anisotropic diffusivity tensor.
#  - The `divergence` may be computed from the trace of tensor `∇ρu`.
function compute_gradient_flux!(
    m::BurgersEquation{FT},
    diffusive::Vars,
    ∇transform::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
) where {FT}
    param_set = parameter_set(m)
    ∇ρu = ∇transform.ρu
    ẑ = vertical_unit_vector(m.orientation, param_set, aux)
    divergence = tr(∇ρu) - ẑ' * ∇ρu * ẑ
    diffusive.α∇ρcT = Diagonal(SVector(m.αh, m.αh, m.αv)) * ∇transform.ρcT
    diffusive.μ∇u = Diagonal(SVector(m.μh, m.μh, m.μv)) * ∇transform.u
    diffusive.νd∇D =
        Diagonal(SVector(m.νd, m.νd, FT(0))) *
        Diagonal(SVector(divergence, divergence, FT(0)))
end;

# Introduce Rayleigh friction towards a target profile as a source.
# Note that:
# - Rayleigh damping is only applied in the horizontal using the `projection_tangential` method.
function source!(
    m::BurgersEquation{FT},
    source::Vars,
    state::Vars,
    diffusive::Vars,
    aux::Vars,
    args...,
) where {FT}
    param_set = parameter_set(m)
    ẑ = vertical_unit_vector(m.orientation, param_set, aux)
    z = aux.coord[3]
    ρ̄ū =
        state.ρ * SVector{3, FT}(
            0.5 - 2 * (z - m.zmax / 2)^2,
            0.5 - 2 * (z - m.zmax / 2)^2,
            0.0,
        )
    ρu_p = state.ρu - ρ̄ū
    source.ρu -=
        m.γ * projection_tangential(m.orientation, param_set, aux, ρu_p)
end;

# Compute advective flux.
# Note that:
# - `state.ρu` is available here because we've specified `ρu` in `vars_state(m, state, FT)`.
function flux_first_order!(
    m::BurgersEquation,
    flux::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
    _...,
)
    flux.ρ = state.ρu

    u = state.ρu / state.ρ
    flux.ρu = state.ρu * u'
    flux.ρcT = u * state.ρcT
end;

# Compute diffusive flux (e.g. ``F(μ, \mathbf{u}, t) = -μ∇\mathbf{u}`` in the original PDE).
# Note that:
# - `diffusive.μ∇u` is available here because we've specified `μ∇u` in `vars_state(m, ::GradientFlux, FT)`.
# - The divergence gradient can be written as a diffusive flux using a divergence diagonal tensor.
function flux_second_order!(
    m::BurgersEquation,
    flux::Grad,
    state::Vars,
    diffusive::Vars,
    hyperdiffusive::Vars,
    aux::Vars,
    t::Real,
)
    flux.ρcT -= diffusive.α∇ρcT
    flux.ρu -= diffusive.μ∇u
    flux.ρu -= diffusive.νd∇D
end;

# ### Boundary conditions

# Second-order terms in our equations, ``∇⋅(G)`` where ``G = μ∇\mathbf{u}``, are
# internally reformulated to first-order unknowns.
# Boundary conditions must be specified for all unknowns, both first-order and
# second-order unknowns which have been reformulated.

struct TopBC <: BoundaryCondition end;
struct BottomBC <: BoundaryCondition end;
boundary_conditions(::BurgersEquation) = (BottomBC(), TopBC());

# The boundary conditions for `ρ`, `ρu` and `ρcT` (first order unknowns)
function boundary_state!(
    nf,
    bc::BottomBC,
    m::BurgersEquation,
    state⁺::Vars,
    aux⁺::Vars,
    n⁻,
    _...,
)
    state⁺.ρ = 1
    state⁺.ρu = SVector(0, 0, 0)
    state⁺.ρcT = state⁺.ρ * m.c * m.T_bottom
end;
function boundary_state!(
    nf,
    bc::TopBC,
    m::BurgersEquation,
    state⁺::Vars,
    aux⁺::Vars,
    n⁻,
    _...,
)
    state⁺.ρ = 1
    state⁺.ρu = SVector(0, 0, 0)
end;

# The boundary conditions for `ρ`, `ρu` and `ρcT` are specified here for
# second-order unknowns

function boundary_state!(
    nf,
    bc::BottomBC,
    m::BurgersEquation,
    state⁺::Vars,
    diff⁺::Vars,
    hyperdiff⁺::Vars,
    aux⁺::Vars,
    n⁻,
    _...,
)
    state⁺.ρ = 1
    state⁺.ρu = SVector(0, 0, 0)
    state⁺.ρcT = state⁺.ρ * m.c * m.T_bottom
end;

function boundary_state!(
    nf,
    bc::TopBC,
    m::BurgersEquation,
    state⁺::Vars,
    diff⁺::Vars,
    hyperdiff⁺::Vars,
    aux⁺::Vars,
    n⁻,
    _...,
)
    state⁺.ρ = 1
    state⁺.ρu = SVector(0, 0, 0)
    diff⁺.α∇ρcT = -n⁻ * m.flux_top
end;

# # Spatial discretization

# Prescribe polynomial order of basis functions in finite elements
N_poly = 5;

# Specify the number of vertical elements
nelem_vert = 10;

# Specify the domain height
zmax = m.zmax;

# # Temporal discretization

# This initializes the ODE
# solver, the horizontal explicit vertical implicit scheme
# with ARK2GiraldoKellyConstantinescu method.

ode_solver_type = ClimateMachine.HEVISolverType(
    FT;
    solver_method = ARK2GiraldoKellyConstantinescu,
    linear_max_subspace_size = Int(30),
    linear_atol = FT(-1.0),
    linear_rtol = FT(1e-5),
    nonlinear_max_iterations = Int(10),
    nonlinear_rtol = FT(1e-4),
    nonlinear_ϵ = FT(1.e-10),
    preconditioner_update_freq = Int(50),
)

# Establish a `ClimateMachine` single stack configuration
driver_config = ClimateMachine.SingleStackConfiguration(
    "BurgersEquation",
    N_poly,
    nelem_vert,
    zmax,
    param_set,
    m,
    numerical_flux_first_order = CentralNumericalFluxFirstOrder(),
);

# # Time discretization

# Specify simulation time (SI units)
t0 = FT(0);
timeend = FT(1);

# We'll define the time-step based on the Fourier
# number and the [Courant number](https://en.wikipedia.org/wiki/Courant–Friedrichs–Lewy_condition)
# of the flow
Δ = min_node_distance(driver_config.grid)

given_Fourier = FT(0.5);
Fourier_bound = given_Fourier * Δ^2 / max(m.αh, m.μh, m.νd);
Courant_bound = FT(0.5) * Δ;

# We define the time step 50 times larger than that defined by CFL law
dt = FT(50.0) * min(Fourier_bound, Courant_bound)

# # Configure a `ClimateMachine` solver.

# This initializes the state vector and allocates memory for the solution in
# space (`dg` has the model `m`, which describes the PDEs as well as the
# function used for initialization).

solver_config = ClimateMachine.SolverConfiguration(
    t0,
    timeend,
    driver_config,
    ode_dt = dt,
    ode_solver_type = ode_solver_type,
);

# ## Inspect the initial conditions for a single nodal stack

# Let's export plots of the initial state
output_dir = @__DIR__;

mkpath(output_dir);

z_scale = 100 # convert from meters to cm
z_key = "z"
z_label = "z [cm]"
z = get_z(driver_config.grid; z_scale = z_scale)
state_vars = get_vars_from_nodal_stack(
    driver_config.grid,
    solver_config.Q,
    vars_state(m, Prognostic(), FT),
);

# Create an array to store the solution:
state_data = Dict[state_vars]  # store initial condition at ``t=0``
time_data = FT[0]                                      # store time data

# Generate plots of initial conditions for the southwest nodal stack
export_plot(
    z,
    time_data,
    state_data,
    ("ρcT",),
    joinpath(output_dir, "initial_condition_T_nodal_bjfnk.png");
    xlabel = "ρcT at southwest node",
    ylabel = z_label,
);
export_plot(
    z,
    time_data,
    state_data,
    ("ρu[1]",),
    joinpath(output_dir, "initial_condition_u_nodal_bjfnk.png");
    xlabel = "ρu at southwest node",
    ylabel = z_label,
);
export_plot(
    z,
    time_data,
    state_data,
    ("ρu[2]",),
    joinpath(output_dir, "initial_condition_v_nodal_bjfnk.png");
    xlabel = "ρv at southwest node",
    ylabel = z_label,
);

# ![](initial_condition_T_nodal_bjfnk.png)
# ![](initial_condition_u_nodal_bjfnk.png)

# ## Inspect the initial conditions for the horizontal averages

# Horizontal statistics of variables

state_vars_var = get_horizontal_variance(
    driver_config.grid,
    solver_config.Q,
    vars_state(m, Prognostic(), FT),
);

state_vars_avg = get_horizontal_mean(
    driver_config.grid,
    solver_config.Q,
    vars_state(m, Prognostic(), FT),
);

data_avg = Dict[state_vars_avg]
data_var = Dict[state_vars_var]

export_plot(
    z,
    time_data,
    data_avg,
    ("ρu[1]",),
    joinpath(output_dir, "initial_condition_avg_u_bjfnk.png");
    xlabel = "Horizontal mean of ρu",
    ylabel = z_label,
);
export_plot(
    z,
    time_data,
    data_var,
    ("ρu[1]",),
    joinpath(output_dir, "initial_condition_variance_u_bjfnk.png");
    xlabel = "Horizontal variance of ρu",
    ylabel = z_label,
);

# ![](initial_condition_avg_u_bjfnk.png)
# ![](initial_condition_variance_u_bjfnk.png)

# # Solver hooks / callbacks

# Define the number of outputs from `t0` to `timeend`
const n_outputs = 5;
const every_x_simulation_time = timeend / n_outputs;

# Create a dictionary for `z` coordinate (and convert to cm) NCDatasets IO:
dims = OrderedDict(z_key => collect(z));

# Create dictionaries to store outputs:
data_var = Dict[Dict([k => Dict() for k in 0:n_outputs]...),]
data_var[1] = state_vars_var

data_avg = Dict[Dict([k => Dict() for k in 0:n_outputs]...),]
data_avg[1] = state_vars_avg

data_nodal = Dict[Dict([k => Dict() for k in 0:n_outputs]...),]
data_nodal[1] = state_vars

# The `ClimateMachine`'s time-steppers provide hooks, or callbacks, which
# allow users to inject code to be executed at specified intervals. In this
# callback, the state variables are collected, combined into a single
# `OrderedDict` and written to a NetCDF file (for each output step `step`).
step = [0];
callback = GenericCallbacks.EveryXSimulationTime(every_x_simulation_time) do
    state_vars_var = get_horizontal_variance(
        driver_config.grid,
        solver_config.Q,
        vars_state(m, Prognostic(), FT),
    )
    state_vars_avg = get_horizontal_mean(
        driver_config.grid,
        solver_config.Q,
        vars_state(m, Prognostic(), FT),
    )
    state_vars = get_vars_from_nodal_stack(
        driver_config.grid,
        solver_config.Q,
        vars_state(m, Prognostic(), FT),
        i = 1,
        j = 1,
    )
    step[1] += 1
    push!(data_var, state_vars_var)
    push!(data_avg, state_vars_avg)
    push!(data_nodal, state_vars)
    push!(time_data, gettime(solver_config.solver))
    nothing
end;

# # Solve

# This is the main `ClimateMachine` solver invocation. While users do not have
# access to the time-stepping loop, code may be injected via `user_callbacks`,
# which is a `Tuple` of [`GenericCallbacks`](@ref ClimateMachine.GenericCallbacks).
ClimateMachine.invoke!(solver_config; user_callbacks = (callback,))

# # Post-processing

# Our solution has now been calculated and exported to NetCDF files in
# `output_dir`.

# Let's plot the horizontal statistics of `ρu` and `ρcT`, as well as the evolution of
# `ρu` for the southwest nodal stack:
export_plot(
    z,
    time_data,
    data_avg,
    ("ρu[1]"),
    joinpath(output_dir, "solution_vs_time_u_avg_bjfnk.png");
    xlabel = "Horizontal mean of ρu",
    ylabel = z_label,
);
export_plot(
    z,
    time_data,
    data_var,
    ("ρu[1]"),
    joinpath(output_dir, "variance_vs_time_u_bjfnk.png");
    xlabel = "Horizontal variance of ρu",
    ylabel = z_label,
);
export_plot(
    z,
    time_data,
    data_avg,
    ("ρcT"),
    joinpath(output_dir, "solution_vs_time_T_avg_bjfnk.png");
    xlabel = "Horizontal mean of ρcT",
    ylabel = z_label,
);
export_plot(
    z,
    time_data,
    data_var,
    ("ρcT"),
    joinpath(output_dir, "variance_vs_time_T_bjfnk.png");
    xlabel = "Horizontal variance of ρcT",
    ylabel = z_label,
);
export_plot(
    z,
    time_data,
    data_nodal,
    ("ρu[1]"),
    joinpath(output_dir, "solution_vs_time_u_nodal_bjfnk.png");
    xlabel = "ρu at southwest node",
    ylabel = z_label,
);
# ![](solution_vs_time_u_avg_bjfnk.png)
# ![](variance_vs_time_u_bjfnk.png)
# ![](solution_vs_time_T_avg_bjfnk.png)
# ![](variance_vs_time_T_bjfnk.png)
# ![](solution_vs_time_u_nodal_bjfnk.png)

# Rayleigh friction returns the horizontal velocity to the objective
# profile on the timescale of the simulation (1 second), since `γ`∼1. The horizontal viscosity
# and 2D divergence damping act to reduce the horizontal variance over the same timescale.
# The initial Gaussian noise is propagated to the temperature field through advection.
# The horizontal diffusivity acts to reduce this `ρcT` variance in time, although in a longer
# timescale.

# To run this file, and
# inspect the solution, include this tutorial in the Julia REPL
# with:

# ```julia
# include(joinpath("tutorials", "Atmos", "burgers_single_stack.jl"))
# ```


using Distributions
import ClusteringGA: fitness

function fruspini()
    ruspini = [
        4	53;
        5	63;
        10	59;
        9	77;
        13	49;
        13	69;
        12	88;
        15	75;
        18	61;
        19	65;
        22	74;
        27	72;
        28	76;
        24	58;
        27	55;
        28	60;
        30	52;
        31	60;
        32	61;
        36	72;
        28	147;
        32	149;
        35	153;
        33	154;
        38	151;
        41	150;
        38	145;
        38	143;
        32	143;
        34	141;
        44	156;
        44	149;
        44	143;
        46	142;
        47	149;
        49	152;
        50	142;
        53	144;
        52	152;
        55	155;
        54	124;
        60	136;
        63	139;
        86	132;
        85	115;
        85	96;
        78	94;
        74	96;
        97	122;
        98	116;
        98	124;
        99	119;
        99	128;
        101	115;
        108	111;
        110	111;
        108	116;
        111	126;
        115	117;
        117	115;
        70	4;
        77	12;
        83	21;
        61	15;
        69	15;
        78	16;
        66	18;
        58	13;
        64	20;
        69	21;
        66	23;
        61	25;
        76	27;
        72	31;
        64	30
    ]
    return [ruspini[i, :] for i = 1:lastindex(ruspini, 1)]
end

function rand200()
    nd1 = MvNormal([0, 10], [1.7, 1.7])
    nd2 = MvNormal([20, 12], [0.7, 0.7])
    nd3 = MvNormal([10, 20], [1.0, 1.0])
    s = hcat(rand(nd1, 120), rand(nd2, 60), rand(nd3, 20))
    t = s'
    tobj = [t[i, :] for i = 1:lastindex(t, 1)]
    return tobj
end

@testset "Clustering Genetic Algorithms" begin
    d, r = cga(fruspini())
    @test counts(r) == [20, 23, 17, 15]
    @test 1.73 < fitness(assignments(r), d) < 1.74
    d, r = cga(rand200())
    @test counts(r) == [120, 60, 20]
    @test 1.75 < fitness(assignments(r), d) < 1.85
end