@testset "Rastrigin" begin

    function test_result(result::Evolutionary.EvolutionaryOptimizationResults, N::Int, tol::Float64)
        fitness = minimum(result)
        extremum = Evolutionary.minimizer(result)
        if round(fitness) == 0
            @test extremum ≈ zeros(N) atol=tol
            @test fitness ≈ 0.0 atol=tol
        else
            # @warn("Found local minimum!!!")
            @test sum(abs, extremum) < N
        end
    end

    # Objective function
    function rastrigin(x::AbstractVector{T}) where {T <: AbstractFloat}
        n = length(x)
        return 10n + sum([ x[i]^2 - 10cos(convert(T,2π*x[i])) for i in 1:n ])
    end

    # Parameters
    N = 3
    P = 100
    initState = ()->rand(N)

    # Testing: (μ/μ_I,λ)-σ-Self-Adaptation-ES
    # with non-isotropic mutation operator y' := y + (σ_1 N_1(0, 1), ..., σ_N N_N(0, 1))
    m = ES(mu = 15, lambda = P)
    @test m.μ == 15
    @test m.λ == P
    result = Evolutionary.optimize( rastrigin,
        initState,
        ES(
            initStrategy = AnisotropicStrategy(N),
            recombination = average, srecombination = average,
            mutation = gaussian, smutation = gaussian,
            selection=:comma,
            μ = 15, λ = P
        ),Evolutionary.Options(iterations=1000, show_trace=false)
    )
    println("(15,$(P))-σ-SA-ES => F: $(minimum(result)), C: $(Evolutionary.iterations(result))")
    test_result(result, N, 1e-1)

    # Testing: CMA-ES
    result = Evolutionary.optimize(rastrigin, initState, CMAES(mu = 15, lambda = P))
    println("(15/15,$(P))-CMA-ES => F: $(minimum(result)), C: $(Evolutionary.iterations(result))")
    test_result(result, N, 1e-1)

    # Testing: GA
    m = GA(epsilon = 10.0)
    @test m.ɛ == 10

    selections = [:roulette=>rouletteinv, :sus=>susinv, :rank=>ranklinear(1.5)]
    crossovers = [:discrete=>discrete, :intermediate0=>intermediate(0.), :intermediate0_25=>intermediate(0.5), :line=>line(0.2)]
    mutations = [:domrng0_5=>domainrange(fill(0.5,N)), :uniform=>uniform(3.0), :gaussian=>gaussian(0.6)]

    @testset "GA settings" for (sn,ss) in selections, (xn,xovr) in crossovers, (mn,ms) in mutations
        xn == :discrete && (mn == :uniform || mn == :gaussian) && continue # bad combination
        xn == :line && mn == :gaussian && continue # bad combination
        result = Evolutionary.optimize( rastrigin, initState,
            GA(
                populationSize = P,
                ɛ = 0.1,
                selection = ss,
                crossover = xovr,
                mutation = ms
            ), Evolutionary.Options(iterations=1000, successive_f_tol=25)
        )
        println("GA:$(sn):$(xn):$(mn)(N=$(N),P=$(P),x=.8,μ=.1,ɛ=0.1) => F: $(minimum(result)), C: $(Evolutionary.iterations(result))")
        test_result(result, N, 1e-1)
    end

    # Testing: DE
    selections = [:rand=>random, :perm=>permutation, :rndoff=>randomoffset, :best=>best]
    mutations = [:exp=>exponential(0.5), :bin=>uniformbin(0.5)]

    @testset "DE settings" for (sn,ss) in selections, (mn,ms) in mutations, n in 1:2
        result = Evolutionary.optimize( rastrigin, initState,
            DE(
                populationSize = P,
                n=n,
                selection = ss,
                recombination = ms,
                F = 0.9
            )
        )
        println("DE/$sn/$n/$mn(F=0.9,Cr=0.5) => F: $(minimum(result)), C: $(Evolutionary.iterations(result))")
        test_result(result, N, 1e-2)
    end

end


@testset "RayClassGroup" begin

  #include(joinpath(Hecke.pkgdir, "examples", "RayClass.jl"))

  @testset "Big prime" begin
    k,  = wildanger_field(2, 13)
    Zk = maximal_order(k)
    p = next_prime(fmpz(2)^100)
    r, mr = ray_class_group(ideal(Zk, p), n_quo = 2)
    @test order(domain(mr)) == 2
  end

  @testset "quadratic fields" begin

    Qx,x=PolynomialRing(FlintQQ,"x")
    K,a=NumberField(x^2+199,"a")
    O=maximal_order(K)
    C,mC=class_group(O)

    for i=9:13
      R1,mR1=ray_class_group(ideal(O,i), n_quo=3)
      for r in R1
        @test r== mR1\(mR1(r))
      end
      R,mR=ray_class_group(ideal(O,i))
      q,mq=quo(R,3)
      @test Hecke.isisomorphic(R1,q)
    end

    K,a=NumberField(x^2-5,"a")
    O=maximal_order(K)
    C,mC=class_group(O, redo = true)
    inf_plc=real_places(K)

    for i=9:13
      R1,mR1=ray_class_group(ideal(O,i),inf_plc)
      for r in R1
        @test r== mR1\(mR1(r))
      end
      R2, mR2=ray_class_group(ideal(O,i), inf_plc, n_quo=2)
      for r in R2
        @test r== mR2\(mR2(r))
      end
      q1,mq1=quo(R1,2)
      @test Hecke.isisomorphic(q1,R2)
    end
  end

  @testset "infinite places" begin

     Qx, x = PolynomialRing(FlintQQ, "x");
     K, a = NumberField(x^4-4*x^3-11*x^2+30*x+20, cached = false)
    O = maximal_order(K)
    r, mr = ray_class_group(ideal(O,4), real_places(K), n_quo=2)
    @test order(r) == 2^5
    for el in r
      @test el == mr\(mr(el))
    end
  end

  @testset "stable subgroups" begin

     Qx,x=PolynomialRing(FlintQQ,"x");
    f=x^2+1;
    K,a=NumberField(f,"a");
    auts = automorphisms(K)
    auts = small_generating_set(auts, *)
    O = maximal_order(K);
    C, mC = class_group(O);
    r, mr = ray_class_group(ideal(O,3*5*7), n_quo=8);
    act = Hecke.induce_action(mr, auts);
    x=Hecke.stable_subgroups(r,act,op = quo, quotype = [8]);
    y=subgroups(r, quotype=[8])
    i=0
    for el in y
      if Hecke.isstable(act,el[2])
        i+=1
      end
    end
    @test length(x)==i

    x=Hecke.stable_subgroups(r, act, op = quo, quotype = [2,4]);
    y=subgroups(r, quotype=[2,4])
    i=0
    for el in y
      if Hecke.isstable(act,el[2])
        i+=1
      end
    end
    @test length(x)==i

    r,mr=ray_class_group(ideal(O,9*19*29), n_quo=9);
    act=Hecke.induce_action(mr, auts);
    x=Hecke.stable_subgroups(r, act, op = quo, quotype = [9]);
    y=subgroups(r, quotype=[9])
    i=0
    for el in y
      if Hecke.isstable(act,el[2])
        i+=1
      end
    end
    @test length(x)==i

    x=Hecke.stable_subgroups(r, act, op = quo, quotype = [3,9]);
    y=subgroups(r, quotype=[3,9])
    i=0
    for el in y
      if Hecke.isstable(act,el[2])
        i+=1
      end
    end
    @test length(x)==i
  end

end


#=
A regular log-λ grid is characterized by:
log10(λ_i+1) - log10(λ_i)  =  log10(λ_i+1 / λ_i)  =  costant step

hence:
λ_i+1 / λ_i  =  cost.
λ_i+1 / λ_i - 1  =  cost.
(λ_i+1 - λ_i) / λ_i  =  Δλ / λ  =  1/R   is constant ∀i

i.e. it has a constant spectral resolution given by:
R = c / σ_kms

The step of the grid is:
log10(σ_kms / c + 1)
=#

function conv_gauss(λ, data, σ_kms; oversampling::Int=1)
    @assert oversampling >= 1

    # Interpolate input data on the regular log-λ grid with proper resolution
    # check with: extrema([(grid_λ[i+1] - grid_λ[i]) / grid_λ[i] * 3e5 for i in 1:length(grid_λ)-1]) .* oversampling, σ_kms
    step = log10(σ_kms / 3e5 + 1) / oversampling
    grid_λ = 10. .^ range(log10(minimum(λ)), log10(maximum(λ)), step=step)
    grid_y = Spline1D(λ, data, k=1)(grid_λ)

    # Convolution with Gaussian kernel and extraction of proper subset
    kernel = gauss.(-5:1. / oversampling:5, 0., 1.)
    i1 = div(length(kernel), 2) + 1
    i2 = i1 + length(grid_y) - 1
    out = conv(grid_y, kernel)[i1:i2] ./ oversampling

    # Interpolate back to original domain
    out = Spline1D(grid_λ, out, k=1, bc="extrapolate")(λ)

    # Replace data close to the edges
    ee = (1 + 2σ_kms / 3e5) * minimum(λ)
    i = findall(λ .< ee)
    out[i] = Spline1D(grid_λ, grid_y, k=1, bc="extrapolate")(λ[i])

    ee = (1 - 2σ_kms / 3e5) * maximum(λ)
    i = findall(λ .> ee)
    out[i] = Spline1D(grid_λ, grid_y, k=1, bc="extrapolate")(λ[i])

    # @gp    λ data "w lp t 'Data'"
    # @gp :- λ out  "w lp t 'Convolution'"
    return out
end

function delta_conv_gauss(λ, data, σ_kms; oversampling::Int=1)
    step = log10(σ_kms / 3e5 + 1) / oversampling
    grid_λ = 10. .^ range(log10(minimum(λ) * (1 - 5 * σ_kms/3e5)),
                          log10(maximum(λ) * (1 + 5 * σ_kms/3e5)),
                          step=step)
    out = grid_λ .* 0.

    for ii in 1:length(data)
        out .+= data[ii] .* gauss(grid_λ, λ[ii], σ_kms / 3e5 * λ[ii])
    end

    return grid_λ, out
end

#=

x, y = QSFit.delta_conv_gauss([1e3, 3e3, 1e4, 3e4], [0., 1., 1., 0.], 3000.)
QSFit.int_tabulated(x, y)  # should be ~ 2
@gp xlog=true x y "w lp"

y2 = QSFit.conv_gauss(x, y, 10000.)
@gp :- x y2 "w lp"
QSFit.int_tabulated(x, y)  # still ~ 2

=#


using Plots, Statistics, MLDatasets, LinearAlgebra, Clustering, Printf, Dates, Random
path1=realpath(dirname(@__FILE__)*"/../..")
include(string(path1,"/src/ANNBN.jl"))

# Polynomial
function ff(x)
    -x[1]^3+5x[1]^2+5x[1]*x[2]
end

i_train=1000
vars=2
xx_train=0.8*rand(i_train,vars).+0.1
yy_train=zeros(i_train); for i=1:i_train yy_train[i]=ff(xx_train[i,:]) end
# mi=minimum(yy_train);yy_train.-=mi;ma=maximum(yy_train);yy_train./=ma;yy_train.*=0.8;yy_train.+=0.1
i_test=1000
xx_test=0.8*rand(i_test,vars).+0.1
yy_test=zeros(i_test); for i=1:i_test yy_test[i]=ff(xx_test[i,:]) end
# yy_test.-=mi;yy_test./=ma;yy_test.*=0.8;yy_test.+=0.1

neurons=Int64(floor(i_train/(vars+1)))
neurons=100
inds_all,n_per_part=ann_by_parts.___clustering(neurons,xx_train,200)
cc1=1.0
# Use # 2 r^4
a_all,a_layer1,layer1=ann_by_parts.train_layer_1_rbf(neurons,vars,i_train,n_per_part,inds_all,xx_train,yy_train,cc1)
predl1=ann_by_parts.predict_new_rbf(a_all,a_layer1,xx_train,i_train,neurons,n_per_part,inds_all,xx_train,vars,cc1)
predl1_test=ann_by_parts.predict_new_rbf(a_all,a_layer1,xx_test,i_test,neurons,n_per_part,inds_all,xx_train,vars,cc1)
maetr=mean(abs.(yy_train-predl1))
maete=mean(abs.(yy_test-predl1_test))

##### APPROXIMATION OF DERIVATIVES
# first variable
yy_train_deriv=-3xx_train[:,1].^2+10xx_train[:,1]+5xx_train[:,2]
yy_train_deriv_deriv=-6xx_train[:,1]+10ones(i_train)
yy_test_deriv=-3xx_test[:,1].^2+10xx_test[:,1]+5xx_test[:,2]
yy_test_deriv_deriv=-6xx_test[:,1]+10ones(i_test)

predl1_deriv=ann_by_parts.predict_new_rbf_deriv(a_all,a_layer1,xx_train,i_train,neurons,
            n_per_part,inds_all,xx_train,vars,cc1,1)
mean(abs.(predl1_deriv-yy_train_deriv))
predl1_deriv_deriv=ann_by_parts.predict_new_rbf_deriv_deriv(a_all,a_layer1,xx_train,i_train,neurons,
            n_per_part,inds_all,xx_train,vars,cc1,1)
mean(abs.(predl1_deriv_deriv-yy_train_deriv_deriv))
predl1_test_deriv=ann_by_parts.predict_new_rbf_deriv(a_all,a_layer1,xx_test,i_test,neurons,
            n_per_part,inds_all,xx_train,vars,cc1,1)
mean(abs.(predl1_test_deriv-yy_test_deriv))
predl1_test_deriv_deriv=ann_by_parts.predict_new_rbf_deriv_deriv(a_all,a_layer1,xx_test,i_test,neurons,
            n_per_part,inds_all,xx_train,vars,cc1,1)
mean(abs.(predl1_test_deriv_deriv-yy_test_deriv_deriv))

# second variable
yy_train_deriv=5xx_train[:,1]
yy_train_deriv_deriv=zeros(i_train)
yy_test_deriv=5xx_test[:,1]
yy_test_deriv_deriv=zeros(i_test)

predl1_deriv=ann_by_parts.predict_new_rbf_deriv(a_all,a_layer1,xx_train,i_train,neurons,
            n_per_part,inds_all,xx_train,vars,cc1,2)
mean(abs.(predl1_deriv-yy_train_deriv))
predl1_deriv_deriv=ann_by_parts.predict_new_rbf_deriv_deriv(a_all,a_layer1,xx_train,i_train,neurons,
            n_per_part,inds_all,xx_train,vars,cc1,2)
mean(abs.(predl1_deriv_deriv-yy_train_deriv_deriv))

predl1_test_deriv=ann_by_parts.predict_new_rbf_deriv(a_all,a_layer1,xx_test,i_test,neurons,
            n_per_part,inds_all,xx_train,vars,cc1,2)
