# A quadratic quadrilateral in 2D space.
struct Quadrilateral8_2D <: Face_2D
    # The points are assumed to be ordered  in counter clockwise order as follows
    # p₁ = vertex A
    # p₂ = vertex B
    # p₃ = vertex C
    # p₄ = vertex D
    # p₅ = point on the quadratic segment from A to B
    # p₆ = point on the quadratic segment from B to C
    # p₇ = point on the quadratic segment from C to D
    # p₈ = point on the quadratic segment from D to A
    points::SVector{8, Point_2D}
end

# Constructors
# -------------------------------------------------------------------------------------------------
Quadrilateral8_2D(p₁::Point_2D, p₂::Point_2D, p₃::Point_2D, p₄::Point_2D,
                  p₅::Point_2D, p₆::Point_2D, p₇::Point_2D, p₈::Point_2D
                 ) = Quadrilateral8_2D(SVector(p₁, p₂, p₃, p₄, p₅, p₆, p₇, p₈))

# Methods
# -------------------------------------------------------------------------------------------------
function (quad8::Quadrilateral8_2D)(r::Real, s::Real)
    # See The Visualization Toolkit: An Object-Oriented Approach to 3D Graphics, 4th Edition
    # Chapter 8, Advanced Data Representation, in the interpolation functions section
    ξ = 2*Float64(r) - 1; η = 2*Float64(s) - 1
    return (1 - ξ)*(1 - η)*(-ξ - η - 1)/4*quad8[1] +
           (1 + ξ)*(1 - η)*( ξ - η - 1)/4*quad8[2] +
           (1 + ξ)*(1 + η)*( ξ + η - 1)/4*quad8[3] +
           (1 - ξ)*(1 + η)*(-ξ + η - 1)/4*quad8[4] +
                      (1 - ξ^2)*(1 - η)/2*quad8[5] +
                      (1 - η^2)*(1 + ξ)/2*quad8[6] +
                      (1 - ξ^2)*(1 + η)/2*quad8[7] +
                      (1 - η^2)*(1 - ξ)/2*quad8[8]
end

# Interpolation with a point, instead of (r,s)
function (quad8::Quadrilateral8_2D)(p::Point_2D)
    r, s = p
    ξ = 2r - 1; η = 2s - 1
    return (1 - ξ)*(1 - η)*(-ξ - η - 1)/4*quad8[1] +
           (1 + ξ)*(1 - η)*( ξ - η - 1)/4*quad8[2] +
           (1 + ξ)*(1 + η)*( ξ + η - 1)/4*quad8[3] +
           (1 - ξ)*(1 + η)*(-ξ + η - 1)/4*quad8[4] +
                      (1 - ξ^2)*(1 - η)/2*quad8[5] +
                      (1 - η^2)*(1 + ξ)/2*quad8[6] +
                      (1 - ξ^2)*(1 + η)/2*quad8[7] +
                      (1 - η^2)*(1 - ξ)/2*quad8[8]
end

area(quad8::Quadrilateral8_2D) = area(quad8, Val(2))

function area(quad8::Quadrilateral8_2D, ::Val{N}) where {N}
    # Numerical integration required. Gauss-Legendre quadrature over a quadrilateral is used.
    # Let Q(r,s) be the interpolation function for quad8,
    #                             1  1
    # A = ∬ ||∂Q/∂r × ∂Q/∂s||dA = ∫  ∫ ||∂Q/∂r × ∂Q/∂s|| ds dr
    #      D                      0  0
    #
    #       N   N
    #   =   ∑   ∑  wᵢwⱼ||∂Q/∂r(rᵢ,sⱼ) × ∂Q/∂s(rᵢ,sⱼ)||
    #      i=1 j=1
    # N is the square root of the number of points used in the quadrature.
    # See tuning/Quadrilateral8_2D_area.jl for more info on how N = 3 was chosen.
    w, r = gauss_legendre_quadrature(Val(N))
    a = 0.0
    for i = 1:N, j = 1:N
        ∂Q_∂r, ∂Q_∂s = gradient(quad8, r[i], r[j])
        a += w[i]*w[j]*abs(∂Q_∂r × ∂Q_∂s)
    end
    return a
end

centroid(quad8::Quadrilateral8_2D) = centroid(quad8, Val(3))
function centroid(quad8::Quadrilateral8_2D, ::Val{N}) where {N}
    # Numerical integration required. Gauss-Legendre quadrature over a quadrilateral is used.
    # Let Q(r,s) be the interpolation function for quad8,
    #                             1  1
    # A = ∬ ||∂Q/∂r × ∂Q/∂s||dA = ∫  ∫ ||∂Q/∂r × ∂Q/∂s|| ds dr
    #      D                      0  0
    #
    #       N   N
    #   =   ∑   ∑  wᵢwⱼ||∂Q/∂r(rᵢ,sⱼ) × ∂Q/∂s(rᵢ,sⱼ)||
    #      i=1 j=1
    #
    # C_x = (∫∫ x dA)/A, C_y = (∫∫ y dA)/A
    #         D                  D
    w, r = gauss_legendre_quadrature(Val(N))
    A = 0.0
    C = Point_2D()
    for i = 1:N, j = 1:N
        ∂Q_∂r, ∂Q_∂s = gradient(quad8, r[i], r[j])
        weighted_val = w[i]*w[j]*abs(∂Q_∂r × ∂Q_∂s)
        A += weighted_val
        C += weighted_val * quad8(r[i], r[j])
    end
    return C/A
end

function gradient(quad8::Quadrilateral8_2D, r::Real, s::Real)
    # Chain rule
    # ∂Q   ∂Q ∂ξ     ∂Q      ∂Q   ∂Q ∂η     ∂Q
    # -- = -- -- = 2 -- ,    -- = -- -- = 2 --
    # ∂r   ∂ξ ∂r     ∂ξ      ∂s   ∂η ∂s     ∂η
    ξ = 2Float64(r) - 1; η = 2Float64(s) - 1
    ∂Q_∂ξ = (1 - η)*(2ξ + η)/4*quad8[1] +
            (1 - η)*(2ξ - η)/4*quad8[2] +
            (1 + η)*(2ξ + η)/4*quad8[3] +
            (1 + η)*(2ξ - η)/4*quad8[4] +
                    -ξ*(1 - η)*quad8[5] +
                   (1 - η^2)/2*quad8[6] +
                    -ξ*(1 + η)*quad8[7] +
                  -(1 - η^2)/2*quad8[8]

    ∂Q_∂η = (1 - ξ)*( ξ + 2η)/4*quad8[1] +
            (1 + ξ)*(-ξ + 2η)/4*quad8[2] +
            (1 + ξ)*( ξ + 2η)/4*quad8[3] +
            (1 - ξ)*(-ξ + 2η)/4*quad8[4] +
                   -(1 - ξ^2)/2*quad8[5] +
                     -η*(1 + ξ)*quad8[6] +
                    (1 - ξ^2)/2*quad8[7] +
                     -η*(1 - ξ)*quad8[8]

    return 2*∂Q_∂ξ, 2*∂Q_∂η
end

function in(p::Point_2D, quad8::Quadrilateral8_2D)
    # If the point is to the left of every edge
    #  4<-----3
    #  |      ^
    #  | p    |
    #  |      |
    #  |      |
    #  v----->2
    #  1
    return isleft(p, QuadraticSegment_2D(quad8[1], quad8[2], quad8[5])) &&
           isleft(p, QuadraticSegment_2D(quad8[3], quad8[4], quad8[7])) &&
           isleft(p, QuadraticSegment_2D(quad8[2], quad8[3], quad8[6])) &&
           isleft(p, QuadraticSegment_2D(quad8[4], quad8[1], quad8[8]))
end

# function in(p::Point_2D{F}, quad8::Quadrilateral8_2D{F}, N::Int64) where {F <: AbstractFloat}
#     # Determine if the point is in the triangle using the Newton-Raphson method
#     # N is the max number of iterations of the method.
#     p_rs = real_to_parametric(p, quad8, N)
#     ϵ = parametric_coordinate_ϵ
#     if (-ϵ ≤ p_rs[1] ≤ 1 + ϵ) &&
#        (-ϵ ≤ p_rs[2] ≤ 1 + ϵ)
#         return true
#     else
#         return false
#     end
# end

function intersect(l::LineSegment_2D, quad8::Quadrilateral8_2D)
    # Create the 3 quadratic segments that make up the triangle and intersect each one
    edges = SVector(QuadraticSegment_2D(quad8[1], quad8[2], quad8[5]),
                    QuadraticSegment_2D(quad8[2], quad8[3], quad8[6]),
                    QuadraticSegment_2D(quad8[3], quad8[4], quad8[7]),
                    QuadraticSegment_2D(quad8[4], quad8[1], quad8[8]))
    ipoints = MVector(Point_2D(), Point_2D(), Point_2D(),
                      Point_2D(), Point_2D(), Point_2D())
    n_ipoints = 0x00000000
    # We need to account for 6 points returned
    for k ∈ 1:4
        npoints, points = l ∩ edges[k]
        for i = 1:npoints
            n_ipoints += 0x00000001
            ipoints[n_ipoints] = points[i]
        end
    end
    return n_ipoints, SVector(ipoints)
end

function jacobian(quad8::Quadrilateral8_2D, r::Real, s::Real)
    # Return the 2 x 2 Jacobian matrix
    ∂Q_∂r, ∂Q_∂s = gradient(quad8, r, s)
    return SMatrix{2, 2}(∂Q_∂r.x, ∂Q_∂r.y,
                         ∂Q_∂s.x, ∂Q_∂s.y)
end

function real_to_parametric(p::Point_2D, quad8::Quadrilateral8_2D)
    return real_to_parametric(p, quad8, 30)
end

function real_to_parametric(p::Point_2D, quad8::Quadrilateral8_2D, N::Int64)
    # Convert from real coordinates to the triangle's local parametric coordinates using the
    # the Newton-Raphson method. N is the max number of iterations
    # If a conversion doesn't exist, the minimizer is returned.
    r = 0.5 # Initial guess at centroid
    s = 0.5
    for i = 1:N
        err = p - quad8(r, s)
        # Inversion is faster for 2 by 2 than \
        Δr, Δs = inv(jacobian(quad8, r, s)) * err
        r += Δr
        s += Δs
        if abs(Δr) + abs(Δs) < 1e-6
            break
        end
    end
    return Point_2D(r, s)
end

function triangulate(quad8::Quadrilateral8_2D, N::Int64)
    # N is the number of divisions of each edge
    triangles = Vector{Triangle_2D}(undef, 2*(N+1)*(N+1))
    if N === 0
        triangles[1] = Triangle_2D(quad8[1], quad8[2], quad8[3])
        triangles[2] = Triangle_2D(quad8[3], quad8[4], quad8[1])
    else
        for j = 0:N, i = 0:N
            triangles[2*(N+1)*j + 2i + 1] = Triangle_2D(quad8(    i/(N+1),     j/(N+1)),
                                                        quad8((i+1)/(N+1),     j/(N+1)),
                                                        quad8(    i/(N+1), (j+1)/(N+1)))
            triangles[2*(N+1)*j + 2i + 2] = Triangle_2D(quad8(    i/(N+1), (j+1)/(N+1)),
                                                        quad8((i+1)/(N+1),     j/(N+1)),
                                                        quad8((i+1)/(N+1), (j+1)/(N+1)))
        end
    end
    return triangles
end

# Plot
# -------------------------------------------------------------------------------------------------
if enable_visualization
    function convert_arguments(LS::Type{<:LineSegments}, quad8::Quadrilateral8_2D)
        q₁ = QuadraticSegment_2D(quad8[1], quad8[2], quad8[5])
        q₂ = QuadraticSegment_2D(quad8[2], quad8[3], quad8[6])
        q₃ = QuadraticSegment_2D(quad8[3], quad8[4], quad8[7])
        q₄ = QuadraticSegment_2D(quad8[4], quad8[1], quad8[8])
        qsegs = [q₁, q₂, q₃, q₄]
        return convert_arguments(LS, qsegs)
    end

    function convert_arguments(LS::Type{<:LineSegments}, Q::Vector{Quadrilateral8_2D})
        point_sets = [convert_arguments(LS, quad8) for quad8 ∈ Q]
        return convert_arguments(LS, reduce(vcat, [pset[1] for pset ∈ point_sets]))
    end

    function convert_arguments(M::Type{<:Mesh}, quad8::Quadrilateral8_2D)
        triangles = triangulate(quad8, 13)
        return convert_arguments(M, triangles)
    end

    function convert_arguments(M::Type{<:Mesh}, Q::Vector{Quadrilateral8_2D})
        triangles = reduce(vcat, triangulate.(Q, 13))
        return convert_arguments(M, triangles)
    end
end


### A Pluto.jl notebook ###
# v0.18.0

using Markdown
using InteractiveUtils

# ╔═╡ f32b184a-ba70-4e80-8f8d-e12def369128
using Parameters

# ╔═╡ 456df07b-1be8-4adb-9302-9290314311e3
begin
	@enum  CartesianDirection XDirection = 1 YDirection ZDirection
	@enum  CartesianAxis XAxis = 1 YAxis ZAxis
	@enum  FieldComponent XComponent = 1 YComponent ZComponent
end

# ╔═╡ 4f7a7193-e75e-4756-9423-7d81e2f7e4db
begin

mutable struct StructuredGrid{nD} <: AbstractArray{Real,nD}
	nodes::NTuple{nD,Vector{Float64}}
end

function StructuredGrid(nodes...)
        StructuredGrid{length(nodes)}(nodes)
end

Base.size(grid::StructuredGrid)	= length.(grid.nodes)

nodes(sgrid::StructuredGrid) = sgrid.nodes
xnodes(grid::StructuredGrid)	= grid.nodes[1]
ynodes(grid::StructuredGrid)	= grid.nodes[2]
znodes(grid::StructuredGrid)	= grid.nodes[3]

xnodes(grid::StructuredGrid,i::Int)	= getindex(xnodes(grid),i)
ynodes(grid::StructuredGrid,j::Int)	=  getindex(ynodes(grid),j)
znodes(grid::StructuredGrid,k::Int)	= getindex(znodes(grid),k)

Base.getindex(grid::StructuredGrid{1},i::Int) = (xnodes(grid,i))
Base.getindex(grid::StructuredGrid{2},i::Int,j::Int) = (xnodes(grid,i),ynodes(grid,j))
Base.getindex(grid::StructuredGrid{3},i::Int,j::Int,k::Int) =  (xnodes(grid,i),ynodes(grid,j), znodes(grid,k))

cell_sizes(sgrid::StructuredGrid) = StructuredGrid(diff.(nodes(sgrid)))
cell_volumes(sgrid::StructuredGrid) = map(prod,cell_sizes(sgrid))
_avg(nodes::AbstractVector) = 0.5(nodes[2:end] + nodes[1:end-1])
cell_centroids(sgrid::StructuredGrid) = StructuredGrid(_avg.(nodes(sgrid)))

abstract type AbstractNodeType end
struct 	Primal <: AbstractNodeType  end
struct 	Dual <: AbstractNodeType  end

function primal_nodes(ax::FieldComponent, reference_grid::StructuredGrid{Dim}) where {Dim}
    StructuredGrid(ntuple(i -> i == Int(ax) ? _avg(nodes(reference_grid)[i]) : nodes(reference_grid)[i], Dim))
end
function dual_nodes(ax::FieldComponent, reference_grid::StructuredGrid{Dim}) where {Dim}
    StructuredGrid(ntuple(i -> i == Int(ax) ? nodes(reference_grid)[i] :  _avg(nodes(reference_grid)[i])  , Dim))
