### A Pluto.jl notebook ###
# v0.15.1

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 5295694c-344f-11ec-2842-d7207c32a925
begin
	using Pkg
	Pkg.activate("../Project.toml")
end

# ╔═╡ 57884ba0-1209-4ebe-aaaf-b03edd5814d8
begin
	using Flux
	using Flux.Data: MNIST
	using Statistics
	using Flux: throttle, params
	using Flux: @epochs, @functor
	using Flux: onehotbatch, argmax
	using Flux.Losses: logitbinarycrossentropy, crossentropy, mse
	using Flux: chunk
	using ProgressMeter: Progress, next!
	using Base.Iterators: repeated, partition
	using Plots
	using Colors
	using Images
end

# ╔═╡ e5ae0713-8bb0-4568-8eda-3cf8fdc5cef0
using PlutoUI

# ╔═╡ f47b82e5-1110-4bfa-9ac1-637076c6faf2
md"""
## Some Examples of using Flux (reference: EECS505, other code from this forked repo)
"""

# ╔═╡ 09fea98d-aaf0-4971-a81f-90b50ef04b47
md"""
## Classifier for data belonging to two circles
"""

# ╔═╡ 3e79f9cc-f01c-4622-ade9-92d6aaa954e1
md"""
Inspect a few images.
"""

# ╔═╡ 6ebac32e-ad7d-4321-be2f-b61127c47771
# [img(X[:, i]) for i in 6:10]

# ╔═╡ f74b47d4-7dff-40a1-b086-26bf8b5c88d7
function generatedata_circle(r1, r2, N, σ=0.1)
	φ1 = LinRange(0, 2 * π, N)
	φ2 = LinRange(0, 2 * π, N)
	rx1 = r1 .+ σ * randn(N)
	rx2 = r2 .+ σ * randn(N)
	X1 = [rx1 .* cos.(φ1) rx1 .* sin.(φ1)]
	X2 = [rx2 .* cos.(φ2) rx2 .* sin.(φ2)]
	return X1', X2'
end

# ╔═╡ b0986b17-bd42-4fb4-a3f1-54a58e52593c
X1c, X2c = generatedata_circle(2, 0.5, 100)

# ╔═╡ 5b13fe75-6636-4ee2-9396-98ba53cedadd
begin
	p1 = scatter(
				X1c[1, :],
				X1c[2, :];
				color="red",
				label="Class 1",
				aspectratio=:equal,
				)
	scatter!(
			X2c[1, :],
			X2c[2, :];
			color="blue",
			label="Class 2",
			legend=:bottomright,
			)
end

# ╔═╡ c1b9423c-c79b-46b5-8f19-c252aaa7473b
X = [X1c X2c]

# ╔═╡ 100900ac-2b1f-492f-9a50-14f4fc5506cb
Y = [zeros(1, size(X1c, 2)) ones(1, size(X2c, 2))]

# ╔═╡ d8ebaeef-27a9-4530-a19f-7987919178df
@bind nHidden Slider(2:10, default=3, show_value=true)

# ╔═╡ b2b1c6a3-acbe-4a39-8f1c-378de73a96ee
nHidden

# ╔═╡ 24d87786-904c-4c2e-9c97-e3574521f49e
activeFunction = relu

# ╔═╡ 123965a9-a4e3-4194-b079-d32bd941aa17
m = Chain(
		Dense(size(X,1), nHidden, activeFunction),
		Dense(nHidden, nHidden, activeFunction),
		Dense(nHidden, 1, activeFunction)
		)

# ╔═╡ bdde6e5e-72ad-4071-ab98-15724f72ba86
md"Output has size $(size(m(X), 2))"

# ╔═╡ 8f5068ce-6301-4234-b284-9ec03b32a8e0
loss(x, y) = mse(m(x), y)

# ╔═╡ d3b9eaa8-be71-4b83-91fd-2662ce842259
iters = 10000

# ╔═╡ f373f536-9929-40a4-ad0d-b2adb606376a
md"""
We package the dataset into `(X, Y)` tuples.
"""

# ╔═╡ ea9abbbe-36a8-46d6-8987-0f4840e4b804
dataset = repeated((X, Y), iters)

# ╔═╡ 83b8153e-8b06-4f25-ad4d-b92c5ed21884
begin
	opt = ADAM()
	evalcb = () -> @show([loss(X, Y)])
end

# ╔═╡ 2c8025b6-cedc-463a-9625-551c4e402052
with_terminal() do
	Flux.train!(loss, params(m), dataset, opt; cb = throttle(evalcb, 0.5))
end

# ╔═╡ b36530a1-ad38-4178-a914-51cc59c1d294
function display_decision_boundaries(X1c, X2c, m, x1range, x2range, τ=0.0)
    D = [m([x1; x2])[1] for x2 in x2range, x1 in x1range]
    heatmap(
			x1range,
			x2range,
			sign.(D .- τ);
			color=:grays,
			xlim = [minimum(x1range),maximum(x1range)],
			ylim = [minimum(x2range),maximum(x2range)],
			)
    scatter!(
			X1c[1, :],
			X1c[2, :],
			color="red",
			label="Class 1",
			aspectratio=:1.0,
			)
    scatter!(
			X2c[1, :],
			X2c[2, :],
			color="blue",
			label="Class 2",
			)
end

# ╔═╡ 09081010-4a09-42fc-8658-f02d7c35a4e9
x1range = range(-3; stop=3, length=200)

# ╔═╡ a184440d-621b-42ff-b4a9-a596b248acbc
x2range = range(-3; stop=3, length=200)

# ╔═╡ e7f6169f-aa66-4a43-9831-227507047156
begin
	display_decision_boundaries(X1c, X2c, m, x1range, x2range)
	plot!(; title="Loss = $(round(loss(X,Y), digits=4))")
end

# ╔═╡ 6eccef88-ac15-4538-93f7-7e807e4891fd
with_terminal() do
	@show m.layers[1].W
	@show m.layers[1].b
	@show m.layers[1].σ;
end

# ╔═╡ b22cf1c6-aed7-43b4-9687-ad04dcd983d4
md"""
## Classifier for MNIST Data
"""

# ╔═╡ 6159bf70-b4d8-4bcb-ba24-4d3502d98e7c
imgs = MNIST.images()

# ╔═╡ 1df18a22-46e5-4c8a-a54e-88fe342ac39c
typeof(imgs)

# ╔═╡ 5a752624-23ad-41ab-8c53-2caf0ded9aa9

# ╔═╡ 91b6edde-c595-4bd4-99c8-dd07752c7a7f

# ╔═╡ 08bb5a70-39e0-46d3-81f7-71e27a87a4e0
function convert_to_image(x, y_size)
    Gray.(permutedims(vcat(reshape.(chunk(x |> cpu, y_size), 28, :)...), (2, 1)))
end

# ╔═╡ 54bc078c-8069-4d2c-a578-7b9d1d2fc367

# ╔═╡ 2cd4d120-f02d-4667-8977-a40e800e9a3a

# ╔═╡ e1f0af51-2519-4846-82d5-76e02f1bd938

# ╔═╡ a8c4ae86-2b50-4f7d-bfd3-763e2d5d2321

# ╔═╡ 10de7adb-52ec-4e24-bfa5-e5652cb3f226

# ╔═╡ Cell order:
# ╟─f47b82e5-1110-4bfa-9ac1-637076c6faf2
# ╠═5295694c-344f-11ec-2842-d7207c32a925
# ╠═57884ba0-1209-4ebe-aaaf-b03edd5814d8
# ╠═e5ae0713-8bb0-4568-8eda-3cf8fdc5cef0
# ╠═09fea98d-aaf0-4971-a81f-90b50ef04b47
# ╟─3e79f9cc-f01c-4622-ade9-92d6aaa954e1
# ╠═6ebac32e-ad7d-4321-be2f-b61127c47771
# ╠═f74b47d4-7dff-40a1-b086-26bf8b5c88d7
# ╠═b0986b17-bd42-4fb4-a3f1-54a58e52593c
# ╟─5b13fe75-6636-4ee2-9396-98ba53cedadd
# ╠═c1b9423c-c79b-46b5-8f19-c252aaa7473b
# ╠═100900ac-2b1f-492f-9a50-14f4fc5506cb
# ╠═d8ebaeef-27a9-4530-a19f-7987919178df
# ╠═b2b1c6a3-acbe-4a39-8f1c-378de73a96ee
# ╠═24d87786-904c-4c2e-9c97-e3574521f49e
# ╠═123965a9-a4e3-4194-b079-d32bd941aa17
# ╠═bdde6e5e-72ad-4071-ab98-15724f72ba86
# ╠═8f5068ce-6301-4234-b284-9ec03b32a8e0
# ╠═d3b9eaa8-be71-4b83-91fd-2662ce842259
# ╟─f373f536-9929-40a4-ad0d-b2adb606376a
# ╠═ea9abbbe-36a8-46d6-8987-0f4840e4b804
# ╠═83b8153e-8b06-4f25-ad4d-b92c5ed21884
# ╠═2c8025b6-cedc-463a-9625-551c4e402052
# ╠═b36530a1-ad38-4178-a914-51cc59c1d294
# ╠═09081010-4a09-42fc-8658-f02d7c35a4e9
# ╠═a184440d-621b-42ff-b4a9-a596b248acbc
# ╠═e7f6169f-aa66-4a43-9831-227507047156
# ╠═6eccef88-ac15-4538-93f7-7e807e4891fd
# ╟─b22cf1c6-aed7-43b4-9687-ad04dcd983d4
# ╠═6159bf70-b4d8-4bcb-ba24-4d3502d98e7c
# ╠═1df18a22-46e5-4c8a-a54e-88fe342ac39c
# ╠═5a752624-23ad-41ab-8c53-2caf0ded9aa9
# ╠═91b6edde-c595-4bd4-99c8-dd07752c7a7f
# ╠═08bb5a70-39e0-46d3-81f7-71e27a87a4e0
# ╠═54bc078c-8069-4d2c-a578-7b9d1d2fc367
# ╠═2cd4d120-f02d-4667-8977-a40e800e9a3a
# ╠═e1f0af51-2519-4846-82d5-76e02f1bd938
# ╠═a8c4ae86-2b50-4f7d-bfd3-763e2d5d2321
# ╠═10de7adb-52ec-4e24-bfa5-e5652cb3f226


using Test, StatsModels
using DataFrames
using StatsBase

using Unfold
include("test_utilities.jl")

data,evts = loadtestdata("test_case_3a") #
data_r = reshape(data,(1,:))

data_e,times = Unfold.epoch(data=data_r,tbl=evts,τ=(0.,1),sfreq=20) # cut the data into epochs
basisfunction = firbasis(τ=(0.,1),sfreq=20,name="basisA")

# f  = @formula 0~ 1 * beta_0 + contA*beta_1
# fit => beta_0 = -1.5, beta_1 = 2
# predict => (contA=0.5) y_hat = 1 * -1.5 + 0.5 * 2 = -0.5

f  = @formula 0~1 # 1
m_mul,m_mul_results = fit(UnfoldLinearModel,f,evts,data_e,times)
m_tul,m_tul_results = fit(UnfoldLinearModel,f,evts,data_r,basisfunction)

conditionA  = [0,1.]
continuousA = [-1.,0,1.]

