using D91PVInversion
using IterativeSolvers
using Random
import PyPlot; const plt = PyPlot

# Define functional form of PV anomaly
function q′(x, y, z, A, L, H, p::Params)
    πc = p.π0 - 0.5
    Π = p.Π
    return A*exp(-(x^2 + y^2)/L^2)*exp(-(z-πc)^2/H^2)
end

# Set up inversion domain
params = Params(Float64)
domain = Domain(params, size = (3, 3, 3), x = (-1, 1)./2, y = (-1, 1)./2)
println(params)
println(domain)

# Create solvers
sψ = Solver(params)
sϕ = Solver(params)

# Allocate storage for fields and vectors
ψ, ϕ, q = allocate_fields(domain)
xψ, ∂ψ, bψ, xϕ, ∂ϕ, bϕ = allocate_rhs(domain)

# Set initial values
set_background_ψ!(ψ, domain, params)
set_background_ϕ!(ϕ, domain, params)

# Set PV field
# Set PV field
A = 0
L = 0.1
H = 0.1
set_q!(q, domain, params, q′, A, L, H, params)

# Plot fields
plt.rc("font", size = 8)
fig, axes = plt.subplots(
    figsize = (9.5, 2), nrows = 1, ncols = 5,
    sharey = true, sharex = true, constrained_layout = true
)
for i = 1:5
    im = axes[i].imshow(ψ[:,i-1,:]')
    plt.colorbar(im, ax = axes[i], location = "bottom", aspect = 5, shrink = 0.7)
    axes[i].set_title(string("ψ, j = ", i))
end
axes[1].invert_yaxis()
plt.show()

# Construct linear operators
Lψ = generate_Lψ(ϕ, domain)
Lϕ = generate_Lϕ(ψ, domain)

# Plot operators
plt.figure()
plt.imshow(Matrix(Lψ))
plt.show()
plt.figure()
plt.imshow(Matrix(Lϕ))
plt.show()

# Fill halos
fill_ψ_halos!(ψ, domain, params)
fill_ϕ_halos!(ϕ, domain, params)

# Re-plot fields
plt.rc("font", size = 8)
fig, axes = plt.subplots(
    figsize = (9.5, 2), nrows = 1, ncols = 5,
    sharey = true, sharex = true, constrained_layout = true
)
for i = 1:5
    im = axes[i].imshow(ψ[:,i-1,:]')
    plt.colorbar(im, ax = axes[i], location = "bottom", aspect = 5, shrink = 0.7)
    axes[i].set_title(string("ψ, j = ", i))
end
axes[1].invert_yaxis()
plt.show()

# Re-plot operators
plt.figure()
plt.imshow(Matrix(Lψ))
plt.show()
plt.figure()
plt.imshow(Matrix(Lϕ))
plt.show()

# BEGIN ITERATION HERE
# Set boundary contributions to ψ RHS
set_∂ψ!(∂ψ, ϕ, domain, params)

# Plot boundary contributions to ψ RHS
f∂ψ = field_from_rhs(∂ψ, domain)
fig, axes = plt.subplots(
    figsize = (9.5, 2), nrows = 1, ncols = 5,
    sharey = true, sharex = true, constrained_layout = true
)
for i = 1:5
    im = axes[i].imshow(f∂ψ[:,i-1,:]')
    plt.colorbar(im, ax = axes[i], location = "bottom", aspect = 5, shrink = 0.7)
    axes[i].set_title(string("∂ψ, j = ", i))
end
axes[1].invert_yaxis()
plt.show()

# Set ψ RHS
set_bψ!(bψ, ψ, ϕ, q, domain)

# Plot q and RHS
fbψ = field_from_rhs(bψ, domain)
fig, axes = plt.subplots(
    figsize = (9.5, 2), nrows = 1, ncols = 5,
    sharey = true, sharex = true, constrained_layout = true
)
for i = 1:5
    im = axes[i].imshow(q[:,i-1,:]')
    plt.colorbar(im, ax = axes[i], location = "bottom", aspect = 5, shrink = 0.7)
    axes[i].set_title(string("q, j = ", i))
end
axes[1].invert_yaxis()
plt.show()
fig, axes = plt.subplots(
    figsize = (9.5, 2), nrows = 1, ncols = 5,
    sharey = true, sharex = true, constrained_layout = true
)
for i = 1:5
    im = axes[i].imshow(fbψ[:,i-1,:]')
    plt.colorbar(im, ax = axes[i], location = "bottom", aspect = 5, shrink = 0.7)
    axes[i].set_title(string("bψ, j = ", i))
end
axes[1].invert_yaxis()
plt.show()

# Compute residual
rhs_from_field!(xψ, ψ, domain)
rψ = field_from_rhs(compute_residual(Lψ, xψ, ∂ψ, bψ, domain), domain)

# Plot residual
fig, axes = plt.subplots(
    figsize = (9.5, 2), nrows = 1, ncols = 5,
    sharey = true, sharex = true, constrained_layout = true
)
for i = 1:5
    im = axes[i].imshow(rψ[:,i-1,:]')
    plt.colorbar(im, ax = axes[i], location = "bottom", aspect = 5, shrink = 0.7)
    axes[i].set_title(string("ψ residual, j = ", i))
end
axes[1].invert_yaxis()
plt.show()

# Solve for ψ
Random.seed!(1234)
@. bψ = bψ - ∂ψ
bicgstabl!(xψ, Lψ, bψ; log = false, verbose = true)
relax!(ψ, xψ, sψ, domain; verbose = true)
fill_ψ_halos!(ψ, domain, params)
println("is_converged(sψ) = ", is_converged(sψ))

# Plot solution
fig, axes = plt.subplots(
    figsize = (9.5, 2), nrows = 1, ncols = 5,
    sharey = true, sharex = true, constrained_layout = true
)
for i = 1:5
    im = axes[i].imshow(ψ[:,i-1,:]')
    plt.colorbar(im, ax = axes[i], location = "bottom", aspect = 5, shrink = 0.7)
    axes[i].set_title(string("ψ, j = ", i))
end
axes[1].invert_yaxis()
plt.show()

# Set boundary contributions to ϕ RHS
set_∂ϕ!(∂ϕ, ψ, domain, params)

# Plot boundary contributions to ϕ RHS
f∂ϕ = field_from_rhs(∂ϕ, domain)
fig, axes = plt.subplots(
    figsize = (9.5, 2), nrows = 1, ncols = 5,
    sharey = true, sharex = true, constrained_layout = true
)
for i = 1:5
    im = axes[i].imshow(f∂ϕ[:,i-1,:]')
    plt.colorbar(im, ax = axes[i], location = "bottom", aspect = 5, shrink = 0.7)
    axes[i].set_title(string("∂ϕ, j = ", i))
end
axes[1].invert_yaxis()
plt.show()

# Set ψ RHS
set_bϕ!(bϕ, ψ, ϕ, q, domain)

# Plot q and RHS
fbϕ = field_from_rhs(bϕ, domain)
fig, axes = plt.subplots(
    figsize = (9.5, 2), nrows = 1, ncols = 5,
    sharey = true, sharex = true, constrained_layout = true
)
for i = 1:5
    im = axes[i].imshow(q[:,i-1,:]')
    plt.colorbar(im, ax = axes[i], location = "bottom", aspect = 5, shrink = 0.7)
    axes[i].set_title(string("q, j = ", i))
end
axes[1].invert_yaxis()
plt.show()
fig, axes = plt.subplots(
    figsize = (9.5, 2), nrows = 1, ncols = 5,
    sharey = true, sharex = true, constrained_layout = true
)
for i = 1:5
    im = axes[i].imshow(fbϕ[:,i-1,:]')
    plt.colorbar(im, ax = axes[i], location = "bottom", aspect = 5, shrink = 0.7)
    axes[i].set_title(string("bϕ, j = ", i))
end
axes[1].invert_yaxis()
plt.show()

# Compute residual
rhs_from_field!(xϕ, ϕ, domain)
rϕ = field_from_rhs(compute_residual(Lϕ, xϕ, ∂ϕ, bϕ, domain), domain)

# Plot residual
fig, axes = plt.subplots(
    figsize = (9.5, 2), nrows = 1, ncols = 5,
    sharey = true, sharex = true, constrained_layout = true
)
for i = 1:5
    im = axes[i].imshow(rϕ[:,i-1,:]')
    plt.colorbar(im, ax = axes[i], location = "bottom", aspect = 5, shrink = 0.7)
    axes[i].set_title(string("ϕ residual, j = ", i))
end
axes[1].invert_yaxis()
plt.show()

# Solve for ϕ
@. bϕ = bϕ - ∂ϕ
bicgstabl!(xϕ, Lϕ, bϕ; log = false, verbose = true)
relax!(ϕ, xϕ, sϕ, domain; verbose = true)
fill_ϕ_halos!(ϕ, domain, params)
println("is_converged(sϕ) = ", is_converged(sϕ))

# Plot solution
fig, axes = plt.subplots(
    figsize = (9.5, 2), nrows = 1, ncols = 5,
    sharey = true, sharex = true, constrained_layout = true
)
for i = 1:5
    im = axes[i].imshow(ϕ[:,i-1,:]')
    plt.colorbar(im, ax = axes[i], location = "bottom", aspect = 5, shrink = 0.7)
    axes[i].set_title(string("ϕ, j = ", i))
end
axes[1].invert_yaxis()
plt.show()


defaultsamplesregion() = ((60, 60, 60), Rect(Vec(-2,-2,-2), Vec(4,4,4)))
samples, region = defaultsamplesregion()

function gyroid()
    gyroid(v) = cos(v[1])*sin(v[2])+cos(v[2])*sin(v[3])+cos(v[3])*sin(v[1])
    gyroid_shell(v) = max(gyroid(v)-0.4,-gyroid(v)-0.4)

    # generate directly using GeometryBasics API
    # Rect specifies the sampling intervals
    gy_mesh = Mesh(gyroid_shell, Rect(Vec(0,0,0),Vec(pi*4,pi*4,pi*4)),
                    MarchingCubes(), samples=(50,50,50))

    return gy_mesh
end

spheredist(v) = norm(v) - 1.0

function boxdist(v, b)
    q = abs.(v) - b
    return norm(max.(q, 0)) + min.(maximum(q), 0)
end

function sphere()

    # generate directly using GeometryBasics API
    # Rect specifies the sampling intervals
    sphere_mesh = Mesh(spheredist, Rect(Vec(-1.2,-1.2,-1.2),Vec(2.4, 2.4, 2.4)),
                    MarchingCubes(), samples=(10,10,10))

    return sphere_mesh
end

D(f) = x -> ForwardDiff.derivative(f, x)

linecurve(t) = Point3{Float64}(1., 1., 1.) .* t
function curvedistgen(curve::Function, mint, maxt; eps=1e-13)
    function curvedist(p)
        disttocurve(t) = normsq(curve(t) - p)
        Ddisttocurve = D(disttocurve)
        return sqrt(minimum(disttocurve.(find_zeros(Ddisttocurve, mint, maxt))))
    end
    return curvedist
end

function arc_length_parameterization(γ::Function, tspan)
    normtangentγ(t) = norm(ForwardDiff.derivative(γ, t))

    f_t_to_s(s, p, t) = normtangentγ(t)
    s0 = 0
    t_to_s_prob = ODEProblem(f_t_to_s, s0, tspan)
    t_to_s = solve(t_to_s_prob, Tsit5(), reltol=1e-8, abstol=1e-8)

    maxs = t_to_s(1.0)

    f_s_to_t(t, p, s) = 1 / normtangentγ(t)
    t0 = tspan[1]
    sspan = (0.0, maxs)
    s_to_t_prob = ODEProblem(f_s_to_t, t0, sspan)
    s_to_t = solve(s_to_t_prob, Tsit5(), reltol=1e-8, abstol=1e-8)

    reparamγ(s) = γ(s_to_t(s))
    t_to_s, s_to_t, reparamγ, (0., maxs)
end

function curveofshapes(curve::AbstractArray{Tuple{Point3{Float64}, Float64}, 1}, shapefunc::Function)
    function innerdistance(v)
        distances = map(p->shapefunc((v - p[1]) .* (1 / p[2])), curve)
        return minimum(distances)
    end
    return innerdistance
end

function curvedistgenkdtree(γ::Function, mint, maxt; arc_length_sampling_freq=0.01)
    tspan = (mint, maxt)

    t_to_s, s_to_t, γs, sspan = arc_length_parameterization(γ, tspan)
    mins, maxs = sspan[1:2]

    γsamples = (mins:arc_length_sampling_freq:maxs)
    γsampled_points = map(γs, γsamples)

    γkdtree = KDTree(γsampled_points)

    function curvedist(p)
        #disttocurvet(t) = normsq(γ(t) - p)
        #Ddisttocurvet = D(disttocurvet)
        #Ddisttocurves = Ddisttocurvet ∘ s_to_t

        s_nearest_index, distance = knn(γkdtree, p, 1)
        #s_nearest = γsamples[s_nearest_index][1]
        #t_nearest = s_to_t(s_nearest)
        #s_search_range = (s_nearest - arc_length_sampling_freq, s_nearest + arc_length_sampling_freq)
        #@show s_search_range
        #try
            #t_min = find_zero(Ddisttocurvet, t_nearest; maxevals=100, atol=1e-4, rtol=1e-3)
            #return sqrt(disttocurvet(t_nearest))
        distance[1]
            #=
        catch e
            @show t_nearest, s_nearest, s_nearest_index
            @show p
            @show disttocurvet(t_nearest)
            @show e
            throw(e)
        end
        =#

    end
    curvedist, γsampled_points
end

struct Curve
    γ::Function
    mint::Float64
    maxt::Float64
end

function curvesdistgenkdtree(γs::Vector{Curve}; arc_length_sampling_freq=0.01)
    tspan = (mint, maxt)

    t_to_s, s_to_t, γs, sspan = arc_length_parameterization(γ, tspan)
    mins, maxs = sspan[1:2]

    γsamples = (mins:arc_length_sampling_freq:maxs)
    γsampled_points = map(γs, γsamples)

    γkdtree = KDTree(γsampled_points)

    function curvedist(p)
        #disttocurvet(t) = normsq(γ(t) - p)
        #Ddisttocurvet = D(disttocurvet)
        #Ddisttocurves = Ddisttocurvet ∘ s_to_t

        s_nearest_index, distance = knn(γkdtree, p, 1)
        #s_nearest = γsamples[s_nearest_index][1]
        #t_nearest = s_to_t(s_nearest)
        #s_search_range = (s_nearest - arc_length_sampling_freq, s_nearest + arc_length_sampling_freq)
        #@show s_search_range
        #try
            #t_min = find_zero(Ddisttocurvet, t_nearest; maxevals=100, atol=1e-4, rtol=1e-3)
            #return sqrt(disttocurvet(t_nearest))
        distance[1]
            #=
        catch e
            @show t_nearest, s_nearest, s_nearest_index
            @show p
            @show disttocurvet(t_nearest)
            @show e
            throw(e)
        end
        =#

    end
    curvedist, γsampled_points
end

function boundary_triangles(tetrahedra::Array{SimplexFace{4, Int32}, 1})
    tritonum = Dict{SVector{3, UInt32}, Int}()
    @show tritonum
    #tritonum[sv(1, 2, 3)] = 1
    @show tritonum

    tricombinations = [sv(1,2,3), sv(1,2,4), sv(1,3,4), sv(2,3,4)]
    @show tricombinations
    @show tricombinations[1]
    alltris = Array{NgonFace{3, UInt32}, 1}(undef, 0)
    for (i, tet) in enumerate(tetrahedra)
        for combination in tricombinations
            index = sortsv(tet[combination]...)
            if !haskey(tritonum, index)
                tritonum[index] = 1
            else
                tritonum[index] += 1
            end
            tri = NgonFace{3, UInt32}(index...)
            push!(alltris, tri)
        end
    end
    @show length(tritonum)

    boundarytris = Array{NgonFace{3, UInt32}, 1}(undef, 0)
    for (key, val) in pairs(tritonum)
        if val == one(UInt32)
            tri = NgonFace{3, UInt32}(key...)
            push!(boundarytris, tri)
        end
    end

    return boundarytris, alltris
end

function rectGBtoGT(rect::GeometryBasics.HyperRectangle{N, T}) where {N, T <: Number}
    return GeometryTypes.HyperRectangle{N, T}(origin(rect), widths(rect))
end

function rectGTtoGB(rect::GeometryTypes.HyperRectangle{N, T}) where {N, T <: Number}
    return GeometryBasics.HyperRectangle{N, T}(rect.origin, rect.widths)
end

function origin(rect::GeometryTypes.HyperRectangle{N, T}) where {N, T <: Number}
    return rect.origin
end

function widths(rect::GeometryTypes.HyperRectangle{N, T}) where {N, T <: Number}
    return rect.widths
end

function spheresdf(samples::NTuple{3, T}, region::Rect) where {T <: Integer}
    regionsamples = sampleregion(samples, region)
    return spheredist.(regionsamples)
end

function sampleregion(samples::NTuple{3, T}, region::Rect) where {T <: Integer}
    regionsamples = zeros(Point3{Float64}, samples[1], samples[2], samples[3])
    regionwidths = widths(region)
    regionorigin = origin(region)
    regionincrements = regionwidths ./ (samples .- 1)
    for i = 1:samples[1], j = 1:samples[2], k = 1:samples[3]
        samplepositions = ((i, j, k) .- 1) .* regionincrements .+ regionorigin
        regionsamples[i, j, k] = Point3{Float64}(samplepositions...)
    end
    return regionsamples
end

function facearray(mesh)
    mesh_faces = faces(mesh)
    m = length(mesh_faces)
    face_arr = Array{NgonFace{3, Int64}, 1}(undef, m)
    for i = 1:m
        face_arr[i] = Int64.(mesh_faces[i])
    end
    return face_arr
end

function singletri()
    verts = Array{Point{3, Float64}, 1}(undef, 3)
    pointtype = Point{3, Float64}
    verts[1] = pointtype([0.0, 0.0, 1.0])
    verts[2] = pointtype([0.0, 1.0, 0.0])
    verts[3] = pointtype([1.0, 0.0, 0.0])
    faces = Array{GLTriangleFace, 1}(undef, 1)
    faces[1] = GLTriangleFace(1, 2, 3)
    geom_mesh = Mesh(verts, faces)
    return geom_mesh
end

function connectedvertices(faces::AbstractArray{NgonFace{3, N}}) where {N <: Integer}
    Ntype = eltype(N)
    connectedvertexsets = Vector{Set{Ntype}}(undef, 0)
    #@show connectedvertexsets
    #@show length(connectedvertexsets)
    for i =  1:length(faces)
    #for i =  1:10
        if i % 100 == 0
            @show i
        end
        #@show i
        curset = Set{Ntype}([Ntype(faces[i][v]) for v = 1:3])
        #@show curset
        #@show typeof(curset)
        for v = 1:3
            for k = length(connectedvertexsets):-1:1
            #@show k
                vert = Ntype(faces[i][v])
                if vert in connectedvertexsets[k]
                    #@show typeof(curset)
                    #@show connectedvertexsets[k]
                    vertexset = popat!(connectedvertexsets, k)
                    #@show vertexset
                    curset = union!(vertexset, curset)
                    #@show typeof(curset)
                end
            end
        end
        #@show typeof(connectedvertexsets)
        push!(connectedvertexsets, curset)
    end
    return connectedvertexsets
end

function removesmallcomponents(mesh, connectedvertexsets; sizetokeep=1)
    sizessorted = sort(length.(connectedvertexsets))
    componentsizetokeep = sizessorted[sizetokeep]
    biggestcomponent = collect(filter(x -> length(x) == componentsizetokeep, connectedvertexsets)[1])
    oldtonew = Dict{UInt32, UInt32}()
    numnewverts = length(biggestcomponent)
    newverts = Array{Point3{Float64}, 1}(undef, numnewverts)
    oldverts = coordinates(mesh)
    #for i = 1:4
    for i = 1:numnewverts
        oldindex = biggestcomponent[i]
        oldtonew[oldindex] = i
        #@show oldverts[oldindex]
        newverts[i] = oldverts[oldindex]
    end
    @show length(newverts)

    oldfaces = faces(mesh)
    newfaces = Array{NgonFace{3, UInt32}, 1}(undef, 0)
    @show newfaces
    for i = 1:length(oldfaces)
    #for i = 1:100
        #@show i
        #@show oldfaces[i]
        newindices = Array{UInt32, 1}(undef, 0)
        for j = 1:3
            oldindex = UInt32(oldfaces[i][j])
            #@show oldindex
            if haskey(oldtonew, oldindex)
                push!(newindices, oldtonew[oldindex])
            end
        end
        #@show newindices
        if length(newindices) == 3
            newface = NgonFace{3, UInt32}(newindices...)
            push!(newfaces, newface)
        end
        #newindices = [oldtonew[] for j = 1:3]
        #@show newindices
        #newface = NgonFace{3, UInt32}([])

    end
    #@show newfaces
    newmesh = GeometryBasics.Mesh(newverts, newfaces)
    return newmesh
    #return biggestcomponent

end

function mandelbulbmesh(samples::NTuple{3, T}, region::Rect, numfractaliter::Integer) where {T <: Integer}
    mandelbulbdist = mandelbulbdistfunc(numfractaliter)
    return Mesh(mandelbulbdist, region, samples, MarchingTetrahedra(), pointtype=Point3{Float64})
end

function meshview(inmesh, color)
    return Makie.mesh(inmesh, color=color, show_axis=false)
end

function meshview(inmesh)
    return Makie.mesh(inmesh, color=[norm(v) for v in coordinates(inmesh)], show_axis=false)
end

function normsq(p::AbstractArray{T}) where {T <: Number}
    return sum(p .* p)
end

function inversesqrt(num::T) where {T <: Number}
    return 1.0 / sqrt(num)
end

function mandelbulbdistfunc(numfractaliter::Integer)
    function mandelbulbdist(p::Point{3, T}) where {T <: Real}
        px, py, pz = p[1], p[2], p[3]
        dz = 1.0
        w = Array{T, 1}(undef, 3)
        w .= 0

        stop_iter = 0

        w .+= p
        #@show typeof(w)
        m = normsq(w)
        #trap = vcat(abs.(w), m)

        for i = 1:numfractaliter

            w2 = w .* w
            w4 = w2 .* w2

            x, y, z = w[1], w[2], w[3]
            x2, y2, z2 = w2[1], w2[2], w2[3]
            x4, y4, z4 = w4[1], w4[2], w4[3]

            m2 = m * m
            m4 = m2 * m2

            dz *= 8.0 * sqrt(m4 * m2 * m)
            dz += 1.0

            k3 = x2 + z2
            k2 = inversesqrt(k3 * k3 * k3 * k3 * k3 * k3 * k3)
            k1 = x4 + y4 + z4 -
                6.0 * y2 * z2 -
                6.0 * x2 * y2 +
                2.0 * z2 * x2
            k4 = x2 - y2 + z2

            neww = [
                px +
                64.0 * x * y * z * (x2 - z2) * k4 *
                (x4 - 6.0 * x2 * z2 + z4) * k1 * k2,
                py +
                -16.0 * y2 * k3 * k4 * k4 + k1 * k1,
                pz +
                -8.0 * y * k4 * (x4 * x4 -
                28.0 * x4 * x2 * z2 +
                70.0 * x4 * z4 -
                28.0 * x2 * z2 * z4 +
                z4 * z4) * k1 * k2
                ]
            w[:] = neww[:]

            #trap[:, :, :] = min.(trap, vcat(abs.(w), m)) .* (1.0 .- stop_iter) .+ stop_iter .* trap
            m = normsq(w)
            if m > 256.0
                return -0.0005
            end
            #m[:, :, :] = min.(normsq(w), 256.0)
            #m[map(isnan, m)] .= 256.0
            #dz[map(isnan, dz)] .= 1.0

            #@show i
            #@show w
            #@show m
            #@show stop_iter
            #@show dz
            stop_iter = i

        end

        dist = 0.25 * log(m) * sqrt(m) / dz
        return dist
    end
    return mandelbulbdist
end

function init()
    blob2 = mandelbulbmesh(samples, region, 2)
    connecteds = connectedvertices(faces(blob2))
    #newblob2 = removesmallcomponents(blob2, connecteds)
    return blob2, connecteds #, newblob2
end

function loadbunny()
    bunnymeshpre = load("meshes/original/small_bunny.ply")
    vertspre, tris = coordinates(bunnymeshpre), faces(bunnymeshpre)
    transformmat = [1 0 0; 0 0 1; 0 1 0]
    verts = map(vert->Point3{Float32}(transformmat * vert), vertspre)

    return Mesh(verts, tris)
end

function bunnyeyepot(t::Float64)
    eyepot = Array(sparsevec(
        # eyeleft, eyeright, max size, booty, ear left, ear right
        [1311, 1015, 1430, 680, 568, 747],
         [Float64(3 * cos(t)), Float64(3 * cos(t)), Float64(0), Float64(-2), Float64(0.5), Float64(0.5)]))
    return eyepot
end

#begin
    #p1 = solvepoissonproblem(bs, bunnyeyepot(0.))
    #for i in 1:100
        #p1 = solveheatstep(bs, p1, 0.04)
        #meshview(bunny, p1)
    #end
#end


module Combinatorics

export AbstractCombinatorics
export Combinations, DulCombinations, Permutations, DulPermutations

"""
    AbstractCombinatorics{M, C}

Abstract combinatorial algorithms.
"""
abstract type AbstractCombinatorics{M, C} end
@inline Base.eltype(::Type{<:AbstractCombinatorics{M, C}}) where {M, C} = NTuple{M, eltype(C)}
@inline @generated gettuple(contents, indices, ::Val{M}) where {M} = Expr(:tuple, [:(contents[indices[$i]]) for i = 1:M]...)

"""
    Combinations{M}(contents::C) where {M, C}

Combinations of M elements from contents. Duplicates are not allowed.
"""
struct Combinations{M, C} <: AbstractCombinatorics{M, C}
    contents::C
    N::Int
    Combinations{M}(contents::C) where {M, C} = new{M, C}(contents, length(contents))
end
@inline Base.length(c::Combinations{M}) where {M} = binomial(c.N, M)
Base.iterate(c::Combinations{M}) where {M} = (M > c.N) ? nothing : (M == 0) ? ((), [c.N+2]) : (gettuple(c.contents, 1:M, Val(M)), nextmstate!(collect(1:M), c.N, M))
Base.iterate(c::Combinations{M}, state) where M = (state[1] > c.N-M+1) ? nothing : (gettuple(c.contents, state, Val(M)), nextmstate!(state, c.N, M))
function nextmstate!(state::Vector{Int}, N::Int, M::Int)
    for i = M:-1:1
        state[i] += 1
        (state[i] > N-(M-i)) && continue
        for j = i+1:M
            state[j] = state[j-1] + 1
        end
        break
    end
    state
end

"""
    DulCombinations{M}(contents::C) where {M, C}

Combinations of M elements from contents. Duplicates are allowed.
"""
struct DulCombinations{M, C} <: AbstractCombinatorics{M, C}
    contents::C
    N::Int
    DulCombinations{M}(contents::C) where {M, C} = new{M, C}(contents, length(contents))
end
@inline Base.length(c::DulCombinations{M}) where {M} = binomial(c.N+M-1, c.N-1)
function Base.iterate(c::DulCombinations{M}) where M
    (M == 0) ? ((), [c.N+1]) : (gettuple(c.contents, ntuple(i->1, M), Val(M)), nextdmstate!(collect(ntuple(i->1, M)), c.N, M))
end
Base.iterate(c::DulCombinations{M}, state) where {M} = (state[1] > c.N) ? nothing : (gettuple(c.contents, state, Val(M)), nextdmstate!(state, c.N, M))
function nextdmstate!(state::Vector{Int}, N::Int, M::Int)
    for i = M:-1:1
        state[i] += 1
        (state[i] > N) && continue
        for j = i+1:M
            state[j] = state[j-1]
        end
        break
    end
    state
end

"""
    Permutations{M}(contents::C) where {M, C}

Permutations of M elements from contents. Duplicates are not allowed.
"""
struct Permutations{M, C} <: AbstractCombinatorics{M, C}
    contents::C
    N::Int
    Permutations{M}(contents::C) where {M, C} = new{M, C}(contents, length(contents))
end
@inline Base.length(p::Permutations{M}) where {M} = (0 <= M <= p.N) ? prod((p.N-M+1):p.N) : 0
function Base.iterate(p::Permutations{M}) where M
    ((p.N == 0) && (M > 0) || (0 < p.N < M)) ? nothing : (state = collect(1:p.N); (gettuple(p.contents, state, Val(M)), nextpstate!(state, p.N, M)))
end
function Base.iterate(p::Permutations{M}, state) where M
    ((p.N == 0) && (M > 0) || (0 < p.N < max(state[1], M))) ? nothing : (gettuple(p.contents, state, Val(M)), nextpstate!(state, p.N, M))
end
function nextpstate!(state::Vector{Int}, N::Int, M::Int)
    (M <= 0) && return [N+1]
    (M < N) && (j = M + 1; while (j <= N) && (state[M] >= state[j]); j += 1; end)
    if (M < N) && (j <= N)
        state[M], state[j] = state[j], state[M]
    else
        (M < N) && reverse!(state, M+1)
        i = M - 1
        while (i >= 1) && (state[i] >= state[i+1]); i -= 1; end
        if i > 0
            j = N
            while (j > i) && (state[i] >= state[j]); j -= 1; end
            state[i], state[j] = state[j], state[i]
            reverse!(state, i+1)
        else
            state[1] = N + 1
        end
    end
    return state
end

"""
    DulPermutations{M}(contents::C) where {M, C}

Permutations of M elements from contents. Duplicates are allowed.
"""
struct DulPermutations{M, C} <: AbstractCombinatorics{M, C}
    contents::C
    N::Int
    DulPermutations{M}(contents::C) where {M, C} = new{M, C}(contents, length(contents))
end
@inline Base.length(p::DulPermutations{M}) where M = p.N^M
function Base.iterate(p::DulPermutations{M}) where M
    indices = CartesianIndices(ntuple(i->p.N, M|>Val))
    index = iterate(indices)
    isnothing(index) && return nothing
    return gettuple(p.contents, reverse(index[1].I), M|>Val), (indices, index[2])
end
function Base.iterate(p::DulPermutations{M}, state) where M
    index = iterate(state[1], state[2])
    isnothing(index) && return nothing
    return gettuple(p.contents, reverse(index[1].I), M|>Val), (state[1], index[2])
end

end # module


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    NormBall(radius, [metric])

A norm ball with `radius` and `metric`.
Default metric is `Euclidean()`.
"""
struct NormBall{T,M} <: MetricBall
  radius::T
  metric::M
end

function NormBall(radius)
  metric = Euclidean()
  NormBall{typeof(radius),typeof(metric)}(radius, metric)
end

metric(ball::NormBall) = ball.metric

"""
    radius(ball)

Return the radius of the norm `ball`.
"""
radius(ball::NormBall) = ball.radius

function Base.show(io::IO, ball::NormBall{T}) where {T}
  r, d = ball.radius, ball.metric
  print(io, "NormBall{$T}($r, $d)")
end


#=
Copyright 2021 BlackRock, Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
=#

"""
    optimize(params::AdmmParams[, settings::Settings[, vars::Vars]])

Carries out the ADMM algorithm by calling `admm_step` until termination conditions are met.
"""
function optimize(params::AdmmParams, settings::Settings, vars::Vars)
    assert_valid(params, settings, vars)
    kkt = factorize_kkt(params, settings)
    pc = ProxCache(params, settings)
    stats = Stats(settings)
    term_cache = get_term_cond_cache(params, settings)
    iter = 0
    while iter < settings.max_iters
        iter += 1
        admm_step!(vars, params, kkt, pc, settings, iter)
        if settings.compute_stats
            compute_stats(vars, params, stats, iter)
        end
        if check_term_conds!(term_cache, vars, params, settings, iter)
            break
        end
    end
    return vars, stats
end
optimize(params::AdmmParams, settings::Settings) = optimize(params, settings, Vars(params))
optimize(params::AdmmParams) = optimize(params, Settings(params))

@doc raw"""
    factorize_kkt(params::AdmmParams, settings::Settings)

Factorize the KKT matrix representing the optimzality equations of the provided problem.

# Extended help
Factorize the coefficient matrix of that KKT system representing the optimality equations
that will be used (and reused) for the constrained least squares problem:

```math
\begin{array}{ll}
\text{minimize} & \displaystyle
\frac12 x^T P x  + q^T x + \sum_{i=1}^n g_i(x_i) \\
\text{subject to} & Ax = b.
\end{array}
```

We can formulate this equivalently as

```math
\begin{array}{ll}
\text{minimize} & \displaystyle
\frac12 \tilde x^T P \tilde x  + q^T \tilde x + \mathcal{I}_\mathcal{A}(\tilde x, \tilde z)
        + g(x) + \mathcal{I}_{\mathcal B}(z) \\
\text{subject to} & \tilde x = x \\
                  & \tilde z = z.
\end{array}
```

where ``\mathcal I_{\mathcal A}`` is the indicator function over the affine
constraints of the original form, *i.e.*,
```math
\mathcal{I_{\mathcal A}}(x, z) = \begin{cases}
0 & Ax = z \\
\infty & {\rm otherwise}.
\end{cases}
```

and ``\mathcal I_{\mathcal B}`` is the indicator function over the singleton
set of vectors equal to ``b``, *i.e.*,

```math
\mathcal{I_{\mathcal B}}(x) = \begin{cases}
0 & x = b \\
\infty & {\rm otherwise}.
\end{cases}
```

We will denote the objective in this forumulation as ``f(x, \tilde x, z, \tilde z)``.

The augmented Lagrangian for this problem is

```math
L_{S,R} (x, \tilde x, z, \tilde z, w, y) =
f(x, \tilde x, z, \tilde z)
+ \frac12 \| \tilde x - x + S^{-1} w \|_S^2
+ \frac12 \| \tilde z - z + R^{-1} y \|_R^2,
```

where ``f`` is the previously mentioned objective function.
The diagonal, positive definite matrices ``S\in \bf S{++}^l`` and ``R\in \bf S_{++}^m``
are algorithm parameters.
We define ``\|\cdot \|_S`` is the ``S``-norm, *i.e.*,
``\|x\|_S = \sqrt{x^T S x}``.

Solving for the optimality conditions, we can find optimal ``\tilde x`` and ``\tilde z`` by solving the linear system

```math
\begin{bmatrix} P + S & A^T \\ A & -R^{-1} \end{bmatrix}
\begin{bmatrix} \tilde x \\ \nu \end{bmatrix}
=
\begin{bmatrix} Sx - q - w \\ z - R^{-1} y \end{bmatrix}.
```

where ``\tilde z`` can be recovered as ``\tilde z = z + R^{-1}(\nu - y)``

In this function, we factor the coefficient matrix on the left hand side so that we repeatedly
solve the linear system efficiently.
"""
function factorize_kkt(params::AdmmParams, settings::Settings)
    AT = sparse(params.A')
    Dσ = sparse(Diagonal(settings.σ))
    Dρ = sparse(Diagonal(1 ./ settings.ρ))

    # construct the block matrix K
    K = [params.P+Dσ AT;
         params.A -Dρ]
    return ldl(K)
end

"""
    solve_kkt!(vars::Vars, kkt::KKT, params::AdmmParams, settings::Settings)

Solve the KKT system set up in the description of `factorize_kkt` for a particular right hand side.
"""
function solve_kkt!(vars::Vars, kkt::KKT, params::AdmmParams, settings::Settings)
    m, n = size(params.A)
    rhs = Vector{Float64}(undef, n + m)
    for i in 1:n
        rhs[i] = settings.σ[i] * vars.x[i] - vars.w[i] - params.q[i]
    end
    for i in 1:m
        rhs[i + n] = vars.z[i] - (vars.y[i] / settings.ρ[i])
    end
    lhs = kkt \ rhs
    xt = lhs[1:n]
    ν = lhs[(n + 1):end]
    zt = vars.z + (ν - vars.y) ./ settings.ρ
    return xt, zt
end

"""
    admm_step!(vars::Vars, params::AdmmParams, kkt::KKT, pc::ProxCache, settings::Settings, iter::Int64; polish::Bool=false)

Carry out a single ADMM iteration.

# Steps:
1. Solve the KKT system to carry out the `xt` and `zt` updates.
2. Use a mixture of the result of (1) and the previous `x` iterate to solve `n` univariate
   subproblems by minimizing the Lagrangian in the description of `factorize_kkt` w.r.t.
   `x`. (Requires evaluation of a proximal operator.)
3. Update `w` and `y` using the residuals associated with the constraints `x = xt` and `z = zt`.

# Note:
* Note that for the entire run of the algorithm, the `z` update is just `z = b`, so we don't need
to do anything here.
"""
function admm_step!(vars::Vars, params::AdmmParams, kkt::KKT, pc::ProxCache,
                    settings::Settings, iter::Int64; polish::Bool=false)
    # update xt and zt by solving the KKT system
    xt, zt = solve_kkt!(vars, kkt, params, settings)

    # get the next x value by using the proximal operators of the separable function pieces
    prox_target = settings.α * xt + (1 - settings.α) * vars.x + vars.w ./ settings.σ
    x = prox_step(pc, params, settings, prox_target)
    @assert all(isfinite.(x))

    # update the dual variables
    z = params.b
    w = vars.w + settings.σ .* (settings.α * xt + (1 - settings.α) * vars.x - x)
    y = vars.y + settings.ρ .* (settings.α * zt + (1 - settings.α) * vars.z - z)

    # save in vars struct
    vars.xt = xt
    vars.zt = zt
    vars.x = x
    vars.w = w
    vars.z = z
    return vars.y = y
end

"""
    compute_stats(vars::Vars, params::AdmmParams, stats::Stats, iter::Int64)

Update the objective, residual, and iteration fields of the Stats object.
"""
function compute_stats(vars::Vars, params::AdmmParams, stats::Stats, iter::Int64)
    x = vars.x
    P = params.P
    q = params.q
    A = params.A
    b = params.b
    g = params.g

    fx = 0.5 * x' * P * x + q' * x
    gx = 0.0
    for i in 1:get_num_vars(params)
        if !isfinite(g[i, 1](x[i]))
        else
            gx += g[i, 1](x[i])
        end
    end
    stats.obj[iter] = fx + gx
    stats.iters = iter

    res = norm(A * x - b)
    return stats.res[iter] = res
end


export initboard, gen_new_tile!, sim_one_move!, simulate_game, simulate_games, DIRS

using Base.Threads

"""intiialise the game board"""
function initboard(random2=true)
    board = zeros(Int8, 4, 4)
    if random2
        gen_new_tile!(board)
        gen_new_tile!(board)
    end
    board
end

const PROB4 = 0.1 # probability of seeing a 4

function gen_new_tile!(board, pos = rand(findall(==(0), board)))
    board[pos] = rand() > PROB4 ? 1 : 2
    board
end

# it moves but doesn't generate new tile
function sim_one_move!(board)
    dirs = collect(DIRS)
    size = 4
    pts = 0

    updated = false

    while !updated && size >= 1
        pos = rand(1:size)
        pts, updated = move!(board, dirs[pos])
        dirs[pos] = dirs[size]
        size -= 1
    end
    pts, updated
end

function simulate_game(board = initboard(), score::Int = 0, delta=1)
    hehe = []
    pts, update_possible =  sim_one_move!(board)
    rounds = 0

    while update_possible
        score += pts*delta^rounds
        gen_new_tile!(board)
        pts, update_possible =  sim_one_move!(board)

        rounds += 1
    end

    board, score, rounds
end

function simulate_games(n, board = initboard(), score::Int = 0)
    res = Vector{Int}(undef, n)
    @threads for i in 1:n
        _, pts = simulate_game(deepcopy(board), score)
        res[i] = pts
    end
    res
end


#---------------------------------------------------------------------------------# Beta
"""
    FitBeta(; weight)

Online parameter estimate of a Beta distribution (Method of Moments).

# Example
    o = fit!(FitBeta(), rand(1000))
"""
struct FitBeta{V<:Variance} <: OnlineStat{Number}
    var::V
end
FitBeta(;kw...) = FitBeta(Variance(;kw...))
nobs(o::FitBeta) = nobs(o.var)
_fit!(o::FitBeta, y::Real) = fit!(o.var, y)
function value(o::FitBeta)
    if o.var.n > 1
        m = mean(o.var)
        v = var(o.var)
        α = m * (m * (1 - m) / v - 1)
        β = (1 - m) * (m * (1 - m) / v - 1)
        return α, β
    else
        return 1.0, 1.0
    end
end
_merge!(o::FitBeta, o2::FitBeta) = _merge!(o.var, o2.var)

#---------------------------------------------------------------------------------# Cauchy
"""
    FitCauchy(; alg, rate)

Approximate parameter estimation of a Cauchy distribution.  Estimates are based on
quantiles, so that `alg` will be passed to [`Quantile`](@ref).

# Example
    o = fit!(FitCauchy(), randn(1000))
"""
mutable struct FitCauchy{T} <: OnlineStat{Number}
    q::Quantile{T}
end
FitCauchy(alg = OMAS(), kw...) = FitCauchy(Quantile([.25, .5, .75]; alg=alg, kw...))
nobs(o::FitCauchy) = nobs(o.q)
_fit!(o::FitCauchy, y) = _fit!(o.q, y)
function value(o::FitCauchy)
    if nobs(o) > 1
        return o.q.value[2], 0.5 * (o.q.value[3] - o.q.value[1])
    else
        return 0.0, 1.0
    end
end
_merge!(o::FitCauchy, o2::FitCauchy) = _merge!(o.q, o2.q)

#---------------------------------------------------------------------------------# Gamma
"""
    FitGamma(; weight)

Online parameter estimate of a Gamma distribution (Method of Moments).

# Example
    using Random
    o = fit!(FitGamma(), randexp(10^5))
"""
struct FitGamma <: OnlineStat{Number}
    v::Variance
end
FitGamma() = FitGamma(Variance())
nobs(o::FitGamma) = nobs(o.v)
_fit!(o::FitGamma, y) = _fit!(o.v, y)
function value(o::FitGamma)
    if nobs(o) > 1
        m = mean(o.v)
        θ = var(o.v) / m
        α = m / θ
        return α, θ
    else
        return 1.0, 1.0
    end
end
_merge!(o::FitGamma, o2::FitGamma) = _merge!(o.v, o2.v)

#---------------------------------------------------------------------------------# LogNormal
"""
    FitLogNormal()

Online parameter estimate of a LogNormal distribution (MLE).

# Example
    o = fit!(FitLogNormal(), exp.(randn(10^5)))
"""
struct FitLogNormal <: OnlineStat{Number}
    v::Variance
    FitLogNormal() = new(Variance())
end
nobs(o::FitLogNormal) = nobs(o.v)
_fit!(o::FitLogNormal, y) = _fit!(o.v, log(y))
function value(o::FitLogNormal)
    if nobs(o) > 1
        return mean(o.v), std(o.v)
    else
        return 0.0, 1.0
    end
end
_merge!(o::FitLogNormal, o2::FitLogNormal) = _merge!(o.v, o2.v)

#---------------------------------------------------------------------------------# Normal
"""
    FitNormal()

Calculate the parameters of a normal distribution via maximum likelihood.

# Example
    o = fit!(FitNormal(), randn(1000))
"""
struct FitNormal{V <: Variance} <: OnlineStat{Number}
    v::V
end
FitNormal(;kw...) = FitNormal(Variance(;kw...))
_fit!(o::FitNormal, y::Real) = _fit!(o.v, y)
nobs(o::FitNormal) = nobs(o.v)
function value(o::FitNormal)
    if nobs(o) > 1
        return mean(o.v), std(o.v)
    else
        return 0.0, 1.0
    end
end
_merge!(o::FitNormal, o2::FitNormal) = _merge!(o.v, o2.v)
Statistics.mean(o::FitNormal) = mean(o.v)
Statistics.var(o::FitNormal) = var(o.v)

function pdf(o::FitNormal, x::Number)
    σ = std(o)
    1 / (sqrt(2π) * σ) * exp(-(x - mean(o))^2 / 2σ^2)
end
function cdf(o::FitNormal, x::Number)
    .5 * (1.0 + erf_approx((x - mean(o)) / (std(o) * √2)))
end
# https://en.wikipedia.org/wiki/Error_function#Approximation_with_elementary_functions
function erf_approx(x)
    s = sign(x)
    p = 0.3275911
    a1 = 0.254829592
    a2 = -0.284496736
    a3 = 1.421413741
    a4 = -1.453152027
    a5 = 1.061405429
    t = 1 / (1 + p * s * x)
    s * (1 - (a1*t + a2*t^2 + a3*t^3 + a4*t^4 + a5*t^5) * exp(-x^2))
end

#-----------------------------------------------------------------------# Multinomial
"""
    FitMultinomial(p)

Online parameter estimate of a Multinomial distribution.  The sum of counts does not need
to be consistent across observations.  Therefore, the `n` parameter of the Multinomial
distribution is returned as 1.

# Example
    x = [1 2 3; 4 8 12]
    fit!(FitMultinomial(3), x)
"""
mutable struct FitMultinomial{T} <: OnlineStat{VectorOb}
    grp::Group{T}
end

FitMultinomial(p::Int=0) = FitMultinomial(p * Mean())
_fit!(o::FitMultinomial, y) = _fit!(o.grp, y)
nobs(o::FitMultinomial) = nobs(o.grp)
function value(o::FitMultinomial)
    m = value.(o.grp.stats)
    p = length(o.grp)
    outvec = all(x-> x==0.0, m) ? ones(p) ./ p : collect(m) ./ sum(m)
    return 1, outvec
end
_merge!(o::FitMultinomial, o2::FitMultinomial) = _merge!(o.grp, o2.grp)

#---------------------------------------------------------------------------------# MvNormal
"""
    FitMvNormal(d)

Online parameter estimate of a `d`-dimensional MvNormal distribution (MLE).

# Example

    y = randn(100, 2)
    o = fit!(FitMvNormal(2), eachrow(y))
"""
struct FitMvNormal <: OnlineStat{VectorOb}
    cov::CovMatrix{Float64}
    FitMvNormal(p::Integer) = new(CovMatrix(p))
end
nvars(o::FitMvNormal) = nvars(o.cov)
nobs(o::FitMvNormal) = nobs(o.cov)
_fit!(o::FitMvNormal, y) = _fit!(o.cov, y)
function value(o::FitMvNormal)
    c = cov(o.cov)
    if isposdef(c)
        return mean(o.cov), c
    else
        return zeros(nvars(o)), Matrix(1.0I, nvars(o), nvars(o))
    end
end
_merge!(o::FitMvNormal, o2::FitMvNormal) = _merge!(o.cov, o2.cov)


# Kalman Filter tests

# Test construction with absolute state
linear_sys = LinearSystem(0.5*eye(2), 0.001*eye(2))
linear_obs = LinearObserver(eye(2), 0.001*eye(2))
initial_est = DiscreteState([1.0, 2.0])
kf = KalmanFilter(linear_sys, linear_obs, initial_est)
@test typeof(kf.estimate) <: UncertainDiscreteState
@test kf.estimate.P == zeros(2,2)

# Test consider constructors
kf = KalmanFilter(linear_sys, linear_obs, initial_est, [2])
kf = KalmanFilter(linear_sys, linear_obs, make_uncertain(initial_est), [2])
@test_throws DimensionMismatch KalmanFilter(linear_sys, linear_obs, initial_est, [3])
@test_throws ArgumentError KalmanFilter(linear_sys, linear_obs, initial_est, [3,3])

# Test discrete kalman filter
srand(2);
linear_sys = LinearSystem(0.5*eye(2), 0.001*eye(2))
linear_obs = LinearObserver(eye(2), 0.001*eye(2))
initial_est = UncertainDiscreteState([1.0, 2.0], 0.1*eye(2))
kf = KalmanFilter(linear_sys, linear_obs, initial_est)
simulator = make_simulator(kf)
for i = 1:10
    true_state, measurement = simulate(simulator, i)
    process!(kf, measurement)
    @test mahalanobis(true_state, kf.estimate) < 3
end

# Test discrete consider kalman filter
srand(2);
A = [1.0 0.5 0.01; -0.5 1.0 0.01; 0.0 0.0 1.0]
linear_sys = LinearSystem(A, 0.001*eye(3))
linear_obs = LinearObserver(eye(2,3), 0.1*eye(2))
initial_est = UncertainDiscreteState([1.0, 2.0, 3.0], 0.1*eye(3))
consider_states = [3]
kf = KalmanFilter(linear_sys, linear_obs, initial_est, consider_states)
simulator = make_simulator(kf)
for i = 1:10
    true_state, measurement = simulate(simulator, i)
    process!(kf, measurement)
    @test mahalanobis(true_state, kf.estimate) < 3
end
@test kf.estimate.P[3,3] == initial_est.P[3,3]

# Test discrete kalman filter with archiving
srand(2);
linear_sys = LinearSystem(0.5*eye(2), 0.001*eye(2))
linear_obs = LinearObserver(eye(2), 0.01*eye(2))
initial_est = UncertainDiscreteState([1.0, 2.0], 0.1*eye(2))
kf = KalmanFilter(linear_sys, linear_obs, initial_est)
simulator = make_simulator(kf)
archive = EstimatorHistory(kf)
for i = 1:10
    true_state, measurement = simulate(simulator, i)
    process!(kf, measurement, archive)
    @test mahalanobis(true_state, kf.estimate) < 3
end

# Test continuous kalman filter
srand(2);
linear_sys = LinearSystem(-eye(3), 0.001*eye(3))
linear_obs = LinearObserver(eye(3), 0.01*eye(3))
initial_est = UncertainContinuousState([1.0, 2.0, 3.0], 0.1*eye(3))
kf = KalmanFilter(linear_sys, linear_obs, initial_est)
simulator = make_simulator(kf)
for i = 1:10
    true_state, measurement = simulate(simulator, kf.estimate.t+0.1)
    process!(kf, measurement)
    @test mahalanobis(true_state, kf.estimate) < 3
end

# Test continuous kalman filter with archiving
srand(2);
linear_sys = LinearSystem(-eye(3), 0.001*eye(3))
linear_obs = LinearObserver(eye(3), 0.01*eye(3))
initial_est = UncertainContinuousState([1.0, 2.0, 3.0], 0.1*eye(3))
kf = KalmanFilter(linear_sys, linear_obs, initial_est)
simulator = make_simulator(kf)
archive = EstimatorHistory(kf)
for i = 1:10
    true_state, measurement = simulate(simulator, kf.estimate.t+0.1)
    process!(kf, measurement, archive)
    @test mahalanobis(true_state, kf.estimate) < 3
end

# Test continuous kalman filter convenience function
srand(2);
linear_sys = LinearSystem(-eye(3), 0.001*eye(3))
linear_obs = LinearObserver(eye(3), 0.001*eye(3))
estimate = UncertainContinuousState([1.0, 2.0, 3.0], 0.1*eye(3))
simulator = make_simulator(KalmanFilter(linear_sys, linear_obs, estimate))
for i = 1:5
    true_state, measurement = simulate(simulator, kf.estimate.t+0.1)
    estimate = kalman_filter(estimate, linear_sys, linear_obs, measurement)
end

# Test continuous kalman filter with archiving
srand(2);
linear_sys = LinearSystem(-eye(3), 0.001*eye(3))
linear_obs = LinearObserver(eye(3), 0.001*eye(3))
estimate = UncertainContinuousState([1.0, 2.0, 3.0], 0.1*eye(3))
kf = KalmanFilter(linear_sys, linear_obs, estimate)
simulator = make_simulator(kf)
archive = EstimatorHistory(kf)
for i = 1:5
    true_state, measurement = simulate(simulator, kf.estimate.t+0.1)
    estimate =
        kalman_filter(estimate, linear_sys, linear_obs, measurement, archive)
end


#' Boosting the Hodrick-Prescott Filter
#'
#' Coded by Mei Ziwei
#' Documented by Shi Zhentao
#'
#' All in one function of conducting the boosted HP-filter.
#' tesing example is at the end of the script
##############################################################################
#' Parameters:
#'
#' x - a raw time series to be filtered.
#' lam - turning parameter, default value is 1600,
#'   as recommended by Hodrick and Prescott (1997) for quarterly data.
#' iter - logical, True (default) to conduct the boosted HP filter.
#'   False does not iterated, which is exactly the original HP filter.
#' stopping - stopping criterion.  "BIC" (default), or  "adf", or  "nonstop"  means keeping
#'    iteration until the maximum number of iteration, specified by Max_Iter is reached.
#' sig_p - a threshold of the p-value for the ADF test, with default value 0.050.
#'    Only effective when stopping = "adf".
#' Max_Iter - maximal number of iterations. The default is 100.
##############################################################################
#' The function returns a dictionary containing the following items:
#'
#' cycle - The cyclical component in the final iteration.
#' trend - The trend component in the final iteration.
#' trend_hist - The estimated trend in each iteration.
#' iter_num - The total number of iterations when it stops.
#' BIC_hist - The path of the BIC up to the final iterations.
#' adf_p_hist - The path of the ADF test p-value up to the final iteration.
##############################################################################
#' Details:
#'
#' This is the main function of implementing the boosted HP filter (Phillisp and
#' Shi, 2019). The arguments accommendate the orginal HP filter (iter =
#' False), the boosted HP filter with the BIC stopping criterion (stopping = "BIC"),
#' or ADF test stopping criterion (stopping = "adf"),
#' or keep going until the maximum number of iterations is reached (stopping = "nonstop").
#'
#' Either the original HP filter or the bHP filter requires lambda to
#' control the strength of the weak learner for in-sample fitting.
#' The default is lambda = 1600,
#' which is recommended by Hodrick and Prescott (1997) for quarterly data.
#' lambda should be adjusted for different frequencies.
#' For example, lambda = 129600 for monthly data and
#' lambda = 100 or 6.25 for annual data.
#'
#' See the vignette with a brief introduction of the idea of bHP.
#'
##############################################################################
#' References:
#'
#' Phillips, Peter CB, and Zhentao Shi. "Boosting: Why you can use the hp
#' filter." arXiv: 1905.00175, Cowles Foundation Discussion Paper No.2192,
#' (2019).
#'
##############################################################################

using LinearAlgebra
using Statistics
using HypothesisTests
function BoostedHP(x,λ=1600,iter=true,stopping="BIC",sig_p=0.050,Max_Iter=100)

    ## Examine data format
    ## generating trend operator matrix "S"
    raw_x = x;
    n = length(x);

    I_n = Matrix(I,n,n);
    D_temp = [zeros(1,n);Matrix(I,n-1,n)];
    D_temp = (I_n-D_temp)*(I_n-D_temp);
    D = D_temp[3:n, :]';

    S = inv(I_n+λ*(D*D'));
    mS = I_n - S ;

    # the Simple HP Filter
    if ~iter

        # get the trend and cycle
        x_f = S*x;
        x_c = x - x_f;

        # Results
        result = Dict("cycle" => x_c, "trend_hist" => x_f,
                   "stopping" => "nonstop", "trend" => x - x_c, "raw_data" => raw_x);
    end

    # the Boosted HP-filter
    if iter

        ## ADF test as the stopping criterion
        if stopping == "adf"

            r = 1;
            stationary = false;
            x_c = x;

            x_f = zeros(n,Max_Iter);
            adf_p = zeros(Max_Iter);

            while (r <= Max_Iter) & ~stationary

                x_c = mS * x_c; # update
                x_f[:,r] = x - x_c;

                adf_p_r = pvalue(ADFTest(x_c,:trend,Int(floor((n-1)^(1/3)))));
                # x_c is the residual after the mean and linear trend being removed by HP filter
                # we use the critical value for the ADF distribution with
                # the intercept and linear trend specification

                adf_p[r] = adf_p_r;
                stationary = (adf_p_r <= sig_p);
                # Truncate the storage matrix and vectors

                if stationary
                    R = r;
                    x_f = x_f[:,1:R];
                    adf_p = adf_p[1:R];
                    break
                end

                r += 1
            end  # end the while loop

            if r > Max_Iter
                R = Max_Iter;
                warning("The number of iterations exceeds Max_Iter.
                The residual cycle remains non-stationary.")

            end
            result = Dict("cycle" => x_c, "trend_hist" => x_f,  "stopping" => stopping,
                 "signif_p" => sig_p, "adf_p_hist" => adf_p, "iter_num" => R,
                 "trend" => x - x_c, "raw_data" => raw_x);

        else # either

            #assignment
            r = 0;
            x_c_r = x;
            x_f = zeros(n,Max_Iter);
            IC = zeros(Max_Iter);
            # IC_decrease = true;

            I_S_0 = I_n - S;
            c_HP = I_S_0 * x;
            I_S_r = I_S_0;

            while r < Max_Iter

                r = r + 1;

                x_c_r = I_S_r * x;# this is the cyclical component after m iterations
                x_f[:,r] = x - x_c_r;
                B_r = I_n - I_S_r;
                IC[r] = var(x_c_r)/var(c_HP) + log(n)/(n-sum(diag(S)))*sum(diag(B_r));

                I_S_r = I_S_0 * I_S_r; # update for the next round

                if (r>=2) & (stopping == "BIC")
                    if IC[r-1] < IC[r]
                        break
                    end
                end

            end # end the shile loop

            # Final assignment
            R = r - 1;
            x_f = x_f[:, 1:R]
            x_c = x - x_f[:,R]

            if stopping == "BIC"
                # save the path of BIC till iter+1 times to keep the "turning point" of BIC history.
                result = Dict("cycle" => x_c, "trend_hist" => x_f,  "stopping" => stopping,
                       "BIC_hist" => IC[1:(R+1)], "iter_num" => R, "trend" =>  x - x_c, "raw_data" => raw_x);
            end

            if stopping == "nonstop"

                result = Dict("cycle" => x_c, "trend_hist" => x_f,  "stopping" => stopping,
                       "BIC_hist" => IC, "iter_num" => Max_Iter-1, "trend" => x - x_c, "raw_data" => raw_x);
            end

        end

    end  # end the Boosted HP

    return result

end  ## end function
###############################################################################
## Examples

using CSV
using DelimitedFiles
IRE = readdlm("IRE.csv", ',')[:]; # laod the data "IRE"
λ = 100;  # tuning the parameter for annumal data
#
#' # raw HP filter
bx_HP = BoostedHP(IRE, λ, false);
# bx_HP_cycle = bx_HP["cycle"]  # The cyclical component
# bx_HP_trend = bx_HP["trend"]  # The trend component
#'
#' # by BIC
bx_BIC = BoostedHP(IRE, λ, true,"BIC");
#'
#' # by ADF
#' # CAVEAT: Results may be slightly different from other languages
#' #    due to different critical values of ADF test.
bx_ADF = BoostedHP(IRE, λ, true,"adf");
#'
#' # If stopping = "nonstop",
#' # Iterated HP filter until Max_Iter and keep the path of BIC.
bx_nonstop = BoostedHP(IRE, λ, true,"nonstop");


using StatsBase, Distances, Hungarian, Printf

export MatchingDistance, MatchDist
export FastMatchingDistance, FastMatchDist
export FixedPenaltyMatchingDistance, FixPenMatchDist
export MinDistMatchingDistance, MinDistMatchDist
export AvgSizeMatchingDistance, AvgSizeMatchDist
export get_cost_matrix_dynamic, get_cost_matrix_fixed

struct MatchingDistance{T<:Metric} <: Metric
    ground_dist::T
end

const MatchDist = MatchingDistance

function get_cost_matrix_dynamic(
    d::MatchingDistance,
    X::Vector{T}, Y::Vector{T}
    ) where {T}

    if length(X) < length(Y)
        return get_cost_matrix_dynamic(d,Y,X)
    elseif length(X)==length(Y)
        return pairwise_inbounds(d.ground_dist, X, Y)
    else
        C = pairwise_inbounds(d.ground_dist, X, Y)
        null_dists = [d.ground_dist(nothing, p) for p in X]
        size_diff = length(X)-length(Y)
        return [C [x for x∈null_dists, j=1:size_diff]]
    end
end

function get_cost_matrix_fixed(
    d::MatchingDistance,
    X::Vector{T}, Y::Vector{T}
    ) where {T}

    if length(X) < length(Y)
        C = pairwise_inbounds(d.ground_dist, X, Y)
        pentalty_vec = [d.ground_dist(nothing, p) for p in Y]
        size_diff = length(Y) - length(X)
        return [C ; [y for i=1:size_diff, y ∈ pentalty_vec]]
    elseif length(X)==length(Y)
        return pairwise_inbounds(d.ground_dist, X, Y)
    else
        C = pairwise_inbounds(d.ground_dist, X, Y)
        null_dists = [d.ground_dist(nothing, p) for p in X]
        size_diff = length(X)-length(Y)
        return [C [x for x∈null_dists, j=1:size_diff]]
    end

end

struct FastMatchingDistance{T<:Metric} <: Metric
    ground_dist::T
    C::Matrix{Float64}
    function FastMatchingDistance(ground_dist::S, K::Int) where {S<:Metric}
        new{S}(ground_dist, zeros(K,K))
    end
end

const FastMatchDist = FastMatchingDistance

Base.show(io::IO, d::FastMatchingDistance) = print(io, "$(typeof(d))(K=$(size(d.C,1)))")

function get_cost_matrix_dynamic(
    d::FastMatchingDistance,
    X::Vector{T}, Y::Vector{T}
    ) where {T}

    N,M = (length(X), length(Y))
    if N < M
        return get_cost_matrix_dynamic(d,Y,X)
    elseif N==M
        C = view(d.C, 1:N, 1:N)
        pairwise_inbounds!(C, d.ground_dist, X, Y)
    else
        C = view(d.C, 1:N, 1:N)
        pairwise_inbounds!(C, d.ground_dist, X, Y)
        for i in 1:N
            C[i,M+1] = d.ground_dist(nothing, X[i])
        end
        for j in (M+2):N
            for i in 1:N
                C[i,j] = C[i,j-1]
            end
        end
    end
    return C
end

function get_cost_matrix_fixed(
    d::FastMatchingDistance,
    X::Vector{T}, Y::Vector{T}
    ) where {T}

    N,M = (length(X), length(Y))
    if N < M
        C = view(d.C, 1:M, 1:M)
        pairwise_inbounds!(C, d.ground_dist, X, Y)
        for j in 1:M
            C[N+1,j] = d.ground_dist(nothing, Y[j])
        end
        for i in (N+2):N
            for j in 1:M
                C[i,j] = C[i-1,j]
            end
        end
    elseif N==M
        C = view(d.C, 1:N, 1:N)
        pairwise_inbounds!(C, d.ground_dist, X, Y)
    else
        C = view(d.C, 1:N, 1:N)
        pairwise_inbounds!(C, d.ground_dist, X, Y)
        for i in 1:N
            C[i,M+1] = d.ground_dist(nothing, X[i])
        end
        for j in (M+2):M
            for i in 1:N
                C[i,j] = C[i,j-1]
            end
        end
    end
    return C

end

function (d::Union{MatchDist,FastMatchDist})(X::Nothing, Y::Vector{T})::Float64 where {T}
    return sum(p->d.ground_dist(nothing,p), Y)
end
function (d::Union{MatchDist,FastMatchDist})(X::Vector{T}, Y::Nothing)::Float64 where {T}
    return d(Y, X)
end
function (d::Union{MatchDist,FastMatchDist})(X::Nothing, Y::Nothing)::Float64 where {T}
    return 0.0
end

struct FixedPenaltyMatchingDistance{T<:Metric} <: Metric
    ground_dist::T
    penalty::Float64
end

const FixPenMatchDist = FixedPenaltyMatchingDistance
# const FixPenMatchDist{T} = FixedPenaltyMatchingDistance{T}
# FixPenMatchDist(d_g::Metric, penalty::Float64) = FixedPenaltyMatchingDistance(d_g, penalty)

Base.show(io::IO, d::FixPenMatchDist) = print(io, "$(typeof(d))(ρ=$(d.penalty))")

function get_cost_matrix_dynamic(
    d::FixPenMatchDist,
    X::Vector{T}, Y::Vector{T}
    ) where {T}

    d_g = d.ground_dist
    N, M = (length(X), length(Y))
    if N < M
        return get_cost_matrix_dynamic(d, Y, X)
    elseif N == M
        return pairwise_inbounds(d_g, X, Y)
    else
        C = pairwise_inbounds(d_g, X, Y)
        size_diff = (N - M)
        return [C [d.penalty for i ∈ eachindex(X), j=1:size_diff]]
    end
end

function get_cost_matrix_fixed(
    d::FixPenMatchDist,
    X::Vector{T}, Y::Vector{T}
    ) where {T}

    d_g = d.ground_dist
    N, M = (length(X), length(Y))
    C = pairwise_inbounds(d_g, X, Y)
    if N < M
        size_diff = (M - N)
        return [C ; [d.penalty for i in 1:size_diff, j ∈ eachindex(Y)]]
    elseif N > M
        size_diff = (N - M)
        return [C [d.penalty for i ∈ eachindex(X), j in 1:size_diff]]
    else
        return pairwise_inbounds(d_g, X, Y)
    end

end

# function (d::FixPenMatchDist)(
#     S1::Vector{T}, S2::Vector{T}
#     ) where {T}

#     length(S1) < length(S2) ? (N_min, N_max)=(length(S1),length(S2)) : (N_min,N_max)=(length(S2),length(S1))

#     C = fill(d.penalty, N_max, N_max)
#     @views Distances.pairwise!(C, d.ground_dist, S1, S2)

#     # If we have max_dist > 2*ρ we must extend C to allow for interactions in BOTH observations being un-matched.
#     if maximum(view(C, 1:length(S1), 1:length(S2))) > 2*d.penalty
#         C = [
#             C fill(d.penalty/2, N_max, N_min);
#             fill(d.penalty/2, N_min, N_max) fill(0.0, N_min, N_min)
#             ]
#     end
#     return hungarian(C)[2]
# end

function (d::FixPenMatchDist)(X::Nothing, Y::Vector{T})::Float64 where {T}
    return d.penalty * length(Y)
end
function (d::FixPenMatchDist)(X::Vector{T}, Y::Nothing)::Float64 where {T}
    d(Y, X)
end
function (d::FixPenMatchDist)(X::Nothing, Y::Nothing)::Float64 where {T}
    return 0.0
end

struct AvgSizeMatchingDistance{T<:Metric} <: Metric
    ground_dist::T
    penalty::Float64
end

const AvgSizeMatchDist = AvgSizeMatchingDistance
# const AvgSizeMatchDist{T} = AvgSizeMatchingDistance{T}

Base.show(io::IO, d::AvgSizeMatchDist) = print(io, "$(typeof(d))(ρ=$(d.penalty))")

function get_cost_matrix_dynamic(
    d::AvgSizeMatchDist,
    X::Vector{T}, Y::Vector{T}
    ) where {T}

    d_g = d.ground_dist
    N, M = (length(X), length(Y))
    if N < M
        return get_cost_matrix_dynamic(d, Y, X)
    elseif N == M
        return pairwise_inbounds(d_g, X, Y)
    else
        C = pairwise_inbounds(d_g, X, Y)
        mean_size = mean(y -> d_g(y, nothing), Y) # Find mean size in smaller set
        pentalty_vec = [abs(d_g(x, nothing) - mean_size) + d.penalty for x in X]
        size_diff = (N - M)
        return [C [x for x ∈ pentalty_vec, j=1:size_diff]]
    end
end

function get_cost_matrix_fixed(
    d::AvgSizeMatchDist,
    X::Vector{T}, Y::Vector{T}
    ) where {T}

    d_g = d.ground_dist
    N, M = (length(X), length(Y))
    C = pairwise_inbounds(d_g, X, Y)
    if N < M
        mean_size = mean(x->d_g(x,nothing), X)
        penalty_vec = [abs(d_g(y,nothing) - mean_size) + d.penalty for y in Y]
        size_diff = (M - N)
        return [C ; [y for i in 1:size_diff, y ∈ penalty_vec]]
    elseif N > M
        mean_size = mean(y->d_g(y,nothing), Y)
        penalty_vec = [abs(d_g(x,nothing) - mean_size) + d.penalty for x in X]
        size_diff = (N - M)
        return [C [x for x ∈ penalty_vec, i in 1:size_diff]]
    else
        return pairwise_inbounds(d_g, X, Y)
    end

end

function (d::AvgSizeMatchDist)(X::Nothing, Y::Vector{T})::Float64 where {T}
    return (d.penalty * length(Y)) + sum(x->dist.ground_dist(x,nothing), Y)
end
function (d::AvgSizeMatchDist)(X::Vector{T}, Y::Nothing)::Float64 where {T}
    d(Y, X)
end
function (d::AvgSizeMatchDist)(X::Nothing, Y::Nothing)::Float64 where {T}
    return 0.0
end

struct MinDistMatchingDistance{T<:Metric} <: SemiMetric
    ground_dist::T
    penalty::Float64
end

const MinDistMatchDist{T} = MinDistMatchingDistance where {T}

Base.show(io::IO, d::MinDistMatchDist) = print(io, "$(typeof(d))(ρ=$(d.penalty))")

function get_cost_matrix_dynamic(
    d::MinDistMatchDist,
    X::Vector{T}, Y::Vector{T}
    ) where {T}

    d_g = d.ground_dist
    N, M = (length(X), length(Y))
    if N < M
        return get_cost_matrix_dynamic(d, Y, X)
    elseif N == M
        return pairwise_inbounds(d_g, X, Y)
    else
        C = pairwise_inbounds(d_g, X, Y)
        penalty_vec = map(minimum, eachrow(C))
        size_diff = (N - M)
        return [C [x for x ∈ penalty_vec, j=1:size_diff]]
    end
end

function get_cost_matrix_fixed(
    d::MinDistMatchDist,
    X::Vector{T}, Y::Vector{T}
    ) where {T}

    d_g = d.ground_dist
    N, M = (length(X), length(Y))
    C = pairwise_inbounds(d_g, X, Y)
    if N < M
        penalty_vec = map(minimum, eachcol(C))
        size_diff = (M - N)
        return [C ; [y for i in 1:size_diff, y ∈ penalty_vec]]
    elseif N > M
        penalty_vec = map(minimum, eachrow(C))
        size_diff = (N - M)
        return [C [x for x ∈ penalty_vec, i in 1:size_diff]]
    else
        return pairwise_inbounds(d_g, X, Y)
    end

end

# Define type union for all these complete matching distances

const CompleteMatchingDistance = Union{MatchingDistance, FastMatchingDistance, FixedPenaltyMatchingDistance, AvgSizeMatchingDistance, MinDistMatchingDistance}


using LinearAlgebra
using FermiCG
using Printf
using Test
using Arpack

@testset "fci" begin
    atoms = []
    push!(atoms,Atom(1,"H",[0,0,0]))
    push!(atoms,Atom(2,"H",[0,0,1]))
    push!(atoms,Atom(3,"H",[0,0,2]))
    push!(atoms,Atom(4,"H",[0,0,3]))
    push!(atoms,Atom(5,"H",[0,0,4]))
    push!(atoms,Atom(6,"H",[0,0,5]))
    #basis = "6-31g"
    basis = "sto-3g"
    mol     = Molecule(0,1,atoms,basis)

    mf = FermiCG.pyscf_do_scf(mol)
    nbas = size(mf.mo_coeff)[1]
    ints = FermiCG.pyscf_build_ints(mol,mf.mo_coeff, zeros(nbas,nbas));
    e_fci, d1_fci, d2_fci = FermiCG.pyscf_fci(ints,3,3)
    @printf(" FCI Energy: %12.8f\n", e_fci)

    n_elec_a = 3
    n_elec_b = 3

    norb = size(ints.h1)[1]
    problem = FermiCG.StringCI.FCIProblem(norb, n_elec_a, n_elec_b)

    display(problem)

    @time Hmat = FermiCG.StringCI.build_H_matrix(ints, problem)
    @time e,v = eigs(Hmat, nev = 10, which=:SR)
    e = real(e)
    for ei in e
        @printf(" Energy: %12.8f\n",ei+ints.h0)
    end
    @test isapprox(e[1], e_fci , atol=1e-10)
end


### A Pluto.jl notebook ###
# v0.17.3

using Markdown
using InteractiveUtils

# ╔═╡ 89001295-1da7-496b-8c72-475520512b8a
begin
	import Pkg
	Pkg.activate(Base.current_project())
	using CairoMakie
	using GeometryBasics
	using PlutoUI
	using Rotations
	using FringeExplorer
	using LinearAlgebra
	using ERFA # For sepp until we add a composable version to FringeExplorer.
	md"""# Coordinate Explorations"""
end

# ╔═╡ 3639b720-f3c6-49e9-9106-84398cc332df
begin
	f=Figure()
	a=Axis(f[1,1], limits=((-5,5),(-5,5)), aspect=1)
	dy = sin(π/3)
	for y in -4:4
		xmin = -4+abs(y/2)
		xmax = - xmin
		for x in xmin:xmax
			lines!(Circle(Point2f0(x,y*dy), 0.5))
		end
	end
	f
end

# ╔═╡ e1e798aa-259a-435b-baea-857bf893cde3
begin
	f2=Figure()
	a2=Axis(f2[1,1], limits=((-5,5),(-5,5)), aspect=1)
	lines!(Circle(Point2f(0), 0.5))
	for r in 1:4
		nx = 6r
		for i in 0:nx-1
			θ = i*2π/nx
			x = r*cos(θ)
			y = r*sin(θ)
			lines!(Circle(Point2f(x,y), 0.5))
		end
	end
	f2
end

# ╔═╡ 3319a555-62b0-44ac-a2e9-01ebeca5a07b
fwhm = deg2rad(4)

# ╔═╡ 62fa9f61-3aa4-4542-9629-67e93198d5b1
α0=5π/4

# ╔═╡ 48266d0e-c873-4597-a981-e86c47250a8b
δ0=π/6

# ╔═╡ 34b142aa-0b8f-4311-b922-88fddc7fdc15
centers=beamrings(α0, δ0; nrings=4, dϕ=fwhm)

# ╔═╡ f657d533-e777-4ec0-8367-c035279d9d20
centerc = FringeExplorer.s2c(centers)

# ╔═╡ 10fb9c7d-bcc3-45ce-8cf2-eca33ab3a5d0
# Cross beam center vectors with Z axis to get rotation axes
crosses = cross.(Ref([0.0,0.0,1.0]), eachcol(centerc))

# ╔═╡ 0417c07f-47d3-4c12-8960-333c78f3288a
beamseps = sepp.(Ref([0.0,0.0,1.0]), eachcol(centerc))

# ╔═╡ 8c965934-14da-4bf8-b30c-8ad94ee601cd
begin
	f3=Figure()
	a3=Axis3(f3[1,1], limits=((-1,1),(-1,1),(-1,1)),
		              aspect=(1,1,1),
					  azimuth=deg2rad(210))
	s3=Sphere(Point3f(0), 1.0)
	c3=Circle(Point2f(0), 1.0)#tan(fwhm/2))
	#arrows!([0,0,0], [0,0,0], [0,0,0], [1,0,0], [0,1,0], [0,0,1])
	mesh!(s3, color=RGBAf(1.0, 1.0, 1.0, 0.4))
	lines!(c3, color=RGBAf(0.5, 0.5, 0.5, 0.5))
	for ax in (Vec3f(0,1,0), Vec3f(1,0,0))
		transformation=Transformation(;rotation=qrotation(ax, π/2))
		lines!(c3; color=RGBAf(0.5, 0.5, 0.5, 0.5), transformation)
	end
	for (i, beam) in enumerate(eachcol(centerc))
		trans = Transformation(;scale=Vec3f(tan(fwhm/2)),
								translation=beam,
								rotation=qrotation(Vec3f(crosses[i]), beamseps[i]))
		lines!(c3, transformation=trans)
	end
	f3
end

# ╔═╡ Cell order:
# ╟─89001295-1da7-496b-8c72-475520512b8a
# ╟─3639b720-f3c6-49e9-9106-84398cc332df
# ╟─e1e798aa-259a-435b-baea-857bf893cde3
# ╠═3319a555-62b0-44ac-a2e9-01ebeca5a07b
# ╠═62fa9f61-3aa4-4542-9629-67e93198d5b1
# ╠═48266d0e-c873-4597-a981-e86c47250a8b
# ╠═34b142aa-0b8f-4311-b922-88fddc7fdc15
# ╠═f657d533-e777-4ec0-8367-c035279d9d20
# ╠═10fb9c7d-bcc3-45ce-8cf2-eca33ab3a5d0
# ╠═0417c07f-47d3-4c12-8960-333c78f3288a
# ╠═8c965934-14da-4bf8-b30c-8ad94ee601cd


include("trapezoid_rule.jl");

 function ordertolen(order::Int64)
    #Three dimensional tensor product basis.
    expsz=3^order;

    #Using vector of three components.
    expsz=3*expsz;

    return expsz;
end

 function evaluate_tensor_basis(i::Int64,j::Int64,k::Int64,x::Float64,y::Float64,z::Float64)
    return exp(im*(i*x+j*y+k*z));
end

 function evaluate_tensor_basis_derivative(i::Int64,j::Int64,k::Int64,d::Int64,x::Float64,y::Float64,z::Float64)
    out=exp(im*(i*x+j*y+k*z));
    is=[i,j,k];
    out=im*is[d]*out;
    return out;
end
 function evaluate_tensor_basis_derivative(i::Int64,j::Int64,k::Int64,d1::Int64,d2::Int64,x::Float64,y::Float64,z::Float64)
    out=exp(im*(i*x+j*y+k*z));
    is=[i,j,k];
    out=(im*is[d2])*(im*is[d1])*out;
    return out;
end

 function evaluate_curl_basis_gradient(order::Int64,n::Int64,x::Float64,y::Float64,z::Float64)
    ncomps=3;
    m=div(n,ncomps);
    r=n%ncomps;
    (id1,id2,id3)=ind2sub((order,order,order),m);
    out=Array(Complex128,(ncomps,ncomps));
    if r==0
        dyx=evaluate_tensor_basis_derivative(id1,id2,id3,2,1,x,y,z);
        dzx=evaluate_tensor_basis_derivative(id1,id2,id3,3,1,x,y,z);
        dyy=evaluate_tensor_basis_derivative(id1,id2,id3,2,2,x,y,z);
        dzy=evaluate_tensor_basis_derivative(id1,id2,id3,3,2,x,y,z);
        dyz=evaluate_tensor_basis_derivative(id1,id2,id3,2,3,x,y,z);
        dzz=evaluate_tensor_basis_derivative(id1,id2,id3,3,3,x,y,z);
        out[:,1]=[0.0*im,dzx,-dyx];
        out[:,2]=[0.0*im,dzy,-dyy];
        out[:,3]=[0.0*im,dzy,-dyy];
    elseif r==1
        dxx=evaluate_tensor_basis_derivative(id1,id2,id3,1,1,x,y,z);
        dzx=evaluate_tensor_basis_derivative(id1,id2,id3,3,1,x,y,z);
        dxy=evaluate_tensor_basis_derivative(id1,id2,id3,1,2,x,y,z);
        dzy=evaluate_tensor_basis_derivative(id1,id2,id3,3,2,x,y,z);
        dxz=evaluate_tensor_basis_derivative(id1,id2,id3,1,3,x,y,z);
        dzz=evaluate_tensor_basis_derivative(id1,id2,id3,3,3,x,y,z);

        out[:,1]=[-dzx,0.0*im,dxx];
        out[:,2]=[-dzy,0.0*im,dxy];
        out[:,3]=[-dzz,0.0*im,dxz];
    else
        dxx=evaluate_tensor_basis_derivative(id1,id2,id3,1,1,x,y,z);
        dyx=evaluate_tensor_basis_derivative(id1,id2,id3,2,1,x,y,z);
        dxy=evaluate_tensor_basis_derivative(id1,id2,id3,1,2,x,y,z);
        dyy=evaluate_tensor_basis_derivative(id1,id2,id3,2,2,x,y,z);
        dxz=evaluate_tensor_basis_derivative(id1,id2,id3,1,3,x,y,z);
        dyz=evaluate_tensor_basis_derivative(id1,id2,id3,2,3,x,y,z);
        out[:,1]=[dyx,-dxx,0.0*im];
        out[:,2]=[dyy,-dxy,0.0*im];
        out[:,3]=[dyz,-dxz,0.0*im];
    end
    return out;
end

 function evaluate_curl_basis(order::Int64,n::Int64,x::Float64,y::Float64,z::Float64)
    ncomps=3;
    m=div(n,ncomps);
    r=n%ncomps;
    (id1,id2,id3)=ind2sub((order,order,order),m);
    out=(0.0*im,0.0*im,0.0*im);
    if r==0
        dy=evaluate_tensor_basis_derivative(id1,id2,id3,2,x,y,z);
        dz=evaluate_tensor_basis_derivative(id1,id2,id3,3,x,y,z);
        out=(0.0*im,dz,-dy);
    elseif r==1
        dx=evaluate_tensor_basis_derivative(id1,id2,id3,1,x,y,z);
        dz=evaluate_tensor_basis_derivative(id1,id2,id3,3,x,y,z);
        out=(-dz,0.0*im,dx);
    else
        dx=evaluate_tensor_basis_derivative(id1,id2,id3,1,x,y,z);
        dy=evaluate_tensor_basis_derivative(id1,id2,id3,2,x,y,z);
        out=(dy,-dx,0.0*im);
    end
    return out;
end

 function curl_basis_inner_product(order::Int64,m::Int64,n::Int64)
    ncomps=3;
    rm=m%ncomps;
    rn=n%ncomps;
    ordm=div(m,ncomps);
    ordn=div(n,ncomps);

    (id1,id2,id3)=ind2sub((order,order,order),ordm);
    idsm=[id1,id2,id3];
    (id1,id2,id3)=ind2sub((order,order,order),ordn);
    idsn=[id1,id2,id3];

    val=Float64((2*pi)*(2*pi)*(2*pi));

    rm=rm+1;
    rn=rn+1;
    if idsm==idsn
        if rm==rn
            if rm==1
                val=val*(idsn[3]*idsn[3]+idsn[2]*idsn[2])
            elseif rm==2
                val=val*(idsn[3]*idsn[3]+idsn[1]*idsn[1]);
            else
                val=val*(idsn[2]*idsn[2]+idsn[1]*idsn[1]);
            end
        else
            val=-val*idsn[rn]*idsm[rm];
        end
    else
        val=0.0;
    end

    return val;
end

 function curl_basis_mass_matrix(order::Int64,quadlen::Int64)
    n=ordertolen(order);
    M=Array(Complex128,(n,n));
    (x,y,z,W)=trapezoid_tensor_product(-pi,pi,quadlen);
    V1=copy(W)*im;
    V2=copy(V1);
    V3=copy(V1);
    U1=copy(V1);
    U2=copy(V1);
    U3=copy(V1);

    for i = 1 : n
        for j = 1 : n
            #Evaluate basis at quadrature nodes.
            for u1 = 1 : length(x)
                for u2 = 1 : length(y)
                    for u3 = 1 : length(z)
                        (vx,vy,vz)=evaluate_curl_basis(order,i,x[u1],y[u2],z[u3]);
                        (ux,uy,uz)=evaluate_curl_basis(order,j,x[u1],y[u2],z[u3]);
                        V1[u1,u2,u3]=vx;
                        V2[u1,u2,u3]=vy;
                        V3[u1,u2,u3]=vz;

                        U1[u1,u2,u3]=ux;
                        U2[u1,u2,u3]=uy;
                        U3[u1,u2,u3]=uz;

                    end
                end
            end

            S=conj(V1).*U1+conj(V2).*U2+conj(V3).*U3;

            M[i,j]=sum(S[:].*W[:]);
        end
    end

    return M;
end

 function filter_dependent_modes(order::Int64,quadlen::Int64)
    M=curl_basis_mass_matrix(order,quadlen);
    r=rank(M);
    (Q,R,p)=qr(M,Val{true});
    return p[1:r];
end

 function filter_dependent_modes(order::Int64)
    M=full(curl_basis_mass_matrix(order));
    r=rank(M);
    (Q,R,p)=qr(M,Val{true});
    return p[1:r];
end

 function filter_dependent_modes_new(order::Int64)
    M=curl_basis_mass_matrix(order);
    (m,n)=size(M);
    S=IntSet([1:m;]);
    for s in S
      for t in S
	if t!=s
	  if M.rowval[M.colptr[t]:M.colptr[t+1]-1]==M.rowval[M.colptr[s]:M.colptr[s+1]-1]
	    v=(norm(M[:,s]'*M[:,t]))/(norm(M[:,s])*norm(M[:,t]));
	    if abs(v-1.0)<1e-3
	      print("eliminated pair=($s,$t)\n");
	      print("diff=($(s-t))\n");
	      print("rel diff=($((s-t)/n))\n");
	      delete!(S,t);
	    end
	  end
	end
      end
    end
    return [Int64(p) for p in S];
end

 function filter_dependent_modes_new(order::Int64,alpha)
    M=curl_basis_mass_matrix(order);
    (m,n)=size(M);
    S=IntSet([1:m;]);
    for s in S
      for t in S
	if t!=s
	  v=(norm(M[:,s]'*M[:,t]))/(norm(M[:,s])*norm(M[:,t]));
	  if abs(v-1.0)<alpha
	    print("eliminated pair=($s,$t)\n");
	    print("diff=($(s-t))\n");
	    print("rel diff=($((s-t)/n))\n");
	    delete!(S,t);
	  end
	end
      end
    end
    return [Int64(p) for p in S];
end

 function curl_basis_mass_matrix(p::Array{Int64},order::Int64,quadlen::Int64)
    n=length(p);
    M=Array(Complex128,(n,n));
    (x,y,z,W)=trapezoid_tensor_product(-pi,pi,quadlen);
    V1=copy(W)*im;
    V2=copy(V1);
    V3=copy(V1);
    U1=copy(V1);
    U2=copy(V1);
    U3=copy(V1);

    for i = 1 : n
        for j = 1 : n
            #Evaluate basis at quadrature nodes.
            for u1 = 1 : length(x)
                for u2 = 1 : length(y)
                    for u3 = 1 : length(z)
                        (vx,vy,vz)=evaluate_curl_basis(order,p[i],x[u1],y[u2],z[u3]);
                        (ux,uy,uz)=evaluate_curl_basis(order,p[j],x[u1],y[u2],z[u3]);
                        V1[u1,u2,u3]=vx;
                        V2[u1,u2,u3]=vy;
                        V3[u1,u2,u3]=vz;

                        U1[u1,u2,u3]=ux;
                        U2[u1,u2,u3]=uy;
                        U3[u1,u2,u3]=uz;

                    end
                end
            end

            S=conj(V1).*U1+conj(V2).*U2+conj(V3).*U3;

            M[i,j]=sum(S[:].*W[:]);
        end
    end

    return M;
end

 function curl_basis_mass_matrix(p::Array{Int64},order::Int64)
    n=length(p);
    M=Array(Complex128,(n,n));
    for i = 1 : n
        for j = 1 : n
           M[i,j]=curl_basis_inner_product(order,p[i],p[j]);
        end
    end

    return M;
end

 function curl_basis_mass_matrix(order::Int64)
    n=ordertolen(order);
    M=spzeros(Complex128,n,n);
    for i = 1 : n
        for j = 1 : n
           M[i,j]=curl_basis_inner_product(order,i,j);
        end
    end
    return M;
end

 function eval_curl_soln(
    u::Array{Complex128},
    p::Array{Int64},
    order::Int64,
    x,y,z)

    outx::Complex128=0.0;
    outy::Complex128=0.0;
    outz::Complex128=0.0;
    for i = 1 : length(u)
        (ux,uy,uz)=evaluate_curl_basis(order,p[i],x,y,z);
        outx=outx+u[i]*ux;
        outy=outy+u[i]*uy;
        outz=outz+u[i]*uz;
    end
    return (outx,outy,outz);
end

function curl_basis_vandermondes(
    order::Int64,
    p::Array{Int64},
    qx::Array{Float64},
    qy::Array{Float64},
    qz::Array{Float64},
    qw::Array{Float64})

    n=length(p);
    nq=length(qw);
    V=Array(Array{Complex128},(3,));
    V[1]=Array(Complex128,(nq,n));
    V[2]=Array(Complex128,(nq,n));
    V[3]=Array(Complex128,(nq,n));

    for i = 1 : n
        for iq = 1 : nq
            (i1,i2,i3)=ind2sub(size(qw),iq);
            (ux,uy,uz)=evaluate_curl_basis(order,p[i],qx[i1],qy[i2],qz[i3]);
            V[1][iq,i]=ux;
            V[2][iq,i]=uy;
            V[3][iq,i]=uz;
        end
    end
    return V;
end

function curl_basis_gradient_vandermondes(
    order::Int64,
    p::Array{Int64},
    qx::Array{Float64},
    qy::Array{Float64},
    qz::Array{Float64},
    qw::Array{Float64})

    n=length(p);
    nq=length(qw);
    GV  = Array(Array{Complex128},(3,3));
    for i = 1 : 3
        for j = 1 : 3
            GV[i,j]=Array(Complex128,(nq,n));
        end
    end

    for i = 1 : n
        for iq = 1 : nq
            (i1,i2,i3)=ind2sub(size(qw),iq);
            U=evaluate_curl_basis_gradient(order,p[i],qx[i1],qy[i2],qz[i3]);
            for di=1:3
                for dj=1:3
                    GV[di,dj][iq,i]=U[di,dj];
                end
            end
        end
    end
    return GV;
end


module Shapes # like your Geometry module

export Sphere, Lens, Retina

using Vector3D

abstract Sphere

type Lens <: Sphere # the lens type
    c::Vec # center
    r::Float64 # radius
    ri::Function # the Refractive Index Gradient (RIG) that describes the distribution of refractive indices inside the lens. We restrict ourselves to spherically symmetrical distributions: so at a set distance form the lens center, the refractive index is equal in all directions (kind of like an onion?).
    ε::Float64 # the step
end
Lens(a::Vector{Float64},b::Float64,c::Function,d::Float64) = Lens(Vec(a),b,c,d) # this is a constructor function

type Retina <: Sphere # the retina type
    c::Vec # retina center
    r::Float64 # radius
end
Retina(a::Vector{Float64},b::Float64) = Retina(Vec(a),b)

end


module AdditionTest

using Test
using ForneyLab
using ForneyLab: outboundType, isApplicable
using ForneyLab: SPAdditionOutNGG, SPAdditionOutNGP, SPAdditionOutNPG, SPAdditionIn1GNG, SPAdditionIn1PNG, SPAdditionIn2GGN, SPAdditionIn2PGN, SPAdditionIn1GNP, SPAdditionIn2GPN, SPAdditionOutNPP, SPAdditionIn1PNP, SPAdditionIn2PPN, MAdditionNGG

@testset "Addition node construction through + syntax" begin
    g = FactorGraph()
    @RV x ~ GaussianMeanVariance(constant(0.0), constant(1.0))
    @RV y ~ GaussianMeanVariance(constant(0.0), constant(1.0))
    @RV z = x + y
    @test isa(z, Variable)
    @test isa(g.nodes[:addition_1], Addition)

    g = FactorGraph()
    @RV x ~ GaussianMeanVariance(constant(0.0), constant(1.0))
    @RV z = x + 1.0
    @test isa(z, Variable)
    @test isa(g.nodes[:addition_1], Addition)
end

@testset "Addition node construction through - syntax" begin
    g = FactorGraph()
    @RV x ~ GaussianMeanVariance(constant(0.0), constant(1.0))
    @RV y ~ GaussianMeanVariance(constant(0.0), constant(1.0))
    @RV z = x - y
    @test isa(z, Variable)
    @test isa(g.nodes[:addition_1], Addition)
    # This syntax works by changing the order interfaces are attached in. Below checks whether the pairings are succesful by
    # matching out and in2 interfaces of the corresponding nodes
    @test g.nodes[:addition_1].i[:out] == g.nodes[:gaussianmeanvariance_1].i[:out].partner
    @test g.nodes[:addition_1].i[:in2] == g.nodes[:gaussianmeanvariance_2].i[:out].partner

    g = FactorGraph()
    @RV x ~ GaussianMeanVariance(constant(0.0), constant(1.0))
    @RV z = x - 1.0
    @test isa(z, Variable)
    @test isa(g.nodes[:addition_1], Addition)
end

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

@testset "SPAdditionOutNGG" begin
    @test SPAdditionOutNGG <: SumProductRule{Addition}
    @test outboundType(SPAdditionOutNGG) == Message{GaussianMeanVariance}
    @test isApplicable(SPAdditionOutNGG, [Nothing, Message{Gaussian}, Message{Gaussian}])
    @test !isApplicable(SPAdditionOutNGG, [Nothing, Message{PointMass}, Message{PointMass}])
    @test !isApplicable(SPAdditionOutNGG, [Message{Gaussian}, Message{Gaussian}, Nothing])

    @test ruleSPAdditionOutNGG(nothing, Message(Univariate, GaussianMeanVariance, m=1.0, v=2.0), Message(Univariate, GaussianMeanVariance, m=3.0, v=4.0)) == Message(Univariate, GaussianMeanVariance, m=4.0, v=6.0)
    @test ruleSPAdditionOutNGG(nothing, Message(Multivariate, GaussianMeanVariance, m=[1.0], v=mat(2.0)), Message(Multivariate, GaussianMeanVariance, m=[3.0], v=mat(4.0))) == Message(Multivariate, GaussianMeanVariance, m=[4.0], v=mat(6.0))
end

@testset "SPAdditionIn2GGN" begin
    @test SPAdditionIn2GGN <: SumProductRule{Addition}
    @test outboundType(SPAdditionIn2GGN) == Message{GaussianMeanVariance}
    @test isApplicable(SPAdditionIn2GGN, [Message{Gaussian}, Message{Gaussian}, Nothing])

    @test ruleSPAdditionIn2GGN(Message(Univariate, GaussianMeanVariance, m=3.0, v=4.0), Message(Univariate, GaussianMeanVariance, m=1.0, v=2.0), nothing) == Message(Univariate, GaussianMeanVariance, m=2.0, v=6.0)
    @test ruleSPAdditionIn2GGN(Message(Multivariate, GaussianMeanVariance, m=[3.0], v=mat(4.0)), Message(Multivariate, GaussianMeanVariance, m=[1.0], v=mat(2.0)), nothing) == Message(Multivariate, GaussianMeanVariance, m=[2.0], v=mat(6.0))
end

@testset "SPAdditionIn1GNG" begin
    @test SPAdditionIn1GNG <: SumProductRule{Addition}
    @test outboundType(SPAdditionIn1GNG) == Message{GaussianMeanVariance}
    @test isApplicable(SPAdditionIn1GNG, [Message{Gaussian}, Nothing, Message{Gaussian}])

    @test ruleSPAdditionIn1GNG(Message(Univariate, GaussianMeanVariance, m=3.0, v=4.0), nothing, Message(Univariate, GaussianMeanVariance, m=1.0, v=2.0)) == Message(Univariate, GaussianMeanVariance, m=2.0, v=6.0)
    @test ruleSPAdditionIn1GNG(Message(Multivariate, GaussianMeanVariance, m=[3.0], v=mat(4.0)), nothing, Message(Multivariate, GaussianMeanVariance, m=[1.0], v=mat(2.0))) == Message(Multivariate, GaussianMeanVariance, m=[2.0], v=mat(6.0))
end

@testset "SPAdditionOutNGP" begin
    @test SPAdditionOutNGP <: SumProductRule{Addition}
    @test outboundType(SPAdditionOutNGP) == Message{GaussianMeanVariance}
    @test isApplicable(SPAdditionOutNGP, [Nothing, Message{Gaussian}, Message{PointMass}])

    @test ruleSPAdditionOutNGP(nothing, Message(Univariate, GaussianMeanVariance, m=1.0, v=2.0), Message(Univariate, PointMass, m=3.0)) == Message(Univariate, GaussianMeanVariance, m=4.0, v=2.0)
    @test ruleSPAdditionOutNGP(nothing, Message(Multivariate, GaussianMeanVariance, m=[1.0], v=mat(2.0)), Message(Multivariate, PointMass, m=[3.0])) == Message(Multivariate, GaussianMeanVariance, m=[4.0], v=mat(2.0))
end

@testset "SPAdditionOutNPG" begin
    @test SPAdditionOutNPG <: SumProductRule{Addition}
    @test outboundType(SPAdditionOutNPG) == Message{GaussianMeanVariance}
    @test isApplicable(SPAdditionOutNPG, [Nothing, Message{PointMass}, Message{Gaussian}])

    @test ruleSPAdditionOutNPG(nothing, Message(Univariate, PointMass, m=3.0), Message(Univariate, GaussianMeanVariance, m=1.0, v=2.0)) == Message(Univariate, GaussianMeanVariance, m=4.0, v=2.0)
    @test ruleSPAdditionOutNPG(nothing, Message(Multivariate, PointMass, m=[3.0]), Message(Multivariate, GaussianMeanVariance, m=[1.0], v=mat(2.0))) == Message(Multivariate, GaussianMeanVariance, m=[4.0], v=mat(2.0))
end

@testset "SPAdditionIn1PNG" begin
    @test SPAdditionIn1PNG <: SumProductRule{Addition}
    @test outboundType(SPAdditionIn1PNG) == Message{GaussianMeanVariance}
    @test isApplicable(SPAdditionIn1PNG, [Message{PointMass}, Nothing, Message{Gaussian}])

    @test ruleSPAdditionIn1PNG(Message(Univariate, PointMass, m=3.0), nothing, Message(Univariate, GaussianMeanVariance, m=1.0, v=2.0)) == Message(Univariate, GaussianMeanVariance, m=2.0, v=2.0)
    @test ruleSPAdditionIn1PNG(Message(Multivariate, PointMass, m=[3.0]), nothing, Message(Multivariate, GaussianMeanVariance, m=[1.0], v=mat(2.0))) == Message(Multivariate, GaussianMeanVariance, m=[2.0], v=mat(2.0))
end

@testset "SPAdditionIn2PGN" begin
    @test SPAdditionIn2PGN <: SumProductRule{Addition}
    @test outboundType(SPAdditionIn2PGN) == Message{GaussianMeanVariance}
    @test isApplicable(SPAdditionIn2PGN, [Message{PointMass}, Message{Gaussian}, Nothing])

    @test ruleSPAdditionIn2PGN(Message(Univariate, PointMass, m=3.0), Message(Univariate, GaussianMeanVariance, m=1.0, v=2.0), nothing) == Message(Univariate, GaussianMeanVariance, m=2.0, v=2.0)
    @test ruleSPAdditionIn2PGN(Message(Multivariate, PointMass, m=[3.0]), Message(Multivariate, GaussianMeanVariance, m=[1.0], v=mat(2.0)), nothing) == Message(Multivariate, GaussianMeanVariance, m=[2.0], v=mat(2.0))
end

@testset "SPAdditionIn1GNP" begin
    @test SPAdditionIn1GNP <: SumProductRule{Addition}
    @test outboundType(SPAdditionIn1GNP) == Message{GaussianMeanVariance}
    @test isApplicable(SPAdditionIn1GNP, [Message{Gaussian}, Nothing, Message{PointMass}])

    @test ruleSPAdditionIn1GNP(Message(Univariate, GaussianMeanVariance, m=3.0, v=4.0), nothing, Message(Univariate, PointMass, m=1.0)) == Message(Univariate, GaussianMeanVariance, m=2.0, v=4.0)
    @test ruleSPAdditionIn1GNP(Message(Multivariate, GaussianMeanVariance, m=[3.0], v=mat(4.0)), nothing, Message(Multivariate, PointMass, m=[1.0])) == Message(Multivariate, GaussianMeanVariance, m=[2.0], v=mat(4.0))
end

@testset "SPAdditionIn2GPN" begin
    @test SPAdditionIn2GPN <: SumProductRule{Addition}
    @test outboundType(SPAdditionIn2GPN) == Message{GaussianMeanVariance}
    @test isApplicable(SPAdditionIn2GPN, [Message{Gaussian}, Message{PointMass}, Nothing])

    @test ruleSPAdditionIn2GPN(Message(Univariate, GaussianMeanVariance, m=3.0, v=4.0), Message(Univariate, PointMass, m=1.0), nothing) == Message(Univariate, GaussianMeanVariance, m=2.0, v=4.0)
    @test ruleSPAdditionIn2GPN(Message(Multivariate, GaussianMeanVariance, m=[3.0], v=mat(4.0)), Message(Multivariate, PointMass, m=[1.0]), nothing) == Message(Multivariate, GaussianMeanVariance, m=[2.0], v=mat(4.0))
end

@testset "SPAdditionOutNPP" begin
    @test SPAdditionOutNPP <: SumProductRule{Addition}
    @test outboundType(SPAdditionOutNPP) == Message{PointMass}
    @test isApplicable(SPAdditionOutNPP, [Nothing, Message{PointMass}, Message{PointMass}])

    @test ruleSPAdditionOutNPP(nothing, Message(Univariate, PointMass, m=1.0), Message(Univariate, PointMass, m=3.0)) == Message(Univariate, PointMass, m=4.0)
    @test ruleSPAdditionOutNPP(nothing, Message(Multivariate, PointMass, m=[1.0]), Message(Multivariate, PointMass, m=[3.0])) == Message(Multivariate, PointMass, m=[4.0])
end

@testset "SPAdditionIn2PPN" begin
    @test SPAdditionIn2PPN <: SumProductRule{Addition}
    @test outboundType(SPAdditionIn2PPN) == Message{PointMass}
    @test isApplicable(SPAdditionIn2PPN, [Message{PointMass}, Message{PointMass}, Nothing])

    @test ruleSPAdditionIn2PPN(Message(Univariate, PointMass, m=3.0), Message(Univariate, PointMass, m=1.0), nothing) == Message(Univariate, PointMass, m=2.0)
    @test ruleSPAdditionIn2PPN(Message(Multivariate, PointMass, m=[3.0]), Message(Multivariate, PointMass, m=[1.0]), nothing) == Message(Multivariate, PointMass, m=[2.0])
end

@testset "SPAdditionIn1PNP" begin
    @test SPAdditionIn1PNP <: SumProductRule{Addition}
    @test outboundType(SPAdditionIn1PNP) == Message{PointMass}
    @test isApplicable(SPAdditionIn1PNP, [Message{PointMass}, Nothing, Message{PointMass}])

    @test ruleSPAdditionIn1PNP(Message(Univariate, PointMass, m=3.0), nothing, Message(Univariate, PointMass, m=1.0)) == Message(Univariate, PointMass, m=2.0)
    @test ruleSPAdditionIn1PNP(Message(Multivariate, PointMass, m=[3.0]), nothing, Message(Multivariate, PointMass, m=[1.0])) == Message(Multivariate, PointMass, m=[2.0])
end

@testset "MAdditionNGG" begin
    @test MAdditionNGG <: MarginalRule{Addition}
    @test isApplicable(MAdditionNGG, [Nothing, Message{Gaussian}, Message{Gaussian}])

    @test ruleMAdditionNGG(Message(Univariate, GaussianWeightedMeanPrecision, xi=1.0, w=2.0), Message(Univariate, GaussianWeightedMeanPrecision, xi=3.0, w=4.0), Message(Univariate, GaussianWeightedMeanPrecision, xi=5.0, w=6.0)) == ProbabilityDistribution(Multivariate, GaussianWeightedMeanPrecision, xi=[4.0, 6.0], w=[6.0 2.0; 2.0 8.0])
    @test ruleMAdditionNGG(Message(Multivariate, GaussianWeightedMeanPrecision, xi=[1.0], w=mat(2.0)), Message(Multivariate, GaussianWeightedMeanPrecision, xi=[3.0], w=mat(4.0)), Message(Multivariate, GaussianWeightedMeanPrecision, xi=[5.0], w=mat(6.0))) == ProbabilityDistribution(Multivariate, GaussianWeightedMeanPrecision, xi=[4.0, 6.0], w=[6.0 2.0; 2.0 8.0])
end

end # module


### A Pluto.jl notebook ###
# v0.14.5

using Markdown
using InteractiveUtils

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

# ╔═╡ fb56c064-b90d-11eb-124f-bd2478c9a5bf
using WAV, SampledSignals, Plots,PlutoUI

# ╔═╡ 56d13e40-49bc-4f45-9033-0a4003070da8
using UnicodePlots

# ╔═╡ a6fcf1e7-e477-41d3-b33a-08f521db740c
@bind a Slider(1:10)

# ╔═╡ 6813d955-d54d-47ac-8cec-5be4ea1847fe
begin
# a =  1   # Amplitude
f = 1# Frequency # in hz # No of oscillations of medium particles per unit time
t = 1/f   # Time Period
l = 100   # WaveLength
w = 2*pi*f # omega
v = f * l # Velocity
k = 1 # 2* pi / l # Wave number
padiff = 10
phdiff = k * padiff # phase difference
padiff = phdiff / k # path difference
phdiff = 10
m = 100 # mass I think
# e = 1/2 * (m * (w * a)^2) # energy of wave
fs = 8000 # sample rate
i = 0.0:1/fs:prevfloat(1.0)
alpha = pi # alpha

end

# ╔═╡ c5413acd-7828-4da0-9009-65fd61264637
begin

# some examples from WaveForm.jl
squarewave(x::Real) = ifelse(mod2pi(x) < π, 1.0, -1.0)

#Compute ``2\pi``-periodic square wave of `x` with a duty cycle `θ`and a peak amplitude ``1``.

function squarewave(x::Real, θ::Real)
    0 ≤ θ ≤ 1 || throw(DomainError(θ, "squwarewave(x, θ) is only defined for 0 ≤ θ ≤ 1."))
    ifelse(mod2pi(x) < 2π * θ, 1.0, -1.0)
end

#Compute ``1``-periodic square wave of `x` with a peak amplitude ``1``.
squarewave1(x::Real) = ifelse(mod(x, 1) < 1/2, 1.0, -1.0)

# Compute ``1``-periodic square wave of `x` with a duty cycle `θ` and a peak amplitude ``1``.
function squarewave1(x::Real, θ::Real)
    0 ≤ θ ≤ 1 || throw(DomainError(θ, "squwarewave1(x, θ) is only defined for 0 ≤ θ ≤ 1."))
    ifelse(mod(x, 1) < θ, 1.0, -1.0)
end

#Compute ``2\pi``-periodic triangle wave of `x` with a peak amplitude ``1``.
function trianglewave(x::Real)
    modx = mod2pi(x + π/2)
    ifelse(modx < π, 2modx/π - 1, -2modx/π + 3)
end

#Compute ``1``-periodic triangle wave of `x` with a peak amplitude ``1``.
function trianglewave1(x::Real)
    modx = mod(x + 1/4, 1.0)
    ifelse(modx < 1/2, 4modx - 1, -4modx + 3)
end
#Compute ``2\pi``-periodic sawtooth wave of `x` with a peak amplitude ``1``.
sawtoothwave(x::Real) = rem2pi(x, RoundNearest) / π
#Compute ``1``-periodic sawtooth wave of `x` with a peak amplitude ``1``.
sawtoothwave1(x::Real) = rem(x, 1.0, RoundNearest) * 2
end

# ╔═╡ df0fbb6c-bad6-435d-a931-2a4cd9ea935c
begin

vpart = a*w*cos.(w*i .- k*phdiff .+alpha) #velocity of a particle
apart = -a*(w^2)*sin.(w*i .- k*phdiff .+alpha) #velocity of a particle
# y = a*sin(w*t-k*phdiff +alpha) # general equation
# y1 = a*sin(w*(i - phdiff/v)) # 1st Equation
# y2 = a*sin((2*pi/l)*(v*i-phdiff)) # 2nd Equation
# y3 = a*sin((2*pi)*(i/t - phdiff/l)) # different forms of same equation

wave = a*sin.(w*i .- k*phdiff .+alpha)
wave1 = a*sin.(w*i .- k*phdiff .+alpha) + a*sin.(w*i .- k*phdiff .+alpha .+pi )
wave2 = a*sin.(w*i .- k*phdiff .+alpha) + a*sin.(w*i .- k*phdiff .+alpha .+pi ) .+ a*squarewave.(w*i .- k*phdiff .+alpha .+pi )
wave3 = a*squarewave.(w*i .- k*phdiff .+alpha .+pi )
# Wave Source # Lossless medium
# fs = 8e3
# t = 0.0:1/fs:prevfloat(1.0)
# f = 1e3
y = sin.(2pi * f * i) * 0.1
wavplay(y, fs)
end

# ╔═╡ 683fa835-7339-4b7c-89fe-6511da652e2d
length(wave)

# ╔═╡ 17af4d9f-687a-48c8-9a68-10d9bb713187
lineplot(i, wave)

# ╔═╡ 263973d5-9d5c-427b-94cb-bd74073a7a29
lineplot(i, wave1)

# ╔═╡ 8c401e62-30b2-492d-a896-b0dceb2c8508
lineplot(i, wave2)

# ╔═╡ aa1d8f3a-09cd-4fee-b6a0-e959650b0faa
lineplot(i, wave3)

# ╔═╡ a5298c48-a277-4c1c-bb35-e129cb205df2
lineplot(i, y)

# ╔═╡ e4b55da5-2cba-46b0-a084-2ebbb93a31f9
# plot(i, y, seriestype = :scatter, title = "My Scatter Plot")

# ╔═╡ 5ffbc4bf-cc16-474e-8985-2ebb81dc821c
# plot(i, wave, seriestype = :scatter, title = "My Scatter Plot")

# ╔═╡ 3ee38394-fce7-456f-86f8-0f889b2c4f1a
# plot(i, wave1, seriestype = :scatter, title = "My Scatter Plot")

# ╔═╡ 37c49652-2500-4236-8ffb-0906ff0248e5
# plot(i, wave2, seriestype = :scatter, title = "My Scatter Plot")

# ╔═╡ 1df63790-0d76-4c41-aaf9-5cf39ec6cb98
@bind fs1 Slider(100:100:10000,default=8e3, show_value=true)

# ╔═╡ 446e0ba3-4c3a-4e15-b1a1-5a8046c02c3c
@bind f1 Slider(0:100:10000, default=1e3, show_value=true)

# ╔═╡ dd194f36-68c0-440b-86f8-9f3ec2a78647
begin
t1 = 0.0:1/fs1:prevfloat(1.0)
#y1 = squarewave.(f1*t1) * 1
#y1 = (sin.(2pi * f1* t1) .*0.1)
#y1 = trianglewave.(f1* t1) * 1
y1 = (sin.(2pi * f1* t1) .*t1) .+ (sin.(2pi * f1*2* t1 .+pi) .*t1/2)
wavplay(y1, fs1)
end

# ╔═╡ b9a462bd-67b0-426a-a912-9422c3524885
lineplot(t1, y1)

# ╔═╡ 545eb713-70f5-4e04-adf5-bbb6100a453f
@bind a1 Slider(0:0.1:2,default=0.1, show_value=true)

# ╔═╡ 8577ca66-db1e-4f9b-a775-adb4b404e4e8
@bind x Slider(0:100:10000,default=100, show_value=true)

# ╔═╡ 2d064c8f-d7f4-48ed-985d-c61172088a62
@bind L Slider(0:100:10000,default=100, show_value=true)

# ╔═╡ ee1eaf9e-1459-4184-85e6-0736053304ef
@bind fw Slider(0:100:10000,default=100, show_value=true)

# ╔═╡ ed916679-4bca-41cc-9fae-c6b8362855a0
@bind n Slider(1:1:10,default=2, show_value=true)

# ╔═╡ 236f96dd-8ed0-4332-9e78-41bd0b57c5ac
@bind vibfs Slider(0:100:10000,default=100, show_value=true)

# ╔═╡ 58f1d291-6cd3-42a2-b48a-44a4fd590dd4
# making vibrations
begin
tn = 0.0:1/vibfs:prevfloat(1.0)
vib =  2 * a1 * sin.(n*pi*x/L).*cos.(2*pi*fw*tn)
end

# ╔═╡ d8396378-e8b5-413f-87b4-d77479e130e7
lineplot(tn, vib)

# ╔═╡ d8d86fdc-8adf-43db-ae77-a04dabb2c63a
# wavplay(vib, vibfs)

# ╔═╡ fd76e076-2b61-4e55-9cd3-e59340168f4d

# ╔═╡ Cell order:
# ╠═fb56c064-b90d-11eb-124f-bd2478c9a5bf
# ╠═a6fcf1e7-e477-41d3-b33a-08f521db740c
# ╠═6813d955-d54d-47ac-8cec-5be4ea1847fe
# ╠═df0fbb6c-bad6-435d-a931-2a4cd9ea935c
# ╠═c5413acd-7828-4da0-9009-65fd61264637
# ╠═56d13e40-49bc-4f45-9033-0a4003070da8
# ╠═683fa835-7339-4b7c-89fe-6511da652e2d
# ╠═17af4d9f-687a-48c8-9a68-10d9bb713187
# ╠═263973d5-9d5c-427b-94cb-bd74073a7a29
# ╠═8c401e62-30b2-492d-a896-b0dceb2c8508
# ╠═aa1d8f3a-09cd-4fee-b6a0-e959650b0faa
# ╠═a5298c48-a277-4c1c-bb35-e129cb205df2
# ╠═e4b55da5-2cba-46b0-a084-2ebbb93a31f9
# ╠═5ffbc4bf-cc16-474e-8985-2ebb81dc821c
# ╠═3ee38394-fce7-456f-86f8-0f889b2c4f1a
# ╠═37c49652-2500-4236-8ffb-0906ff0248e5
# ╠═1df63790-0d76-4c41-aaf9-5cf39ec6cb98
# ╠═446e0ba3-4c3a-4e15-b1a1-5a8046c02c3c
# ╠═dd194f36-68c0-440b-86f8-9f3ec2a78647
# ╠═b9a462bd-67b0-426a-a912-9422c3524885
# ╠═545eb713-70f5-4e04-adf5-bbb6100a453f
# ╠═8577ca66-db1e-4f9b-a775-adb4b404e4e8
# ╠═2d064c8f-d7f4-48ed-985d-c61172088a62
# ╠═ee1eaf9e-1459-4184-85e6-0736053304ef
# ╠═ed916679-4bca-41cc-9fae-c6b8362855a0
# ╠═236f96dd-8ed0-4332-9e78-41bd0b57c5ac
# ╠═58f1d291-6cd3-42a2-b48a-44a4fd590dd4
# ╠═d8396378-e8b5-413f-87b4-d77479e130e7
# ╠═d8d86fdc-8adf-43db-ae77-a04dabb2c63a
# ╠═fd76e076-2b61-4e55-9cd3-e59340168f4d


export ilog2

"""
    ilog2(a::Integer)::Int

For nonzero integers `a` return Int(floor(log2(abs(a))))
For zero return `-1`
"""
function ilog2(a::Integer)
    ndigits(a, base=2, pad=0) - 1
end
#=
function ilog2(a::BigInt)
    Limb = Base.GMP.Limb
    sol = sizeof(Limb)
    bpl = sol * 8
    siz = max(abs(a.size), 1)
    v = unsafe_wrap(Vector{Limb}, a.d + (siz - 1) * sol, 1)
    msl = v[1]
    siz * bpl - 1 - leading_zeros(msl)
end
=#
function ilog2(a::Base.BitInteger)
    bpl = sizeof(a) * 8
    bpl - 1 - leading_zeros(abs(a))
end


mutable struct pbsrtfProblem{T<:Real}
    y       :: Vector{T}              # response
    x       :: Vector{T}              # design
    k       :: Int
    restriction :: String             # 'increasing','convex','decreasing','concave'
                                      # 'inc-convex','inc-concave','dec-convex','dec-concave'
    lb      :: Union{Vector{T},Nothing} # lower bound vector for β
    ub      :: Union{Vector{T},Nothing} # upper bound vector for β
    λ       :: Union{T,Nothing}       # Moreau-Yosida envelope parameter
    s       :: T                      # IG(s, r) prior for σ²
    r       :: T
    μ       :: Union{T,Nothing}       # exp(μ) prior for α
    m       :: Int                    # number of observations
    n       :: Int                    # number of grid locations
    xgrid   :: Vector{T}              # strictly increasing grid locations (length n)
    w       :: Vector{T}              # number of observations at each location (length n)
    ybar    :: Vector{T}              # the average of y at each location (length n)
    sse     :: T                      # sum of squared errors
    D¹      :: SparseMatrixCSC{T,Int}
    D²      :: SparseMatrixCSC{T,Int}
    D       :: SparseMatrixCSC{T,Int}
    GRB_ENV :: Any                    # gurobi environment
end

function pbsrtfProblem(y::Vector{T}, x::Vector{T}, k::Int, restriction::String;
                       lb::Union{Vector{T},Nothing}=nothing,
                       ub::Union{Vector{T},Nothing}=nothing,
                       thinning::Bool=false,nbins::Int=100,λ::Union{T,Nothing}=nothing,
                       s::T=1e-3,r::T=1e-3,μ::Union{T,Nothing}=nothing) where T<: Real
    if (k!=1) & (k!=2)
        throw(ArgumentError("Currenly only k=1 and k=2 are accepted for pbsrtf."))
    end
    m = length(x)
    xgrid = sort(unique(x))
    n = length(xgrid)
    w = zeros(T,n)
    ybar = zeros(T,n)
    sse = 0.0
    for i in 1:n
        w[i] = length(x[x.==xgrid[i]])
        ybar[i] = mean(y[x.==xgrid[i]])
        sse += sum(abs2, y[x.==xgrid[i]].-ybar[i])
    end
    if thinning
        xgrid,w,ybar,sse = thin(xgrid,w,ybar,sse,nbins)
        n = nbins
    end
    if λ == nothing
        λ = var(y)*0.25*1e-3
    end
    if μ == nothing
       μ = 3.0
    end
    D¹ = getD(0,n,xgrid)
    D² = getD(1,n,xgrid)
    D  = getD(k,n,xgrid)
    GRB_ENV = Env()
    GRBsetdblparam(GRB_ENV,"FeasibilityTol",1e-6)
    GRBsetdblparam(GRB_ENV,"OptimalityTol",1e-6)
    pbsrtfProblem{T}(y,x,k,restriction,lb,ub,λ,s,r,μ,m,n,xgrid,w,ybar,sse,D¹,D²,D,GRB_ENV)
end

dimension(problem::pbsrtfProblem) = problem.n + 2
capabilities(::Type{<:pbsrtfProblem}) = LogDensityOrder{1}()

function logdensity_and_gradient(problem::pbsrtfProblem, z)
    @unpack y,x,k,restriction,lb,ub,λ,s,r,μ,m,n,xgrid,w,ybar,sse,D¹,D²,D,GRB_ENV=problem
    β        = view(z,1:n)
    logσ²    = z[n+1]
    σ²       = exp(logσ²)
    logα     = z[n+2]
    α        = exp(logα)
    η = Variable(n)
    t = Variable()
    p_gurobi = minimize(sumsquares(β-η)+sumsquares(α-t),norm(D*η,1)<=t)
    if restriction=="increasing"
    elseif restriction=="decreasing"
        p_gurobi.constraints += D¹*η<=0
    elseif restriction=="convex"
        p_gurobi.constraints += D²*η>=0
    elseif restriction=="concave"
        p_gurobi.constraints += D²*η<=0
    elseif restriction=="inc-convex"
        p_gurobi.constraints += [D¹*η>=0,D²*η>=0]
    elseif restriction=="dec-convex"
        p_gurobi.constraints += [D¹*η<=0,D²*η>=0]
    elseif restriction=="inc-concave"
        p_gurobi.constraints += [D¹*η>=0,D²*η<=0]
    elseif restriction=="dec-concave"
        p_gurobi.constraints += [D¹*η<=0,D²*η<=0]
    end
    if !isnothing(lb)
        p_gurobi.constraints += η>=lb
    end
    if !isnothing(ub)
        p_gurobi.constraints += η<=ub
    end
    solve!(p_gurobi,Optimizer(GRB_ENV),silent_solver=true,warmstart=true,verbose=false)
    #while p_gurobi.status!=Convex.MOI.OPTIMAL
        #β = β .+ 1e-3 .*randn(n)
        #p_gurobi.objective = sumsquares(β-η)+sumsquares(α-t)
        #solve!(p_gurobi,Gurobi.Optimizer(GRB_ENV),silent_solver=true,warmstart=true)
        #println("Warning, solution failed!")
    #end
    β_proj = evaluate(η)
    α_proj = evaluate(t)
    res      = ybar .- β
    ssm      = sum(abs2.(res).*w)
    qf       = (ssm+sse+2*r)/(2*σ²)
    logl     = -qf-(m/2+s)*logσ²+logα-μ*α
    dist²    = norm(β.-β_proj)^2+(α-α_proj)^2
    logl    -= dist²/(2*λ)
    ∇        = similar(z)
    ∇[1:n]   = (w.*res)./σ².-((β.-β_proj)./λ)
    ∇[n+1]   = qf-(m/2+s)
    ∇[n+2]   = 1-α/λ*(α-α_proj)-μ*α
    logl, ∇
end


using SparseArrays, NamedArrays

function _make_incidence_matrix(applications::Vector{Application}, levelfun)
    class = reduce(vcat, classification.(applications)) |> unique
    class = levelfun.(class) |> unique
    classdict = Dict(c => i for (i, c) in enumerate(class))
    incmat = falses(length(applications), length(class))
    for (i, a) in enumerate(applications)
        for c in classification(a)
            incmat[i, classdict[levelfun(c)]] = true
        end
    end
    return incmat, class
end

function coclassification(applications::Vector{Application}; level=maingroup)
    incmat, class = _make_incidence_matrix(applications, level)
    incmat = sparse(incmat)
    coclass = incmat' * incmat
    coclass = NamedArray(Matrix(coclass))
    setnames!(coclass, string.(class), 1)
    setnames!(coclass, string.(class), 2)
    return coclass
end