using Parameters

abstract type AbstractThermo end
export AbstractThermo

@inline function getGibbs(th::P,T::N) where {N<:Number,P<:AbstractThermo}
    return @fastmath getEnthalpy(th,T)-T*getEntropy(th,T)
end

@with_kw struct NASApolynomial <: AbstractThermo
    coefs::Array{Float64,1}
    Tmin::Float64
    Tmax::Float64
end
export NASApolynomial

@inline function calcHSCpdless(poly::NASApolynomial,T::X) where {X<:Real}
    if length(poly.coefs) != 7
        Tpoly0 = T
        Tpoly1 = T*T
        Tpoly2 = Tpoly1*T
        Tpoly3 = Tpoly2*T
        Tpoly4 = 1.0/T
        Tpoly5 = Tpoly4*Tpoly4
        Tpoly6 = log(T)

        ct0 = poly.coefs[1]*Tpoly5
        ct1 = poly.coefs[2]*Tpoly4
        ct2 = poly.coefs[3]
        ct3 = poly.coefs[4]*Tpoly0
        ct4 = poly.coefs[5]*Tpoly1
        ct5 = poly.coefs[6]*Tpoly2
        ct6 = poly.coefs[7]*Tpoly3

        cpdivR = ct0 + ct1 + ct2 + ct3 + ct4 + ct5 + ct6

        hdivRT = -ct0+Tpoly6*ct1+ct2+0.5*ct3+0.33333333333*ct4+0.25*ct5+0.2*ct6+poly.coefs[8]*Tpoly4
        sdivR = -0.5*ct0 - ct1 + Tpoly6*ct2 + ct3 + 0.5*ct4 + 0.33333333333*ct5 + 0.25*ct6 + poly.coefs[9]
    else
        Tpoly0 = T
        Tpoly1 = T*T
        Tpoly2 = Tpoly1*T
        Tpoly3 = Tpoly2*T
        Tpoly4 = 1.0/T
        #Tpoly5 = Tpoly4*Tpoly4
        Tpoly6 = log(T)

        #ct0 = 0.0
        #ct1 = 0.0
        ct2 = poly.coefs[1]
        ct3 = poly.coefs[2]*Tpoly0
        ct4 = poly.coefs[3]*Tpoly1
        ct5 = poly.coefs[4]*Tpoly2
        ct6 = poly.coefs[5]*Tpoly3

        cpdivR = ct2 + ct3 + ct4 + ct5 + ct6

        hdivRT =ct2+0.5*ct3+0.33333333333*ct4+0.25*ct5+0.2*ct6+poly.coefs[6]*Tpoly4
        sdivR = Tpoly6*ct2 + ct3 + 0.5*ct4 + 0.33333333333*ct5 + 0.25*ct6 + poly.coefs[7]
    end
    return (cpdivR,hdivRT,sdivR)::Tuple{Float64,Float64,Float64}
end
export calcHSCpdless

@inline function getHeatCapacity(poly::NASApolynomial,T::N) where {N<:Number}
    if length(poly.coefs) == 9
        return @views @inbounds @fastmath evalpoly(T,poly.coefs[1:7])/T^2*R
    elseif length(poly.coefs) == 7
        return @views @inbounds @fastmath evalpoly(T,poly.coefs[1:5])*R
    else
        throw(error("NASA polynomial has a number of coefficients not equal to 9 or 7"))
    end
end

@inline function getEntropy(poly::NASApolynomial,T::N) where {N<:Number}
    if length(poly.coefs) == 9
        return @views @inbounds @fastmath ((-poly.coefs[1]/(2*T)-poly.coefs[2])/T+poly.coefs[3]*log(T)+T*evalpoly(T,poly.coefs[4:end-2]./(1:4))+poly.coefs[end])*R
    elseif length(poly.coefs) == 7
        return @views @inbounds @fastmath (poly.coefs[1]*log(T)+T*evalpoly(T,poly.coefs[2:end-2]./(1:4))+poly.coefs[end])*R
    else
        throw(error("NASA polynomial has a number of coefficients not equal to 9 or 7"))
    end
end

@inline function getEnthalpy(poly::NASApolynomial,T::N) where {N<:Number}
    if length(poly.coefs) == 9
        return @views @inbounds @fastmath ((-poly.coefs[1]/T+poly.coefs[2]*log(T))/T+evalpoly(T,poly.coefs[3:end-2]./(1:5)))*R*T+poly.coefs[end-1]*R
    elseif length(poly.coefs) == 7
        return @views @inbounds @fastmath evalpoly(T,poly.coefs[1:end-2]./(1:5))*R*T+poly.coefs[end-1]*R
    else
        throw(error("NASA polynomial has a number of coefficients not equal to 9 or 7"))
    end
end

@with_kw struct NASA{T<:AbstractThermoUncertainty} <: AbstractThermo
    polys::Array{NASApolynomial,1}
    unc::T = EmptyThermoUncertainty()
end
export NASA

@inline function selectPoly(nasa::NASA,T::N) where {N<:Real}
    """
    retrieve the nasa polynomial corresponding to the T range
    """
    for p in nasa.polys
        if T<=p.Tmax
            return p
        end
    end
    return nasa.polys[end]
end
export selectPoly

@inline getHeatCapacity(nasa::NASA,T::N) where {N<:Number} = getHeatCapacity(selectPoly(nasa,T),T)
@inline getEntropy(nasa::NASA,T::N) where {N<:Number} = getEntropy(selectPoly(nasa,T),T)
@inline getEnthalpy(nasa::NASA,T::N) where {N<:Number} = getEnthalpy(selectPoly(nasa,T),T)
@inline getGibbs(nasa::NASA,T::N) where {N<:Number} = getGibbs(selectPoly(nasa,T),T)
@inline calcHSCpdless(nasa::NASA,T::N) where {N<:Real}= calcHSCpdless(selectPoly(nasa,T),T)

@with_kw struct Wilhoit{N,Q,T,P,U,R<:Number,M<:AbstractThermoUncertainty} <: AbstractThermo
    Cp0::N
    Cpinf::T
    coefs::Array{Q,1}
    H0::P
    S0::U
    B::R
    unc::M = EmptyThermoUncertainty()
end
export Wilhoit

@inline function getHeatCapacity(w::Wilhoit,T::N) where {N<:Number}
    @fastmath y = T/(T+w.B)
    return @fastmath w.Cp0 + (w.Cpinf-w.Cp0)*y^2*(1+(y-1)*evalpoly(y,w.coefs))
end

@inline function getEnthalpy(w::Wilhoit,T::N) where {N<:Number}
    @fastmath y = T/(T+w.B)
    return @views @fastmath @inbounds w.H0 + w.Cp0 * T - (w.Cpinf - w.Cp0) * T * (
            y * y * ((3 * w.coefs[1] + sum(w.coefs[2:end])) / 6. +
                     (4 * w.coefs[2] + sum(w.coefs[3:end])) * y / 12. +
                     (5 * w.coefs[3] + w.coefs[4]) * y^2 / 20. +
                     w.coefs[4] * y^3 / 5.) +
            (2 + sum(w.coefs)) * (y / 2. - 1 + (1.0 / y - 1.) * log(w.B + T))
        )
end

@inline function getEntropy(w::Wilhoit,T::N) where {N<:Number}
    @fastmath y = T/(T+w.B)
    return @fastmath w.S0 + w.Cpinf*log(T)-(w.Cpinf-w.Cp0)*(log(y)+y*(1+y*evalpoly(y,w.coefs./(2:5))))
end

export getGibbs, getEntropy, getEnthalpy, getHeatCapacity


# Practical Machine learning
# Bayesian learning - Naive Bayes example
# Chapter 9

using NaiveBayes
using RDatasets

iris = dataset("datasets", "iris")

# observations in columns and variables in rows
X = array(iris[:, 1:4])'
p, n = size(X)
# by default species is a PooledDataArray,
y = [species for species in iris[:, 5]]

# how much data use for training
train_frac = 0.9
k = int(floor(train_frac * n))
idxs = randperm(n)
train = idxs[1:k]
test = idxs[k+1:end]

model = GaussianNB(unique(y), p)
fit(model, X[:, train], y[train])

accuracy = countnz(predict(model, X[:,test]) .== y[test]) / countnz(test)

println("Accuracy: $accuracy")


module LloydsMirror
function lloydsmirror_singlereflection(c, f, r_src, r_hyd, z_src, z_hyd)
	λ = c/f
	k = 2π/λ

	s_dir = sqrt((r_src - r_hyd)^2 + (z_src - z_hyd)^2)
	s_rfl = sqrt((r_src - r_hyd)^2 + (z_src + z_hyd)^2)

	pressure(s) = exp(im*k*s)/s
	p_dir = pressure(s_dir)
	p_rfl = -pressure(s_rfl)

	p = p_dir + p_rfl
	TL = -20log10(abs(p))
end
end


function (::Type{QuadratureRule{2,RefTetrahedron}})(quad_type::Strang, order::Int)
    if order == 0 || order == 1
        points = [Vec{2}([1.0/3.0, 1.0/3.0])]
        weigths = [0.5]
    elseif order == 2
        points = [Vec{2}([1.0/6.0, 1.0/6.0]),
                  Vec{2}([1.0/6.0, 2.0/3.0]),
                  Vec{2}([2.0/3.0, 1.0/6.0])]
        weigths = [1/6,1/6,1/6]
    elseif order == 3
        points = [Vec{2}([0.659027622374092, 0.231933368553031]),
                  Vec{2}([0.659027622374092, 0.109039009072877]),
                  Vec{2}([0.231933368553031, 0.659027622374092]),
                  Vec{2}([0.231933368553031, 0.109039009072877]),
                  Vec{2}([0.109039009072877, 0.659027622374092]),
                  Vec{2}([0.109039009072877, 0.231933368553031])]
        weigths = 1/12*ones(Float64,6)
    elseif order == 4
        points = [Vec{2}([0.816847572980459, 0.091576213509771]),
                  Vec{2}([0.091576213509771, 0.816847572980459]),
                  Vec{2}([0.091576213509771, 0.091576213509771]),
                  Vec{2}([0.108103018168070, 0.445948490915965]),
                  Vec{2}([0.445948490915965, 0.108103018168070]),
                  Vec{2}([0.445948490915965, 0.445948490915965])]
        weigths = 1/2*ones(Float64,6)
        weigths[1:3] = 0.109951743655322*weigths[1:3]
        weigths[4:6] = 0.223381589678011*weigths[4:6]
    elseif order == 5
        points = [Vec{2}([0.33333333333333333, 0.33333333333333333]),
                  Vec{2}([0.79742698535308720, 0.10128650732345633]),
                  Vec{2}([0.10128650732345633, 0.79742698535308720]),
                  Vec{2}([0.10128650732345633, 0.10128650732345633]),
                  Vec{2}([0.05971587178976981, 0.47014206410511505]),
                  Vec{2}([0.47014206410511505, 0.05971587178976981]),
                  Vec{2}([0.47014206410511505, 0.47014206410511505])]
        weigths = 1/2*ones(Float64,7)
        weigths[1] = 0.22500000000000000*weigths[1]
        weigths[2:4] = 0.12593918054482717*weigths[2:4]
        weigths[5:7] = 0.13239415278850616*weigths[5:7]
    elseif order == 6
        points =[Vec{2}([0.873821971016996, 0.063089014491502]),
                 Vec{2}([0.063089014491502, 0.873821971016996]),
                 Vec{2}([0.063089014491502, 0.063089014491502]),
                 Vec{2}([0.501426509658179, 0.249286745170910]),
                 Vec{2}([0.249286745170910, 0.501426509658179]),
                 Vec{2}([0.249286745170910, 0.249286745170910]),
                 Vec{2}([0.636502499121399, 0.310352451033785]),
                 Vec{2}([0.636502499121399, 0.053145049844816]),
                 Vec{2}([0.310352451033785, 0.636502499121399]),
                 Vec{2}([0.310352451033785, 0.053145049844816]),
                 Vec{2}([0.053145049844816, 0.636502499121399]),
                 Vec{2}([0.053145049844816, 0.310352451033785])]
        weigths = 1/2*ones(Float64,12)
        weigths[1:3] = 0.050844906370207*weigths[1:3]
        weigths[4:6] = 0.116786275726379*weigths[4:6]
        weigths[7:12] = 0.082851075618374*weigths[7:12]
    else
        throw(ArgumentError("Strang rule of order $order not available"))
    end
    return QuadratureRule{2,RefTetrahedron,Float64}(weigths, points)
end


function evalPsiFromBEM(BEMesh,quadRule,f,q,xr,xz; derivative=0)
#=
Given a BE contour BEMesh, a quadrature rule on this contour, the Dirichlet P1 dof and the Neumann P0 dof,
evaluate psi at any point (xr,xz) outside the contour resulting from BEM
q is a vector of size n edges
f is a vector of size n nodes (the BEMesh contour is closed)
=#

	#constant
	mu0=4*pi*10^-7; mu0inv=1/mu0

	# dimension
	n=size(BEMesh.edges,1)

	#quadrature rule
	xGauss=quadRule.x
	wGauss=quadRule.w
	nQuad=size(quadRule.w,1)

	#in order to remap the quadrature we need the midpoints and length of edges
	e=collect(1:n)
	edgeToNode=BEMesh.edges
	h=BEMesh.length

    P=0.5*(BEMesh.coordinates[edgeToNode[:,1],:]+BEMesh.coordinates[edgeToNode[:,2],:])
	#@show P

	#normal
	normal=BEMesh.normalsIn

	#quadrature points on all edges
	PG=zeros(size(P,1),2,nQuad)
	for i in 1:nQuad
		PG[:,:,i]=P + xGauss[i]*0.5*(BEMesh.coordinates[edgeToNode[:,2],:]-BEMesh.coordinates[edgeToNode[:,1],:])
	end

	#evaluate Green function at all y=Gauss points PG and x=(xr,xz)
	#need to broadcast on the number of edges

	DE=zeros(n,2); DI=zeros(n,2); DJ=zeros(n,2)
	NE=zeros(n); NI=zeros(n); NJ=zeros(n)

	dxrDE=zeros(n,2);
	dxzDE=zeros(n,2);
	dxrNE=zeros(n)
	dxzNE=zeros(n)

	for i in 1:nQuad
		#Evaluate Green function and first derivative (in y=B)
		G = @. Green(xr,xz,PG[:,1,i],PG[:,2,i],derivative=derivative+1)
		#G is a vector of named tuples, we need the broadcasted getindex function to access a value (with .() or @. see below)
		#@show G

		#Green function value
		G_PXPG = getindex.(G,:G_xy)
		#@show G_PXPG

		#Neumann trace
		dnG_PXPG = @.getindex(G,:dyrG_xy) .* normal[:,1] +  @.getindex(G,:dyzG_xy) .* normal[:,2]
		#@show dnG_P1GP2G

		# Matrices. Matrix.vector product to be summed afterwards for integral values
		# for psi(x)
		NE=NE+0.5*wGauss[i]*h.*G_PXPG

		#DE=DE-0.5*wGauss[i]*[ h.*dnG_PXPG*(1-(0.5+xGauss[i]*0.5))  h.*dnG_PXPG*((0.5+xGauss[i]*0.5)) ]
		DE=DE-0.5*wGauss[i]*mu0inv*[ h.*(1.0./PG[:,1,i]).*dnG_PXPG*(1-(0.5+xGauss[i]*0.5))  h.*(1.0./PG[:,1,i]).*dnG_PXPG*((0.5+xGauss[i]*0.5)) ]

		# for d_xr_psi(x) and d_xz_psi(x)
		if(derivative>0)
			dxrG= @.getindex(G,:dxrG_xy);
			dxzG= @.getindex(G,:dxzG_xy);
			dxrNE=dxrNE+0.5*wGauss[i]*h.*dxrG
			dxzNE=dxzNE+0.5*wGauss[i]*h.*dxzG

			dxrdyn = @.getindex(G,:dxrdyrG_xy) .* normal[:,1] + @.getindex(G,:dxrdyzG_xy) .* normal[:,2]
			dxzdyn = @.getindex(G,:dxzdyrG_xy) .* normal[:,1] + @.getindex(G,:dxzdyzG_xy) .* normal[:,2]

			dxrDE=dxrDE-0.5*wGauss[i]*mu0inv*[ h.*(1.0./PG[:,1,i]).*dxrdyn*(1-(0.5+xGauss[i]*0.5))  h.*(1.0./PG[:,1,i]).*dxrdyn*((0.5+xGauss[i]*0.5)) ]
			dxzDE=dxzDE-0.5*wGauss[i]*mu0inv*[ h.*(1.0./PG[:,1,i]).*dxzdyn*(1-(0.5+xGauss[i]*0.5))  h.*(1.0./PG[:,1,i]).*dxzdyn*((0.5+xGauss[i]*0.5)) ]
		end

	end

	NI=e; NJ=e
	DI=[e e]; DJ=edgeToNode

	DI=reshape(DI,size(DI,1)*size(DI,2))
	DJ=reshape(DJ,size(DJ,1)*size(DJ,2))
	DE=reshape(DE,size(DE,1)*size(DE,2))

	Dspmat=sparse(DI,DJ,DE,n,n)
	Nspmat=sparse(NI,NJ,NE,n,n)
	Ival=Dspmat*f+Nspmat*q
	Ival=sum(Ival)

	if(derivative==0)
		return (psi=Ival,)
	end

	if(derivative>0)
		dxrDE=reshape(dxrDE,size(dxrDE,1)*size(dxrDE,2))
		dxzDE=reshape(dxzDE,size(dxzDE,1)*size(dxzDE,2))
		dxrDspmat=sparse(DI,DJ,dxrDE,n,n)
		dxzDspmat=sparse(DI,DJ,dxzDE,n,n)
		dxrNspmat=sparse(NI,NJ,dxrNE,n,n)
		dxzNspmat=sparse(NI,NJ,dxzNE,n,n)

		dxrIval=dxrDspmat*f+dxrNspmat*q
		dxrIval=sum(dxrIval)

		dxzIval=dxzDspmat*f+dxzNspmat*q
		dxzIval=sum(dxzIval)

		return (psi=Ival, drpsi=dxrIval, dzpsi=dxzIval)
	end

