using SafeTestsets

@safetestset "SurrogatesFlux" begin

using Surrogates
using Surrogates: SobolSample
using Flux
using Flux: @epochs
using SurrogatesFlux
using LinearAlgebra
using Zygote

#1D
a = 0.0
b = 10.0
obj_1D = x -> 2*x+3
x = sample(10,0.0,10.,SobolSample())
y = obj_1D.(x);
my_model = Chain(Dense(1,1), first)
my_loss(x, y) = Flux.mse(my_model(x), y)
my_opt = Descent(0.01)
n_echos = 1
my_neural = NeuralSurrogate(x,y,a,b,model = my_model,loss = my_loss,opt = my_opt,n_echos=1)
my_neural_kwargs = NeuralSurrogate(x,y,a,b)
add_point!(my_neural,8.5,20.0)
add_point!(my_neural,[3.2,3.5],[7.4,8.0])
val = my_neural(5.0)

#ND

lb = [0.0,0.0]
ub = [5.0,5.0]
x = sample(5,lb,ub, SobolSample())
obj_ND_neural(x) = x[1]*x[2];
y = obj_ND_neural.(x)
my_model = Chain(Dense(2,1), first)
my_loss(x, y) = Flux.mse(my_model(x), y)
my_opt = Descent(0.01)
n_echos = 1
my_neural = NeuralSurrogate(x,y,lb,ub,model=my_model,loss=my_loss,opt=my_opt,n_echos=1)
my_neural_kwargs = NeuralSurrogate(x,y,lb,ub)
my_neural((3.5, 1.49))
my_neural([3.4,1.4])
add_point!(my_neural,(3.5,1.4),4.9)
add_point!(my_neural,[(3.5,1.4),(1.5,1.4),(1.3,1.2)],[1.3,1.4,1.5])

# Multi-output #98
f  = x -> [x^2, x]
lb = 1.0
ub = 10.0
x  = sample(5, lb, ub, SobolSample())
push!(x, 2.0)
y  = f.(x)
my_model = Chain(Dense(1,2))
my_loss(x, y) = Flux.mse(my_model(x), y)
surrogate = NeuralSurrogate(x,y,lb,ub,model = my_model,loss = my_loss ,opt=my_opt,n_echos=1)
surr_kwargs = NeuralSurrogate(x,y,lb,ub)

f  = x -> [x[1], x[2]^2]
lb = [1.0, 2.0]
ub = [10.0, 8.5]
x  = sample(20, lb, ub, SobolSample())
push!(x, (1.0, 2.0))
y  = f.(x)
my_model = Chain(Dense(2,2))
my_loss(x, y) = Flux.mse(my_model(x), y)
surrogate = NeuralSurrogate(x,y,lb,ub,model=my_model,loss=my_loss,opt=my_opt,n_echos=1)
surrogate_kwargs = NeuralSurrogate(x,y,lb,ub)
surrogate((1.0, 2.0))
x_new = (2.0, 2.0)
y_new = f(x_new)
add_point!(surrogate, x_new, y_new)

#Optimization
lb = [1.0,1.0]
ub = [6.0,6.0]
x = sample(5,lb,ub,SobolSample())
objective_function_ND = z -> 3*norm(z)+1
y = objective_function_ND.(x)
model = Chain(Dense(2,1), first)
loss(x, y) = Flux.mse(model(x), y)
opt = Descent(0.01)
n_echos = 1
my_neural_ND_neural = NeuralSurrogate(x,y,lb,ub)
surrogate_optimize(objective_function_ND,SRBF(),lb,ub,my_neural_ND_neural,SobolSample(),maxiters=15)

# AD Compatibility
lb = 0.0
ub = 3.0
n = 10
x = sample(n,lb,ub,SobolSample())
f = x -> x^2
y = f.(x)
#NN
@testset "NN" begin
    my_model = Chain(Dense(1,1), first)
    my_loss(x, y) = Flux.mse(my_model(x), y)
    my_opt = Descent(0.01)
    n_echos = 1
    my_neural = NeuralSurrogate(x,y,lb,ub,model=my_model,loss=my_loss,opt=my_opt,n_echos=1)
    g = x->my_neural'(x)
    g(3.4)
end

lb = [0.0,0.0]
ub = [10.0,10.0]
n = 5
x = sample(n,lb,ub,SobolSample())
f = x -> x[1]*x[2]
y = f.(x)

#NN
@testset "NN ND" begin
    my_model = Chain(Dense(2,1), first)
    my_loss(x, y) = Flux.mse(my_model(x), y)
    my_opt = Descent(0.01)
    n_echos = 1
    my_neural = NeuralSurrogate(x,y,lb,ub,model=my_model,loss=my_loss,opt=my_opt,n_echos=1)
    g = x -> Zygote.gradient(my_neural, x)
    g((2.0,5.0))
end

# ###### ND -> ND ######

lb = [0.0, 0.0]
ub = [10.0, 2.0]
n = 5
x = sample(n,lb,ub,SobolSample())
f = x -> [x[1]^2, x[2]]
y = f.(x)

#NN
@testset "NN ND -> ND" begin
    my_model = Chain(Dense(2,2))
    my_loss(x, y) = Flux.mse(my_model(x), y)
    my_opt = Descent(0.01)
    n_echos = 1
    my_neural = NeuralSurrogate(x,y,lb,ub,model=my_model,loss=my_loss,opt=my_opt,n_echos=1)
    Zygote.gradient(x -> sum(my_neural(x)), (2.0, 5.0))

    my_rad = RadialBasis(x,y,lb,ub,rad = linearRadial())
    Zygote.gradient(x -> sum(my_rad(x)), (2.0, 5.0))

    my_p = 1.4
    my_inverse = InverseDistanceSurrogate(x,y,lb,ub,p=my_p)
    my_inverse((2.0, 5.0))
    Zygote.gradient(x -> sum(my_inverse(x)), (2.0, 5.0))

    my_second = SecondOrderPolynomialSurrogate(x,y,lb,ub)
    Zygote.gradient(x -> sum(my_second(x)), (2.0, 5.0))
end
end


@testset "SymbolicExpressionDeriv1" begin
    env = SymbolicEnv()
    cache = SymbolicCache()
    x = 5.6
    env[:x] = 5.6
    expr = symbolic(:(4.0 * x + 1.6))
    y = seval(expr, :x, env, cache)
    @test y == 4.0
end

@testset "SymbolicExpressionDeriv2" begin
    env = SymbolicEnv()
    cache = SymbolicCache()
    x = 5.6
    env[:x] = 5.6
    expr = symbolic(:(-4.0 * x - 1.6 * x))
    y = seval(expr, :x, env, cache)
    @test y == -4-1.6
end

@testset "SymbolicExpressionDeriv3" begin
    env = SymbolicEnv()
    cache = SymbolicCache()
    x = 5.6
    env[:x] = 5.6
    expr = symbolic(:(-4.0 * x * x))
    println(expr)
    y = seval(expr, :x, env, cache)
    @test y == -8 * x
end

@testset "SymbolicExpressionDeriv4" begin
    env = SymbolicEnv()
    cache = SymbolicCache()
    x = 5.6
    y = 10.1
    env[:x] = x
    env[:y] = y
    expr = symbolic(:(-4.0 * x / y))
    res = seval(expr, :x, env, cache)
    @test res == -4 / y
end

@testset "SymbolicExpressionDeriv5" begin
    env = SymbolicEnv()
    cache = SymbolicCache()
    x = 5.6
    y = 10.1
    env[:x] = x
    env[:y] = y
    expr = symbolic(:(-4.0 * x / y))
    res = seval(expr, :y, env, cache)
    @test res == 4 * x / y^2
end

@testset "SymbolicExpressionDeriv6" begin
    env = SymbolicEnv()
    cache = SymbolicCache()
    x = 5.6
    y = 10.1
    env[:x] = x
    env[:y] = y
    expr = symbolic(:(-4 * x * y))
    res = seval(expr, :x, env, cache)
    @test res == -4 * y
end

@testset "SymbolicExpressionDeriv7" begin
    env = SymbolicEnv()
    cache = SymbolicCache()
    x = 5.6
    env[:x] = x
    expr = symbolic(:(-4 * x ^ 3))
    res = seval(expr, :x, env, cache)
    @test res == -12 * x^2
end

@testset "SymbolicExpressionDeriv8" begin
    env = SymbolicEnv()
    cache = SymbolicCache()
    x = 5.6
    y = 10.1
    env[:x] = x
    env[:y] = y
    expr = -log(symbolic(:(x*y))) * symbolic(:y)
    res = seval(expr, :x, env, cache)
    @test res == - y / (x * y) * y
    res = seval(expr, :y, env, cache)
    @test res == - y / (x * y) * x - log(x * y)
    println(cache)
end

@testset "SymbolicExpressionDeriv2_1" begin
    env = SymbolicEnv()
    cache = SymbolicCache()
    x = 5.6
    y = 10.1
    env[:x] = x
    env[:y] = y
    expr = symbolic(:(x^2 + 10*x*y + 4*y^2))
    res = seval(expr, (:x,:y), env, cache)
    @test res == 10.0
end

@testset "lam1" begin
    env = SymbolicEnv()
    cache = SymbolicCache()
    @bind env begin
        t = rand()
        lam1 = rand()
        lam2 = rand()
    end
    expr = exp(-lam1*t) - (lam1/lam2) * exp(-lam1*t) * (exp(-lam2*t) - 1)
    res1 = seval(expr, lam1, env, cache)
    @test res1 ≈ seval(-t*exp(-lam1*t) - (1/lam2)*exp(-lam1*t)*(exp(-lam2*t)-1) + t*(lam1/lam2)*exp(-lam1*t)*(exp(-lam2*t)-1), env)
    res2 = seval(expr, lam2, env, cache)
    @test res2 ≈ seval((lam1/lam2^2)*exp(-lam1*t)*(exp(-lam2*t)-1) + t*(lam1/lam2)*exp(-lam1*t)*exp(-lam2*t), env)
    res3 = seval(expr, (lam2,lam1), env, cache)
    res4 = seval(expr, (:lam1,:lam2), env, cache)
    @test isapprox(res3, res4)
    @test isapprox(seval((1/lam2^2)*exp(-lam1*t)*(exp(-lam2*t)-1) - t*(lam1/lam2^2)*exp(-lam1*t)*(exp(-lam2*t)-1) + t*(1/lam2)*exp(-lam1*t)*exp(-lam2*t) - t^2*(lam1/lam2)*exp(-lam1*t)*exp(-lam2*t), env),
                    seval((1/lam2^2)*exp(-lam1*t)*(exp(-lam2*t)-1) + t*(1/lam2)*exp(-lam1*t)*exp(-lam2*t) - t*(lam1/lam2^2)*exp(-lam1*t)*(exp(-lam2*t)-1) - t^2*(lam1/lam2)*exp(-lam1*t)*exp(-lam2*t), env))
    @test res3 ≈ seval((1/lam2^2)*exp(-lam1*t)*(exp(-lam2*t)-1) - t*(lam1/lam2^2)*exp(-lam1*t)*(exp(-lam2*t)-1) + t*(1/lam2)*exp(-lam1*t)*exp(-lam2*t) - t^2*(lam1/lam2)*exp(-lam1*t)*exp(-lam2*t), env)
    @test res4 ≈ seval((1/lam2^2)*exp(-lam1*t)*(exp(-lam2*t)-1) + t*(1/lam2)*exp(-lam1*t)*exp(-lam2*t) - t*(lam1/lam2^2)*exp(-lam1*t)*(exp(-lam2*t)-1) - t^2*(lam1/lam2)*exp(-lam1*t)*exp(-lam2*t), env)
end


@testset "Bernoulli sampler" begin

	cc = dummyChoiceContext()

	@testset "getparams, numparams, paramranges, sample" begin

		@testset "using default construction" begin

		    s = DataGenerators.BernoulliSampler()

		    @testset "numparams and paramranges" begin
		        @test DataGenerators.numparams(s) == 1
		        prs = DataGenerators.paramranges(s)
		        @test typeof(prs) <: Vector{Tuple{Float64,Float64}}
		        @test prs == [(0.0,1.0)]
		    end

		    @testset "default params" begin
		        @test DataGenerators.getparams(s) == [0.5]
		    end

		    @mtestset "default sampling" reps=Main.REPS alpha=Main.ALPHA begin
		        x, trace = DataGenerators.sample(s, (0,1), cc)
		        @test typeof(x) <: Int
		        @mtest_values_are [0,1] x
				@mtest_distributed_as Bernoulli(0.5) x
		    end

		end

		@testset "using non-default construction" begin

		    s = DataGenerators.BernoulliSampler([0.3])
		    @test DataGenerators.getparams(s) == [0.3]

		    @mtestset "consistent with Bernoulli" reps=Main.REPS alpha=Main.ALPHA begin
		    	x, trace = DataGenerators.sample(s, (0,1), cc)
		        @mtest_distributed_as Bernoulli(0.3) x
		    end

		end

	end

	@testset "setparams" begin

	    s = DataGenerators.BernoulliSampler()
	    prs = DataGenerators.paramranges(s)
	    midparams = map(pr->robustmidpoint(pr[1],pr[2]), prs)

	    @testset "valid parameters $params" for params in [[0.8], [0.5], [0.3],]

	        DataGenerators.setparams!(s, params)

			@mtestset "consistent with Bernoulli" reps=Main.REPS alpha=Main.ALPHA begin
	        	x, trace = DataGenerators.sample(s, (0,1), cc)
		        @mtest_distributed_as Bernoulli(params[1]) x
			end

	    end

	    @testset "boundary values for parameter index $pidx" for pidx in 1:length(prs)

	        pr = prs[pidx]
	        params = copy(midparams)

			@testset "bound index $bidx" for bidx in 1:2

	            params[pidx] = pr[bidx]
	            DataGenerators.setparams!(s, params)

				@mtestset "consistent with Bernoulli" reps=Main.REPS alpha=Main.ALPHA begin
	            	x, trace = DataGenerators.sample(s, (0,1), cc)
			        @mtest_distributed_as Bernoulli(params[1]) x
				end

				@testset "range check exception" begin
		            params[pidx] = bidx == 1 ? prevfloat(pr[bidx]) : nextfloat(pr[bidx])
		            @test_throws ErrorException DataGenerators.setparams!(s, params)
				end

			end

	    end

	    @testset "wrong number of parameters" begin

	        @test_throws ErrorException DataGenerators.setparams!(s, midparams[1:end-1])
	        @test_throws ErrorException DataGenerators.setparams!(s, [midparams; 0.5])

	    end

	end

	@testset "estimateparams" begin

	    otherparams = [0.5]

	    @testset "from parameters $params" for params in [[0.2], [0.7], [0.0], [1.0],]

	        s1 = DataGenerators.BernoulliSampler(params)
	        s2 = DataGenerators.BernoulliSampler(otherparams)
	        traces = map(1:100) do i
	            x, trace = DataGenerators.sample(s1, (0,1), cc)
	            trace
	        end
	        estimateparams!(s2, traces)

			@mtestset "consistent with Bernoulli" reps=Main.REPS alpha=Main.ALPHA begin
	        	x, trace = DataGenerators.sample(s2, (0,1), cc)
		        @mtest_distributed_as Bernoulli(params[1]) x
			end

	    end

	    @testset "too few traces" begin

	        params = [0.2]
	        s1 = DataGenerators.BernoulliSampler(params)
	        s2 = DataGenerators.BernoulliSampler(otherparams)
	        traces = map(1:0) do i
	            x, trace = DataGenerators.sample(s1, (0,1), cc)
	            trace
	        end
	        estimateparams!(s2, traces)

			@mtestset "consistent with Bernoulli" reps=Main.REPS alpha=Main.ALPHA begin
	        	x, trace = DataGenerators.sample(s2, (0,1), cc)
		        @mtest_distributed_as Bernoulli(otherparams[1]) x
			end

	    end

	end