tmp = reshape([ [x,y]  for x=conditionA, y=continuousA ],length(conditionA)*length(continuousA))
evts_grid = DataFrame(hcat(tmp...)')
rename!(evts_grid,["conditionA","continuousA"])

yhat_tul = predict(m_tul,evts_grid)
yhat_mul = predict(m_mul,evts_grid,times)

@test unique(yhat_mul.times) == times

@test size(m_tul_results)[1] * size(evts_grid)[1] == size(yhat_tul)[1]
@test size(m_mul_results)[1] * size(evts_grid)[1] == size(yhat_mul)[1]

@test yhat_mul[end-3,:yhat] ≈ mean(data[data.!=0])
@test yhat_tul[end-3,:yhat] ≈ mean(data[data.!=0])

@test yhat_mul.times[1] == 0.

f  = @formula 0~1+conditionA+continuousA# 1
m_tul,m_tul_results = fit(UnfoldLinearModel,f,evts,data_r,basisfunction)
m_mul,m_mul_results = fit(UnfoldLinearModel,f,evts,data_e,times)
yhat_tul = predict(m_tul,evts_grid)
yhat_mul = predict(m_mul,evts_grid,times)
#@test predict(m_mul,evts).yhat ≈ yhat.yhat

@test isapprox(sum(m_tul.beta[[5,25,45]]),  yhat_tul.yhat[end-3],atol=0.001)

@test isapprox(yhat_tul[yhat_tul.times.==0.5,:yhat], yhat_mul[yhat_mul.times.==0.5,:yhat],atol=0.001)
@test isapprox(yhat_tul[yhat_tul.times.==0.5,:yhat], [-2,1,2,5,6,9.],atol=0.0001)

if 1 == 0
    #for visualization
    using AlgebraOfGraphics
    yhat_mul.conditionA = categorical(yhat_mul.conditionA)
    m = mapping(:times,:yhat,color=:continuousA,linestyle=:conditionA)
    df = yhat_mul
    AlgebraOfGraphics.data(df) * visual(Lines) * m  |> draw
end


# knot.jl
# File for creating knots in the physical space in 1D, 2D or 3D.
@doc """
### calcKnot

Computes a uniform knot vector for the mapping object. It is assumed that the
number of control points is known from before. It is user specified. The knot
vector is computed in the parametric space [0,1] in all the 3 directions

**Arguments**

*  `map` : Object of mapping type

"""->

function calcKnot(map::AbstractMappingType)

  # Assertions
  for i = 1:3
    @assert length(map.edge_knot[i]) >= 2*map.order[i] "Length of knot vector
    not sufficient to ensure nonperiodic/clamped/open knot vector in $i
    direction"
  end

  for di = 1:3
    order = map.order[di]
    nctl = map.nctl[di]
    # Zero out the knot vector
    map.edge_knot[di][:] = 0.0
    step = 1 / (nctl-order+1)
    low = 0.0
    for i = (order+1):nctl
      map.edge_knot[di][i] = low + step
      low += step
    end
    # Fill the end of the knot vector
    map.edge_knot[di][nctl+1:nctl+order] = 1.0
  end

  return nothing
end

@doc """
### calcKnotLMS

Computes the knot vector in the parametric space [0,1]. This is for a least
squares curve approximation.

REFERENCE: Equation 9.68 & eqation 9.69, Page 412 The NURBS book, 2nd Edition

**Inputs**

*  `order` : Order of B-spline
*  `nctl`  : Number of control points for that knot vector
*  `X`     : ordinates used to make the knot vector in parametric space
*  `U`     : Knot vector to be populated

**Outputs**

*  None

"""->

function calcKnotLMS(order, nctl, X, U)
  # REFERENCE : Gaetan's pySpline knots.f90, knots_lms

  @assert length(U) = nctl + order "Length of specified knot vector is not equal
  to the sum of number of control points and order"

  lengthX = length(X)

  # Populate the extremas of the knot vector
  for i = 1:order
    U[i] = 0.0
    U[nctl+i] = 1.0
  end  # End for i = 1:order

  # Populate the middle of the knot vector. Recall the structure of knot vectors
  if isodd(lengthX)
    d = lengthX/(nctl-order+1)
    for j = 1:nctl-order
      i = convert(Int, floor(j*d)) # Conversion for integer based indexing
      alpha = j*d-i
      U[order+j] = (1-alpha)*X[i] + alpha*X[i+2]
    end
  else
    d = lengthX/(nctl-order+1)
    for j = 1:nctl-order
      i = convert(Int, floor(j*d)) # Conversion for integer based indexing
      alpha = j*d-i+0.5
      U[order+j] = (1-alpha)*X[i] + alpha*X[i]
    end
  end  # End if isodd(lengthX)

  return nothing
end

@doc """
### calcKnotInterp

Create a knot vector for an interpolating spline. This is again done in the
parametric space between [0,1]

**Inputs**

*  `X`  : Ordinates used to make the knot vector in the parametric space
*  `nd` : Number of derivatives specified
*  `order` : Order of the B-spline curve
*  `U`  : Knot vector to be computed

**Outputs**

*  None

"""->
function calcKnotInterp(X, nd, order, U)

  lengthX = length(X)
  nctl = lengthX + nd

  @assert nctl+order == length(U) # Sanity check

  # Populate the ends of the knot vector in the parametric space
  for j = 1:order
    U[j] = 0.0
    U[nctl+j] = 1.0
  end

  if nd == lengthX
    if order == 3
      for i = 1:(lengthX-2)
        U[order+2*i-1] = 0.5*(X[i] + X[i+1])
        U[order+2*i] = X[i+1]
      end # end for i = 1:(lengthX-2)
      U[nctl] = 0.5*(X[lengthX-1] + 1)
    elseif k == 4
      U[5] = 0.5 * X[2]
      U[nctl] = 0.5 * (X[lengthX-1] + 1)
      for i = 1:(lengthX-3)
        U[order+2*i] = (1/3)*(2*X[i+1] + X[i+2])
        U[order+2*i+1] = (1/3)*(X[i+1] + 2*X[i+2])
      end  # End for i = 1:(lengthX-3)
    end  # end if order == 3
  elseif nd == 0
    if order == 2
      for j = 1:lengthX-order
        U[order+j] = X[j+1]
      end
    elseif order == 3
      for j = 1:lengthX-order
        U[order+j] = 0.5*(X[j+1] + X[j+2])
      end
    elseif order == 4
      for j = 1:lengthX-order
        U[order+j] = X[j+2]
      end
    else
      error("Interpolation is only available for k = 2, 3 or 4")
    end
  else
    error("More options for nd not available yet")
  end  # End if nd == lengthX

  return nothing
end


@testset "MPS" begin

D = 10
d = 4
sites = 5
T = ComplexF64

@testset "Random MPS with the same physical dimension" begin

    ψ = randn(MPS{T}, sites, D, d)

    @testset "has correct number of sites" begin
        @test length(ψ) == sites
        @test size(ψ) == (sites, )
    end

    @testset "has correct type" begin
        @test eltype(ψ) == T
    end

    @testset "has correct rank" begin
        @test rank(ψ) == Tuple(fill(d, 1:sites))
    end

    @testset "has correct bonds" begin
        @test bond_dimension(ψ) ≈ D
        @test verify_bonds(ψ) === nothing
    end

    @testset "is equal to itself" begin
        @test ψ == ψ
        @test ψ ≈ ψ
    end

    @testset "is equal to its copy" begin
        ϕ = copy(ψ)
        @test ϕ == ψ
        @test ϕ ≈ ψ
    end
end

@testset "Random MPS with varying physical dimension" begin

    dims = (3, 2, 5, 4)
    ψ = randn(MPS{T}, D, dims)

    @testset "has correct number of sites" begin
        n = length(dims)
        @test length(ψ) == n
        @test size(ψ) == (n, )
    end

    @testset "has correct type" begin
        @test eltype(ψ) == T
    end

    @testset "has correct rank" begin
        @test rank(ψ) == dims
    end

    @testset "has correct bonds" begin
        @test bond_dimension(ψ) ≈ D
        @test verify_bonds(ψ) === nothing
    end

    @testset "is equal to itself" begin
        @test ψ == ψ
        @test ψ ≈ ψ
    end

    @testset "is equal to its copy" begin
        ϕ = copy(ψ)
        @test ϕ == ψ
        @test ϕ ≈ ψ
    end
end

@testset "Random MPO with the same physical dimension" begin

    W = randn(MPO{T}, sites, D, d)

    @testset "has correct number of sites" begin
        @test length(W) == sites
        @test size(W) == (sites, )
    end

    @testset "has correct type" begin
        @test eltype(W) == T
    end

    @testset "is equal to itself" begin
        @test W == W
        @test W ≈ W
    end

    @testset "is equal to its copy" begin
        U = copy(W)
        @test U == W
        @test U ≈ W
    end
end

@testset "Random MPO with varying physical dimension" begin

     dims = (3, 2, 5, 4)
     W = randn(MPO{T}, D, dims)

     @testset "has correct number of sites" begin
         n = length(dims)
         @test length(W) == n
         @test size(W) == (n, )
     end

     @testset "has correct type" begin
         @test eltype(W) == T
     end

     @testset "is equal to itself" begin
         @test W == W
         @test W ≈ W
     end

     @testset "is equal to its copy" begin
         U = copy(W)
         @test U == W
         @test U ≈ W
     end
end

@testset "MPS from tensor" begin
    ϵ = 1E-14

    dims = (2,3,4,3,5)
    sites = length(dims)
    A = randn(T, dims)

    ψ = MPS(A, :right)

    @test norm(ψ) ≈ 1
    @test_nowarn verify_bonds(ψ)
    @test_nowarn verify_physical_dims(ψ, dims)
    @test is_right_normalized(ψ)

    B = randn(T, dims...)
    ϕ = MPS(B, :left)

    @test norm(ϕ) ≈ 1
    @test_nowarn verify_bonds(ϕ)
    @test_nowarn verify_physical_dims(ϕ, dims)
    @test is_left_normalized(ϕ)

    χ = MPS(A, :left)

    @test norm(χ) ≈ 1
    @test abs(1 - abs(dot(ψ, χ))) < ϵ
end

end

@testset "Objects with equal tensors have the same hash" begin
    D = 10
    d = 4
    sites = 5
    T = ComplexF64

    ψ = randn(MPS{T}, sites, D, d)
    ϕ = copy(ψ)

    W = randn(MPO{T}, sites, D, d)
    V = copy(W)

    @testset "Equal MPSs have the same hash" begin
        @test hash(ψ) == hash(ϕ)
    end

    @testset "Equal MPOs have the same hash" begin
        @test hash(W) == hash(V)
    end

    @testset "Equal tuples with MPS and MPO have the same hash" begin
        tuple_1 = (ψ, W, [1, 2, 3])
        tuple_2 = (ϕ, V, [1, 2, 3])
        @test tuple_1 == tuple_2
        @test hash(tuple_1) == hash(tuple_2)
    end
end


module Smoothing
using SparseArrays
using LarSurf

# function eye(sz)
#     arr = spzeros(Int8, sz, sz)
#     for i=1:sz
#         arr[i,i] = 1
#     end
#     return arr
# end
#
# function setIzero!(arr)
#     for i=1:size(arr)[1]
#         arr[i,i] = 0
#     end
#     return arr
# end
#
# function getDiag(arr)
#     sz = size(arr)[1]
#     arro = zeros(eltype(arr), sz)
#     for i=1:sz
#         arro[i] = arr[i,i]
#     end
#     return arro
# end
#
# function make_full(arr)
#     arrout = Array{eltype(arr), 2}(undef, size(arr)...)
#     for i=1:size(arr)[1]
#         for j=1:size(arr)[2]
#             arrout[i,j] = arr[i,j]
#         end
#     end
#     return arrout
# end
# #neighboor vertexes
#
# """
# Smoothing based on EV
# """
# function smoothing_EV(V, EVch::SparseArrays.SparseMatrixCSC, k=0.35)
# #     Matrix(setIzero!(EVch' * EVch))
#
#     neighboorNumber = getDiag(EVch' * EVch)
#     neighboors = setIzero!(EVch' * EVch)
#
# #     targetV =  neighboors * V' ./ neighboorNumber
# #     diff = targetV - V'
# #     newV = (V' + k * diff)'
#
#     targetV = V * neighboors ./ neighboorNumber'
#     # println("targetV shape: ", size(targetV))
#     diff = targetV - V
#     # println("diff shape: ", size(diff))
#     newV = V + k * diff
#
#     return make_full(newV)
# end

function eye(sz)
    arr = spzeros(Int8, sz, sz)
    for i=1:sz
        arr[i,i] = 1
    end
    return arr
end

function setIzero!(arr)
    for i=1:size(arr)[1]
        arr[i,i] = 0
    end
    return arr
end

function getDiag(arr)
    sz = size(arr)[1]
    arro = zeros(eltype(arr), sz)
    for i=1:sz
        arro[i] = arr[i,i]
    end
    return arro
end

function make_full(arr)
    arrout = Array{eltype(arr), 2}(undef, size(arr)...)
    for i=1:size(arr)[1]
        for j=1:size(arr)[2]
            arrout[i,j] = arr[i,j]
        end
    end
    return arrout
end
#neighboor vertexes

function smoothing_EV(V, EVch::SparseMatrixCSC, k=0.35)
#     Matrix(setIzero!(EVch' * EVch))

	EVchPow = EVch' * EVch
    neighboorNumber = getDiag(EVchPow)
    neighboors = setIzero!(EVchPow)
	dropzeros!(neighboors)

#     targetV =  neighboors * V' ./ neighboorNumber
#     diff = targetV - V'
#     newV = (V' + k * diff)'

    targetV = (V * neighboors) ./ neighboorNumber'
    # println("targetV shape: ", size(targetV))
    diff = targetV - V
	# @info "sm V orig   $(V[:,1:5])"
	# @info "sm V target $(targetV[:,1:5])"
	# @info "sm V diff   $(diff[:,1:5])"
    # println("diff shape: ", size(diff))
    newV = V + k * diff

    return make_full(newV)
end

"""
Get EV array from FV of surface of quads. Filtration step does not work for non surface
objects. But implementation is prepared for triangles
"""
function get_EV_quads(FV::Array{Array{Int64,1},1}; return_unfiltered=false, data=nothing)
	@debug "getFV " FV
	# EV = []
	# for f in FV
	# 	push!(EV, [[f[1],f[2]],[f[3],f[4]],  [f[1],f[3]],[f[2],f[4]]])
	# 	# push!(EV, [[f[1],f[2]],[f[3],f[4]],  [f[1],f[3]],[f[2],f[4]]])
	# end
	if size(FV[1],1) == 4

		# couples = [[1,2], [3,4], [1,3], [2,4]]
		couples = [[1,2], [2,3], [3,4], [4,1]]
	else
		couples = [[1,2], [2,3], [3,1]]
	end
	szc = size(couples,1)

	ev_time = @elapsed EV = reshape([sort([f[couples[c][1]], f[couples[c][2]]]) for f in FV, c=1:szc],:)
	# doubleedges = Base.sort(mycat(EV))
	if return_unfiltered
		if data != nothing
			data["smoothing EV construction time [s]"] = ev_time
		end
		return EV
	end
    ev_unique_time = @elapsed EVnew = collect(Set(EV))
	if data != nothing
		data["smoothing EV construction time [s]"] = ev_time
		data["smoothing make EV unique time [s]"] = ev_unique_time
	end
	return EVnew
	# doubleedges = Base.sort(EV)
	# doubleedges = convert(LarSurf.Lar.Cells, doubleedges)
	# newEV = [doubleedges[k] for k=1:2:length(doubleedges)]
	# return newEV
end

function get_EV_quads2(FV::Array{Array{Int64,1},1}; return_unfiltered=false)
	# @info "getFV " FV
	# EV = []
	# for f in FV
	# 	push!(EV, [[f[1],f[2]],[f[3],f[4]],  [f[1],f[3]],[f[2],f[4]]])
	# 	# push!(EV, [[f[1],f[2]],[f[3],f[4]],  [f[1],f[3]],[f[2],f[4]]])
	# end
	if size(FV[1],1) == 4
		couples = [[1,2], [3,4], [1,3], [2,4]]
	else
		couples = [[1,2], [2,3], [3,1]]
	end
	EV = reshape([sort([f[couples[c][1]], f[couples[c][2]]]) for f in FV, c=1:size(couples,1)],:)
	# doubleedges = Base.sort(mycat(EV))
	if return_unfiltered
		return EV
	end
	doubleedges = Base.sort(EV)
	doubleedges = convert(LarSurf.Lar.Cells, doubleedges)
	newEV = [doubleedges[k] for k=1:2:length(doubleedges)]
	return newEV
end

"""
Smoothing with defined k. Works for quads or triangles.
"""
function smoothing_FV(V::Array, FV::Array{Array{Int64,1},1}, k=0.35, n=1)
	@info "smoothing V by FV, size(V) = $(size(V)), size(FV) = $(size(FV))"

	EV = LarSurf.Smoothing.get_EV_quads(FV)
	# kEV = LarSurf.characteristicMatrix(aEV, size(bigV)[2])
	# kEV = LarSurf.characteristicMatrix(aEV, size(V)[2])
	aEV = LarSurf.ll2array(EV)
	kEV = LarSurf.characteristicMatrix(aEV, size(V)[2])
	# kEV = Lar.characteristicMatrix(EV)
	# @info "computing new V" @__MODULE__
	# @info "computing new V" @__FILE__
	newV = V
	for i=1:n
		newV = smoothing_EV(newV, kEV, k)
	end
	return newV
end

function smoothing_FV_taubin(V::Array, FV::Array{Array{Int64,1},1}, k1=0.35, k2=-0.1, n=1, data=nothing)
	@info "smoothing V by FV, size(V) = $(size(V)), size(FV) = $(size(FV))"
	if size(FV[1],1) == 4
		@info "FV are quads"
	else
		@info "FV are triangles"
	end
	EV = get_EV_quads(FV, data=data)

	# LarSurf
	aEV = LarSurf.ll2array(EV)

	# kEV = LarSurf.characteristicMatrix(aEV, size(bigV)[2])
	kEV = LarSurf.characteristicMatrix(aEV, size(V)[2])
	# kEV = Lar.characteristicMatrix(EV)
	@info "computing new V"
	newV = V
	taubin_time = @elapsed for i=1:n
		@info "taubin iteration $(i)"
		newV = smoothing_EV(newV, kEV, k1)
		newV = smoothing_EV(newV, kEV, k2)
	end

	if data != nothing
		data["smoothing taubin iterations time [s]"] = taubin_time
	end
	return newV
end

# """
# Iterative smoothing
# """
# function smoothing_FV(V::Array, FV::Array{Array{Int64,1},1}, k::Real, n::Integer)
# 	Vtmp = V
# 	for i=1:n
# 		Vtmp = smoothing_FV(Vtmp, FV, k)
# 	end
# 	return Vtmp
#
# end
#
# """
# Iterative smoothing
# """
# function smoothing_FV_taubin(V::Array, FV::Array{Array{Int64,1},1}, k1::Real, k2::Real, n::Integer)
# 	Vtmp = V
# 	for i=1:n
# 		Vtmp = smoothing_FV(Vtmp, FV, k1)
# 		Vtmp = smoothing_FV(Vtmp, FV, k2)
# 	end
# 	return Vtmp
#
# end

end


using LDUFacts: sqrtapp

@testset "lduupdate(A+u*u')" begin
    u = ones(5)
    v = [1.0, 2, 0, -1, 1]
    A = [-2.0 2 2 0 2; 2 1 1 0 1; 2 1 2 -1 1; 0 0 -1 -2 0; 2 1 1 0 -1]
    la = ldu(A, FullPivot())
    f = 0.5
    ldu_update!(la, v, f)
    @test norm(la.P' * (A + v * v' * f) * la.P - la.L * la.D * la.U) <= 1e-13