end


module JointPosteriors

using   LineSearches: BackTracking
using   Optim
using   ForwardDiff
using   SparseQuadratureGrids
using   ConstrainedParameters
using   LogDensities
using   Interpolations
using   Distributions

import  Base: show, quantile, Val
import  LogDensities: log_density
import  SparseQuadratureGrids: SlidingVecFun

export  JointPosterior,
        JointPosteriorRaw,
        fit,
        marginal,
        Data,
        parameter,
        CovarianceMatrix,
        PositiveVector,
        ProbabilityVector,
        RealVector,
        Simplex,
        Model,
        construct,
        log_density,
        log_density!,
        log_jacobian!,
        quad_form,
        inv_det,
        inv_root_det,
        root_det,
        log_root_det,
        trace_inverse,
        lpdf_InverseWishart,
        lpdf_normal,
        Normal,
        Adaptive,
        AdaptiveRaw,
        Smolyak,
        pdf,
        cdf,
        quantile

include("joint_posterior.jl")
include("interp.jl")
include("marginal_posterior.jl")

end # module


# Load X and y variable
using JLD
data = load("outliersData.jld")
(X,y,Xtest,ytest) = (data["X"],data["y"],data["Xtest"],data["ytest"])

# Fit a least squares model
include("leastMax.jl")
model = leastMax(X,y)

# Evaluate training error
yhat = model.predict(X)
trainError = mean((yhat - y).^2)
@printf("Squared train Error with least squares: %.3f\n",trainError)

# Evaluate test error
yhat = model.predict(Xtest)
testError = mean((yhat - ytest).^2)
@printf("Squared test Error with least squares: %.3f\n",testError)

# Plot model
using PyPlot
figure()
plot(X,y,"b.")
Xhat = minimum(X):.01:maximum(X)
yhat = model.predict(Xhat)
plot(Xhat,yhat,"g")


using CayleyDickson
using Base.Test: @test, @test_throws

@test begin
    a = CayleyDickson.randomBigInt()
    b = CayleyDickson.randomBigInt()
    l = Exo1Quaternion(Exo1Real(a, b))
    r = Exo1Quaternion(a, 0, 0, 0, b, 0, 0, 0)
    l == r
end

@test begin
    a = CayleyDickson.randomBigFloat()
    isreal(Exo1Quaternion(a))
end

@test begin
    a = 1
    !isreal(Exo1Quaternion(a, a, a, a))
end

@test begin
    z = random(Exo1Quaternion{Int})
    z == +(z)
end

@test begin
    a = CayleyDickson.randomBigFloat()
    real(Exo1Quaternion(a)) == a
end

@test begin
    io = IOBuffer()
    show(io, Exo1Quaternion(1.0, 2, 3, 4, 5, 6, 7, 8))
    l = String(take!(io))
    r = "[1: 1.0, i: 2.0, j: 3.0, ij: 4.0, a: 5.0, ia: 6.0, ja: 7.0, (ij)a: 8.0]"
    l == r
end

@test begin
    io = IOBuffer()
    show(io, Exo1Quaternion(1.0, 2, 3, 4, 5, 6, 7))
    l = String(take!(io))
    r = "[1: 1.0, i: 2.0, j: 3.0, ij: 4.0, a: 5.0, ia: 6.0, ja: 7.0, (ij)a: 0.0]"
    l == r
end

@test begin
    io = IOBuffer()
    show(io, Exo1Quaternion(1.0, 2, 3, 4, 5, 6))
    l = String(take!(io))
    r = "[1: 1.0, i: 2.0, j: 3.0, ij: 4.0, a: 5.0, ia: 6.0, ja: 0.0, (ij)a: 0.0]"
    l == r
end

@test begin
    io = IOBuffer()
    show(io, Exo1Quaternion(1.0, 2, 3, 4, 5))
    l = String(take!(io))
    r = "[1: 1.0, i: 2.0, j: 3.0, ij: 4.0, a: 5.0, ia: 0.0, ja: 0.0, (ij)a: 0.0]"
    l == r
end

@test begin
    io = IOBuffer()
    show(io, Exo1Quaternion(1.0, 2, 3, 4))
    l = String(take!(io))
    r = "[1: 1.0, i: 2.0, j: 3.0, ij: 4.0, a: 0.0, ia: 0.0, ja: 0.0, (ij)a: 0.0]"
    l == r
end

@test begin
    a = CayleyDickson.randomBigInt()
    b = CayleyDickson.randomBigFloat()
    c = CayleyDickson.randomBigInt()
    d = CayleyDickson.randomBigFloat()
    l = Exo1Quaternion(Quaternion(a, b, c, d))
    r = Exo1Quaternion(a, b, c, d)
    l == r
end

@test begin
    a = CayleyDickson.randomBigInt()
    b = CayleyDickson.randomBigFloat()
    l = Exo1Quaternion(Binion(a, b))
    r = Exo1Quaternion(a, b)
    l == r
end

@test begin
    a = CayleyDickson.randomBigInt()
    b = CayleyDickson.randomBigFloat()
    c = CayleyDickson.randomBigInt()
    l = Exo1Quaternion(a, b, c)
    r = Exo1Quaternion(a, b, c, 0)
    l == r
end

@test begin
    x = random(Exo1Quaternion{BigFloat})
    y = random(Exo1Quaternion{BigFloat})
    x != y
end

@test begin
    a = CayleyDickson.randomBigInt()
    z = random(Exo1Quaternion{BigInt})
    l = a + z
    r = z + a
    l == r
end

@test begin
    a = random(Quaternion{BigInt})
    z = random(Exo1Quaternion{BigInt})
    l = a + z
    r = z + a
    l == r
end

@test begin
    a = CayleyDickson.randomBigInt()
    z = random(Exo1Quaternion{BigInt})
    l = a - z
    r = -(z - a)
    l == r
end

@test begin
    a = random(Quaternion{BigInt})
    z = random(Exo1Quaternion{BigInt})
    l = a - z
    r = -(z - a)
    l == r
end

@test begin
    x = random(Exo1Quaternion{BigInt})
    y = random(Exo1Quaternion{BigInt})
    l = x + y
    r = y + x
    l == r
end

@test begin
    x = random(Exo1Quaternion{BigInt})
    y = random(Exo1Quaternion{BigInt})
    l = x * y
    r = y * x
    l != r
end

@test begin
    x = random(Exo1Quaternion{BigInt})
    y = random(Exo1Quaternion{BigInt})
    l = x - y
    r = -(y - x)
    l == r
end

@test begin
    x = random(Exo1Quaternion{BigInt})
    y = random(Exo1Quaternion{BigInt})
    z = random(Exo1Quaternion{BigInt})
    l = (x + y) + z
    r = x + (y + z)
    l == r
end

@test begin
    x = random(Exo1Quaternion{BigInt})
    y = random(Exo1Quaternion{BigInt})
    z = random(Exo1Quaternion{BigInt})
    l = (x * y) * z
    r = x * (y * z)
    l != r
end

@test begin
    x = random(Exo1Quaternion{BigInt})
    y = random(Exo1Quaternion{BigInt})
    l = alternatorL(x, y)
    r = zero(Exo1Quaternion{BigInt})
    l == r
end

@test begin
    x = random(Exo1Quaternion{BigInt})
    y = random(Exo1Quaternion{BigInt})
    l = alternatorR(x, y)
    r = zero(Exo1Quaternion{BigInt})
    l == r
end

@test begin
    x = random(Exo1Quaternion{BigInt})
    conj(conj(x)) == x
end

@test begin
    x = random(Exo1Quaternion{Rational{BigInt}})
    inv(inv(x)) == x
end

@test begin
    x = random(Exo1Quaternion{BigInt})
    y = random(Exo1Quaternion{BigInt})
    l = conj(x * y)
    r = conj(y) * conj(x)
    l == r
end

@test begin
    x = random(Exo1Quaternion{BigInt})
    y = random(Exo1Quaternion{BigInt})
    l = conj(x + y)
    r = conj(x) + conj(y)
    l == r
end

@test begin
    x = random(Exo1Quaternion{BigInt})
    y = random(Exo1Quaternion{BigInt})
    l = conj(x - y)
    r = conj(x) - conj(y)
    l == r
end

@test begin
    x = random(Exo1Quaternion{Rational{BigInt}})
    y = random(Exo1Quaternion{Rational{BigInt}})
    l = inv(x * y)
    r = inv(y) * inv(x)
    l == r
end

@test begin
    x = random(Exo1Quaternion{Rational{BigInt}})
    y = random(Exo1Quaternion{Rational{BigInt}})
    l = inv(x / y)
    r = y / x
    l == r
end

@test begin
    x = random(Exo1Quaternion{Rational{BigInt}})
    y = random(Exo1Quaternion{Rational{BigInt}})
    l = inv(x \ y)
    r = y \ x
    l == r
end

@test begin
    x = random(Exo1Quaternion{BigInt})
    y = random(Exo1Quaternion{BigInt})
    z = random(Exo1Quaternion{BigInt})
    l = x * (y + z)
    r = (x * y) + (x * z)
    l == r
end

@test begin
    a = CayleyDickson.randomBigInt()
    x = random(Exo1Quaternion{BigInt})
    y = random(Exo1Quaternion{BigInt})
    l = a * (x + y)
    r = (a * x) + (a * y)
    l == r
end

@test begin
    a = CayleyDickson.randomBigInt()
    x = random(Exo1Quaternion{BigInt})
    l = a * x
    r = x * a
    l == r
end

@test begin
    a = random(Quaternion{BigInt})
    x = random(Exo1Quaternion{BigInt})
    l = commutator(a, x)
    r = zero(Exo1Quaternion{BigInt})
    l != r
end

@test begin
    x = random(Exo1Quaternion{BigInt})
    y = random(Exo1Quaternion{BigInt})
    z = random(Exo1Quaternion{BigInt})
    l = (x + y) * z
    r = (x * z) + (y * z)
    l == r
end

@test begin
    x = random(Exo1Quaternion{BigInt})
    y = random(Exo1Quaternion{BigInt})
    z = random(Exo1Quaternion{BigInt})
    l = x * (y - z)
    r = (x * y) - (x * z)
    l == r
end

@test begin
    x = random(Exo1Quaternion{BigInt})
    y = random(Exo1Quaternion{BigInt})
    z = random(Exo1Quaternion{BigInt})
    l = (x - y) * z
    r = (x * z) - (y * z)
    l == r
end

@test begin
    x = random(Exo1Quaternion{BigInt})
    quadrance(x) > big(0)
end

@test begin
    x = random(Exo1Quaternion{BigInt})
    y = random(Exo1Quaternion{BigInt})
    l = quadrance(x * y)
    r = quadrance(x) * quadrance(y)
    l == r
end

@test_throws ErrorException begin
    inv(Exo1Quaternion(0, 0.0))
end

@test_throws ErrorException begin
    random(Exo1Quaternion{Int}) / 0
end

@test_throws ErrorException begin
    0 \ random(Exo1Quaternion{Int})
end

@test begin
    a = CayleyDickson.randomBigRational()
    z = random(Exo1Quaternion{Rational{BigInt}})
    l = z / a
    r = a \ z
    l == r
end

@test begin
    a = CayleyDickson.randomBigRational()
    z = random(Exo1Quaternion{Rational{BigInt}})
    l = a / z
    r = z \ a
    l == r
end

@test begin
    z = random(Exo1Quaternion{BigInt})
    z == z + zero(Exo1Quaternion{BigInt})
end

@test begin
    z = random(Exo1Quaternion{BigInt})
    z == z * one(Exo1Quaternion{BigInt})
end

@test begin
    z = random(Exo1Quaternion{BigInt})
    z == z + zero(z)
end

@test begin
    z = random(Exo1Quaternion{BigInt})
    z == z * one(z)
end


"""
    BoolDistribution

A distribution that provides the probabilities of true or false.
Can construct with `BoolDistribution(p_true)`.
"""
struct BoolDistribution
    p::Float64 # probability of true
end

pdf(d::BoolDistribution, s::Bool) = s ? d.p : 1.0-d.p

rand(rng::AbstractRNG, d::BoolDistribution) = rand(rng) <= d.p

iterator(d::BoolDistribution) = [true, false]

==(d1::BoolDistribution, d2::BoolDistribution) = d1.p == d2.p

Base.hash(d::BoolDistribution) = hash(d.p)

Base.length(d::BoolDistribution) = 2


using BlockDiagonalMatrices
using BlockDiagonalMatrices: svd_blockwise
using LinearAlgebra
using Random
using Test