end


abstract type ExactSolutionType end

export ExactSolutionType, ExactTrig, calcErrorL2Norm

mutable struct ExactTrig <: ExactSolutionType
end
function (obj::ExactTrig)(xy::AbstractArray{Tmsh}, q::AbstractArray{Tsol, 1}) where {Tmsh, Tsol}  # (numDofPerNode))
  k = 2.0
  q[:] = sin(2*k*pi*xy[1])*sin(2*k*pi*xy[2])
  return nothing
end

mutable struct ExactExpTrig <: ExactSolutionType
end
function (obj::ExactExpTrig)(xy::AbstractArray{Tmsh}, q::AbstractArray{Tsol, 1}) where {Tmsh, Tsol}  # (numDofPerNode))
  k = 2.0
  ex = exp(xy[1] + xy[2])
  q[:] = ex * sin(2*k*pi*xy[1])*sin(2*k*pi*xy[2])
  return nothing
end

mutable struct ExactPoly2nd <: ExactSolutionType
end
function (obj::ExactPoly2nd)(xy::AbstractArray{Tmsh}, q::AbstractArray{Tsol, 1}) where {Tmsh, Tsol}
  a = 1.0
  b = 1.0
  c = 1.0
  q[:] =  a*xy[1]*xy[1] + b*xy[1]*xy[2] + c*xy[2]*xy[2]
  return nothing
end

global const ExactDict = Dict{String, ExactSolutionType}(
 "ExactTrig" => ExactTrig(),
 "ExactPoly2nd" => ExactPoly2nd(),
 "ExactExpTrig" => ExactExpTrig()
)

function calcErrorL2Norm(mesh::AbstractMesh{Tmsh},
                         sbp::AbstractOperator,
                         eqn::AbstractEllipticData{Tsol, Tres},
                         opts) where {Tmsh, Tsol, Tres}
  l2norm::Float64 = 0.
  lInfnorm::Float64 = 0.
  qe = Array{Tsol}(mesh.numDofPerNode)
  exactFunc = ExactDict[opts["exactSolution"]]

  lInfnorm = -1.0
  elem_with_max_dq = 0
  for e = 1 : mesh.numEl
    for n = 1 : mesh.numNodesPerElement
      xy = sview(mesh.coords, :, n, e)
      exactFunc(xy, qe)
      q = sview(eqn.q, :, n, e)
      jac = mesh.jac[n, e]
      for v = 1:mesh.numDofPerNode
        # dq = Real(q[v] - qe[v])
        dq = Float64(q[v] - qe[v])
        l2norm += dq*dq*sbp.w[n]/jac
        if lInfnorm < abs(dq)
          lInfnorm = abs(dq)
        end
      end
    end
  end

  return sqrt(l2norm), lInfnorm
end


module JuliaPlot
using Plots, Latexify, LaTeXStrings, Printf, Polynomials, Measures

export juliaPlot,mandelbrotPlot,plotBoth,moreLikelyToBeInterestingPlot,julianimation,juliaMatrix,mandelbrotMatrix,multiMonitorJuliaPlot