end

function nodes(::Type{Primal},ax::FieldComponent, reference_grid::StructuredGrid)
primal_nodes(ax,reference_grid)
end
function nodes(::Type{Dual},ax::FieldComponent, reference_grid::StructuredGrid)
dual_nodes(ax,reference_grid)
end

end

# ╔═╡ 1bc97871-d861-47d4-bde9-e2373b1e903a
begin
struct YeeGrid{T,nD}
	x::StructuredGrid{nD}
	y::StructuredGrid{nD}
	z::StructuredGrid{nD}
end

function YeeGrid{T}(refgrid::StructuredGrid{nD}) where {nD,T}
	YeeGrid{T,nD}(
		nodes(T,XComponent,refgrid),
		nodes(T,YComponent,refgrid),
		nodes(T,ZComponent,refgrid))
end
function YeeGrid{T}(nodes...) where T
	refgrid = StructuredGrid(nodes...)
	YeeGrid{T}(refgrid)
end
Base.size(yeegrid::YeeGrid) = (size(yeegrid.x),size(yeegrid.y),size(yeegrid.z))

end

# ╔═╡ 4b3b120f-52b0-47ae-bec2-207ea12ddd6a
begin
	grid_3d_primal = YeeGrid{Primal}(LinRange(-2,2,10),LinRange(-5,5,100),LinRange(-3,3,40))
	grid_3d_dual = YeeGrid{Dual}(LinRange(-2,2,10),LinRange(-5,5,100),LinRange(-3,3,40))
	# nonuniform grid
	grid_2d_primal = YeeGrid{Primal}([1,4,7,8,9,10,15,20],LinRange(-1,1,15))
	grid_2d_dual = YeeGrid{Primal}([1,4,7,8,9,10,15,20],LinRange(-1,1,15))

	grid_1d_dual = YeeGrid{Primal}([1,4,7,8,9,10,15,20])
end

# ╔═╡ 6ad42c82-25c5-41c3-83f1-0ff002edbafe
begin
struct VectorField{PD,T,N}
	x::Array{T,N}
	y::Array{T,N}
	z::Array{T,N}
end
function VectorField{T}(yeegrid::YeeGrid{PD,N}) where {T,PD,N}
	VectorField{PD,T,N}(zeros.(T,size(yeegrid))...)
end
struct TensorField{PD,T,N}
	xx::Array{T,N}
	xy::Array{T,N}
	xz::Array{T,N}
	yx::Array{T,N}
	yy::Array{T,N}
	yz::Array{T,N}
	zx::Array{T,N}
	zy::Array{T,N}
	zz::Array{T,N}
end
function TensorField{T}(yeegrid::YeeGrid{PD,N}) where {T,PD,N}
	TensorField{PD,T,N}(zeros.(T,size(yeegrid))...,zeros.(T,size(yeegrid))...,zeros.(T,size(yeegrid))...)
end
end

# ╔═╡ 4fd9f13e-8279-41e2-b7f5-e9865ed869be
function FaradayStep!(b,e,m,cb;ref_grid)

(nx,ny,nz) = size(ref_grid)

for k = 1:nz-1
	for j = 1:ny-1
		for i = 1:nx-1
b.x[i,j,k] =
	 b.x[i,j,k] + cb.x[i,j,k] * (e.z[i,j,k]-e.z[i,j+1,k]+e.y[i,j,k+1]-e.y[i,j,k]) -m.x[i,j,k]
b.y[i,j,k] =
	 b.y[i,j,k] + cb.y[i,j,k] * (e.x[i,j,k]-e.x[i,j,k+1]+e.z[i+1,j,k]-e.z[i,j,k]) - m.y[i,j,k]
b.z[i,j,k] =
	 b.z[i,j,k] + cb.z[i,j,k] * (e.y[i,j,k]-e.y[i+1,j,k]+e.x[i,j+1,k]-e.x[i,j,k]) -
	 m.z[i,j,k]
		end
	end
end

for k = 1:nz-1
	for j = 1:ny-1
		for i = nx:nx
b.x[i,j,k] =
	 b.x[i,j,k] + cb.x[i,j,k] * (e.z[i,j,k]-e.z[i,j+1,k]+e.y[i,j,k+1]-e.y[i,j,k]) -m.x[i,j,k]
		end
	end
end
for k = 1:nz-1
	for j = ny:ny
		for i = 1:nx-1
b.y[i,j,k] =
	 b.y[i,j,k] + cb.y[i,j,k] * (e.x[i,j,k]-e.x[i,j,k+1]+e.z[i+1,j,k]-e.z[i,j,k]) - m.y[i,j,k]
		end
	end
end
for k = nz:nz
	for j = 1:ny-1
		for i = 1:nx-1
b.z[i,j,k] =
	 b.z[i,j,k] + cb.z[i,j,k] * (e.y[i,j,k]-e.y[i+1,j,k]+e.x[i,j+1,k]-e.x[i,j,k]) -
	 m.z[i,j,k]
		end
	end
end

end

# ╔═╡ fcb5f651-3e20-4181-9e72-9310f823adf0
function AmpereStep!(d,h,J,cb;ref_grid)

(nx,ny,nz) = size(ref_grid)

for k = 2:nz-1
	for j = 2:ny-1
		for i = 2:nx-1
d.x[i,j,k] = d.x[i,j,k] +
cb.x[i,j,k] * (h.z[i,j,k]-h.z[i,j-1,k]-h.y[i,j,k]+h.y[i,j,k-1]) -J.x[i,j,k]
d.y[i,j,k] = d.y[i,j,k] +
cb.y[i,j,k] * (h.x[i,j,k]-h.x[i,j,k-1]-h.z[i,j,k]+h.z[i-1,j,k]) -J.y[i,j,k]
d.z[i,j,k] = d.z[i,j,k] +
cb.z[i,j,k] * (h.y[i,j,k]-h.y[i-1,j,k]-h.x[i,j,k]+h.z[i,j-1,k]) - J.z[i,j,k]
		end
	end
end

for k = 2:nz-1
	for j = 2:ny-1
		for i = 1:1
d.x[i,j,k] = d.x[i,j,k] +
cb.x[i,j,k] * (h.z[i,j,k]-h.z[i,j-1,k]-h.y[i,j,k]+h.y[i,j,k-1]) -J.x[i,j,k]		end
	end
end
for k = 2:nz-1
	for j = 1:1
		for i = 2:nx-1
d.y[i,j,k] = d.y[i,j,k] +
cb.x[i,j,k] * (h.x[i,j,k]-h.x[i,j,k-1]-h.z[i,j,k]+h.z[i-1,j,k]) -J.y[i,j,k]		end
	end
end
for k = 1:1
	for j = 2:ny-1
		for i = 2:nx-1
d.z[i,j,k] = d.z[i,j,k] +
cb.x[i,j,k] * (h.y[i,j,k]-h.y[i-1,j,k]-h.x[i,j,k]+h.z[i,j-1,k]) - J.z[i,j,k]
		end
	end
end
end

# ╔═╡ 7f1e505f-7223-4ae8-a6d4-9ff455553256
function _update_primal_constitutive!(e,d,C,nx,ny,nz)
for k = 2:nz-1
	for j = 2:ny-1
		for i = 2:nx-1
dy =  0.25 * (d.y[i,j,k] + d.y[i+1,j,k] + d.y[i+1,j-1,k] + d.y[i,j-1,k] )
dz =  0.25 * (d.z[i,j,k] + d.z[i+1,j,k] + d.z[i+1,j,k-1] + d.z[i,j,k-1] )
e.x[i,j,k]  = C.xx[i,j,k] * d.x[i,j,k]  + C.xy[i,j,k] * dy  + C.xz[i,j,k] * dz

dx = 0.25*(d.x[i,j,k]+d.x[i,j+1,k]+d.x[i-1,j,k]+d.x[i-1,j+1,k])
dz = 0.25*(d.z[i,j,k] + d.z[i+1,j,k]+ d.z[i+1,j,k-1] + d.z[i,j,k-1])

e.y[i,j,k]  = C.yx[i,j,k] * dx  + C.yy[i,j,k] * d.y[i,j,k]  + C.yz[i,j,k] * dz

dx = 0.25 * (d.x[i,j,k] + d.x[i,j,k+1] + d.x[i-1,j,k+1]+d.x[i-1,j,k])
dy = 0.25 * (d.y[i,j,k] + d.y[i,j,k+1] + d.y[i,j-1,k+1]+d.y[i,j-1,k])

e.z[i,j,k]  = C.zx[i,j,k] * dx  + C.zy[i,j,k] * dy  + C.yz[i,j,k] *   d.z[i,j,k]

		end
	end
end

for k = 2:nz-1
	for j = 2:ny-1
		for i = 1:1

dy =  0.25 * (d.y[i,j,k] + d.y[i+1,j,k] + d.y[i+1,j-1,k] + d.y[i,j-1,k] )
dz =  0.25 * (d.z[i,j,k] + d.z[i+1,j,k] + d.z[i+1,j,k-1] + d.z[i,j,k-1] )
e.x[i,j,k]  = C.xx[i,j,k] * d.x[i,j,k]  + C.xy[i,j,k] * dy  + C.xz[i,j,k] * dz

	end
	end
end
for k = 2:nz-1
	for j = 1:1
		for i = 2:nx-1

dx = 0.25*(d.x[i,j,k]+d.x[i,j+1,k]+d.x[i-1,j,k]+d.x[i-1,j+1,k])
dz = 0.25*(d.z[i,j,k] + d.z[i,j+1,k]+ d.z[i,j+1,k-1] + d.z[i,j,k-1])

e.y[i,j,k]  = C.yx[i,j,k] * dx  + C.yy[i,j,k] * d.y[i,j,k]  + C.yz[i,j,k] * dz
		end
	end
end
for k = 1:1
	for j = 2:ny-1
		for i = 2:nx-1
dx = 0.25 * (d.x[i,j,k] + d.x[i,j,k+1] + d.x[i-1,j,k+1]+d.x[i-1,j,k])
dy = 0.25 * (d.y[i,j,k] + d.y[i,j,k+1] + d.y[i,j-1,k+1]+d.y[i,j-1,k])

e.z[i,j,k]  = C.zx[i,j,k] * dx  + C.zy[i,j,k] * dy  + C.yz[i,j,k] *   d.z[i,j,k]
		end
	end
end

end

# ╔═╡ 8176e6dd-0768-4d9b-a2b7-957236e4c55b
function _update_dual_constitutive!(h,b,C,nx,ny,nz)
	for k = 2:nz-1
	for j = 2:ny-1
		for i = 2:nx-1
by = 0.25 * (b.y[i,j,k] + b.y[i-1,j,k] + b.y[i-1,j+1,k] + b.y[i,j+1,k])
bz = 0.25 * (b.z[i,j,k] + b.z[i,j,k+1] + b.z[i-1,j,k+1] + b.z[i-1,j,k])
h.x[i,j,k]  = C.xx[i,j,k] * b.x[i,j,k]  + C.xy[i,j,k] * by  + C.xz[i,j,k] * bz

bx = 0.25*(b.x[i,j,k]+b.x[i,j-1,k]+b.x[i+1,j-1,k]+b.x[i+1,j,k])
bz = 0.25*(b.z[i,j,k] + b.z[i,j-1,k]+ b.z[i,j-1,k+1] + b.z[i,j,k+1])

h.y[i,j,k]  = C.yx[i,j,k] * bx  + C.yy[i,j,k] * b.y[i,j,k]  + C.yz[i,j,k] * bz

bx = 0.25*(b.x[i,j,k]+b.x[i+1,j,k]+b.x[i+1,j,k-1]+b.x[i,j,k-1])
by = 0.25*(b.y[i,j,k]+b.y[i,j+1,k]+b.y[i,j+1,k-1]+b.y[i,j,k-1])

h.z[i,j,k]  = C.zx[i,j,k] * bx  + C.zy[i,j,k] * by  + C.zz[i,j,k] * b.z[i,j,k]

		end
	end
end

for k = 1:1
	for j = 1:1
		for i = 2:nx-1
by = 0.25 * (b.y[i,j,k] + b.y[i-1,j,k] + b.y[i-1,j+1,k] + b.y[i,j+1,k])
bz = 0.25 * (b.z[i,j,k] + b.z[i,j,k+1] + b.z[i-1,j,k+1] + b.z[i-1,j,k])
h.x[i,j,k]  = C.xx[i,j,k] * b.x[i,j,k]  + C.xy[i,j,k] * by  + C.xz[i,j,k] * bz
		end
	end
end
for k = 1:1
	for j = 2:ny-1
		for i = 1:1
bx = 0.25*(b.x[i,j,k]+b.x[i,j-1,k]+b.x[i+1,j-1,k]+b.x[i+1,j,k])
bz = 0.25*(b.z[i,j,k] + b.z[i,j-1,k]+ b.z[i,j-1,k+1] + b.z[i,j,k+1])
h.y[i,j,k]  = C.yx[i,j,k] * bx  + C.yy[i,j,k] * b.y[i,j,k]  + C.yz[i,j,k] * bz
		end
	end
end
for k = 2:ny-1
	for j = 1:1
		for i = 1:1
bx = 0.25*(b.x[i,j,k]+b.x[i+1,j,k]+b.x[i+1,j,k-1]+b.x[i,j,k-1])
by = 0.25*(b.y[i,j,k]+b.y[i,j+1,k]+b.y[i,j+1,k-1]+b.y[i,j,k-1])

h.z[i,j,k]  = C.zx[i,j,k] * bx  + C.zy[i,j,k] * by  + C.zz[i,j,k] * b.z[i,j,k]
		end
	end
end

end

# ╔═╡ d87ef4f9-e4e9-4b1d-907d-054505d3ff54
begin
	ref_grid = StructuredGrid(LinRange(-5,15,30),LinRange(-5,17,10),LinRange(-5,1,101))

	e = VectorField{Float64}(YeeGrid{Primal}(ref_grid))
	d = VectorField{Float64}(YeeGrid{Primal}(ref_grid))
	ϵ = TensorField{Float64}(YeeGrid{Primal}(ref_grid))

	b = VectorField{Float64}(YeeGrid{Dual}(ref_grid))
	h = VectorField{Float64}(YeeGrid{Dual}(ref_grid))
	μ = TensorField{Float64}(YeeGrid{Dual}(ref_grid))

	j =  VectorField{Float64}(YeeGrid{Primal}(ref_grid))
	σₑ = TensorField{Float64}(YeeGrid{Primal}(ref_grid))

	m =  VectorField{Float64}(YeeGrid{Dual}(ref_grid))
	σₘ = TensorField{Float64}(YeeGrid{Dual}(ref_grid))

	cf = VectorField{Float64}(YeeGrid{Dual}(ref_grid))
	ca = VectorField{Float64}(YeeGrid{Primal}(ref_grid))

end

# ╔═╡ 03ac2bca-403e-478e-b947-208427fdf579
begin
	FaradayStep!(b,e,m,cf, ref_grid = ref_grid)
	_update_dual_constitutive!(h,b,μ,size(ref_grid)...)
	_update_dual_constitutive!(m,h,σₘ,size(ref_grid)...)

	AmpereStep!(d,h,j,ca, ref_grid = ref_grid)
	_update_primal_constitutive!(e,d,ϵ,size(ref_grid)...)
	_update_primal_constitutive!(j,e,σₑ,size(ref_grid)...)

end