end

@testset "lduupdate(A+u*v'+v*u')" begin
    u = ones(5)
    v = [1.0, 2, 0, -1, 1]
    A = [-2.0 2 2 0 2; 2 1 1 0 1; 2 1 2 -1 1; 0 0 -1 -2 0; 2 1 1 0 -1]
    la = ldu(A, FullPivot())
    ldu_update!(la, u, v, 0.1)
    @test norm(la.P' * (A + (u * v' + v * u') * 0.1) * la.P - la.L * la.D * la.U) <= 1e-13
    # note: test fails when f = -1 - check!!!
    la = ldu(A, FullPivot())
    f = -1
    @test_throws SingularException ldu_update!(la, u, v, f)
    #@test_broken norm(la.P' * (A + (u * v' + v * u') * f) * la.P - la.L * la.D * la.U) <= 1e-13
end

@testset "sqrtapp" begin
    @test abs(sqrtapp(0.5) / sqrt(0.5) - 1) < 0.05
    @test abs(sqrtapp(3//2) / sqrt(1.5) - 1) < 0.05
end


using Base.Test
import WaterData
import WaterData.testhelpers: between, noerror

@testset "Regions and tessellations" begin
    @testset "Bounding boxes" begin
        bb = WaterData.BoundingBox(1,2,0,3)
        @test (1.5, 1.5) ∈ bb
        @test (0, 0) ∉ bb
        infbox = WaterData.BoundingBox(-2.5, 2.5, -Inf, Inf)
        @test (1, 4) ∈ infbox
        @test (-2, 3e9) ∈ infbox
        @test (-8, 0) ∉ infbox
    end

    @testset "Polygon inclusion testing" begin
        xs = [1, 5, 5, 4, 3, 2, 1]
        ys = [1, 1, 3, 2, 3, 2, 3]
        poly = WaterData.Polygon(xs, ys)

        # outside bounding box
        @test (0, 0) ∉ poly
        @test (6, 2) ∉ poly
        # clearly inside, non-pathological
        @test (1.5, 1.5) ∈ poly
        @test (3, 2.5) ∈ poly
        @test (4.5, 1.5) ∈ poly
        # outside, but within bounding box, non-pathological
        @test (2, 2.5) ∉ poly
        @test (4, 2.5) ∉ poly
        # outside, pathogical cases
        @test (0, 1) ∉ poly
        @test (0, 2) ∉ poly
        @test (0, 3) ∉ poly
        @test (6, 1) ∉ poly
        @test (6, 2) ∉ poly
        @test (6, 3) ∉ poly
        # inside, pathological cases
        @test (3, 2) ∈ poly
        @test (1.5, 2) ∈ poly

        # one more to convince myself it works
        xs = [1, 2, 3, 4, 5, 5, 4, 3, 2, 1]
        ys = [3, 3, 4, 3, 3, 2, 2, 1, 2, 2]
        poly = WaterData.Polygon(xs, ys)

        # outside bounding box
        @test (-3, -2) ∉ poly
        # clearly inside, non-pathological
        @test (2.5, 2.5) ∈ poly
        # outside, but within bounding box, non-pathological
        @test (1.5, 3.5) ∉ poly
        # outside, pathogical cases
        @test (1, 4) ∉ poly
        @test (0, 3) ∉ poly
        # inside, pathological cases
        @test (3, 2) ∈ poly
        @test (3, 3) ∈ poly
    end

    @testset "Normalisation functions" begin
        norm12 = WaterData.norm12
        unnorm = WaterData.unnorm
        lognorm12 = WaterData.lognorm12
        unlognorm = WaterData.unlognorm
        min_coord = WaterData.min_coord
        max_coord = WaterData.max_coord

        @testset "Normalisation to (1,2)" begin
            @test norm12(4, 4, 6) == min_coord
            @test norm12(6, 4, 6) == max_coord
            @test norm12(5, 4, 6) ≈ 1.5
            @test norm12(0, 4, 6) < min_coord
            @test norm12(10, 4, 6) > max_coord

            r = collect(linspace(20, 30))
            @test all(norm12(r) .≥ min_coord)
            @test all(norm12(r) .≤ max_coord)

            @test lognorm12(10^1.5, 10^1, 10^2) ≈ 1.5
            @test lognorm12(5, 10, 100) < min_coord
            @test lognorm12(500, 10, 100) > max_coord
            @test lognorm12(10, 10, 100) == min_coord
            @test lognorm12(100, 10, 100) == max_coord

            r2 = collect(logspace(3, 4))
            @test all(lognorm12(r2) .≥ min_coord)
            @test all(lognorm12(r2) .≤ max_coord)
        end

        @testset "De-normalising from (1,2)" begin
            @test unnorm(min_coord, 40, 60) == 40
            @test unnorm(max_coord, 40, 60) == 60
            @test unnorm(norm12(50, 40, 60), 40, 60) ≈ 50
            @test unnorm(0.5, 40, 60) < 40
            @test unnorm(2.5, 40, 60) > 60

            @test unlognorm(min_coord, 10, 100) ≈ 10
            @test unlognorm(max_coord, 10, 100) ≈ 100
            @test unlognorm(lognorm12(50, 10, 100), 10, 100) ≈ 50
            @test unlognorm(0.5, 10, 100) < 10
            @test unlognorm(500, 10, 100) > 500
        end
    end

    @testset "Tessellation and interpolation" begin
        geta = WaterData.geta
        getb = WaterData.getb
        getc = WaterData.getc
        getx = WaterData.getx
        gety = WaterData.gety

        xs = [0, 1, 0, 1]
        ys = [0, 0, 1, 1]
        zs = [0, 1, 1, 2]
        @test WaterData.get_tessellation(xs, ys) |> noerror
        tess = WaterData.get_tessellation(xs, ys)
        @test WaterData.findtriangle(tess, 1.2, 1.2) |> noerror
        @test WaterData.findtriangle(tess, 1.8, 1.8) |> noerror
        tri1 = WaterData.findtriangle(tess, 1.2, 1.2)
        tri2 = WaterData.findtriangle(tess, 1.8, 1.8)
        p1, p2, p3 = geta(tri1), getb(tri1), getc(tri1)
        q1, q2, q3 = geta(tri2), getb(tri2), getc(tri2)

        for f in [getx, gety], p in [p1, p2, p3, q1, q2, q3]
            @test f(p) ∈ (WaterData.min_coord, WaterData.max_coord)
        end

        λ1 = WaterData.barycoords(tri1, 1.2, 1.2)
        @test sum(λ1) ≈ 1
        λ2 = WaterData.barycoords(tri2, 1.8, 1.8)
        @test sum(λ2) ≈ 1

        x1 = [0, 1, 2]
        x2 = [8, 9, 10]
        @test WaterData.lininterp(λ1, x1) |> between(0, 2)
        @test WaterData.lininterp(λ2, x2) |> between(8, 10)
    end

    @testset "Testing for points in a tessellation" begin
        # it's a parallelogram
        tess = WaterData.get_tessellation([0., 1, 1, 2], [0., 0, 1, 1])

        # any point in (1,2)x(1,2) should be inside the bounding box
        @test in(1.5, 1.5, WaterData.BoundingBox(tess))
        @test in(0.5, 1.5, WaterData.BoundingBox(tess)) == false

        # only points actually in the parallelogram should test true here
        @test in(1.1, 1.9, tess) == false  # top left corner
        @test in(1.5, 1.5, tess)    # the middle
        @test in(1.9, 1.1, tess) == false  # bottom right corner
        @test in(1.1, 1.1, tess)    # bottom left corner
        @test in(1.9, 1.9, tess)    # bottom right corner
    end
end


## @lc code=start
using LeetCode
count_even(num::Int) = sum(iseven(sum(j-'0' for j in string(i))) for i in 1:num)

## @lc code=end

## @lc test=start
@testset "2180.count-integers-with-even-digit-sum.jl" begin
    @test count_even(4) == 2
    @test count_even(30) == 14
end
## @lc test=end

## @lc info=start
# link: [solution to 2180](https://leetcode-cn.com/problems/count-integers-with-even-digit-sum/solution/python-yi-xing-dai-ma-by-rex-rs-1xty/)
## @lc info=end


function fdist!(f::Vector{Float64},x::Vector{Float64},D::Vector{Float64},C::Float64,
    R0::Float64,R1::Float64,Ws::Float64,W1::Float64,rho::Float64,beta::Float64,
    c0::Float64,a1::Float64,a2::Float64,Pn::Float64,h::Float64)
    f[1] = D[1]*(a1/(a2*Ws^5)*(0.125*(W1-Ws*x[1])+c0)^2 - (W1^5/(8192*Ws^5) -
        3*W1^4*x[1]/(8192*Ws^4) + W1^4*c0/(256*Ws^5) + W1^3*x[1]^2/(4096*Ws^3) -
        W1^3*x[1]*c0/(128*Ws^4) + 3*W1^3*c0^2/(64*Ws^5) + W1^2*x[1]^3/(4096*Ws^2) -
        3*W1^2*x[1]*c0^2/(64*Ws^4) + W1^2*c0^3/(4*Ws^5) - 3*W1*x[1]^4/(8192*Ws) +
        W1*x[1]^3*c0/(128*Ws^2) - 3*W1*x[1]^2*c0^2/(64*Ws^3) + W1*c0^4/(2*Ws^5) +
        x[1]^5/8192 - x[1]^4*c0/(256*Ws) + 3*x[1]^3*c0^2/(64*Ws^2) -
        x[1]^2*c0^3/(4*Ws^3) + x[1]*c0^4/(2*Ws^4)) + Pn/(a2*Ws^5));
end


export sub2ind,
        statistics,
        prior,
        gaussian_kernel,
        kernel_density_estimate

"""
    function sub2ind(siz, x)

Convert a length n cartesian coordinate for indexing an n dimensional array into its linear (scalar) equivalent.

# Arguments

* siz::NTuple{Int}
  * Result of calling size(A) on the matrix of interest A.

* x::NTuple{Int}
  * length(siz)-length cartesian index to convert from
"""
function sub2ind(siz, x)
    k = vcat(1, cumprod(siz[1:end-1]))
    return dot(k, x .- 1) + 1
end

"""
    function statistics(vars, G, D::Matrix{Int})

Extract the statistics (ie counts) from a discrete dataset D assuming a Bayesian network with variables vars and structure G.  The dataset is an _n x m_ matrix, where n is the number of variables, and m is the number of data points.

# Returns

Returns an array *M* of length _n_.  The _i_th component consists of a _qᵢ x rᵢ_ matrix of counts.
"""
function statistics(vars, G, D)
    # n::number of variables
    n = size(D, 1)
    # rᵢ::for each variable i, number of discrete values it can take
    r = [vars[i].m for i in 1:n]

    # qᵢ: number of unique value combinations variable i's direct parents can take
    q = [ prod(Int64[r[j] for j in inneighbors(G,i)]) for i in 1:n ]

    # initialize counts matrix
    M = [zeros(q[i], r[i]) for i in 1:n]

    # for each data point (one per column of D)
    for o in eachcol(D)
        # for each variable in each sample
        for i in 1:n
            # actual category observed in this data point
            k = o[i]
            # parents of variable i
            parents = inneighbors(G,i)
            j=1

            # Find linear index of parents
            if !isempty(parents)
                j = sub2ind(r[parents], o[parents])
            end
            #increment mᵢⱼₖ
            M[i][j,k] += 1.0
        end
    end
    return M
end

"""
    function prior(vars, G)

Generate a matrix of Dirichlet priors {αᵢⱼₖ = 1 ∀ i, j, k} corresponding to a uniform prior distribution over the parameters of a bayesian network.

# Arguments

    * vars::Vector{Variables}
        The Variables of the discrete BayesianNetwork
    * G::DAG
        The bayesian directed acyclic graph structure
"""
function prior(vars, G)

    n = length(vars)

    r = [vars[i].m for i in 1:n]

    q = [prod([r[j] for j in inneighbors(G,i)]) for i in 1:n]

    return [ones(q[i], r[i]) for i in 1:n]
end

"""
    gaussian_kernel(b)

Get a scalar probability density function of a zero-mean, b standard deviation gaussian.
"""
gaussian_kernel(b) = x->pdf(Normal(0,b), x)

"""
    function kernel_density_estimate(φ, O)

"""
function kernel_density_estimate(φ, O)
    return x -> sum([φ(x - o) for o in O])/length(O)
end


"""
    getreference(blocks; axis=["X","Y","Z"])

Extract the mid-surface points that split the informed block model in two.
Returns a coordinate matrix with the reference points for unfolding.

## Parameters:

* `blocks` - coordinate matrix of the regular blocks centroids.
* `axis`   - axis or list of axis that will be scanned to extract reference
  points. Defaults to ["X","Y","Z"].
"""
function getreference(blocks::AbstractMatrix; axis=["X","Y","Z"])
	typeof(axis) in (String, Symbol) && (axis = [axis])
	axis[1] isa Symbol && (axis = String.(axis))
	@assert axis ⊆ ["X","Y","Z"] "invalid axis: $axis"
	coord_ids = Dict("X"=>1, "Y"=>2, "Z"=>3)

	# get cell sizes
	cells = zeros(Float64,3)
	for c in 1:3
		coords   = sort!(unique(blocks[c,:]))
		cells[c] = minimum([coords[x]-coords[x-1] for x in 2:length(coords)])
	end

	# scan sections and save skeleton points
	ref = Float64[]
	for c in axis
		ax = coord_ids[c]
	    axis_coords = sort!(unique(blocks[ax,:]))
		sec = setdiff([1,2,3],ax)

	    for s in axis_coords
	        section = view(blocks,:,view(blocks,ax,:) .== s)

	        min_i = minimum(section[sec[1],:])
	        min_j = minimum(section[sec[2],:])
	        i = section[sec[1],:] .- min_i
	        j = section[sec[2],:] .- min_j

			i = Int.(i/cells[sec[1]]) .+ 1
			j = Int.(j/cells[sec[2]]) .+ 1

	        img = zeros(Bool,maximum(i)+1,maximum(j)+1)

	        for x in 1:length(i)
	            img[i[x],j[x]]=true
	        end

	        img = thinning(img)

	        for x in 1:length(i)
	            if img[i[x],j[x]]==true
	                io = (i[x]-1)*cells[sec[1]]+min_i
	                jo = (j[x]-1)*cells[sec[2]]+min_j
					order = sortperm([ax,sec[1],sec[2]])
					append!(ref, [s, io, jo][order])
	            end
	        end
	    end
	end

	reshape(ref, 3, Int(length(ref)/3))
end

# Average duplicate points
function remove_duplicates(coords::AbstractMatrix; tol=0.01)
	groups = [union(x.core_indices,x.boundary_indices) for x in dbscan(coords, tol)]
	outcoords = zeros(Float64,3,length(groups))
	for g in 1:length(groups)
		idx = groups[g]
		if length(idx)==1
			outcoords[:,g] .= dropdims(view(coords,:,idx),dims=2)
		else
			outcoords[:,g] .= mean!(view(outcoords,:,g),view(coords,:,idx))
		end
	end
	outcoords
end

function get_resolution(ref_pts::AbstractMatrix)
	idxs, dists = get_neighbors(ref_pts, :knn, 2)
	closest_pt = [sum(x) for x in dists]
	quantile(closest_pt,0.75)
end


## lattice.jl : unit testing for lattice.jl
#
# This file is part of MultilevelEstimators.jl - A Julia toolbox for
# Multilevel Monte Carlo Methods (c) Pieterjan Robbe, 2019

@testset "LatticeRule32                " begin

lattice_rule = LatticeRule32([0x00000001,0x00000022], 2, 0x00000037)
point = get_point(lattice_rule, 2)
@test lattice_rule isa LatticeRule32
@test length(point) == 2
@test all(0 .≤ point .< 1)

z_file_1 = joinpath(dirname(pathof(MultilevelEstimators)), "generating_vectors", "K_3600_32.txt")
z_file_2 = joinpath(dirname(pathof(MultilevelEstimators)), "generating_vectors", "CKN_250_20.txt")
lattice_rule = LatticeRule32(z_file_1, 101)
point = get_point(lattice_rule, 100_000)
@test lattice_rule isa LatticeRule32
@test length(point) == 101
@test all(0 .≤ point .< 1)

function approx_π(n::Integer)
    lattice_rule = LatticeRule32(z_file_2, 2)
    count = 0
    for i in 0:n-1
        x = get_point(lattice_rule, i)
        if x[1]*x[1] + x[2]*x[2] < 1
            count += 1
        end
    end
    return 4*count/n
end

for (n, ϵ) in zip(2 .^(0:20), exp10.([0 0 0 0 0 0 1 1 2 2 2 3 3 3 3 3 3 4 4 5]))
    @test approx_π(n) ≈ π atol=ϵ
end

@test LatticeRule32(3600) isa LatticeRule32
@test_throws ArgumentError LatticeRule32(3601)
@test LatticeRule32(z_file_2, 250) isa LatticeRule32
@test_throws ArgumentError LatticeRule32(z_file_2, 251)
@test_throws ArgumentError LatticeRule32(-2)
@test LatticeRule32(z_file_2, 1) isa LatticeRule32
@test_throws ArgumentError LatticeRule32(0)
@test_throws ArgumentError LatticeRule32(z_file_1, 24, 4294967297)

end


A = -5:0.2:5
B = -5:5
@testset "Indicator" begin
    @testset "above, nonstrict $y, $b" for y in A, b in B
        if y ≥ b
            @test indicator(y, b) == 1
        else
            @test indicator(y, b) == 0
        end
    end
    @testset "above, strict $y, $b" for y in A, b in B
        if y > b
            @test indicator(y, b, ineq=:strict) == 1
        else
            @test indicator(y, b, ineq=:strict) == 0
        end
    end
    @testset "below, nonstrict $y, $b" for y in A, b in B
        if y ≤ b
            @test indicator(y, b, indregion=:below) == 1
        else
            @test indicator(y, b, indregion=:below) == 0
        end
    end
    @testset "below, strict $y, $b" for y in A, b in B
        if y < b
            @test indicator(y, b, indregion=:below, ineq=:strict) == 1
        else
            @test indicator(y, b, indregion=:below, ineq=:strict) == 0
        end
    end
    @testset "Unknown input error checks" begin
        @test_throws ErrorException("Unknown indicator region: test") indicator(1, 0, indregion=:test)
        @test_throws ErrorException("Unknown inequality type: test") indicator(1, 0, ineq=:test)
    end
end


using Test
using PolynomialBasis
using ImplicitDomainQuadrature
# using Revise
using CutCell
include("useful_routines.jl")

vectosymmconverter = CutCell.vector_to_symmetric_matrix_converter()
E1 = [
    1.0 0.0
    0.0 0.0
    0.0 1.0
]

E2 = [
    0.0 0.0
    0.0 1.0
    1.0 0.0
]

@test allapprox(vectosymmconverter[1], E1)
@test allapprox(vectosymmconverter[2], E2)

matrix = CutCell.make_row_matrix(E1, [1, 2, 3])
expE1 = hcat(E1, 2 * E1, 3 * E1)
@test allapprox(matrix, expE1)

matrix = CutCell.interpolation_matrix([1.0, 2.0, 3.0], 2)
testm = [
    1.0 0.0 2.0 0.0 3.0 0.0
    0.0 1.0 0.0 2.0 0.0 3.0
]
@test allapprox(matrix, testm)

basis = TensorProductBasis(2, 1)
cellmap = CutCell.CellMap([0.0, 0.0], [3.0, 1.0])
grad = CutCell.transform_gradient(gradient(basis, 0.0, 0.0), CutCell.jacobian(cellmap))
@test allapprox(grad[1, :], [-1 / 6, -1 / 2])
@test allapprox(grad[2, :], [-1 / 6, 1 / 2])
@test allapprox(grad[3, :], [1 / 6, -1 / 2])
@test allapprox(grad[4, :], [1 / 6, 1 / 2])

l, m = 1.0, 2.0
stiffness = plane_strain_voigt_hooke_matrix(1.0, 2.0)
cellmap = CutCell.CellMap([0.0, 0.0], [1.0, 1.0])
quad = tensor_product_quadrature(2, 2)
bf = CutCell.bilinear_form(basis, quad, stiffness, CutCell.jacobian(cellmap))
@test size(bf) == (8, 8)

rhsfunc(x) = [1.0, 1.0]
rhs = CutCell.linear_form(rhsfunc, basis, quad, cellmap)
@test allapprox(rhs, 0.25 * ones(8))

basis = TensorProductBasis(2, 2)
quad = tensor_product_quadrature(2, 4)
cellmap = CutCell.CellMap([-1, -1.0], [1.0, 1.0])
M = CutCell.mass_matrix(basis, quad, cellmap, 2)
rhsfunc(x) = [x[1]^2 + x[2]^2, 2x[1] * x[2]]
R = CutCell.linear_form(rhsfunc, basis, quad, cellmap)

sol = reshape(M \ R, 2, :)
coords = hcat([cellmap(basis.points[:, i]) for i = 1:9]...)
exactsol = hcat([rhsfunc(coords[:, i]) for i = 1:9]...)
@test allapprox(sol, exactsol, 1e3eps())

K = CutCell.bilinear_form(basis, quad, stiffness, CutCell.jacobian(cellmap))
bodyforce(x) = -4 * (l + 2m) * [1.0, 0.0]
R = CutCell.linear_form(bodyforce, basis, quad, cellmap)

boundarynodeids = [1, 4, 7, 8, 9, 6, 3, 2]
bcvals = exactsol[:, boundarynodeids]
CutCell.apply_dirichlet_bc!(K, R, boundarynodeids, bcvals)

sol = reshape(K \ R, 2, :)
@test allapprox(sol, exactsol, 1e3eps())

K = CutCell.bilinear_form(basis, quad, stiffness, cellmap)
bodyforce(x) = -4 * (l + 2m) * [1.0, 0.0]
R = CutCell.linear_form(bodyforce, basis, quad, cellmap)

boundarynodeids = [1, 4, 7, 8, 9, 6, 3, 2]
bcvals = exactsol[:, boundarynodeids]
CutCell.apply_dirichlet_bc!(K, R, boundarynodeids, bcvals)

sol = reshape(K \ R, 2, :)
@test allapprox(sol, exactsol, 1e3eps())

polyorder = 1
numqp = 2
basis = TensorProductBasis(2, polyorder)
quad = tensor_product_quadrature(2, numqp)
lambda, mu = 1.0, 2.0
penalty = 1e0
dx = 0.1
e11 = dx / 1.0
e22 = -lambda / (lambda + 2mu) * e11
dy = e22

cellmap = CutCell.CellMap([0.0, 0.0], [1.0, 1.0])
stiffness = plane_strain_voigt_hooke_matrix(lambda, mu)
facequads = CutCell.face_quadratures(numqp)
facedetjac = CutCell.face_determinant_jacobian(cellmap)
bf = CutCell.bilinear_form(basis, quad, stiffness, cellmap)

ftop1 = CutCell.face_traction_component_operator(
    basis,
    facequads[1],
    [0.,1.],
    [0.,-1.],
    stiffness,
    facedetjac[1],
    cellmap,
)
ftop2 = CutCell.face_traction_component_operator(
    basis,
    facequads[2],
    [1.,0.],
    [1.,0.],
    stiffness,
    facedetjac[2],
    cellmap,
)
ftop4 = CutCell.face_traction_component_operator(
    basis,
    facequads[4],
    [1.,0.],
    [-1.0, 0],
    stiffness,
    facedetjac[4],
    cellmap,
)

bcop1 = penalty*CutCell.component_mass_matrix(basis,facequads[1],[0.,1.],facedetjac[1])
bcop2 = penalty*CutCell.component_mass_matrix(basis,facequads[2],[1.,0.],facedetjac[2])
bcop4 = penalty*CutCell.component_mass_matrix(basis,facequads[4],[1.,0.],facedetjac[4])
bcrhs2 = penalty*CutCell.component_linear_form(dx,basis,facequads[2],[1.,0.],facedetjac[2])

K = bf - ftop1 - ftop2 - ftop4 + bcop1 + bcop2 + bcop4
rhs = bcrhs2

sol = K\rhs
disp = reshape(sol,2,:)
testdisp = [0. 0. dx dx
            0. dy 0. dy]
@test allapprox(disp,testdisp,1e3eps())


@testset "Gauging type" begin

    # Test the construction of the couple (level, discharge)
    G = Gauging(0.0,1.0)
    @test G.level ≈ 0.0 atol = sqrt(eps())
    @test G.discharge ≈ 1.0 atol = sqrt(eps())

    # Test the promotion of level and discharge
    G = Gauging(0,1)
    @test typeof(G.level) == Float64
    @test typeof(G.discharge) == Float64

    # Test the promotion of level
    G = Gauging(0,1.1)
    @test typeof(G.level) == Float64
    @test typeof(G.discharge) == Float64

    # Test the empty gauing construction
    G = Gauging()
    @test G.level ≈ 0.0 atol = sqrt(eps())
    @test G.discharge ≈ 0.0 atol = sqrt(eps())

    # Test the broadcast
    G = Gauging.([1,2],[3,4])
    @test G[2].level ≈ 2
    @test G[2].discharge ≈ 4

end

@testset "Gauging methods" begin

    # Test `level()`for a single gauging
    G = Gauging(0.0,1.0)
    @test level(G) ≈ 0.0

    # Test the broadcasting of `level()`
    G = Gauging.([0.0,1.0],[2.0,3.0])
    @test level.(G) ≈ [0.0, 1.0]

    # Test `discharge()`for a single gauging
    G = Gauging(0.0,1.0)
    @test discharge(G) ≈ 1.0

    # Test the broadcasting of `discharge()`
    G = Gauging.([0.0,1.0],[2.0,3.0])
    @test discharge.(G) ≈ [2.0, 3.0]

end

@testset "bootstrap Gauging" begin
    G = Gauging.(1:3, 1:3)

    Gbs = RatingCurves.bootstrap(G)

    @test minimum(G) ∈ Gbs
    @test length(G) == length(Gbs)

end

@testset "crcfit" begin
    G = Gauging.([1/4, 1/3, 1/2, 2,3,4],[1/32, 2/27, 1/4, 8, 18, 32])
    crc = crcfit(G)

    rc₁ = crc.component[1]
    rc₂ = crc.component[2]

    @test crc.threshold[] ≈ 1.0

    @test rc₁.a ≈ 2.0
    @test rc₁.b ≈ 0.0 atol=1e-8
    @test rc₁.c ≈ 3.0

    @test rc₂.a ≈ 2.0
    @test rc₂.b ≈ 0.0 atol=1e-8
    @test rc₂.c ≈ 2.0
end

@testset "crcfit with given k" begin
    G = Gauging.([1/4, 1/3, 1/2, 2,3,4],[1/32, 2/27, 1/4, 8, 18, 32])
    crc = crcfit(G, 1)

    rc₁ = crc.component[1]
    rc₂ = crc.component[2]

    @test crc.threshold[] ≈ 1.0

    @test rc₁.a ≈ 2.0
    @test rc₁.b ≈ 0.0 atol=1e-8
    @test rc₁.c ≈ 3.0

    @test rc₂.a ≈ 2.0
    @test rc₂.b ≈ 0.0 atol=1e-8
    @test rc₂.c ≈ 2.0
end

@testset "getinitialvalues" begin

    G = Gauging.([1,2,3], [4,5,6])
    rc = RatingCurves.getinitialvalues(G)

    @test rc.a ≈ 126.4911 atol = 1e-4
    @test rc.b ≈ .9
    @test rc.c ≈ 3/2

end

@testset "minimum of Gauging" begin
    G = Gauging.(1:3, 1:3)

    @test minimum(G).level ≈ 1.0
    @test minimum(G).discharge ≈ 1.0

end

@testset "maximum of Gauging" begin
    G = Gauging.(1:3, 1:3)

    @test maximum(G).level ≈ 3.0
    @test maximum(G).discharge ≈ 3.0

end

@testset "rcfit with given b" begin
    G = Gauging.([2,3,4],[3,12,27])
    rc = rcfit(G, 1)
    @test rc.a ≈ 3.0
    @test rc.b ≈ 1.0
    @test rc.c ≈ 2.0
end

@testset "rcfit" begin
    G = Gauging.([2,3,4],[3,12,27])
    rc = rcfit(G)
    @test rc.a ≈ 3.0
    @test rc.b ≈ 1.0
    @test rc.c ≈ 2.0
end

@testset "rcfit with given b s.t. constraint" begin
    G = Gauging.([2,3,4],[3,12,27])
    rc = rcfit(G, 1, [5, 48])
    @test rc.a ≈ 3.0
    @test rc.b ≈ 1.0
    @test rc.c ≈ 2.0
end

@testset "rcfit s.t. constraint" begin
    G = Gauging.([2,3,4],[3,12,27])
    rc = rcfit(G, [5, 48])
    @test rc.a ≈ 3.0
    @test rc.b ≈ 1.0
    @test rc.c ≈ 2.0
end

@testset "sort Gauging" begin
    G = Gauging.(1:3, 1:3)

    G_ascend = sort(G)
    @test level.(G_ascend) ≈ [1, 2, 3]

    G_descend = sort(G, rev=true)
    @test level.(G_descend) ≈ [3, 2, 1]

end


module Mazes
using SimpleGraphs, SimplePartitions, SimpleDrawing, Plots

_vtx = Tuple{Int,Int}
_edge = Tuple{_vtx,_vtx}

import SimpleDrawing: draw
import Base: show, size
export Maze, draw_ans, draw, amaze

"""
`M = Maze(r,c)` creates an `r`-by-`c` maze. To see the maze,
use `draw(M)`. The maze is meant to be solved from the upper left
to the lower right. This is indicated by stars at those locations.
To draw the maze without those stars, use `draw(M,false)`.

To draw the solution to the maze, use `draw_ans(M)`. The solutions is
always drawn without stars.
"""
struct Maze
    r::Int  # number of rows
    c::Int  # number of cols
    T::SimpleGraph{_vtx}  # the tree for this maze
    function Maze(nr::Int, nc::Int)
        G = _gen_tree(nr, nc)
        new(nr, nc, G)
    end
end

Maze(n::Int) = Maze(n, n)

function show(io::IO, M::Maze)
    print(io, "Maze($(M.r),$(M.c))")
end

"""
`size(M::Maze)` returns a tuple `(r,c)` where `r` is the number of
rows and `c` is the number of columns in the maze.
"""
size(M::Maze) = (M.r, M.c)

function _gen_tree(nr, nc)
    @assert nr >= 2 && nc >= 2 "Inputs must both be at least 2"
    G = Grid(nr, nc)

    EE = elist(G)  # all possible edges
    # assign random weights to edges
    wt = Dict{_edge,Float64}()
    for e in EE
        wt[e] = rand()
    end

    # the upper left and lower right squares should be leaves
    # these weights should enforce that

    e1 = ((1, 1), (1, 2))
    e2 = ((1, 1), (2, 1))
    if rand() > 0.5
        wt[e1] = 0
        wt[e2] = 1
    else
        wt[e1] = 1
        wt[e2] = 0
    end

    e1 = (nr - 1, nc), (nr, nc)
    e2 = (nr, nc - 1), (nr, nc)
    if rand() > 0.5
        wt[e1] = 0
        wt[e2] = 1
    else
        wt[e1] = 1
        wt[e2] = 0
    end

    wt_list = [wt[e] for e in EE]
    p = sortperm(wt_list)
    E_list = [EE[j] for j in p] # EE sorted by weight

    # copy vertices from G
    T = SimpleGraph{_vtx}()
    for v in G.V
        add!(T, v)
    end
    P = Partition(G.V)

    for e in E_list
        a = e[1]
        b = e[2]
        if !in_same_part(P, a, b)
            add!(T, a, b)
            merge_parts!(P, a, b)
        end
    end

    d = Dict{_vtx,Array{Int,1}}()
    for v in T.V
        x = v[1]
        y = v[2]
        d[v] = [y, -x]
    end
    embed(T, d)
    return T
end

function _puzzle_draw(M::Maze)
    # get the tree's embedding
    xy = M.T.cache[:xy]

    newdraw()
    # draw(M.T)  # DEBUG #

    # draw outside rectangle
    ul = (0.5, -0.5)
    ur = (M.c + 0.5, -0.5)
    ll = (0.5, -M.r - 0.5)
    lr = (M.c + 0.5, -M.r - 0.5)

    draw_segment(ul..., ur..., color = :black)
    draw_segment(ul..., ll..., color = :black)
    draw_segment(ur..., lr..., color = :black)
    draw_segment(ll..., lr..., color = :black)

    G = Grid(M.r, M.c)

    non_edge_list = [e for e in G.E if !has(M.T, e[1], e[2])]

    for e in non_edge_list
        a = e[1]
        b = e[2]
        mid = 0.5 * (xy[a] + xy[b])

        if a[1] == b[1]  # vertical segment
            x = mid[1]
            y1 = mid[2] - 0.5
            y2 = mid[2] + 0.5
            draw_segment(x, y1, x, y2, color = :black)
        else
            x1 = mid[1] - 0.5
            x2 = mid[1] + 0.5
            y = mid[2]
            draw_segment(x1, y, x2, y, color = :black)
        end

    end
end

"""
`draw(M::Maze)` draws the maze `M` on the screen. Circles mark the start
(upper left) and end (lower right). The size is printed above the maze.

Full syntax: `draw(M::Maze, markers::Bool=true, title::Bool=true)`
"""
function draw(M::Maze, markers::Bool = true, title::Bool = true)
    _puzzle_draw(M)

    if markers
        r = 0.125
        draw_circle(1, -1, r, color = :black)
        draw_circle(M.c, -M.r, r, color = :black)
    end

    if title
        plot!(title = "$(M.r)-by-$(M.c)")
    end

    finish()
end

function _ans_draw(M::Maze, s::_vtx, t::_vtx)
    xy = M.T.cache[:xy]
    @assert has(M.T, s) "$s not a square in this maze"
    @assert has(M.T, t) "$t not a square in this maze"
    if s == t
        @warn "Two locations are the same, no path will be drawn"
    end
    P = find_path(M.T, s, t)
    n = length(P)
    for i = 1:n-1
        a = P[i]
        xa, ya = xy[a]
        b = P[i+1]
        xb, yb = xy[b]
        draw_segment(xa, ya, xb, yb, color = :red, style = :dot)
    end
end

"""
`draw_ans(M::Maze)` draws the maze `M` and the solution (as a dotted red line)
from the start (upper left) to the end (lower right).

To draw a different path (solution) use `draw(M,s,t)` where `s` and `t` are
integer 2-tuples. Note that squares in the maze are indexed like a matrix,
so `(2,3)` refers to the cell in row 2, column 3.
"""
function draw_ans(M::Maze, s::_vtx, t::_vtx)
    _puzzle_draw(M)
    _ans_draw(M, s, t)
    finish()
end

draw_ans(M::Maze) = draw_ans(M, (1, 1), size(M))

"""
`amaze(r,c)` creates and draws a maze.
"""
function amaze(r::Int, c::Int)
    M = Maze(r, c)
    draw(M)
end

end # end of module


function distplot1D(df, key; quantiles=[0.25, 0.50, 0.75], kwargs...)
    p = @df df stephist(cols(key), fill=true; kwargs...) # Histogram
    qs = @df df quantile(cols(key), quantiles) # Quantiles
    vline!(p, qs)
    return p
end

function distplot2D(df, k1, k2; kwargs...)
    dens = @df df kde((cols(Symbol(k1)), cols(Symbol(k2))))
    plot(dens; fill=true, levels=5, kwargs...)
end

function corner(
        df,
        vars=nothing;
        quantiles=[0.25,0.50,0.75],
        bandwidthx=100,
        bandwidthy=100,
        kwargs...,
    )

    # valid variables
    validvars = propertynames(df)

    # plot all variables by default
    isnothing(vars) && (vars = validvars)
    @assert vars ⊆ validvars "invalid variable name"

    plts = []
    n = length(vars)
    for i in 1:n, j in 1:n
        # Only label axes of first column and last row or plots
        xaxis = i == n
        yaxis = j == 1

        xticks = i == n
        xguide = i == n ? vars[j] : ""
        yticks = i > 1 && j == 1
        yguide = i > 1 && j == 1 ? vars[i] : ""

        # 1D histograms
        if i == j
            #p = plot(xticks=xticks, yticks=yticks)
            p = distplot1D(
                df,
                vars[i];
                quantiles = quantiles,
                xaxis = xaxis,
                yaxis = yaxis,
                xrotation = 45,
                yrotation = 45,
                title = vars[i],
                xticks = xticks,
                yticks = yticks,
                #xguide=xguide,
                #yguide=yguide,
            )
        # 2D histograms
        elseif i > j
            #p = plot(xticks=xticks, yticks=yticks)
            p = distplot2D(
                df,
                vars[j],
                vars[i],
                quantiles = quantiles,
                xaxis = xaxis,
                yaxis = yaxis,
                xrotation = 45,
                yrotation = 45,
                xticks = xticks,
                yticks = yticks,
                #xguide=xguide,
                #yguide=yguide,
            )
        else
            p = plot(framestyle=:none)
        end

        push!(plts, p)
    end

    return plot(plts...; layout=(n, n), leg=:none, colorbar=:none, kwargs...)
end


function build_linear_formula(factors::Array{Symbol, 1})
    linear_formula = Expr(:call)
    linear_formula.args = vcat(:+, factors)
    return Formula(:y, linear_formula)
end

function get_model_variables(formula::Formula)
    variables = Array{Any, 1}()
    push!(variables, Terms(formula).terms...)
    return variables
end

"""
    scale_orthogonal!(design::Array{Float64, 2}, factors::Array{T, 1}) where T <: Any

Orthogonally scale and center factors of a design using design and factor
limits.

# Examples

```jldoctest
julia> using ExperimentalDesign, DataStructures, StatsModels

julia> A = [5. 2. -1.; 5. 3. 0.; -5. 1. -2.]
3×3 Array{Float64,2}:
  5.0  2.0  -1.0
  5.0  3.0   0.0
 -5.0  1.0  -2.0

julia> factors = OrderedDict([(:f1, [-5., 0., 5.]), (:f2, [1., 2., 3.]), (:f3, [-2., -1., 0.])])
DataStructures.OrderedDict{Symbol,Array{Float64,1}} with 3 entries:
  :f1 => [-5.0, 0.0, 5.0]
  :f2 => [1.0, 2.0, 3.0]
  :f3 => [-2.0, -1.0, 0.0]

julia> scale_orthogonal!(A, collect(values(factors)))
3×3 Array{Float64,2}:
  1.0   0.0   0.0
  1.0   1.0   1.0
 -1.0  -1.0  -1.0

julia> A
3×3 Array{Float64,2}:
  1.0   0.0   0.0
  1.0   1.0   1.0
 -1.0  -1.0  -1.0

```

# Formula

For a design ``D_{n,p}`` with ``n`` experiments or rows and ``p`` factors
or columns, scale each factor ``\\mathbf{x}_i`` to ``\\mathbf{x}_{i}^{s}``
according to:

```math
\\mathbf{x}_{i}^{s} = \\dfrac{\\mathbf{x}_i - \\bar{M}}{\\bar{M}_{def}}
```

Where ``\\mathbf{x}_{def}`` is the factor defined in the `factors` parameter and:

```math
\\bar{M} = (max(\\mathbf{x}_i) + min(\\mathbf{x}_i)) / 2
```

```math
\\bar{M}_{def} = (max(\\mathbf{x}_{def}) - min(\\mathbf{x}_{def})) / 2
```
"""
function scale_orthogonal!(design::Array{Float64, 2},
                           factors::Array{T, 1}) where T <: Any
    for i = 1:size(design, 2)
        design_range = (max(design[:, i]...) + min(design[:, i]...)) / 2
        factor_range = (max(factors[i]...) - min(factors[i]...)) / 2
        design[:, i] = (design[:, i] .- design_range) ./ factor_range
    end

    return design
end

"""
    scale_boxdraper_encoding!(design::Array{Float64, 2},
                              factors::Array{T, 1};
                              scale_denominator = true) where T <: Any

Scale factors of a design using the Box and Draper's coding convention from
"Response Surfaces, Mixtures, and Ridge Analyses".

# Examples

With denominator scaling:

```jldoctest
julia> using ExperimentalDesign, DataStructures, StatsModels

julia> A = float(plackett_burman(4))
4×3 Array{Float64,2}:
  1.0   1.0   1.0
  1.0  -1.0  -1.0
 -1.0  -1.0   1.0
 -1.0   1.0  -1.0

julia> factors = OrderedDict([(:f1, [-1., 1.]), (:f2, [-1., 1.]), (:f3, [-1., 1.])])
DataStructures.OrderedDict{Symbol,Array{Float64,1}} with 3 entries:
  :f1 => [-1.0, 1.0]
  :f2 => [-1.0, 1.0]
  :f3 => [-1.0, 1.0]

julia> scale_boxdraper_encoding!(A, collect(values(factors)))
4×3 Array{Float64,2}:
  1.0   1.0   1.0
  1.0  -1.0  -1.0
 -1.0  -1.0   1.0
 -1.0   1.0  -1.0

julia> all([isapprox(1.0, sqrt(sum(A[:, i] .^ 2.0) / size(A, 1))) for i = 1:size(A, 2)])
true

```

Without denominator scaling:

```jldoctest
julia> using ExperimentalDesign, DataStructures, StatsModels

julia> A = float(plackett_burman(4))
4×3 Array{Float64,2}:
  1.0   1.0   1.0
  1.0  -1.0  -1.0
 -1.0  -1.0   1.0
 -1.0   1.0  -1.0

julia> factors = OrderedDict([(:f1, [-1., 1.]), (:f2, [-1., 1.]), (:f3, [-1., 1.])])
DataStructures.OrderedDict{Symbol,Array{Float64,1}} with 3 entries:
  :f1 => [-1.0, 1.0]
  :f2 => [-1.0, 1.0]
  :f3 => [-1.0, 1.0]

julia> scale_boxdraper_encoding!(A, collect(values(factors)), scale_denominator = false)
4×3 Array{Float64,2}:
  0.5   0.5   0.5
  0.5  -0.5  -0.5
 -0.5  -0.5   0.5
 -0.5   0.5  -0.5

julia> all([isapprox(1.0, sqrt(sum(A[:, i] .^ 2.0))) for i = 1:size(A, 2)])
true

```

# Formula

For a design ``D_{n,p}`` with ``n`` experiments or rows and ``p`` factors
or columns, scale each factor ``\\mathbf{x}_i`` to ``\\mathbf{x}_{i}^{s}``
according to:

```math
\\mathbf{x}_{i}^{s} = \\dfrac{\\mathbf{x}_i - \\bar{\\mathbf{x}}_i}{S_i}
```

Where ``\\bar{\\mathbf{x}}_i`` is mean of factor definition values in the
`factors` parameter and:

```math
S_{i}^{2} = \\dfrac{1}{n} \\sum\\limits_{j = 1}^{n}{(x_{ij} - \\bar{\\mathbf{x}}_i)^{2}}
```

If we pass `scale_denominator = false`, ``S_i`` becomes:

```math
S_{i}^{2} = \\sum\\limits_{j = 1}^{n}{(x_{ij} - \\bar{\\mathbf{x}}_i)^{2}}
```
"""
function scale_boxdraper_encoding!(design::Array{Float64, 2},
                                   factors::Array{T, 1};
                                   scale_denominator = true) where T <: Any
    for i = 1:size(design, 2)
        factor_mean = mean(factors[i])
        denominator = sum((design[:, i] .- factor_mean) .^ 2.0)

        if scale_denominator
            denominator /= size(design, 1)
        end

        denominator = sqrt(denominator)

        numerator = design[:, i] .- factor_mean

        if !iszero(denominator) && !iszero(numerator)
            numerator ./= denominator
        end

        design[:, i] = numerator
    end

    return design
end

"""
    get_expanded_values(factors::OrderedDict)

Return a mapping from an original level of a factor to a column
in an expanded design. Does not expand factors.

# Examples

```jldoctest
julia> using ExperimentalDesign, DataStructures, StatsModels

julia> A = OrderedDict(:A => [1., 2.], :B => [1, 2, 3, 4], :C => ["A", "B"], :D => [1.2, 2.2, 3.4], :E => [:a, :b, :c, :d])
DataStructures.OrderedDict{Symbol,Any} with 5 entries:
  :A => [1.0, 2.0]
  :B => [1, 2, 3, 4]
  :C => String["A", "B"]
  :D => [1.2, 2.2, 3.4]
  :E => Symbol[:a, :b, :c, :d]

julia> get_expanded_values(A)
DataStructures.OrderedDict{Any,Any} with 2 entries:
  :C => DataStructures.OrderedDict{Any,Symbol}("B"=>:C_1)
  :E => DataStructures.OrderedDict{Any,Symbol}(:b=>:E_1,:c=>:E_2,:d=>:E_3)

```
"""
function get_expanded_values(factors::OrderedDict)
    expanded_values = OrderedDict()
    for factor in factors
        if !(eltype(factor[2]) <: Real)
            expanded_values[factor[1]] = OrderedDict{Any, Symbol}()
            for l = 2:length(factor[2])
                expanded_values[factor[1]][factor[2][l]] = Symbol(factor[1], "_$(l - 1)")
            end
        end
    end
    return expanded_values
end

function expand_design(design::Array{T, 2}, factors::OrderedDict) where T <: Any
    small_design = DataFrame(design)
    rename!(small_design, OrderedDict(zip(names(small_design), keys(factors))))

    exp_values = get_expanded_values(factors)
    expanded_design = DataFrame()

    for name in names(small_design)
        if name in keys(exp_values)
            for key in keys(exp_values[name])
                expanded_design[exp_values[name][key]] = [small_design[i, name] == key ? 1.0 : 0. for i = 1:length(small_design[name])]
            end
        else
            expanded_design[name] = small_design[name]
        end
    end

    return expanded_design
end

"""
    expand_factors(factors::OrderedDict)

Return an `OrderedDict` with true factors expanded to `[0., 1.]` 2-level numerical factors.

For a true factor with ``n`` levels, creates ``n - 1`` 2-level factors. Only
one of these new factors can be at level `1.` at each row of a design, encoding
each level of the original factor. The first level of the original factor is
encoded by all new factors being at level `0.`.

# Examples
```jldoctest
julia> using ExperimentalDesign, DataStructures, StatsModels

julia> A = OrderedDict(:A => [1., 2.], :B => [1, 2, 3, 4], :C => ["A", "B"], :D => [1.2, 2.2, 3.4], :E => [:a, :b, :c, :d])
DataStructures.OrderedDict{Symbol,Any} with 5 entries:
  :A => [1.0, 2.0]
  :B => [1, 2, 3, 4]
  :C => String["A", "B"]
  :D => [1.2, 2.2, 3.4]
  :E => Symbol[:a, :b, :c, :d]

julia> expand_factors(A)
DataStructures.OrderedDict{Symbol,Any} with 7 entries:
  :A   => [1.0, 2.0]
  :B   => [1, 2, 3, 4]
  :C_1 => [0.0, 1.0]
  :D   => [1.2, 2.2, 3.4]
  :E_1 => [0.0, 1.0]
  :E_2 => [0.0, 1.0]
  :E_3 => [0.0, 1.0]

```
"""
function expand_factors(factors::OrderedDict)
    expanded_factors = OrderedDict{Symbol, Any}()
    for factor in factors
        if !(eltype(factor[2]) <: Real)
            for l = 1:(length(factor[2]) - 1)
                expanded_factors[Symbol(factor[1], "_$l")] = [0., 1.]
            end
        else
            expanded_factors[factor[1]] = factor[2]
        end
    end
    return expanded_factors
end

"""
    generate_model_matrix(formula::Formula,
                          design::Array{Float64, 2},
                          factors::OrderedDict;
                          scale::Function = scale_boxdraper_encoding!)

Generate a `DataFrame` with a scaled model matrix for a given formula, design and factors.

Assumes that `formula` is a linear relationship between all the factors in `factors`.

# Examples

```jldoctest
julia> using ExperimentalDesign, DataStructures, StatsModels

julia> A = plackett_burman(4)
4×3 Array{Int64,2}:
  1   1   1
  1  -1  -1
 -1  -1   1
 -1   1  -1

julia> factors = OrderedDict([(:f1, [-1., 1.]), (:f2, [-1., 1.]), (:f3, [-1., 1.])])
DataStructures.OrderedDict{Symbol,Array{Float64,1}} with 3 entries:
  :f1 => [-1.0, 1.0]
  :f2 => [-1.0, 1.0]
  :f3 => [-1.0, 1.0]

julia> m = generate_model_matrix(@formula(y ~ f1 + f2 + f3), float(A), factors)
4×4 Array{Float64,2}:
 1.0   1.0   1.0   1.0
 1.0   1.0  -1.0  -1.0
 1.0  -1.0  -1.0   1.0
 1.0  -1.0   1.0  -1.0

```
"""
function generate_model_matrix(formula::Formula,
                               input_design::DataFrame,
                               factors::OrderedDict;
                               scale::Function = scale_boxdraper_encoding!)
    design = deepcopy(input_design)
    contrasts = []

    for factor in keys(factors)
        if !(eltype(factors[factor]) <: Real)
            push!(contrasts,
                  (factor, ContrastsMatrix(DummyCoding(), factors[factor])))
        end
    end

    design[:y] = ones(size(design, 1))

    model_frame = ModelFrame(formula, design, contrasts = Dict(contrasts))

    return ModelMatrix(model_frame).m
end


include("../src/interface.jl")
include("../src/henon_heiles.jl")
using LinearAlgebra
using Plots
using Plots.PlotMeasures
function getmindif(tab::Vector{Vector{BigFloat}})
    nmmin = Inf
    ret = (0, 0)
    for i = 1:size(tab, 1)
        for j = (i+1):size(tab, 1)
            nm = norm(tab[i] - tab[j], Inf)
            if nm < nmmin
                nmmin = nm
                ret = i, j
            end
        end
    end
    return ret
end
function fctmain(n_tau, prec)
    setprecision(prec)
    u0 = BigFloat.([90, -44, 83, 13] // 100)
    t_max = big"1.0"
    epsilon = big"0.0017"
    println("epsilon=$epsilon")
    nbmaxtest = 12
    ordmax = 17
    debord = 3
    pasord = 2
    y = ones(Float64, nbmaxtest, div(ordmax - debord, pasord) + 1)
    res_gen = Array{Array{BigFloat,1},2}(undef, nbmaxtest, div(ordmax - debord, pasord) + 1)
    x = zeros(Float64, nbmaxtest)
    A = [0 0 1 0; 0 0 0 0; -1 0 0 0; 0 0 0 0]
    prob = HiOscODEProblem(henon_heiles, u0, (big"0.0", t_max), missing, A, epsilon)
    tabsol = Array{Array{BigFloat,1},1}(undef, 0)
    println("ordmax=$ordmax")
    indref = 1
    ind = 1
    for order = debord:pasord:ordmax
        println("eps=$epsilon order=$order")
        nb = 100
        indc = 1
        labels = Array{String,2}(undef, 1, div(order - debord, pasord) + 1)
        resnorm = 0
        resnormprec = 1
        sol = undef
        println("preparation ordre $order + 2")
        par_u0 = missing
        borindc = (order < ordmax) ? nbmaxtest : nbmaxtest + 1
        while indc <= borindc
            if indc > nbmaxtest
                nb = 51 * 2^nbmaxtest
            end
            res = solve(
                prob,
                nb_tau = n_tau,
                order = order,
                nb_t = nb,
                par_u0 = par_u0,
                dense = false,
            )
            par_u0 = res.par_u0
            sol = res[end]
            push!(tabsol, sol)
            if indc <= nbmaxtest
                res_gen[indc, ind] = sol
                x[indc] = 1.0 / nb
            end
            (a, b) = getmindif(tabsol)
            if a != 0
                indref = a
                solref = tabsol[b]
                for i = 1:ind
                    borne = (i < ind) ? size(res_gen, 1) : min(indc, size(res_gen, 1))
                    for j = 1:borne
                        nm = norm(res_gen[j, i] - solref, Inf)
                        if nm == 0
                            nm = norm(res_gen[j, i] - tabsol[a], Inf)
                        end
                        y[j, i] = (nm < 1) ? nm : NaN
                    end
                end
            end
            println("result=$y")
            println("log2(y)=$(log2.(y))")
            nb *= 2
            indc += 1
        end
        for i = debord:pasord:order
            labels[1, (i-debord)÷pasord+1] = " order = $i "
        end
        p = Plots.plot(
            x,
            view(y, :, 1:ind),
            xlabel = "Δt",
            xaxis = :log,
            ylabel = "error",
            yaxis = :log,
            legend = :bottomright,
            label = labels,
            marker = 2,
        )
        pp = Plots.plot(
            x,
            view(y, :, 1:ind),
            xlabel = "Δt",
            xaxis = :log,
            ylabel = "error",
            yaxis = :log,
            legend = :bottomright,
            label = labels,
            marker = 2,
            bottom_margin = 30px,
            left_margin = 60px,
        )
        prec_v = precision(BigFloat)
        eps_v = convert(Float32, epsilon)
        Plots.savefig(p, "out/res4p_$(prec_v)_$(eps_v)_$(order)_$(n_tau)_henon_heiles.pdf")
        Plots.savefig(p, "out/res4p_$(prec_v)_$(eps_v)_$(order)_$(n_tau)_henon_heiles.png")
        Plots.savefig(
            pp,
            "out/res4pp_$(prec_v)_$(eps_v)_$(order)_$(n_tau)_henon_heiles.pdf",
        )
        Plots.savefig(
            pp,
            "out/res4pp_$(prec_v)_$(eps_v)_$(order)_$(n_tau)_henon_heiles.png",
        )
        ind += 1
    end
end

# testODESolver()

# for i=3:9
#     fctMain(2^i)
# end
fctmain(32, 512)


"""
    localmaxmin!(x, maxes, mins)

Detect the local extrema of x.
Push the maxima into maxes and the minima into mins.
"""
function localmaxmin!(y, maxes::Vector{Int}, mins::Vector{Int})
    empty!(maxes)
    empty!(mins)
    for i in 2:length(y)-1
        if y[i+1] < y[i] > y[i-1]
            push!(maxes,i)
        elseif y[i+1] > y[i] < y[i-1]
            push!(mins, i)
        end
    end
end

"""
    zerocrossing!(y, crosses)

Compute the indices of zero crossings of vector `y`.
It searches for elements which are either zero or near a signflip and pushes the
indices into crosses.
"""
function zerocrossing!(y, crosses)
    empty!(crosses)
    for i ∈ 1:(length(y)-1)
        if y[i] == zero(y[i]) || sign(y[i]) * sign(y[i+1]) == -1
            push!(crosses, i)
        end
    end
end

"""
    get_edgepoint(y, xvec, extremas, pos, comp)

Compute the edgepoint which should be used as the extrema on the edge for the
spline computation.
"""
function get_edgepoint(y, xvec, extremas, pos, comp)
    if pos == first
        index = [1,2]
    elseif pos == last
        index = [length(extremas) - 1, length(extremas)]
    end
    # the x values must be embedded into a tuple
    knots = (xvec[extremas[index]],)
    itp = Interpolations.interpolate(knots, y[extremas[index]], Gridded(Linear()))
    expf = extrapolate(itp, Line())
    edgepoint = expf(pos(xvec))
    if comp(edgepoint, pos(y))
        edgepoint
    else
        edgepoint = pos(y)
    end
end

"""
    ismonotonic(x::AbstractVector)

Check whether `x` is monotonic.
"""
ismonotonic(x::AbstractVector) = issorted(x) || issorted(x, rev=true)

abstract type InterpMethod end
struct DierckXInterp <: InterpMethod end
#immutable InterpolationsInterp <: InterpMethod end

function interpolate(knotxvals::Vector, knotyvals::Vector, predictxvals::AbstractVector, m::DierckXInterp, k=3)
    spl = Dierckx.Spline1D(knotxvals, knotyvals, k=k)
    #@show spl, predictxvals
    Dierckx.evaluate(spl,predictxvals)
end

struct HaltingIterable{I, F}
    iter::I
    fun::F
end

function iterate(iter::HaltingIterable)
    next = iterate(iter.iter)
    return dispatch(iter, next)
end

function iterate(iter::HaltingIterable, (instruction, state))
    if instruction == :halt return nothing end
    next = iterate(iter.iter, state)
    return dispatch(iter, next)
end

function dispatch(iter::HaltingIterable, next)
    if next === nothing return nothing end
    return next[1], (iter.fun(next[1]) ? :halt : :continue, next[2])
end

halt(iter::I, fun::F) where {I,F} = HaltingIterable{I,F}(iter, fun)


# Contains the rocket set up and structs

#=
Now we want to set up the rocket dynamics. We have a linear system:
xk+1 = Axk + Buk + G

where xk = [sk; vk] (the position and velocity, respectively). If sk is n×1,
A is 2n×2n, B is 2n×n, and G is a 2n×1 vector.

Thus, we get the system:
[A B -I] * [xk; uk; xk+1] = G
which we can stack as

[A B -I 0 0 0 0 0 0]
[0 0 A B -I 0 0 0 0]
[0 0 0 0 A B -I 0 0]
[0 0 0 0 0 0 A B -I]
=#

using SparseArrays, LinearAlgebra

# include("trajectory-setup.jl")

"""
    rocket_simple(mass, isp, grav, deltaTime)

Barebones "spherical" rocket. This is to say, this struct only holds the mass
and the isp (specific impulse) of the rocket.
"""
struct rocket_simple
    mass
    isp
    grav
    deltaTime
end

@doc raw"""
    rocketDynamics(r::rocket_simple, nDim::Int64)

We have a linear dynamics system, that when continuous is
```math
\frac{dx}{dt} = Ax + Bu + g
```
If we have `y = [x; u; g]`, then,
```math
\frac{dy}{dt} = [A B I; 0 0 0; 0 0 0] y = W y
```
which has solution
```math
y(t) = e^{W t} y0
```
where W is square. In a discrete system,
```math
yk+1 = e^{W Δt} yk = E yk
```

Note that due to all the zeros in W, we get:
```math
xk+1 = E11 xk + E12 uk + E13 g
```
```math
uk+1 = E21 xk + E22 uk + E23 g = uk
```
```math
g = E31 xk + E32 uk + E33 g = g
```

Thus, once we calculate E, we extract E11, E12, and E13 to a system of the
form:
```math
x_{k+1} = Aₑ x_k + Bₑ u_k + Gₑ
```
We compute and return Aₑ, Bₑ, and Gₑ for a [`rocket_simple`](@ref).

For a simple rocket, we have

```math
A = [0 \ I; 0 \ 0] \quad B = [0; \frac{1}{m} I] \quad G = [0; -g]
```

We then take the exponent and return a resulting sparse matrices

"""
function rocketDynamics(r::rocket_simple, nDim::Int64)

    # Start off with a dense representation bc it is needed for the
    # exponentiation.
    z0 = zeros(nDim, nDim)
    aMat = [z0 I; z0 z0]
    bMat = [z0; -(1/r.mass) * I]
    gMat = [z0; I]
    # gVec = [spzeros(nDim, 1); r.grav]

    # Top Row
    abg = [aMat bMat gMat]
    abg_size = size(abg)

    # We need to pad the remaining rows to make a square matrix
    remainingRows = abg_size[2] - abg_size[1]
    abgTot = [abg; zeros(remainingRows, abg_size[2])]

    # Take the exponent e^(W Δt)
    eMat = exp(abgTot * r.deltaTime)

    # Extract the effective matrices
    rowsA = size(aMat, 1)
    colsA = size(aMat, 2)
    colsB = size(bMat, 2)

    aEffMat = sparse(eMat[1:rowsA, 1:colsA])
    bEffMat = sparse(eMat[1:rowsA, colsA + 1:colsA + colsB])
    gEffMat = sparse(eMat[1:rowsA, colsA + colsB + 1:end])

    return aEffMat, bEffMat, gEffMat
end

@doc raw"""
    rocketDynamicsStack(r::rocket_simple, nDim::Int64, NSteps::Int64)

We want to take the per discretization dynamics and write it as a full matrix.

Specifically, we get the system:
`[A B -I] * [xk; uk; xk+1] = G`
which we can stack as

```math
[A \ B \ -I \ 0 \ 0 \ 0 \ 0 \ 0 \ 0]
```

```math
[0 \ 0 \ A \ B \ -I \ 0 \ 0 \ 0 \ 0]
```

```math
[0 \ 0 \ 0 \ 0 \ A \ B \ -I \ 0 \ 0]
```

```math
[0 \ 0 \ 0 \ 0 \ 0 \ 0 \ A \ B \ -I]
```

We return the above matrix and the stacked `G` matrix.

See: [`rocketDynamicsFull`](@ref) and [`rocketDynamics`](@ref)
"""
function rocketDynamicsStack(r::rocket_simple, nDim::Int64, NSteps::Int64)
    # First calculate the A and B matrices. The C vector is not needed in this
    # case.
    Ak, Bk, Gk = rocketDynamics(r::rocket_simple, nDim::Int64)
    ABI_unit = [Ak Bk -I]
    colLen = size(ABI_unit, 2)

    # At each step we go down 2 * n rows and move right 3 * n columns.
    sizeCols = 3 * NSteps * nDim + 2 * nDim
    sizeRows = 2 * nDim * NSteps

    # Initialize and empty sparse A matrix and G vector
    AStacked = spzeros(sizeRows, sizeCols)
    GStacked = spzeros(sizeRows, 1)

    rStart = 1
    cStart = 1
    for k in 1:NSteps
        rEnd = rStart + (2 * nDim - 1)
        cEnd = cStart + colLen - 1

        # First, we calculate that stacked A Matrix
        # println("Accessing $rStart:$rEnd and $cStart:$cEnd")
        AStacked[rStart:rEnd, cStart:cEnd] = ABI_unit

        # Second, we calculate the stacked G vector
        GStacked[rStart:rEnd] = -Gk * r.grav

        # Prepare for next iteration.
        rStart += 2 * nDim
        cStart += 3 * nDim
    end

    return AStacked, GStacked
end

"""
    rocketDynamicsFull(r::rocket_simple, x0::Array{Float64, 1},
                             xN::Array{Float64, 1}, NSteps::Int64,
                             noEnd::Bool = false)

We combine the stack at each `xk` with the initial and final conditions. Namely,

`[I0; AStack; IN] * XU = [x0; GStack; xN]`

where `XU` is the stacked trajectory vector (x states and u controls). We
return `[I0; AStack; IN]` and `[x0; GStack; xN]`

See [`rocketDynamicsStack`](@ref), [`rocketDynamics`](@ref),
[`rocket_simple`](@ref)
"""
function rocketDynamicsFull(r::rocket_simple, x0::Array{Float64, 1},
                                 xN::Array{Float64, 1}, NSteps::Int64,
                                 noEnd::Bool = false)
     if size(x0) != size(xN)
         error("x0 and xN must have the same size. " *
                 "Currently: x0 = $x0 and xN = $xN")
     end
     if mod(size(x0, 1), 2) != 0
         error("x0 and xN must have an even number of variables. " *
               "Example, x0 = [sx0; sy0; vx0; vy0]. Currently x0 = $x0.")
     end
    # The number of columns must be the same for all three. (Stacked vertically)
    nDim = Int64(size(x0, 1) / 2)
    sizeCols = 3 * NSteps * nDim + 2 * nDim
    sizeRowsI = 2 * nDim

    # Use the above function to determine the stacked matrix in the middle.
    AStacked, GStacked = rocketDynamicsStack(r, nDim, NSteps)
    I0 = [I spzeros(sizeRowsI, sizeCols - sizeRowsI)]

    if noEnd
        AFull = [I0; AStacked]
        BFull = [x0; GStacked]
    else
        IN = [spzeros(sizeRowsI, sizeCols - sizeRowsI) I]

        AFull = [I0; AStacked; IN]
        BFull = [x0; GStacked; xN]
    end

    return AFull, BFull
end

function rocketDynamicsFull(r::rocket_simple, x0::Array{Int64, 1},
                                 xN::Array{Int64, 1}, NSteps::Int64)
    return rocketDynamicsFull(r, float.(x0), float.(xN), NSteps)
end


# Function for running simplified manifold Metropolis adjusted Langevin
# algorithm (SMMALA)
function smmala(model::Model, opts::SmmalaOpts)
  mcmc = Array(Float64, opts.mcmc.nPostBurnin, model.nPars)
  z = Array(Float64, opts.mcmc.nPostBurnin, model.nPars)
  proposed, accepted = 0., 0.

  println("Running burn-in iterations...")

  currentPars = model.randPrior()
  currentLogPosterior = model.logPosterior(currentPars)
  currentGradLogPosterior = model.gradLogPosterior(currentPars)
  currentG = model.tensor(currentPars)
  currentInvG = inv(currentG)
  currentFirstTerm = currentInvG*currentGradLogPosterior
  driftStep = opts.setDriftStep(1, 0., opts.mcmc.n, opts.mcmc.nBurnin, 0.)

  for i = 1:opts.mcmc.n
    proposed += 1

    proposedPars = copy(currentPars)
    parameterMeans = proposedPars+(driftStep/2)*currentFirstTerm
    cholCurrentInvG = chol(driftStep*currentInvG)
    proposedPars = parameterMeans+cholCurrentInvG'*randn(model.nPars)
    proposedLogPosterior = model.logPosterior(proposedPars)
    probNewGivenOld = (-sum(log(diag(cholCurrentInvG)))
      -(0.5*(parameterMeans-proposedPars)'
      *(currentG/driftStep)*(parameterMeans-proposedPars))[1])
    proposedG = model.tensor(proposedPars)
    proposedInvG = inv(proposedG)
    proposedGradLogPosterior = model.gradLogPosterior(proposedPars)
    proposedFirstTerm = proposedG\proposedGradLogPosterior
    parameterMeans = proposedPars+(driftStep/2)*proposedFirstTerm
    probOldGivenNew =
      (-sum(log(diag(chol(driftStep*eye(model.nPars)*proposedInvG))))
      -(0.5*(parameterMeans-currentPars)'
      *(proposedG/driftStep)*(parameterMeans-currentPars))[1])

   ratio =
     proposedLogPosterior+probOldGivenNew-currentLogPosterior-probNewGivenOld

    if ratio > 0 || (ratio > log(rand()))
      accepted += 1

      currentPars = copy(proposedPars)
      currentLogPosterior = copy(proposedLogPosterior)
      currentGradLogPosterior = copy(proposedGradLogPosterior)
      currentG = copy(proposedG)
      currentInvG = copy(proposedInvG)
      currentFirstTerm  = copy(proposedFirstTerm)
    end

    if i > opts.mcmc.nBurnin
      mcmc[i-opts.mcmc.nBurnin, :] = currentPars;
      z[i-opts.mcmc.nBurnin, :] = -currentGradLogPosterior/2;
    end

    if mod(i, opts.mcmc.monitorRate) == 0
      acceptanceRatio = accepted/proposed

    driftStep = opts.setDriftStep(i, acceptanceRatio, opts.mcmc.n,
      opts.mcmc.nBurnin, driftStep)

      println("Iteration $i of $(opts.mcmc.n): ", round(100*acceptanceRatio, 2),
        " % acceptance ratio")

      proposed, accepted = 0., 0.
    end

    if i == opts.mcmc.nBurnin
      println("Burn-in completed...\n\nRunning post burn-in MCMC...");
    end
  end

  return mcmc, z
end


using Combinatorics

# Part 1
data = parse.(Int64, readlines("day09_input.txt"))
valid = [d ∈ Set(sum.(combinations(data[i-25:i-1], 2))) for (i, d) in enumerate(data) if i > 25]
target = data[argmin(valid)+25]
println("Part 1: first invalid = $target")

# Part 2
function find_weakness(data, target)
    for start in 1:lastindex(data)
        total = data[start]
        for finish in (start+1):lastindex(data)
            total += data[finish]
            if total == target return minimum(data[start:finish]) + maximum(data[start:finish]) end
            if total > target break end
        end
    end
end
println("Part 2: weakness = $(find_weakness(data, target))")


@testset "PointCoud transforms" begin

    for inplace in [true, false]
        @testset "ScalePointCloud inplace=$(inplace)" begin
            p = rand(Float32, 3, 8, 2) |> gpu
            t1 = ScalePointCloud(2.0; inplace = inplace) |> gpu
            t2 = ScalePointCloud(0.5; inplace = inplace) |> gpu
            pc1 = PointCloud(p) |> gpu
            pc2 = t2(t1(pc1))
            @test pc1.points isa CuArray
            @test pc2.points isa CuArray
            @test all(isapprox.(Array(p), Array(pc2.points), rtol = 1e-5, atol = 1e-5))
        end
    end

    for inplace in [true, false]
        @testset "RotatePointCloud inplace=$(inplace)" begin
            p = rand(Float32, 3, 8, 2)
            rotmat = 2 .* one(rand(Float32, 3, 3))
            rotmat_inv = inv(rotmat)

            t1 = RotatePointCloud(rotmat; inplace = inplace) |> gpu
            t2 = RotatePointCloud(rotmat_inv; inplace = inplace) |> gpu
            pc1 = PointCloud(p) |> gpu
            pc2 = t2(t1(pc1)) |> gpu
            @test t1.rotmat isa CuArray
            @test t2.rotmat isa CuArray
            @test pc1.points isa CuArray
            @test pc2.points isa CuArray
            @test all(isapprox.(Array(p), Array(pc2.points), rtol = 1e-5, atol = 1e-5))
        end
    end

    for inplace in [true, false]
        @testset "ReAlignPointCloud inplace=$(inplace)" begin
            p1 = rand(Float32, 3, 8, 2)
            p2 = rand(Float32, 3, 8, 1)
            tgt = PointCloud(p2)
            src = PointCloud(p1) |> gpu
            t = ReAlignPointCloud(tgt; inplace = inplace) |> gpu
            src_1 = t(src) |> gpu
            @test t.t_min isa CuArray
            @test t.t_max isa CuArray
            @test src.points isa CuArray
            @test src_1.points isa CuArray

            # Transformed PointCloud should be inside the bounding box defined by `tgt` PointCloud
            @test all(
                reshape(maximum(Array(tgt.points), dims = 2), (:, 1)) .>=
                Array(src_1.points) .>=
                reshape(minimum(Array(tgt.points), dims = 2), (:, 1)),
            )
        end
    end

    for inplace in [true, false]
        @testset "NormalizePointCloud inplace=$(inplace)" begin
            p = rand(Float32, 3, 8, 2)
            t = NormalizePointCloud(; inplace = inplace)
            pc1 = PointCloud(p) |> gpu
            pc2 = t(pc1) |> gpu
            @test pc1.points isa CuArray
            @test pc2.points isa CuArray
            @test all(
                isapprox.(mean(Array(pc2.points); dims = 2), 0, rtol = 1e-5, atol = 1e-5),
            )
            @test all(
                isapprox.(std(Array(pc2.points); dims = 2), 1, rtol = 1e-5, atol = 1e-5),
            )
        end
    end

    @testset "Chain" begin
        p = rand(Float32, 3, 8, 2)
        t =
            Chain(
                ScalePointCloud(0.5),
                RotatePointCloud(rand(3, 3)),
                NormalizePointCloud(),
            ) |> gpu
        pc1 = PointCloud(p) |> gpu
        pc1 = t(pc1)
        @test pc1.points isa CuArray
        @test all(isapprox.(mean(Array(pc1.points); dims = 2), 0, rtol = 1e-5, atol = 1e-5))
        @test all(isapprox.(std(Array(pc1.points); dims = 2), 1, rtol = 1e-5, atol = 1e-4))
    end

    _mesh =
        load_trimesh(
            joinpath(@__DIR__, "../assets/teapot.obj"),
            joinpath(@__DIR__, "../assets/sphere.obj"),
        ) |> gpu

    for inplace in [true, false]
        @testset "ScaleTriMesh inplace=$(inplace)" begin
            m = deepcopy(_mesh)
            t1 = ScaleTriMesh(2.0; inplace = inplace) |> gpu
            t2 = ScaleTriMesh(0.5; inplace = inplace) |> gpu
            m2 = t2(t1(m))

            if inplace
                @test m2 === m
            else
                @test m2 !== m
            end
            @test get_verts_packed(m2) isa CuArray{Float32,2}
            @test all(
                isapprox.(
                    cpu(get_verts_packed(_mesh)),
                    cpu(get_verts_packed(m2)),
                    rtol = 1e-5,
                    atol = 1e-5,
                ),
            )
        end
    end

    for inplace in [true, false]
        @testset "RotateTriMesh inplace=$(inplace)" begin
            m = deepcopy(_mesh)
            rotmat = 2 .* one(rand(Float32, 3, 3))
            rotmat_inv = inv(rotmat)
            t1 = RotateTriMesh(rotmat; inplace = inplace) |> gpu
            t2 = RotateTriMesh(rotmat_inv; inplace = inplace) |> gpu
            m2 = t2(t1(m))
            if inplace
                @test m2 === m
            else
                @test m2 !== m
            end
            @test get_verts_packed(m2) isa CuArray{Float32,2}
            @test all(
                isapprox.(
                    cpu(get_verts_packed(_mesh)),
                    cpu(get_verts_packed(m2)),
                    rtol = 1e-5,
                    atol = 1e-5,
                ),
            )
        end
    end

    for inplace in [true, false]
        @testset "ReAlignTriMesh inplace=$(inplace)" begin
            m = deepcopy(_mesh)
            src = Flux3D.scale(m, 2.0)
            tgt = Flux3D.scale(m, 1.0)
            t = ReAlignTriMesh(tgt, 1; inplace = inplace) |> gpu
            src_1 = t(src)
            if inplace
                @test src === src_1
            else
                @test src !== src_1
            end
            @test get_verts_packed(src_1) isa CuArray{Float32,2}
            @test all(
                maximum(get_verts_list(_mesh)[1]; dims = 2) .>=
                get_verts_list(src_1)[1] .>=
                minimum(get_verts_list(_mesh)[1]; dims = 2),
            )
            @test all(
                maximum(get_verts_list(_mesh)[1]; dims = 2) .>=
                get_verts_list(src_1)[2] .>=
                minimum(get_verts_list(_mesh)[1]; dims = 2),
            )
        end
    end

    for inplace in [true, false]
        @testset "NormalizeTriMesh inplace=$(inplace)" begin
            m = deepcopy(_mesh)
            t = NormalizeTriMesh(; inplace = inplace) |> gpu
            m2 = t(m)
            if inplace
                @test m2 === m
            else
                @test m2 !== m
            end
            @test get_verts_packed(m2) isa CuArray{Float32,2}
            @test all(
                isapprox.(
                    mean(cpu(get_verts_list(m2))[1]; dims = 2),
                    0.0,
                    rtol = 1e-5,
                    atol = 1e-5,
                ),
            )
            @test all(
                isapprox.(
                    std(cpu(get_verts_list(m2))[1]; dims = 2),
                    1.0,
                    rtol = 1e-4,
                    atol = 1e-5,
                ),
            )
            @test all(
                isapprox.(
                    mean(cpu(get_verts_list(m2))[2]; dims = 2),
                    0.0,
                    rtol = 1e-5,
                    atol = 1e-5,
                ),
            )
            @test all(
                isapprox.(
                    std(cpu(get_verts_list(m2))[2]; dims = 2),
                    1.0,
                    rtol = 1e-4,
                    atol = 1e-5,
                ),
            )
        end
    end

    for inplace in [true, false]
        @testset "TranslateTriMesh inplace=$(inplace)" begin
            m = deepcopy(_mesh)
            t1 = TranslateTriMesh(1.0; inplace = inplace) |> gpu
            t2 = TranslateTriMesh(-1.0; inplace = inplace) |> gpu
            m2 = t2(t1(m))
            if inplace
                @test m2 === m
            else
                @test m2 !== m
            end
            @test get_verts_packed(m2) isa CuArray{Float32,2}
            @test all(
                isapprox.(
                    cpu(get_verts_packed(_mesh)),
                    cpu(get_verts_packed(m2)),
                    rtol = 1e-5,
                    atol = 1e-5,
                ),
            )
        end
    end

    for inplace in [true, false]
        @testset "OffsetTriMesh inplace=$(inplace)" begin
            m = deepcopy(_mesh)
            _offset = ones(size(get_verts_packed(m)))
            t1 = OffsetTriMesh(_offset; inplace = inplace) |> gpu
            t2 = OffsetTriMesh(_offset .* -1; inplace = inplace) |> gpu
            m2 = t2(t1(m))

            if inplace
                @test m2 === m
            else
                @test m2 !== m
            end
            @test get_verts_packed(m2) isa CuArray{Float32,2}
            @test all(
                isapprox.(
                    cpu(get_verts_packed(_mesh)),
                    cpu(get_verts_packed(m2)),
                    rtol = 1e-5,
                    atol = 1e-5,
                ),
            )
        end
    end

    _voxels = zeros(Float32, 32, 32, 32, 2) |> gpu
    _voxels[1:15, 2:10, 18:32, :] .= 1

    _v = VoxelGrid(_voxels) |> gpu
    _m =
        load_trimesh(
            joinpath(@__DIR__, "../assets/teapot.obj"),
            joinpath(@__DIR__, "../assets/sphere.obj"),
        ) |> gpu
    _p = PointCloud(sample_points(_m, 1024)) |> gpu

    res = 28
    points = 512
    thresh = 0.9

    @testset "TriMeshToVoxelGrid" begin
        m = deepcopy(_m)
        t = TriMeshToVoxelGrid(res)
        v = t(m)
        @test v isa VoxelGrid{Float32}
        @test v.voxels isa CUDA.CuArray
        @test size(v.voxels) == (res, res, res, 2)
    end

    @testset "PointCloudToVoxelGrid" begin
        p = deepcopy(_p)
        t = PointCloudToVoxelGrid(res)
        v = t(p)
        @test v isa VoxelGrid{Float32}
        @test v.voxels isa CUDA.CuArray
        @test size(v.voxels) == (res, res, res, 2)
    end

    for algo in [:Exact, :MarchingCubes, :MarchingTetrahedra, :NaiveSurfaceNets]
        @testset "VoxelGridToTriMesh algo=$(algo)" begin
            v = deepcopy(_v)
            t = VoxelGridToTriMesh(thresh = thresh, algo = algo)
            m = t(v)
            @test m isa TriMesh{Float32,UInt32,CuArray}
        end
    end

    for algo in [:Exact, :MarchingCubes, :MarchingTetrahedra, :NaiveSurfaceNets]
        @testset "PointCloudToTriMesh algo=$(algo)" begin
            p = deepcopy(_p)
            t = PointCloudToTriMesh(res, algo = algo)
            m = t(p)
            @test m isa TriMesh{Float32,UInt32,CuArray}
        end
    end

    @testset "TriMeshToPointCloud" begin
        m = deepcopy(_m)
        t = TriMeshToPointCloud(points)
        p = t(m)
        @test p isa PointCloud{Float32}
        @test p.points isa CUDA.CuArray
        @test size(p.points) == (3, points, 2)
    end

    for algo in [:Exact, :MarchingCubes, :MarchingTetrahedra]#, :NaiveSurfaceNets]
        @testset "VoxelGridToPointCloud algo=$(algo)" begin
            v = deepcopy(_v)
            t = VoxelGridToPointCloud(points, thresh = thresh, algo = algo)
            p = t(v)
            @test p isa PointCloud{Float32}
            @test p.points isa CUDA.CuArray
            @test size(p.points) == (3, points, 2)
        end
    end
end # Transforms


using Plots
using Random
using LaTeXStrings
include("../cmd.jl")
include("../simmd.jl")
include("../extramd.jl")
include("../pextramd.jl")
include("../AutoDiffUtils.jl")
include("../Potentials.jl")

lightblue = colorant"rgb(63%,74%,78%)"
orange = colorant"rgb(85%,55%,13%)"
silver = colorant"rgb(69%,67%,66%)"
rust = colorant"rgb(72%,26%,6%)"

m = 1
n = 1

true_x = [0.1]
# true_x[randperm(m)[1:20]] .= 0.0
true_y = [0.1]
# true_y[randperm(m)[1:20]] .= 0.0

A = ones(m, n)

l3 = 3.0

f(x, y) = 0.1 * 4^l3 * sum((x .- true_x)' * A * (y .- true_y))
g(x, y) = -f(x, y)

x0 = [0.4]
y0 = [0.3]

let
    max_iter = 5000
    global scatter_cmd = plot()
    global xHist_cmd = zeros(m, max_iter + 1)
    global yHist_cmd = zeros(n, max_iter + 1)
    global m
    global n
    global A
    global f
    global g
    global x0
    global true_x
    global y0
    global true_y

    α = 1.0

    ψ_x = shannon_entropy(α)
    ψ_y = shannon_entropy(α)

    x = ψ_x[2](x0)
    y = ψ_y[2](y0)

    opt = CMD(x, ψ_x, f, y, ψ_y, g, cg!)

    grad_computations = zeros(Int, max_iter + 1)
    for k = 1 : max_iter
      xHist_cmd[:, k], yHist_cmd[:, k] = iterate_primal(opt)
      # grad_computations[k+1] = grad_computations[k] + step!(opt)
      step!(opt)

      opt.x̄ .= ψ_x[2](max.(iterate_primal(opt)[1], 0.0))
      opt.ȳ .= ψ_y[2](max.(iterate_primal(opt)[2], 0.0))
      # @show grad_computations[k+1] - grad_computations[k]
    end
    xHist_cmd[:, end], yHist_cmd[:, end]= iterate_primal(opt)

    @show grad_computations[end]/max_iter

    scatter!(scatter_cmd, xHist_cmd[1, :], yHist_cmd[1, :], xlims=(0.0, 0.6), ylims=(0.0, 0.6))
end

let
    max_iter = 5000
    global scatter_extramd = plot()
    global xHist_extramd = zeros(m, max_iter + 1)
    global yHist_extramd = zeros(n, max_iter + 1)
    global m
    global n
    global A
    global f
    global g
    global x0
    global true_x
    global y0
    global true_y

    α = 1.0

    ψ_x = positive_projection(α)
    ψ_y = positive_projection(α)

    x = ψ_x[2](x0)
    y = ψ_y[2](y0)

    opt = ExtraMD(x, ψ_x, f, y, ψ_y, g)

    iters = zeros(Int, max_iter + 1)
    for k = 1 : max_iter
      xHist_extramd[:, k], yHist_extramd[:, k] = iterate_primal(opt)
      # iters[k+1] = iters[k] + step!(opt)
      step!(opt)

      opt.x̄ .= ψ_x[2](max.(iterate_primal(opt)[1], 0.0))
      opt.ȳ .= ψ_y[2](max.(iterate_primal(opt)[2], 0.0))
    end
    @show xHist_extramd[:, end], yHist_extramd[:, end]= iterate_primal(opt)

     scatter!(scatter_extramd, xHist_extramd[1, :], yHist_extramd[1, :], xlims=(0.0, 0.6), ylims=(0.0, 0.6))
end

out_plot = scatter(xHist_cmd[1, :], yHist_cmd[1, :], xlims=(0.0, 0.6), ylims=(0.0, 0.6), markercolor=silver, aspect_ratio=:equal, size=(500,500), label=L"\mathrm{CMW}", thickness_scaling=2.00, markerstrokewidth=0.1)
scatter!(out_plot, xHist_extramd[1, :], yHist_extramd[1, :], xlims=(0.0, 0.6), ylims=(0.0, 0.6), markercolor=rust, aspect_ratio=:equal, label=L"\mathrm{Extra gradient}", markerstrokewidth=0.1)

savefig(out_plot, "./out/figures/bilinear_positive_3to$(Int(l3)).pdf")
run(`pdfcrop --margins '5 5 5 5' ./out/figures/bilinear_positive_3to$(Int(l3)).pdf ./out/figures/bilinear_positive_3to$(Int(l3)).pdf`)


"""
    federated_averaging(round_weights)

Implement the federated averaging algorithm descibed in:
- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://arxiv.org/abs/1602.05629)

`round_weights` must be an array containing the weights from each clients
participating in the round.
"""
function federated_averaging(round_weights)
    mapreduce(w -> w[:], +, round_weights) / length(round_weights)
end


module mqpit1
using Elfel
using Elfel.RefShapes: npts, param_coords, weights
using Elfel.FElements: FE, FEData, refshape
using Elfel.FElements: bfun, bfungradpar, FEH1_T3
using Elfel.FESpaces: FESpace, ndofs, setebc!, nunknowns
using MeshSteward: Mesh, load, baseincrel
using Elfel.QPIterators: QPIterator, bfun, bfungradpar, weight
using Test
# using BenchmarkTools
function test()
    fe = FEH1_T3()
    mesh = load(Mesh(), "qmesh.mesh")
    fesp = FESpace(Float64, mesh, fe)

    qpit = QPIterator(fesp, (kind = :default,))
    @test length(qpit) == 1
    # @test isapprox(param_coords(quadrule(idom)), [0.3333333333333333 0.3333333333333333])
    # @test isapprox(weights(quadrule(idom)), [0.5])

    refNs = [[0.3333333333333334, 0.3333333333333333, 0.3333333333333333]]
    k = 1
    for qp in qpit
        Ns = bfun(qp)
        @test isapprox(Ns, refNs[k])
        k = k + 1
    end

    qpit = QPIterator(fesp, (kind = :default, npts = 3))
    @test length(qpit) == 3
    # @test isapprox(param_coords(quadrule(idom)), [0.3333333333333333 0.3333333333333333])
    # @test isapprox(weights(quadrule(idom)), [0.5])

    refNs = [[0.1666666666666667, 0.6666666666666666, 0.16666666666666666],
    [0.16666666666666674, 0.16666666666666666, 0.6666666666666666],
    [0.6666666666666667, 0.16666666666666666, 0.16666666666666666]]
    k = 1
    for qp in qpit
        Ns = bfun(qp)
        @test isapprox(Ns, refNs[k])
        k = k + 1
    end
    k = 1
    for qp in qpit
        Ns = bfun(qp)
        @test isapprox(Ns, refNs[k])
        k = k + 1
    end

    refgNs = [[-1.0 -1.0], [1.0 0.0], [0.0 1.0]]
    for qp in qpit
        gradNs = bfungradpar(qp)
        for i in 1:length(gradNs)
            @test isapprox(gradNs[i], refgNs[i])
        end
    end

    for qp in qpit
        w = weight(qp)
        @test w == 0.16666666666666666
    end
    # Ns, gradNparams = bfundata(idom)
    # @test length(Ns[1]) == 3

    # geom = geomattr(femesh)
    # conn = connectivity(femesh)

    # it = FEIterator(fesp)
    # @show c = iterate(it)
    # Ns, gradNparams = bfundata(idom)
    # J = jac(geom, c, gradNparams[1])
    # # @btime J = jac($geom, retrieve($ir, 1), $gradNparams[1])
    # @test isapprox(J, [0.5 0.5; 0.0 0.3333333333333333])
end
end
using .mqpit1
mqpit1.test()