function juliaPlot(
        ;
        c::Complex = complex(0.6,0.7), # function offset
        res::Integer = 200, # x resolution
        R::Float64 = 2.0, # Escape radius
        aleph::Float64 = 0.8, # Proportion of escape radius to draw
        I::Integer = 100, # maximum iteration number
        f::AbstractVector = [0.0,0.0,1.0], # function to iterate
        size::Tuple{Integer,Integer} = (1920,1080),
        filename::Union{String,Nothing} = nothing,
        color::Symbol = :hawaii,
        labelplatecolor::Union{Symbol,Nothing} = :white,
        pl = heatmap(
                     aspect_ratio=:equal,
                     axis=false,
                     legend=false,
                     ticks=[],
                     bg=:black,
                    )
    )
    g = Polynomial{Complex{Float64}}(f)+c
    x = range(-aleph*R,aleph*R,length=res)
    y = range(-aleph*R/size[1]*size[2],aleph*R/size[1]*size[2],
        length=round(Int,res/size[1]*size[2]))
    values = juliaValue.(g,complex.(x',y),R,I)
    heatmap!(pl,
        x,y,values,
        color=color,
       )
    #= Annotation (disable for now)
    io = IOBuffer()
    print(io,"\$")
    printpoly(io,g,MIME"text/latex"(),descending_powers=true)
    print(io,"\$")
    annotation = LaTeXString(String(take!(io)))
    annotation_size = Plots.text_size(annotation,14)
    annotcoords=(0.9-0.011*annotation_size[1]/annotation_size[2],1.0,-0.95,-0.85) # x1, x2, y1, y2
    if !isnothing(labelplatecolor)
    plot!(pl,
    Shape(
    [annotcoords[[1,2,2,1]]...]*aleph*R, # x1, x2, x2, x1
    [annotcoords[[3,3,4,4]]...]*aleph*R/size[1]*size[2] # y1, y1, y2, y2
    ),
    opacity=0.1,color=labelplatecolor,linecolor=nothing,
    padding=0mm,margins=0mm,
    axis=false,
    bg=:black,
    )
    end
    annotate!(pl,
    0.9*aleph*R,
    -0.9*aleph*R/size[1]*size[2],
    text(annotation,14,:right),
    size=size,
    padding=0mm,margins=0mm,
    axis=false,
    )
    =#

    if !isnothing(filename)
        #closeall() # workaround for darkening issue
        savefig(pl,filename)
    end
    return pl
end # function juliaPlot

function mandelbrotPlot(
        ;
        res::Integer = 200,
        R::Float64 = 2.0,
        aleph::Float64 = 0.8,
        I::Integer = 100,
        f::AbstractVector = [0.0,0.0,1.0],
        size::Tuple{Integer,Integer} = (1920,1080),
        filename::Union{String,Nothing} = nothing,
        color::Symbol = :hawaii,
        pl = heatmap(
                     aspect_ratio=:equal,
                     axis=false,
                     legend=false,
                     ticks=[],
                     bg=:black,
                    ),
    )::Complex{Float64}
    g = Polynomial{Complex{Float64}}(f)
    x = range(-aleph*R,aleph*R,length=res)
    y = range(-aleph*R/size[1]*size[2],aleph*R/size[1]*size[2],
        length=round(Int,res/size[1]*size[2]))
    values = juliaValue.(g.+complex.(x',y),complex(0.0,0.0),R,I)
    (y_i,x_i) = div.(size,2)
    try
        (y_i,x_i) = Tuple(rand(findall( (values .> 0.5*I) .& (values.< 0.9*I))))
    catch
    end
    c = complex(x[x_i],y[y_i])
    if !isnothing(pl)
        heatmap!(pl,
                 x,y,values,
                 color=color,
                )
        plot!(pl,
              [x[x_i]],[y[y_i]],
              seriestype=:scatter,
              marker = (:circle, size[1]/250, :white, stroke(0.2, 0.2, :black)),
             )
        if !isnothing(filename)
            savefig(pl,filename)
        end
    end
    return c
end # mandelbrotPlot

function plotBoth(
        ;
        f=[0.0,0.0,1.0],
        res=2000,
        sizes=(2000,1000),
        filename=nothing,
        color=:hawaii,
    )
    sizes = eltype(sizes)<:Integer ? (sizes,sizes) : sizes
    res = typeof(res)<:Tuple ? res : (res,res)
    filenames = eltype(filename)<:String ? filename : (nothing,nothing)
    pl1 = heatmap(
        aspect_ratio=:equal,
        size=sizes[1],
        axis=false,
        legend=false,
        ticks=[],
        bg=:black,
    )
    pl2 = heatmap(
        aspect_ratio=:equal,
        size=sizes[2],
        axis=false,
        legend=false,
        ticks=[],
        bg=:black,
    )
    c = mandelbrotPlot(
        ;
        f=f,
        res=res[1],
        size=sizes[1],
        pl=pl1,
        filename=filenames[1],
        color=color,
    )
    juliaPlot(
        ;
        c=c,
        f=f,
        res=res[2],
        size=sizes[2],
        pl=pl2,
        filename=filenames[2],
        color=color,
    )
    l = @layout([
                 a;
                 b
                ])
    pl = plot(pl1,pl2,
        layout=( # TODO I think you have to specify size here...
                l
               ),
        size=(maximum(getindex.(sizes,1)),sum(getindex.(sizes,2))))
    if eltype(filename)<:String
        savefig(pl,filename)
    end
    return pl
end # plotBoth

function moreLikelyToBeInterestingPlot(
        ;
        f=[0.0,0.0,1.0],
        res=(300,2000),
        sizes=((300,200),(2000,1000)),
        filename=nothing,
        color=:hawaii,
        labelplatecolor=:white,
    )
    sizes = eltype(sizes)<:Integer ? (sizes,sizes) : sizes
    res = typeof(res)<:Tuple ? res : (res,res)
    c = mandelbrotPlot(
        res=res[1],
        f=f,
        size=sizes[1],
        color=color,
        pl=nothing,
    )
    juliaPlot(
        c=c, # function offset
        res=res[2], # x resolution
        f=f, # function to iterate
        size=sizes[2],
        filename=filename,
        color=color,
        labelplatecolor=labelplatecolor,
    )
end

function juliaValue(
        f::Polynomial, # Function to evaluate
        z_0::Complex, # Starting point
        R::Real, # Escape radius
        I::Integer=100, # Max i
    )::Int64
    i = 0
    z = z_0
    R = R^2
    while( abs2(z) < R && i<I )
        z = f(z)
        i+=1
    end
    return i
end # function juliaValue

function juliaValues(
        f::Union{Function,Polynomial}, # Function to evaluate
        z_0::Complex, # Starting point
        R::Real, # Escape radius
        I::Integer=100, # Max i
    )::Vector{Complex{Float64}}
    R = R^2
    return collect(Complex{Float64},takewhile(x->abs2(x)<R,IterTools.take(iterated(f,z_0),I)))
end # function juliaValues

function julianimation(
        ;
        f=[0.0,0.0,1.0],
        res=500,
        N_t=10,
        filename=nothing,
        color=:rainbow,
        R::Float64 = 2.0, # Escape radius
        aleph::Float64 = 0.8, # Proportion of escape radius to draw
        I::Integer = 100, # maximum iteration number
        size::Tuple{Integer,Integer} = (1920,1080),
    )

    x = range(-aleph*R,aleph*R,length=res)
    y = range(-aleph*R/size[1]*size[2],aleph*R/size[1]*size[2],length=round(Int,res/size[1]*size[2]))
    mandelbrotvalues=juliaValue.(Polynomial{Complex{Float64}}(f).+complex.(x',y),complex(0.0,0.0),R,I)

    # Choose one coordinate inside and one outside the set
    extr = extrema(mandelbrotvalues)
    lims = [0.9 0.1;0.1 0.9]*[extr...]
    i_small = Tuple(rand(findall( (mandelbrotvalues .< lims[1]) )))
    c_small = Complex(x[i_small[2]],y[i_small[1]])
    i_large = Tuple(rand(findall( (mandelbrotvalues .> lims[2]) )))
    c_large = Complex(x[i_large[2]],y[i_large[1]])
    print("$c_small\t$c_large\t$extr\n")

    @animate for c in range(c_small,c_large,length=N_t)
        print("$c\n")
        juliavalues = juliaValue.(Polynomial{Complex{Float64}}(f)+c,complex.(x',y),R,I)
        pl= heatmap(x,y,juliavalues,ticks=[],colorbar=false)
        heatmap!(pl,x,y,mandelbrotvalues,inset=(1,bbox(0.05,0.05,0.5,0.5*size[2]/size[1],:bottom,:right)),ticks=[],colorbar=false,subplot=2)
        plot!(pl,real.([c_small,c_large]),imag.([c_small,c_large]),subplot=2)
        scatter!(pl,[real(c)],[imag(c)],subplot=2)
    end
end

function mandelbrotMatrix(;
        res::Integer = 200,
        R::Float64 = 2.0,
        aleph::Float64 = 0.8,
        I::Integer = 100,
        f::AbstractVector = [0.0,0.0,1.0],
        size::Tuple{Integer,Integer} = (1920,1080),
    )
    g = Polynomial{Complex{Float64}}(f)
    x = range(-aleph*R,aleph*R,length=res)
    y = range(-aleph*R/size[1]*size[2],aleph*R/size[1]*size[2],
        length=round(Int,res/size[1]*size[2]))
    values = juliaValue.(g.+complex.(x',y),complex(0.0,0.0),R,I)
end # mandelbrotMatrix

function juliaMatrix(;
        c::Complex = complex(0.6,0.7), # function offset
        res::Integer = 200, # x resolution
        R::Float64 = 2.0, # Escape radius
        aleph::Float64 = 0.8, # Proportion of escape radius to draw
        I::Integer = 100, # maximum iteration number
        f::AbstractVector = [0.0,0.0,1.0], # function to iterate
        size::Tuple{Integer,Integer} = (1920,1080),
    )
    g = Polynomial{Complex{Float64}}(f)+c
    x = range(-aleph*R,aleph*R,length=res)
    y = range(-aleph*R/size[1]*size[2],aleph*R/size[1]*size[2],
        length=round(Int,res/size[1]*size[2]))
    values = juliaValue.(g,complex.(x',y),R,I)
end # juliaMatrix

function multiMonitorJuliaPlot(;
        pos = ((0,0,1920,1080),(1920,0,2560,1440)), # tuple of x, y, w, h tuples
        filename = "~/Pictures/Julia/background.png",
        f = vcat(0,rand(-1:0.05:1,rand(2:10))),
        color = :copper,
    )
    if length(pos) == 1
        p = first(pos)
        pos = (p,(p[1]+p[3]÷20, p[2]+p[4]-p[4]÷20-p[4]÷5, p[3]÷5, p[4]÷5))
    end
    sz = (
          maximum(p[1] + p[3] for p in pos) - minimum(p[1] for p in pos),
          maximum(p[2] + p[4] for p in pos) - minimum(p[2] for p in pos)
    );
    b = Tuple( # collection of inset boxes
              bbox(p[1]/sz[1], (1-(p[2]+p[4])/sz[2]), p[3]/sz[1], p[4]/sz[2], :bottom, :left) for p in pos
             );
    p = pos[end]
    insetpos = ((p[1]+p[3]/20)/sz[1],1-(p[2]+0.9*p[4])/sz[2],(p[3]/5)/sz[1],(p[4]/40)/sz[2])
    gr(leg=false,framestyle=:none,ticks=nothing,margin=0mm,bg=:black);
    pl = plot(size=sz);
    for k = 1:length(pos)
        plot!(pl,[missing],inset=b[k],subplot=k+1)
    end
    plot!(pl,[missing],inset=bbox(insetpos..., :bottom,:left),subplot=length(pos)+2)

    c = mandelbrotPlot(
        res=length(pos)>1 ? pos[2][3] : 200,
        f=f,
        size=length(pos)>1 ? pos[2][3:4] : (200,200),
        color=color,
        pl=length(pos)>1 ? pl.subplots[3] : nothing,
    )
    juliaPlot(
        c=c, # function offset
        res=pos[1][3], # x resolution
        f=f, # function to iterate
        size=pos[1][3:4],
        filename=nothing,
        color=color,
        labelplatecolor=:white,
        pl=pl.subplots[2],
    )
    heatmap!(pl.subplots[length(pos)+2],f',color=:seaborn_icefire_gradient,clim=(-1,1))

    savefig(pl,"~/Pictures/Julia/background.png")

end # mulitMonitorJuliaPlot

end # module


module VarianceComponentTypeTest

using VarianceComponentModels, MathProgBase, Ipopt, LinearAlgebra, Random, Test

Random.seed!(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)
  copyto!(Σ[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)
  copyto!(V[i], Vi * Vi')
end
copyto!(V[m], Matrix(I, n, n))
# form Ω
Ω = zeros(n*d, n*d)
for i = 1:m
  global Ω += kron(Σ[i], V[i])
end
Ωchol = cholesky(Ω)
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 = eigen(Σ[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 norm(vcmodel.B) ≈ 0.0
@test norm(vcmodel.Σ[1] - Matrix(I, d, d)) ≈ 0.0
@test norm(vcmodel.Σ[2] - Matrix(I, d, 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)
rmul!(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


using Plots, StatsBase, Statistics, LaTeXStrings, JLD, StatsPlots

function FindBorder(dim, Network)
    border0 = []
    for i in 1:dim
        for j in 1:dim
            check = false
            if Network[i,j] == 0
                Neighbors = ((i-1,j),(i,j-1),(i+1,j),(i,j+1))
                for nei in Neighbors
                    if dim+1 ∉ nei && 0 ∉ nei && Network[nei...] == 1
                        check = true
                        break
                    end
                end
                if check == true
                    push!(border0,(i,j))
                end
            end
        end
    end
    return border0
end

function FindCorrelationLength(Network,dim)
    FindOnSpots = findall(x->x==1,Network)
    SpotNums = length(FindOnSpots)
    ΣR² = 0.0; iMC = 0.0; jMC = 0.0
    for indx in FindOnSpots
        iMC += indx[1]/SpotNums
        jMC += indx[2]/SpotNums
    end
    for indx in FindOnSpots
        ΣR² += ((indx[1]-iMC)^2 + (indx[2]-jMC)^2)/SpotNums
    end
    return sqrt(ΣR²)
end

function ClusterGrowth(dim,P)
    TotalNetworks = []
    Network = zeros(Int,dim,dim)
    # Network[rand(1:dim^2)] = 1
    Network[round(Int, dim/2), round(Int, dim/2)] = 1
    Borders = FindBorder(dim,Network)
    while length(Borders) != 0
        push!(TotalNetworks, copy(Network))
        Borders = FindBorder(dim,Network)
        for spot in Borders
            if P > rand()
                Network[spot...] = 1
            else
                Network[spot...] = -1
            end
        end
    end
    return TotalNetworks
end

PList = [0.1,0.3,0.5,0.7,0.9]
for P in PList
    save("../../Data/Q3/Q3-anim-P=$P.jld","Networks",Networks)
    Networks = ClusterGrowth(150,0.6)
    anim = Animation()
    for i in 1:length(Networks)
        plt = heatmap(Networks[i], legend = false, border = :none, title = L"Growing\ Cluster\ for\ L= 150\ , P\ =\ %$P")
        frame(anim, plt)
    end
    gif(anim, "../../Figs/Q3/Q3-anim-P=$P.gif")
end


"""
    ILP(optimizer, max_iter)

ILP iteratively solves a linearized primal optimization to compute maximum allowable disturbance.

# Problem requirement
1. Network: any depth, ReLU activation
2. Input: hyperrectangle
3. Output: halfspace

# Return
`AdversarialResult`

# Method
Iteratively solve a linear encoding of the problem.
Default `optimizer` is `GLPKSolverMIP()`. Default `max_iter` is `10`.

# Property
Sound but not complete.

# Reference
[O. Bastani, Y. Ioannou, L. Lampropoulos, D. Vytiniotis, A. Nori, and A. Criminisi,
"Measuring Neural Net Robustness with Constraints,"
in *Advances in Neural Information Processing Systems*, 2016.](https://arxiv.org/abs/1605.07262)
"""
@with_kw struct ILP{O<:AbstractMathProgSolver}
    optimizer::O    = GLPKSolverMIP()
    max_iter::Int64 = 10
end

function solve(solver::ILP, problem::Problem)
    x = problem.input.center
    i = 0
    while i < solver.max_iter
        model = Model(solver = solver.optimizer)
        act_pattern = get_activation(problem.network, x)

        neurons = init_neurons(model, problem.network)
        add_complementary_set_constraint!(model, problem.output, last(neurons))
        encode_relaxed_lp!(model, problem.network, neurons, act_pattern)
        o = max_disturbance!(model, first(neurons) - problem.input.center)

        status = solve(model, suppress_warnings = true)
        if status != :Optimal
            return AdversarialResult(:Unknown)
        end
        x = getvalue(first(neurons))
        if match_activation(problem.network, x, act_pattern)
            radius = getvalue(o)
            if radius >= minimum(problem.input.radius)
                return AdversarialResult(:SAT, radius)
            else
                return AdversarialResult(:UNSAT, radius)
            end
        end
        i += 1
    end
    return AdversarialResult(:Unknown)
end

function match_activation(nnet::Network, x::Vector{Float64}, act_pattern::Vector{Vector{Bool}})
    curr_value = x
    for (i, layer) in enumerate(nnet.layers)
        curr_value = layer.weights * curr_value + layer.bias
        for (j, val) in enumerate(curr_value)
            act = act_pattern[i][j]
            act  && val < 0.0 && return false
            !act && val > 0.0 && return false
        end
        curr_value = layer.activation(curr_value)
    end
    return true
end


#!/usr/bin/env julia

using Luxor

using Test

using Random
Random.seed!(42)

function test_intersections(fname)
    Drawing(2000, 2000, fname)
    origin()
    sethue("black")
    setline(2)
    setopacity(0.95)

    tiles = Tiler(2000, 2000, 5, 5)

    for (pos, n) in tiles
        gsave()
        translate(pos)
        box(O, tiles.tilewidth, tiles.tileheight, :clip)
        p1 = Point(rand(-tiles.tilewidth/2:tiles.tilewidth/2), rand(-tiles.tileheight/2:tiles.tilewidth/2))
        p2 = Point(rand(-tiles.tilewidth/2:tiles.tilewidth/2), rand(-tiles.tileheight/2:tiles.tilewidth/2))
        p3 = Point(rand(-tiles.tilewidth/2:tiles.tilewidth/2), rand(-tiles.tileheight/2:tiles.tilewidth/2))
        map(pt -> circle(pt, 2, :fill), [p1, p2, p3])

        gsave()
        randomhue()
        rad = rand(20:tiles.tilewidth/3)
        circle(p2, rad, :fill)
        sethue("black")
        line(p1, p3, :stroke)

        sethue("green")
        setdash("dot")
        setline(0.5)
        line(between(p1, p3, -5), between(p1, p3, 5), :stroke)
        grestore()

        sethue("red")
        n, int1, int2 = intersectionlinecircle(p1, p3, p2, rad)
        if n == 1
            circle(int1, 10, :fill)
        elseif n == 2
            circle(int1, 10, :fill)
            circle(int2, 10, :fill)
        end
        sethue("black")
        text("$n intersection(s)")
        box(O, tiles.tilewidth, tiles.tileheight, :stroke)
        grestore()
    end
    l1 = Luxor.Point(-100.0, -100.0)
    l2 = Luxor.Point(300.0, 200.0)
    cpoint = Point(0, 0)
    @test intersectionlinecircle(l1, l2, cpoint, 15)[1] == 0
    @test intersectionlinecircle(l1, l2, cpoint, 20)[1] == 1
    @test intersectionlinecircle(l1, l2, cpoint, 25)[1] == 2
    @test finish() == true
end

test_intersections("intersection_line_circle_test.pdf")
println("...finished test: output in $(fname)")


function genLPM(lammbda::AbstractVector, T::Int, X::AbstractMatrix)
    # generating the Linear Probability Model
    # lammbda is a vector (n x 1)
    # X is the matrix (T x n)

    y = BitArray(undef, T)
    for i = 1:T
        y[i] = dot(X[i,:], lammbda) > randn()
    end

    return y

end

genMC(PIt::Array{Float64,3}, T::Int; Tdrop::Int = 0, S0::Int = 1) = genEndoMC(PIt, T; Tdrop = Tdrop, S0 = S0)

function genEndoMC(PIt::Array{Float64,3}, T::Int; Tdrop::Int = 0, S0::Int = 1)
    S  = [S0]; # fix starting point
    for i = 2:T + Tdrop
        cumu_p = cumsum(PIt[S[i - 1], :,i])
        push!(S, sum(rand() .> cumu_p) + 1)
    end
    return S[Tdrop + 1:end]
end

simMC(PI::Array{Float64,2}, T::Int; Tdrop::Int = 0, S0::Int = 1) = genMC(PI, T, Tdrop = Tdrop, S0 = S0)

function genMC(PI::Array{Float64,2}, T::Int; Tdrop::Int = 0, S0::Int = 1)
    S  = [S0] # fix starting point
    cumPI = cumsum(PI;dims = 2)
    for i = 2:T
        push!(S, sum(rand() .> cumPI[S[i - 1], :]) + 1);
    end
    return S[Tdrop + 1:end]
end


module NumericalMethodsforEngineers

using IterativeSolvers
using Reexport
@reexport using Symata, SparseArrays, LinearAlgebra, Compat

#using BandedMatrices

# package code goes here
### Imports ###
#import Symata: @sym, setsymata, getsymata, symeval

### Includes ###
# Type definitions in nmlib
include("nmlib/SkylineMatrices.jl")
include("nmlib/SymmetricBandedMatrices.jl")
# Direct methods in nmlib
include("nmlib/checkit.jl")
include("nmlib/nmex.jl")
include("nmlib/chobac.jl")
include("nmlib/cholin.jl")
include("nmlib/ldlfor.jl")
include("nmlib/ldlt.jl")
include("nmlib/lufac.jl")
include("nmlib/sparin.jl")
include("nmlib/spabac.jl")
include("nmlib/subbac.jl")
include("nmlib/subfor.jl")
include("nmlib/bicgstab.jl")
# Iterative methods in ch02
#include("ch02", "bicgstabl.jl"))
# Initial Value methods
include("ch05/lagrangianpolynomial.jl")
# Initial Value methods
include("ch07/ivp.jl")
# Boundary Value methods
include("ch07/bvp.jl")
# Symata functions
include("ch07/sym_fcns.jl")

### Exports ###

export
  # nmlib (types)
  SkylineMatrix,
  fromskyline,
  SymmetricBandedMatrix,
  tosymmetricbandedmatrix,
  # nmlib (functions)
  nmex,
  checkit,
  bicgstab,
  chobac!,
  cholin!,
  ldlt!,
  ldltfac!,
  ldlfor!,
  lufac,
  sparin!,
  spabac!,
  subbac!,
  subfor!,
  # ch05
  lagrangianpolynomial,
  # ch07
  euler,
  modified_euler,
  mid_point_euler,
  runga_kutta_4,
  shootingmethod,
  bigcstab,

  # Symata
  @sym,
  setsymata,
  getsymata,
  symeval

### Deprecated ###

  include("deprecated.jl")

end # module


import DomainSets: convert_space, restrict_space, IntegerSpace, RationalSpace, RealSpace, ComplexSpace,
                VectorSpace, AnySpace, EmptySpace, FullSpace, GeometricSpace, promote_space_type, promote_space

# test_spaces.jl

######################################
# Helper functions and generic tests
######################################

nonzero_element(::Type{T}) where {T <: Number} = one(T)
nonzero_element(::Type{SVector{N,T}}) where {N,T<:Number} = ones(SVector{N,T})
nonzero_element(::Type{Tuple{T,S}}) where {T,S} = (nonzero_element(T),nonzero_element(S))
nonzero_element(::Type{Tuple{T,S,U}}) where {T,S,U} = (nonzero_element(T),nonzero_element(S),nonzero_element(U))
nonzero_element(::Type{Tuple{T,S,U,V}}) where {T,S,U,V} = (nonzero_element(T),nonzero_element(S),nonzero_element(U),nonzero_element(V))

# Generic tests for isomorphisms
function test_isomorphism(A, B)
    @test isomorphic(A(), B())
    @test isomorphic(B, A)
    test_embedding(A, B)
    test_embedding(B, A)
    T = eltype(A)
    S = eltype(B)
    x = nonzero_element(T)
    y = nonzero_element(S)
    # Is conversion invertible?
    @test convert_space(A, convert_space(B, x)) == x
    @test convert_space(B, convert_space(A, y)) == y
end

# Generic tests for embeddings
function test_embedding(A, B)
    @test embedded(A(), B())
    T = eltype(A)
    S = eltype(B)
    x = nonzero_element(T)
    y = nonzero_element(S)
    # Does conversion work? We verify if it yields the right type.
    @test typeof(convert_space(B, x)) == S
    # And we test the left inverse
    @test restrict_space(A, convert_space(B, x)) == x
end

##################
# Specific tests
##################

@testset "Spaces" begin
    @testset "basic spaces" begin
        Z = IntegerSpace{Int}
        Q = RationalSpace{Int}
        R = RealSpace{Float64}
        C = ComplexSpace{Float64}
        R1 = VectorSpace{1,Float64}
        R2 = VectorSpace{2,Float64}
        R3 = VectorSpace{3,Float64}
        @test ℤ == Z
        @test ℚ == Q
        @test ℝ == R
        @test ℝ1 == R1
        @test ℝ2 == R2
        @test ℝ3 == R3

        # The basic spaces are not isomorphic to each other
        @test !isomorphic(Z, Q)
        @test !isomorphic(Q, R)
        @test !isomorphic(R1, R2)
        @test !isomorphic(R2, R3)
        test_embedding(Z, Q)
        test_embedding(Z, R)
        test_embedding(Z, C)
        test_embedding(Q, R)
        test_embedding(Q, C)
        test_embedding(R, C)
        test_embedding(R, R1)
        test_embedding(R1, R2)
        test_embedding(R1, R3)

        @test embedded(AnySpace, AnySpace)
        @test embedded(Z, AnySpace)
        @test !embedded(AnySpace, Z)

        # Isomorphism between T and SVector{1,T}
        test_isomorphism(Z, VectorSpace{1,Int})
        test_isomorphism(Q, VectorSpace{1,Rational{Int}})
        test_isomorphism(R, R1)

        # Isomorphism between C and R2
        test_isomorphism(C, R2)

        # Some tests with BigFloat and BigInt
        @test widen(Z) == IntegerSpace{widen(Int)}
        @test widen(Q) == RationalSpace{widen(Int)}
        @test widen(C) == ComplexSpace{widen(Float64)}
        @test widen(R) == RealSpace{widen(Float64)}
        @test widen(R3) == VectorSpace{3,widen(Float64)}
        test_embedding(Z, widen(Z))
        test_embedding(Q, widen(Q))
        test_embedding(C, widen(C))
        test_embedding(R, widen(R))
        test_embedding(R1, widen(R3))

        test_isomorphism(widen(R), widen(R1))
        test_isomorphism(widen(C), widen(R2))

        @testset "Combination of isomorphism and widening" begin
            test_embedding(R, widen(R1))
            test_embedding(R1, widen(R))
            test_embedding(C, widen(R2))
            test_embedding(R2, widen(C))
            @test !embedded(widen(R), R1)
            @test !embedded(widen(C), R2)
            @test !embedded(widen(R1), R2)
            @test !embedded(widen(R1), R3)
        end

        @testset "Some promotions" begin
            @test promote_space() == ()
            @test typeof(promote_space(1.0)) == Tuple{typeof(1.0)}
            @test typeof(promote_space(1, 1.0)) == Tuple{typeof(1.0),typeof(1.0)}
        end

        @test promote_space_type(RealSpace{Float64}, RealSpace{BigFloat}) == RealSpace{BigFloat}
        @test promote_space_type(VectorSpace{1,BigFloat}, VectorSpace{2,Float64}) == VectorSpace{2,BigFloat}
        @test promote_space_type(VectorSpace{1,Float64}, VectorSpace{2,BigFloat}) == VectorSpace{2,BigFloat}

        # Preference for scalars over 1D vectors
        @test promote_space_type(RealSpace{Float64}, VectorSpace{1,Float64}) == RealSpace{Float64}
        @test promote_space_type(VectorSpace{1,Float64}, RealSpace{Float64}) == RealSpace{Float64}
        # - even when one space is bigger
        @test promote_space_type(RealSpace{Float64}, VectorSpace{1,BigFloat}) == RealSpace{BigFloat}
        @test promote_space_type(RealSpace{BigFloat}, VectorSpace{1,Float64}) == RealSpace{BigFloat}

        # Preference for 2D vectors over complex numbers
        @test promote_space_type(ComplexSpace{Float64}, VectorSpace{2,Float64}) == VectorSpace{2,Float64}
        @test promote_space_type(VectorSpace{2,Float64}, ComplexSpace{Float64}) == VectorSpace{2,Float64}
        # - except when one space is bigger
        @test promote_space_type(ComplexSpace{BigFloat}, VectorSpace{2,Float64}) == VectorSpace{2,BigFloat}
        @test promote_space_type(ComplexSpace{Float64}, VectorSpace{2,BigFloat}) == VectorSpace{2,BigFloat}

        @test promote_space_type(AnySpace, VectorSpace{2,Float64}) == AnySpace
        @test promote_space_type(ComplexSpace{Float64}, AnySpace) == AnySpace
        @test promote_space_type(AnySpace, AnySpace) == AnySpace

        # some basic tests

        @test_throws InexactError convert_space(RealSpace{Float64}, RealSpace)
        @test_throws InexactError restrict_space(RealSpace{Float64}, RealSpace)

        @test DomainSets._promote_via_embedding_reduction(AnySpace, AnySpace, AnySpace, ℤ, 1) == AnySpace
        @test DomainSets._promote_via_embedding_reduction(AnySpace, AnySpace, 1, AnySpace, ℤ) == AnySpace

        @test !(zero(Z) ∈ R())
        @test !(zero(Z) ∈ R())
        @test zero(Z) ∈ Z

        @test origin(Z()) == zero(Z)
        @test typeof(origin(Z())) == typeof(zero(Z))

        @test spaceof(1) == ℤ
        @test superspaceof(1) == AnySpace

        @test issubspace(Z, Z)
        @test issubspace(Z, AnySpace)
        @test !issubspace(AnySpace, Z)

        @test DomainSets.result(DomainSets.isomorphism_reduction_result(GeometricSpace{Float64}, GeometricSpace{Float64}, GeometricSpace{Float64}, GeometricSpace{Float64}, GeometricSpace{Float64}, GeometricSpace{Float64}))
        @test DomainSets.result(DomainSets.isomorphism_reduction_result(GeometricSpace{Float64}, GeometricSpace{Float64}, GeometricSpace{SVector{2,Float64}}, GeometricSpace{ComplexF64}, GeometricSpace{ComplexF64}, GeometricSpace{SVector{2,Float64}}))
        @test DomainSets.result(DomainSets.isomorphism_reduction_result(GeometricSpace{Float64}, GeometricSpace{Float64}, GeometricSpace{ComplexF64}, GeometricSpace{SVector{2,Float64}}, GeometricSpace{ComplexF64}, GeometricSpace{SVector{2,Float64}}))
        @test_throws ErrorException  DomainSets.isomorphism_reduction_result(GeometricSpace{Float64}, GeometricSpace{Float64}, GeometricSpace{Float64}, GeometricSpace{Float64}, GeometricSpace{SVector{1,Float64}}, GeometricSpace{Float64})
    end

    @testset "product spaces" begin
        T = Float64
        z = T(0)
        Z = IntegerSpace{Int}
        R = RealSpace{T}
        P = Z × R
        @test eltype(P) == Tuple{Int,T}

        R2 = R × R
        @test eltype(R2) == Tuple{T,T}
        @test zero(R2) == (z,z)
        test_isomorphism(R2, VectorSpace{2,T})

        R3 = cartesianproduct(R, R, R)
        @test eltype(R3) == Tuple{T,T,T}
        @test zero(R3) == (z,z,z)
        test_isomorphism(R3, VectorSpace{3,T})

        R4 = cartesianproduct(R, R, R, R)
        @test eltype(R4) == Tuple{T,T,T,T}
        @test zero(R4) == (z,z,z,z)
        test_isomorphism(R4, VectorSpace{4,T})

        RR = R2 × R
        @test eltype(RR) == Tuple{Tuple{T,T},T}
        test_isomorphism(RR, VectorSpace{3,T})

        @test R == cartesianproduct(R)

        @test cartesianproduct() == nothing
        @test eltype(cross(R())) == T
        @test eltype(cross(R(),R())) == NTuple{2,T}
        @test eltype(cross(R(),R(),R())) == NTuple{3,T}
        @test eltype(cross(R(),R(),R(),R())) == NTuple{4,T}
    end
end


# Grids

"""
	$(SIGNATURES)

Abstract grid type.
"""
abstract type AbstractGrid end

"""
	$(SIGNATURES)

Linear, equally spaced grid.
"""
struct LinearGrid{T <: AbstractFloat, I <: Integer} <: AbstractGrid
    xMin :: T
    xMax :: T
    nPoints :: I
end

"""
	$(SIGNATURES)

Grid with increasing or decreasing intervals. Interval(j) / interval(j-1) is a constant.
Constructed from:
- xMin, xMax: range
- nPoints: no of grid points
- ratio of interval(j) / interval(j-1)
"""
struct PowerSpacedGrid{T <: AbstractFloat, I <: Integer} <: AbstractGrid
    xMin :: T
    xMax :: T
    nPoints :: I
    intervalRatio :: T
end

"""
	$(SIGNATURES)

Grid with intervals that increase or decrease by constants.
"""
struct LinSpacedGrid{T <: AbstractFloat, I <: Integer} <: AbstractGrid
    xMin :: T
    xMax :: T
    nPoints :: I
    intervalDiff :: T
end

## -------------  Common functions

Base.eltype(g :: AbstractGrid) = Base.eltype(minimum(g));
Base.length(g :: AbstractGrid) = g.nPoints;
Base.minimum(g :: AbstractGrid) = g.xMin;
Base.maximum(g :: AbstractGrid) = g.xMax;
n_intervals(g :: AbstractGrid) = length(g) - 1;
intervals(g :: AbstractGrid) = diff(grid(g));
Base.firstindex(g :: AbstractGrid) = 1;
Base.lastindex(g :: AbstractGrid) = length(g);
Base.getindex(g :: AbstractGrid, idx) = grid(g)[idx];

function grid(g :: AbstractGrid)
    return minimum(g) .+ vcat(zero(eltype(g)), cumsum(intervals(g)))
end

function Base.iterate(g :: AbstractGrid, idx)
    if idx > length(g)
        return nothing
    else
        return g[idx], idx+1
    end
end

Base.iterate(g :: AbstractGrid) = Base.iterate(g, 1);

## --------------  Linear Grid

grid(g :: LinearGrid) = range(minimum(g), maximum(g), length = length(g));

function Base.getindex(g :: LinearGrid, idx)
    dx = (maximum(g) - minimum(g)) / n_intervals(g);
    return minimum(g) .+ dx .* (collect(idx) .- 1)
end

## --------------  Power Spaced Grid

function first_interval(g :: PowerSpacedGrid{T,I}) where {T, I}
    factor1 = (g.intervalRatio ^ (length(g) - 1) - one(T)) / (g.intervalRatio - 1);
    return (maximum(g) - minimum(g)) / factor1
end

function intervals(g :: PowerSpacedGrid{T,I}) where {T, I}
    intV = Vector{T}(undef, length(g) - 1);
    intV[1] = first_interval(g);
    for j = 2 : (length(g) - 1)
        intV[j] = intV[j-1] * g.intervalRatio;
    end
    return intV
end

# function grid(g :: PowerSpacedGrid{T,I}) where {T,I}
#     return minimum(g) .+ vcat(zero(eltype(g)), cumsum(intervals(g)))
# end

## ----------------  Linearly changing intervals

function first_interval(g :: LinSpacedGrid{T,I}) where {T, I}
    n = length(g);
    dX = maximum(g) - minimum(g);

    return dX / (n-1) - g.intervalDiff * (n-2) / 2.0;
end

interval(g :: LinSpacedGrid{T,I}, idx) where {T, I} =
    first_interval(g) .+ (idx .- 1) .* g.intervalDiff;

intervals(g :: LinSpacedGrid{T,I}) where {T, I} =
    interval(g, 1 : n_intervals(g))

# -------------


#=  Here I compare a wide variety of tiebreaking rules, some of which (STB, MTB, and
    HTB) are described by Ashlagi and Nikzad (2020), and others (XTB, WTB) of my own
    creation. Testing the tiebreaking rules over a hybrid market in which the students
    uniformly prefer a subset of the schools allows us to examine the properties of
    these mechanisms in over- and underdemanded markets at a glance. See the readme
    for further discussion.                     =#

using Permutations
using Plots
using DeferredAcceptance

"""
Plots outcomes for various tiebreaking rules for a market with
the given characteristics.
"""
function plotter_more(n, m_pop, m_unp, samp)

    m = m_pop + m_unp

    descr = "$n students, $m schools ($m_pop popular, $m_unp unpopular), $samp samples"
    println(descr)

    # All schools have unit capacities
    capacities = ones(Int64, m)

    # Blank rank dists
    cdf_STB = zeros(Float64, m)
    cdf_MTB = zeros(Float64, m)
    cdf_HTB = zeros(Float64, m)
    cdf_XTB = zeros(Float64, m)
    cdf_WHTB = zeros(Float64, m)
    cdf_WXTB = zeros(Float64, m)

    # For HTB and WTB
    blend = ones(Float64, 1, m) # Use MTB in all schools except
    blend[1:m_pop] .= 0            # the popular schools, which use STB

    for i in 1:samp
        # Students unilaterally prefer popular schools to unpopular ones
        # Random ranking otherwise
        students = vcat(hcat((randperm(m_pop) for i in 1:n)...),
                   m_pop .+ hcat((randperm(m_unp) for i in 1:n)...))

        # Schools place all students in single priority category
        schools = ones(Int64, n, m)

        # Break ties
        schools_STB = singletiebreaking(schools)
        schools_MTB = multipletiebreaking(schools)
        schools_HTB = hybridtiebreaking(schools, blend)
        schools_XTB = hybridtiebreaking(schools, 0.5)
        schools_WHTB = welfaretiebreaking(students, schools, blend)
        schools_WXTB = welfaretiebreaking(students, schools, 0.5)

        # Update rank risks
        # println("Starting STB $i")
        cdf_STB += DA_rank_dist(students, schools_STB, capacities)
        # println("Starting MTB $i")
        cdf_MTB += DA_rank_dist(students, schools_MTB, capacities)
        # println("Starting HTB $i")
        cdf_HTB += DA_rank_dist(students, schools_HTB, capacities)
        # println("Starting XTB $i")
        cdf_XTB += DA_rank_dist(students, schools_XTB, capacities)
        # println("Starting WHTB $i")
        cdf_WHTB += DA_rank_dist(students, schools_WHTB, capacities)
        # println("Starting WXTB $i")
        cdf_WXTB += DA_rank_dist(students, schools_WXTB, capacities)

        # To display plot as it updates
        # display(plot([cdf_STB, cdf_MTB, cdf_HTB, cdf_XTB, cdf_WHTB, cdf_WXTB],
        #         label = ["DA-STB" "DA-MTB" "DA-HTB" "DA-XTB, λ=0.5" "DA-WHTB" "DA-WXTB, λ=0.5"],
        #         lc = [:dimgray :dimgray :dimgray :crimson :dodgerblue :olivedrab],
        #         ls = [:dot :dash :dashdot :dot :dash :dashdot],
        #         legend = :bottomright,
        #         title = descr, titlefontsize=11,
        #         xlabel = "rank", ylabel= "average number of students"))

    end

    # Norm rank dists against sample size
    for i in (cdf_STB, cdf_MTB, cdf_HTB, cdf_XTB, cdf_WHTB, cdf_WXTB)
        i ./= samp
    end

    return plot([cdf_STB, cdf_MTB, cdf_HTB, cdf_XTB, cdf_WHTB, cdf_WXTB],
                label = ["DA-STB" "DA-MTB" "DA-HTB" "DA-XTB, λ=0.5" "DA-WHTB" "DA-WXTB, λ=0.5"],
                lc = [:dimgray :dimgray :dimgray :crimson :dodgerblue :olivedrab],
                ls = [:dot :dash :dashdot :dot :dash :dashdot],
                legend = :bottomright,
                title = descr, titlefontsize=11,
                xlabel = "rank", ylabel= "average number of students")
end

(n, m_pop, m_unp, samp) = (50, 17, 33, 200)

p = plotter_more(n, m_pop, m_unp, samp)

# savefig(p, string("examples/hybrid/plot.pdf"))
# savefig(p, string("examples/hybrid/plot.png"))


function ql(A::Matrix)
    Q,R=qr(A[end:-1:1,end:-1:1])
    Q[end:-1:1,end:-1:1],R[end:-1:1,end:-1:1]
end

joukowsky(z) = 0.5*(z+1/z)

function jacobimatrix(a,b,t0,t1,N)
    J = BandedMatrix(Float64,N,N,1,1)

    for i = 1:min(length(a),N)
        J[i,i] = a[i]
    end
    for i=length(a)+1:N
        J[i,i] = t0
    end

    for i = 1:min(length(b),N-1)
        J[i,i+1] = J[i+1,i] = b[i]
    end

    for i=length(b)+1:N-1
        J[i,i+1] = J[i+1,i] = t1
    end

    J
end

jacobimatrix(a,b,N) = jacobimatrix(a,b,0,.5,N)

function jacobioperator(a,b,t0,t1)
    n = max(length(a),length(b)+1)
    a = [a;zeros(n-length(a))]; b = [b;.5+zeros(n-length(b))]
    SymTriPertToeplitz(ToeplitzOperator([t1],[t0,t1]),SymTriOperator(a-t0,b-t1))
end

jacobioperator(a,b) = jacobioperator(a,b,0,.5)

freejacobioperator() = SymTriPertToeplitz([0.],[.5],0.,.5)

## tridiagonal ql

function tridql!(L::Matrix)
    n=size(L,1)

  # Now we do QL for the compact part in the top left
    cc=Array{eltype(L)}(n)
    ss=Array{eltype(L)}(n-1)

    for i = n:-1:2
        nrm=sqrt(L[i-1,i]^2+L[i,i]^2)
        c,s = L[i,i]/nrm, -L[i-1,i]/nrm
        if i > 2
            L[i-1:i,i-2:i] = [c s; -s c]*L[i-1:i,i-2:i]
            L[i-1,i]=0
        else
            L[i-1:i,i-1:i] = [c s; -s c]*L[i-1:i,i-1:i]
            L[i-1,i]=0
        end
        cc[i]=c
        ss[i-1]=s
    end
    cc[1]=sign(L[1,1])
    L[1,1]=abs(L[1,1])
    cc,ss,L
end

function tridql!(J::BandedMatrix)
    n=size(J,1)
    L=BandedMatrix(copy(J.data),J.m,2,0)

  # Now we do QL for the compact part in the top left
    cc=Array{eltype(J)}(n)
    ss=Array{eltype(J)}(n-1)

    for i = n:-1:2
        nrm=sqrt(J[i-1,i]^2+J[i,i]^2)
        c,s = J[i,i]/nrm, -J[i-1,i]/nrm

        for j=max(i-2,1):i
            L[i,j]=-s*J[i-1,j]+c*J[i,j]
            J[i-1,j]=c*J[i-1,j]+s*J[i,j]
        end
        cc[i]=c
        ss[i-1]=s
    end
    cc[1]=sign(J[1,1])
    L[1,1]=abs(J[1,1])
    cc,ss,L
end

#Finds NxN truncation of C such that C'(Q_k(s)) =  (P_k(s)),
# where P_k has Jacobi coeffs a,b and Q_k has Jacobi coeffs c,d
function connectioncoeffsmatrix(a::AbstractVector, b::AbstractVector, c::AbstractVector, d::AbstractVector, N)
  if N>max(length(a),length(b)+1,length(c),length(d)+1)
    a = [a;zeros(N-length(a))]; b = [b;.5+zeros(N-length(b))]
    c = [c;zeros(N-length(c))]; d = [d;.5+zeros(N-length(d))]
  end

  C = zeros(eltype(a),N,N)
  C[1,1] = 1
  C[1,2] = (c[1]-a[1])/b[1]
  C[2,2] = d[1]/b[1]
  for j = 3:N
    C[1,j] = ((c[1]-a[j-1])*C[1,j-1] + d[1]*C[2,j-1] - b[j-2]*C[1,j-2])/b[j-1]
    for i = 2:j-1
      C[i,j] = (d[i-1]*C[i-1,j-1] + (c[i]-a[j-1])*C[i,j-1] + d[i]*C[i+1,j-1] - b[j-2]*C[i,j-2])/b[j-1]
    end
    C[j,j] = d[j-1]*C[j-1,j-1]/b[j-1]
  end
  C
end

#Makes the matrix C which transforms the coefficients of an expansion in "right" matrix orthonormal polynomial with recurrence coeffs a,b (J) to those of an expansion in the "right" MOPs with recurrence coeffs c,d (D). So C is block-upper-triangular and CJ = DC (modulo the final column). The types of a,b,c,d must be BlockArrays which are like column vectors of length N whose entries are kxk blocks. The blocks of a and c should be symmetric.
function connectioncoeffsmatrix(a::AbstractBlockArray, b::AbstractBlockArray, c::AbstractBlockArray, d::AbstractBlockArray)
  k = blocksize(a,1,1)[1]
  N = nblocks(a,1)
  C = BlockArray(zeros(k*N,k*N),k*ones(Int64,N),k*ones(Int64,N))
  C[Block(1,1)] = eye(k,k)
  C[Block(1,2)] = (c[Block(1,1)]-a[Block(1,1)])*pinv(b[Block(1,1)]')
  C[Block(2,2)] = d[Block(1,1)]'/b[Block(1,1)]'
  for j = 3:N
    C[Block(1,j)] = (c[Block(1,1)]*C[Block(1,j-1)]-C[Block(1,j-1)]*a[Block(j-1,1)] + d[Block(1,1)]*C[Block(2,j-1)]-C[Block(1,j-2)]*b[Block(j-2,1)])/b[Block(j-1,1)]'
    for i = 2:j-1
        C[Block(i,j)] = (d[Block(i-1,1)]'*C[Block(i-1,j-1)] + c[Block(i,1)]*C[Block(i,j-1)]-C[Block(i,j-1)]*a[Block(j-1,1)] +d[Block(i,1)]*C[Block(i+1,j-1)] - C[Block(i,j-2)]*b[Block(j-2,1)])/b[Block(j-1,1)]'
    end
    C[Block(j,j)] = (d[Block(j-1,1)]'*C[Block(j-1,j-1)])/b[Block(j-1,1)]'
  end
  C
end

# This is for Chebyshev U
connectioncoeffsmatrix(a,b,N) = connectioncoeffsmatrix(a,b,[],[],N)

# Converts coefficients a^J to coefficients a^D using Clenshaw
function applyconversion(J::SymTriPertToeplitz,D::SymTriPertToeplitz,v::Vector)
  N = length(v)
  T = eltype(b)
  b = zeros(T,N); b1 = zeros(T,N); b2 = zeros(T,N)
  for k = N:-1:1
    # before: b = b_k+1, b1 = b_k+2, (and b2 = b_k+3 is to be forgotten)
    b2 = pad((D-J[k,k]*I)*b,N)/J[k,k+1]-b1*(J[k,k+1]/J[k+1,k+2])
    b2[1] += v[k]
    b2, b1, b = b1, b, b2
    # after: b = b_k, b1 = b_k+1, b2 = b_k+2
  end
  b
end


struct Normal{T}
    μ::T
    σ::T
    function Normal(m,s)
        μ,σ = promote(m,s)
        T = typeof(μ)
        new{T}(μ, σ)
    end
end
Normal(;μ,σ) = Normal(μ,σ)
Base.show(io::IO, o::Normal) = print(io, "Normal(μ=$(o.μ), σ=$(o.σ))")

struct NormalRegressor end

function (o::NormalRegressor)(arr::AbstractMatrix)
    @argcheck mod(size(arr,1), 2) == 0
    nc, nb = size(arr)
    b = Int(nc/2)
    cmu    = view(arr, 1:b)
    csigma = view(arr, b+1:2b)
    return Normal_from_2channels.(cmu, csigma)
end
function Normal_from_2channels(cmu, csigma)
    return Normal(cmu, softplus(csigma))
end

function Distributions.logpdf(o::Normal, x)
    μ = o.μ
    σ = o.σ
    T = typeof(μ)
    half = T(1) / T(2)
    B = -1/(2*σ^2) * (μ - x)^2
    A = -half * log2π - log(σ)
    return A + B
end

sampletype(::Type{Normal{T}}) where {T} = T
Distributions.mean(o::Normal) = o.μ
Distributions.std(o::Normal)  = o.σ
Distributions.Normal(o::Normal) = Distributions.Normal(o.μ, o.σ)


using BenchmarkTools
using Plots
using Base.Threads

f(x, p) = p*x*(one(x)-x)

function calc_attractor!(out, f, p, num_attract, warmup)

    x = 0.25
    for i in 1:warmup
        x = f(x, p)
    end
    @inbounds out[1] = x

    @inbounds for i in 1:num_attract-1
        out[i+1] = f(out[i], p)
    end
end

x₀ = 0.25
out = Vector{typeof(x₀)}(undef, 150)
out[1] = x₀
@btime calc_attractor!(out, f, 2.9, 150, 400)

@code_llvm calc_attractor!(out, f, 2.9, 150, 400)
@code_warntype calc_attractor!(out, f, 2.9, 150, 400)

function bifurcation_data!(data, f, r, n_cols)

    Threads.@threads for i in 1:n_cols
        calc_attractor!(@view(data[:, i]), f, r[i], 150, 400)
    end
end

r = 2.9:0.001:4
n_cols = length(r)
n_rows = 150
pl_data = Array{Float64}(undef, n_rows, n_cols)
pl_data[1, :] .= 0.25
@btime bifurcation_data!(pl_data, f, r, n_cols)

p = plot(xlims=(2.9,4), ylims=(0, 1), title="Bifurcation Plot", label=false)
for i in 1:n_rows
    scatter!(r, pl_data[i, :], label=false, color=:black, alpha=0.1)
end
p


#
# The Computer Language Benchmarks Game
# binary-trees benchmark
# http://shootout.alioth.debian.org/u32/performance.php?test=binarytrees
#
# Ported from an OCaml version
#

abstract BTree

type Empty <: BTree
end

type Node <: BTree
    info
    left::BTree
    right::BTree
end

function make(val, d)
    if d == 0
        Node(val, Empty(), Empty())
    else
        nval = val * 2
        Node(val, make(nval-1, d-1), make(nval, d-1))
    end
end

check(t::Empty) = 0
check(t::Node) = t.info + check(t.left) - check(t.right)

function loop_depths(d, min_depth, max_depth)
    for i = 0:div(max_depth - d, 2)
        niter = 1 << (max_depth - d + min_depth)
        c = 0
        for j = 1:niter
            c += check(make(i, d)) + check(make(-i, d))
        end
#        @printf("%i\t trees of depth %i\t check: %i\n", 2*niter, d, c)
        d += 2
    end
end

function binary_trees(N::Int=10)
    const min_depth = 4
    const max_depth = N
    const stretch_depth = max_depth + 1

    # create and check stretch tree
    let c = check(make(0, stretch_depth))
#        @printf("stretch tree of depth %i\t check: %i\n", stretch_depth, c)
    end

    long_lived_tree = make(0, max_depth)

    loop_depths(min_depth, min_depth, max_depth)
#    @printf("long lived tree of depth %i\t check: %i\n", max_depth, check(long_lived_tree))
end


using Test
using SimpleTools

sq = Dict{Int,Int}()
for k=1:10
    sq[k] = k*k
end

rt = Dict{Int,Int}()
for k=1:10
    rt[k*k] = k
end

f = rt*sq
for k=1:10
    @test f[k] == k
end

y1 = [x^2 for x=0:.01:3]
y2 = [mod(y,1) for y in y1 ]
make_continuous!(y2,1)
@test y1==y2

@test flush_print(23,5) == "   23"
@test flush_print(23,5,false) == "23   "

A = dcat(1,1,1)
M = [1 0 0 ; 0 1 0 ; 0 0 1]
@test A==M

@test sage(A)==nothing

@test sum(eye(5)) == 5
@test cofactor_det(2eye(Int,3))==8

using Polynomials
x = Poly([0,1])
A = [1 1; 3 2]
@test char_poly(A) == (x-1)*(x-2)-3


function smoothing_factor(t_e::R, cutoff::R) where R <: Real
    r = 2 * π * cutoff * t_e
    r / (r + 1)
end

function exponential_smoothing(a::R, x::R, x_prev::R) where R <: Real
    a * x + (1 - a) * x_prev
end

function one_euro_filter(t, x, t_prev, x_prev, dx_prev, min_cutoff, beta, d_cutoff)
    # Time difference
    t_e = t - t_prev

    # The filtered derivative of the signal.
    a_d = smoothing_factor(t_e, d_cutoff)
    dx = @. (x - x_prev) / t_e
    dx_hat = exponential_smoothing.(a_d, dx, dx_prev)

    # The filtered signal.
    cutoff = @. min_cutoff + beta * abs(dx_hat)
    a = smoothing_factor.(t_e, cutoff)
    x_hat = exponential_smoothing.(a, x, x_prev)

    return t, x_hat, dx_hat
end

function one_euro_filter(t::R, x::R, t_prev::R, x_prev::R, dx_prev::R; min_cutoff::R = R(1), beta::R = R(0), d_cutoff::R = R(1)) where R <: Real
    one_euro_filter(t, x, t_prev, x_prev, dx_prev, min_cutoff, beta, d_cutoff)
end

function one_euro_filter(t::R, x::R, t_prev::R, x_prev::R; min_cutoff::R = R(1), beta::R = R(0), d_cutoff::R = R(1)) where R <: Real
    one_euro_filter(t, x, t_prev, x_prev, R(0), min_cutoff, beta, d_cutoff)
end

function one_euro_filter(t::R, x::A, t_prev::R, x_prev::A, dx_prev::A; min_cutoff::R = R(1), beta::R = R(0), d_cutoff::R = R(1)) where A <: AbstractVector{R} where R <: Real
    one_euro_filter(t, x, t_prev, x_prev, dx_prev, min_cutoff, beta, d_cutoff)
end

function one_euro_filter(t::R, x::A, t_prev::R, x_prev::A; min_cutoff::R = R(1), beta::R = R(0), d_cutoff::R = R(1)) where A <: AbstractVector{R} where R <: Real
    one_euro_filter(t, x, t_prev, x_prev, zeros(R, length(x)), min_cutoff, beta, d_cutoff)
end


# Drive train example
# ============================

using HybridSystems, MathematicalSystems, LazySets, Plots
import SX:is_halfspace, is_hyperplane, convert, parse_sxmath
import LazySets.HalfSpace

expr_p = parse_sxmath("0.2<=x & x<=0.3 & -0.1<=y & y<=0.1 & z==0 & x1==0 & x2==0 & x3==0")
#expr = [x,y,x1,x2,x3,z]
x0sets = []
vars = Basic[:x, :y, :x1, :x2, :x3, :z]
for expr_i in expr_p
    if is_halfspace(expr_i)
        push!(x0sets, convert(HalfSpace, expr_i, vars=vars))
    elseif is_hyperplane(expr_i)
        push!(x0sets, convert(Hyperplane, expr_i, vars=vars))
    end
end

AFFINE_SYSTEM = ConstrainedLinearControlContinuousSystem
HS = readsxmodel("models/SpaceEx/FilteredOscillator/filtered_oscillator_flattened.xml", ST=AFFINE_SYSTEM)

# calculate reachable states up to time T
prob = InitialValueProblem(HS, x0sets)
input_options = Options(:mode=>"reach");

problem_options = Options(:vars=>1:system_dimension, :T=>5.0, :δ=>0.05,
                          :plot_vars=>[1, 2], :max_jumps=>1, :verbosity=>1, :init_loc=>3);
options_input = merge(problem_options, input_options)
using Polyhedra
sol = solve(prob, options_input);

plot(sol, indices=1:2:length(sol.Xk))


using PyPlot

	# 2D approximation example
function plot2D()

	ndims = 2

	doplot = false

	# bounds
	lb = [-2.1,-2.5]
	ub = [2.2,2.6]

	# number of eval points and basis functions:
	# we require square basis matrices!
	npoints = [11,12]

	# number of basis funcs
	nbasis = npoints

	# splien degrees
	degs = [3,3]

	# implies a number of knots for each spline
	# remember the restriction that nknots == ncoefs
	nknots = Dict(i => nbasis[i] - degs[i] + 1 for i=1:ndims)

	# eval points
	points = Dict(i => collect(linspace(lb[i],ub[i],npoints[i])) for i=1:ndims)

	# set up BSplines
	bsp = Dict(i => BSpline(nknots[i],degs[i],lb[i],ub[i]) for i=1:ndims)

	# set of basis functions
	d = Dict{Integer,Array{Float64,2}}()
	for i=1:ndims
		d[i] = full(getBasis(points[i],bsp[i]))
	end

	# set of INVERSE basis functions
	id = Dict{Integer,Array{Float64,2}}()
	for k in collect(keys(d))
		id[k] = inv(d[k])
	end

	#  get a function
	function f(x,y)
		sin(sqrt(x^2+y^2))
	end

	# compute function values so that
	# k is fastest varying index
	y = Float64[f(i,j) for i in points[1], j in points[2]]

	yvec = y[:]

	# get coefs using the function
	mycoef = getTensorCoef(id,yvec)

	 # testing tolerance
	 tol = 1e-2

	 # make a plot
	 # ===========

	new_npoints = [17,16]
	# new_npoints = npoints
	new_points = Dict(i => linspace(lb[i],ub[i],new_npoints[i]) for i=1:ndims)
	# set of new basis functions
	nd = Dict{Integer,Array{Float64,2}}()
	for i=1:ndims
		nd[i] = full(getBasis(new_points[i],bsp[i]))
	end

	pred = BSplines.evalTensor2(nd[2],nd[1],mycoef)

	BSplines.figure()
	BSplines.subplot(121,projection="3d")
	BSplines.mesh(points[2],points[1],y)
	BSplines.title("true values on grid")

	BSplines.subplot(122,projection="3d")
	pp = reshape(pred,new_npoints[1],new_npoints[2])
	BSplines.mesh(new_points[2],new_points[1],pp)
	BSplines.title("Approximation off grid")

	BSplines.suptitle("2D TEST")
end

function plot3D()
	# 3D approximation example

	ndims = 3
	doplot = false

	# bounds
	lb = [-2.1,-2.5,0.9]
	ub = [2.2,2.6,5.0]

	# number of eval points and basis functions:
	# we require square basis matrices!
	npoints = [12,16,13]

	# number of basis funcs
	nbasis = npoints

	# splien degrees
	degs = [3,3,3]

	# implies a number of knots for each spline
	# remember the restriction that nknots == ncoefs
	nknots = Dict(i => nbasis[i] - degs[i] + 1 for i=1:ndims)

	# eval points
	points = Dict(i => linspace(lb[i],ub[i],npoints[i]) for i=1:ndims)

	# set up BSplines
	bsp = Dict(i => BSpline(nknots[i],degs[i],lb[i],ub[i]) for i=1:ndims)

	# set of basis functions
	d = Dict{Integer,Array{Float64,2}}()
	for i=1:ndims
		d[i] = full(getBasis(points[i],bsp[i]))
	end

	# set of INVERSE basis functions
	id = Dict{Integer,Array{Float64,2}}()
	for k in collect(keys(d))
		id[k] = inv(d[k])
	end

	#  get a function
	function f(x,y,z)
		sin(sqrt(x^2+y^2)) - z^2
	end

	# compute function values so that
	# k is fastest varying index
	y = Float64[f(i,j,k) for i in points[1], j in points[2], k in points[3]]

	yvec = y[:]

	# get coefs using the function
	mycoef = getTensorCoef(id,yvec)

		 # make a plot
		 # ===========

		new_npoints = [17,16,5]
		# new_npoints = npoints
		new_points = Dict(i => linspace(lb[i],ub[i],new_npoints[i]) for i=1:ndims)
		# set of new basis functions
		nd = Dict{Integer,Array{Float64,2}}()
		for i=1:ndims
			nd[i] = full(getBasis(new_points[i],bsp[i]))
		end

		pred = BSplines.evalTensor3(nd[3],nd[2],nd[1],mycoef)

		BSplines.figure()
		BSplines.subplot(221,projection="3d")
		BSplines.mesh(points[2],points[1],y[:,:,1])
		BSplines.title("true values at lowest 3D state")

		BSplines.subplot(222,projection="3d")
		pp = reshape(pred,new_npoints[1],new_npoints[2],new_npoints[3])
		BSplines.mesh(new_points[2],new_points[1],pp[:,:,1])
		BSplines.title("Approximation")

		BSplines.subplot(223,projection="3d")
		BSplines.mesh(points[2],points[1],y[:,:,npoints[3]])
		BSplines.title("true values at highest 3D state")

		BSplines.subplot(224,projection="3d")
		pp = reshape(pred,new_npoints[1],new_npoints[2],new_npoints[3])
		BSplines.mesh(new_points[2],new_points[1],pp[:,:,new_npoints[3]])
		BSplines.title("Approximation")
		BSplines.suptitle("3D TEST")

end

function plot4D()
	ndims = 4

	# bounds
	lb = [-1.1,-1.5,-0.9,-1.0]
	ub = [1.2,1.6,0.9,1]

	# number of eval points and basis functions:
	# we require square basis matrices!
	npoints = [13,13,13,13]

	# number of basis funcs
	nbasis = npoints

	# splien degrees
	degs = [3,3,3,3]

	# implies a number of knots for each spline
	# remember the restriction that nknots == ncoefs
	nknots = Dict(i => nbasis[i] - degs[i] + 1 for i=1:ndims)

	# eval points
	points = Dict(i => linspace(lb[i],ub[i],npoints[i]) for i=1:ndims)

	# set up BSplines
	bsp = Dict(i => BSpline(nknots[i],degs[i],lb[i],ub[i]) for i=1:ndims)

	# set of basis functions
	d = Dict{Integer,Array{Float64,2}}()
	for i=1:ndims
		d[i] = full(getBasis(points[i],bsp[i]))
	end

	# set of INVERSE basis functions
	id = Dict{Integer,Array{Float64,2}}()
	for k in collect(keys(d))
		id[k] = inv(d[k])
	end

	#  get a function
	function f(x,y,z,w)
		sin(sqrt(x^2+y^2)) + (z-w)^3
	end

	# compute function values so that
	# k is fastest varying index
	y = Float64[f(i,j,k,w) for i in points[1], j in points[2], k in points[3], w in points[4]]

	yvec = y[:]

	# get coefs using the function
	mycoef = getTensorCoef(id,yvec)

	 # testing tolerance
	 tol = 5e-3

		 # make a plot
		 # ===========

		new_npoints = [17,16,5,7]
		# new_npoints = npoints
		new_points = Dict(i => linspace(lb[i],ub[i],new_npoints[i]) for i=1:ndims)
		# set of new basis functions
		nd = Dict{Integer,Array{Float64,2}}()
		for i=1:ndims
			nd[i] = full(getBasis(new_points[i],bsp[i]))
		end

		# predict everywhere
		pred = BSplines.evalTensor4(nd[4],nd[3],nd[2],nd[1],mycoef)

		BSplines.figure()
		BSplines.subplot(221,projection="3d")
		BSplines.mesh(points[2],points[1],y[:,:,1,1])
		BSplines.title("true values at lowest state (:,:,1,1)")

		BSplines.subplot(222,projection="3d")
		pp = reshape(pred,new_npoints[1],new_npoints[2],new_npoints[3],new_npoints[4])
		BSplines.mesh(new_points[2],new_points[1],pp[:,:,1,1])
		BSplines.title("Approximation")

		BSplines.subplot(223,projection="3d")
		BSplines.mesh(points[2],points[1],y[:,:,npoints[3],npoints[4]])
		BSplines.title("true values at state (:,:,end,end)")

		BSplines.subplot(224,projection="3d")
		pp = reshape(pred,new_npoints[1],new_npoints[2],new_npoints[3],new_npoints[4])
		BSplines.mesh(new_points[2],new_points[1],pp[:,:,new_npoints[3],new_npoints[4]])
		BSplines.title("Approximation")
		BSplines.suptitle("4D TEST")
end


@testset "Simplex" begin
    s = Simplex(1, 2, 3)
    @test s == Simplex(1, 2, 3)
    @test sort!(indices(s)) == [1,2,3]
    @test eltype(s) == Int
    @test eltype(Simplex{1,Float64}) == Float64

    s = Simplex([1,2,3])
    @test dim(s) == 2
    @test sort!(indices(s)) == [1,2,3]
    allfaces = [Simplex(2, 3), Simplex(1, 3), Simplex(1, 2)]
    @testset "Simplex faces" for o in faces(s)
        idx = findfirst(isequal(o), allfaces)
        @test indices(o) == indices(allfaces[idx])
    end

    ss = Set{Simplex}()
    push!(ss, s)
    push!(ss, s)
    @test length(ss) == 1

    X = hcat(zeros(3), Matrix(I,3,3))
    S = Simplex(1,2,3,4)
    @test volume(S,X) == 1/6

    s = Simplex(1,3)
    fs = collect(faces(s))
    @test dim(s) == 1
    @test all(f -> f ∈ [Simplex(1),Simplex(3)], fs)

    @test fs[2] ∪ fs[1]  == s

end

@testset "Cube" begin
    c = Cube([0,0,0]) # 0-cube
    @test dim(c) == 0
    c = Cube([0,0,0],[0,0,1]) # 1-cube
    @test dim(c) == 1
    c = Cube([1,0,1]) # 2-cube
    @test dim(c) == 2
    c = Cube([1,1,1]) # 3-cube
    @test dim(c) == 3
    @test hash(c) == hash([0,0,0,1,1,1])
    @test c == Cube([0,0,0],[1,1,1])

    c = Cube([.5,0,.5]) # 2-cube
    @testset "Cube faces" for f in faces(c)
        @test hash(f) ∈ map(hash, [Cube([0,0,.5]), Cube([.5,0,0])])
    end
    boundary(Float64, c)

    c = Cube([.5,.0,.5]) # 2-cube
    @test volume(c) == 0.5^2
    c = Cube([.5,.5,.5]) # 3-cube
    @test volume(c) == 0.5^3
end

@testset "CW" begin
    ComputationalHomology.resetCellId!()
    a = Cell()
    b = Cell()
    @test dim(a) == 0
    @test hash(a) == 0x02011ce34bce797f
    @test a == a
    @test eltype(a) == Int
    @test eltype(Cell{Float64}) == Float64

    @test_throws AssertionError Cell(0, a)

    c = Cell(1, a, b)
    @test dim(c) == 1
    @test vertices(c)[1] == b
    @test vertices(c)[2] == a
    @test c.boundaryMap[hash(a)] == 1
    @test c.boundaryMap[hash(b)] == -1

    d = Cell(Int, 1)
    push!(d, a)
    push!(d, b)
    @test_throws AssertionError push!(d, c)
    @test dim(d) == 1
    @test vertices(d)[1] == b
    @test vertices(d)[2] == a
    @test d.boundaryMap[hash(a)] == 1
    @test d.boundaryMap[hash(b)] == -1

    c2 = Cell()
    d = Cell(1, a, b, c2)
    e = c ∪ d
    @test dim(e) == 2

    @testset for (f1, f2) in zip(vertices(c), [b, a])
        @test f1 == f2
    end

    @testset for (f1, f2) in zip(faces(e), [c, d])
        @test f1 == f2
    end

    @test vertices(e)[1] == b
    @test vertices(e)[2] == c2
    @test vertices(e)[3] == a

    f = Cell(1, [a, b], [1, 1])
    @test f.boundaryMap[hash(a)] == 1
    @test f.boundaryMap[hash(b)] == 1

    g = Cell(Int, 1)
    push!(g, a, 1)
    push!(g, b, 2)
    @test_throws AssertionError push!(g, f, 3)
    @test g.boundaryMap[hash(a)] == 1
    @test g.boundaryMap[hash(b)] == 2
end


"""
    project!(fun, grid, i, cⁱ)

Project the function `fun` onto the finite element of `grid` with index
`i`, writing the results in the vector of expansion coefficients `cⁱ`.
"""
function project!(fun, grid::Grid, i::Integer,
                  cⁱ::AbstractVector)
    n = order(grid)
    cⁱ[1] = grid.N[i,1]*grid.W[i,1]*fun(grid.X[i,1])
    for m = 2:n-1
        cⁱ[m] = √(grid.W[i,m])*fun(grid.X[i,m])
    end
    cⁱ[end] = grid.N[i,end]*grid.W[i,end]*fun(grid.X[i,end])
end

"""
    project(fun, basis)

Project the function `fun` onto the FEDVR `basis`.
"""
function project(fun, basis::Basis)
    grid = basis.grid
    C = zero(grid.X)
    for i = elems(grid)
        project!(fun, grid, i, view(C, i, :))
    end
    # The bridge functions overlap the function both in element i and
    # element i+1; thus the contributions need to be added into one
    # expansion coefficient.
    C[2:elcount(grid),1] += C[1:elcount(grid)-1,end]

    boundary_sel(grid, [C[:,1:end-1]'[:]..., C[end]])
end

export project


" Implementation of a separating axis test for the intersection of two convex shapes.
# Arguments
- `shape1::Dict` - Convex shape with entry `vertices`
- `shape2::Dict` - Convex shape with entry `vertices`
"
function do_shapes_intersect(shape1::Dict, shape2::Dict; axes=nothing)
    int_result = true
    min_distances = [-1.]

    if isnothing(axes)
        axes = get_axes([shape1, shape2])
    end

    num_axes = size(axes)[2]
    for i=1:num_axes
        axis = axes[:, i:i]
        p1 = project(shape1, axis)
        p2 = project(shape2, axis)
        overlap = (p2[1] <= p1[1] <= p2[2] || p2[1] <= p1[2] <= p2[2] ||
                   p1[1] <= p2[1] <= p1[2] || p1[1] <= p2[2] <= p1[2] )
        push!(min_distances, minimum([abs(p1[2] - p2[1]), abs(p2[2] - p1[1])]))
        if !overlap
            # push!(min_distances, minimum([abs(p1[2] - p2[1]), abs(p2[2] - p1[1])]))
            int_result = false
        end
    end

    return int_result, maximum(min_distances)
end

" Checks if shape2 contains shape1.
# Arguments
- `shape1::Dict` - Convex shape with entry `vertices`
- `shape2::Dict` - Convex shape with entry `vertices`
"
function does_shape2_contain_shape1(shape1::Dict, shape2::Dict; axes=nothing)
    int_result = false
    min_distance = Inf

    if isnothing(axes)
        axes = get_axes([shape1, shape2])
    end

    num_axes = size(axes)[2]
    for i=1:num_axes
        axis = axes[:, i:i]
        p1 = project(shape1, axis)
        p2 = project(shape2, axis)
        # overlap = (p2[1] <= p1[1] <= p2[2] || p2[1] <= p1[2] <= p2[2] ||
        #            p1[1] <= p2[1] <= p1[2] || p1[1] <= p2[2] <= p1[2] )
        contain = p2[1] <= p1[1] <= p2[2] && p2[1] <= p1[2] <= p2[2]
        if contain
            min_distance = minimum([abs(p1[1] - p2[1]), abs(p1[2] - p2[2]), min_distance])
            int_result = true
        else
            int_result = false
            break
        end
    end

    return int_result, min_distance
end

" Projects the shape vertices onto the axis.
# Arguments
- `shape::Dict` - Shape with entry `vertices`
- `axis` - Axis of projection
"
function project(shape::Dict, axis)
    dots = [(vertex*axis)[1] for vertex in shape["vertices"]]
    min = minimum(dots)
    max = maximum(dots)
    return [min, max]
end

" Gets the principal normals of each shape.
# Arguments
- `shapes::Vector` - Vector with shape definitions
"
function get_axes(shapes::Vector)
    axes = []
    for shape in shapes
        num_v = length(shape["vertices"])
        for i=1:num_v
            normals = nullspace(shape["vertices"][i]-(i==num_v ? shape["vertices"][1] : shape["vertices"][i+1]))
            if isempty(axes)
                axes = normals
            else
                # TODO: remove redundant axes
                axes = hcat([axes normals])
            end
        end
    end
    return axes
end


###############################
##  Markov chain simulation  ##
###############################

function tr_prioH(θ, N_H, N_E, p_H, p_E)
    if rand() < θ #incoming H
        if rand() > (1 - p_H^2)^N_H #can match to another H
            N_H -= 1
        elseif rand() > (1 - p_E*p_H)^N_E #can match to an E
            N_E -= 1
        else
            N_H += 1
        end
    else #incoming E
         if rand() > (1 - p_E*p_H)^N_H #can match to another H
            N_H -= 1
        elseif rand() > (1 - p_E^2)^N_E #can match to an E
            N_E -= 1
        else
            N_E += 1
        end
    end
    return (N_H, N_E)
end

function tr_random_prio(θ, N_H, N_E, p_H, p_E)
    if rand() < θ #incoming H
      n_E = N_E
      n_H = N_H
      while n_E > 0 || n_H > 0
        if rand() < n_E / (n_E + n_H) # probability that we pick an E next
          if rand() < p_E*p_H # can match to the E
            N_E -= 1
            break
          else
            n_E -= 1 # decrease pool of available matches
          end
        else
          if rand() < p_H^2
            N_H -= 1
            break
          else
            n_H -= 1 # decrease pool of available matches.
          end
        end
      end
      if n_E == 0 && n_H == 0
        N_H += 1
      end
    else #incoming E
      n_E = N_E
      n_H = N_H
      while n_E > 0 || n_H > 0
        if rand() < n_E / (n_E + n_H) # probability that we pick an E next
          if rand() < p_E^2 # can match to the E
            N_E -= 1
            break
          else
            n_E -= 1 # decrease pool of available matches
          end
        else
          if rand() < p_H * p_E
            N_H -= 1
            break
          else
            n_H -= 1 # decrease pool of available matches.
          end
        end
      end
      if n_E == 0 && n_H == 0
        N_E += 1
      end
    end
    return (N_H, N_E)
end

function tr_prioE(θ, N_H, N_E, p_H, p_E)
    if rand() < θ #incoming H
        if rand() > (1 - p_E*p_H)^N_E #can match to an E
            N_E -= 1
        elseif rand() > (1 - p_H^2)^N_H #can match to another H
            N_H -= 1
        else
            N_H += 1
        end
    else #incoming E
        if rand() > (1 - p_E^2)^N_E #can match to an E
            N_E -= 1
        elseif rand() > (1 - p_E*p_H)^N_H #can match to another H
            N_H -= 1
        else
            N_E += 1
        end
    end
    return (N_H, N_E)
end

function tr_chains(θ, N_H, N_E, p_H, p_E, d)
    continue_chain = false
    if rand() < θ #incoming H
        if rand() < 1 - (1 - p_H)^d #can match to another H
            continue_chain = true
        else
            N_H += 1
        end
    else #incoming E
         if rand() < 1 - (1 - p_E)^d
            continue_chain = true
        else
            N_E += 1
        end
    end
    while continue_chain
        if rand() < 1 - (1 - p_H)^N_H
            N_H -= 1
        elseif rand() < 1 - (1 - p_E)^N_E
            N_E -= 1
        else
            continue_chain = false
        end
    end
    return (N_H, N_E)
end

function tr_chains_prioE(θ, N_H, N_E, p_H, p_E, d)
    continue_chain = false
    if rand() < θ #incoming H
        if rand() < 1 - (1 - p_H)^d #can match to another H
            continue_chain = true
        else
            N_H += 1
        end
    else #incoming E
         if rand() < 1 - (1 - p_E)^d
            continue_chain = true
        else
            N_E += 1
        end
    end
    while continue_chain
        if rand() < 1 - (1 - p_E)^N_E
            N_E -= 1
        elseif rand() < 1 - (1 - p_H)^N_H
            N_H -= 1
        else
            continue_chain = false
        end
    end
    return (N_H, N_E)
end


module TestMacrosAggregateVector

import Test: @testset, @test
import Volcanito: @group_by, @aggregate_vector, materialize
import DataFrames: DataFrame
import Statistics: mean

@testset "@aggregate_vector" begin
    df = DataFrame(
        a = [1, 2, 3, 4],
        b = [2, 3, 4, 5],
        c = [0, 0, 1, 1],
    )

    plan1 = @group_by(df, c)
    gdf = materialize(plan1)

    plan2 = @aggregate_vector(gdf, m = mean(a + b))
    res = materialize(plan2)
    @test res == DataFrame(
        c = [0, 1],
        m = [mean([1 + 2, 2 + 3]), mean([3 + 4, 4 + 5])],
    )
end

end


using Plots, LaTeXStrings
gr()
fontFamily = "Times"
default(titlefont=(18, fontFamily), guidefont=(18, fontFamily),
      tickfont=(16, fontFamily),legendfont=(13, fontFamily),
      bg_legend = nothing,framestyle = :semi, grid = true)
      # another helpful keyword is:  size= (600, 400)
# ...............................................................#
x = LinRange(0,2π,100)
plot(x, sin.(x), c = :red,label = L"sin(x)", title = "Sine",
    yticks = ([-1,0,1], [L"-1", L"0", L"1"]), ylab = L"f(x)",
    xlab = L"x", xticks = ([0,π/2, π, 3π/2, 2π],
    [L"0", L"\pi/2", L"\pi", L"3\pi/2", L"2\pi"]))
savefig("figSin2.pdf")


using Polynomials
using SemiLagrangian: getpoly
function test_cmplx(n, decref::CartesianIndex{2}=CartesianIndex(0,0))
#    fctref = Polynomial(big.(rationalize.(rand(n^2-1), tol=0.001)+ im*rationalize.(rand(n^2-1), tol=0.001)))
    tabref = [big(i+decref.I[1]+im*(j+decref.I[2])) + 0.1*(rand()+im*rand()-0.5-0.5*im) for i=1:n, j=1:n]
    moy = sum(tabref)/length(tabref)
    tabref = tabref .- moy

    p = getpoly(tabref)
    @show p
    for i=1:n, j=1:n
        res = p(tabref[i,j]) - i -im*j
#        @show n,decref,i,j,res
        @test isapprox(p( tabref[i,j] ), i +im*j, atol=1e-27)
    end
end

@testset "Lagrange complexe" begin
    for n=3:2:11
        println("")
        test_cmplx(n)
        test_cmplx(n, CartesianIndex(35,47))
        test_cmplx(n, CartesianIndex(-3557,4702))
    end
end


"""
    GaborKernel(; ell::Real=1.0, p::Real=1.0)

Gabor kernel with lengthscale `ell` and period `p`.

# Definition

For inputs ``x, x' \\in \\mathbb{R}^d``, the Gabor kernel with lengthscale ``l_i > 0``
and period ``p_i > 0`` is defined as
```math
k(x, x'; l, p) = \\exp\\bigg(- \\cos\\bigg(\\pi\\sum_{i=1}^d \\frac{x_i - x'_i}{p_i}\\bigg)
                             \\sum_{i=1}^d \\frac{(x_i - x'_i)^2}{l_i^2}\\bigg).
```
"""
struct GaborKernel{K<:Kernel} <: Kernel
    kernel::K
    function GaborKernel(; ell=nothing, p=nothing)
        k = _gabor(; ell=ell, p=p)
        return new{typeof(k)}(k)
    end
end

@functor GaborKernel

(κ::GaborKernel)(x, y) = κ.kernel(x, y)

function _gabor(; ell=nothing, p=nothing)
    if ell === nothing
        if p === nothing
            return SqExponentialKernel() * CosineKernel()
        else
            return SqExponentialKernel() * transform(CosineKernel(), 1 ./ p)
        end
    elseif p === nothing
        return transform(SqExponentialKernel(), 1 ./ ell) * CosineKernel()
    else
        return transform(SqExponentialKernel(), 1 ./ ell) *
               transform(CosineKernel(), 1 ./ p)
    end
end

function Base.getproperty(k::GaborKernel, v::Symbol)
    if v == :kernel
        return getfield(k, v)
    elseif v == :ell
        kernel1 = k.kernel.kernels[1]
        if kernel1 isa TransformedKernel
            return 1 ./ kernel1.transform.s[1]
        else
            return 1.0
        end
    elseif v == :p
        kernel2 = k.kernel.kernels[2]
        if kernel2 isa TransformedKernel
            return 1 ./ kernel2.transform.s[1]
        else
            return 1.0
        end
    else
        error("Invalid Property")
    end
end

function Base.show(io::IO, κ::GaborKernel)
    return print(io, "Gabor Kernel (ell = ", κ.ell, ", p = ", κ.p, ")")
end

kernelmatrix(κ::GaborKernel, x::AbstractVector) = kernelmatrix(κ.kernel, x)

function kernelmatrix(κ::GaborKernel, x::AbstractVector, y::AbstractVector)
    return kernelmatrix(κ.kernel, x, y)
end

kernelmatrix_diag(κ::GaborKernel, x::AbstractVector) = kernelmatrix_diag(κ.kernel, x)


#
# Do work here ...
#

v = sin(1.234567) + cos(0.122122)
println("Julia side: I will return a Float64: ", v)

return v


function mle_for_θ(y::AbstractVector, μ::AbstractVector, wts::AbstractVector;
                   maxiter=30, tol=1.e-6)
    function first_derivative(θ::Real)
        tmp(yi, μi) = (yi+θ)/(μi+θ) + log(μi+θ) - 1 - log(θ) - digamma(θ+yi) + digamma(θ)
        unit_weights ? sum(tmp(yi, μi) for (yi, μi) in zip(y, μ)) :
                       sum(wti * tmp(yi, μi) for (wti, yi, μi) in zip(wts, y, μ))
    end
    function second_derivative(θ::Real)
        tmp(yi, μi) = -(yi+θ)/(μi+θ)^2 + 2/(μi+θ) - 1/θ - trigamma(θ+yi) + trigamma(θ)
        unit_weights ? sum(tmp(yi, μi) for (yi, μi) in zip(y, μ)) :
                       sum(wti * tmp(yi, μi) for (wti, yi, μi) in zip(wts, y, μ))
    end

    unit_weights = length(wts) == 0
    if unit_weights
        n = length(y)
        θ = n / sum((yi/μi - 1)^2 for (yi, μi) in zip(y, μ))
    else
        n = sum(wts)
        θ = n / sum(wti * (yi/μi - 1)^2 for (wti, yi, μi) in zip(wts, y, μ))
    end
    δ, converged = one(θ), false

    for t = 1:maxiter
        θ = abs(θ)
        δ = first_derivative(θ) / second_derivative(θ)
        if abs(δ) <= tol
            converged = true
            break
        end
        θ = θ - δ
    end
    if !converged
        info_msg = "Estimating dispersion parameter failed, which may " *
            "indicate Poisson distributed data."
        throw(ConvergenceException(maxiter, NaN, NaN, info_msg))
    end
    θ
end

"""
    negbin(formula,
           data,
           link;
           initialθ::Real=Inf,
           maxiter::Integer=30,
           atol::Real=1e-6,
           rtol::Real=1.e-6,
           verbose::Bool=false,
           kwargs...)

Fit a negative binomial generalized linear model to data, while simultaneously
estimating the shape parameter θ. Extra arguments and keyword arguments will be
passed to [`glm`](@ref).

# Keyword Arguments
- `initialθ::Real=Inf`: Starting value for shape parameter θ. If it is `Inf`
  then the initial value will be estimated by fitting a Poisson distribution.
- `maxiter::Integer=30`: See `maxiter` for [`glm`](@ref)
- `atol::Real=1.0e-6`: See `atol` for [`glm`](@ref)
- `rtol::Real=1.0e-6`: See `rtol` for [`glm`](@ref)
- `verbose::Bool=false`: See `verbose` for [`glm`](@ref)
"""
function negbin(F,
                D,
                args...;
                initialθ::Real=Inf,
                maxiter::Integer=30,
                atol::Real=1e-6,
                rtol::Real=1.e-6,
                verbose::Bool=false,
                kwargs...)
    if haskey(kwargs, :maxIter)
        Base.depwarn("'maxIter' argument is deprecated, use 'maxiter' instead", :negbin)
        maxiter = kwargs[:maxIter]
    end
    if haskey(kwargs, :minStepFac)
        Base.depwarn("'minStepFac' argument is deprecated, use 'minstepfac' instead", :negbin)
        minstepfac = kwargs[:minStepFac]
    end
    if haskey(kwargs, :convTol)
        Base.depwarn("`convTol` argument is deprecated, use `atol` and `rtol` instead", :negbin)
        rtol = kwargs[:convTol]
    end
    if !issubset(keys(kwargs), (:maxIter, :minStepFac, :convTol))
        throw(ArgumentError("unsupported keyword argument"))
    end
    if haskey(kwargs, :tol)
        Base.depwarn("`tol` argument is deprecated, use `atol` and `rtol` instead", :negbin)
        rtol = kwargs[:tol]
    end

    maxiter >= 1 || throw(ArgumentError("maxiter must be positive"))
    atol > 0  || throw(ArgumentError("atol must be positive"))
    rtol > 0  || throw(ArgumentError("rtol must be positive"))
    initialθ > 0 || throw(ArgumentError("initialθ must be positive"))

    # fit a Poisson regression model if the user does not specify an initial θ
    if isinf(initialθ)
        regmodel = glm(F, D, Poisson(), args...;
                       maxiter=maxiter, atol=atol, rtol=rtol, verbose=verbose, kwargs...)
    else
        regmodel = glm(F, D, NegativeBinomial(initialθ), args...;
                       maxiter=maxiter, atol=atol, rtol=rtol, verbose=verbose, kwargs...)
    end

    μ = regmodel.model.rr.mu
    y = regmodel.model.rr.y
    wts = regmodel.model.rr.wts
    lw, ly = length(wts), length(y)
    if lw != ly && lw != 0
        throw(ArgumentError("length of wts must be either $ly or 0 but was $lw"))
    end

    θ = mle_for_θ(y, μ, wts; maxiter=maxiter, tol=rtol)
    d = sqrt(2 * max(1, deviance(regmodel)))
    δ = one(θ)
    ll = loglikelihood(regmodel)
    ll0 = ll + 2 * d

    converged = false
    for i = 1:maxiter
        if abs(ll0 - ll)/d + abs(δ) <= rtol
            converged = true
            break
        end
        verbose && println("[ Alternating iteration ", i, ", θ = ", θ, " ]")
        regmodel = glm(F, D, NegativeBinomial(θ), args...;
                       maxiter=maxiter, atol=atol, rtol=rtol, verbose=verbose, kwargs...)
        μ = regmodel.model.rr.mu
        prevθ = θ
        θ = mle_for_θ(y, μ, wts; maxiter=maxiter, tol=rtol)
        δ = prevθ - θ
        ll0 = ll
        ll = loglikelihood(regmodel)
    end
    converged || throw(ConvergenceException(maxiter))
    regmodel
end


struct kNNRegression
    t
    y::Vector
    kernel::Function
    h::Float64
end

function kNNRegression(X::Matrix, y::Vector; kernel::Symbol = :epanechnikov, bandwidth::Real = 1.0, metric::Metric = Euclidean())
    return kNNRegression(KDTree(X, metric), y,SmoothingKernels.kernels[kernel], bandwidth)
end

function StatsBase.predict(m::kNNRegression, x::Matrix, kernel, bandwidth, k::Integer = 1)
    inds, dists = NearestNeighbors.knn(m.t, x, k)
    map(zip(inds,dists)) do z
        i, d = z
        w = bandwidth .* kernel.(bandwidth .* d)
        sum(w .* m.y[i])/sum(w)
    end
end

StatsBase.predict(m::kNNRegression, x::Matrix, k::Integer = 1) = predict(m, x, m.kernel, m.h, k)


"""
We distinguish between models objects and fitted model objects generally,
for OLSModel specifically:

We separate OLSModel {LinearModel} from its fitted counterpart FittedOLSModel
 {LinearModelFit}  - OLSModel and FittedOLSModel are immutable structs. However,
 FittedOLSModel, whilst immutable, contains a mutable attribute
 modelfit::FitOutput. fit!(model::OLSModel) returns an object of type FitOutput
 in place but fit(model::OLSModel) returns an object of type FittedOLSModel with
 updated modelfit::FitOutput.

 This means in simulations we can save memory allocation by just calling fit!
 without having to pass around y's, X's etc. It also means struct attributes
 are immutable so we can't accidentally overwrite data whilst also leaving
 FitOutput accessible to manipulation after estimating models etc.
"""

"""
    OLSModel <: LinearModel

Object describing data and variance covariance matrix to be computed using an
Ordinary Least Squares Model. Q, R are the QR decomposition pre-calculated and
carried around by OLSModel.
"""
struct OLSModel <: LinearModel
    y::Vector
    X::Matrix
    vcov::vcov
    Q::Matrix # TODO: Create a lighweight version that doesn't carry data
    R::Matrix
    X_names::Vector
    N::Int64
    K::Int64
    # Default standard errors homoscedastic
    function OLSModel(y::Vector,
                      X::Matrix;
                      vcov::vcov = vcovIID(),
                      X_names = nothing)
        N = size(X, 1)
        K = size(X, 2)
        length(y) == N || error("y and x have differing numbers
        of observations.")
        Q, R = qr(X)

        if isnothing(X_names)
            X_names = "x_" .* string.(1:size(X, 2))
        end
        new(y, X, vcov, Q, R, X_names, N, K)
    end
end

# Adding DataFrames functionality
function OLSModel(y::Vector, X::DataFrame; vcov::vcov = vcovIID())
    X_names = names(X)
    X = Matrix(X) # Backcompability warning requires Juila > 0.7
    return OLSModel(y, X, vcov = vcov, X_names = X_names)
end

mutable struct FitOutput
    β::Vector
    se_β::Vector
    pval::Vector
    resid::Vector
    σ_sq::Float64
    vcov_matrix::Matrix
end

struct FittedOLSModel <: LinearModelFit
    y::Vector
    X::Matrix
    vcov::vcov
    Q::Matrix
    R::Matrix
    X_names::Vector
    N::Int64
    K::Int64
    modelfit::FitOutput
end

"""
    fit!(model::OLSModel)

Use the QR decomposition to estimate y = X β and return FitOutput.
"""
function fit!(model::OLSModel)
    y = model.y
    X = model.X
    # QR decomposition for numerical stability
    Q = model.Q
    R = model.R
    # Find beta
    β = inv(R) * Q' * y
    # SEs
    resid = y - X*β
    se_β, pval, σ_sq, vcov_matrix = inference(resid, β, model, model.vcov)
    return FitOutput(β, se_β, pval, resid, σ_sq, vcov_matrix)
end

"""
    fit(model::OLSModel)

Instead of returning FitOutput, returns a FittedOLSModel.
"""
function fit(model::OLSModel)
    FittedOLSModel(
        model.y,
        model.X,
        model.vcov,
        model.Q,
        model.R,
        model.X_names,
        model.N,
        model.K,
        fit!(model))
end


using LazySets, MathematicalPredicates, SparseArrays

function beam_specification()
    # initial set: xᵢ = 0 if i ≤ 300 and xᵢ ∈ [0.0015, 0.002] otherwise
    X0 = Hyperrectangle(low=[zeros(300); fill(0.0015, 48)],
                        high=[zeros(300); fill(0.002, 48)])

    # safety property: x89 ≤ 2100
    property = is_contained_in(HalfSpace(sparsevec([89], [1.0], 348), 2100.0))

    # time horizon: 20 time units
    time_horizon = 20.0

    # specification
    O = Dict(:T => time_horizon, :property => property)

    return X0, O
end


using LinearAlgebra
# positions
x⃗ = [-1, 0, 1] .* 1.0

# position operator
x̂ = Diagonal(x⃗)

# construct differentiation matrix / position operator
X⃗ = [x⃗[i]^(j-1) for i in eachindex(x⃗), j in eachindex(x⃗)]
dX⃗ = [(j-1) * x⃗[i]^(j-2) for i in eachindex(x⃗), j in eachindex(x⃗)]
dX⃗[:,1] .= 0.0 # get rid of NaNs
# p̂ * X⃗ = dX⃗, so just invert to find p̂
p̂ = dX⃗ / X⃗

# fake identity
# won't be identiy for polynomial orders ≥ length(x⃗)-1
I_s =  p̂ * x̂ - x̂ * p̂

# and check that its correct to machine precision
check = 3 .+ x⃗
I_s * check - check


#= Given a number N check if it is a power of 2 in O(1). We will be
using bitwise operator to get the calculation.=#

## Function

function isPowerOfTwo(n)
    return (n & (n - 1) == 0)
end

## Input

println("Enter the number : ")
n = readline()
n = parse(Int64, n)

## Calling the function

if (isPowerOfTwo(n))
    println("Its a power of 2")
else
    println("Its not a power of 2")
end

#=
Sample Test Case
Input:
    Enter a number : 4294967296
Output:
    Its a power of 2
Time complexity: O( 1 )
=#


# ---
# title: 43. Multiply Strings
# id: problem43
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Math, String
# link: <https://leetcode.com/problems/multiply-strings/description/>
# hidden: true
# ---
#
# Given two non-negative integers `num1` and `num2` represented as strings,
# return the product of `num1` and `num2`, also represented as a string.
#
# **Note:**  You must not use any built-in BigInteger library or convert the
# inputs to integer directly.
#
#
#
# **Example 1:**
#
#
#
#     Input: num1 = "2", num2 = "3"
#     Output: "6"
#
#
# **Example 2:**
#
#
#
#     Input: num1 = "123", num2 = "456"
#     Output: "56088"
#
#
#
#
# **Constraints:**
#
#   * `1 <= num1.length, num2.length <= 200`
#   * `num1` and `num2` consist of digits only.
#   * Both `num1` and `num2` do not contain any leading zero, except the number `0` itself.
#
#
## @lc code=start
using LeetCode

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