mean(abs.(predl1_test_deriv-yy_test_deriv))
predl1_test_deriv_deriv=ann_by_parts.predict_new_rbf_deriv_deriv(a_all,a_layer1,xx_test,i_test,neurons,
            n_per_part,inds_all,xx_train,vars,cc1,2)
mean(abs.(predl1_test_deriv_deriv-yy_test_deriv_deriv))


using BenchmarkTools, Filtering

SUITE = BenchmarkGroup()

SUITE["ar_linear"] = BenchmarkGroup(["models","linear",])

β = [0.0;1.0]
δ = 1.0
φ = 0.95
ση = 1.0
σε = 1.0

θ0 = [δ;atanh(φ);ση;σε]

β = [0.0;1.0]

N = 1000

arh = ARH(β)

SUITE["ar_linear"]["simulate"] = @benchmarkable Filtering.simulate(arh,θ0,1000)

X,Y = Filtering.simulate(arh,θ0,1000)

SUITE["ar_linear"]["kalman"] = @benchmarkable kalman_filter(arh,θ0,Y,fill([1.0],length(Y)))
SUITE["ar_linear"]["particle"] = @benchmarkable particle_filter(arh,θ0,Y,N,fill([1.0],length(Y)))


using StaticArrays
using LinearAlgebra
import DifferentialEquations;	const DE = DifferentialEquations

function pauli_matrices(; type=Float64)
	σ1 = SMatrix{2, 2, Complex{type}}([0 1; 1 0])
	σ2 = SMatrix{2, 2, Complex{type}}([0 -1im; 1im 0])
	σ3 = SMatrix{2, 2, Complex{type}}([1 0; 0 -1])
	σ4 = SMatrix{2, 2, Complex{type}}([1 0; 0 1])
	σ = (σ1, σ2, σ3, σ4)
	return σ
end
"""
Cosine envelope function.
# Arguments
- `t::Real`: time.
- `tf::Real`: duration of the snap gate.
- `n::Integer`: harmonic order.
"""
function fc(t, tf, n)
	tau = t/tf
	ft = (1 - cos(2*pi*n*tau))/tf
	return ft
end
"""
Sine envelope function.
# Arguments
- `t::Real`: time.
- `tf::Real`: duration of the snap gate.
- `n::Integer`: harmonic order.
"""
function fs(t, tf, n)
	tau = t/tf
	ft = sin(2*pi*n*tau)/tf
	return ft
end
"""
Definite integral of fc.
# Arguments
- `t::Real`: time.
- `tf::Real`: duration of the snap gate.
- `n::Integer`: harmonic order.
"""
function intfc(t, tf, n)
	tau = t/tf
	c = 2*pi*n
	return tau - sin(c*tau)/c
end
"""
Driving fields in the x and y direction in a given subspace n of the bosonic system.
# Arguments
- `t::Real`: time.
- `tf::Real`: duration of the snap gate.
- `α::Real`: angle of the rotation axis.
"""
function f(t, tf, α)
	#ft = zeros(2)
	theta = pi/2
	T = tf/2
	if t < T
		# ft[1] = theta*fc(t, T, 1)
		fx, fy = theta*fc(t, T, 1), 0.0
	else
		fct = theta*fc(t - T, T, 1)
		# ft[1], ft[2] = fct*cos(α), fct*sin(α)
		fx, fy = fct*cos(α), fct*sin(α)
	end
	#return ft
	return fx, fy
end
"""
Driving fields in the x and y direction in the whole truncated bosonic Hilbert space.
# Arguments
- `t::Real`: time.
- `tf::Real`: duration of the snap gate.
- `α::Array{Real}`: angle of the rotation axes.
- `N::Integer`: size of the truncated bosonic Hilbert space.
"""
function f(t, tf, levels, α, N)
	ft = zeros(2, N)
	theta = pi/2
	index = levels .+ 1
	T = tf/2
	if t < T
		ft[1, index] .= theta*fc(t, T, 1)
	else
		fct = theta*fc(t - T, T, 1)
		ft[1, index] .= fct.*cos.(α)
		ft[2, index] .= fct.*sin.(α)
	end
	return ft
end
"""
Integral of the driving fields in the x and y direction in in a given subspace n of the
bosonic system.
# Arguments
- `t::Real`: time.
- `tf::Real`: duration of the snap gate.
- `α::Array{Real}`: angle of the rotation axes.
"""
function intf(t, tf, α)
	theta = pi/2

	intfct = theta*intfc(t, tf, 1)
	intfx = intfct*cos(α)
	intfy = -intfct*sin(α)

	return intfx, intfy
end
"""
Integral of the driving fields in the x and y direction in the whole truncated bosonic
Hilbert space.
# Arguments
- `t::Real`: time.
- `tf::Real`: duration of the snap gate.
- `levels::Array{Integer}`: bosonic levels that are resonantly driven.
- `α::Array{Real}`: angle of the rotation axes.
- `N::Integer`: size of the truncated bosonic Hilbert space.
"""
function intf(t, tf, levels, α, N)
	intft = zeros(2, N)
	theta = pi/2
	index = levels .+ 1

	intfct = theta*intfc(t, tf, 1)
	intft[1, index] .= intfct.*cos.(α)
	intft[2, index] .= -intfct.*sin.(α)

	return intft
end

function F(t, tf, χ, levels, α, N; ideal_interaction::Bool=false)
	χef = (χ[1], χ[2]/2, χ[3]/6)
	Ft = zeros(4, N)
	for i = 1: length(levels)
		Fj!(Ft, t, tf, χef, levels[i], α[i], N, ideal_interaction)
	end
	return Ft
end

function Fj!(Ft, t, tf, χef, j, α, N, ideal_interaction)
	fx, fy = f(t, tf, α)
	if ideal_interaction == true
		Ft[1, j + 1] = fx
		Ft[2, j + 1] = -fy
	else
		for n = 0: N - 1
			dφn = χef[1]*n + χef[2]*n*(n - 1) + χef[3]*n*(n - 1)*(n - 2)
			dφj = χef[1]*j + χef[2]*j*(j - 1) + χef[3]*j*(j - 1)*(j - 2)
			φ = (dφn - dφj)*t
			cosφ = cos(φ)
			sinφ = sin(φ)
			Ft[1, n + 1] +=  + fx*cosφ - fy*sinφ
			Ft[2, n + 1] +=  - fx*sinφ - fy*cosφ
		end
	end
	return nothing
end

function V(t, tf, χ, levels, α, N)
	χef = (χ[1], χ[2]/2, χ[3]/6)
	Vt = zeros(4, N)
	for i = 1: length(levels)
		Vj!(Vt, t, tf, χef, levels[i], α[i], N)
	end
	return Vt
end

function Vj!(Vt, t, tf, χef, j, α, N)
	fx, fy = f(t, tf, α)
	for n = 0: N - 1
		if n != j
			dφn = χef[1]*n + χef[2]*n*(n - 1) + χef[3]*n*(n - 1)*(n - 2)
			dφj = χef[1]*j + χef[2]*j*(j - 1) + χef[3]*j*(j - 1)*(j - 2)
			φ = (dφn - dφj)*t
			cosφ = cos(φ)
			sinφ = sin(φ)
			Vt[1, n + 1] +=  + fx*cosφ - fy*sinφ
			Vt[2, n + 1] +=  - fx*sinφ - fy*cosφ
		end
	end
	return nothing
end

function VI(t, tf, χ, levels, α, N)
	T = tf/2
	Vt = V(t, tf, χ, levels, α, N)
	v = zeros(4, N)
	if t < T + 10.0*eps()
		v[1: 2, :] .= intf(t, T, levels, α.*0, N)
		unitary_tranformation!(Vt, v, N)
	else
		v[1: 2, :] .= intf(t - T, T, levels, α, N)
		unitary_tranformation!(Vt, v, N)
		v[1: 2, :] .= intf(T, T, levels, α.*0, N)
		unitary_tranformation!(Vt, v, N)
	end
	return Vt
end

function VI2(t, tf, χ, levels, α, N)
	Vt = V(t, tf, χ, levels, α, N)
	Vt = vector2matrix(Vt, σ)
	U0 = U(t, tf, χ, levels, α, N)
	VIt = adjoint(U0)*Vt*U0
	VI2t = zeros(4, N)
	for i = 1: N
		VI2t[:, i] = pauli_decomp(VIt[2*(i - 1) + 1: 2*i, 2*(i - 1) + 1: 2*i])
	end
	return VI2t
end

function H(t, F, args)
    σ = args[1]
	N = args[2]
    Ft = F(t)
	Ht = zeros(Complex{Float64}, 2, 2, N)
	for j = 1: N
		vj = view(Ft,:,j)
		Htj = view(Ht, :, :, j)
		vector2matrix!(Htj, vj, σ)
	end
	return Ht
end

function U(t, tf, χ, levels, α, N, σ)
	index = levels .+ 1
	aux = zeros(Complex{Float64}, N, 4)
	T = tf/2
	if t < T
		aux[1:2,:] .= -intf(t, T, levels, α.*0, N)
		aux .= exp_matrix(aux, N)
		U0 = vector2matrix(aux, σ)
	else
		aux[1: 2, :] .= -intf(T, T, levels, α.*0, N)
		aux .= exp_matrix(aux, N)
		U0 = vector2matrix(aux, σ)

		aux[1: 2, :] .= -intf(t-T, T, levels, α, N)
		aux[3: 4, :] .= zero(Complex{Float64})#zeros(Complex{Float64}, 2, N)
		aux .= exp_matrix(aux, N)
		M = vector2matrix(aux, σ)
		U0 = M*U0
	end
	return U0
end

function U2(t, tf, χ, levels, α, N, σ)
	index = levels .+ 1
	aux = zeros(4, N)
	M = zeros(Complex{Float64}, 2*N, 2*N)
	T = tf/2
	if t<T
		aux[1: 2, :] .= intf(t, T, levels, α.*0, N)
		vector2matrix!(M, aux, σ)
		U0 = exp(-1im*M)
	else
		aux[1: 2, :] .= intf(T, T, levels, α.*0, N)
		vector2matrix!(M, aux, σ)
		U0 = exp(-1im*M)

		aux[1: 2, :] .= intf(t-T, T, levels, α, N)
		vector2matrix!(M, aux, σ)
		U0 = exp(-1im*M)*U0
	end
	return U0
end

function dU!(dU, U, args, t)
	H = args[1]
	N = args[2]::Int
	Ht = args[3]
	M = args[4]
	Ht .= H(t)
	for j = 1: N
		mul!(M, view(Ht,:,:,j), view(U,:,:,j))
		dU[:,:,j] .= (-1im).*M
	end
end

function solveU(H, χ, F, tf, N, σ)
    args = (σ, N)
	fH = @closure t-> H(t, F, args)

	U0 = zeros(Complex{Float64}, 2, 2, N)
	for j = 1: N
		U0[:, :, j] .= Matrix{Complex{Float64}}(I, 2, 2)
	end
	tspan = (zero(tf), tf)
	args2 = (fH, N, zeros(Complex{Float64}, 2, 2, N), zeros(Complex{Float64}, 2, 2))
	prob = DE.ODEProblem(dU!, U0, tspan, args2)
	sol = DE.solve(prob, DE.Vern9(), saveat=[tf], reltol=1e-10, abstol=1e-14, maxiters=1e6)
	return sol(tf)
end

function U(system_params, σ; ideal_dynamics::Bool = false, matrix::Bool = false,
		   frame::String = "qubit")
	tf = system_params["tf"]
	χ = system_params["chi"]
	levels = system_params["levels"]
	α = system_params["alpha"]
	N = system_params["N"]

	if frame == "qubit"
		Faux = @closure t-> F(t, tf, χ, levels, α, N, ideal_interaction=ideal_dynamics)
	elseif frame == "interaction"
		if ideal_dynamics == true
			Faux = @closure t-> zeros(4, N)
		else
			Faux = @closure t-> VI(t, tf, χ, levels, α, N)
		end
	elseif frame == "lab"
		Faux = @closure t-> Flab(t, tf, χ, levels, α, N; ideal_interaction=ideal_dynamics)
	end

	sol = solveU(H, χ, Faux, tf, N, σ)
	if matrix == false
		return sol
	else
		return tensor2matrix(sol)
	end
end

function solveOmega_snap(H, tf, order, N)
	numdims = 4*N
	fH = @closure t-> -reshape(H(t), (numdims))

	Omegaf = solveOmega(tf, order, fH, comm!, numdims)

	reshaped_Omegaf = zeros(order, 4, N)
	for i = 1: order
		reshaped_Omegaf[i, :, :] .= reshape(Omegaf[:, i], (4, N))
	end
	return reshaped_Omegaf

end


module MathToolkit

export recvec
export pslq