@testset "linalg.jl" begin
    rng = MersenneTwister(123456)
    N1, N2, N3 = 3, 4, 5
    N = N1 + N2 + N3
    b1 = BlockDiagonal([rand(rng, N1, N1), rand(rng, N2, N2), rand(rng, N3, N3)])
    b2 = BlockDiagonal([rand(rng, N1, N1), rand(rng, N3, N3), rand(rng, N2, N2)])
    b3 = BlockDiagonal([rand(rng, N1, N1), rand(rng, N2, N2), rand(rng, N2, N2)])
    A = rand(rng, N, N + N1)
    B = rand(rng, N + N1, N + N2)
    A′, B′ = A', B'
    a = rand(rng, N)
    b = rand(rng, N + N1)

    @testset "mul!" begin
        c = similar(b1)
        d = similar(Matrix(b1))
        mul!(c, b1, b1)
        mul!(d, Matrix(b1), Matrix(b1))
        @test c ≈ d
        if VERSION ≥ v"1.3"
            mul!(c, b1, b1, 2.0, 3.0)
            mul!(d, Matrix(b1), Matrix(b1), 2.0, 3.0)
            @test c ≈ d
        end
    end

    @testset "Unary Linear Algebra" begin
        @testset "$f" for f in (adjoint, det, diag, eigvals, inv, pinv, svdvals, transpose, tr)
            @test f(b1) ≈ f(Matrix(b1))
        end

        @testset "permute=$p, scale=$s" for p in (true, false), s in (true, false)
            @testset "$g" for g in (eigmin, eigmax)
                # `b2` has real eigenvals, required for `eigmin`, `b1` has Complex eigenvals
                @test g(b2; permute=p, scale=s) ≈ g(Matrix(b2); permute=p, scale=s)
            end
            @testset "eigvals" begin
                result = eigvals(b2; permute=p, scale=s)
                expected = eigvals(Matrix(b2), permute=p, scale=s)
                @static if VERSION < v"1.2"
                    # the `eigvals` method we hit above did not sort real eigenvalues pre-v1.2
                    # but some `eigvals` methods did, so we always sort real eigenvalues.
                    @test sort!(result) ≈ sort!(expected)
                else
                    @test result ≈ expected
                end
            end
        end

        @testset "eigvals on LinearAlgebra types" begin
            # `eigvals` has different methods for different types, e.g. Hermitian
            b_herm = BlockDiagonal([Hermitian(rand(rng, 3, 3) + I) for _ in 1:3])
            @test eigvals(b_herm) ≈ eigvals(Matrix(b_herm))
            @test eigvals(b_herm, 1.0, 2.0) ≈ eigvals(Hermitian(Matrix(b_herm)), 1.0, 2.0)
        end

        # Requires a postive semi-definite matrix
        @testset "logdet" begin
            b̂1 = b1 * b1'
            @test logdet(b̂1) ≈ logdet(Matrix(b̂1))
        end
    end  # Unary

    @testset "Cholesky decomposition" begin
        X = [  4  12 -16
              12  37 -43
             -16 -43  98]
        U = [ 2.0 6.0 -8.0
              0.0 1.0  5.0
              0.0 0.0  3.0]

        B = BlockDiagonal([X, X])
        C = cholesky(B)
        @test C isa Cholesky{Float64, <:BlockDiagonal{Float64}}
        @test C.U ≈ cholesky(Matrix(B)).U
        @test C.U ≈ BlockDiagonal([U, U])
        @test C.L ≈ BlockDiagonal([U', U'])
        @test C.UL ≈ C.U
        @test C.uplo === 'U'
        @test C.info == 0

        M = BlockDiagonal(map(Matrix, blocks(C.L)))
        C = Cholesky(M, 'L', 0)
        @test C.U ≈ cholesky(Matrix(B)).U
        @test C.U ≈ BlockDiagonal([U, U])
        @test C.L ≈ BlockDiagonal([U', U'])
        @test C.UL ≈ C.L
        @test C.uplo === 'L'
        @test C.info == 0
    end  # Cholesky

    @testset "Singular Value Decomposition" begin
        X = [  4  12 -16
              12  37 -43
             -16 -43  98]
        B = BlockDiagonal([X, X])

        @testset "full=$full" for full in (true, false)

            @testset "svd_blockwise" begin
                U, S, Vt = svd_blockwise(B; full=full)
                F = SVD(U, S, Vt)
                @test B ≈ F.U * Diagonal(F.S) * F.Vt

                # Matrices should be BlockDiagonal
                @test F isa SVD{Float64, Float64, <:BlockDiagonal{Float64}}
                @test F.U isa BlockDiagonal
                @test F.V isa BlockDiagonal
                @test F.Vt isa BlockDiagonal

                # Should have same values, but not sorted so as to keep BlockDiagonal structure
                F_ = svd(Matrix(B), full=full)
                for fname in fieldnames(SVD)
                    @test sort(vec(getfield(F, fname))) ≈ sort(vec(getfield(F_, fname)))
                end
                # Singular values should be block-wise
                s = svd(X).S
                @test F.S == vcat(s, s)
            end

            @testset "svd" begin
                F = svd(B; full=full)
                F_ = svd(Matrix(B), full=full)

                @test F isa SVD
                @test B ≈ F.U * Diagonal(F.S) * F.Vt

                @test F == F_
                for fname in fieldnames(SVD)
                    @test getfield(F, fname) ≈ getfield(F_, fname)
                end

                # Singular values should be sorted in descending order
                @test F.S == sort(F.S, rev=true)
            end
        end
    end  # SVD
end


# ---
# title: 1518. Water Bottles
# id: problem1518
# author: Tian Jun
# date: 2020-10-31
# difficulty: Easy
# categories: Greedy
# link: <https://leetcode.com/problems/water-bottles/description/>
# hidden: true
# ---
#
# Given `numBottles` full water bottles, you can exchange `numExchange` empty
# water bottles for one full water bottle.
#
# The operation of drinking a full water bottle turns it into an empty bottle.
#
# Return the **maximum** number of water bottles you can drink.
#
#
#
# **Example 1:**
#
# **![](https://assets.leetcode.com/uploads/2020/07/01/sample_1_1875.png)**
#
#
#
#     Input: numBottles = 9, numExchange = 3
#     Output: 13
#     Explanation: You can exchange 3 empty bottles to get 1 full water bottle.
#     Number of water bottles you can drink: 9 + 3 + 1 = 13.
#
#
# **Example 2:**
#
# ![](https://assets.leetcode.com/uploads/2020/07/01/sample_2_1875.png)
#
#
#
#     Input: numBottles = 15, numExchange = 4
#     Output: 19
#     Explanation: You can exchange 4 empty bottles to get 1 full water bottle.
#     Number of water bottles you can drink: 15 + 3 + 1 = 19.
#
#
# **Example 3:**
#
#
#
#     Input: numBottles = 5, numExchange = 5
#     Output: 6
#
#
# **Example 4:**
#
#
#
#     Input: numBottles = 2, numExchange = 3
#     Output: 2
#
#
#
#
# **Constraints:**
#
#   * `1 <= numBottles <= 100`
#   * `2 <= numExchange <= 100`
#
#
## @lc code=start
using LeetCode

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


using LinearAlgebraicRepresentation
Lar = LinearAlgebraicRepresentation

"""
    frag_edge_channel(in_chan, out_chan,
        V::Lar.Points, EV::Lar.ChainOp, bigPI)

Parallel fragmentation of edges in `EV` using the spatial index `bigPI`.
"""
function frag_edge_channel(in_chan, out_chan, V, EV, bigPI)
    run_loop = true
    while run_loop
        edgenum = take!(in_chan)
        if edgenum != -1
            put!(out_chan, (edgenum, frag_edge(V, EV, edgenum, bigPI)))
        else
            run_loop = false
        end
    end
end

"""
    frag_edge(V::Lar.Points, EV::Lar.ChainOp, edge_idx::Int, bigPI::Array)

Fragment the edge of index `edge_idx` using the edges indicized in `bigPI`.
"""
function frag_edge(V, EV::Lar.ChainOp, edge_idx::Int, bigPI)
    alphas = Dict{Float64, Int}()
    edge = EV[edge_idx, :]
    verts = V[edge.nzind, :]
    for i in bigPI[edge_idx]
        if i != edge_idx
            intersection = Lar.Arrangement.intersect_edges(
            	V, edge, EV[i, :])
            for (point, alpha) in intersection
                verts = [verts; point]
                alphas[alpha] = size(verts, 1)
            end
        end
    end
    alphas[0.0], alphas[1.0] = [1, 2]
    alphas_keys = sort(collect(keys(alphas)))
    edge_num = length(alphas_keys)-1
    verts_num = size(verts, 1)
    ev = SparseArrays.spzeros(Int8, edge_num, verts_num)
    for i in 1:edge_num
        ev[i, alphas[alphas_keys[i]]] = 1
        ev[i, alphas[alphas_keys[i+1]]] = 1
    end
    return verts, ev
end