# ╔═╡ 00000000-0000-0000-0000-000000000001
PLUTO_PROJECT_TOML_CONTENTS = """
[deps]
Parameters = "d96e819e-fc66-5662-9728-84c9c7592b0a"

[compat]
Parameters = "~0.12.3"
"""

# ╔═╡ 00000000-0000-0000-0000-000000000002
PLUTO_MANIFEST_TOML_CONTENTS = """
# This file is machine-generated - editing it directly is not advised

julia_version = "1.7.0"
manifest_format = "2.0"

[[deps.OrderedCollections]]
git-tree-sha1 = "85f8e6578bf1f9ee0d11e7bb1b1456435479d47c"
uuid = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
version = "1.4.1"

[[deps.Parameters]]
deps = ["OrderedCollections", "UnPack"]
git-tree-sha1 = "34c0e9ad262e5f7fc75b10a9952ca7692cfc5fbe"
uuid = "d96e819e-fc66-5662-9728-84c9c7592b0a"
version = "0.12.3"

[[deps.UnPack]]
git-tree-sha1 = "387c1f73762231e86e0c9c5443ce3b4a0a9a0c2b"
uuid = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
version = "1.0.2"
"""

# ╔═╡ Cell order:
# ╠═456df07b-1be8-4adb-9302-9290314311e3
# ╠═4f7a7193-e75e-4756-9423-7d81e2f7e4db
# ╠═1bc97871-d861-47d4-bde9-e2373b1e903a
# ╠═4b3b120f-52b0-47ae-bec2-207ea12ddd6a
# ╠═6ad42c82-25c5-41c3-83f1-0ff002edbafe
# ╠═f32b184a-ba70-4e80-8f8d-e12def369128
# ╠═4fd9f13e-8279-41e2-b7f5-e9865ed869be
# ╠═fcb5f651-3e20-4181-9e72-9310f823adf0
# ╠═7f1e505f-7223-4ae8-a6d4-9ff455553256
# ╠═8176e6dd-0768-4d9b-a2b7-957236e4c55b
# ╠═d87ef4f9-e4e9-4b1d-907d-054505d3ff54
# ╠═03ac2bca-403e-478e-b947-208427fdf579
# ╟─00000000-0000-0000-0000-000000000001
# ╟─00000000-0000-0000-0000-000000000002


using Catalan
using Base.Test

@test ([5,4,2,2]\[2,1]) == ([5, 4, 2, 2],[2, 1])
@test isrimhook([4,3,2], [2,2,2])
@test !isrimhook([4,3,2], [2,2,1])
@test !isrimhook([4,3,2], [1,1])
λ = [5,4,2,1]
@test partitionsequence(λ) == [1, 0, 1, 0, 1, 1, 0, 1, 0]
@test character(λ, [4,3,2,2,1]) == 0
@test character([1], [1]) == 1


using Test
using QCMaterial

@testset "hamiltonian.generate_ham_1d_hubbard" begin
    nsite = 2
    n_qubit = 2*nsite
    U = 1.0
    t = -0.01
    μ = 0.0

    ham = generate_ham_1d_hubbard(t, U, nsite, μ)

    ham_ref = FermionOperator()
    up1 = up_index(1)
    down1 = down_index(1)
    up2 = up_index(2)
    down2 = down_index(2)

    #クーロン斥力
    ham_ref += FermionOperator("$(up1)^ $(down1)^ $(up1) $(down1)", -U)
    ham_ref += FermionOperator("$(up2)^ $(down2)^ $(up2) $(down2)", -U)

    #ホッピング積分
    ham_ref += FermionOperator("$(up2)^ $(up1)", t)
    ham_ref += FermionOperator("$(up1)^ $(up2)", t)
    ham_ref += FermionOperator("$(down2)^ $(down1)", t)
    ham_ref += FermionOperator("$(down1)^ $(down2)", t)

    #ケミカルポテンシャルの項
    ham_ref += FermionOperator("$(up1)^  $(up1) ", -μ)
    ham_ref += FermionOperator("$(down1)^ $(down1)", -μ)
    ham_ref += FermionOperator("$(up2)^  $(up2) ", -μ)
    ham_ref += FermionOperator("$(down2)^ $(down2)", -μ)

    @test ham == ham_ref
end

@testset "hamiltonian.generate_impurity_dimer" begin
    nsite = 2
    U = 1.0
    V = 1.0
    μ = 1.0
    ε = 1.0

    ham = generate_impurity_ham_with_1imp_multibath(U, V, μ, ε, nsite)

    ham_ref = FermionOperator()
    up1 = up_index(1)
    down1 = down_index(1)
    up2 = up_index(2)
    down2 = down_index(2)

    #Coulomb
    ham_ref += FermionOperator("$(up1)^ $(down1)^ $(up1) $(down1)", -U)

    #hybridization
    ham_ref += FermionOperator("$(up2)^ $(up1)", -V)
    ham_ref += FermionOperator("$(up1)^ $(up2)", -V)
    ham_ref += FermionOperator("$(down2)^ $(down1)", -V)
    ham_ref += FermionOperator("$(down1)^ $(down2)", -V)

    #chemical potential
    ham_ref += FermionOperator("$(up1)^  $(up1) ", -μ)
    ham_ref += FermionOperator("$(down1)^ $(down1)", -μ)

    #bath energy level
    ham_ref += FermionOperator("$(up2)^  $(up2) ", ε)
    ham_ref += FermionOperator("$(down2)^ $(down2)", ε)

    @test ham == ham_ref
end

@testset "hamiltonian.generate_impurity_ham_with_1imp_3bath" begin
    nsite = 4
    U = 1.0
    V = 1.0
    μ = 1.0
    ε = 1.0

    ham = generate_impurity_ham_with_1imp_multibath(U, V, μ, ε, nsite)

    ham_ref = FermionOperator()
    up1 = up_index(1)
    down1 = down_index(1)
    up2 = up_index(2)
    down2 = down_index(2)
    up3 = up_index(3)
    down3 = down_index(3)
    up4 = up_index(4)
    down4 = down_index(4)

    #Coulomb
    ham_ref += FermionOperator("$(up1)^ $(down1)^ $(up1) $(down1)", -U)

    #hybridization
    ham_ref += FermionOperator("$(up2)^ $(up1)", -V)
    ham_ref += FermionOperator("$(up1)^ $(up2)", -V)
    ham_ref += FermionOperator("$(down2)^ $(down1)", -V)
    ham_ref += FermionOperator("$(down1)^ $(down2)", -V)

    ham_ref += FermionOperator("$(up3)^ $(up1)", -V)
    ham_ref += FermionOperator("$(up1)^ $(up3)", -V)
    ham_ref += FermionOperator("$(down3)^ $(down1)", -V)
    ham_ref += FermionOperator("$(down1)^ $(down3)", -V)

    ham_ref += FermionOperator("$(up4)^ $(up1)", -V)
    ham_ref += FermionOperator("$(up1)^ $(up4)", -V)
    ham_ref += FermionOperator("$(down4)^ $(down1)", -V)
    ham_ref += FermionOperator("$(down1)^ $(down4)", -V)

    #chemical potential
    ham_ref += FermionOperator("$(up1)^  $(up1) ", -μ)
    ham_ref += FermionOperator("$(down1)^ $(down1)", -μ)

    #bath energy level
    ham_ref += FermionOperator("$(up2)^  $(up2) ", ε)
    ham_ref += FermionOperator("$(down2)^ $(down2)", ε)

    ham_ref += FermionOperator("$(up3)^  $(up3) ", ε)
    ham_ref += FermionOperator("$(down3)^ $(down3)", ε)

    ham_ref += FermionOperator("$(up4)^  $(up4) ", ε)
    ham_ref += FermionOperator("$(down4)^ $(down4)", ε)

    @test ham == ham_ref
end


type Chain
  x_coords::Array{Float64, 1} #space discretization, first entry must be 0, and then increasing order
  w::Array{Float64, 1} #deviation at x_coords
  w_dot::Array{Float64, 1} #velocities at x_coordinates

  A::AbstractMatrix
  B::AbstractVector
end

#constructor
function Chain(x_coords::Array{Float64, 1})
  @assert(x_coords[1] == 0.0)
  @assert length(x_coords)>3 #3 because 2nd derivatives will be needed
  @assert issorted(x_coords)
  if length(x_coords)<201
    A, B = hc_fdiff_mat(x_coords)
  else
    A, B = hc_fdiff_mat_sparse(x_coords)
  end
  return Chain(x_coords, zeros(length(x_coords)),  zeros(length(x_coords)), A, B)
end

function simulate!(chain::Chain, F::Number, dt::Float64)
  #some parameters
  N = length(chain.x_coords) #number of discrete points in space

  #simulate one step (Runge Kutta 4)
  v = vcat(chain.w_dot, chain.w)
  v1 = chain.A*v+chain.B*F
  v2 = chain.A*(v+dt/2*v1)+chain.B*F
  v3 = chain.A*(v+dt/2*v2)+chain.B*F
  v4 = chain.A*(v+dt*v3)+chain.B*F

  v += dt/6 * (v1 + 2*v2 + 2*v3 + v4)

  #update chain object
  chain.w_dot = v[1:N]
  chain.w = v[N+1:end]
end

function hc_fdiff_mat_old(x_coords)
  const rho = 2.3
  const g = 9.81
  const M = 6.151
  const L = x_coords[end]
  const N = length(x_coords)
  p(x::Number) = rho*g*(L-x)

  #preallocate output
  OutA = zeros(2*N, 2*N)
  OutB = vcat(1/M, zeros(2*N-1))

  #add boundary condition of the car
  OutA[1, N+1] = -p(x_coords[1])/(M*(x_coords[2]-x_coords[1]))
  OutA[1, N+2] = p(x_coords[1])/(M*(x_coords[2]-x_coords[1]))

  #add inner equations
  for row = 2:N-1
    OutA[row, N+row-1] = ( p(x_coords[row])/((x_coords[row+1]-x_coords[row])*(x_coords[row]-x_coords[row-1])) - (p(x_coords[row+1])-p(x_coords[row-1]))/(x_coords[row+1]-x_coords[row-1]) )/rho

    OutA[row, N+row] = (-2*p(x_coords[row])/((x_coords[row+1]-x_coords[row])*(x_coords[row]-x_coords[row-1])) )/rho

    OutA[row, N+row+1] = (p(x_coords[row])/((x_coords[row+1]-x_coords[row])*(x_coords[row]-x_coords[row-1])) + (p(x_coords[row+1])-p(x_coords[row-1]))/(x_coords[row+1]-x_coords[row-1]) )/rho
  end

  #add boundary (double backwards Euler)
  row = N
  OutA[row, N+row-2] = ( p(x_coords[row-1])/((x_coords[row]-x_coords[row-1])*(x_coords[row-1]-x_coords[row-2])) )/rho
  OutA[row, N+row-1] = ( -p(x_coords[row-1])/((x_coords[row]-x_coords[row-1])*(x_coords[row-1]-x_coords[row-2])) - p(x_coords[row])/(x_coords[row]-x_coords[row-1])^2 )/rho
  OutA[row, N+row] = ( p(x_coords[row])/(x_coords[row]-x_coords[row-1])^2 )/rho

  #add I to the bottom left
  for row = 1:N
    OutA[row+N, row] = 1.0
  end

  return OutA, OutB
end

function hc_fdiff_mat(x_coords)
  const rho = 2.3 #almost not needed, got rid of it analytically
  const g = 9.81
  const M = 6.151
  const L = x_coords[end]
  const N = length(x_coords)
  p(x::Number) = g*(L-x)

  #preallocate output
  OutA = zeros(2*N, 2*N)
  OutB = vcat(1/M, zeros(2*N-1))

  #add boundary condition of the car
  OutA[1, N+1] = -p(x_coords[1])*rho/(M*(x_coords[2]-x_coords[1]))
  OutA[1, N+2] = p(x_coords[1])*rho/(M*(x_coords[2]-x_coords[1]))

  #add inner equations
  for row = 2:N-1
    OutA[row, N+row-1] = -g/(x_coords[row+1]-x_coords[row-1]) + p(x_coords[row])/((x_coords[row+1]-x_coords[row])*(x_coords[row]-x_coords[row-1]))

    OutA[row, N+row] = -2*p(x_coords[row])/((x_coords[row+1]-x_coords[row])*(x_coords[row]-x_coords[row-1]))

    OutA[row, N+row+1] = g/(x_coords[row+1]-x_coords[row-1]) + p(x_coords[row])/((x_coords[row+1]-x_coords[row])*(x_coords[row]-x_coords[row-1]))
  end

  #add boundary (same as row above)
  row = N-1
  OutA[row+1, N+row-1] = -g/(x_coords[row+1]-x_coords[row-1]) + p(x_coords[row])/((x_coords[row+1]-x_coords[row])*(x_coords[row]-x_coords[row-1]))
  OutA[row+1, N+row] = -2*p(x_coords[row])/((x_coords[row+1]-x_coords[row])*(x_coords[row]-x_coords[row-1]))
  OutA[row+1, N+row+1] = g/(x_coords[row+1]-x_coords[row-1]) + p(x_coords[row])/((x_coords[row+1]-x_coords[row])*(x_coords[row]-x_coords[row-1]))

  #row = N
  #OutA[row, N+row-1] = g/(x_coords[row]-x_coords[row-1])
  #OutA[row, N+row] = -g/(x_coords[row]-x_coords[row-1])

  #add I to the bottom left
  for row = 1:N-1
    OutA[row+N, row] = 1.0
  end
  OutA[2*N, N-1] = 1.0

  return OutA, OutB
end

function hc_fdiff_mat_sparse_old(x_coords)
  const rho = 2.3
  const g = 9.81
  const M = 6.151
  const L = x_coords[end]
  const N = length(x_coords)
  p(x::Number) = rho*g*(L-x)

  #preallocate output
  n_vals = N*4-1 #N*3-1+N
  I = Array{Int}(n_vals)
  J = Array{Int}(n_vals)
  Val = Array{Float64}(n_vals)

  #add boundary condition of the car
  I[1] = 1
  J[1] = N+1
  Val[1] = -p(x_coords[1])/(M*(x_coords[2]-x_coords[1]))

  I[2] = 1
  J[2] = N+2
  Val[2] = p(x_coords[1])/(M*(x_coords[2]-x_coords[1]))

  #add inner equations
  index = 3
  for row = 2:N-1
    I[index] = row
    J[index] = N+row-1
    Val[index] = ( p(x_coords[row])/((x_coords[row+1]-x_coords[row])*(x_coords[row]-x_coords[row-1])) - (p(x_coords[row+1])-p(x_coords[row-1]))/(x_coords[row+1]-x_coords[row-1]) )/rho
    index += 1

    I[index] = row
    J[index] = N+row
    Val[index] = (-2*p(x_coords[row])/((x_coords[row+1]-x_coords[row])*(x_coords[row]-x_coords[row-1])) )/rho
    index += 1

    I[index] = row
    J[index] = N+row+1
    Val[index] =  (p(x_coords[row])/((x_coords[row+1]-x_coords[row])*(x_coords[row]-x_coords[row-1])) + (p(x_coords[row+1])-p(x_coords[row-1]))/(x_coords[row+1]-x_coords[row-1]) )/rho
    index += 1
  end

  #add boundary (double backwards Euler)
  I[index] = N
  J[index] = 2*N-2
  Val[index] = ( p(x_coords[N-1])/((x_coords[N]-x_coords[N-1])*(x_coords[N-1]-x_coords[N-2])) )/rho
  index += 1

  I[index] = N
  J[index] = 2*N-1
  Val[index] = ( -p(x_coords[N-1])/((x_coords[N]-x_coords[N-1])*(x_coords[N-1]-x_coords[N-2])) - p(x_coords[N])/(x_coords[N]-x_coords[N-1])^2 )/rho
  index += 1

  I[index] = N
  J[index] = 2*N
  Val[index] = ( p(x_coords[N])/(x_coords[N]-x_coords[N-1])^2 )/rho
  index += 1

  #add I to the bottom left
  for row = 1:N
    I[index] = N+row
    J[index] = row
    Val[index] = 1.0
    index += 1
  end

  return sparse(I, J, Val, 2*N, 2*N), sparsevec([1], [1/M], 2*N)