function recvec(v; compute=BigInt, output=Int64)
  n = length(v)
  l = Array(Rational{compute}, n)
  l[:] = v[:]
  z = div(n, 2)
  q1 = Array(Rational{compute}, n)
  q2 = Array(Rational{compute}, n)
  q1[1] = 0
  q2[1] = 1
  sq1 = 1
  sq2 = 1
  q  = Array(Rational{compute}, n)
  m  = Array(Rational{compute}, n)
  b = 1
  while sq2 <= z
    # look for first non zero term in the series taylor(l)
    while l[b] == 0
      if (b += 1) > n
        c = 1
        g = num(q2[1])
        for k = 1:sq2
          c = lcm(c, den(q2[k]))
          g = gcd(g, num(q2[k]))
        end
        if q2[1] < 0
          c = -c
        end
        q = Array(output, sq2)
        q[:] = (c//g)*q2[1:sq2]
        return q
      end
    end
    # reciprocal of the series: x^(b-1) / taylor(l)
    m[1] = 1 / l[b]
    for k = (b+1):n
      c = zero(compute)
      for j = b:(k-1)
        c -= l[j+1]*m[k-j]
      end
      m[ k-b+1 ] = c / l[b]
    end
    l, m = m, l
    n -= b-1
    # compute (only) denominator of Pade approximant
    q[1:sq2] = l[1] * q2[1:sq2]
    q[sq2+1:b+sq1-1] = 0
    q[b:b+sq1-1] += q1[1:sq1]
    sq = max(sq2, b+sq1-1)
    while q[sq] == 0
      sq -= 1
    end
    q1, q2, q = q2, q, q1
    sq1, sq2 = sq2, sq
    # substract term of degree 0 before computing next reciprocal
    b = 2
  end
  return Array(output, 0)
end

function pslq(x, prec; maxiter=256)
  const n = length(x)
  const teps = prec * 16.0
  const gam = 1.2
  s = Array(BigFloat, n)
  s[n] = abs(x[n])
  t = x[n]^2
  for i = (n-1):-1:1
    t += x[i]^2
    s[i] = sqrt(t)
  end
  t = s[1]
  y = x/t
  s = s/t

  h = [ zeros(BigFloat, n-1) for _ = 1:n ]
  for i = 1:n, j = 1:min(i,n-1)
    h[i][j] = (i==j)?s[j+1]/s[j]:-y[i]*y[j]/(s[j]*s[j+1])
  end

  a = [ zeros(Integer, n) for _ = 1:n ]
  for i = 1:n
    a[i][i] = 1
  end
  b = deepcopy(a)

  for i = 2:n, j = (i-1):-1:1
    t = round(h[i][j]/h[j][j])
    y[j] += t*y[i]
    for k =1:j
      h[i][k] -= t*h[j][k]
    end
    t = convert(Integer,t)
    for k = 1:n
      a[i][k] -= t*a[j][k]
      b[j][k] += t*b[i][k]
    end
  end

  m = 0
  for itr =1:maxiter
    mval = -1
    g = gam
    for i = 1:(n-1)
      t = abs(g*h[i][i])
      if t > mval
        mval = t
        m = i
      end
      g *= gam
    end
    y[m], y[m+1] = y[m+1], y[m]
    h[m], h[m+1] = h[m+1], h[m]
    a[m], a[m+1] = a[m+1], a[m]
    b[m], b[m+1] = b[m+1], b[m]
    if m < n-1
      t0 = sqrt(h[m][m]^2 + h[m][m+1]^2)
      t1 = h[m][m]/t0
      t2 = h[m][m+1]/t0
      for i = m:n
        t3 = h[i][m]
        t4 = h[i][m+1]
        h[i][m] = t1*t3 + t2*t4
        h[i][m+1] = t1*t4 - t2*t3
      end
    end
    for i = (m+1):n, j=min(i-1,m+1):-1:1
      t = round(h[i][j] / h[j][j])
      y[j] += t*y[i]
      for k=1:j
        h[i][k] -= t*h[j][k]
      end
      t = convert(Integer,t)
      for k=1:n
        a[i][k] -= t*a[j][k]
        b[j][k] += t*b[i][k]
      end
    end

    mval = 0
    for j=1:(n-1)
      t = abs(h[j][j])
      if t > mval
        mval = t
      end
    end
    for i=1:n
      t = abs(y[i])
      if t < teps
        return b[m]
      end
    end
  end
  return zeros(Integer, n)
end

end # module


""" Graphviz drawing of categories, functors, diagrams, and related structures.
"""
module GraphvizCategories
export to_graphviz

using ...CategoricalAlgebra, ...Graphs, ..GraphvizGraphs
import ..Graphviz
import ..GraphvizGraphs: to_graphviz, to_graphviz_property_graph

function to_graphviz(F::FinFunctor; kw...)
  to_graphviz(to_graphviz_property_graph(F; kw...))
end

function to_graphviz_property_graph(F::FinFunctor; kw...)
  g = graph(dom(F))
  pg = to_graphviz_property_graph(g; kw...)
  for v in vertices(g)
    lᵥ = g[v, :vname]
    tᵥ = ob_map(F, v)
    set_vprop!(pg, v, :label, string(lᵥ, ":", tᵥ))
  end
  for e in edges(g)
    tₑ = hom_map(F, e)
    set_eprop!(pg, e, :label, string(tₑ))
  end
  pg
end

end


"Kamada Kawai Layout"
function kamada_kawai(G::PowerModelsGraph, dist::Union{Nothing,Matrix{Float64}}=nothing, pos::Union{Nothing,Matrix{Float64}}=nothing, weight="weight", scale=1, center=nothing, dim=2)
    graph = G.graph # convert to undirected graph
    nNodes = LightGraphs.nv(graph)
    if nNodes == 0
        return
    end

    if dist===nothing
        dist=Dict()
        for i in 1:nNodes
            dist[i]=LightGraphs.dijkstra_shortest_paths(graph, i).dists
        end
    end
    dist_mtx = 1e6 * ones(nNodes, nNodes) #reform dist into matrix, probably easier??
    for nr in LightGraphs.vertices(graph)
        rdist = dist[nr]
        for nc in LightGraphs.vertices(graph)
            dist_mtx[nr,nc] = rdist[nc]
        end
    end
    if pos===nothing
        if dim >= 3
            pos= 2 .* rand(Float64,dim, nNodes) .- 1 #??? make a matrix n x dim
        elseif dim == 2
            a = LightGraphs.adjacency_matrix(graph)
            pos = convert(Array,RecursiveArrayTools.VectorOfArray(NetworkLayout.Circular.layout(a)))
        else
            pos = [pt for pt in range(0, 1, length=nNodes)]
        end
    end

    pos = _kamada_kawai_solve(dist_mtx, pos, dim, nNodes)
    positions = [[pos[1,i],pos[2,i]] for i in 1:size(pos,2)]
    return positions
end

"""
Anneal node locations based on the Kamada-Kawai cost-function,
using the supplied matrix of preferred inter-node distances,
and starting locations.
"""
function _kamada_kawai_solve(dist_mtx::Array{Float64,2}, pos_arr::Array{Float64,2}, dim::Int, nNodes::Int)::Matrix{Float64}
    pos_vec = pos_arr[:]
    meanwt = 1e-3

    opt = NLopt.Opt(:LD_LBFGS, length(pos_vec))
    opt.xtol_rel = 1e-4
    opt.min_objective = (pos_vec, grad)->_kamada_kawai_costfn(pos_vec,grad,1.0./(dist_mtx + LinearAlgebra.I(nNodes) * 1e-3), meanwt, dim, nNodes)
    (minf,minx,ret) = NLopt.optimize(opt, pos_vec)

    numevals = opt.numevals # the number of function evaluations
    #println("got $minf at $minx after $numevals iterations (returned $ret)")
    return reshape(minx, dim, nNodes)
end

"Cost-function and gradient for Kamada-Kawai layout algorithm"
function _kamada_kawai_costfn(pos_vec::Vector{Float64}, grad::Vector{Float64}, invdist::Array{Float64,2}, meanweight::Float64, dim::Int, nNodes::Int)
    pos_arr = reshape(pos_vec,dim,nNodes)
    delta = zeros(Float64,dim,nNodes,nNodes)
    for i in 1:nNodes
        delta[:,:,i] = pos_arr .- pos_arr[:,i]
    end

    nodesep = zeros(Float64, dim, nNodes)
    nodesep =  reshape(sqrt.(sum(x -> x^2, delta; dims=1)),nNodes,nNodes)
    direction = OMEinsum.ein"ijk,jk -> ijk"(delta, 1.0./(nodesep+LinearAlgebra.I(nNodes)*1e-3))

    offset = nodesep .* invdist .- 1.0
    for i in 1:nNodes
        offset[i,i] = 0.0
    end

    cost = 0.5 * sum(offset.^2)
    gradient = OMEinsum.ein"jk,jk,ijk->ij"(invdist, offset, direction) - OMEinsum.ein"jk,jk,ijk->ik"(invdist, offset, direction)

    # # Additional parabolic term to encourage mean position to be near origin:
    sumpos = sum(pos_arr, dims=2)
    cost += 0.5 .* meanweight .* sum(sumpos.^2)
    origin_penalty = meanweight*sumpos
    for i in 1:nNodes
        gradient[:,i] += origin_penalty
    end
    if length(grad) > 0
        grad[:] = gradient[:]
    end

    return cost
end

function layout_graph_kamada_kawai!(G) #return type must be dictionary
    return positions = kamada_kawai(G)
end


"""Calculates line of sight distance between two points on the earth (earth radius and conversion to miles hardcoded below)"""
function haversine(x, y)
    @inbounds begin
        # longitudes
        Δλ = deg2rad(y[1] - x[1])

        # latitudes
        φ₁ = deg2rad(x[2])
        φ₂ = deg2rad(y[2])
    end

    Δφ = φ₂ - φ₁

    # haversine formula
    a = sin(Δφ/2)^2 + cos(φ₁)*cos(φ₂)*sin(Δλ/2)^2

    # distance on the sphere
    distance_in_meters = round(2 * 6.3781e6 * asin( min(√a, one(a)) )) # take care of floating point errors
    distance_in_miles = round(distance_in_meters / 1_609, 2)
    distance_in_miles
end


using PWS
using Catalyst
using DrWatson
using LinearAlgebra

import Distributions: MvNormal, Poisson

sn = @reaction_network begin
    κ, ∅ --> S
    λ, S --> ∅
end κ λ

rn = @reaction_network begin
    ρ, S --> X + S
    μ, X --> ∅
end ρ μ

function get_system(mean_s, mean_x, tau_s, tau_x, duration=tau_s)
    κ = mean_s / tau_s
    λ = 1 / tau_s
    ρ = (mean_x / mean_s) / tau_x
    μ = 1 / tau_x
    PWS.JumpSystem(sn, rn, [κ, λ], [ρ, μ], mean_s, mean_x, duration)
end

nothing


function solve24(nums)
    length(nums) != 4 && error("Input must be a 4-element Array")
    syms = [+,-,*,/]
    for x in syms, y in syms, z in syms
        for i = 1:24
            a,b,c,d = nthperm(nums,i)
            if round(x(y(a,b),z(c,d)),5) == 24
                return "($a$y$b)$x($c$z$d)"
            elseif round(x(a,y(b,z(c,d))),5) == 24
                return "$a$x($b$y($c$z$d))"
            elseif round(x(y(z(c,d),b),a),5) == 24
                return "(($c$z$d)$y$b)$x$a"
            elseif round(x(y(b,z(c,d)),a),5) == 24
                return "($b$y($c$z$d))$x$a"
            end
        end
    end
    return "0"
end


@testset "Zlattices" begin
  G = matrix(ZZ, 2, 2, [2,1,1,2])
  B = matrix(ZZ, 1, 2, [1,0])
  L = Zlattice(B, gram=G)
  LL = orthogonal_sum(L,L)
  @test rank(LL) == 2*rank(L)
end


# ==================================
# FOM
#
# system type: continuous LTI system
# state dimension: 1006
# input dimension: 1
# ==================================
using ReachabilityBenchmarks, MathematicalSystems, LazySets, MAT

function fom_model()
    file = matopen(@relpath "fom.mat")

    # system matrix
    A = float(read(file, "A"))  # the matrix has Int entries

    # input matrix
    B = read(file, "B")

    # state domain
    X = Universe(1006)

    # input domain
    U = BallInf([0.0], 1.0)

    # continuous LTI system
    S = @system(x' = Ax + Bu, x ∈ X, u ∈ U)

    return S
end


# This is reference to some of the ideas in Applied Numerical Linear Algebra,
# We are investigating a Lanczos the self-correct the lost of orthogonality
# using ritz vectors projections.
# Here is what we need to do:
#   * Testing for the loss of orthgonality during the runtime of the Lanczos algorithm
#   * Once it's detected, start a Eigensolve on SymTridiagonal and then look for well converged Ritz Vectors.
#   * Store all of those Ritz Vectors as a matrix for reorthogonalizations.

mutable struct IterativeLanczosModified{T<:Number}
    A::Function                         # Linear Operator on Tensor as a function.
    T::Dict{Tuple{Int64, Int64}, T}     # Mapping coordinates [i, j] to elements in the Tridiaognal Matrix.
    k::Int64                            # when k=1, the algorithm haven't started yet.
    Aq::Vector{T}                       # For temp storage to reduce GC time.
    Q::Dict{Int64, Vector{T}}           # A Map for storing all the Lanczos Vectors.
    Q_blocks::Vector{Matrix{T}}         # Blocks of Qs from each restarts
    T_blocks::Vector{Dict}              # Blocks of Ts from each restarts

    tensor_shape::Tuple{Int64}          # Shape of the input for the linear operator.
    V::Vector{Union{Matrix{T}, Vector{T}}}
                                        # Eigen system of T.
    Y::Vector{Vector{T}}                # Well converged ritz vectors from last restart.
    ritz_values::Vector{T}              # The list of ritz values.

    function IterativeLanczosModified(A::Function, q1)
        q1 /= norm(q1)
        Aq = A(q1)
        α = dot(q1, reshape(Aq, :))
        D = typeof(Aq[1])        # Generic Type info
        this = new{D}()          # Referenc to instance created here.
        this.tensor_shape = size(Aq)
        this.Aq = reshape(Aq, :)        # flatten it.
        this.A = A
        this.Q = Dict{Int64, Vector{D}}()
        this.T = Dict{Tuple{Int64}, D}()
        this.T[1, 1] = α
        this.T[0, 1] = this.T[1, 0] = 0
        this.T_blocks = Vector{AbstractArray{D}}()
        this.Q[1] = q1
        this.Q[0] = zeros(D, length(this.Aq))
        this.Q_blocks = Vector{Matrix{D}}()
        this.k = 1
        this.V = Vector{Union{Matrix{D}, Vector{D}}}()
        this.Y = Vector{Vector{D}}()
        this.ritz_values = Vector{D}()
        return this
    end

    function IterativeLanczosModified(A::AbstractMatrix, q1)
        IterativeLanczosModified((x) -> A*x, q1)
    end

end

# Override the () operator.
"""
    It performs one iteration of the Lanczos Orthogonalization.
"""
function (this::IterativeLanczosModified)()
    Aq = this.Aq
    k = this.k
    Q = this.Q
    T = this.T

    q̃ = Aq - T[k - 1, k]*Q[k - 1]
    q̃ -= T[k, k]*Q[k]

    for y in this.Y
        q̃ .-= dot(y, q̃)*y
    end

    β = norm(q̃)
    q = q̃/β

    this.Aq = this.A(q)
    α = dot(q, this.Aq)
    T[k + 1, k + 1] = α
    T[k, k + 1] = T[k + 1, k] = β
    Q[k + 1] = q
    this.k += 1
    return q
end

"""
    Restart and figure out the Ritz Vectors from the SymTridiaognal Matrix,
    get the converged ones and put them into the V vectors of matrices.
"""
function Restart(this::IterativeLanczosModified)
    if this.k == 1                      #  TODO: Investigate Edge Case.
        return
    end

    T = GetPreviousT(this, true)        # T is partial
    Q = GetPreviousQ(this)              # Q is full
    D = typeof(this.Aq[1])
    k = this.k

    this.V = Vector{Union{Matrix{D}, Vector{D}}}()
    this.Y = Vector{Vector{D}}()
    this.ritz_values = Vector{D}()
    # Θ, V = eigen(T, max(k - 8, 1):k - 1)
    Θ, V = eigen(T)
    converged = 0
    for (θ, j) in zip(Θ, 1:length(Θ))
        v = V[:, j]
        error = abs(this.T[k, k - 1]*v[k - 1])
        if error < sqrt(eps(Float32))*norm(T)    # store the ritz values and the ritz vector.
            y = Q[:, 1:end - 1]*v
            ŷ = y/norm(y)
            projVal = 0
            for y in this.Y
                projVal += dot(y, ŷ)
            end

            push!(this.Y, ŷ)
            push!(this.ritz_values, θ)
            converged += 1
        end
    end
    sort!(this.ritz_values)
    return converged
end

"""
    Get the orthogonality error between the newest lanczos vectors
    and the most recent lanczos vector after the last restart.
"""
function GetOrthogonalError(this::IterativeLanczosModified)
    return abs(dot(this.Q[1], this.Q[this.k]))
end

"""
    Get the Q matrix generated from the Iterative Lanczos Process.
"""
function GetQMatrix(this::IterativeLanczosModified)
    return hcat([this.Q_blocks..., GetPreviousQ(this)]...)
end

"""
    Get the Q Matrix since the last restart of the Lanczos Algorithm.
"""
function GetPreviousQ(this::IterativeLanczosModified)
    return hcat([this.Q[idx] for idx in 1: this.k]...)
end

"""
    Get the Symmetric Tridiagonal matrix from the Iterative Lanczos Process.
"""
function GetTMatrix(this::IterativeLanczosModified)
    αs = Vector{Number}()
    βs = Vector{Number}()
    for (T, j) in zip([this.T_blocks..., this.T], 1:length(this.T_blocks) + 1)
        blockSize = maximum(tuple[1] for tuple in keys(T))
        α = [T[idx, idx] for idx in 1:blockSize]
        middle = j > 1 && j <= length(this.T_blocks) + 1
        middleStart = middle ? 0 : 1
        β = [T[idx + 1, idx] for idx in middleStart: blockSize - 1]
        append!(αs, α)
        append!(βs, β)
    end
    return SymTridiagonal(αs, βs)
end

"""
    Get the T matrix since last restart of Lanczos Algorithm.
    parameter:
        `ignore_last`: ignore the last row of matrix T_k.
"""
function GetPreviousT(this::IterativeLanczosModified, ignore_last=false)
    endAt = ignore_last ? this.k - 1 : this.k
    α = [this.T[idx, idx] for idx in 1: endAt]
    β = [this.T[idx + 1, idx] for idx in 1: endAt]
    return SymTridiagonal(α, β)
end

# Basic Testing here -----------------------------------------------------------
using LinearAlgebra, Plots
n = 512
eigenValues = collect(LinRange(1e-3, 1, n)).^2
A = Diagonal(eigenValues);
b = rand(n)
ilm = IterativeLanczosModified(A, b)
converged = 0
for itr in 1: n - 1
    ilm()
    # if itr % 1 == 0
    if GetOrthogonalError(ilm) > 1e-10
        converged = Restart(ilm)
        println("Restarts at itr=$itr, converged = $converged")
        # if converged >= n
        #     break
        # end
    end
end

Q = GetQMatrix(ilm)
T = GetTMatrix(ilm)
println()
heatmap(Q'*Q .|> abs .|> log2)

# LANSO,


function calculateLikelihood(matrix, beta, gammas)
    llh = 0
    D, V = size(matrix) # number of documents

    for d = 1:D
        g = gammas[:, d]  / sum(gammas[:, d]) # normalize
        for j = 1:V
            if matrix[d, j] != 0
                llh += matrix[d, j] * log(dot(g, beta[:, j]))
            end
        end
    end

    return llh
end

function EStep(K, alpha, beta, docVector, maxIter=10)
    # doing E-step
    # K is number of topic
    # alpha, beta are parameters
    # docVector: term vector of doc 1 x N
    # docVector[i] = j means i-th word has j occurrence in corpus

    # init phi and gamma
    wIds = find(docVector) # word id of term in document
    N = length(wIds)
    phi = (ones(N, K) * 1 / K)'
    gammas = copy(alpha) + N / K # size of gamma K x 1

    # loop until convergence
    convergent = false
    threshold = 0.01 # parameter for checking convergence
    counter = 1

    while ~convergent
        oldPhi = copy(phi)
        oldGamma = copy(gammas)

        # update phi
        e_dig = exp(digamma(gammas))
        phi = beta[:, wIds] .* e_dig
        phi = phi ./ sum(phi, 1)

        # update gamma
        gammas = alpha + sum(phi, 2)

        # checking convergence
        m = max(sum(abs(phi - oldPhi)), sum(abs(gammas - oldGamma)))
        if m < threshold || counter == 10
            convergent = true
        end
        counter += 1
    end

    # return result
    return (phi, gammas)
end

function updateBeta(K, phi, matrix)
    D, N = size(matrix)
    beta = zeros(K, N)

    # update beta
    for d=1:D
        wIds = find(matrix[d, :])
        beta[:, wIds] += (diagm(vec(matrix[d, wIds])) * phi[d]')'
    end
    #normalize beta
    beta = broadcast( * , beta, 1./sum(beta, 1))
    return beta
end

function updateAlpha(K, gammas, maxIter=10)
    # K is number of topic
    # D is total number of documents
    # N is total number of words
    # phi is
    # gammas is K x D matrix
    # matrix is (

    #alpha = zeros(K, 1) + 1
    (K, D) = size(gammas)
    alpha = mean(gammas, 2) ./ K
    # update alpha
    convergent = false
    threshold = 0.01
    counter = 1
    pg = sum(digamma(gammas),2) - sum(digamma(sum(gammas,1)));

    while ~convergent
        #L = L_α(alpha, gammas)
        #println("MStep: ", counter, ":L:", L, ":alpha:", alpha)
        oldAlpha = copy(alpha)

        alpha0 = sum(alpha)
        g = D * (digamma(alpha0) - digamma(alpha)) + pg
        h = -1 ./ trigamma(alpha)
        hgz = h' * g / (1 / trigamma(alpha0) + sum(h))

        for i=1:K
            alpha[i] = alpha[i] - h[i] * ( g[i] - hgz[1, 1]) / D
        end

        # check convergence
        if sum(abs(alpha - oldAlpha)) < threshold || counter == maxIter
            convergent = true
        end
        counter +=1
    end

    return alpha
end

function L_α(α, γs)
    (K, D) = size(γs)
    L = D * (  log(gamma(sum(α))) - sum(log(gamma(α))))
    for d=1:D
        for i=1:K
            L += ( α[i] - 1 ) * ( digamma(γs[i, d]) - digamma(sum(γs[:, d])))
        end
    end
    return L
end

function doingEM(K, vocFile="../data/testvocab.txt", dataFile="../data/test.dat")
    # K is number of topic
    # D number of documents
    # N total number of words
    # alpha is parameter of Dirichlet distribution (size K x 1)
    # beta is matrix K x D
    corpus = readFile(dataFile)
    D = length(corpus) # number of doc in dataset
    println(D)
    vocab = readVoc(vocFile)
    N = length(vocab) # number of vocabulary in dataset
    println(N)
    matrix = convertToMatrix(corpus, vocab) # size D x N

    convergent = false
    # create phi and gamma to allocate memory only
    phi = Dict()
    gammas = zeros(K, D)
    # init alpha and beta
    alpha = rand(K, 1)
    alpha = alpha ./ sum(alpha) # normalize alpha
    b = rand(K, N)
    beta = broadcast(*, b, 1 ./ sum(b, 1)) # normalize

    counter = 1

    pllh = 0 # previous likelihood
    threshold = 0.01

    while ~convergent
        # loop until convergence

        # E-step
        #println("E step")
        for d = 1:D
            # doing for each document in corpus
            (subPhi, subGamma) = EStep(K, alpha, beta, matrix[d, :])
            phi[d] = subPhi
            gammas[:, d] = subGamma
        end

        # M-step
        #println("M Step")
        beta = updateBeta(K, phi, matrix)
        alpha = updateAlpha(K, gammas)

        # checking convergence
        llh = calculateLikelihood(matrix, beta, gammas)
        println("iter:", counter,"\tLLH:", llh)
        #if abs(pllh - llh) < threshold
        if counter == 10
            convergent = true
        #else
        #    pllh = llh
        #    println(llh)
        end
        counter += 1
    end
end


module VarianceComponentTypeTest

using VarianceComponentModels, MathProgBase, Ipopt, Base.Test

srand(123)

# generate data from a d-variate response variane component model
n = 100   # no. observations
d = 2     # no. categories
m = 2     # no. variance components
p = 2     # no. covariates
X = randn(n, p)
B = ones(p, d)
Σ = ntuple(x -> zeros(d, d), m)
for i in 1:m
  Σi = randn(d, d)
  copy!(Σ[i], Σi' * Σi)
end
## make the first variance component 0 matrix
#fill!(Σ[1], 0.0)
V = ntuple(x -> zeros(n, n), m)
for i = 1:m-1
  Vi = randn(n, 50)
  copy!(V[i], Vi * Vi')
end
copy!(V[m], eye(n))
# form Ω
Ω = zeros(n*d, n*d)
for i = 1:m
  Ω += kron(Σ[i], V[i])
end
Ωchol = cholfact(Ω)
Y = X * B + reshape(Ωchol[:L] * randn(n*d), n, d)

info("VarianceComponentModel type")
#@code_warntype VarianceComponentModel(B, Σ)
@inferred VarianceComponentModel(B, Σ)
#@code_warntype VarianceComponentModel(Σ)
@inferred VarianceComponentModel(Σ)
vcmodel = VarianceComponentModel(Σ)
@test isempty(vcmodel.B)

info("TwoVarCompModelRotate type")
F = eigfact(Σ[1], Σ[2])
#@code_warntype TwoVarCompModelRotate(B * F.vectors, F.values, F.vectors, logdet(Σ[2]))
@inferred TwoVarCompModelRotate(B * F.vectors, F.values, F.vectors, logdet(Σ[2]))

info("VarianceComponentVariate type")
#@code_warntype VarianceComponentVariate(Y, V)
@inferred VarianceComponentVariate(Y, V)
vcdata = VarianceComponentVariate(Y, V)
@test isempty(vcdata.X)

info("Forming VarianceComponentModel from data")
vcdata = VarianceComponentVariate(Y, X, V)
#@code_warntype VarianceComponentVariate(Y, X, V)
@inferred VarianceComponentVariate(Y, X, V)
vcdata = VarianceComponentVariate(Y, X, V)
#@code_warntype VarianceComponentModel(vcdata)
@inferred VarianceComponentModel(vcdata)
vcmodel = VarianceComponentModel(vcdata)
@test vecnorm(vcmodel.B) ≈ 0.0
@test vecnorm(vcmodel.Σ[1] - eye(d)) ≈ 0.0
@test vecnorm(vcmodel.Σ[2] - eye(d)) ≈ 0.0

info("Pre-compute eigen-decomposition and rotate data")
#@code_warntype TwoVarCompVariateRotate(vcdata)
@inferred TwoVarCompVariateRotate(vcdata)
vcdatarot = TwoVarCompVariateRotate(vcdata)
#@code_warntype TwoVarCompModelRotate(vcmodel)
@inferred TwoVarCompModelRotate(vcmodel)
vcmodelrot = TwoVarCompModelRotate(vcmodel)

info("Test generalized eigen-decomposition")
vcdatacopy = deepcopy(vcdata)
scale!(vcdatacopy.V[2], 2.0) # V[2] = 2I
vcdatacopyrot = TwoVarCompVariateRotate(vcdatacopy)
@test vcdatacopyrot.logdetV2 ≈ n * log(2.0) # det(V[2]) = 2^n

info("VarianceComponentModel from TwoVarCompVariateRotate")
@inferred VarianceComponentModel(vcdatarot)
vcmodelfromrot = VarianceComponentModel(vcdatarot)
@test size(vcmodelfromrot.B, 1) == p
@test size(vcmodelfromrot.B, 2) == d
@test length(vcmodelfromrot.Σ) == m

info("VarianceComponentAuxData from VarianceComponentVariate")
vcobsaux1 = VarianceComponentAuxData(vcdata)
vcobsaux2 = VarianceComponentAuxData(vcdatarot)
@test size(vcobsaux1.res) == size(vcobsaux2.res)
@test size(vcobsaux1.Xwork) == size(vcobsaux2.Xwork)
@test size(vcobsaux1.ywork) == size(vcobsaux2.ywork)
@test size(vcobsaux1.obswt) == size(vcobsaux2.obswt)

info("Query functions")
@test eltype(vcmodel) == eltype(B)
@test eltype(vcdata) == eltype(Y)
@test eltype(vcmodelrot) == eltype(vcmodel)
@test eltype(vcdatarot) == eltype(vcdata)
@test length(vcmodel) == d
@test length(vcdata) == d
@test length(vcmodelrot) == d
@test length(vcdatarot) == d
@test size(vcdata) == (n, d)
@test size(vcdatarot) == (n, d)
@test nvarcomps(vcmodel) == m
@test nvarcomps(vcdata) == m
@test nvarcomps(vcmodelrot) == m
@test nvarcomps(vcdatarot) == m
@test nmeanparams(vcmodel) == p * d
@test nmeanparams(vcmodelrot) == p * d
@test nmeanparams(vcdata) == p * d
@test nmeanparams(vcdatarot) == p * d
@test nvarparams(vcmodel) == m * binomial(d + 1, 2)
@test nvarparams(vcmodelrot) == m * binomial(d + 1, 2)
@test nparams(vcmodel) == p * d + m * binomial(d + 1, 2)
@test nparams(vcmodelrot) == p * d + m * binomial(d + 1, 2)

info("Mean, covariance, and residual of model")
#@code_warntype mean(vcmodel, vcdata)
@inferred mean(vcmodel, vcdata)
μ = mean(vcmodel, vcdata)
@test size(μ, 1) == n
@test size(μ, 2) == d
#@code_warntype cov(vcmodel, vcdata)
@inferred cov(vcmodel, vcdata)
Ω = cov(vcmodel, vcdata)
@test size(Ω, 1) == n * d
@test size(Ω, 2) == n * d
#@code_warntype residual(vcmodel, vcdata)
@inferred residual(vcmodel, vcdata)
res = residual(vcmodel, vcdata)
@test size(res, 1) == n
@test size(res, 2) == d
@inferred residual(vcmodelrot, vcdatarot)
resrot = residual(vcmodelrot, vcdatarot)
@test size(resrot, 1) == n
@test size(resrot, 2) == d

end # module VarianceComponentTypeTest


@testset "geometry.coords2d" begin

@testset "cartesian2d" begin
    c2d = cartesian2d(float.([1 2; 3 4; 5 6; 7 8]))
    @test size(c2d.coords) == (4, 2)
    @test length(point(c2d, 2).pos) == 2
    @test sum(x_components(c2d)) == 16
    @test sum(y_components(c2d)) == 20
end

@testset "point2d" begin
    c2d = cartesian2d(float.([1 2; 3 4; 5 6; 7 8]))
    p1 = point(c2d, 3)
    p2 = point(1.2, -4.6)
    @test x(p1) == 5
    @test y(p2) == -4.6
    @test isapprox((p1 + p2).pos, [6.2, 1.4])
    @test isapprox((p1 - p2).pos, [3.8, 10.6])
    @test isapprox((p2 * 2.0).pos, [2.4, -9.2])
    @test isapprox(cross(p1, point(c2d, 4)), -2)

    v = point(3.0, 0.0)
    @test isapprox(interiorangle(v, point(5.0, 0.0)), 0)
    @test isapprox(interiorangle(v, point(-20.0, 0.0)), pi)
    @test isapprox(interiorangle(v, point(3.0, -3sqrt(3))), pi / 3)
    @test isapprox(interiorangle(v, point(-3.0, 3sqrt(3))), 2pi / 3)
    @test isnan(interiorangle(v, point(0.0, 0.0)))
end

@testset "segment2d" begin
    c2d = cartesian2d(float.([1 2; 3 4; 5 6; 7 8]))
    seg1 = segment(c2d, 2, 4)
    seg2 = Segment2D([1.2 4.5; 3.4 7.9])
    @test size(seg2.coords) == (2, 2)
    @test u(seg1).pos == [3, 4]
    @test v(seg2).pos == [3.4, 7.9]
    @test ux(seg2) == 1.2
    @test uy(seg1) == 4
    @test vx(seg1) == 7
    @test vy(seg2) == 7.9

    @test isapprox(vector(seg2).pos, [2.2, 3.4])
    @test isapprox(midpoint(seg2).pos, [2.3, 6.2])
    tl = translate(seg1, -pi / 2, sqrt(2) / 2)
    @test isapprox(tl.coords, [3.5 3.5; 7.5 7.5])

    seg3 = Segment2D(float.([1 1; 5 5]))
    @test isapprox(trim_u(seg3, 1 / 4).coords, [2 2; 5 5])
    @test isapprox(trim_v(seg3, 1 / 4).coords, [1 1; 4 4])
    @test isapprox(trim_uv(seg3, 1 / 4).coords, [1.5 1.5; 4.5 4.5])

    @test isapprox(cross(u(seg1), v(seg1)), -4)
end

@testset "reference" begin
    c2d = cartesian2d(float.([1 2; 3 4; 5 6; 7 8]))
    seg = segment(c2d, 3, 4)
    @test seg.coords isa SubArray
    @test u(seg).pos == [5, 6]
    setcoord!(seg, 2, point(7.8, 8.3))
    @test v(seg).pos == [7.8, 8.3]
    @test point(c2d, 4).pos == [7.8, 8.3]
end

@testset "isclockwise" begin
    @test isclockwise(cartesian2d(float.([0 1; 1 0; 0 -1; -1 0])))
    @test !isclockwise(cartesian2d(float.([0 1; -1 0; 0 -1; 1 0])))
    @test isclockwise(cartesian2d(float.([1 1; 1 -1; -1 1; -1 -1]))) === nothing
    @test isclockwise(
        cartesian2d(float.([-1 1; 1 1; 1 2; 2 2; 2 -2; -2 -2; -2 2; -1 2])))
end

end # geometry.coords2d


import Base: sum, push!

"""
This is a binary tree where the leaves are values
and the nodes are sums of those values. It is meant to make it
easier to find the leaf such that the sum of it and all previous
leaves is greater than a given value.
"""
struct PrefixSearchTree{T<:Real}
	array::Array{T,1}
	depth::Int
	offset::Int
end

"""
    PrefixSearchTree(ElementType::DataType, count::Int)

Constructor for an empty tree of type ElementType and length count.
"""
function PrefixSearchTree(ElementType::DataType, count::Int)
	new(zeros(ElementType, count), 0, 0)
end

"""
    PrefixSearchTree(summables)

Constructor of a prefix search tree from an iterable list of real numbers.
"""
function PrefixSearchTree(summables)
	a = Array(summables)
	@assert length(a) > 0
	depth = Int(ceil(log(2, length(a)))) + 1
	offset = 2^(depth - 1)
	b = zeros(eltype(summables), 2^depth - 1)
	b[offset:(length(a) + offset - 1)] = a
	#print("depth $depth offset $offset length $(length(b))\n")
	pst = PrefixSearchTree{eltype(summables)}(b, depth, offset)
	calculate_prefix!(pst)
	pst
end

"""
    choose(pst::PrefixSearchTree, value)

Find the minimum index such that the prefix is greater than the given
value.

Precondition: The value must be strictly less than the total for
the tree.
"""
function choose(pst::PrefixSearchTree, value)
	if value >= pst.array[1]
		error("Value $value not less than total $(pst.array[1])")
	end

	requested_value = value
    index = 1
    for level = 0:(pst.depth - 2)
        left_child = 2 * index
        if pst.array[left_child] > value
            index = left_child
        else
            index = left_child + 1
            value -= pst.array[left_child]
        end
    end
    index - pst.offset + 1, pst.array[index] - value
end

sum(pst::PrefixSearchTree) = pst.array[1]

"""
If there are multiple values to enter, then present them
at once as pairs of tuples, (index, value).
"""
function push!(pst::PrefixSearchTree, pairs)
	modify=Set{Int}()
	for (pos, value) in pairs
		index = pos + pst.offset - 1
		pst.array[index] = value
		push!(modify, div(index, 2))
	end

	# everything at depth-1 is correct, and changes are in modify.
    for depth = (pst.depth - 2):-1:0
        parents = Set{Int}()
        for node_idx in modify
            pst.array[node_idx]=pst.array[2 * node_idx] + pst.array[2 * node_idx + 1]
            push!(parents, div(node_idx, 2))
        end
        modify = parents
    end
end

function push!(pst::PrefixSearchTree{T}, index::Int, value::T) where T
    push!(pst, [(index, value)])
end

# Private
function calculate_prefix!(pst::PrefixSearchTree)
	for d = (pst.depth - 2):-1:0
		for node_idx = (2^d):(2^(d + 1) - 1)
			pst.array[node_idx] = pst.array[2 * node_idx] + pst.array[2 * node_idx + 1]
		end
	end
end


#=
# 6. Force and the added mass

## Force

This part introduces the calculation of the force and moment on the body through the negative rate of change of impulse in the fluid. The continuous expressions for linear and angular impulse (about the origin) are, in two dimensions,

$\begin{align}
    \boldsymbol{P} &= \int_{V_f} \boldsymbol{x}\times \boldsymbol{\omega}\,\mathrm{d}V + \int_{S_b} \boldsymbol{x}\times \left( \boldsymbol{n}\times\boldsymbol{v}\right)\,\mathrm{d}S \\
    \boldsymbol{\Pi}_0 &= \frac{1}{2}\int_{V_f} \boldsymbol{x}\times \left(\boldsymbol{x}\times\boldsymbol{\omega}\right)\,\mathrm{d}V + \frac{1}{2} \int_{S_b} \boldsymbol{x}\times \left[ \boldsymbol{x}\times\left( \boldsymbol{n} \times \boldsymbol{v} \right) \right] \,\mathrm{d}S
\end{align}$

If there is only a single body, then the force and moment (about the origin) exerted by the fluid on that body are given by

$\boldsymbol{F} = -\rho \frac{\mathrm{d}\boldsymbol{P}}{\mathrm{d}t}, \qquad \boldsymbol{M}_{0} = -\rho \frac{\mathrm{d}\boldsymbol{\Pi}_{0}}{\mathrm{d}t},$

where $\rho$ is the fluid density. In the two-dimensional applications of this package, angular impulse and the moment have only a single component, e.g., $\boldsymbol{\Pi}_0 =\Pi_0\boldsymbol{e}_z$, where $\boldsymbol{e}_z$ is the unit vector out of the plane.
=#

#=
#### Point vortices past a cylinder

To verify our method of calculating the impulse, we create a model of two point vortices of equal and opposite circulation positioned at both sides of the $x$-axis, with the $x$-axis as axis of symmetry, and left of a circular cylinder positioned at the origin. If the top vortex has a positive circulation, the vortices propel each other towards the cylinder and convect past it. This configuration has an analytical solution for the trajecory and impulse.
=#

#md # See the notebook in the examples folder for the analytical solution.

# We use a rectangular grid.
#md # ```@setup 6.-Force-and-the-added-mass
#md # using GridPotentialFlow
#md # using Plots
#md # ```
#!md using GridPotentialFlow
#!md using Plots
Lx = 10.0
Ly = 4.0
xlim = (-Lx/2,Lx/2)
ylim = (-Ly/2,Ly/2)
g = PhysicalGrid(xlim, ylim, 0.04)

#
Rc = 1
circle = Circle(Rc,2*cellsize(g))
pfb = PotentialFlowBody(circle)

Δs = dlength(circle);

# The initial spacing between the vortices is `y∞`. After the vortices pass the cylinder, they should return to this spacing.
y∞ = Rc/2
v1 = Vortex(-Lx/2+10*cellsize(g),y∞,1.0);
v2 = Vortex(-Lx/2+10*cellsize(g),-y∞,-1.0);

#md # ```@setup 6.-Force-and-the-added-mass
#md # # Analytical trajectory
#md # x(r) = sqrt(r^2-(y∞^2*(r^2-1)^2)/((r^2-1)^2-4*y∞^2))
#md # y(r) = y∞*(r^2-1)/sqrt((r^2-1)^2-4*y∞^2)
#md # s = 0:0.001:π/2
#md # rmin = 1.4375649
#md # c1 = rmin/(-Lx/2+rmin)
#md # c2 = -Lx/2*c1
#md # r = c2 ./ (cos.(s).-c1)
#md # x_trajectory = -reverse(x.(r)); append!(x_trajectory,x.(r))
#md # y_trajectory_upper = reverse(y.(r)); append!(y_trajectory_upper,y.(r))
#md # y_trajectory_lower = -y_trajectory_upper;
#md # # Analytical impulse
#md # Px_func(x,y) = y - y/(x^2+y^2) + y - y/(x^2+y^2);
#md # Py_func(x,y) = -x + x/(x^2+y^2) + x - x/(x^2+y^2);
#md # Px_exact = Px_func.(x_trajectory,y_trajectory_upper);
#md # Py_exact = Py_func.(x_trajectory,y_trajectory_upper);
#md # ```

#!md # Analytical trajectory
#!md x(r) = sqrt(r^2-(y∞^2*(r^2-1)^2)/((r^2-1)^2-4*y∞^2))
#!md y(r) = y∞*(r^2-1)/sqrt((r^2-1)^2-4*y∞^2)
#!md s = 0:0.001:π/2
#!md rmin = 1.4375649
#!md c1 = rmin/(-Lx/2+rmin)
#!md c2 = -Lx/2*c1
#!md r = c2 ./ (cos.(s).-c1)
#!md x_trajectory = -reverse(x.(r)); append!(x_trajectory,x.(r))
#!md y_trajectory_upper = reverse(y.(r)); append!(y_trajectory_upper,y.(r))
#!md y_trajectory_lower = -y_trajectory_upper;
#!md # Analytical impulse
#!md Px_func(x,y) = y - y/(x^2+y^2) + y - y/(x^2+y^2);
#!md Py_func(x,y) = -x + x/(x^2+y^2) + x - x/(x^2+y^2);
#!md Px_exact = Px_func.(x_trajectory,y_trajectory_upper);
#!md Py_exact = Py_func.(x_trajectory,y_trajectory_upper);

# We create the vortex model with these two point vortices and the circle and advance the position of the point vortices over some time. We use `OrdinaryDiffEq.RK4()` for the time stepping and we compute the impulse history with `impulse` in a post-processing step.

model = VortexModel(g,bodies=[pfb],vortices=[v1,v2]);

function rhs(X,model,t)
    setvortexpositions!(model,X)
    Ẋ = vortexvelocities!(model)
    return Ẋ
end

import OrdinaryDiffEq
X = getvortexpositions(model)
prob = OrdinaryDiffEq.ODEProblem(rhs,X,(0.0,64.0),model);
sol = OrdinaryDiffEq.solve(prob,dt=0.1,OrdinaryDiffEq.RK4(),dense=false,adaptive=false);

# Using the history of the vortex positions in `sol.u`, we can calculate the impulse history.
Px_numerical_hist = []
Py_numerical_hist = []
for u in sol.u
    setvortexpositions!(model,u)
    Px, Py = impulse(model)
    push!(Px_numerical_hist,Px)
    push!(Py_numerical_hist,Py)
end

# When we compare the trajectories and impulse history, the numerical and analytical solution should match closely, which is indeed the case.
plot(circle,fillcolor=:black,fillrange=0,fillalpha=0.25,linecolor=:black,linewidth=2,xlabel="x",ylabel="y")
scatter!(model.vortices.x,model.vortices.y,color=:red)
plot!(x_trajectory,y_trajectory_upper,linecolor=:red,label="exact")
plot!(x_trajectory,y_trajectory_lower,linecolor=:red,label="")
plot!(map(s->s.u[1],sol.u),map(s->s.v[1],sol.u),color=:blue,linestyle=:dash,label="simulated")
plot!(map(s->s.u[2],sol.u),map(s->s.v[2],sol.u),color=:blue,linestyle=:dash,label="")

#
plot(x_trajectory,Px_exact,color=:red,label="exact",xlabel="x",ylabel="Px")
plot!(map(s->s.u[1],sol.u),Px_numerical_hist,color=:blue,linestyle=:dash,label="simulated")

#=
#### Flat plate

We can apply this method also to our flat plate example and compare it to the Biot-Savart method from the `PotentialFlow.jl` package.
=#

#md # See the notebook in the examples folder for the Biot-Savart solution.

# Because we create the `PotentialFlow.jl` model with a moving plate instead of a uniform flow, we specify a translational velocity. The `GridPotentialFlow.jl` model will still use a moving body coordinate system and thus a uniform flow that equals the negative translational velocity.

c = 1.0;   # chord length

ċ = -1.0;  # translational velocity
α = -π/3;   # angle of attack

σLE = 0.0;
σTE = 0.0;

Δt = 5e-2;
tf = 1.0;
T = 0.0:Δt:tf;

#md # ```@setup 6.-Force-and-the-added-mass
#md # # Potential Flow
#md # using PotentialFlow
#md # function compute_ẋ!(ẋ, x, t)
#md #     plate, ambient_sys = x
#md #     motion = ẋ[1]
#md #     motion.ċ, motion.c̈, motion.α̇ = motion.kin(t)
#md #     Plates.enforce_no_flow_through!(plate, motion, ambient_sys, t)
#md #     reset_velocity!(ẋ, x)
#md #     self_induce_velocity!(ẋ, x, t)
#md # end
#md # function shed_new_vorticity!(blobs, plate, motion, t, lesp = 0.0, tesp = 0.0)
#md #     z₊ = (blobs[end-1].z + 2plate.zs[end])/3
#md #     z₋ = (blobs[end].z + 2plate.zs[1])/3
#md #     blob₊ = PotentialFlow.Vortex.Blob(z₊, 1.0, δ)
#md #     blob₋ = PotentialFlow.Vortex.Blob(z₋, 1.0, δ)
#md #     Plates.enforce_no_flow_through!(plate, motion, blobs, t)
#md #     Γ₊, Γ₋, _, _ = Plates.vorticity_flux!(plate, blob₊, blob₋, t, lesp, tesp);
#md #     push!(blobs, PotentialFlow.Vortex.Blob(z₊, Γ₊, blobs[1].δ), PotentialFlow.Vortex.Blob(z₋, Γ₋, blobs[1].δ))
#md # end
#md # N = 128 # number of plate control points (distributed along a extrema Chebyshev grid)
#md # δ = 0.01
#md # plate = PotentialFlow.Plate(N, c, zero(ComplexF64), α)
#md # motion = Plates.RigidBodyMotion(ċ, 0.0);
#md # # Initial step
#md # Δz₀ = im*3Δt*exp(im*plate.α) # vectors perpendicular to the plate
#md # z₋, z₊ = plate.zs[[1,N]] # LE and TE
#md # blobs = PotentialFlow.Vortex.Blob.(Δz₀ .+ [z₊, z₋], 1.0, δ) # First two point vortices are placed close to the LE and TE with unit strength
#md # Plates.enforce_no_flow_through!(plate, motion, (), 0)
#md # Γ₊, Γ₋, _, _ = Plates.vorticity_flux!(plate, blobs[1], blobs[2], 0.0, σLE, σTE); # Determine strength of first two vortices
#md # blobs = PotentialFlow.Vortex.Blob.(Δz₀ .+ [z₊, z₋], [Γ₊, Γ₋], δ) # Create first two point vortices now with calculated strengths
#md # sys₀ = (plate, blobs)
#md # sys = deepcopy(sys₀)
#md # sys₊ = deepcopy(sys₀) # Used for storage during time-marching
#md # ẋs = (motion, allocate_velocity(blobs))
#md # imp = ComplexF64[];
#md # t_hist = Float64[]
#md # global t = 0
#md # push!(t_hist,t)
#md # # Time stepping
#md # for tloc in T[2:end]
#md #     global t += Δt
#md #     global sys
#md #     global sys₊
#md #     push!(imp,Elements.impulse(sys))
#md #     push!(t_hist,t)
#md #     local plate, ambient_ω = sys
#md #     local motion, ambient_u = ẋs
#md #     resize!(sys₊[2], length(sys[2]))
#md #     forward_euler!(sys₊, sys, t, Δt, compute_ẋ!, advect!, ẋs)
#md #     sys, sys₊ = sys₊, sys
#md #     shed_new_vorticity!(sys[2], sys[1], ẋs[1], t, σLE, σTE)
#md # end
#md # force = -diff(imp)/Δt;
#md # ```

#!md # Potential Flow
#!md using PotentialFlow
#!md function compute_ẋ!(ẋ, x, t)
#!md     plate, ambient_sys = x
#!md     motion = ẋ[1]
#!md     motion.ċ, motion.c̈, motion.α̇ = motion.kin(t)
#!md     Plates.enforce_no_flow_through!(plate, motion, ambient_sys, t)
#!md     reset_velocity!(ẋ, x)
#!md     self_induce_velocity!(ẋ, x, t)
#!md end
#!md function shed_new_vorticity!(blobs, plate, motion, t, lesp = 0.0, tesp = 0.0)
#!md     z₊ = (blobs[end-1].z + 2plate.zs[end])/3
#!md     z₋ = (blobs[end].z + 2plate.zs[1])/3
#!md     blob₊ = PotentialFlow.Vortex.Blob(z₊, 1.0, δ)
#!md     blob₋ = PotentialFlow.Vortex.Blob(z₋, 1.0, δ)
#!md     Plates.enforce_no_flow_through!(plate, motion, blobs, t)
#!md     Γ₊, Γ₋, _, _ = Plates.vorticity_flux!(plate, blob₊, blob₋, t, lesp, tesp);
#!md     push!(blobs, PotentialFlow.Vortex.Blob(z₊, Γ₊, blobs[1].δ), PotentialFlow.Vortex.Blob(z₋, Γ₋, blobs[1].δ))
#!md end
#!md N = 128 # number of plate control points (distributed along a extrema Chebyshev grid)
#!md δ = 0.01
#!md plate = PotentialFlow.Plate(N, c, zero(ComplexF64), α)
#!md motion = Plates.RigidBodyMotion(ċ, 0.0);
#!md # Initial step
#!md Δz₀ = im*3Δt*exp(im*plate.α) # vectors perpendicular to the plate
#!md z₋, z₊ = plate.zs[[1,N]] # LE and TE
#!md blobs = PotentialFlow.Vortex.Blob.(Δz₀ .+ [z₊, z₋], 1.0, δ) # First two point vortices are placed close to the LE and TE with unit strength
#!md Plates.enforce_no_flow_through!(plate, motion, (), 0)
#!md Γ₊, Γ₋, _, _ = Plates.vorticity_flux!(plate, blobs[1], blobs[2], 0.0, σLE, σTE); # Determine strength of first two vortices
#!md blobs = PotentialFlow.Vortex.Blob.(Δz₀ .+ [z₊, z₋], [Γ₊, Γ₋], δ) # Create first two point vortices now with calculated strengths
#!md sys₀ = (plate, blobs)
#!md sys = deepcopy(sys₀)
#!md sys₊ = deepcopy(sys₀) # Used for storage during time-marching
#!md ẋs = (motion, allocate_velocity(blobs))
#!md imp = ComplexF64[];
#!md t_hist = Float64[]
#!md global t = 0
#!md push!(t_hist,t)
#!md # Time stepping
#!md for tloc in T[2:end]
#!md     global t += Δt
#!md     global sys
#!md     global sys₊
#!md     push!(imp,Elements.impulse(sys))
#!md     push!(t_hist,t)
#!md     local plate, ambient_ω = sys
#!md     local motion, ambient_u = ẋs
#!md     resize!(sys₊[2], length(sys[2]))
#!md     forward_euler!(sys₊, sys, t, Δt, compute_ẋ!, advect!, ẋs)
#!md     sys, sys₊ = sys₊, sys
#!md     shed_new_vorticity!(sys[2], sys[1], ẋs[1], t, σLE, σTE)
#!md end
#!md force = -diff(imp)/Δt;

Δx = 0.01
xlim = (-0.5,2)
ylim = (-1,1)
g = PhysicalGrid(xlim,ylim,Δx);

#

dsdx = 2
plate = RigidBodyTools.Plate(c,dsdx*Δx)
pfb = PotentialFlowBody(plate,edges=[1,length(plate)],σ=[SuctionParameter(σLE),SuctionParameter(σTE)])
transform = RigidTransform((0.0,0.0),-π/3)
transform(plate);
Δs = dlength(plate)
maximum(Δs/Δx);

# As before, we create the initial vortices based on the time step and uniform flow and define a function that create the new point vortices for the subsequent time steps.

firstvLE = GridPotentialFlow.Vortex(plate.x[1]+3Δt*(-ċ)*cos(plate.α+π/2),plate.y[1]+3Δt*(-ċ)*sin(plate.α+π/2),0.0);
firstvTE = GridPotentialFlow.Vortex(plate.x[end]+3Δt*(-ċ)*cos(plate.α+π/2),plate.y[end]+3Δt*(-ċ)*sin(plate.α+π/2),0.0);

#

function createsheddedvortices(plate,oldvortices)

    vLE = GridPotentialFlow.Vortex(2/3*plate.x[1]+1/3*oldvortices[end-1].x,2/3*plate.y[1]+1/3*oldvortices[end-1].y,0.0)
    vTE = GridPotentialFlow.Vortex(2/3*plate.x[end]+1/3*oldvortices[end].x,2/3*plate.y[end]+1/3*oldvortices[end].y,0.0)

    return vLE,vTE
end

# The model free stream should be the negative of the translational velocity of the plate.

model = VortexModel(g,bodies=[pfb],vortices=[firstvLE,firstvTE],U∞=(-ċ,0.0));

#=
Then we enter a time stepping loop and record the impulse every time we create new vortices. Note that because we release new vortices at every time step, it is not straightforward to use a multi-step time integration method like RK4 to advance the positions of the point vortices. Instead, we use forward Euler and the size of the positions vector increases every time step by the number of new vortices we insert. To advance the system one time step, we first need to solve the saddle-point system, set the strengths of the new vortices, subtract that circulation from the bodies, and calculate the vortex velocities with the saddle-point system solution. For convenience, these four operations are carried out internally if we would call `vortexvelocities!(model)`. However, if we want to have access to the solution of the saddle-point system, we can explicitly perform each operation as is done in the following time stepping loop. Note that we have to set the strengths of the new vortices before calling the `impulse` function, otherwise the new vortices won't be included in the calculation.
=#

Px_hist = Float64[];
Py_hist = Float64[];
sol = solve(model);
for tloc in T[2:end]
    X = getvortexpositions(model) # gets bigger every time step because we add vortices
    Ẋ = deepcopy(X)

    solve!(sol, model)
    setvortexstrengths!(model, sol.δΓ_vec, length(X.u)-1:length(X.u))
    subtractcirculation!(model.bodies, sol.δΓ_vec)
    Px, Py = impulse(model)

    vortexvelocities!(Ẋ, model, sol.ψ)
    X .= X .+ Ẋ*Δt
    setvortexpositions!(model, X)

    vLE, vTE = createsheddedvortices(plate,model.vortices)
    pushvortices!(model,vLE,vTE)
    push!(Px_hist,Px)
    push!(Py_hist,Py)
end

# The force history can be obtained by taking finite differences of the impulse history.

Fx_hist = -diff(Px_hist)/Δt;
Fy_hist = -diff(Py_hist)/Δt;

# We can now compare the positions of the point vortices by shifting the `PotentialFlow.jl` solution by `-tf*ċ` such that the origin of the frame of reference coincides with the center plate. Superimposing the `GridPotentialFlow.jl` solution then shows that the positions of the point vortices agree very well.

plot(plate,fillcolor=:black,fillrange=0,fillalpha=0.25,linecolor=:black,linewidth=2,xlim=xlim,ylim=ylim,xlabel="x",ylabel="y")
scatter!(real.((v->v.z).(sys[2])).-tf*ċ,imag.((v->v.z).(sys[2])),color=:red,markersize=4,label="PotentialFlow.jl")
scatter!(model.vortices.x,model.vortices.y,color=:blue,markersize=2,label="GridPotentialFlow.jl")

# The vertical impulse and the vertical force (lift) can also be compared and show good agreement as well.
plot(xlabel="t",ylabel="Py")
plot!(T[1:end-1],imag.(imp),color=:blue,label="PotentialFlow.jl")
plot!(T[1:end-1],Py_hist,color=:red,label="GridPotentialFlow.jl")

#
plot(xlabel="t",ylabel="Fy")
plot!(T[2:end-1],imag.(force),color=:blue,label="PotentialFlow.jl")
plot!(T[2:end-1],Fy_hist,color=:red,label="GridPotentialFlow.jl")

#=
## Added mass
=#

# The added mass tensor provides a measure of the inertial influence of the fluid on the body in response to changes in the body's translational or rotational motion. The coefficients of the added mass tensor of a body are obtained by computing the impulse components associated with a unit-valued component of motion. The motion's influence is both direct, via the surface velocity, and indirect, in the bound vortex sheet that develops on the surface.

#=
The added mass for simple geometries can easiliy be calculated in an analytical way. For example, the entries of the translational added mass tensor for an ellipse with semi-major axis $a$ and semi-minor axis $b$ are $m_{xx} = \rho \pi b^2$ for motion in the $x$ direction, $m_{yy} = \rho \pi a^2$ for motion in the $y$-direction, with the diagonal entries $m_{xy}=m_{yx}=0$. By using `addedmass` from this package we can approximate these results numerically. In this case, we have one body and the method will return a matrix with the entries

$\begin{bmatrix}
m_{xx} & m_{xy} \\
m_{yx} & m_{yy}
\end{bmatrix}.$

=#

#md # ```@setup 6.-Force-and-the-added-mass
#md # Lx = 2.0
#md # xlim = (-Lx/2,Lx/2)
#md # ylim = (-Lx/2,Lx/2)
#md # g = PhysicalGrid(xlim,ylim,0.01)
#md # Δx = cellsize(g);
#md # ```
#!md Lx = 2.0
#!md xlim = (-Lx/2,Lx/2)
#!md ylim = (-Lx/2,Lx/2)
#!md g = PhysicalGrid(xlim,ylim,0.01)
#!md Δx = cellsize(g);

a = 0.5
b = 0.25
ellipse = Ellipse(a,b,Δx)
pfb = PotentialFlowBody(ellipse)
model = VortexModel(g,bodies=[pfb])
M = GridPotentialFlow.addedmass(model)

# We can compare the values using the `@test` macro.
using Test
@test isapprox(M[1,1], π*b^2, rtol=1e-1)
#
@test isapprox(M[2,2], π*a^2, atol=1e-1)

# In the case of $N$ bodies, `addedmass` returns the translational added mass tensor of size $2N$-by-$2N$. As an example, we will compute the added mass tensor for an array of cylinders. To compare it with an analytically obtained solution, we will calculate the added mass matrix tensor for the following gap-to-radius ratios.

GRratios = [0.1,0.2,0.3,0.4,0.5,0.6,0.8,1.0,1.2,1.4,1.6];

#md # ```@setup 6.-Force-and-the-added-mass
#md # function rectangulararray(rows,columns,spacing)
#md #     centers = VectorData(rows*columns)
#md #     for j in 1:columns
#md #         for i in 1:rows
#md #             xc = (i-1)*spacing
#md #             yc = (j-1)*spacing
#md #             centers.u[(j-1)*rows+i] = xc
#md #             centers.v[(j-1)*rows+i] = yc
#md #         end
#md #     end
#md #     centers.u .-= mean(centers.u)
#md #     centers.v .-= mean(centers.v)
#md #     return centers
#md # end
#md # ```

#!md function rectangulararray(rows,columns,spacing)
#!md      centers = VectorData(rows*columns)
#!md      for j in 1:columns
#!md         for i in 1:rows
#!md             xc = (i-1)*spacing
#!md             yc = (j-1)*spacing
#!md             centers.u[(j-1)*rows+i] = xc
#!md             centers.v[(j-1)*rows+i] = yc
#!md         end
#!md     end
#!md     centers.u .-= mean(centers.u)
#!md     centers.v .-= mean(centers.v)
#!md     return centers
#!md  end

# For this case, the array consists of three rows of three cylinders.

n = 100
rows = 3
columns = 3
N = rows*columns
R = 0.20
𝒱 = π*R^2
bodies = fill(Circle(R,2*cellsize(g)),N);
Δs = minimum(dlength(bodies[1]));

# We loop over the gap-to-radius ratios, position the bodies, and compute the ratio of the largest eigenvalue to the largest diagonal element of the translational added mass tensor. To position the bodies, we defined the method `rectangulararray` (see notebook in the examples folder).
using Statistics: mean
using LinearAlgebra: eigen
#
λoverMratios = zeros(length(GRratios))
for idx in 1:length(GRratios)
    global bodies
    gap = GRratios[idx]*R
    spacing = 2*R+gap
    bodycenters = rectangulararray(rows,columns,spacing)
    for i in 1:N
        Tf = RigidTransform((bodycenters.u[i],bodycenters.v[i]),0.0)
        global bodies[i] = Tf(deepcopy(bodies[i]))
    end
    body_list = PotentialFlowBody.(bodies)
    vm = VortexModel(g,bodies=body_list);
    M = GridPotentialFlow.addedmass(vm)/𝒱
    eigM = eigen(M);
    max_eig_value_coef = maximum(real(eigM.values))
    max_self_added_mass_coef = maximum(M)
    λoverMratios[idx] = max_eig_value_coef/max_self_added_mass_coef
end

# The same ratios, obtained in an analytical way, are available in literature [^1] and can be used to verify the numerical values.
nineRodArrayChen = [2.3637,2.2092,2.1007,2.0120,1.9350,1.8665,1.7494,1.6531,1.5732,1.5066,1.4508];
plot(xlabel="GR",ylabel="λ/max(Mij)")
plot!(GRratios,nineRodArrayChen,label="Chen1975")
plot!(GRratios,λoverMratios,label="GridPotentialFlow.jl")

# [^1]: Chen, S. S. (1975) "Vibration of nuclear fuel bundles," *Nuclear Engineering and Design*, 35 (3), 399-–422.


# ---
# title: 1043. Partition Array for Maximum Sum
# id: problem1043
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Dynamic Programming
# link: <https://leetcode.com/problems/partition-array-for-maximum-sum/description/>
# hidden: true
# ---
#
# Given an integer array `arr`, you should partition the array into (contiguous)
# subarrays of length at most `k`. After partitioning, each subarray has their
# values changed to become the maximum value of that subarray.
#
# Return _the largest sum of the given array after partitioning._
#
#
#
# **Example 1:**
#
#
#
#     Input: arr = [1,15,7,9,2,5,10], k = 3
#     Output: 84
#     Explanation: arr becomes [15,15,15,9,10,10,10]
#
#
# **Example 2:**
#
#
#
#     Input: arr = [1,4,1,5,7,3,6,1,9,9,3], k = 4
#     Output: 83
#
#
# **Example 3:**
#
#
#
#     Input: arr = [1], k = 1
#     Output: 1
#
#
#
#
# **Constraints:**
#
#   * `1 <= arr.length <= 500`
#   * `0 <= arr[i] <= 109`
#   * `1 <= k <= arr.length`
#
#
## @lc code=start
using LeetCode

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


abstract type Strength end
struct Classical{T} <: Strength
    θ::T
end
Classical(;θ = 0.25) = Classical(θ)

function (c::Classical)(At::SparseMatrixCSC{Tv,Ti}) where {Ti,Tv}

    θ = c.θ

    m, n = size(At)
    T = deepcopy(At)

    for i = 1:n
        _m = find_max_off_diag(T, i)
        threshold = θ * _m
        for j in nzrange(T, i)
            row = T.rowval[j]
            val = T.nzval[j]

            if row != i
                if abs(val) >= threshold
                    T.nzval[j] = abs(val)
                else
                    T.nzval[j] = 0
                end
            end

        end
    end

    dropzeros!(T)

    scale_cols_by_largest_entry!(T)

    adjoint(T), T
end

function find_max_off_diag(A, i)
    m = zero(eltype(A))
    for j in nzrange(A, i)
        row = A.rowval[j]
        val = A.nzval[j]
        if row != i
            m = max(m, abs(val))
        end
    end
    m
end

function find_max(A, i)
    m = zero(eltype(A))
    for j in nzrange(A, i)
        row = A.rowval[j]
        val = A.nzval[j]
        m = max(m, val)
    end
    m
end

function scale_cols_by_largest_entry!(A::SparseMatrixCSC)
    n = size(A, 1)
    for i = 1:n
        _m = find_max(A, i)
        for j in nzrange(A, i)
            A.nzval[j] /= _m
        end
    end
    A
end

struct SymmetricStrength{T} <: Strength
    θ::T
end
SymmetricStrength() = SymmetricStrength(0.)

function (s::SymmetricStrength{T})(A, bsr_flag = false) where {T}

    θ = s.θ

    if bsr_flag && θ == 0
        S = SparseMatrixCSC(size(A)...,
                    A.colptr, A.rowval, ones(eltype(A), size(A.rowval)))
        return S, S
    else
        S = deepcopy(A)
    end
    n = size(A, 1)
    diags = Vector{eltype(A)}(undef, n)

    for i = 1:n
        diag = zero(eltype(A))
        for j in nzrange(A, i)
            row = A.rowval[j]
            val = A.nzval[j]
            if row == i
                diag += val
            end
        end
        diags[i] = norm(diag)
    end

    for i = 1:n
        eps_Aii = θ * θ * diags[i]
        for j in nzrange(A, i)
            row = A.rowval[j]
            val = A.nzval[j]
            if row != i
                if val*val < eps_Aii * diags[row]
                    S.nzval[j] = 0
                end
            end
        end
    end

    dropzeros!(S)

    S.nzval .= abs.(S.nzval)
    scale_cols_by_largest_entry!(S)

    S, S
end


# MonoidAlternative
# =================

a = [1,2,3]
b = [4,5,6]

@test a ⊕ b == [a; b]

# FunctorApplicativeMonad
# =======================

@test eltype(Vector{Int}) == Int

@test map(x -> x*x, a) == [1, 4, 9]

product_ab = [5, 6, 7, 6, 7, 8, 7, 8, 9]
@test mapn(a, b) do x, y
  x + y
end == product_ab

h = @syntax_flatmap begin
  x = a
  y = b
  @pure x + y
end
@test h == product_ab

@test pure(Vector{Int}, 3) == [3]

# flatten combinations with other Monads
# --------------------------------------

@test flatten([Identity(3), Identity(4)]) == [3, 4]
@test flatten([Option(3), Option(nothing), Option(4)]) == [3, 4]
@test flatten([Try(4), (@Try error("hi")), Try(5)]) == [4, 5]
@test flatten([Either{String}(4), either("left", false, 3), Either{Int, String}("right")]) == [4, "right"]

h1 = @syntax_flatmap begin
 a, b = [(1, 2), (2, 3), (3, 4), (4, 5), (5, 6)]
 iftrue(a % 2 == 0) do
   a + b
 end
end
@test h1 == [5, 9]

# FlipTypes
# ========

v = [Option(:a),
     Option(:b),
     Option(:c)]

@test flip_types(v) == Option([:a, :b, :c])


export AffineMap

"""
    AffineMap{N<:Real, S<:LazySet{N}, NM, MAT<:AbstractMatrix{NM},
              VN<:AbstractVector{NM}} <: AbstractAffineMap{N, S}

Type that represents an affine transformation ``M⋅X ⊕ v`` of a convex set ``X``.

### Fields

- `M` -- matrix/linear map
- `X` -- convex set
- `v` -- translation vector

### Notes

An affine map is the composition of a linear map and a translation. This type is
parametric in the coefficients of the linear map, `NM`, which may be different from
the numeric type of the wrapped set (`N`). However, the numeric type of the
translation vector should be `NM`.

### Examples

For the examples we create a ``3×2`` matrix, a two-dimensional unit square, and
a three-dimensional vector.
Then we combine them in an `AffineMap`.

```jldoctest constructors
julia> A = [1 2; 1 3; 1 4]; X = BallInf([0, 0], 1); b2 = [1, 2]; b3 = [1, 2, 3];

julia> AffineMap(A, X, b3)
AffineMap{Int64,BallInf{Int64,Array{Int64,1}},Int64,Array{Int64,2},Array{Int64,1}}([1 2; 1 3; 1 4], BallInf{Int64,Array{Int64,1}}([0, 0], 1), [1, 2, 3])
```

For convenience, `A` does not need to be a matrix but we also allow to use
`UniformScaling`s resp. scalars (interpreted as a scaling, i.e., a scaled
identity matrix).
Scaling by ``1`` is ignored and simplified to a pure `Translation`.

```jldoctest constructors
julia> using LinearAlgebra

julia> am = AffineMap(2I, X, b2)
AffineMap{Int64,BallInf{Int64,Array{Int64,1}},Int64,Diagonal{Int64,Array{Int64,1}},Array{Int64,1}}([2 0; 0 2], BallInf{Int64,Array{Int64,1}}([0, 0], 1), [1, 2])

julia> AffineMap(2, X, b2) == am
true

julia> AffineMap(1, X, b2)
Translation{Int64,Array{Int64,1},BallInf{Int64,Array{Int64,1}}}(BallInf{Int64,Array{Int64,1}}([0, 0], 1), [1, 2])
```

Applying a linear map to an `AffineMap` object combines the two maps into a new
`AffineMap` instance.
Again we can make use of the conversion for convenience.

```jldoctest constructors
julia> B = [2 0; 0 2]; am2 = B * am
AffineMap{Int64,BallInf{Int64,Array{Int64,1}},Int64,Array{Int64,2},Array{Int64,1}}([4 0; 0 4], BallInf{Int64,Array{Int64,1}}([0, 0], 1), [2, 4])

julia> 2 * am == am2
true
```

The application of an `AffineMap` to a `ZeroSet` or an `EmptySet` is simplified
automatically.

```jldoctest constructors
julia> AffineMap(A, ZeroSet{Int}(2), b3)
Singleton{Int64,Array{Int64,1}}([1, 2, 3])

julia> AffineMap(A, EmptySet{Int}(2), b3)
EmptySet{Int64}(2)
```
"""
struct AffineMap{N<:Real, S<:LazySet{N}, NM, MAT<:AbstractMatrix{NM},
                 VN<:AbstractVector{NM}} <: AbstractAffineMap{N, S}
    M::MAT
    X::S
    v::VN

    # default constructor with dimension match check
    function AffineMap(M::MAT, X::S, v::VN) where {N<:Real, S<:LazySet{N}, NM,
                                                   MAT<:AbstractMatrix{NM},
                                                   VN<:AbstractVector{NM}}

        @assert dim(X) == size(M, 2) "a matrix of size $(size(M)) cannot be " *
            "applied to a set of dimension $(dim(X))"

        @assert size(M, 1) == length(v) "a map with output dimension " *
            "$(size(M, 1)) is incompatible with a translation vector of " *
            "dimension $(length(v))"

        return new{N, S, NM, MAT, VN}(M, X, v)
    end
end

isoperationtype(::Type{<:AffineMap}) = true
isconvextype(::Type{<:AffineMap{N, S}}) where {N, S} = isconvextype(S)

# convenience constructor from a UniformScaling
function AffineMap(M::UniformScaling, X::LazySet, v::AbstractVector)
    return AffineMap(M.λ, X, v)
end

# convenience constructor from a scalar
function AffineMap(α::N, X::LazySet, v::AbstractVector) where {N<:Real}
    if α == one(N)
        return Translation(X, v)
    end
    return AffineMap(Diagonal(fill(α, length(v))), X, v)
end

# simplification for a LinearMap
for MAP in (:AbstractMatrix, :Real)
    @eval begin
        function LinearMap(map::$MAP, am::AffineMap)
            return AffineMap(map * am.M, am.X, map * am.v)
        end
    end
end

# ZeroSet is "almost absorbing" for the linear map (only the dimension changes)
# such that only the translation vector remains
function AffineMap(M::AbstractMatrix{N}, Z::ZeroSet{N}, v::AbstractVector{N}
                  ) where {N<:Real}
    @assert dim(Z) == size(M, 2) "a matrix of size $(size(M)) cannot be " *
        "applied to a set of dimension $(dim(Z))"

    @assert size(M, 1) == length(v) "a map with output dimension " *
        "$(size(M, 1)) is incompatible with a translation vector of " *
        "dimension $(length(v))"

    return Singleton(v)
end

# EmptySet is absorbing for AffineMap
function AffineMap(M::AbstractMatrix{N}, ∅::EmptySet{N}, v::AbstractVector{N}
                  ) where {N<:Real}
    return ∅
end

# --- AbstractAffineMap interface functions ---

function matrix(am::AffineMap)
    return am.M
end

function vector(am::AffineMap{N}) where {N<:Real}
    return am.v
end

function set(am::AffineMap)
    return am.X
end

# --- LazySet interface functions ---

function concretize(am::AffineMap)
    return affine_map(am.M, concretize(am.X), am.v)
end


struct InstructionList
    instructions::Vector{Tuple{Symbol,Tuple{Symbol,Any,Any}}}
    var::Symbol
    n::Base.RefValue{Int}
end

function InstructionList(; var::Symbol = :_ι, n::Base.RefValue{Int} = Ref(0))
    InstructionList(Tuple{Symbol,Tuple{Symbol,Any,Any}}[], var, n)
end

function instruction_list(
    exprs::Vector{Expression};
    subexpressions::Dict{Expression,Expression} = Dict{Expression,Expression}(),
    perform_cse::Bool = true,
)
    v = InstructionList()
    PSE = Set{Symbol}()
    if perform_cse
        exprs, CSE = cse(exprs)
        merge!(CSE, subexpressions)
    else
        CSE = subexpressions
    end
    v, map(ex -> flat_expr!(v, ex, CSE, PSE), exprs)
end

Base.length(v::InstructionList) = v.n[]
function Base.show(io::IO, v::InstructionList)
    for (id, arg) in v.instructions
        println(io, :($id = $(Expr(:call, arg...))))
    end
end

function Base.push!(v::InstructionList, x)
    id = Symbol(v.var, v.n[] += 1)
    push!(v.instructions, (id, x))
    id
end
function Base.push!(v::InstructionList, x::Pair{Symbol,<:Tuple{Symbol,Any,Any}})
    push!(v.instructions, (first(x), last(x)))
    first(x)
end

function add!(list::InstructionList, @nospecialize(a), @nospecialize(b))
    if a === nothing
        b === nothing && return nothing
        return b
    else
        b === nothing && return a
        return push!(list, (:+, a, b))
    end
end
function sub!(list::InstructionList, @nospecialize(a), @nospecialize(b))
    if a === nothing
        b === nothing && return nothing
        return push!(list, (:-, b, nothing))
    else
        b === nothing && return a
        return push!(list, (:-, a, b))
    end
end

function mul!(list::InstructionList, @nospecialize(a), @nospecialize(b))
    (a === nothing || b === nothing) && return nothing
    if a == 1
        b == 1 && return a
        return b
    else
        b == 1 && return a
        return push!(list, (:*, a, b))
    end
end

function muladd!(
    list::InstructionList,
    @nospecialize(a),
    @nospecialize(b),
    @nospecialize(c),
)
    add!(list, c, mul!(list, a, b))
end

function div!(list::InstructionList, @nospecialize(a), @nospecialize(b))
    a === nothing && return nothing
    b == 1 && return a
    push!(list, (:/, a, b))
end

function pow!(list, @nospecialize(a), b::Int)
    @assert b ≥ 0
    b == 1 && return a
    b == 0 && return 1
    a == 1 && return a
    a === nothing && return nothing
    push!(list, (:^, a, b))
end

function flat_expr!(
    v,
    ex::Basic,
    # common sub expression
    CSE::Dict{Expression,Expression},
    # processed sub expression
    PSE::Set{Symbol} = Set{Symbol}(),
)
    t = class(ex)
    if t == :Symbol
        s = Symbol(to_string(ex))
        if haskey(CSE, ex) && s ∉ PSE
            push!(PSE, s)
            val = flat_expr!(v, CSE[ex], CSE, PSE)
            if val isa Symbol
                v.instructions[end] = (s, last(v.instructions[end]))
            else
                push!(v, s => val)
            end
        end
        return s
    elseif t == :Pow
        x, k = ModelKit.args(ex)
        return push!(v, (:^, flat_expr!(v, x, CSE, PSE), convert(Int, k)))
    elseif t == :Add || t == :Mul
        op = t == :Add ? :+ : :*
        vec = ModelKit.args(ex)
        op_arg = flat_expr!(v, vec[1], CSE, PSE)
        for i = 2:length(vec)
            op_arg = push!(v, (op, op_arg, flat_expr!(v, vec[i], CSE, PSE)))
        end
        return op_arg
    elseif t == :Div
        x, y = ModelKit.args(ex)
        return push!(v, (:/, flat_expr!(v, x, CSE, PSE), flat_expr!(v, y, CSE, PSE)))
    else
        return to_number(ex)
    end
end

struct DiffMap
    D::Dict{Tuple{Symbol,Int},Any}
end
DiffMap() = DiffMap(Dict{Tuple{Symbol,Int},Any}())
function Base.getindex(D::DiffMap, s::Symbol, i::Int)
    i == 0 && return s
    get(D.D, (s, i), nothing)
end
function Base.getindex(D::DiffMap, s, i)
    i == 0 && return s
    nothing
end
Base.setindex!(D::DiffMap, v::Any, s::Symbol, i::Int) = D.D[(s, i)] = v
Base.setindex!(D::DiffMap, v::Nothing, s::Symbol, i::Int) = nothing

function Base.diff(list::InstructionList, vars::Vector{Symbol}, f::Vector{Symbol})
    diff_map = DiffMap()
    for (i, v) in enumerate(vars)
        diff_map[v, i] = 1
    end
    dlist = diff!(list, length(vars), diff_map)

    J = Matrix{Union{Nothing,Symbol,Number}}(undef, length(f), length(vars))
    for (j, v) in enumerate(vars), (i, fi) in enumerate(f)
        J[i, j] = diff_map[fi, j]
    end
    dlist, J
end

"""
    diff!(list::InstructionList, N::Int, diff_map)

Returns an `InstructionList` list computing the derivatives of `N` variables.
The derivatives of the instructions are stored in `diff_map`.
"""
function diff!(list::InstructionList, N::Int, D::DiffMap)
    n = length(list)
    v = InstructionList(n = Ref(n))
    for (id, el) in list.instructions
        (op, arg1, arg2) = el
        if op == :^
            p2 = :NONE
            instr_added = false
            for ∂i = 1:N
                exp::Int = arg2
                arg1_∂i = D[arg1, ∂i]
                if arg1_∂i !== nothing
                    if p2 == :NONE
                        p1 = pow!(v, arg1, exp - 1)
                        p2 = mul!(v, exp, p1)
                    end
                    if !instr_added
                        if exp == 2
                            push!(v, id => (:^, arg1, 2))
                        else
                            push!(v, id => (:*, p1, arg1))
                        end
                        instr_added = true
                    end
                    D[id, ∂i] = mul!(v, p2, arg1_∂i)
                elseif p2 != :NONE && !instr_added
                    push!(v, id => el)
                    instr_added = true
                end

            end
            if !instr_added
                push!(v, id => el)
            end
        elseif op == :*
            for ∂i = 1:N
                ∂i == 1 && push!(v, id => el)
                a = mul!(v, D[arg1, ∂i], arg2)
                b = mul!(v, arg1, D[arg2, ∂i])
                D[id, ∂i] = add!(v, a, b)
            end
        elseif op == :/
            for ∂i = 1:N
                ∂i == 1 && push!(v, id => el)

                D[id, ∂i] = div!(
                    v,
                    sub!(v, mul!(v, D[arg1, ∂i], arg2), mul!(v, arg1, D[arg2, ∂i])),
                    pow!(v, arg2, 2),
                )
            end
        elseif op == :+
            for ∂i = 1:N
                ∂i == 1 && push!(v, id => el)
                D[id, ∂i] = add!(v, D[arg1, ∂i], D[arg2, ∂i])
            end
        end
    end
    v
end

## Higher order AD. This follows Chapter 13 of Griewank and Walther
"""
    univariate_diff!(list::InstructionList, K::Int, diff_map)

Returns an `InstructionList` list computing the first `K` derivatives.
The derivatives of the instructions are stored in `diff_map`.
"""
function univariate_diff!(list::InstructionList, K::Int, D::DiffMap)
    n = length(list)
    v = InstructionList(n = Ref(n))
    for (id, el) in list.instructions
        (op, arg1, arg2) = el
        if op == :^
            r = arg2
            arg1_1 = D[arg1, 1]
            if arg1_1 !== nothing && r == 2
                push!(v, id => el)
                for k = 1:K
                    w_k = nothing
                    # Compute ∑_{j=0}^k u_j U_{k - j} =
                    #         2(∑_{j=0}^div(k - 1,2) u_j U_{k - j}) +
                    #         iseven(k) * u_{k-1}^2
                    for j = 0:div(k - 1, 2)
                        w_k = muladd!(v, D[arg1, j], D[arg1, k-j], w_k)
                    end
                    w_k = mul!(v, 2, w_k)
                    if iseven(k)
                        w_k = add!(v, w_k, pow!(v, D[arg1, div(k, 2)], 2))
                    end
                    D[id, k] = w_k
                end
            elseif arg1_1 !== nothing
                w_0 = push!(v, id => (:^, arg1, r))
                u_0_inv = push!(v, (:inv_not_zero, arg1, nothing))
                # w_k is v_k in Griewank
                w = Symbol[]
                for k = 1:K
                    s = nothing
                    for j = 1:k
                        ũ_j = mul!(v, j, D[arg1, j])
                        w_kj = k == j ? w_0 : w[k-j]
                        s = muladd!(v, w_kj, ũ_j, s)
                    end
                    s = mul!(v, r, s)

                    t = nothing
                    for j = 1:k-1
                        u_kj = D[arg1, k-j]
                        w̃_j = mul!(v, j, w[j])
                        t = muladd!(v, u_kj, w̃_j, t)
                    end
                    w̃_k = mul!(v, u_0_inv, sub!(v, s, t))
                    w_k = div!(v, w̃_k, k)
                    push!(w, w_k)
                    D[id, k] = w_k
                end
            else
                push!(v, id => el)
            end
        elseif op == :*
            push!(v, id => el)
            for k = 1:K
                c_k = nothing
                for j = 0:k
                    c_k = muladd!(v, D[arg1, j], D[arg2, k-j], c_k)
                end
                # c_k === nothing && break
                D[id, k] = c_k
            end
        elseif op == :+
            push!(v, id => el)
            for k = 1:K
                D[id, k] = add!(v, D[arg1, k], D[arg2, k])
            end
        elseif op == :/
            push!(v, id => el)
            for k = 1:K
                s = nothing
                for j = 0:(k-1)
                    s = muladd!(v, D[id, j], D[arg2, k-j], s)
                end
                D[id, k] = div!(v, sub!(v, D[arg1, k], s), arg2)
            end
        end
    end
    v
end

@inline sqr(x) = x^2
@inline function sqr(z::Complex)
    x, y = reim(z)
    Complex((x + y) * (x - y), (x + x) * y)
end

"""
    inv_not_zero(x)

Invert x unless it is 0, then return 0.
"""
@inline inv_not_zero(x) = ifelse(iszero(x), x, inv(x))

function unroll_pow(var, n)
    n == 0 && return :(one($var))
    n == 1 && return var
    n == 2 && return :(sqr($var))
    n == 3 && return :($var * sqr($var))
    n == 4 && return :(sqr(sqr($var)))
    n == 5 && return :($var * sqr(sqr($var)))

    # base to expansion shows up which power it is needed to compute
    d = digits(n, base = 2)
    x = :x
    exprs = map(2:length(d)) do i
        :(local $(Symbol(x, 1 << (i - 1))) = sqr($(Symbol(x, 1 << (i - 2)))))
    end
    prods = Symbol[]
    for (i, di) in enumerate(d)
        if !iszero(di)
            push!(prods, Symbol(x, 1 << (i - 1)))
        end
    end
    if length(prods) > 1
        push!(exprs, Expr(:call, :*, prods...))
    end
    Expr(:let, :($(Symbol(x, 1)) = $var), Expr(:block, exprs...))
end

quick_div(a, b::Complex) = Base.FastMath.div_fast(a, b)
quick_div(a, b) = a / b

function to_expr(
    list::InstructionList,
    var_map = Dict{Symbol,Union{Expr,Symbol}}(),
    assignements = Dict{Symbol,Vector{Expr}}(),
)
    block = Expr(:block)
    exprs = block.args
    for (id, (op, arg1, arg2)) in list.instructions
        if op == :^
            x::Union{Symbol,Expr} = arg1
            if arg2 isa Int
                k::Int = arg2
                if k < 0
                    push!(exprs, :($id = inv($(unroll_pow(get(var_map, x, x), -k)))))
                else
                    push!(exprs, :($id = $(unroll_pow(get(var_map, x, x), k))))
                end
            else
                push!(exprs, :($id = $(get(var_map, x, x))^$arg2))
            end
        elseif op == :/
            a = get(var_map, arg1, arg1)
            b = get(var_map, arg2, arg2)
            push!(exprs, :($id = quick_div($a, $b)))
        else
            a = get(var_map, arg1, arg1)
            if arg2 !== nothing
                b = get(var_map, arg2, arg2)
                push!(exprs, :($id = $(Expr(:call, op, a, b))))
            else
                push!(exprs, :($id = $(Expr(:call, op, a))))
            end
        end

        if haskey(assignements, id)
            append!(exprs, assignements[id])
        end
    end
    block
end

function untuple(varname, tuplename, M)
    :($(Expr(:tuple, varname, (Symbol(varname, k) for k = 1:M-1)...)) = $tuplename)
end

function _impl_taylor_bivariate(K, M, N, op; inline = true)
    list = ModelKit.InstructionList()
    id = push!(list, (op, :x, :y))
    diff_map = ModelKit.DiffMap()
    for i = 1:M-1
        diff_map[:x, i] = Symbol(:x, i)
    end
    for j = 1:N-1
        diff_map[:y, j] = Symbol(:y, j)
    end
    dlist = ModelKit.univariate_diff!(list, K, diff_map)
    quote
        $(inline ? :(Base.@_inline_meta) : :())
        $(untuple(:x, :tx, M))
        $(untuple(:y, :ty, N))
        $(ModelKit.to_expr(dlist))
        $(Expr(:tuple, id, ((d = diff_map[id, k];
        d === nothing ? :(zero(x)) : d) for k = 1:K)...))
    end
end

@generated function taylor(
    ::Type{Val{op}},
    ::Type{Val{K}},
    tx::NTuple{M},
    ty::NTuple{N},
) where {op,K,M,N}
    _impl_taylor_bivariate(K, M, N, op)
end

make_ntuple(t) = t
make_ntuple(t::Number) = (t,)
make_ntuple(t::Tuple{A,B}) where {A,B} = convert(NTuple{2,promote_type(A, B)}, t)
make_ntuple(t::Tuple{A,B,C}) where {A,B,C} = convert(NTuple{3,promote_type(A, B, C)}, t)
make_ntuple(t::Tuple{A,B,C,D}) where {A,B,C,D} =
    convert(NTuple{4,promote_type(A, B, C, D)}, t)
make_ntuple(t::Tuple{A,B,C,D,E}) where {A,B,C,D,E} =
    convert(NTuple{5,promote_type(A, B, C, D, E)}, t)

taylor(op::Type{T1}, v::Type{T2}, tx, ty) where {T1,T2} =
    taylor(op::Type{T1}, v::Type{T2}, make_ntuple(tx), make_ntuple(ty))
function taylor(op::Type{T1}, v::Type{T2}, tx) where {T1,T2}
    taylor(op::Type{T1}, v::Type{T2}, make_ntuple(tx))
end

function _impl_taylor_pow(K, M)
    list = ModelKit.InstructionList()
    id = push!(list, (:^, :x, :r))
    diff_map = ModelKit.DiffMap()
    for i = 1:M-1
        diff_map[:x, i] = Symbol(:x, i)
    end
    dlist = ModelKit.univariate_diff!(list, K, diff_map)
    quote
        $(untuple(:x, :tx, M))
        $(ModelKit.to_expr(dlist))
        $(Expr(:tuple, id, ((d = diff_map[id, k];
        d === nothing ? :(zero(x)) : d) for k = 1:K)...))
    end
end

@generated function taylor(
    ::Type{Val{:^}},
    ::Type{Val{K}},
    tx::NTuple{M},
    r::Integer,
) where {K,M}
    _impl_taylor_pow(K, M)
end

function _impl_taylor_sqr(K, M)
    list = ModelKit.InstructionList()
    id = push!(list, (:^, :x, 2))
    diff_map = ModelKit.DiffMap()
    for i = 1:M-1
        diff_map[:x, i] = Symbol(:x, i)
    end
    dlist = ModelKit.univariate_diff!(list, K, diff_map)
    quote
        $(untuple(:x, :tx, M))
        $(ModelKit.to_expr(dlist))
        $(Expr(:tuple, id, ((d = diff_map[id, k];
        d === nothing ? :(zero(x)) : d) for k = 1:K)...))
    end
end

@generated function taylor(::Type{Val{:sqr}}, ::Type{Val{K}}, tx::NTuple{M}) where {K,M}
    _impl_taylor_sqr(K, M)
end