"""
    intersect_edges(V::Lar.Points, edge1::Lar.Cell, edge2::Lar.Cell)

Intersect two 2D edges (`edge1` and `edge2`).
"""
function intersect_edges(V::Lar.Points, edge1::Lar.Cell, edge2::Lar.Cell)
    err = 10e-8

    x1, y1, x2, y2 = vcat(map(c->V[c, :], edge1.nzind)...)
    x3, y3, x4, y4 = vcat(map(c->V[c, :], edge2.nzind)...)
    ret = Array{Tuple{Lar.Points, BigFloat}, 1}()

    v1 = [x2-x1, y2-y1];
    v2 = [x4-x3, y4-y3];
    v3 = [x3-x1, y3-y1];
    ang1 = dot(normalize(v1), normalize(v2))
    ang2 = dot(normalize(v1), normalize(v3))
    parallel = 1-err < abs(ang1) < 1+err
    colinear = parallel && (1-err < abs(ang2) < 1+err || -err < norm(v3) < err)
    if colinear
        o = [x1 y1]
        v = [x2 y2] - o
        alpha = 1/dot(v,v')
        ps = [x3 y3; x4 y4]
        for i in 1:2
            a = alpha*dot(v',(reshape(ps[i, :], 1, 2)-o))
            if 0 < a < 1
                push!(ret, (ps[i:i, :], a))
            end
        end
    elseif !parallel
        denom = (v2[2])*(v1[1]) - (v2[1])*(v1[2])
        a = ((v2[1])*(-v3[2]) - (v2[2])*(-v3[1])) / denom
        b = ((v1[1])*(-v3[2]) - (v1[2])*(-v3[1])) / denom

        if -err < a < 1+err && -err <= b <= 1+err
            p = [(x1 + a*(x2-x1))  (y1 + a*(y2-y1))]
            push!(ret, (p, a))
        end
    end
    return ret
end

"""
    merge_vertices!(V::Lar.Points, EV::Lar.ChainOp, edge_map, err=1e-4)

Merge congruent vertices and edges in `V` and `EV`.
"""
function merge_vertices!(V::Lar.Points, EV::Lar.ChainOp, edge_map, err=1e-4)
    vertsnum = size(V, 1)
    edgenum = size(EV, 1)
    newverts = zeros(Int, vertsnum)
    # KDTree constructor needs an explicit array of Float64
    V = Array{Float64,2}(V)
    kdtree = KDTree(permutedims(V))

    # merge congruent vertices
    todelete = []
    i = 1
    for vi in 1:vertsnum
        if !(vi in todelete)
            nearvs = Lar.inrange(kdtree, V[vi, :], err)
            newverts[nearvs] .= i
            nearvs = setdiff(nearvs, vi)
            todelete = union(todelete, nearvs)
            i = i + 1
        end
    end
    nV = V[setdiff(collect(1:vertsnum), todelete), :]

    # merge congruent edges
    edges = Array{Tuple{Int, Int}, 1}(undef, edgenum)
    oedges = Array{Tuple{Int, Int}, 1}(undef, edgenum)
    for ei in 1:edgenum
        v1, v2 = EV[ei, :].nzind
        edges[ei] = Tuple{Int, Int}(sort([newverts[v1], newverts[v2]]))
        oedges[ei] = Tuple{Int, Int}(sort([v1, v2]))
    end
    nedges = union(edges)
    nedges = filter(t->t[1]!=t[2], nedges)
    nedgenum = length(nedges)
    nEV = spzeros(Int8, nedgenum, size(nV, 1))
    # maps pairs of vertex indices to edge index
    etuple2idx = Dict{Tuple{Int, Int}, Int}()
    # builds `edge_map`
    for ei in 1:nedgenum
        nEV[ei, collect(nedges[ei])] .= 1
        etuple2idx[nedges[ei]] = ei
    end
    for i in 1:length(edge_map)
        row = edge_map[i]
        row = map(x->edges[x], row)
        row = filter(t->t[1]!=t[2], row)
        row = map(x->etuple2idx[x], row)
        edge_map[i] = row
    end
    # return new vertices and new edges
    return Lar.Points(nV), nEV
end

"""
    biconnected_components(EV::Lar.ChainOp)

Compute the biconnected components of the `EV` graph, represented
by edges as pairs of vertices.
"""
function biconnected_components(EV::Lar.ChainOp)

    ps = Array{Tuple{Int, Int, Int}, 1}()
    es = Array{Tuple{Int, Int}, 1}()
    todel = Array{Int, 1}()
    visited = Array{Int, 1}()
    bicon_comps = Array{Array{Int, 1}, 1}()
    hivtx = 1

    function an_edge(point) # TODO: fix bug
        # error? : BoundsError: attempt to access 0×0 SparseMatrix ...
        edges = setdiff(EV[:, point].nzind, todel)
        if length(edges) == 0
            edges = [false]
        end
        edges[1]
    end

    function get_head(edge, tail)
        setdiff(EV[edge, :].nzind, [tail])[1]
    end

    function v_to_vi(v)
        i = findfirst(t->t[1]==v, ps)
        # seems findfirst changed from 0 to Nothing
        if typeof(i) == Nothing
            return false
        elseif i == 0
            return false
        else
            return ps[i][2]
        end
    end

    push!(ps, (1,1,1))
    push!(visited, 1)
    exit = false
    while !exit
        edge = an_edge(ps[end][1])
        if edge != false
            tail = ps[end][2]
            head = get_head(edge, ps[end][1])
            hi = v_to_vi(head)
            if hi == false
                hivtx += 1
                push!(ps, (head, hivtx, ps[end][2]))
                push!(visited, head)
            else
                if hi < ps[end][3]
                    ps[end] = (ps[end][1], ps[end][2], hi)
                end
            end
            push!(es, (edge, tail))
            push!(todel, edge)
        else
            if length(ps) == 1
                found = false
                pop!(ps)
                for i in 1:size(EV,2)
                    if !(i in visited)
                        hivtx = 1
                        push!(ps, (i, hivtx, 1))
                        push!(visited, i)
                        found = true
                        break
                    end
                end
                if !found
                    exit = true
                end

            else
                if ps[end][3] == ps[end-1][2]
                    edges = Array{Int, 1}()
                    while true
                        edge, tail = pop!(es)
                        push!(edges, edge)
                        if tail == ps[end][3]
                            if length(edges) > 1
                                push!(bicon_comps, edges)
                            end
                            break
                        end
                    end

                else
                    if ps[end-1][3] > ps[end][3]
                        ps[end-1] = (ps[end-1][1], ps[end-1][2], ps[end][3])
                    end
                end
                pop!(ps)
            end
        end
    end
    bicon_comps = sort(bicon_comps, lt=(x,y)->length(x)>length(y))
    return bicon_comps
end

function get_external_cycle(V::Lar.Points, EV::Lar.ChainOp, FE::Lar.ChainOp)
    FV = abs.(FE)*EV
    vs = sparsevec(mapslices(sum, abs.(EV), dims=1)').nzind
    minv_x1 = maxv_x1 = minv_x2 = maxv_x2 = pop!(vs)
    for i in vs
        if V[i, 1] > V[maxv_x1, 1]
            maxv_x1 = i
        elseif V[i, 1] < V[minv_x1, 1]
            minv_x1 = i
        end
        if V[i, 2] > V[maxv_x2, 2]
            maxv_x2 = i
        elseif V[i, 2] < V[minv_x2, 2]
            minv_x2 = i
        end
    end
    cells = intersect(
        FV[:, minv_x1].nzind,
        FV[:, maxv_x1].nzind,
        FV[:, minv_x2].nzind,
        FV[:, maxv_x2].nzind
    )
    if length(cells) == 1
        return cells[1]
    else
        for c in cells
            if Lar.face_area(V, EV, FE[c, :]) < 0
                return c
            end
        end
    end
end
function pre_containment_test(bboxes)
    n = length(bboxes)
    containment_graph = spzeros(Int8, n, n)

    for i in 1:n
        for j in 1:n
            if i != j && Lar.bbox_contains(bboxes[j], bboxes[i])
                containment_graph[i, j] = 1
            end
        end
    end

    return containment_graph
end
function prune_containment_graph(n, V, EVs, shells, graph)

    for i in 1:n
        an_edge = shells[i].nzind[1]
        origin_index = EVs[i][an_edge, :].nzind[1]
        origin = V[origin_index, :]

        for j in 1:n
            if i != j
                if graph[i, j] == 1
                    shell_edge_indexes = shells[j].nzind
                    ev = EVs[j][shell_edge_indexes, :]

                    if !Lar.point_in_face(origin, V, ev)
                        graph[i, j] = 0
                    end
                end
             end
         end

     end
     return graph
end
function transitive_reduction!(graph)
    n = size(graph, 1)
    for j in 1:n
        for i in 1:n
            if graph[i, j] > 0
                for k in 1:n
                    if graph[j, k] > 0
                        graph[i, k] = 0
                    end
                end
            end
        end
    end
end
function cell_merging(n, containment_graph, V, EVs, boundaries, shells, shell_bboxes)
    function bboxes(V::Lar.Points, indexes::Lar.ChainOp)
        boxes = Array{Tuple{Any, Any}}(undef, indexes.n)
        for i in 1:indexes.n
            v_inds = indexes[:, i].nzind
            boxes[i] = Lar.bbox(V[v_inds, :])
        end
        boxes
    end
    # initiolization
    sums = Array{Tuple{Int, Int, Int}}(undef, 0);
    # assembling child components with father components
    for father in 1:n
        if sum(containment_graph[:, father]) > 0
            father_bboxes = bboxes(V, abs.(EVs[father]')*abs.(boundaries[father]'))
            for child in 1:n
                if containment_graph[child, father] > 0
                    child_bbox = shell_bboxes[child]
                    for b in 1:length(father_bboxes)
                        if Lar.bbox_contains(father_bboxes[b], child_bbox)
                            push!(sums, (father, b, child))
                            break
                        end
                    end
                end
            end
        end
    end
    # offset assembly initialization
    EV = vcat(EVs...)
    edgenum = size(EV, 1)
    facenum = sum(map(x->size(x,1), boundaries))
    FE = spzeros(Int8, facenum, edgenum)
    shells2 = spzeros(Int8, length(shells), edgenum)
    r_offsets = [1]
    c_offset = 1
    # submatrices construction
    for i in 1:n
        min_row = r_offsets[end]
        max_row = r_offsets[end] + size(boundaries[i], 1) - 1
        min_col = c_offset
        max_col = c_offset + size(boundaries[i], 2) - 1
        FE[min_row:max_row, min_col:max_col] = boundaries[i]
        shells2[i, min_col:max_col] = shells[i]
        push!(r_offsets, max_row + 1)
        c_offset = max_col + 1
    end
    # offsetting assembly of component submatrices
    for (f, r, c) in sums
        FE[r_offsets[f]+r-1, :] += shells2[c, :]
    end
    return EV, FE
end

# function decomposition

function componentgraph(V, copEV, bicon_comps)
    # arrangement of isolated components
	n = size(bicon_comps, 1)
   	shells = Array{Lar.Chain, 1}(undef, n)
	boundaries = Array{Lar.ChainOp, 1}(undef, n)
	EVs = Array{Lar.ChainOp, 1}(undef, n)
    # for each component
	for p in 1:n
		ev = copEV[sort(bicon_comps[p]), :]
        # computation of 2-cells
		fe = Lar.Arrangement.minimal_2cycles(V, ev)
        # exterior cycle
		shell_num = Lar.Arrangement.get_external_cycle(V, ev, fe)
        # decompose each fe (co-boundary local to component)
		EVs[p] = ev
		tokeep = setdiff(1:fe.m, shell_num)
		boundaries[p] = fe[tokeep, :]
		shells[p] = fe[shell_num, :]
    end
    # computation of bounding boxes of isolated components
	shell_bboxes = []
	for i in 1:n
    	vs_indexes = (abs.(EVs[i]')*abs.(shells[i])).nzind
   		push!(shell_bboxes, Lar.bbox(V[vs_indexes, :]))
	end
    # computation and reduction of containment graph
	containment_graph = pre_containment_test(shell_bboxes)
	containment_graph = prune_containment_graph(n, V, EVs, shells, containment_graph)
	transitive_reduction!(containment_graph)
	return n, containment_graph, V, EVs, boundaries, shells, shell_bboxes
end

function cleandecomposition(V, copEV, sigma, edge_map)
    # Deletes edges outside sigma area
    todel = []
    new_edges = []
    map(i->new_edges=union(new_edges, edge_map[i]), sigma.nzind)
    ev = copEV[new_edges, :]
    for e in 1:copEV.m
        if !(e in new_edges)

            vidxs = copEV[e, :].nzind
            v1, v2 = map(i->V[vidxs[i], :], [1,2])
            centroid = .5*(v1 + v2)

            if ! Lar.point_in_face(centroid, V, ev)
                push!(todel, e)
            end
        end
    end

    for i in reverse(todel)
        for row in edge_map

            filter!(x->x!=i, row)

            for j in 1:length(row)
                if row[j] > i
                    row[j] -= 1
                end
            end
        end
    end

    V, copEV = Lar.delete_edges(todel, V, copEV)
	return V,copEV
end

function planar_arrangement_1( V, copEV,
		sigma::Lar.Chain=spzeros(Int8, 0),
		return_edge_map::Bool=false,
		multiproc::Bool=false)
	# data structures initialization
	edgenum = size(copEV, 1)
	edge_map = Array{Array{Int, 1}, 1}(undef,edgenum)
	rV = Lar.Points(zeros(0, 2))
	rEV = SparseArrays.spzeros(Int8, 0, 0)
	finalcells_num = 0

	# spaceindex computation
	model = (convert(Lar.Points,V'),Lar.cop2lar(copEV))
	bigPI = Lar.spaceindex(model::Lar.LAR)

    # multiprocessing of edge fragmentation
    if (multiproc == true)
        in_chan = Distributed.RemoteChannel(()->Channel{Int64}(0))
        out_chan = Distributed.RemoteChannel(()->Channel{Tuple}(0))
        ordered_dict = SortedDict{Int64,Tuple}()
        @async begin
            for i in 1:edgenum
                put!(in_chan,i)
            end
            for p in distributed.workers()
                put!(in_chan,-1)
            end
        end
        for p in distributed.workers()
            @async Base.remote_do(frag_edge_channel, p, in_chan, out_chan, V, copEV, bigPI)
        end
        for i in 1:edgenum
            frag_done_job = take!(out_chan)
            ordered_dict[frag_done_job[1]] = frag_done_job[2]
        end
        for (dkey, dval) in ordered_dict
            i = dkey
            v, ev = dval
            newedges_nums = map(x->x+finalcells_num, collect(1:size(ev, 1)))
            edge_map[i] = newedges_nums
            finalcells_num += size(ev, 1)
            rV, rEV = Lar.skel_merge(rV, rEV, v, ev)
        end
    else
        # sequential (iterative) processing of edge fragmentation
        for i in 1:edgenum
            v, ev = frag_edge(V, copEV, i, bigPI)
            newedges_nums = map(x->x+finalcells_num, collect(1:size(ev, 1)))
            edge_map[i] = newedges_nums
            finalcells_num += size(ev, 1)
            rV = convert(Lar.Points, rV)
            rV, rEV = Lar.skel_merge(rV, rEV, v, ev)
        end
    end
    # merging of close vertices and edges (2D congruence)
    V, copEV = rV, rEV
    V, copEV = merge_vertices!(V, copEV, edge_map)
	return V,copEV,sigma,edge_map
end

function planar_arrangement_2(V, copEV,bicon_comps, edge_map,
		sigma::Lar.Chain=spzeros(Int8, 0))

    edges = sort(union(bicon_comps...))
    todel = sort(setdiff(collect(1:size(copEV,1)), edges))

    for i in reverse(todel)
        for row in edge_map

            filter!(x->x!=i, row)

            for j in 1:length(row)
                if row[j] > i
                    row[j] -= 1
                end
            end
        end
    end

	bicon_comps = Lar.Arrangement.biconnected_components(copEV)

	# component graph
	n, containment_graph, V, EVs, boundaries, shells, shell_bboxes=Lar.Arrangement.componentgraph(V,copEV,bicon_comps)

	copEV, FE = Lar.Arrangement.cell_merging(
	   	n, containment_graph, V, EVs, boundaries, shells, shell_bboxes)

	return V,copEV,FE
end

"""
    planar_arrangement(V::Points, EV::ChainOp, [sigma::Chain], [return_edge_map::Bool], [multiproc::Bool])

Compute the arrangement on the given cellular complex 1-skeleton in 2D.

A cellular complex is arranged when the intersection of every possible pair of cell
of the complex is empty and the union of all the cells is the whole Euclidean space.
The basic method of the function without the `sigma`, `return_edge_map` and `multiproc` arguments
returns the full arranged complex `V`, `EV` and `FE`.

## Additional arguments:
- `sigma::Chain`: if specified, `planar_arrangement` will delete from the output every edge and face outside this cell. Defaults to an empty cell.
- `return_edge_map::Bool`: makes the function return also an `edge_map` which maps the edges of the imput to the one of the output. Defaults to `false`.
- `multiproc::Bool`: Runs the computation in parallel mode. Defaults to `false`.
"""
function planar_arrangement(
        V::Lar.Points,
        copEV::Lar.ChainOp,
        sigma::Lar.Chain=spzeros(Int8, 0),
        return_edge_map::Bool=false,
        multiproc::Bool=false)

    #planar_arrangement_1
	V,copEV,sigma,edge_map=Lar.Arrangement.planar_arrangement_1(V,copEV,sigma,return_edge_map,multiproc)
    # cleandecomposition
	if sigma.n > 0
		V,copEV=Lar.Arrangement.cleandecomposition(V, copEV, sigma, edge_map)
	end

    bicon_comps = Lar.Arrangement.biconnected_components(copEV)
    # EV = Lar.cop2lar(copEV)
    # V,bicon_comps = Lar.biconnectedComponent((V,EV))

	if isempty(bicon_comps)
    	println("No biconnected components found.")
    	if (return_edge_map)
    	    return (nothing, nothing, nothing, nothing)
    	else
    	    return (nothing, nothing, nothing)
    	end
	end
    #Planar_arrangement_2
	V,copEV,FE=Lar.Arrangement.planar_arrangement_2(V,copEV,bicon_comps,edge_map,sigma)
    println("TEST EULERO PLANAR START")
    Lar.test_eulero(V,copEV, FE)
    Lar.test_eulero_face4face(V,copEV, FE)
    println("TEST EULERO PLANAR END")
	if (return_edge_map)
	     return V, copEV, FE, edge_map
	else
	     return V, copEV, FE
	end
end


__precompile__()

module ODEFilters

using Reexport
@reexport using DiffEqBase
using OrdinaryDiffEq
import DiffEqBase: check_error!

import Base: copy, copy!, show, size, ndims
stack(x) = copy(reduce(hcat, x)')

using LinearAlgebra
using TaylorSeries
@reexport using StructArrays
using UnPack
using RecipesBase
using ModelingToolkit
using RecursiveArrayTools
using StaticArrays

@reexport using PSDMatrices
import PSDMatrices: X_A_Xt
X_A_Xt(A, X) = X*A*X'
apply_diffusion(Q, diffusion::Diagonal) = X_A_Xt(Q, sqrt.(diffusion))
apply_diffusion(Q, diffusion::Number) = Q*diffusion

# All the Gaussian things
@reexport using GaussianDistributions
using GaussianDistributions: logpdf
const PSDGaussian{T} = Gaussian{Vector{T}, PSDMatrix{T}}
const PSDGaussianList{T} = StructArray{PSDGaussian{T}}
copy(P::Gaussian) = Gaussian(copy(P.μ), copy(P.Σ))
copy!(dst::Gaussian, src::Gaussian) = (copy!(dst.μ, src.μ); copy!(dst.Σ, src.Σ); nothing)
RecursiveArrayTools.recursivecopy(P::Gaussian) = copy(P)
show(io::IO, g::Gaussian) = print(io, "Gaussian($(g.μ), $(g.Σ))")
show(io::IO, ::MIME"text/plain", g::Gaussian{T, S}) where {T, S} =
    print(io, "Gaussian{$T,$S}($(g.μ), $(g.Σ))")
size(g::Gaussian) = size(g.μ)
ndims(g::Gaussian) = ndims(g.μ)

Base.:*(M, g::PSDGaussian) = Gaussian(M * g.μ, X_A_Xt(g.Σ, M))
GaussianDistributions.whiten(Σ::PSDMatrix, z) = Σ.L\z

import Statistics: mean, var, std
var(p::PSDGaussian{T}) where {T} = diag(p.Σ)
std(p::PSDGaussian{T}) where {T} = sqrt.(diag(p.Σ))
mean(s::PSDGaussianList{T}) where {T} = mean.(s)
var(s::PSDGaussianList{T}) where {T} = var.(s)
std(s::PSDGaussianList{T}) where {T} = std.(s)

include("priors.jl")
include("diffusions.jl")

include("algorithms.jl")
export EK0, EK1
include("alg_utils.jl")
include("caches.jl")
include("state_initialization.jl")
include("integrator_utils.jl")
include("filtering.jl")
include("perform_step.jl")
include("preconditioning.jl")
include("smoothing.jl")
include("solution.jl")

# Utils
include("jacobian.jl")
include("numerics_tricks.jl")

# Iterated Extended Kalman Smoother
include("ieks.jl")
export IEKS, solve_ieks

end


#=

  The BUGS Book
  Example 3.3.2 Surgery (continued): beta-binomial analysis using BUGS

  Cf ~/jags/bugs_book_3_3_2.R
  """
  model {
    theta ~ dbeta(a, b)
    y ~ dbin(theta, n)
    Y.pred ~ dbin(theta, n.pred)
    P.crit <- step(Y.pred - n.crit + 0.5)
  }

  Output:
          Mean      SD  Naive SE Time-series SE
P.crit 0.41531 0.49278 0.0017422      0.0019429
Y.pred 1.49477 1.42409 0.0050349      0.0061002
theta  0.07453 0.04117 0.0001456      0.0002249
  """

  Summary Statistics
  parameters      mean       std   naive_se      mcse         ess      rhat   ess_per_sec
      Symbol   Float64   Float64    Float64   Float64     Float64   Float64       Float64

       theta    0.0748    0.0416     0.0004    0.0005   7624.8902    1.0000     3450.1765
      y_pred    1.4931    1.4276     0.0143    0.0144   9178.6177    1.0000     4153.2207
      p_crit    0.4177    0.4932     0.0049    0.0051   9070.0308    1.0001     4104.0863

  Cf ~/webppl/bugs_book_3_3_2.wppl

=#

using Turing, StatsPlots, DataFrames
include("jl_utils.jl")

@model function bugs_book_3_3_2(y=0)
    a = 3
    b = 27
    n = 10
    n_pred = 20
    n_crit = 2

    theta ~ Beta(a,b)
    y ~ Binomial(n,theta)
    y_pred ~ Binomial(n_pred,theta)
    p_crit ~ Dirac(y_pred + 0.5 > n_crit) #  toReal(y_pred - n_crit) + 0.5 > 0.0;

end

model = bugs_book_3_3_2()

num_chains = 4

# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 100_000)
# chns = sample(model, PG(15), 10_000)
chns = sample(model, SMC(10_000), 10_000)
# chns = sample(model, IS(), 10_000)

# chns = sample(model, NUTS(), 10_000)
# chns = sample(model, HMC(0.1,5), 10_000)

display(chns)
# display(plot(chns))


using DimensionalData, Test

A = zeros(X(4.0:7.0), Y(10.0:12.0))
di = DimIndices(A)

ci = CartesianIndices(A)
@test val.(collect(di)) == Tuple.(collect(ci))

di[4, 3] == (X(4), Y(3))
@test di[X(1)] == [(Y(1),), (Y(2),), (Y(3),)]
@test map(ds -> A[ds...] + 2, di) == fill(2.0, 4, 3)
@test map(ds -> A[ds...], di[X(At(7.0))]) == [fill(0.0, 4) for i in 1:3]
@test_throws ArgumentError DimIndices(zeros(2, 2))

di = DimKeys(A)
di[4, 3] == (X(At(7.0)), Y(At(12.0)))
@test di[X(1)] == [(Y(At(10.0)),), (Y(At(11.0)),), (Y(At(12.0)),)]
@test map(ds -> A[ds...] + 2, di) == fill(2.0, 4, 3)
@test map(ds -> A[ds...], di[X(At(7.0))]) == [fill(0.0, 4) for i in 1:3]
@test_throws ArgumentError DimIndices(zeros(2, 2))


#!/usr/local/bin/julia

using PyPlot, JLD
# Plot the learning curve for the MNIST tutorial to verify
# the solver is progressing toward convergence

stats = load("snapshots/statistics.jld")
# println(typeof(stats))

tables = stats["statistics"]
ov = tables["obj_val"]
xy = sort(collect(ov))
datapoints = size(xy[:])
println("XY contains $datapoints datapoints")
x = [i for (i,j) in xy]
y = [j for (i,j) in xy]
x = convert(Array{Int64}, x)
y = convert(Array{Float64}, y)

raw = plot(x, y, linewidth=1, label="Raw")
xlabel("Iterations")
ylabel("Objective Value")
title("MNIST Learning Curve")
grid("on")

function low_pass{T <: Real}(x::Vector{T}, window::Int)
  len = length(x)
  y = Vector{Float64}(len)
  for i in 1:len
      # I want the mean of the first i terms up to width of window
      # Putting some numbers to this with window 4
      # i win lo  hi
      # 1  4  1   1
      # 2  4  1   2
      # 3  4  1   3
      # 4  4  1   4
      # 5  4  1   5
      # 6  4  2   6  => window starts to slide
      lo = max(1, i - window)
      hi = i
      y[i] = mean(x[lo:hi])
  end
  return y
end

window = Int64(round(length(xy)/4.0))
y_avg = low_pass(y, window)
avg = plot(x, y_avg, linewidth=2, label="Low Pass")
legend(handles=[raw; avg])
show()  #Shows the final plot


#    Copyright (c) 2021 J.A. Duffek
#    Copyright (c) 2000 D.M. Spink
#
#    This program is free software; you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation; either version 2 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program; if not, see <http://www.gnu.org/licenses/>.

"""
    vecscale(svec::Vector{F})::Matrix{F} where {F<:AbstractFloat}

Transformation matrix for a scaling.

Returns a (4x4) Transformation matrix for scaling.

The matrix is:

    [ sx   0   0   0]
    [  0  sy   0   0]
    [  0   0  sz   0]
    [  0   0   0   1]

# Arguments:
- `svec`: A vectors defining the scaling along the x,y and z axes.
  i.e. `[sx; sy; sy]`

# Output:
- `st`: Scaling Transformation Matrix

# Examples:
```julia
julia> ss = vecscale(svec)
```

Scale up the NURBS line `[0.0;0.0;0.0]` - `[1.0;1.0;1.0]` by 3 along the
x-axis, 2 along the y-axis and 4 along the z-axis.

```julia
julia> line = nrbline(vec([0.0 0.0 0.0]),vec([1.0 1.0 1.0]));
julia> scale = vecscale(vec([3.0 2.0 4.0]));
julia> sline = nrbtform(line, scale)
```

# Reference
- see also: [`nrbtform`](@ref)
"""
function vecscale(svec::Vector{F})::Matrix{F} where {F<:AbstractFloat}
if length(svec) < 1 || length(svec) > 3
  throw(ArgumentError("Translation vector required"));
end
s = [svec;1;1];
ss = [s[1]    0    0   0;
      0    s[2]    0   0;
      0       0 s[3]   0;
      0       0    0   1];
return ss;
end # vecscale


using JLBoost, RDatasets
iris = dataset("datasets", "iris")

iris[!, :is_setosa] = iris[!, :Species] .== "setosa"
target = :is_setosa

features = setdiff(names(iris), ["Species", "is_setosa"])

# fit one tree
# ?jlboost for more details
xgtreemodel = jlboost(iris, target)

typeof(trees(xgtreemodel))

typeof(xgtreemodel.loss)

typeof(xgtreemodel.target)

xgtreemodel2 = jlboost(iris, target; nrounds = 2, max_depth = 2)

xgtreemodel3 = jlboost(iris, target; nrounds = 1, max_leaves = 4, max_depth = 0)

iris.pred1 = predict(xgtreemodel, iris)
iris.pred2 = predict(xgtreemodel2, iris)
iris.pred1_plus_2 = predict(vcat(xgtreemodel, xgtreemodel2), iris)

AUC(-iris.pred1, iris.is_setosa)

gini(-iris.pred1, iris.is_setosa)

feature_importance(xgtreemodel, iris)

nrow(df) # returns the number of rows
ncol(df)
view(df, rows, cols)

using DataFrames
using JLBoost
df = DataFrame(x = rand(100) * 100)

df[!, :y] = 2*df.x .+ rand(100)

target = :y
features = [:x]
warm_start = fill(0.0, nrow(df))

using LossFunctions: L2DistLoss
loss = L2DistLoss()
jlboost(df, target, features, warm_start, loss; max_depth=2) # default max_depth = 6

JLBoost.save(xgtreemodel, "model.jlb")
JLBoost.save(trees(xgtreemodel), "model_tree.jlb")

JLBoost.load("model.jlb")
JLBoost.load("model_tree.jlb")

using JLBoost, RDatasets, JDF
iris = dataset("datasets", "iris")

iris[!, :is_setosa] = iris[!, :Species] .== "setosa"
target = :is_setosa

features = setdiff(names(iris), [:Species, :is_setosa])

savejdf("iris.jdf", iris)
irisdisk = JDFFile("iris.jdf")

# fit using on disk JDF format
xgtree1 = jlboost(irisdisk, target, features)
xgtree2 = jlboost(iris, target, features; nrounds = 2, max_depth = 2)

# predict using on disk JDF format
iris.pred1 = predict(xgtree1, irisdisk)
iris.pred2 = predict(xgtree2, irisdisk)

# AUC
AUC(-predict(xgtree1, irisdisk), irisdisk[!, :is_setosa])

# gini
gini(-predict(xgtree1, irisdisk), irisdisk[!, :is_setosa])

# clean up
rm("iris.jdf", force=true, recursive=true)


## general data and market data =================================
seed = 2            ## seed for random generator
n_scen = 1_000_000  ## number of monte carlo scenarios
α = 0.99            ## confidence level for expected shortfall
coc_rate = 0.06     ## cost of capital rate for Mindestbetrag
spot = [0.00155, 0.00013, -0.00011]  ## spot rates as of 2012
stock_increase = 0.05   ## expected rel. increase of stock index

## risk factors =================================================

## CHF spot rates
x0_spot = [0.0, 0.0, 0.0]  ## as of 2012
h_spot = [0.01, 0.01, 0.01]
add_spot = [true, true, true]
index_spot = collect(1:3) # Arg. must always be of the form 1:T
## stock index
x0_stock = 1.0
h_stock = 0.1
add_stock = false
index_stock = [4]
## mortality
x0_mort = 1.0
h_mort = 0.1
add_mort = false
index_mort = [5]

## data provied by FINMA (for 2012) -----------------------------
## components link to index_spot, index_stock, index_mort

## standard deviations of risk factors
σ = [0.0060292643195745,         # spot_chf[1]
     0.00605990680656078,        # spot_chf[2]
     0.00633152336734695,        # spot_chf[3]
     0.150517788474455,          # stock
     0.05 ]                      # q

## correlations for risk factors
##       spot_chf[1] spot_chf[2] spot_chf[3] stock       mort
corr =  [1.000000000 0.721560179 0.545556472 0.402239567 0;
         0.721560179 1.000000000 0.953194547 0.433388849 0;
         0.545556472 0.953194547 1.000000000 0.413041139 0;
         0.402239567 0.433388849 0.413041139 1.000000000 0;
         0.000000000 0.000000000 0.000000000 0.000000000 1]

## stress scenarios

## data frame that describes the stresst scenarios to be added
## n:       number of scenarios
## name:    stress scenario name / description
## target:  true: include, false: do not include
## prob:    probability for this stress scenario occuring
## Δx:      effect of this scenario on risk factor
## Δrtk:    efect of this scenario on RTK (to be calculated)
stress =
  Stress(8,
         ["SZ01: Equity drop -60%",
          "SZ03: Stock mkt crash (87)",
          "SZ04: Nikkei crash (89/90)",
          "SZ05: European currency crisis (92)",
          "SZ06: US interest crisis (94)",
          "SZ07: Russian crisis/LTCM (98)",
          "SZ08: Stock mkt crash (00/01)",
          "SZ11: Financial crisis (08)"],
         [true, true, true, true, true, true, true, true],
         [0.001, 0.001, 0.001, 0.001, 0.001, 0.001, 0.001, 0.001],
         [ 0.00000  0.00000 0.00000 -0.6000 0;
          -0.00155 -0.00013 0.00000 -0.2323 0;
           0.01563  0.01098 0.01177 -0.2643 0;
          -0.00155 -0.00013 0.00000 -0.0580 0;
           0.01109  0.01406 0.01509 -0.1852 0;
          -0.00155 -0.00013 0.00000 -0.2841 0;
          -0.00155 -0.00013 0.00000 -0.3567 0;
          -0.00155 -0.00013 0.00000 -0.3881 0],
         zeros(Float64, 8))

## portfolio data  ==============================================
invest =
  DataFrame(maturity = [NaN, 1, 2, 3],
            index = [4, 1, 2, 3],
            kind = ["stock", "zerobond", "zerobond", "zerobond"],
            nominal = [300.0, 0.0, 300.0, 500.0])

n_ins = 10                         ## number insurance contracts
ins_sum = 100.0                    ## insured sum per contract
qx_flat = 0.02                     ## mortality probability
qx = [qx_flat, qx_flat, qx_flat]   ## mortality probabilities
B_PX  =[0.0, 0.0, n_ins * ins_sum] ## benefit profile (survival)
T = length(B_PX)                   ## Time until run off


using LinearAlgebra
using SparseArrays
using Laplacians
using Random

function lmdmin_val(L,n)
	I=spzeros(n,n);
	for i=1:n
		I[i,i]=n;
	end
	for i=1:n
		for j=1:n
			L[i,j]=I[i,j]-L[i,j];
		end
	end
	u=randn(n);
	u=u./norm(u);
	y=L*u;
	y=y./norm(y);
	while norm(y-u)>1e-2
		u=y;
		y=L*u;
		y=y./norm(y);
	end
	return n-u'*L*y;
end

function lmdedge(L,s,k)
	for i=1:k
		L[Int(s[i]),Int(s[i])]+=1;
	end
	f=eigen(L);
	return f.values[1];
end

function lmd_LG(L,s,n)
	S=union(1:n);
	setdiff!(S,s);
	LG=L[S,S];
	f=eigvals(LG);
	k=1;
	while f[k]<1e-8
		k+=1;
	end
	return f[k];
end

function lap(G :: Graph)
    F = zeros(G.n, G.n);
    for i = 1 : G.m
        F[G.u[i], G.v[i]] -= 1
        F[G.v[i], G.u[i]] -= 1
        F[G.u[i], G.u[i]] += 1
        F[G.v[i], G.v[i]] += 1
    end
    return F
end

function lapsp(G :: Graph)
    F = spzeros(G.n, G.n)
    for i = 1 : G.m
        F[G.u[i], G.v[i]] -= 1
        F[G.v[i], G.u[i]] -= 1
        F[G.u[i], G.u[i]] += 1
        F[G.v[i], G.v[i]] += 1
    end
    return F
end

function getB(G :: Graph, mm)
    B :: SparseMatrixCSC{Float64} = spzeros(mm,G.n)
	jishu=0;
    for i = 1 : G.m
    #    if (G.u[i]>G.n0+G.n1 && G.v[i]>G.n0+G.n1)
			jishu+=1;
            B[jishu, G.u[i]] = 1
            B[jishu, G.v[i]] = -1
    #    end
    end
    return B
end

function getD(G :: Graph)
    D = spzeros(G.n-G.n0-G.n1,G.n-G.n0-G.n1);
    for i = 1 : G.m
        if (G.u[i]<=G.n0+G.n1) && (G.v[i]>G.n0+G.n1)
            D[G.v[i]-G.n0-G.n1] += 1;
        end
        if (G.v[i]<=G.n0+G.n1) && (G.u[i]>G.n0+G.n1)
            D[G.u[i]-G.n0-G.n1] += 1;
        end
    end
    return D;
end

function calc(G :: Graph ,i ,Linv, sel)
    e=spzeros(G.n-G.n0-G.n1,1)
    e[i,1]=1
    vect = Linv*e-Linv*e*(e'*Linv*sel);
    ans = sum(vect)/(1+Linv[i,i]);
    return ans
end

function findmaxexa(G :: Graph, Linv ,sel,eta)
    bestv = 0;
    tmp = 0
    maxinc = 0
    for i = G.n0+G.n1+1 : G.n
        if sel[i-(G.n0+G.n1)]< G.n1
            tmp = calc(G, i-G.n0-G.n1 , Linv ,sel)
            if tmp > maxinc && (Linv*sel)[i-G.n0-G.n1]<eta
                maxinc = tmp
                bestv = i
            end
			#print(tmp," ")
        end
    end
    return bestv
end

function exa(G :: Graph,k,eta,fout)
    L = lap(G);
    sel = zeros(G.n - G.n0- G.n1);
    for i = G.n0+G.n1+1 : G.n
        sel[i-(G.n0+G.n1)] =  sum(L[i,1:G.n0]) - sum(L[i,1:G.n0+G.n1]);
    end
	T1=time();
    Linv = inv(L[G.n0+G.n1+1:G.n,G.n0+G.n1+1:G.n]);
	firans = sum(Linv * sel);
	j=0;
	T2=time();
	#println(" ");
    for i = 1 : k
        j = findmaxexa(G, Linv ,sel,eta);
		#print(j-G.n0-G.n1," ")
        sel[j-G.n0-G.n1] += 1;
        e=spzeros(G.n-G.n0-G.n1,1);
        e[j-G.n0-G.n1,1]=1;
        Linv = Linv - Linv*e*e'*Linv/(1+Linv[j-G.n0-G.n1,j-G.n0-G.n1]);
    end
	#println("");
	T3=time();
	println(fout,"exact basetime=",T2-T1,",",k," times=",T3-T2);
    ansv = Linv * sel;
    return sum(ansv);
end

function gre(G :: Graph, k ,ep,eta,fout)
	finans= 0 ;
    d0=ep/(3*sqrt(G.n));
	#d0=eps/(3);
    #d=4*eps/(3*n^2);
	d=4*ep/(3*G.n^2)/(G.n^2);
	JLfac=1.0;
    #t= round(Int64,12*12*24*log2(G.n)/eps^2);
	t= round(Int64,JLfac*log(G.n)/(ep^2))+1;
	L = lapsp(G);
	nn = G.n-G.n0-G.n1;
    LL = L[G.n0+G.n1+1:G.n,G.n0+G.n1+1:G.n];
    mm = 0;
    for i = 1: G.m
        if ((G.u[i]>G.n0+G.n1) && (G.v[i]>G.n0+G.n1))
            mm += 1;
        end
    end
	B = getB(G,mm);
    D = getD(G);
	Dadd = spzeros(G.n-G.n0-G.n1,G.n-G.n0-G.n1);
    #sqrt.(D);
    rng = MersenneTwister(Int(round(time() * 10000.0)));
    #f = approxchol_sddm(LL, tol=d0)

    sel = zeros(nn);
    for i = G.n0+G.n1+1 : G.n
        sel[i-(G.n0+G.n1)] =  sum(L[i,1:G.n0]) - sum(L[i,1:G.n0+G.n1]);
    end
    o1 = ones(nn);
	maxx = 0;
	z1 = zeros(nn, 1)
	z2 = zeros(nn, 1)
	f = approxchol_sddm(LL, tol=1e-8);
	firans =sum(f(sel));
	#r2 = randn(rng, nn, 1);
	#r1 = randn(rng, mm, 1);

	T1=time();

    for i = 1 : k

##################

####f = Laplacians.approxchol_sddm(LL+Dadd, tol=d0);
f = Laplacians.approxchol_sddm(LL, tol=d0,verbose=false);
        h1=f(o1);
        h2=f(sel);

#这里就是用求解器的步骤，f = approxchol_sddm(L, tol=d0);L是拉普拉斯矩阵，tol=d0是误差，选取1e-6完全够了，
#h1=f(o1)就是返回L的伪逆*o1向量
#####################

		Xe = zeros(nn);
		Ye = zeros(nn);
        for j = 1 : t
            r2 = randn(rng, nn, 1);
			r1 = randn(rng, mm, 1);
			xx = B'*r1;
			yy = sqrt.(D)*r2;
            f = approxchol_sddm(LL, tol=d);
			z1[:, 1] = f(xx[:, 1]);
			z2[:, 1] = f(yy[:, 1]);
			for p = 1 : nn
				Xe[p] += z1[p,1]^2;
				Ye[p] += z2[p,1]^2;
			end
		end
		add = zeros(G.n-G.n0-G.n1);
		for j = 1 : G.n-G.n0-G.n1
			add[j] += h1[j]*(1-h2[j])/(1+Xe[j]/t+Ye[j]/t);
		end
		maxx=0;
		bestv=0;
		for j = 1 : nn
			if (sel[j]<G.n1 && add[j]>maxx && h2[j]<eta)
				maxx = add[j]
				bestv = j
			end
		end
		#println(" ");
		#for j =1 : G.n-G.n0-G.n1
		#	print(add[j]," ")
		#end
		sel[bestv]+=1;
		#LL[bestv,bestv]+=1;
		D[bestv,bestv]+=1;
		Dadd[bestv,bestv]+=1;
		finans += maxx;
		#print(bestv," ")
	end
	T2=time();
	#println(" ");
	f = approxchol_sddm(LL+Dadd, tol=1e-8);
	finans =sum(f(sel));
	println(fout,"greed ",k," times=",T2-T1," eps=",ep);
	return finans;
end

function appxInvTrace(parL::SparseMatrixCSC{Float64}, parB::SparseMatrixCSC{Float64})

  f = approxchol_sddm(parL,tol=1e-6, verbose=false)
  m = size(parB, 1)
  n = size(parL, 1)
  k = round(Int, log(n)) # number of dims for JL
  print(n)
  rst = 0

  for i = 1:k
  #  V[:,i] .= f(UR[:,i])
    r = rand([1,-1],m,1)
    #ur = U'*r
    v = r'*parB
    print(v)
    v = f(v')
    #rst += sum(sqrt(v.*v))
  end

  return rst/k;

end


module NonlinearSamplingTest

using Test
using Random
using ForneyLab
using ForneyLab: outboundType, isApplicable, prod!, logPdf, unsafeMean, unsafeVar, Sampling, requiresBreaker, breakerParameters
using ForneyLab: SPNonlinearSInGX, SPNonlinearSOutNGX, SPNonlinearSOutNM, SPNonlinearSIn1MN, SPNonlinearSOutNMX, SPNonlinearSInMX, MNonlinearSInMGX, gradientOptimization
using ForwardDiff

Random.seed!(1234)

f(x) = x
g(x) = x^2
h(x, y) = x + y

@testset "gradientOptimization" begin
    f(x) = -x[1]^2 - x[2]^2
    grad(x) = ForwardDiff.gradient(f, x)
    res = gradientOptimization(f, grad, [4.0, 1.0], 0.001)

    # Optimum is [0.0, 0.0]
    @test isapprox(res, [0.0, 0.0], atol=1e-16)
end

@testset "requiresBreaker and breakerParameters" begin
    fg = FactorGraph()
    x = Variable()
    y = Variable()
    nd = GaussianMeanVariance(x, 0.0, 1.0)
    Nonlinear{Sampling}(y, x, g=g)

    @test !requiresBreaker(nd.i[:out]) # Single-input Nonlinear{Sampling} does not require breaker

    fg = FactorGraph()
    x = Variable()
    z = Variable()
    y = Variable()
    GaussianMeanVariance(z, 0.0, 1.0)
    nd = GaussianMeanVariance(x, 0.0, 1.0)
    Nonlinear{Sampling}(y, x, z, g=h)

    @test requiresBreaker(nd.i[:out]) # Multi-input Nonlinear{Sampling} requires breaker
    @test_throws Exception breakerParameters(nd.i[:out].partner)
    @test breakerParameters(nd.i[:out]) == (Message{GaussianMeanVariance, Univariate}, ())
end

#-------------
# Update rules
#-------------

@testset "SPNonlinearSOutNM" begin
    @test SPNonlinearSOutNM <: SumProductRule{Nonlinear{Sampling}}
    @test outboundType(SPNonlinearSOutNM) == Message{SampleList}
    @test isApplicable(SPNonlinearSOutNM, [Nothing, Message{Gaussian}])
    @test isApplicable(SPNonlinearSOutNM, [Nothing, Message{Bernoulli}])

    msg = ruleSPNonlinearSOutNM(f, nothing, Message(Univariate, GaussianMeanVariance, m=2.0, v=tiny), n_samples=1)
    @test isapprox(msg.dist.params[:s][1], 2.0, atol=1e-4)
    @test msg.dist.params[:w] == [1.0]
    msg = ruleSPNonlinearSOutNM(f, nothing, Message(Multivariate, GaussianMeanVariance, m=[2.0], v=mat(tiny)), n_samples=1)
    @test isapprox(msg.dist.params[:s][1][1], 2.0, atol=1e-4)
    @test msg.dist.params[:w] == [1.0]
end

@testset "SPNonlinearSIn1MN" begin
    @test SPNonlinearSIn1MN <: SumProductRule{Nonlinear{Sampling}}
    @test outboundType(SPNonlinearSIn1MN) == Message{Function}
    @test isApplicable(SPNonlinearSIn1MN, [Message{Union{Bernoulli, Beta, Categorical, Dirichlet, Gaussian, Gamma, LogNormal, Poisson, Wishart}}, Nothing])

    log_pdf(x) = ruleSPNonlinearSIn1MN(f, Message(Univariate, GaussianMeanVariance, m=2.0, v=1.0), nothing, n_samples=0).dist.params[:log_pdf](x)
    @test log_pdf(1.5) == logPdf(ProbabilityDistribution(Univariate, GaussianMeanVariance, m=2.0, v=1.0), 1.5)
end

@testset "SPNonlinearSOutNGX" begin
    @test SPNonlinearSOutNGX <: SumProductRule{Nonlinear{Sampling}}
    @test outboundType(SPNonlinearSOutNGX) == Message{SampleList}
    @test !isApplicable(SPNonlinearSOutNGX, [Nothing, Message{Poisson}])
    @test isApplicable(SPNonlinearSOutNGX, [Nothing, Message{Gaussian}, Message{Gaussian}])

    msg = ruleSPNonlinearSOutNGX(h, nothing, Message(GaussianMeanVariance, m=1.0, v=tiny), Message(GaussianMeanVariance, m=2.0, v=tiny), n_samples=1)
    @test isapprox(msg.dist.params[:s][1], 3.0, atol=1e-4)
    @test msg.dist.params[:w] == [1.0]
end

@testset "SPNonlinearSInGX" begin
    @test SPNonlinearSInGX <: SumProductRule{Nonlinear{Sampling}}
    @test outboundType(SPNonlinearSInGX) == Message{GaussianWeightedMeanPrecision}
    @test !isApplicable(SPNonlinearSInGX, [Nothing, Message{Gamma}])
    @test isApplicable(SPNonlinearSInGX, [Message{Gaussian}, Message{Gaussian}, Nothing])

    msg_out = Message(Univariate, GaussianMeanVariance, m=2.0, v=1.0)
    msg_in1 = Message(Univariate, GaussianMeanVariance, m=0.0, v=1.0)
    msg_in2 = Message(Univariate, GaussianMeanVariance, m=1.0, v=1.0)

    res = ruleSPNonlinearSInGX(h, 1, msg_out, msg_in1, msg_in2, n_samples=1000)

    @test isapprox(mean(res.dist), 0.9998084187686186, atol=0.1)
    @test isapprox(var(res.dist), 1.9999999999999982, atol=0.1)
end

@testset "prod!" begin
    d = prod!(ProbabilityDistribution(Multivariate, Function, log_pdf=(s)->s), ProbabilityDistribution(Multivariate, Function, log_pdf=(s)->s))
    @test isa(d, ProbabilityDistribution{Multivariate,Function})
    @test d.params[:log_pdf](1) == 2
end

@testset "SPNonlinearSOutNMX" begin
    @test SPNonlinearSOutNMX <: SumProductRule{Nonlinear{Sampling}}
    @test outboundType(SPNonlinearSOutNMX) == Message{SampleList}
    @test !isApplicable(SPNonlinearSOutNMX, [Nothing, Message{Gamma}])
    @test !isApplicable(SPNonlinearSOutNMX, [Message{Gaussian}, Message{Gaussian}, Nothing])
    @test !isApplicable(SPNonlinearSOutNMX, [Message{Gaussian}, Message{Gamma}, Nothing])
    @test isApplicable(SPNonlinearSOutNMX, [Nothing, Message{Gaussian}, Message{Gamma}])
end

@testset "SPNonlinearSInMX" begin
    @test SPNonlinearSInMX <: SumProductRule{Nonlinear{Sampling}}
    @test outboundType(SPNonlinearSInMX) == Message{Function}
    @test !isApplicable(SPNonlinearSInMX, [Message{FactorFunction}, Message{Gamma}])
    @test !isApplicable(SPNonlinearSInMX, [Message{Gaussian}, Message{Gaussian}, Nothing])
    @test !isApplicable(SPNonlinearSInMX, [Nothing, Message{Gaussian}, Message{Gamma}])
    @test isApplicable(SPNonlinearSInMX, [Message{FactorFunction}, Nothing, Message{Gaussian}, Message{Gamma}])
end

@testset "MNonlinearSInMGX" begin
    @test MNonlinearSInMGX <: MarginalRule{Nonlinear{Sampling}}
    @test isApplicable(MNonlinearSInMGX, [Nothing, Message{Gaussian}, Message{Gaussian}])
    @test !isApplicable(MNonlinearSInMGX, [Nothing, Message{Gaussian}])
    @test !isApplicable(MNonlinearSInMGX, [Message{Gaussian}, Message{Gaussian}, ProbabilityDistribution])
end

#------------
# Integration
#------------

@testset "Nonlinear integration via sampling" begin
    fg = FactorGraph()

    @RV x ~ GaussianMeanVariance(2.0, 1.0)
    @RV y ~ GaussianMeanVariance(2.0, 3.0)
    n = Nonlinear{Sampling}(y, x, g=g, n_samples=2000)

    pfz = PosteriorFactorization(fg)
    algo = messagePassingAlgorithm(y)
    code = algorithmSourceCode(algo)
    @test occursin("ruleSPNonlinearSOutNM(g, nothing, messages[2], n_samples=2000)", code)
end

@testset "Nonlinear integration via sampling" begin
    fg = FactorGraph()

    @RV x ~ GaussianMeanVariance(2.0, 1.0)
    @RV y ~ GaussianMeanVariance(2.0, 3.0)
    @RV z ~ GaussianMeanVariance(2.0, 3.0)
    n = Nonlinear{Sampling}(z, x, y, g=g)

    # Forward; g_inv should not be present in call
    pfz = PosteriorFactorization(fg)
    algo = messagePassingAlgorithm(y)
    code = algorithmSourceCode(algo)

    @test occursin("ruleSPNonlinearSInGX(g, 2, messages[3], messages[2], messages[1])", code)
end

@testset "Nonlinear integration via sampling with specified variate types" begin
    fg = FactorGraph()

    @RV x ~ GaussianMeanVariance([2.0], mat(1.0))
    @RV y ~ GaussianMeanVariance([2.0], mat(3.0))
    @RV z ~ GaussianMeanVariance([2.0], mat(3.0))
    n = Nonlinear{Sampling}(z, x, y, g=g, dims=[(1,), (1,), (1,)])

    # Forward; g_inv should not be present in call
    pfz = PosteriorFactorization(fg)
    algo = messagePassingAlgorithm(y)
    code = algorithmSourceCode(algo)

    @test occursin("ruleSPNonlinearSInGX(g, 2, messages[3], messages[2], messages[1], dims=(1,))", code)
end

end # module


"""
$(TYPEDEF)
$(TYPEDFIELDS)

"""
struct QuantLogNormal <: LogNormalQuantDist
    mu::NestedFloat
    sigma::NestedFloat
    q::NestedFloat
end

function quantlognormal(mu::Float64, sigma::Float64, q::Float64)::Float64
    return round(exp(normal(mu, sigma)) / q) * q
end

lognormalquant(x::QuantLogNormal, mu, sigma, q) = quantlognormal(mu, sigma, q)


using DataFrames

export solve_poisson!

"""
   solve_poisson!( efield, particle_group, kernel_smoother, maxwell_solver, rho )

Accumulate rho and solve Poisson
 - `particle_group` : Particles
 - `maxwell_solver` : Maxwell solver (FEM 1D)
 - `kernel_smoother_0` : Particle-Mesh method
 - `rho` : preallocated array for Charge density
 - `efield_dofs` : spline coefficients of electric field (1D)
"""
function solve_poisson!(
    efield_dofs::Vector{Float64},
    particle_group::AbstractParticleGroup,
    kernel_smoother_0::AbstractParticleMeshCoupling,
    maxwell_solver::Maxwell1DFEM,
    rho::Vector{Float64},
)
    fill!(rho, 0.0)

    for i_part in 1:(particle_group.n_particles)
        xi = particle_group.array[1, i_part]
        wi = get_charge(particle_group, i_part)
        add_charge!(rho, kernel_smoother_0, xi, wi)
    end

    return compute_e_from_rho!(efield_dofs, maxwell_solver, rho)
end

"""
    pic_diagnostics_transfer( particle_group, kernel_smoother_0,
                            kernel_smoother_1, efield_dofs, transfer)

Compute ``\\sum_{particles} w_p ( v_1,p e_1(x_p) + v_2,p e_2(x_p)) ``

- `particle_group`
- `kernel_smoother_0`  : Kernel smoother (order p+1)
- `kernel_smoother_1`  : Kernel smoother (order p)
- `efield_dofs` : coefficients of efield

"""
function pic_diagnostics_transfer(
    particle_group::ParticleGroup{1,2}, kernel_smoother_0, kernel_smoother_1, efield_dofs
)
    transfer = 0.0::Float64

    for i_part in 1:(particle_group.n_particles)
        xi = particle_group.array[1, i_part]
        wi = get_charge(particle_group, i_part)
        v1 = particle_group.array[2, i_part]
        v2 = particle_group.array[3, i_part]

        efield_1 = evaluate(kernel_smoother_1, xi, efield_dofs[1])
        efield_2 = evaluate(kernel_smoother_0, xi, efield_dofs[2])

        transfer += (v1 * efield_1 + v2 * efield_2) * wi
    end

    return transfer
end

"""
    pic_diagnostics_vvb( particle_group, kernel_smoother_1, bfield_dofs )

Compute ``\\sum_{particles} ( w_p v_1, p b(x_p) v_2, p )``

- `particle_group`    : particle group object
- `kernel_smoother_1` : Kernel smoother (order p)
- `bfield_dofs` : coefficients of bfield

"""
function pic_diagnostics_vvb(particle_group, kernel_smoother_1, bfield_dofs)
    vvb = 0.0::Float64

    for i_part in 1:(particle_group.n_particles)
        xi = particle_group.array[1, i_part]
        v1 = particle_group.array[2, i_part]
        v2 = particle_group.array[3, i_part]
        wi = particle_group.array[4, i_part]
        wi *= particle_group.charge
        wi *= particle_group.common_weight

        bfield = evaluate(kernel_smoother_1, xi[1], bfield_dofs)

        vvb += wi * v1 * v2 * bfield
    end

    return vvb
end

"""
    pic_diagnostics_poynting( maxwell_solver, degree, efield_dofs, bfield_dofs,
                              rho )

Compute ``e^T M_0^{-1}  R^T b``

- `maxwell_solver` : maxwell solver object
- `degree` : degree of finite element
- `efield_dofs` : coefficients of `efield`
- `bfield_dofs` : coefficients of `bfield`

"""
function pic_diagnostics_poynting(maxwell_solver, degree, efield_dofs, bfield_dofs)
    scratch = similar(bfield_dofs)
    # Multiply B by M_0^{-1}  R^T
    compute_e_from_b!(scratch, maxwell_solver, 1.0, bfield_dofs)

    return inner_product(maxwell_solver, efield_dofs, scratch, degree)
end

export TimeHistoryDiagnostics

"""
    TimeHistoryDiagnostics( particle_group, maxwell_solver,
                            kernel_smoother_0, kernel_smoother_1 )

Context to save and plot diagnostics

- `particle_group` : Particles data
- `maxwell_solver` : Maxwell solver
- `kernel_smoother_0` : Mesh coupling operator
- `kernel_smoother_1` : Mesh coupling operator
- `data` : DataFrame containing time history values
"""
struct TimeHistoryDiagnostics
    particle_group::ParticleGroup
    maxwell_solver::Maxwell1DFEM
    kernel_smoother_0::ParticleMeshCoupling1D
    kernel_smoother_1::ParticleMeshCoupling1D
    data::DataFrame
    diagnostics::Vector{Float64}
    potential_energy::Vector{Float64}

    function TimeHistoryDiagnostics(
        particle_group::ParticleGroup,
        maxwell_solver::Maxwell1DFEM,
        kernel_smoother_0::ParticleMeshCoupling1D,
        kernel_smoother_1::ParticleMeshCoupling1D,
    )
        data = DataFrame(;
            Time=Float64[],
            KineticEnergy=Float64[],
            Momentum1=Float64[],
            Momentum2=Float64[],
            PotentialEnergyE1=Float64[],
            PotentialEnergyE2=Float64[],
            PotentialEnergyB3=Float64[],
            Transfer=Float64[],
            VVB=Float64[],
            Poynting=Float64[],
            ErrorPoisson=Float64[],
        )

        diagnostics = zeros(Float64, 3)
        potential_energy = zeros(Float64, 3)

        return new(
            particle_group,
            maxwell_solver,
            kernel_smoother_0,
            kernel_smoother_1,
            data,
            diagnostics,
            potential_energy,
        )
    end
end

export write_step!

"""
    write_step!( thdiag, time, degree, efield_dofs, bfield_dofs,
                 efield_dofs_n, efield_poisson)

write diagnostics for PIC
- `time` : Time
- `efield_dofs` : Electric field
- `efield_dofs_n` : Electric field at half step
- `efield_poisson` : Electric field compute from Poisson equation
- `bfield_dofs` : Magnetic field
- `degree` : Spline degree
"""
function write_step!(
    thdiag::TimeHistoryDiagnostics,
    time,
    degree,
    efield_dofs,
    bfield_dofs,
    efield_dofs_n,
    efield_poisson,
)
    fill!(thdiag.diagnostics, 0.0)
    fill!(thdiag.potential_energy, 0.0)

    for i_part in 1:(thdiag.particle_group.n_particles)
        v1 = thdiag.particle_group.array[2, i_part]
        v2 = thdiag.particle_group.array[3, i_part]
        wi = thdiag.particle_group.array[4, i_part]
        wi *= thdiag.particle_group.mass
        wi *= thdiag.particle_group.common_weight

        # Kinetic energy
        thdiag.diagnostics[1] += (v1^2 + v2^2) * wi
        # Momentum 1
        thdiag.diagnostics[2] += v1 * wi
        # Momentum 2
        thdiag.diagnostics[3] += v2 * wi
    end

    transfer = pic_diagnostics_transfer(
        thdiag.particle_group,
        thdiag.kernel_smoother_0,
        thdiag.kernel_smoother_1,
        efield_dofs,
    )

    vvb = pic_diagnostics_vvb(thdiag.particle_group, thdiag.kernel_smoother_1, bfield_dofs)

    poynting = pic_diagnostics_poynting(
        thdiag.maxwell_solver, degree, efield_dofs[2], bfield_dofs
    )

    thdiag.potential_energy[1] = inner_product(
        thdiag.maxwell_solver, efield_dofs[1], efield_dofs_n[1], degree - 1
    )

    thdiag.potential_energy[2] = inner_product(
        thdiag.maxwell_solver, efield_dofs[2], efield_dofs_n[2], degree
    )

    thdiag.potential_energy[3] = l2norm_squared(
        thdiag.maxwell_solver, bfield_dofs, degree - 1
    )

    return push!(
        thdiag.data,
        (
            time,
            thdiag.diagnostics...,
            thdiag.potential_energy...,
            transfer,
            vvb,
            poynting,
            maximum(abs.(efield_dofs[1] .- efield_poisson)),
        ),
    )
end

"""
    evaluate( kernel_smoother, field_dofs,  xi, n_dofs )

Evaluate the field at points xi

- `field_dofs` : field value on dofs
- `xi` : positions where the field is evaluated
"""
function evaluate(
    kernel_smoother::ParticleMeshCoupling1D, field_dofs::AbstractArray, x::AbstractArray
)
    field_grid = similar(x)
    for j in eachindex(x)
        field_grid[j] = evaluate(kernel_smoother, x[j], field_dofs)
    end
    return field_grid
end

"""
    pic_diagnostics_hpi( particle_group,  index, kinetic )

compute v(index)-part of kinetic energy

- `particle_group`
- `index` : velocity component
- `kinetic` : value of index part of kinetic energy

"""
function pic_diagnostics_hpi(particle_group, index, kinetic)
    kinetic = 0.0
    for i_part in 1:(particle_group.n_particles)
        vi = get_v(particle_group, i_part)
        wi = get_mass(particle_group, i_part)
        kinetic += kinetic_local + (vi[index]^2) * wi[1]
    end

    return kinetic
end

"""
    eval_derivative_spline( position, xmin, delta_x, n_grid,
                            field_dofs, degree, derivative )

Compute the spline coefficient of the derivative of some given spline expansion

- `position` : particle position
- `xmin` : lower boundary of the domain
- `delta_x` : step
- `n_grid` : number of grid points
- `field_dofs` : coefficients of spline representation of the field
- `degree` : degree of spline
- `derivative` : value of the derivative
"""
function eval_derivative_spline(position, xmin, delta_x, n_grid, field_dofs, degree)
    der_degree = degree - 1

    xi = (position[1] - xmin) / delta_x
    index = ceil(xi)
    xi = xi - (index - 1)
    index = index - der_degree

    spline_val = uniform_bsplines_eval_basis(der_degree, xi)

    derivative = 0.0

    for i1 in 1:degree
        ind = mod(index + i1 - 2, n_grid) + 1
        derivative +=
            spline_val[i1] * (field_dofs[ind] - field_dofs[mod(ind - 2, n_grid) + 1])
    end

    return derivative / delta_x
end


function factors(n)
  f = [one(n)]
  for (p,e) in factor(n)
    f = reduce(vcat, f, [f*p^j for j in 1:e])
  end
  return length(f) == 1 ? [one(n), n] : sort!(f)
end


"""
    `build_blp_model(bp::BilevelLP, solver; [comp_method])`

Build a JuMP model (constraints and objective) based
on the data from bp and with a solver.

Returns a tuple `(m, x, y, λ, s)`.

An optional keyword `comp_method` can be passed for
handling complementarity constraints, default is `SOS1Complementarity`.
"""
function build_blp_model(bp::BilevelLP, optimizer; comp_method = SOS1Complementarity())
    m = JuMP.Model(optimizer)
    @variable(m, bp.xl[j] <= x[j=1:bp.nu] <= bp.xu[j])
    @variable(m, y[j=1:bp.nl])
    @constraint(m, uppercons[i=1:bp.mu],
        sum(bp.G[i,j]*x[j] for j in 1:bp.nu) +
        sum(bp.H[i,j]*y[j] for j in 1:bp.nl) <= bp.q[i])
    @objective(m, Min, sum(bp.cx .* x) + dot(bp.cy, y))
    for j in bp.Jx
        JuMP.set_integer(x[j])
    end
    # adding SOS1 constraints
    return build_blp_model(m, bp, x, y, comp_method = comp_method)
end

"""
    `build_blp_model(m::JuMP.Model, bp::BilevelLP, x, y; [comp_method])`

Adds the lower-level constraints and optimality conditions to
an existing JuMP model. This assumes the upper-level feasibility
constraints and objective have already been set.

Returns a tuple `(m, x, y, λ, s)`.

An optional keyword `comp_method` can be passed for
handling complementarity constraints, default is `SOS1Complementarity`.
"""
function build_blp_model(m::JuMP.Model, bp::BilevelLP, x, y; comp_method = SOS1Complementarity())
    @variable(m, s[1:bp.ml] >= 0) # lower-level slack variables
    @constraint(m, lowercons[i=1:bp.ml],
        sum(bp.A[i,j]*x[j] for j in 1:bp.nu) +
        sum(bp.B[i,j]*y[j] for j in 1:bp.nl) + s[i] == bp.b[i]
    )
    @variable(m, λ[1:bp.ml] >= 0)
    @variable(m, σ[1:bp.nl] >= 0) # dual of lower-level lower bound
    for j in Base.OneTo(bp.nl)
        if !bp.yl[j]
            # sigma at 0 if free variable
            @constraint(m, σ[j] == 0)
        else
            JuMP.set_lower_bound(y[j], 0)
        end
    end
    @constraint(m, [j in 1:bp.nl], bp.d[j] + dot(bp.F[:,j], x) + dot(bp.B[:,j], λ)  - σ[j] == 0)
    add_complementarity_constraint(m, comp_method, s, λ, y, σ)
    return (m, x, y, λ, s)
end

"""
    `build_blp_model(m, B, d, x, y, s, F; [comp_method])`

Build the bilevel JuMP model from the data without
grouping everything into a `BilevelLP`.
An optional keyword `comp_method` can be passed for
handling complementarity constraints, default is `SOS1Complementarity`.
"""
function build_blp_model(m::JuMP.Model, B::M, d, x, y, s, F; comp_method = SOS1Complementarity()) where {M<:MT}
    @variable(m, λ[1:bp.ml] >= 0)
    @variable(m, σ[1:bp.nl] >= 0.) # dual of lower-level lower bound
    for j in Base.OneTo(bp.nl)
        if !bp.yl[j]
            # sigma at 0 if free variable
            @constraint(m, σ[j] == 0.)
        else
            JuMP.set_lower_bound(y[j], 0.)
        end
    end
    @constraint(m, d .+ F' * x .+ B' * λ .- σ .== 0)
    add_complementarity_constraint(m, comp_method, s, λ, y, σ)
    return (m, λ, s)
end

function build_blp_model(m::JuMP.Model, B::M, d, s; comp_method = SOS1Complementarity()) where {M<:MT}
    (ml, nl) = size(B)
    @variable(m, λ[1:ml] >= 0)
    @constraint(m, kkt[j=1:nl], d[j] + dot(B[:,j], λ) == 0)
    add_complementarity_constraint(m, comp_method, s, λ, [], [])
    return (m, λ, s)
end


using PyPlot
using JLD
using Optim

# type Measurements{T}
#     i::Int64                # measurement counter
#     t::Array{Float64}       # time data (when it was received by this recorder)
#     t_msg::Array{Float64}   # time that the message was sent
#     z::Array{T}             # measurement values
# end

function main(code::AbstractString)
    #log_path_record = ("$(homedir())/open_loop/output-record-8bbb.jld","$(homedir())/open_loop/output-record-3593.jld")
    log_path_record = "$(homedir())/open_loop/output-record-$code.jld"
    d_rec = load(log_path_record)

    cmd_pwm_log = d_rec["cmd_pwm_log"]
    vel_est     = d_rec["vel_est"]
    pos_info    = d_rec["pos_info"]

    t0 = max(cmd_pwm_log.t[1],vel_est.t[1])
    t_end = min(cmd_pwm_log.t[end],vel_est.t[end])

    t = t0+0.1:.02:t_end-0.1
    v = zeros(length(t))
    cmd = zeros(length(t))

    for i=1:length(t)
        v[i] = vel_est.z[t[i].>vel_est.t,1][end]
        #v[i] = pos_info.z[t[i].>pos_info.t,15][end]
        cmd[i] = cmd_pwm_log.z[t[i].>cmd_pwm_log.t,1][end]
    end
    v_opt   = v[1:end]
    cmd_opt = cmd[1:end]
    v_opt   = v[cmd.>90]
    cmd_opt = cmd[cmd.>90]

    res = optimize(c->cost(cmd_opt,v_opt,c,0),[0.001,0,0.1,0.001,0])
    c = Optim.minimizer(res)
    c2 = [1/c[1], c[2]/c[1], 1/c[4], c[5]/c[4]]
    cost(cmd_opt,v_opt,c,1)
    println("c = $c")
    println("and c_2 = $c2")
end

function v_over_u(code::AbstractString)
    log_path_record = "$(homedir())/open_loop/output-record-$code.jld"
    d_rec = load(log_path_record)

    cmd_pwm_log = d_rec["cmd_pwm_log"]
    vel_est     = d_rec["vel_est"]
    pos_info    = d_rec["pos_info"]
    imu_meas    = d_rec["imu_meas"]

    t0 = max(cmd_pwm_log.t[1],vel_est.t[1])
    t_end = min(cmd_pwm_log.t[end],vel_est.t[end])

    L_b = 0.125
    t = t0+0.1:.02:t_end-0.1
    v = zeros(length(t))
    cmd = zeros(length(t),2)
    psiDot = zeros(length(t))

    for i=1:length(t)
        v[i] = vel_est.z[t[i].>vel_est.t,1][end]
        #v[i] = pos_info.z[t[i].>pos_info.t,15][end]
        cmd[i,:] = cmd_pwm_log.z[t[i].>cmd_pwm_log.t,:][end,:]
        psiDot[i] = imu_meas.z[t[i].>imu_meas.t,3][end]
    end
    v_x = real(sqrt(complex(v.^2-psiDot.^2*L_b^2)))
    v_y = L_b*psiDot
    delta = atan2(psiDot*0.25,v_x)
    c = ((cmd[:,1]-91)/6.5-0.5*v)./delta.^2
    c = c[abs(c).<100]
    plot(c)
    grid("on")
    figure()
    plot(v)
    grid("on")
    #plot(cmd[:,1],v)
end

function cost(cmd,v,c,p)
    cost = 0
    v_s = zeros(length(v))
    v_s[1] = v[1]
    v_s[2] = v[2]
    for i=2:length(cmd)-1
        if cmd[i]-cmd[i-1] != 0
            v_s[i] = v[i]
        end
        if cmd[i] > 90
            v_s[i+1] = v_s[i] + 0.02*(c[1]*cmd[i] + c[2] - c[3]*v_s[i])
        else
            v_s[i+1] = max(v_s[i] + 0.02*(c[4]*cmd[i] + c[5] - c[3]*v_s[i]),0)
        end
    end
    if p==1
        plot((1:length(v_s))*0.02,v_s,(1:length(v))*0.02,v)
        grid("on")
        legend(["v_sim","v_meas"])
        title("Curve fitting acceleration")
        xlabel("t [s]")
        ylabel("v [m/s]")
    end
    cost = norm(v_s-v)
    return cost
end


using Plots
using Random
Random.seed!(1)

include_ddp()

# Model
include_model("double_integrator")

function f(model::DoubleIntegratorContinuous, x, u, w)
    [x[2] + w[1]; (1.0 + w[3]) * u[1] + w[2]]
end

function fd(model::DoubleIntegratorContinuous{Midpoint, FixedTime}, x, u, w, h, t)
	if t == 1
		x1 = u[2:3]
    	return x1 + h * f(model, x1 + 0.5 * h * f(model, x1, u, w), u, w)
	else
		return x + h * f(model, x + 0.5 * h * f(model, x, u, w), u, w)
	end
end

model = DoubleIntegratorContinuous{Midpoint, FixedTime}(2, 1, 3)
n = [model.n for t = 1:T]
m = [t == 1 ? model.m + model.n : model.m for t = 1:T]

# Time
T = 101
h = 0.1

z = range(0.0, stop = 3.0 * 2.0 * π, length = T)
p_ref = 1.0 * cos.(1.0 * z)
plot(z, p_ref)

# Initial conditions, controls, disturbances
x1 = [p_ref[1]; 0.0]
x1_alt = [0.0; 0.0]
xT = [[p_ref[t]; 0.0] for t = 1:T]
ū = [rand(m[t]) for t = 1:T-1]
w = [zeros(model.d) for t = 1:T-1]

# Rollout
x̄ = rollout(model, x1, ū, w, h, T)

# Objective
Q = [(t < T ? h : 1.0) * Diagonal([100.0; 0.1]) for t = 1:T]
q = [-2.0 * Q[t] * xT[t] for t = 1:T]
R = [h * Diagonal(0.01 * ones(m[t])) for t = 1:T-1]
r = [zeros(m[t]) for t = 1:T-1]
obj = StageCosts([QuadraticCost(Q[t], q[t],
	t < T ? R[t] : nothing, t < T ? r[t] : nothing) for t = 1:T], T)

function g(obj::StageCosts, x, u, t)
	T = obj.T
    if t < T
		Q = obj.cost[t].Q
		q = obj.cost[t].q
	    R = obj.cost[t].R
		r = obj.cost[t].r
        return x' * Q * x + q' * x + u' * R * u + r' * u
    elseif t == T
		Q = obj.cost[T].Q
		q = obj.cost[T].q
        return x' * Q * x + q' * x
    else
        return 0.0
    end
end

g(obj, x̄[1], ū[1], 1)
g(obj, x̄[2], ū[2], 2)
g(obj, x̄[T], nothing, T)
objective(obj, x̄, ū)

# Constraints
p = [t == 1 ? n[1] : 0 for t = 1:T]
info_t = Dict(:x1 => x1_alt)#:ul => [-5.0], :uu => [5.0], :inequality => (1:2 * m))
info_T = Dict()#:xT => xT)
con_set = [StageConstraint(p[t], t < T ? info_t : info_T) for t = 1:T]

function c!(c, cons::StageConstraints, x, u, t)
	T = cons.T
	if t == 1
		c .= view(u, 2:3) - cons.con[t].info[:x1]
	end
end

prob = problem_data(model, obj, con_set, copy(x̄), copy(ū), w, h, T, n = n, m = m)

# Solve
@time constrained_ddp_solve!(prob,
    max_iter = 1000, max_al_iter = 6,
	ρ_init = 1.0, ρ_scale = 10.0)

# # Solve
# @time ddp_solve!(prob,
#     max_iter = 100, verbose = true)

x, u = current_trajectory(prob)
x̄, ū = nominal_trajectory(prob)

x[1] = u[1][2:3]
x̄[1] = ū[1][2:3]

# Visualize
using Plots
plot(hcat([[p_ref[t]; 0.0] for t = 1:T]...)',
    width = 2.0, color = :black, label = "")
plot!(hcat(x...)[:, 1:T]', color = :magenta, label = "")
# plot(hcat(u..., u[end])', linetype = :steppost)


"""
    expval1(basis::SingleBlockBasis, op::AbstractMatrix, i::Int, wf::AbstractVector{Float64})

Calculate the expectation value of the one-body operator `op` at site `i` in
the state `wf`.
"""
function expval1(basis::SingleBlockBasis{N}, op::AbstractMatrix, i::Int, wf::AbstractVector{Float64}) where {N}
    w = Array{Int}(undef, 2N)

    result = 0.0

    for col in 1:basis.block.size
        v = basis.block.vectors[col]

        li = v[2i-1]
        mi = v[2i]
        coli = li*(li+1) + mi + 1

        rowi = 0
        for lip in 0:basis.l_max
            for mip in (-lip):lip
                rowi += 1
                op[rowi, coli] == 0.0 && continue

                w .= v
                w[2i-1] = lip
                w[2i] = mip

                haskey(basis.block.lookup, w) || continue

                result += wf[basis.block.lookup[w]] * wf[col] * op[rowi, coli]
            end
        end
    end

    result
end

"""
    expval2(basis::SingleBlockBasis, op_i::AbstractMatrix, i::Int, op_j::AbstractMatrix, j::Int, wf::AbstractVector{Float64})

Calculate the expectation value of the two-body operator composed of `op_i` at
site `i` and `op_j` at site `j` in the state `wf`.
"""
function expval2(basis::SingleBlockBasis{N}, op_i::AbstractMatrix, i::Int, op_j::AbstractMatrix, j::Int, wf::AbstractVector{Float64}) where {N}
    w = Array{Int}(undef, 2N)

    result = 0.0

    for col in 1:basis.block.size
        v = basis.block.vectors[col]

        li = v[2i-1]
        mi = v[2i]
        coli = li*(li+1) + mi + 1
        lj = v[2j-1]
        mj = v[2j]
        colj = lj*(lj+1) + mj + 1

        rowi = 0
        for lip in 0:basis.l_max
            for mip in (-lip):lip
                rowi += 1
                op_i[rowi, coli] == 0.0 && continue

                rowj = 0
                for ljp in 0:basis.l_max
                    for mjp in (-ljp):ljp
                        rowj += 1
                        op_j[rowj, colj] == 0.0 && continue

                        w .= v
                        w[2i-1] = lip
                        w[2i] = mip
                        w[2j-1] = ljp
                        w[2j] = mjp

                        haskey(basis.block.lookup, w) || continue

                        result += wf[basis.block.lookup[w]] * wf[col] * op_i[rowi, coli] * op_j[rowj, colj]
                    end
                end
            end
        end
    end

    result
end


"""
$(TYPEDEF)

Represents a simple control system with linear dynamics.

# Parameters:

- `ΔT`: the sampling time of the system (se)
- `nx`: the number of states
- `nu`: the number of inputs

# Fields

$(TYPEDFIELDS)
"""
struct LinearSystem{ΔT, nx, nu, TA<:SMatrix{nx, nx}, TB<:SMatrix{nx, nu}} <: ControlSystem{ΔT, nx, nu}
    "The state transition matrix"
    A::TA
    "The control input matrix"
    B::TB
end

LinearSystem{ΔT}(A::TA, B::TB) where {ΔT, nx, nu, TA<:SMatrix{nx, nx}, TB<:SMatrix{nx, nu}} = LinearSystem{ΔT, nx, nu, TA, TB}(A, B)
LinearizationStyle(::LinearSystem) = TrivialLinearization()

dx(ls::LinearSystem, x::SVector, u::SVector, t::AbstractFloat)  = begin @assert !issampled(ls); ls.A*x + ls.B*u end
next_x(ls::LinearSystem, x::SVector, u::SVector) = begin @assert issampled(ls); ls.A*x + ls.B*u end
linearize(ls::LinearSystem, x::SVector, u::SVector, t::AbstractFloat) = ls
function discretize(ls::LinearSystem, vt::Val{ΔT}) where {ΔT}
    @assert !issampled(ls) "Can't discretize a discrete system."
    @assert ΔT > 0 "Discrtization requires finite sampling time ΔT."
    # use the euler approximation to avoid expensive matrix exponential
    return LinearSystem{ΔT}(I+ΔT*ls.A, ΔT*ls.B)
end

# TODO: this should probably know the absolute time (like the `SystemTrajectory`)
struct LTVSystem{ΔT, nx, nu, TD<:AbstractVector{<:LinearSystem{ΔT, nx, nu}}} <: ControlSystem{ΔT, nx, nu}
    "The discrete time series of linear systems."
    dyn::TD
    LTVSystem(dyn::TD) where {ΔT, nx, nu, TD<:AbstractVector{<:LinearSystem{ΔT, nx, nu}}} = begin
        @assert ΔT > 0 "LTVSystem require finite discretization steps."
        new{ΔT, nx, nu, TD}(dyn)
    end
end
LinearizationStyle(::LTVSystem) = TrivialLinearization()
Base.eltype(::Type{<:LTVSystem{ΔT, nx, nu, TD}}) where {ΔT, nx, nu, TD} = eltype(TD)
Base.getindex(ds::LTVSystem, i) = getindex(ds.dyn, i)
Base.setindex!(ds::LTVSystem, v, i) = setindex!(ds.dyn, v, i)
next_x(cs::LTVSystem, xₖ::SVector, uₖ::SVector, k::Int) = next_x(cs.dyn[k], xₖ, uₖ)

# TODO: we could also get away without this type but it seems more clean (?)
struct LTISystem{ΔT,nx,nu,TL<:LinearSystem{ΔT,nx,nu},
                 TXY,TXI}<:ControlSystem{ΔT,nx,nu}
    dyn::TL
    xyids::TXY
    xids::TXI
end
LTISystem(dyn::LinearSystem, xyids) = LTISystem(dyn, xyids, nothing)

LinearizationStyle(::LTISystem) = TrivialLinearization()
xyindex(cs::LTISystem) = cs.xyids
xindex(cs::LTISystem) = cs.xids
Base.eltype(::Type{<:LTISystem{ΔT,nx,nu,TL}}) where {ΔT,nx,nu,TL} = TL
Base.getindex(cs::LTISystem, i) = cs.dyn
next_x(cs::LTISystem, xₖ::SVector, uₖ::SVector, ::Float64) = next_x(cs.dyn, xₖ, uₖ)


include("graph.jl")
include("BFS.jl")
M = [0 1 0 0 0 1 1;
     0 0 0 0 0 1 0;
     0 0 0 0 0 0 1;
     0 0 1 0 1 0 0;
     0 0 0 0 0 0 0;
     0 0 0 1 1 0 1;
     0 0 0 0 1 0 0]
s = 1;
f = 3;
G = nonEmptyGraph(M)
BFS!(G, s)
path = getPath(G, s, f)
println("Path = $path")
printGraph(G)


@testset "Jacobian" begin
    graph = Compgraph()
    c = [1.1, 1.2, 1.3, 1.4, 1.5, 1.6]
    add_lincomb!(graph, :B1, c[1], :I, c[2], :A)
    add_lincomb!(graph, :B2, c[3], :I, c[4], :A)
    add_mult!(graph, :N, :B1, :B2)
    add_lincomb!(graph, :D, c[5], :I, c[6], :A)
    add_ldiv!(graph, :R, :D, :N)
    add_output!(graph, :R)

    x = randn(5)
    @test eval_graph(graph, x) ≈
          (c[1] .+ c[2] * x) .* (c[3] .+ c[4] * x) ./ (c[5] .+ c[6] * x)

    dg_dc3 = eval_jac(graph, x, (:B2, 1))
    @test dg_dc3 ≈ (c[1] .+ c[2] * x) ./ (c[5] .+ c[6] * x)

    dg_dc6 = eval_jac(graph, x, (:D, 2))
    @test dg_dc6 ≈
          -x .* (c[1] .+ c[2] * x) .* (c[3] .+ c[4] * x) ./
          (c[5] .+ c[6] * x) .^ 2

    dg_dc2c5 = eval_jac(graph, x, [(:B1, 2), (:D, 1)])
    @test dg_dc2c5 ≈ hcat(
        x .* (c[3] .+ c[4] * x) ./ (c[5] .+ c[6] * x),
        -(c[1] .+ c[2] * x) .* (c[3] .+ c[4] * x) ./ (c[5] .+ c[6] * x) .^ 2,
    )
end