end

function hc_fdiff_mat_sparse(x_coords)
  return hc_fdiff_mat_sparse_old(x_coords)
end


#using Distributions

function jfpGMM{T1<:Number, T2<:Number}(x::Array{Float64, 2},
  prop::Vector{Float64},
  mu::Array{Float64, 2},
  sigma::Array{Array, 1},
  k::Int64,
  df::T1,
  lambda::T2,
  itermax::Int64 = 300,
  tol::Float64 = 1e-06
  )
  # mu a k by d matrix
  # sigma a [d,d,k] array
  const n::Int64 = size(x, 1)
  const d::Int64 = size(x, 2)

  delta::Float64 = 1.0
  prop_old = copy(prop)
  mu_old = deepcopy(mu)
  sigma_old = deepcopy(sigma)

  stp::Int64 = 0
  pdf_est = Array{Float64,2}(n,k)
  prob0 = Array{Float64,2}(n,k)
  while stp < itermax
    # E step
    for i in 1:k
      @inbounds tmp_mu = mu_old[i,:]
      @inbounds tmp_sigma = sigma_old[i]
      @inbounds pdf_est[:,i] = pdf(MvNormal(tmp_mu, PDMat(cholfact(Hermitian(tmp_sigma)))),x')
      @inbounds prob0[:,i] = pdf_est[:,i].*prop_old[i]
    end
    h_est = prob0./sum(prob0,2)

    # M step
    mu_new = *(h_est',x)./sum(h_est,1)'
    sigma_new = deepcopy(sigma_old)
    for i in 1:k
      @inbounds sigma_new[i] = *((x-reshape(kron(mu[i,:], rep(1,times = n)), (n,d)))',
                      diagm(h_est[:,i]),
                      (x-reshape(kron(mu[i,:], rep(1,times = n)), (n,d))))./sum(h_est[:,i])
    end
    prop_new = vec((sum(h_est,1).-lambda*df)./(n-k*lambda*df))

    if any(prop_new .< 1e-04)
      prop_new[prop_new .< 1e-04] = 0
    end
    prop_new = prop_new./sum(prop_new)

    delta = sum(abs(prop_new .- prop_old))
    if delta < tol
      break
    end

    if any(prop_new .== 0)
        keepidx = find(x -> x > 0, prop_new)
        k = length(keepidx)

        prop_old = prop_new[keepidx]
        mu_old = mu_new[keepidx,:]
        sigma_old = sigma_new[keepidx]
        pdf_est = pdf_est[:,keepidx]
        prob0 = h_est[:,keepidx]
        h_est = h_est[:,keepidx]
        delta = 1
      else
        prop_old = prop_new
        mu_old = mu_new
        sigma_old = sigma_new
    end

    if k <= 1
      break
    end
    stp += 1
  end
  tag = mapslices(indmax, h_est, 2)
  return Dict("k" => k,
              "prop" => prop_old,
              "mu" => mu_old,
              "sigma" => sigma_old,
              "pdf_est" => pdf_est,
              "ll" => sum(log.(sum(prob0,2))),
              "cluster" => tag)
end


@testset "dfs" begin
    A = load_matrix_network("dfs_example")
    dfs(MatrixNetwork(A),1)

    n = 10
    O = ones(Int64,n-1)
    Z = zeros(Int64,n)
    A_ref = Tridiagonal(O,Z,O)
    B = sparse(Matrix(A_ref))
    @testset "SparseMatrixCSC" begin
        dfs(A,1)

        (d,dt,ft,pred) = dfs(B,1)
        @test d == collect(0:n-1)
        @test dt == collect(0:n-1)
        @test ft == collect(2*n-1:-1:n)
        @test pred == collect(0:n-1)
    end
    @testset "CSR" begin
        dfs(sparse_to_csr(A)..., 1)

        (d,dt,ft,pred) = dfs(sparse_to_csr(B)..., 1)
        @test d == collect(0:n-1)
        @test dt == collect(0:n-1)
        @test ft == collect(2*n-1:-1:n)
        @test pred == collect(0:n-1)
    end
    @testset "triplet" begin
        dfs(findnz(A)[1], findnz(A)[2], 1)

        (d,dt,ft,pred) = dfs(findnz(B)[1], findnz(B)[2], 1)
        @test d == collect(0:n-1)
        @test dt == collect(0:n-1)
        @test ft == collect(2*n-1:-1:n)
        @test pred == collect(0:n-1)
    end
end


"""
    solve_tsp(g, w; cutoff=Inf)

Given a graph `g` and an edgemap `w` containing the weight associated to each edge,
solves the Travelling Salesman Problem, returning the tour among all the cities with
minimal total weight.

Edges `e` with no associated value in `w` or with `w[e] > cutoff`
will not be considered for the optimal tour.

The algorithm uses a generic Integer Programming solver combined with a lazy
constraint augmentation procedure.

The package JuMP.jl and one of its supported solvers are required.

Returns a tuple `(status, W, tour)` containing:
- a solve `status` (indicating whether the problem was solved to optimality)
- the tototal weight `W` of the tour
- a vector `tour` containing the sequence of vertices in the optimal tour.

**Example**
```julia
pos = rand(30, 2)
g = CompleteGraph(30)
w = EdgeMap(g, e -> norm(pos[src(e)] - pos[dst(e)]))
status, W, tour = solve_tsp(g, w)
```
"""
function solve_tsp(g::G, w::AEdgeMap; cutoff=Inf, verb=false) where {G<:AGraph}
    h = G(nv(g))
    for e in edges(g)
        haskey(w, e) && w[e] <= cutoff && add_edge!(h, e)
    end
    return _solve_tsp(h, w, verb)
end

function _solve_tsp(g::AGraph, w::AEdgeMap, verb::Bool)
    elist = collect(edges(g))
    model = Model(with_optimizer(MIP_OPTIMIZER))
    @variable(model, y[elist], Bin)
    @objective(model, Min, sum(y[e]*w[e] for e in edges(g)))
    @constraint(model, c1[i=1:nv(g)], sum(y[sort(e)] for e in edges(g, i)) == 2)

    ## TODO add lazy callback when avalaible again in JuMP
    ## (see https://github.com/JuliaOpt/JuMP.jl/pull/1849)
    # addlazycallback(model,
    #     cb -> begin
    #         sol = getvalue(y)
    #         alltours = gettours(g, sol)
    #         verb && println("Found $(length(alltours)) tours of lengths $(length.(alltours)). Adding lazy constraints.")
    #         for tour in alltours
    #             length(tour) == nv(g) && break
    #             add_tour_constraint(cb, y, g, tour)
    #         end
    #     end)
    has_short_tours = true
    while has_short_tours
        optimize!(model)
        sol = Dict(e => value(y[e]) for e in elist)
        alltours = gettours(g, sol)
        if length(alltours) == 1
            @assert length(alltours[1]) == nv(g)
            has_short_tours = false
        else
            verb && println("@ Found $(length(alltours)) tours.")
            for tour in alltours
                verb && println("@@ Add constraint on tour of length $(length(tour)).")
                add_tour_constraint(model, y, g, tour)
            end
        end
    end
    status = termination_status(model)
    sol = Dict(e => value(y[e]) for e in elist)
    alltours = gettours(g, sol)
    cost = objective_value(model)
    return status, cost, alltours[1]
end

function gettours(g, sol; atol=1e-8)
    alltours = Vector{Vector{Int}}()
    notintour = [1:nv(g);]
    while !isempty(notintour)
        root = first(notintour)
        tour = Int[root]
        u  = root
        prev = -1
        while u != root || length(tour) == 1
            for e in out_edges(g, u)
                abs(sol[sort(e)]) < atol && continue
                v = dst(e)
                v == prev && continue
                push!(tour, v)
                u, prev = v, u
                break
            end
        end
        resize!(tour, length(tour)-1)
        @assert length(tour) > 2
        notintour = setdiff(notintour, tour)
        @assert length(notintour) > 2 || length(notintour) == 0 "notintour=$notintour"
        push!(alltours, tour)
    end
    @assert sum(length.(alltours)) == nv(g)
    return alltours
end

function add_tour_constraint(model, y, g, tour)
    aff = AffExpr()
    for (k, v) in enumerate(tour)
        vprev = k == 1 ? tour[end] : tour[k-1]
        vnext = k == length(tour) ? tour[1] : tour[k+1]
        for e in edges(g, v)
            if dst(e) ∉ [vprev,vnext]
                aff += y[sort(e)]
            end
        end
    end
    @constraint(model, aff >= 2)
end


# An implementation of DQGMRES for the solution of the square linear system Ax = b.
#
# This method is described in
#
# Y. Saad, Iterative methods for sparse linear systems.
# PWS Publishing Company, Boston, USA, 1996.
#
# Y. Saad and K. Wu, DQGMRES: a quasi minimal residual algorithm based on incomplete orthogonalization.
# Numerical Linear Algebra with Applications, Vol. 3(4), pp. 329--343, 1996.
#
# Alexis Montoison, <alexis.montoison@polymtl.ca>
# Montreal, August 2018.

export dqgmres

"""Solve the consistent linear system Ax = b using DQGMRES method.

DQGMRES algorithm is based on the incomplete Arnoldi orthogonalization process
and computes a sequence of approximate solutions with the quasi-minimal residual property.

This implementation allows a right preconditioner M.
"""
function dqgmres(A :: AbstractLinearOperator, b :: AbstractVector{T};
                 M :: AbstractLinearOperator=opEye(size(A,1)), atol :: Float64=1.0e-8,
                 rtol :: Float64=1.0e-6, itmax :: Int=0, memory :: Int=20,
                 verbose :: Bool=false) where T <: Number

  m, n = size(A)
  m == n || error("System must be square")
  length(b) == m || error("Inconsistent problem size")
  verbose && @printf("DQGMRES: system of size %d\n", n)

  # Initial solution x₀ and residual r₀.
  x = zeros(T, n)
  # Compute β.
  rNorm = @knrm2(n, b) # rNorm = ‖r₀‖₂
  rNorm == 0.0 && return x, SimpleStats(true, false, [rNorm], [], "x = 0 is a zero-residual solution")

  iter = 0
  itmax == 0 && (itmax = 2*n)

  rNorms = [rNorm;]
  ε = atol + rtol * rNorm
  verbose && @printf("%5d  %7.1e\n", iter, rNorm)

  # Set up workspace.
  mem = min(memory, itmax) # Memory.
  V = [zeros(n) for i = 1 : mem] # Preconditioned Krylov vectors, orthogonal basis for {r₀, AM⁻¹r₀, (AM⁻¹)²r₀, ..., (AM⁻¹)ᵐ⁻¹r₀}.
  P = [zeros(n) for i = 1 : mem] # Directions for x : Pₘ = Vₘ(Rₘ)⁻¹.
  s = zeros(mem)    # Last mem Givens sines used for the factorization QₘRₘ = Hₘ.
  c = zeros(mem)    # Last mem Givens cosines used for the factorization QₘRₘ = Hₘ.
  H = zeros(mem+2)  # Last column of the band hessenberg matrix Hₘ.
  # Each column has at most mem + 1 nonzero elements. hᵢ.ₘ is stored as H[m-i+2].
  # m-i+2 represents the indice of the diagonal where hᵢ.ₘ is located.
  # In addition of that, the last column of Rₘ is also stored in H.

  # Initial γ₁ and V₁.
  γₘ = rNorm # γₘ and γₘ₊₁ are the last components of gₘ, right-hand of the least squares problem min ‖ Hₘyₘ - gₘ ‖₂.
  @. V[1] = b / rNorm

  # The following stopping criterion compensates for the lag in the
  # residual, but usually increases the number of iterations.
  # solved = sqrt(max(1, iter-mem+1)) * |γₘ₊₁| ≤ ε
  solved = rNorm ≤ ε # less accurate, but acceptable.
  tired = iter ≥ itmax
  status = "unknown"

  while !(solved || tired)

    # Update iteration index.
    iter = iter + 1

    # Set position in circulars stacks.
    pos = mod(iter-1, mem) + 1 # Position corresponding to pₘ and vₘ in circular stacks P and V.
    next_pos = mod(iter, mem) + 1 # Position corresponding to vₘ₊₁ in the circular stack V.

    # Incomplete Arnoldi procedure.
    z = M * V[pos] # Forms pₘ
    w = A * z # Forms vₘ₊₁
    for i = max(1, iter-mem+1) : iter
      ipos = mod(i-1, mem) + 1 # Position corresponding to vᵢ in the circular stack V.
      diag = iter - i + 2
      H[diag] = @kdot(n, w, V[ipos]) # hᵢ.ₘ = < A * M⁻¹vₘ , vᵢ >
      @kaxpy!(n, -H[diag], V[ipos], w) # w ← w - hᵢ.ₘ * vᵢ
    end
    # Compute hₘ₊₁.ₘ and vₘ₊₁.
    H[1] = @knrm2(n, w) # hₘ₊₁.ₘ = ‖vₘ₊₁‖₂
    if H[1] ≠ 0.0 # hₘ₊₁.ₘ = 0 ⇒ "lucky breakdown"
      @. V[next_pos] = w / H[1] # vₘ₊₁ = w / hₘ₊₁.ₘ
    end
    # rₘ₋ₘₑₘ.ₘ ≠ 0 when m ≥ mem + 1
    if iter ≥ mem + 2
      H[mem+2] = 0.0 # hₘ₋ₘₑₘ.ₘ = 0
    end

    # Update the QR factorization of H.
    # Apply mem previous Givens reflections Ωᵢ.
    for i = max(1,iter-mem) : iter-1
      irot_pos = mod(i-1, mem) + 1 # Position corresponding to cᵢ and sᵢ in circular stacks c and s.
      diag = iter - i + 1
      next_diag = diag + 1
      H_aux        = c[irot_pos] * H[next_diag] + s[irot_pos] * H[diag]
      H[diag]      = s[irot_pos] * H[next_diag] - c[irot_pos] * H[diag]
      H[next_diag] = H_aux
    end

    # Compute and apply current Givens reflection Ωₘ.
    # [cₘ  sₘ] [ hₘ.ₘ ] = [ρₘ]
    # [sₘ -cₘ] [hₘ₊₁.ₘ]   [0 ]
    (c[pos], s[pos], H[2]) = sym_givens(H[2], H[1])
    γₘ₊₁ = s[pos] * γₘ
    γₘ   = c[pos] * γₘ

    # Compute the direction pₘ, the last column of Pₘ = Vₘ(Rₘ)⁻¹.
    for i = max(1,iter-mem) : iter-1
      ipos = mod(i-1, mem) + 1 # Position corresponding to pᵢ in the circular stack P.
      diag = iter - i + 2
      if ipos == pos
        # pₐᵤₓ ← -hₘ₋ₘₑₘ.ₘ * pₘ₋ₘₑₘ
        @kscal!(n, -H[diag], P[pos])
      else
        # pₐᵤₓ ← pₐᵤₓ - hᵢ.ₘ * pᵢ
        @kaxpy!(n, -H[diag], P[ipos], P[pos])
      end
    end
    # pₐᵤₓ ← pₐᵤₓ + M⁻¹vₘ
    @kaxpy!(n, 1.0, z, P[pos])
    # pₘ = pₐᵤₓ / hₘ.ₘ
    @. P[pos] = P[pos] / H[2]

    # Compute solution xₘ.
    # xₘ ← xₘ₋₁ + γₘ * pₘ
    @kaxpy!(n, γₘ, P[pos], x)

    # Update residual norm estimate.
    # ‖ Axₘ - b ‖₂ ≈ |γₘ₊₁|
    rNorm = abs(γₘ₊₁)
    push!(rNorms, rNorm)

    # Update γₘ.
    γₘ = γₘ₊₁

    # Update stopping criterion.
    solved = rNorm ≤ ε
    tired = iter ≥ itmax
    verbose && @printf("%5d  %7.1e\n", iter, rNorm)
  end
  verbose && @printf("\n")

  status = tired ? "maximum number of iterations exceeded" : "solution good enough given atol and rtol"
  stats = SimpleStats(solved, false, rNorms, T[], status)
  return (x, stats)
end


using ADSeismic
using ADCME
using PyPlot
using PyCall
using DelimitedFiles
using MAT
using Statistics
np = pyimport("numpy")
matplotlib.use("Agg")
use_gpu()

################### Generate synthetic data #####################
output_dir = "data/acoustic"
if !ispath(output_dir)
  mkpath(output_dir)
end
figure_dir = "figure/acoustic/rupture-time/"
if !ispath(figure_dir)
  mkpath(figure_dir)
end
result_dir = "result/acoustic/rupture-time/"
if !ispath(result_dir)
  mkpath(result_dir)
end
# if isfile(joinpath(result_dir, "loss.txt"))
#   rm(joinpath(result_dir, "loss.txt"))
# end

################### Generate gaussian source #####################
# reset_default_graph()

ap_sim = load_acoustic_model("models/layer-model-rupture.mat"; IT_DISPLAY=0)
rcv_sim = load_acoustic_receiver("models/layer-model-rupture.mat")

ix = matread("models/layer-model-rupture.mat")["source"][1]["ix"][:]
iy = matread("models/layer-model-rupture.mat")["source"][1]["iy"][:]
s =  matread("models/layer-model-rupture.mat")["source"][1]["shift"][:] .* ap_sim.param.DELTAT
amp = matread("models/layer-model-rupture.mat")["source"][1]["amp"][:]
a = 1/pi

rs = []
for i = 1:length(s)
  gauss = Gauss(ap_sim.param, a, s[i], amp[i])
  push!(rs, gauss)
end
rs = hcat(rs...)
src = [AcousticSource(ix, iy, rs)]

sess = Session(); init(sess)

srcv = run(sess, src[1].srcv)
vmax = maximum(srcv)
vmin = -vmax
t = collect(0:ap_sim.param.NSTEP)* ap_sim.param.DELTAT
n = collect(1:length(src[1].srci)).-0.5
n, t = np.meshgrid(n, t)
close("all")
pcolormesh(n, t, srcv, vmax=vmax, vmin=vmin, cmap="seismic")
xlabel("Location index")
ylabel("Time (s)")
savefig(joinpath(figure_dir, "ground_truth.png"))
writedlm(joinpath(result_dir, "ground_truth.txt"), srcv)

################### Generate synthetic data #####################
# reset_default_graph()

ap_sim = load_acoustic_model("models/layer-model-rupture.mat"; IT_DISPLAY=0)
# src = load_acoustic_source("models/layer-model-rupture.mat")
rcv_sim = load_acoustic_receiver("models/layer-model-rupture.mat")
Rs_ = compute_forward_GPU(ap_sim, src, rcv_sim)

sess = Session(); init(sess)
Rs = run(sess, Rs_)

for i = 1:length(src)
    writedlm(joinpath(output_dir, "rupture-time-r$i.txt"), Rs[i])
end

u = run(sess, ap_sim(src[1]).u)
visualize_file(u, ap_sim.param, dirs=figure_dir)

################### Inversion using Automatic Differentiation #####################
reset_default_graph()

ap_sim = load_acoustic_model("models/layer-model-rupture.mat"; IT_DISPLAY=0)
rcv_sim = load_acoustic_receiver("models/layer-model-rupture.mat")

ix = matread("models/layer-model-rupture.mat")["rupture"]["ix"][:]
iy = matread("models/layer-model-rupture.mat")["rupture"]["iy"][:]
nsrc = length(ix)

# create variables
s_ = Variable(-1.5*ones(nsrc))
s = ap_sim.param.NSTEP*ap_sim.param.DELTAT*sigmoid(s_)
amp_ = Variable(0.1*ones(nsrc))
amp = 1e7*sigmoid(amp_)
if ! @isdefined a
  a = 1/pi
end

rs = []
for i = 1:length(s)
  gauss = Gauss(ap_sim.param, a, s[i], amp[i])
  push!(rs, gauss)
end
rs = hcat(rs...)
src_ = [AcousticSource(ix, iy, rs)]

Rs = Array{Array{Float64,2}}(undef, length(src_))
for i = 1:length(src_)
    Rs[i] = readdlm(joinpath(output_dir, "rupture-time-r$i.txt"))
end

sess = Session(); init(sess)

vars = [s_, amp_]
losses, gs = compute_loss_and_grads_GPU(ap_sim, src_, rcv_sim, Rs, vars)
g = Array{PyObject}(undef, length(gs[1]))
for i = 1:length(gs[1])
  g[i] = sum([pp[i] for pp in gs])
end

loss = sum(losses)
sess = Session(); init(sess)

function sigmoid_(x)
  return 1 / (1 + exp(-x))
end

function callback(v, iter, loss)
    if iter%5==0
      # @show sigmoid_.(v[2])*700.0
      # @show sigmoid_.(v[3])*1e7
      close("all")
      t = collect(0:ap_sim.param.NSTEP)* ap_sim.param.DELTAT
      n = collect(1:length(ix)).-0.5
      n, t = np.meshgrid(n, t)
      pcolormesh(n, t, v[1], vmax=vmax, vmin=vmin, cmap="seismic")
      xlabel("Location index")
      ylabel("Time (s)")
      # if iter>=0
      #   title("Iteration = $iter")
      # end
      savefig(joinpath(figure_dir, "inv_$(lpad(iter,4,"0")).png"))
      writedlm(joinpath(result_dir, "inv_$(lpad(iter,5,"0")).txt"), v[1])
      open(joinpath(result_dir, "loss.txt"), "a") do io
        writedlm(io, loss)
      end
    else
      return
    end
end

vars_ = [src_[1].srcv, vars...]
# callback(run(sess, vars_), 0, run(sess, loss))
LBFGS!(sess, loss, vars=vars_, callback=callback)


# module optimization
# TODO:
# (1) make a wrapper for Optim package

export optimize!, StochasticGradient, RProp

abstract Optimizer

@doc doc"Loop over parameters and call the optimizer for each parameter set with non null gradient" ->
function optimize!(s::Optimizer, m::Layer)
  for v in parameters(m)
    if !isnull(v.gradient)
        optimize!(s, v)
    end
  end
end

@doc doc"Stochatic gradient" ->
type StochasticGradient <: Optimizer
    learningRate::Float64
    # learningRateDecay::Float64

    # weithDecay::Float64
    # momentum::Float64

    # @doc doc"dampening for momentum" ->
    # dampening::Float64

    # @doc doc"Whether we use Nesterov momentum"
    # nesterov::Bool

    StochasticGradient(learningRate) = new(learningRate)
end

function optimize!(s::StochasticGradient, p::ArrayParameters)
    # TODO : Implement momentum
    axpy!(-s.learningRate, get(p.gradient), p.values)
end

#
# RProp optimizer
#

@doc doc"RProp optimizer
- stepsize    : initial step size, common to all components
- etaplus     : multiplicative increase factor, > 1 (default 1.2)
- etaminus    : multiplicative decrease factor, < 1 (default 0.5)
- stepsizemax : maximum stepsize allowed (default 50)
- stepsizemin : minimum stepsize allowed (default 1e-6)
- niter       : number of iterations (default 1)
" ->
immutable RProp <: Optimizer
    stepsize::Float64

    etaplus::Float64
    etaminus::Float64

    stepsizemax::Float64
    stepsizemin::Float64

    function RProp()
        self = new()
        self.stepsize = 0.1
        self.etaplus = 1.2
        self.etaminus = 0.5
        self.stepsizemax = 50.0
        self.stepsizemin = 1e-6
        self
    end

    RProp(stepsize) = new(stepsize, 1.2, 0.5, 50., 1e-6)
    RProp() = new(0.1)
end

type RPropState
    delta
    stepsize

    sign::BitArray{2}

    function RPropState(l::Int, stepsize::Float64)
        self = new()
        self.delta = zeros(l)
        self.stepsize = fill(stepsize, l)
        self.sign = BitArray(l, 2)
        self
    end
end

const RProp_POSITIVE = bitpack([false, false])
const RProp_NEGATIVE = bitpack([false, true])
const RProp_ZERO = bitpack([true, true])

function optimize!(s::RProp, p::ArrayParameters)
    # initialize auxiliary storage
    if isnull(p.optimization_state)
        p.optimization_state = Nullable(RPropState(length(p.values), s.stepsize))
    end

    state = get(p.optimization_state)::RPropState
    gradient = get(p.gradient)

    @inbounds for i in eachindex(gradient)
        # Compute the new step size
        sign = gradient[i] > 0 ? RProp_POSITIVE : (gradient[i] < 0 ? RProp_NEGATIVE : RProp_ZERO)

        if sign != RProp_ZERO
            if sign == state.sign[i]
                state.stepsize[i] = min(state.stepsize[i] * s.etaplus, s.stepsizemax)
            else
                state.stepsize[i] = max(state.stepsize[i] * s.etaminus, s.stepsizemin)
            end
        end

        # Update weight
        p.values[i] -= state.stepsize[i] * gradient[i]
    end

end


### A Pluto.jl notebook ###
# v0.14.2

using Markdown
using InteractiveUtils

# ╔═╡ 794a5356-6e05-44b4-ad9d-375a1fa3613b
begin
	using PlutoUI
	using DataFrames
	using Statistics
	using CategoricalArrays
	using Test
end

# ╔═╡ b3dcf3f4-9d68-11eb-01a6-317fe0e0cd7e
md"""
## Introduction to DataFrames / 3 - Missing values

ref. [Introduction to DataFrames, Bogumił Kamiński](https://github.com/bkamins/Julia-DataFrames-Tutorial/blob/master/03_missingvalues.ipynb)

$(html"
© Pascal, April 2021
")
"""

# ╔═╡ d6993a5b-8158-4ce3-accb-8e8c514229b3
PlutoUI.TableOfContents(indent=true, depth=4, aside=true)

# ╔═╡ de97ded0-3952-4188-9a43-a8f383e1ae65
md"""
### Handling missing values

A `Julia` singleton type `Missing` allows us to deal with missing values:
"""

# ╔═╡ 624c9b0a-e7e8-4ec7-8802-dcf12a71117e
missing, typeof(missing)

# ╔═╡ 0921e011-b191-4b27-ac9e-17ea3bc0105f
md"""
Arrays automatically create an appropriate union type.
"""

# ╔═╡ 1fb36dcf-0b6a-4ab0-8c0f-0b2c7f23c3c9
begin
	x = [1, 2, missing, 3];
	typeof(x)
end

# ╔═╡ c64cd43a-8a17-4d2f-8d69-0db1073b7968
md"""
`ismissing` (predicate) function checks if a passed value is missing.
"""

# ╔═╡ 58aeff5c-3f12-4bbb-82b9-84bddb2159dd
ismissing.(x)   ## using broadcasting

# ╔═╡ 60f16698-3f4a-4a15-9366-10e7f801d9f5
md"""
We can extract the type combined with `Missing` from a `Union` via `nonmissingtype`

This is useful for arrays, see:
"""

# ╔═╡ 94ccb18d-4ea6-40f7-aa90-314d2c5704cc
eltype(x), nonmissingtype(eltype(x))

# ╔═╡ dcac6464-a916-4c57-9f8a-991b98baa4db
md"""
missing comparisons produce missing.
"""

# ╔═╡ 6e2620ac-736e-4531-b580-fd4bbf6112e7
missing == missing, missing != missing , missing < missing

# ╔═╡ a1f09f3b-e49d-49e2-8b82-aed03374b896
md"""
This is also true when missings are compared with values of other types.
"""

# ╔═╡ 123870b9-a79f-4fa7-b8b8-24e81a989ad9
1 == missing, π != missing , missing < √2

# ╔═╡ 92eccb60-fbf3-464f-a207-60dc2e23ebe6
md"""
`isequal`, `isless`, and `===` produce results of type `Bool`.

**Note that `missing` is considered greater than any numeric value.**
"""

# ╔═╡ 673e843d-0ce8-4cae-bc38-37aed282e4e6
isequal(missing, missing), missing === missing, isequal(1, missing), isless(1, missing)

# ╔═╡ cf1e4d30-2c35-49b9-a917-55df01522788
md"""
In the next few examples, we see that many (*not all!*) functions handle missing.
"""

# ╔═╡ ee9e58dd-2d41-4fc8-81f1-8b9c348df8bc
map(x -> x(missing), [sin, cos, zero, sqrt])   ## part 1

# ╔═╡ 1f71ca99-9b15-43e4-8afc-15ffc5d4a1a1
map(x -> x(missing, 1), [+, - , *, /, div])    ## part 2 - arithmetic ops

# ╔═╡ d6c88a1c-80b3-49c0-84af-716c28974f42
map(x -> x([1, 2, missing]),
	[minimum, maximum, extrema, mean, float])  ## part 3

# ╔═╡ fb6deb6c-c503-4a9d-8a4c-7f1c4c24ca76
md"""
`skipmissing` returns iterator skipping missing values. We can use `collect` and `skipmissing` to create an array that excludes these missing values."
"""

# ╔═╡ 9d7df124-4942-44c2-807b-f5c39d14ec14
skipmissing([1, missing, 2, missing]) |> collect

# ╔═╡ e8ffd09a-1944-4d0f-926d-0fb1f9dda7a4
md"""
Here we use `replace` to create a new array that replaces all missing values with some value (`NaN` in this case).
"""

# ╔═╡ a108cd54-45f7-48ec-a130-441dc7cf0f5e
replace([1.0, missing, 2.0, missing], missing => NaN)

# ╔═╡ 4049c78a-2250-42f1-b548-6fd0925711bf
md"""
Another way to do this:
"""

# ╔═╡ aeacb14c-db6f-4363-948c-00588ccb4583
coalesce.([1.0, missing, 2.0, missing], NaN)

# ╔═╡ 7e385236-ff48-4f6d-af46-924c727c25d7
md"""
We can also use `recode` from `CategoricalArrays.jl` if we have a default output value.
"""

# ╔═╡ 8fa0b5c6-22a4-46ba-ac06-4c1c71a0304a
recode([1.0, missing, 2.0, missing], false, missing => true)

# ╔═╡ e9eb3ad3-9278-4d5c-8de8-4821939080c6
md"""
There are also `replace!` and `recode!` functions that work in place.

Here is an example how you can to missing inputation in a data frame.
"""

# ╔═╡ 56a979fe-dc5e-4416-8f85-47f1f9db0285
df = DataFrame(a=[1, 2, missing], b=["a", "b", missing])

# ╔═╡ d32c9670-c8f9-46a5-ac54-4f9639325a5c
md"""
we change `df.a` vector in place:
"""

# ╔═╡ 9c8574fe-c181-4017-aada-00508d8360e3
replace!(df.a, missing => 999)

# ╔═╡ cf55fe7e-c5b9-46f0-815b-785f1fc64a95
md"""
Now we overwrite `df.b` with a new vector, because the replacement type is different than what `eltype(df.b)` accepts:
"""

# ╔═╡ 9f1fb12a-39bb-479a-b115-5364c7ea999d
df.b = coalesce.(df.b, 100)

# ╔═╡ 53c20047-c44d-44b9-93a6-9b093c326c0b
df

# ╔═╡ 38f2f006-3d80-4190-bbf8-e87c8400dbb4
md"""
We can use `unique` or `levels` to get unique values with or without missings, respectively.
"""

# ╔═╡ 2284ff98-9303-4d98-9409-34b2ac28117d
unique([1, missing, 2, missing]), levels([1, missing, 2, missing])

# ╔═╡ 47687957-622a-4ca4-9b45-361b011d4af9
md"""
In this next example, we convert x to y with `allowmissing`, where y has a type that accepts missings.
"""

# ╔═╡ 82a421c5-2b80-44d5-9bfb-242877ea624d
begin
	x₁ = [1, 2, 3]
	y₁ = allowmissing(x₁)
	typeof(y₁)
end

# ╔═╡ 7080eca8-faba-4c70-a7ba-a843d3ccb459
md"""
Then, we convert back with `disallowmissing`. This would fail if `y` contained missing values!
"""

# ╔═╡ 29010c54-cd2f-41e8-81fc-5354e4ca7f6e
begin
	z₁ = disallowmissing(y₁)
	x₁, y₁, z₁
end

# ╔═╡ 7021180c-01a1-4565-be9d-53e72ffee9bc
md"""
`disallowmissing` has an `error` keyword argument that can be used to decide how it should behave when it encounters a column that actually contains a missing value.
"""

# ╔═╡ ae893d9e-9ea5-4ea9-94d8-d646b8699088
df₁ = allowmissing(DataFrame(ones(2, 3), :auto))

# ╔═╡ 8d7b3298-8f14-4d21-a569-04c61ade9bda
df₁[1, 1] = missing

# ╔═╡ 9f8c6697-33bf-4362-96b1-98ab5802dd9e
df₁

# ╔═╡ e1a6c188-2299-494e-92a7-ab4556a59c85
@test_throws MethodError disallowmissing(df₁)    ## an error is thrown

# ╔═╡ 140fe595-4d87-4282-b880-b7d73dd2ecdb
disallowmissing(df₁, error=false) # column :x1 is left untouched as it contains missing

# ╔═╡ 1dd4fba9-c80d-44c0-8a6d-ff88a6dd9545
md"""
In this next example, we show that the type of each column in x is initially `Int64`. After using `allowmissing!` to accept missing values in columns 1 and 3, the types of those columns become `Union{Int64,Missing}`.
"""

# ╔═╡ 3bcca454-4077-4714-b6c9-a2c50f0a7748
with_terminal() do
	x₂ = DataFrame(rand(Int, 2,3), :auto)
	println("Before: ", eltype.(eachcol(x₂)))

	allowmissing!(x₂, 1)       # make first column accept missings
	allowmissing!(x₂, :x3)     # make :x3 column accept missings

	println("After: ", eltype.(eachcol(x₂)))
end

# ╔═╡ 5d4237f4-8bcc-47cf-a21d-8fe3c56e5ec1
md"""
In this next example, we will use `completecases` to find all the rows of a DataFrame that have complete data.
"""

# ╔═╡ 6a3921f9-ceb7-47e2-a042-574dc01d7951
with_terminal() do
	x₂ = DataFrame(A=[1, missing, 3, 4], B=["A", "B", missing, "C"])
	println("Complete cases:\n", completecases(x₂))
end

# ╔═╡ 99efee9e-d7b5-4bdb-ad60-f6613545bbd9
md"""
We can use `dropmissing` or `dropmissing!` to remove the rows with incomplete data from a DataFrame and either create a new DataFrame or mutate the original in-place.
"""

# ╔═╡ bbffa430-3ae8-42f8-80cc-04c646c3ae89
begin
	x₂ = DataFrame(A=[1, missing, 3, 4], B=["A", "B", missing, "C"])
	y₂ = dropmissing(x₂)
	dropmissing!(x₂), describe(x₂)
end

# ╔═╡ 5fe7ddc1-b36d-4757-a213-d7a5a437ef14
md"""
Alternatively you can pass `disallowmissing` keyword argument to `dropmissing` and `dropmissing!`
"""

# ╔═╡ 5bb6d14f-af90-4a46-ab25-168cf7e57670
x₃ = DataFrame(A=[1, missing, 3, 4], B=["A", "B", missing, "C"])

# ╔═╡ 47bb179b-f865-4d48-8e4f-317af7fbf590
dropmissing!(x₃, disallowmissing=false)

# ╔═╡ ca9927e5-3ff2-4c73-8206-4d7495668312
md"""
### Making functions missing-aware

If we have a function that does not handle missing values we can wrap it using `passmissing` function so that if any of its positional arguments is missing we will get a missing value in return.

In the example below we change how string function behaves:
"""

# ╔═╡ 13e1f76a-baac-4a79-b150-3be8af1d2d54
string(missing)

# ╔═╡ a6d55bb4-8c4e-4eee-8452-df6510d3dcc4
string(missing, " ", missing)

# ╔═╡ e6d047c6-40d4-4bb4-b566-1878d59f3172
string(1, 2, 3)

# ╔═╡ 8c7187f7-ad7d-47e2-8152-26969ed66f58
lift_string = passmissing(string)

# ╔═╡ d8a357b2-a6f3-4aa4-978b-b7d5967efae9
lift_string(missing)  ## No longer string "missing" but instead the value missing

# ╔═╡ 1d4e483c-4217-49fa-bee1-6c4f729dd823
lift_string(missing, " ", missing)

# ╔═╡ adc9610d-01f6-4060-97d2-9b3a4d0970d5
lift_string(1, 2, 3)

# ╔═╡ a4615e9b-d3d9-443a-9146-e411f9bbd749
md"""
### Aggregating rows containing missing values

Create an example DataFrame containing missing values:
"""

# ╔═╡ 19982031-e9a0-4a17-aad2-18d816096f45
df₂ = DataFrame(a=[1, missing, missing], b=[1, 2, missing])

# ╔═╡ 609834b7-c4fc-4f5e-9ec3-f442e084c94c
md"""
If we just run `sum` on the rows we get two missing entries:
"""

# ╔═╡ a27323a5-d651-41d0-a588-5d5fd9395522
sum.(eachrow(df₂))

# ╔═╡ 9c886a50-b929-4ef5-bbae-a727d7b691b1
md"""
We can apply `skipmissing` on the rows to avoid this problem
"""

# ╔═╡ 1fc8067a-da59-4e6b-a8e6-7e32fcad49dd
@test_throws ArgumentError sum.(skipmissing.(eachrow(df₂)))  ## However this triggers another problem

# ╔═╡ 372b242d-81eb-4b99-84e1-7cc79f244a1c
md"""
We get an error. The problem is that the last row of `df₂` contains only missing values, and since eachrow is type unstable the `eltype` of the result of `skipmissing` is unknown (so it is marked `Any`), see:
"""

# ╔═╡ e6a28a9e-942c-407d-9b99-da9bf1e594ff
skipmissing(eachrow(df₂)[end]) |> collect

# ╔═╡ 85063cfc-9992-4530-bf15-5d101b64e170
md"""
In such cases it is useful to switch to `Tables.namedtupleiterator` which is type stable as discussed in `01_constructors.ipynb` notebook.
"""

# ╔═╡ 77330f55-8ec6-492b-9e08-00dfcbb004eb
Tables.namedtupleiterator(df₂) .|>
	skipmissing .|>
    sum

# ≡ sum.(skipmissing.(Tables.namedtupleiterator(df)))

# ╔═╡ c54cdbf1-7787-4f0f-b9ff-d4389f4a7c5e
html"""

"""

# ╔═╡ Cell order:
# ╟─b3dcf3f4-9d68-11eb-01a6-317fe0e0cd7e
# ╠═794a5356-6e05-44b4-ad9d-375a1fa3613b
# ╟─d6993a5b-8158-4ce3-accb-8e8c514229b3
# ╟─de97ded0-3952-4188-9a43-a8f383e1ae65
# ╠═624c9b0a-e7e8-4ec7-8802-dcf12a71117e
# ╟─0921e011-b191-4b27-ac9e-17ea3bc0105f
# ╠═1fb36dcf-0b6a-4ab0-8c0f-0b2c7f23c3c9
# ╟─c64cd43a-8a17-4d2f-8d69-0db1073b7968
# ╠═58aeff5c-3f12-4bbb-82b9-84bddb2159dd
# ╟─60f16698-3f4a-4a15-9366-10e7f801d9f5
# ╠═94ccb18d-4ea6-40f7-aa90-314d2c5704cc
# ╟─dcac6464-a916-4c57-9f8a-991b98baa4db
# ╠═6e2620ac-736e-4531-b580-fd4bbf6112e7
# ╟─a1f09f3b-e49d-49e2-8b82-aed03374b896
# ╠═123870b9-a79f-4fa7-b8b8-24e81a989ad9
# ╟─92eccb60-fbf3-464f-a207-60dc2e23ebe6
# ╠═673e843d-0ce8-4cae-bc38-37aed282e4e6
# ╟─cf1e4d30-2c35-49b9-a917-55df01522788
# ╠═ee9e58dd-2d41-4fc8-81f1-8b9c348df8bc
# ╠═1f71ca99-9b15-43e4-8afc-15ffc5d4a1a1
# ╠═d6c88a1c-80b3-49c0-84af-716c28974f42
# ╟─fb6deb6c-c503-4a9d-8a4c-7f1c4c24ca76
# ╠═9d7df124-4942-44c2-807b-f5c39d14ec14
# ╟─e8ffd09a-1944-4d0f-926d-0fb1f9dda7a4
# ╠═a108cd54-45f7-48ec-a130-441dc7cf0f5e
# ╟─4049c78a-2250-42f1-b548-6fd0925711bf
# ╠═aeacb14c-db6f-4363-948c-00588ccb4583
# ╟─7e385236-ff48-4f6d-af46-924c727c25d7
# ╠═8fa0b5c6-22a4-46ba-ac06-4c1c71a0304a
# ╟─e9eb3ad3-9278-4d5c-8de8-4821939080c6
# ╠═56a979fe-dc5e-4416-8f85-47f1f9db0285
# ╟─d32c9670-c8f9-46a5-ac54-4f9639325a5c
# ╠═9c8574fe-c181-4017-aada-00508d8360e3
# ╟─cf55fe7e-c5b9-46f0-815b-785f1fc64a95
# ╠═9f1fb12a-39bb-479a-b115-5364c7ea999d
# ╠═53c20047-c44d-44b9-93a6-9b093c326c0b
# ╟─38f2f006-3d80-4190-bbf8-e87c8400dbb4
# ╠═2284ff98-9303-4d98-9409-34b2ac28117d
# ╟─47687957-622a-4ca4-9b45-361b011d4af9
# ╠═82a421c5-2b80-44d5-9bfb-242877ea624d
# ╟─7080eca8-faba-4c70-a7ba-a843d3ccb459
# ╠═29010c54-cd2f-41e8-81fc-5354e4ca7f6e
# ╟─7021180c-01a1-4565-be9d-53e72ffee9bc
# ╠═ae893d9e-9ea5-4ea9-94d8-d646b8699088
# ╠═8d7b3298-8f14-4d21-a569-04c61ade9bda
# ╠═9f8c6697-33bf-4362-96b1-98ab5802dd9e
# ╠═e1a6c188-2299-494e-92a7-ab4556a59c85
# ╠═140fe595-4d87-4282-b880-b7d73dd2ecdb
# ╟─1dd4fba9-c80d-44c0-8a6d-ff88a6dd9545
# ╠═3bcca454-4077-4714-b6c9-a2c50f0a7748
# ╟─5d4237f4-8bcc-47cf-a21d-8fe3c56e5ec1
# ╠═6a3921f9-ceb7-47e2-a042-574dc01d7951
# ╟─99efee9e-d7b5-4bdb-ad60-f6613545bbd9
# ╠═bbffa430-3ae8-42f8-80cc-04c646c3ae89
# ╟─5fe7ddc1-b36d-4757-a213-d7a5a437ef14
# ╠═5bb6d14f-af90-4a46-ab25-168cf7e57670
# ╠═47bb179b-f865-4d48-8e4f-317af7fbf590
# ╟─ca9927e5-3ff2-4c73-8206-4d7495668312
# ╠═13e1f76a-baac-4a79-b150-3be8af1d2d54
# ╠═a6d55bb4-8c4e-4eee-8452-df6510d3dcc4
# ╠═e6d047c6-40d4-4bb4-b566-1878d59f3172
# ╠═8c7187f7-ad7d-47e2-8152-26969ed66f58
# ╠═d8a357b2-a6f3-4aa4-978b-b7d5967efae9
# ╠═1d4e483c-4217-49fa-bee1-6c4f729dd823
# ╠═adc9610d-01f6-4060-97d2-9b3a4d0970d5
# ╟─a4615e9b-d3d9-443a-9146-e411f9bbd749
# ╠═19982031-e9a0-4a17-aad2-18d816096f45
# ╟─609834b7-c4fc-4f5e-9ec3-f442e084c94c
# ╠═a27323a5-d651-41d0-a588-5d5fd9395522
# ╟─9c886a50-b929-4ef5-bbae-a727d7b691b1
# ╠═1fc8067a-da59-4e6b-a8e6-7e32fcad49dd
# ╟─372b242d-81eb-4b99-84e1-7cc79f244a1c
# ╠═e6a28a9e-942c-407d-9b99-da9bf1e594ff
# ╟─85063cfc-9992-4530-bf15-5d101b64e170
# ╠═77330f55-8ec6-492b-9e08-00dfcbb004eb
# ╟─c54cdbf1-7787-4f0f-b9ff-d4389f4a7c5e


@testset "Reynolds_Number" begin
    Tair,pressure,ustar,z0m = 25.0,100-0,0.5,0.5
    R = @inferred Reynolds_Number(Tair,pressure,ustar,z0m)
    @test ≈(R, 15870, rtol=1e-3)
end

@testset "roughness_parameters" begin
    zh = thal.zh
    zr = thal.zr
    LAI = thal.LAI
    keys_exp = (:d, :z0m, :z0m_se)
    rp = @inferred roughness_parameters(RoughnessCanopyHeight(), zh)
    #round.(values(rp); sigdigits = 4)
    @test keys(rp) == keys_exp
    @test all(isapproxm.(values(rp), (18.55, 2.65, missing), rtol=1e-3))
    #
    rp = @inferred roughness_parameters(RoughnessCanopyHeightLAI(), zh, LAI)
    #round.(values(rp); sigdigits = 4)
    @test keys(rp) == keys_exp
    @test all(isapproxm.(values(rp), (21.77, 1.419, missing), rtol=1e-3))
    #
    df = copy(tha48)
    #df.wind[1] = missing # obscures comparing to R
    d=0.7*zh
    psi_m = stability_correction!(df; z=zr, d).psi_m
    # note: must use columntable for type stability - but needs compilation timede
    # not type-stable if some columns allow missings in Julia 1.6
    #rp = @inferred roughness_parameters(Roughness_wind_profile(), df.ustar, df.wind, psi_m; zh, zr)
    rp = roughness_parameters(Roughness_wind_profile(), df.ustar, df.wind, psi_m; zh, zr)
    @test keys(rp) == keys_exp
    #@test all(isapproxm.(values(rp), (18.55, 1.879, 0.3561), rtol=1e-3))
    #from R:
    @test all(isapproxm.(values(rp), (18.55, 1.879402, 0.356108), rtol=1e-3))
    #
    # no stability correction
    # broadcast across Missings is not type-stable
    # rp0 = @inferred roughness_parameters(Roughness_wind_profile(), df.ustar, df.wind,
    #     df.Tair, df.pressure, df.H; zh, zr,
    #     stab_formulation = NoStabilityCorrection())
    dfd = disallowmissing(df[!,Not(:LE)])
    rp0 = @inferred roughness_parameters(Roughness_wind_profile(), dfd.ustar, dfd.wind,
        dfd.Tair, dfd.pressure, dfd.H; zh, zr,
        stab_formulation = NoStabilityCorrection())
    @test keys(rp0) == keys_exp
    # same magnitude as with stability correction
    @test all(isapproxm.(values(rp0), values(rp), rtol=0.5))
    # providing DataFrame is not type stable
    rp0b = roughness_parameters(Roughness_wind_profile(), df; zh, zr,
        stab_formulation = NoStabilityCorrection())
    @test rp0b == rp0
    #
    # estimate psi
    #@code_warntype stability_correction(columntable(df), zr, 0.7*zh)
    #@code_warntype roughness_parameters(Roughness_wind_profile(), columntable(df), zh, zr)
    rp_psiauto = roughness_parameters(Roughness_wind_profile(), df; zh, zr)
    @test rp_psiauto == rp
    #
    # DataFrame with only columns ustar and wind
    rp0c = roughness_parameters(Roughness_wind_profile(), df[!,Cols(:ustar, :wind)]; zh, zr,
        stab_formulation = NoStabilityCorrection())
    @test rp0c == rp0
end

@testset "wind_profile" begin
    datetime, ustar, Tair, pressure, H = values(tha48[1,:])
    z = 30
    d=0.7*thal.zh
    z0m= 2.65
    u30 = @inferred wind_profile(z, ustar, d, z0m, 0.0)
    @test ≈(u30, 1.93, rtol = 1/100 ) # from R
    #
    u30c = @inferred wind_profile(z, ustar, d, z0m, Tair,pressure, H)
    @test ≈(u30c, 2.31, rtol = 1/100 ) # from R
    #
    z0m=1.9 #2.14 #2.65
    u30 = @inferred wind_profile(z, ustar, d, z0m, 0.0) # used below
    u30c = @inferred wind_profile(z, ustar, d, z0m, Tair,pressure, H)
    df = copy(tha48)
    dfd = disallowmissing(df[!,Not(:LE)])
    windz = @inferred wind_profile(z, columntable(dfd), d, z0m; stab_formulation = NoStabilityCorrection())
    windz = wind_profile(z, df, d, z0m; stab_formulation = NoStabilityCorrection())
    @test length(windz) == 48
    @test windz[1] == u30
    windzc = @inferred wind_profile(z, columntable(dfd), d, z0m; stab_formulation = Dyer1970())
    @test windzc[1] == u30c
    #plot(windz)
    #plot!(windz2)
    #
    # with providing psi_m
    psi_m = stability_correction(df; z, d).psi_m
    windzc2 = @inferred wind_profile(z, columntable(dfd), d, z0m; psi_m)
    @test windzc2 == windzc
    #
    # with providing L
    MOL = Monin_Obukhov_length.(df.Tair, df.pressure, df.ustar, df.H)
    windzc3 = @inferred wind_profile(z, columntable(dfd), d, z0m; MOL)
    @test windzc3 == windzc
end


import Base.sum

"""
    sum(N::AbstractEcologicalNetwork)

This function will return the sum of all interactions in the network. For
quantitative networks, this is the sum of interaction strengths. For binary
networks, this is the number of interactions. For probabilistic networks, this
is the expected number of realized interactions.
"""
function sum(N::AbstractEcologicalNetwork)
   return sum(N.A)
end

"""
    links(N::BinaryNetwork)

Number of non-zero interactions in a deterministic network.
"""
function links(N::BinaryNetwork)
    return sum(N)
end

"""
    links(N::QuantitativeNetwork)

Number of non-zero interactions in a quantitative network (use `sum` to get the
sum of interaction strengths).
"""
function links(N::QuantitativeNetwork)
    return sum(N.A .!= zero(eltype(N.A)))
end

"""
    links(N::ProbabilisticNetwork)

Expected number of interactions in a probabilistic network. To get the number of
interactions that have a non-zero probability, use *e.g.* `links(N>0.0)`.
"""
function links(N::ProbabilisticNetwork)
    return sum(N)
end

"""
**Variance in the expected number of links**

    links_var(N::ProbabilisticNetwork)

Expected variance of the number of links for a probabilistic network.
"""
function links_var(N::ProbabilisticNetwork)
   return sum(N.A .* (1 .- N.A))
end

"""
    connectance(N::AbstractEcologicalNetwork)

Number of links divided by the number of possible interactions. In unipartite
networks, this is ``L/S^2``. In bipartite networks, this is ``L/(T × B)``. It is
worth noting that while the maximal connectance is always 1 (i.e. the graph is
complete), the minimum value (assuming that the network is not degenerate) is
*not* 0. Instead, the minimum number of interactions in a unipartite network is
`S-1`, and in a bipartite network it is `max(T,B)`.

Connectance can therefore be transformed between 0 and 1, using the following
approach: let `m` be the minimum number of interactions, and Co be the measured
connectance, then the corrected value is `(Co-m)/(1-m)`. To our best knowledge,
this is not standard practice, and therefore is not suggested as a function in
the package.
"""
function connectance(N::AbstractEcologicalNetwork)
    return links(N) / (richness(N; dims=1)*richness(N; dims=2))
end

"""
    linkage_density(N::AbstractEcologicalNetwork)

Number of links divided by species richness.
"""
function linkage_density(N::AbstractEcologicalNetwork)
    return links(N) / richness(N)
end

"""
    connectance_var(N::ProbabilisticNetwork)

Expected variance of the connectance for a probabilistic matrix, measured as the
variance of the number of links divided by the squared size of the matrix.
"""
function connectance_var(N::ProbabilisticNetwork)
   return links_var(N) / (prod(size(N))^2)
end


using BAT
using Test
using Statistics
using StatsBase
using Distributions
using HypothesisTests

@testset "Funnel Distribution" begin
    #Tests different forms of instantiate Funnel Distribution
    @test @inferred(BAT.FunnelDistribution(1., 2., 3)) isa BAT.FunnelDistribution
    @test @inferred(BAT.FunnelDistribution()) isa BAT.FunnelDistribution
    @test @inferred(BAT.FunnelDistribution(a = 1.0, b = 0.5, n = 3)) isa BAT.FunnelDistribution

    funnel = BAT.FunnelDistribution(a = Float32(1.0), b = Float32(0.5), n = Int32(3))

    #Check Parameters
    @test @inferred(StatsBase.params(funnel)) == (Float32(1.0), Float32(0.5), Int32(3))

    @test @inferred(Base.length(funnel)) == 3 #Number of dimensions
    @test @inferred(Base.eltype(funnel)) == Float32 #Element Type

    #Check mean and covariance
    @test @inferred(Statistics.mean(funnel)) == [0.0, 0.0, 0.0] #Check mean
    @test isapprox(@inferred(Statistics.cov(funnel)), [1.00256 -0.0124585 -0.00373376;
    -0.0124585 7.04822 -0.165097;  -0.00373376  -0.165097  7.1126], rtol = 1e-1, atol = 0)

    #logpdf
    @test isapprox(@inferred(Distributions._logpdf(funnel, [0., 0., 0.])), -2.75681, atol = 1e-5)

    #KS Test
    #Test the constant-variance Gaussian
    funnel = BAT.FunnelDistribution(a = 1., b = 0., n = 1)
    ks_test = HypothesisTests.ExactOneSampleKSTest(rand(funnel, 10^6)[:], Normal(0., 1.))
    @test pvalue(ks_test) > 0.05

    #Test the variable-variance Gaussian
    funnel = BAT.FunnelDistribution(a = 0., b = 0., n = 2)
    ks_test = HypothesisTests.ExactOneSampleKSTest(rand(funnel, 10^6)[2,:], Normal(0., 1.))
    @test pvalue(ks_test) > 0.05

end


let

    vec_A = Array(Matrix{Float64}, 2)
    vec_b = Array(Vector{Float64}, 2)
    vec_A[1] = [0.5,]'
    vec_A[2] = [2.0,]'
    vec_b[1] = [2.0]
    vec_b[2] = [0.5]
    mix_obs = MixtureModel([MvNormal([0.0], [1.0]), MvNormal([1.0], [1.0])],[0.5,0.5])
    mix_AgO = MixtureModel([MvNormal([NaN], [1.0]), MvNormal([NaN], [3.0])])

    gmr = GMR(vec_A, vec_b, mix_obs, mix_AgO)

    show(IOBuffer(), gmr)

    pred = gmr([0.0])
    weights = probs(pred.prior)
    @test isapprox(weights[1], 0.6225, atol=1e-4)
    @test isapprox(weights[2], 0.3775, atol=1e-4)
    @test isapprox(pred.components[1].μ[1], 2.0, atol=1e-8)
    @test isapprox(pred.components[2].μ[1], 0.5, atol=1e-8)
    @test full(pred.components[1].Σ)[1] == 1.0
    @test full(pred.components[2].Σ)[1] == 9.0

    mix = MixtureModel([
            MvNormal([0.0, 0.0], [1.0, 1.0]),
            MvNormal([1.0, 1.0], [1.0, 1.0]),
        ], [0.25,0.75]
        )
    gmr = GMR(mix, 1)
    @test isapprox(gmr.vec_A[1][1], 0.0, atol=1e-8)
    @test isapprox(gmr.vec_A[2][1], 0.0, atol=1e-8)
    @test isapprox(gmr.vec_b[1][1], 0.0, atol=1e-8)
    @test isapprox(gmr.vec_b[2][1], 1.0, atol=1e-8)
    @test isapprox(full(gmr.mixture_Act_given_Obs.components[1].Σ)[1], 1.0, atol=1e-8)
    @test isapprox(full(gmr.mixture_Act_given_Obs.components[2].Σ)[1], 1.0, atol=1e-8)
    @test probs(gmr.mixture_Obs.prior) == [0.25, 0.75]
    @test isapprox(gmr.mixture_Obs.components[1].μ[1], 0.0, atol=1e-8)
    @test isapprox(gmr.mixture_Obs.components[2].μ[1], 1.0, atol=1e-8)
    @test isapprox(full(gmr.mixture_Obs.components[1].Σ)[1], 1.0, atol=1e-8)
    @test isapprox(full(gmr.mixture_Obs.components[2].Σ)[1], 1.0, atol=1e-8)
end


"""
Basic grids including common grids like arbitrary grids, uniform grids, log grids,
and optimized grids like barycheb for interpolation and gausslegendre for integration.

"""
module SimpleG

export AbstractGrid, OpenGrid, ClosedGrid, Uniform, BaryCheb, GaussLegendre, Arbitrary, Log, denseindex

using StaticArrays, FastGaussQuadrature

include("chebyshev.jl")
export barychebinit, barycheb

"""
All Grids are derived from AbstractGrid; ClosedGrid has bound[1], bound[2] == grid[1], grid[end],
while OpenGrid has bound[1]<grid[1]<grid[end]<bound[2]
"""
abstract type AbstractGrid end
abstract type OpenGrid <: AbstractGrid end
abstract type ClosedGrid <: AbstractGrid end

"""

    struct Arbitrary{T<:AbstractFloat} <: ClosedGrid

Arbitrary grid generated from given sorted grid.

#Members:
- `bound` : boundary of the grid
- `size` : number of grid points
- `grid` : grid points
- `weight` : integration weight

#Constructor:
-    function Arbitrary{T}(grid) where {T<:AbstractFloat}
"""
struct Arbitrary{T<:Real} <: ClosedGrid
    bound::SVector{2,T}
    size::Int
    grid::Vector{T}
    weight::Vector{Float64}

"""
    function Arbitrary{T}(grid) where {T<:AbstractFloat}

create Arbitrary from grid.
"""
    function Arbitrary{T}(grid) where {T<:Real}
        bound = [grid[1],grid[end]]
        size = length(grid)
        weight = zeros(Float64, size)
        for i in 1:size
            if i==1
                weight[1] = 0.5*(grid[2]-grid[1])
            elseif i==size
                weight[end] = 0.5*(grid[end]-grid[end-1])
            else
                weight[i] = 0.5*(grid[i+1]-grid[i-1])
            end
        end
        return new{T}(bound, size, grid, weight)
    end
end

"""
    function Base.floor(grid::AbstractGrid, x) #where {T}

use basic searchsorted function to find the index of largest
grid point smaller than x.

return 1 for x<grid[1] and grid.size-1 for x>grid[end].
"""
function Base.floor(grid::AbstractGrid, x) #where {T}
    if x <= grid.grid[1]
        return 1
    elseif x >= grid.grid[end]
        return grid.size-1
    end

    result = searchsortedfirst(grid.grid, x)-1
    return Base.floor(Int, result)
end

Base.length(grid::AbstractGrid) = grid.size
Base.size(grid::AbstractGrid) = grid.size
Base.show(io::IO, grid::AbstractGrid) = print(io, grid.grid)
Base.view(grid::AbstractGrid, inds...) where {N} = Base.view(grid.grid, inds...)
# set is not allowed for grids
Base.getindex(grid::AbstractGrid, i) = grid.grid[i]
Base.firstindex(grid::AbstractGrid) = 1
Base.lastindex(grid::AbstractGrid) = grid.size
# iterator
Base.iterate(grid::AbstractGrid) = (grid.grid[1],1)
Base.iterate(grid::AbstractGrid, state) = (state>=grid.size) ? nothing : (grid.grid[state+1],state+1)

"""
    struct Uniform{T<:AbstractFloat} <: ClosedGrid

Uniform grid generated on [bound[1], bound[2]] with N points

#Members:
- `bound` : boundary of the grid
- `size` : number of grid points
- `grid` : grid points
- `weight` : integration weight

#Constructor:
-    function Uniform{T}(bound, size) where {T<:AbstractFloat}
"""
struct Uniform{T<:AbstractFloat} <: ClosedGrid
    bound::SVector{2,T}
    size::Int
    grid::Vector{T}
    weight::Vector{T}

"""
    function Uniform{T}(bound, N) where {T<:AbstractFloat}

create Uniform grid.
"""
    function Uniform{T}(bound, N) where {T<:AbstractFloat}
        Ntot = N - 1
        interval = (bound[2]-bound[1])/Ntot
        grid = bound[1] .+ Vector(1:N) .* interval .- ( interval )
        weight = similar(grid)
        for i in 1:N
            if i==1
                weight[1] = 0.5*(grid[2]-grid[1])
            elseif i==N
                weight[end] = 0.5*(grid[end]-grid[end-1])
            else
                weight[i] = 0.5*(grid[i+1]-grid[i-1])
            end
        end
        return new{T}(bound, N, grid, weight)
    end
end

"""
    function Base.floor(grid::Uniform{T}, x) where {T}

find the index of largest
grid point smaller than x.

return 1 for x<grid[1] and grid.size-1 for x>grid[end].
"""
function Base.floor(grid::Uniform{T}, x) where {T}
    result = (x-grid.grid[1])/(grid.grid[end]-grid.grid[1])*(grid.size-1)+1
    if result <=0
        return 1
    elseif result >= grid.size
        return grid.size-1
    else
        return Base.floor(Int, result)
    end

end

"""
    struct BaryCheb{T<:AbstractFloat} <: OpenGrid

BaryCheb grid generated on [bound[1], bound[2]] with order N.

#Members:
- `bound` : boundary of the grid
- `size` : number of grid points
- `grid` : grid points
- `weight` : interpolation weight

#Constructor:
-    function BaryCheb{T}(bound, size) where {T<:AbstractFloat}
"""
struct BaryCheb{T<:AbstractFloat} <: OpenGrid
    bound::SVector{2,T}
    size::Int
    grid::Vector{T}
    weight::Vector{T}

"""
    function BaryCheb{T}(bound, N) where {T<:AbstractFloat}

create BaryCheb grid.
"""
    function BaryCheb{T}(bound, N) where {T<:AbstractFloat}
        order = N
        x, w =barychebinit(order)
        grid = zeros(T, N)
        a, b = bound[1], bound[2]
        weight = (b - a) / 2  .* w
        grid = (a + b) / 2 .+ (b - a) / 2 .* x

        return new{T}(bound, N, grid, weight)
    end
end

"""
    struct GaussLegendre{T<:AbstractFloat} <: OpenGrid

GaussLegendre grid generated on [bound[1], bound[2]] with order N.

#Members:
- `bound` : boundary of the grid
- `size` : number of grid points
- `grid` : grid points
- `weight` : integration weight

#Constructor:
-    function GaussLegendre{T}(bound, size) where {T<:AbstractFloat}
"""
struct GaussLegendre{T<:AbstractFloat} <: OpenGrid
    bound::SVector{2,T}
    size::Int
    grid::Vector{T}
    weight::Vector{T}

"""
    function GaussLegendre{T}(bound, N) where {T<:AbstractFloat}

create GaussLegendre grid.
"""
    function GaussLegendre{T}(bound, N) where {T<:AbstractFloat}
        order = N
        x, w = gausslegendre(order)
        grid = zeros(T, N)
        a, b = bound[1], bound[2]
        weight = (b - a) / 2  .* w
        grid = (a + b) / 2 .+ (b - a) / 2 .* x

        return new{T}(bound, N, grid, weight)
    end
end

"""
    struct Log{T<:AbstractFloat} <: ClosedGrid

Log grid generated on [bound[1], bound[2]] with N grid points.
Minimal interval is set to be minterval. Dense to sparse if d2s, vice versa.

On [0, 1], a typical d2s Log grid looks like
[0, λ^(N-1), ..., λ^2, λ, 1].

#Members:
- `bound` : boundary of the grid
- `size` : number of grid points
- `grid` : grid points
- `weight` : integration weight

- `λ` : scale parameter
- `d2s` : dense to sparse or not

#Constructor:
-    function Log{T}(bound, size, minterval, d2s) where {T<:AbstractFloat}
"""
struct Log{T<:AbstractFloat} <: ClosedGrid
    bound::SVector{2,T}
    size::Int
    grid::Vector{T}
    weight::Vector{T}

    λ::T
    d2s::Bool

"""
    function Log{T}(bound, N, minterval, d2s) where {T<:AbstractFloat}

create Log grid.
"""
    function Log{T}(bound, N, minterval, d2s) where {T<:AbstractFloat}
        grid = zeros(T, N)
        M = N-2
        λ = (minterval/(bound[2]-bound[1]))^(1.0/M)

        if d2s
            for i in 1:M
                grid[i+1] = bound[1] + (bound[2]-bound[1])*λ^(M+1-i)
            end
        else
            for i in 2:M+1
                grid[i] = bound[2] - (bound[2]-bound[1])*λ^(i-1)
            end
        end
        grid[1] = bound[1]
        grid[end] = bound[2]
        weight = similar(grid)
        for i in 1:N
            if i==1
                weight[1] = 0.5*(grid[2]-grid[1])
            elseif i==N
                weight[end] = 0.5*(grid[end]-grid[end-1])
            else
                weight[i] = 0.5*(grid[i+1]-grid[i-1])
            end
        end

        return new{T}(bound, N, grid,weight, λ, d2s)
    end
end

"""
    function Base.floor(grid::Log{T}, x) where {T}

find the index of largest
grid point smaller than x.

return 1 for x<grid[1] and grid.size-1 for x>grid[end].
"""
function Base.floor(grid::Log{T}, x) where {T}
    if x <= grid.grid[1]
        return 1
    elseif x >= grid.grid[end]
        return grid.size-1
    end

    a,b=grid.bound[1],grid.bound[2]
    if grid.d2s
        i = Base.floor( log((x-a)/(b-a))/log(grid.λ)  )
        if i > grid.size-2
            result = 1
        else
            result = grid.size-i-1
        end
    else
        i = Base.floor( log((b-x)/(b-a))/log(grid.λ)  )
        if i > grid.size-2
            result = grid.size-1
        else
            result = i+1
        end
    end

    return Base.floor(Int, result)
end

@inline function denseindex(grid::Log)
    return [(grid.d2s) ? 1 : grid.size, ]
end

end


export PeriodicCurve

struct PeriodicCurve{S<:Space,T,VT} <: PeriodicDomain{T}
    curve::Fun{S,T,VT}
end

==(a::PeriodicCurve, b::PeriodicCurve) = a.curve == b.curve
isempty(::PeriodicCurve) = false

points(c::PeriodicCurve, n::Integer) = c.curve.(points(domain(c.curve),n))

checkpoints(d::PeriodicCurve) = fromcanonical.(Ref(d),checkpoints(domain(d.curve)))

for op in (:(leftendpoint),:(rightendpoint),:(rand))
    @eval $op(c::PeriodicCurve) = c.curve($op(domain(c.curve)))
end

canonicaldomain(c::PeriodicCurve) = domain(c.curve)

fromcanonical(c::PeriodicCurve{S,T},x) where {S<:Space,T<:Number} = c.curve(x)
function tocanonical(c::PeriodicCurve,x)
    rts=roots(c.curve-x)
    @assert length(rts)==1
    first(rts)
end

fromcanonicalD(c::PeriodicCurve,x)=differentiate(c.curve)(x)

function indomain(x,c::PeriodicCurve)
    rts=roots(c.curve-x)
    if length(rts) ≠ 1
        false
    else
        in(first(rts),canonicaldomain(c))
    end
end

reverseorientation(d::PeriodicCurve) = PeriodicCurve(reverseorientation(d.curve))

isambiguous(d::PeriodicCurve) = ncoefficients(d.curve)==0 && isambiguous(domain(d.curve))

convert(::Type{PeriodicCurve{S,T}},::AnyDomain) where {S,T}=Fun(S(AnyDomain()),[NaN])

arclength(d::PeriodicCurve) = linesum(ones(d))


##########################################################################
# Helper functions for set partitions.

# A block is just a set of integers.
immutable Block
    elements::IntSet
end
Block() = Block(IntSet())

# Inherit part of the interface from IntSet
Base.show(io::IO, b::Block) = showcompact(io, b.elements)
Base.length(b::Block) = length(b.elements)
Base.push!(b::Block, e::Int) = push!(b.elements, e)

# The Partition type describes a partition of integers into blocks.
immutable Partition
    blocks::Set{Block}
end
Partition() = Partition(Set{Block}())

# Inherit part of the interface from Set
Base.show(io::IO, p::Partition) = showcompact(io, p.blocks)
Base.length(p::Partition) = length(p.blocks)

# Create partition from restricted growth string
function Partition(a::Vector; skip_check::Bool=false)
    I = length(a)
    a_max = maximum(a)
    num_blocks = 1 + a_max

    # Consistency checks
    if !skip_check
        all(sort(unique(a)) .== 0:a_max) || Lumberjack.error("Restricted growth string may not skip values.", @compat Dict( :a => a ))
    end

    # Build blocks and partition
    blocks = [ Block() for n = 1:num_blocks ]
    for l = 1:I
        push!(blocks[a[l] + 1], l)
    end

    return Partition(Set(blocks))
end

# Create partition from assignment matrix
Partition(A::Matrix; skip_check::Bool=false) =
    Partition(restricted_growth_string(A); skip_check=skip_check)

# Convert from partition to restricted growth string
restricted_growth_string(p::Partition) = restricted_growth_string(logical_matrix(p))

# Convert from partition to assignment matrix
function logical_matrix(p::Partition)
    I = maximum([ maximum(block.elements) for block in p.blocks ])
    A = eye(Int, I, I)

    for block in p.blocks
        elements = collect(block.elements)
        A[elements, elements] = 1
    end

    return A
end

# Convert from assignment matrix to restricted growth string
function restricted_growth_string(A::Matrix)
    I = size(A, 1)

    # Consistency checks
    A' == A || Lumberjack.error("Tried to create partition from asymmetric assignment matrix.", @compat Dict(:A => A))
    diag(A) == ones(I) || Lumberjack.error("Elements must belong to their own block.", @compat Dict(:A => A))

    # Build restricted growth string
    Ac = copy(A)
    a = zeros(Int, I)
    block_no = 0
    for i = 1:I
        # Find all integers in the block containing integer i.
        elements = find(Ac[i,:] .== 1)

        if length(elements) > 0
            # Consistency check
            for j in elements[2:end]
                find(Ac[j,:] .== 1) == elements || Lumberjack.error("Incorrect structure in assignment matrix.", @compat Dict(:A => Ac))
            end

            # Add all elements to block and increment number of blocks.
            a[elements] = block_no
            block_no += 1

            # Do not investigate these elements more
            Ac[elements,elements] = 0
        end
    end

    return a
end

# Convert from restricted growth string to assignment matrix
function logical_matrix(a::Vector)
    I = length(a)
    a_max = maximum(a)
    num_blocks = 1 + a_max

    # Consistency checks
    all(sort(unique(a)) .== 0:a_max) || Lumberjack.error("Restricted growth string may not skip values.", @compat Dict(:a => a))

    # Create matrix iteratively
    A = eye(Int, I, I)
    for n = 0:(num_blocks-1)
        elements = find(a .== n)
        A[elements, elements] = 1
    end

    return A
end

# PartitionIterator gives all partitions of the set 1:n. This iterator is
# based on Algorithm H from TAoCP 7.2.1.5. Note that a is zero-based.
immutable PartitionIterator
    n::Int

    PartitionIterator(n) = n >= 2 ? new(n) : error("only applicable for n >= 2.")
end

type PartitionIteratorState
    a::Vector{Int}
    b::Vector{Int}
    m::Int
    terminate::Bool
end
PartitionIteratorState(n) = PartitionIteratorState(zeros(n), ones(n - 1), 1, false)

Base.start(p::PartitionIterator) = PartitionIteratorState(p.n) # H1
Base.done(p::PartitionIterator, s) = s.terminate
function Base.next(p::PartitionIterator, s)
    # H2
    current_partition = Partition(s.a, skip_check=true) # By construction, a is a valid restricted growth string.

    if s.a[p.n] < s.m
        # H3
        s.a[p.n] += 1
    else
        # H4
        j = p.n - 1
        while s.a[j] == s.b[j]
            j -= 1
        end

        # H5
        if j == 1
            s.terminate = true
        else
            s.a[j] += 1

            # H6
            s.m = s.b[j] + (s.a[j] == s.b[j] ? 1 : 0)
            j += 1
            while j < p.n
                s.a[j] = 0
                s.b[j] = s.m
                j += 1
            end
            s.a[p.n] = 0
        end
    end

    # H2 cont'd.
    return current_partition, s
end

##########################################################################
# Stirling and Bell numbers

# Stirling2NumberCoefficientIterator gives all coefficients that when summed
# result in a Stirling number of the second kind. Note that the coefficients
# may be real numbers, but when summed are integers. The coefficients are
# thus calculated as floating point numbers.
immutable Stirling2NumberCoefficientIterator
    n::BigInt
    k::BigInt
    k_fac::BigInt
end
Stirling2NumberCoefficientIterator(n, k) = Stirling2NumberCoefficientIterator(n, k, factorial(k))
Base.start(s::Stirling2NumberCoefficientIterator) = BigInt(0) # state: j
Base.done(s::Stirling2NumberCoefficientIterator, j) = (j > s.k)
Base.next(s::Stirling2NumberCoefficientIterator, j) = ((iseven(s.k - j) ? 1. : -1.)*(binomial(s.k, j)*j^s.n/s.k_fac), j + 1)

# Stirling2NumberIterator gives all Stirling numbers of the second kind
# up to n. The Stirling numbers are integers, but the Stirling number
# coefficients are not. We therefore round the sum of the coefficients,
# in order to end up with an integer.
immutable Stirling2NumberIterator
    n::BigInt
end
Base.start(s::Stirling2NumberIterator) = BigInt(0) # state: k
Base.done(s::Stirling2NumberIterator, k) = (k > s.n)
Base.next(s::Stirling2NumberIterator, k) = (BigInt(round(sum(Stirling2NumberCoefficientIterator(s.n, k)))), k + 1)

# The Bell number can we calculated as the sum of all
# Stirling numbers of the second kind.
bell_number(n) = sum(Stirling2NumberIterator(n))


# let
    q = Quat(1.0,2.0,3.0,4.5)
    @test length(q) == 4
    @test imag(q) == VecE3(1.0,2.0,3.0)
    @test conj(q) == Quat(-1.0,-2.0,-3.0,4.5)
    @test vec(q) == [1.0,2.0,3.0,4.5]
    @test vec(copy(q)) == vec(q)
    @test convert(Quat, [1.0,2.0,3.0,4.5]) == q
    @test norm(q) == norm(vec(q))
    @test vec(normalize(q)) == vec(q) ./ norm(vec(q))

    q = Quat(-0.002, -0.756, 0.252, -0.604)
    sol = push!(vec(get_axis(q)), get_rotation_angle(q))
    @test isapprox(sol, [-0.00250956, -0.948614,  0.316205, mod2pi(-1.84447)], atol=1e-3) ||
          isapprox(sol, [ 0.00250956,  0.948614, -0.316205, mod2pi( 1.84447)], atol=1e-3)

    srand(0)
    for i in 1 : 5
        a = normalize(VecE3(rand(),rand(),rand()))
        b = normalize(VecE3(rand(),rand(),rand()))
        q = quat_for_a2b(a,b)
        @test norm(rot(q, a) - b) < 1e-8
    end

    @test isapprox(angledist(Quat(1,0,0,0), Quat(0,1,0,0)), π, atol=1e-6)
    @test norm(lerp(Quat(1,0,0,0), Quat(0,1,0,0), 0.5) - Quat(√2/2, √2/2, 0, 0)) < 1e-6

    srand(0)
    for i in 1 : 5
        q = rand(Quat)
        @test isapprox(norm(q), 1.0, atol=1e-8)
        @test norm(inv(q)*q - Quat(0,0,0,1)) < 1e-8
        @test norm(q*inv(q) - Quat(0,0,0,1)) < 1e-8
    end
# end