using Makie
using MultipleViewGeometry, Test, Random
using MultipleViewGeometry.ModuleCostFunction
using MultipleViewGeometry.ModuleTypes
using MultipleViewGeometry.ModuleConstraints
using MultipleViewGeometry.ModuleConstruct
using MultipleViewGeometry.ModuleDraw
using MultipleViewGeometry.ModuleMove
using LinearAlgebra
using StaticArrays
using GeometryTypes

# Fix random seed.
Random.seed!(1234)
# Construct two camera matrices and parametrise two planar surfaces.
f = 50
image_width = 640 / 10
image_height = 480 / 10
𝐊₁ = @SMatrix [f 0 0 ;
               0 f 0 ;
               0 0 1 ]
𝐑₁ = SMatrix{3,3,Float64,9}(rotxyz(0, 25*(pi/180), 0))
𝐭₁ = [-30.0, 0.0, -5.0]

# Normals and distance from origin
𝐧₁ = [0.0, 0.0, 1.0]
d₁ = 55.0

𝐗₁ = [30,15,60]
𝐗₂ = [-30,15,60]
𝐗₃ = [30,-15,60]
𝐗₄ = [-30,-15,60]

scale = 20.0f0
x = Vec3f0(0); baselen = 0.2f0 * scale ; dirlen = 1f0 * scale
# create an array of differently colored boxes in the direction of the 3 axes
rectangles = [
    (HyperRectangle(Vec3f0(x), Vec3f0(dirlen, baselen, baselen)), RGBAf0(1,0,0,1)),
    (HyperRectangle(Vec3f0(x), Vec3f0(baselen, dirlen, baselen)), RGBAf0(0,1,0,1)),
    (HyperRectangle(Vec3f0(x), Vec3f0(baselen, baselen, dirlen)), RGBAf0(0,0,1,1))
]
meshes = map(GLNormalMesh, rectangles)

scene = mesh(merge(meshes))
scatter!(scene, 𝐗₁[1,:], 𝐗₁[2,:], 𝐗₁[3,:], markersize = 3, color = :red)
scatter!(scene, 𝐗₂[1,:], 𝐗₂[2,:], 𝐗₂[3,:], markersize = 3, color = :green)
scatter!(scene, 𝐗₃[1,:], 𝐗₃[2,:], 𝐗₃[3,:], markersize = 3, color = :black)
scatter!(scene, 𝐗₄[1,:], 𝐗₄[2,:], 𝐗₄[3,:], markersize = 3, color = :blue)

# Convention I: world_basis = (Vec(1.0, 0.0, 0.0), Vec(0.0, 1.0, 0.0), Vec(0.0, 0.0, 1.0))
world_basis = (Vec(1.0, 0.0, 0.0), Vec(0.0, 1.0, 0.0), Vec(0.0, 0.0, 1.0))
# Convention I: camera_basis = (Point(0.0, 0.0, 0.0), Vec(-1.0, 0.0, 0.0), Vec(0.0, -1.0, 0.0), Vec(0.0, 0.0, 1.0))
camera_basis = (Point(0.0, 0.0, 0.0), Vec(-1.0, 0.0, 0.0), Vec(0.0, -1.0, 0.0), Vec(0.0, 0.0, 1.0))
# Convention I: picture_basis = (Point(0.0, 0.0), Vec(1.0, 0.0), Vec(0.0, 1.0))
picture_basis = (Point(0.0, 0.0), Vec(-1.0, 0.0), Vec(0.0, -1.0))

camera₁ = Pinhole(image_width, image_height, f, camera_basis..., picture_basis...)

relocate!(camera₁, 𝐑₁, 𝐭₁)
draw!(camera₁, scene)

𝐑₁′, 𝐭₁′ = ascertain_pose(camera₁, world_basis... )
𝐊₁′ = obtain_intrinsics(camera₁, CartesianSystem())
#𝐊₁′ = SMatrix{3,3,Float64,9}(f, 0.0, 0.0, 0.0, f, 0.0, 0.0, 0.0 , 1)

𝐏₁ = construct(ProjectionMatrix(), 𝐊₁′, 𝐑₁′, 𝐭₁′)

# Red
ℳ₁ = project(camera₁, 𝐏₁, [SVector(𝐗₁...)])
Base.display(ℳ₁)

# Green
ℳ₂ = project(camera₁, 𝐏₁, [SVector(𝐗₂...)])
Base.display(ℳ₂)

# Black
ℳ₃ = project(camera₁, 𝐏₁, [SVector(𝐗₃...)])
Base.display(ℳ₃)

# Blue
ℳ₄ = project(camera₁, 𝐏₁, [SVector(𝐗₄...)])
Base.display(ℳ₄)

𝐗₁′  = Array([𝐑₁′' -𝐑₁′'*𝐭₁′] * vcat(𝐗₁,1))
𝐗₂′  = Array([𝐑₁′' -𝐑₁′'*𝐭₁′] * vcat(𝐗₂,1))
𝐗₃′  = Array([𝐑₁′' -𝐑₁′'*𝐭₁′] * vcat(𝐗₃,1))
𝐗₄′  = Array([𝐑₁′' -𝐑₁′'*𝐭₁′] * vcat(𝐗₄,1))

scene₂ = mesh(merge(meshes))
scatter!(scene₂, 𝐗₁′[1,:], 𝐗₁′[2,:], 𝐗₁′[3,:], markersize = 3, color = :red)
scatter!(scene₂, 𝐗₂′[1,:], 𝐗₂′[2,:], 𝐗₂′[3,:], markersize = 3, color = :green)
scatter!(scene₂, 𝐗₃′[1,:], 𝐗₃′[2,:], 𝐗₃′[3,:], markersize = 3, color = :black)
scatter!(scene₂, 𝐗₄′[1,:], 𝐗₄′[2,:], 𝐗₄′[3,:], markersize = 3, color = :blue)

relocate!(camera₁, 𝐑₁′, -𝐭₁′)
draw!(camera₁, scene₂)

# Red
hom⁻¹(𝐊₁′*𝐗₁′)
Base.display(hom⁻¹(𝐊₁′*𝐗₁′))

# Green
hom⁻¹(𝐊₁′*𝐗₂′)
Base.display(hom⁻¹(𝐊₁′*𝐗₂′))

# Black
hom⁻¹(𝐊₁′*𝐗₃′)
Base.display(hom⁻¹(𝐊₁′*𝐗₃′))

# Blue
hom⁻¹(𝐊₁′*𝐗₄′)
Base.display(hom⁻¹(𝐊₁′*𝐗₄′))

# Visualize the set of corresponding points
scene = Scene()
M₁ = reshape(reinterpret(Float64,ℳ₁),(2,length(ℳ₁)))
M₂ = reshape(reinterpret(Float64,ℳ₂),(2,length(ℳ₂)))
M₃ = reshape(reinterpret(Float64,ℳ₃),(2,length(ℳ₃)))
M₄ = reshape(reinterpret(Float64,ℳ₄),(2,length(ℳ₄)))

scatter!(scene,M₁[1,:], M₁[2,:], markersize = 1, color = :red, limits = FRect(0, 0, 64, 48.0))
scatter!(scene,M₂[1,:], M₂[2,:], markersize = 1, color = :green, limits = FRect(0, 0, 64, 48.0))
scatter!(scene,M₃[1,:], M₃[2,:], markersize = 1, color = :black, limits = FRect(0, 0, 64, 48.0))
scatter!(scene,M₄[1,:], M₄[2,:], markersize = 1, color = :blue, limits = FRect(0, 0, 64, 48.0))

# scatter!(scene, 𝐗₁′[1,:], 𝐗₁′[2,:], 𝐗₁′[3,:], markersize = 3, color = :red)
# scatter!(scene, 𝐗₂′[1,:], 𝐗₂′[2,:], 𝐗₂′[3,:], markersize = 3, color = :green)
# scatter!(scene, 𝐗₃′[1,:], 𝐗₃′[2,:], 𝐗₃′[3,:], markersize = 3, color = :black)
# scatter!(scene, 𝐗₄′[1,:], 𝐗₄′[2,:], 𝐗₄′[3,:], markersize = 3, color = :blue)


########################################################################
# Observed info -  for 2PL pars and 1-D latent                         #
########################################################################
"""
```julia
_item_observed_information(
    parameters::Parameters1PL,
    latent::Latent1D,
    response_val::Float64,
)
```

# Description

It is equal to the item expected information.

# Arguments

- **`parameters::Parameters1PL`** : Required. A 1-parameter logistic parameters object.
- **`latent::Latent1D`** : Required. A 1-dimensional `Latent1D` latent variable.
- **`response_val::Float64`** : Required. A scalar response.

# Output
A `Float64` scalar observed (expected) informations.
"""
function _item_observed_information(
    parameters::Parameters1PL,
    latent::Latent1D,
    response_val::Float64
)
    return _item_expected_information(parameters, latent)::Float64
end

########################################################################
# Observed info -  for 2PL pars and 1-D latent                         #
########################################################################
"""
```julia
_item_observed_information(
    parameters::Parameters2PL,
    latent::Latent1D,
    response_val::Float64,
)
```

# Description

Theoretically, it should be equal to its expected counterpart.
However, it is used to compute the standard errors of item parameter estimates which are based on the log-likelihood (which is based on the responses).
Look at the example in the documentation of the exported method.

# Arguments

- **`parameters::Parameters2PL`** : Required. A 2-parameter logistic parameters object.
- **`latent::Latent1D`** : Required. A 1-dimensional `Latent1D` latent variable.
- **`response_val::Float64`** : Required. A scalar response.

# Output
A ``2 \times 2`` matrix of the observed informations.
"""
function _item_observed_information(
    parameters::Parameters2PL,
    latent::Latent1D,
    response_val::Float64
)
    #return _item_expected_information(parameters, latent)::Matrix{Float64}
    p = _probability(latent, parameters)
    p_1_p = p * (1 - p)
    i_aa = p_1_p * (latent.val - parameters.b)^2 # v p.49 Kim, Baker: - L_11
    i_ab =
        (
             ((parameters.a * (latent.val - parameters.b) * (- p) * p_1_p) +
             (response_val - p))
        )  # X p.50 Kim, Baker: - L_12 #TODO
    i_bb = - parameters.a^2 * p_1_p # v p.49 Kim, Baker: - L_22
    return [i_aa i_ab; i_ab i_bb]::Matrix{Float64}
end

########################################################################
# Observed info -  for 3PL pars and 1-D latent                         #
########################################################################
"""
```julia
_item_observed_information(
    parameters::Parameters3PL,
    latent::Latent1D,
    response_val::Float64,
)
```

# Description

It computes the observed information (-second derivative of the likelihood) with respect to the 3 parameters of the 3PL model.
It follows the parametrization ``a(θ - b)``.

# Arguments

- **`parameters::Parameters3PL`** : Required. A 3-parameter logistic parameters object.
- **`latent::Latent1D`** : Required. A 1-dimensional `Latent1D` latent variable.
- **`response_val::Float64`** : Required. A scalar response.

# Output
A ``3 \times 3`` matrix of the observed informations.
"""
function _item_observed_information(
    parameters::Parameters3PL,
    latent::Latent1D,
    response_val::Float64,
)
    p = _probability(latent, parameters)
    i = (1 - p) * (p - parameters.c)
    h = (response_val * parameters.c - p^2) * i
    j = response_val * i
    den = ((1 - parameters.c) * p)^2
    i_aa = - h * (latent.val - parameters.b)^2 / den # v p.49 Kim, Baker: - L_11
    i_ab =
        (
            (p - parameters.c) * ((parameters.a * (latent.val - parameters.b) * h) +
            (p * (response_val - p) * (1 - parameters.c)))
        ) / den # X p.50 Kim, Baker: - L_12 #TODO
    i_ac = j * (latent.val - parameters.b) / den # v p.51 Kim, Baker: - L_13
    i_bc = - parameters.a * j / den # v p.51 Kim, Baker: - L_23
    i_bb = - parameters.a^2 * h / den # v p.49 Kim, Baker: - L_22
    i_cc = (response_val - 2 * response_val * p + p^2) / den # X p.50 Kim, Baker: - L_33
    return [i_aa i_ab i_ac; i_ab i_bb i_bc; i_ac i_bc i_cc]::Matrix{Float64}
end


using SymmetricSparseMatrices
using SparseArrays
using Test

@testset "SymmetricSparseMatrix" begin
    # SymmetricSparseMatrix
    @test size(SymmetricSparseMatrix()) == (0,0)
    @test nnz(SymmetricSparseMatrix()) == 0
    @test getindex(SymmetricSparseMatrix(1,[1,2],[1],[1]),1,1) == 1
    @test nnz(SymmetricSparseMatrix(1,[1,2],[1],[1])) == 1
    @test SymmetricSparseMatrix(5,[1,4,6,9,12,14], [1,2,4,1,2,3,4,5,1,3,4,2,5], [1,-1,-3,-2,5,4,6,4,-4,2,7,8,-5])*ones(5) == [-3, 4, 14, 10, -1]

    # Conversions
    a = sparse([1,2,3,4],[1,2,3,4],[1,2,3,4])
    b = Matrix(a)
    c = symsparse(b)
    @test a == b == c

    a = sparse([1,2,3,4,1,2,3,4],[1,2,3,4,4,3,2,1],[1,2,3,4,4,3,2,1])
    b = Matrix(a)
    c = symsparse(b)
    @test a != c != b
end


using OrdinaryDiffEq
import OrdinaryDiffEq: solve
using GeneralizedGenerated

f1(du, u, p, t) = begin
    du[1] = p[1]*1/u[1]
end
u0 = [1.0]
p = [1.0]
t = (0.0,1.0)
du = similar(u0)
f1(du, u0, p, 0)
@show du
prob = ODEProblem(f1, u0, t, p)
solve(prob, Tsit5())
println("The old fashioned way works")

fex = :((du, u, p, t) -> begin
        du[1] = p[1]*1/u[1]
        end)

f2 = eval(fex)
prob = ODEProblem(f2, u0, t, p)
soln = solve(prob, Tsit5())
println("Works fine with eval")

f3   = mk_function(fex)

f3(du, u0, p, 0)
@show du
println("We can call f3")
f4 = (du, u, p, t) -> f3(du,  u, p, t)
prob = ODEProblem(f4, u0, t, p)
soln = solve(prob, Tsit5())


#=
Test the wrapper function.

Reference
---------
1. https://github.com/stevengogogo/FindSteadyStates.jl/issues/29
=#

include("model.jl")

function fixedpoint_gen(func, u0, p, ranges)

    # Define a problem
    de = DEsteady(func=func, p=p, u0= u0, method=SSRootfind())

    j_gen = jacobian(de) # jacobian generator

    # Searching method and domain
    param_gen = ParameterGrid(ranges)

    # Solve
    sols = solve(de, param_gen)

    # Remove similar solutions
    steadies = unique(sols)

    # Jacobian
    jac_ms = j_gen.(steadies)

    # Stability
    stab_modes = StabilityType.(jac_ms)

    return steadies, jac_ms, jac_ms

end

func_ = bistable_model.func
p_ = bistable_model.p
u0_ = bistable_model.u0

fixedpoint_gen(func_, u0_, p_, [1:0.1:2, 1:0.1:2])


using MOCNeutronTransport
using BenchmarkTools
using Plots
pyplot()

import MOCNeutronTransport.area

struct UM_Tuple{P, F, T}
    points::NTuple{P, Point_2D{T}}
    faces::NTuple{F, Tuple{Vararg{Int64}}}
end
struct UM_Vector{P, T}
    points::NTuple{P, Point_2D{T}}
    faces::Vector{Vector{Int64}}
end

function area(mesh::UM_Vector)
    mesh_area = 0.0
    for face in mesh.faces
        if face[1] == 9
            mesh_area += area(
                        Quadrilateral_2D(
                                map(x -> mesh.points[x], Tuple(face[2:5]))
                            )
                        )
        end
    end
    return mesh_area
end

function area(mesh::UM_Tuple)
    mesh_area = 0.0
    for face in mesh.faces
        if face[1] == 9
            mesh_area += area(
                        Quadrilateral_2D(
                                map(x -> mesh.points[x], face[2:5])
                            )
                        )
        end
    end
    return mesh_area
end

tuple_times = Float64[]
vector_times = Float64[]
Nset = Int64.([1E1, 1E2, 1E3, 1E4, 1E5, 1E6])
Nctr = 1
for N = Nset
    # Generate the points
    points = Tuple(reduce(vcat, [[Point_2D(1.0*i, 0.0), Point_2D(1.0*i, 1.0)] for i = 0:N]))
    # Generate the faces for quadrilaterals
    quad_faces_vec = [[9, i, i+2, i+3, i+1] for i = 1:2:2N]
    quad_faces_tuple = Tuple([ Tuple(v) for v in quad_faces_vec ])

    faces_tuple = quad_faces_tuple
    faces_vec = quad_faces_vec

    tuple_mesh = UM_Tuple(points, faces_tuple)
    vector_mesh = UM_Vector(points, faces_vec)

    # Test speed of converting to quadrilaterals and getting the area
    println("N = 1E$Nctr")
    time = @belapsed area($vector_mesh)
    push!(vector_times, time/N)
    println("  Vector: $time")
    time = @belapsed area($tuple_mesh)
    push!(tuple_times, time/N)
    println("  Tuple: $time")
    global Nctr += 1
end
println("tuple_times ./ vector_times = ", tuple_times ./ vector_times)
plot(Nset,
     vector_times,
     label = "Vector",
     title = "Number of mesh cells vs area time per cell",
     xaxis=:log,
     yaxis=:log
    )
plot!(Nset, tuple_times, label = "Tuple")
xlabel!("Mesh cells")
ylabel!("Area computation time per cell")


# Attempt to find all zeros in an interval (a,b)

# Algorithm due to @djsegal in https://github.com/JuliaMath/Roots.jl/pull/113

# A naive approach to find zeros: split (a,b) by n points, look into each for a zero
# * k is oversampling rate for bisection. (It is relatively cheap to check for a bracket so we
#   oversample our intervals looking for brackets
# * assumes f(a) *not* a zero
function _fz(f, a, b, no_pts, k=4)
    zs = Real[]
    _fz!(zs, f, a, b, no_pts, k)
    zs
end

function _fz!(zs, f, a::T, b, no_pts, k=4) where {T}

    pts = range(a, stop=b, length=(no_pts-1)*k+1)
    n::Int = length(pts)

    fs = f.(pts)
    sfs = sign.(fs)

    u::T = first(pts)  # keep track of bigger interval
    found_bisection_zero = false

    for (i,x) in enumerate(pts[1:end])
        q,r = divrem(i-1, k)

        if i > 1 && iszero(r)
            v::T = x
            if !found_bisection_zero
                try
                    p1::T = identify_starting_point(u, v, sfs[(i-k):i])
                    rt::T = dfree(f, p1)
                    if !isnan(rt) && u < rt <= v
                        push!(zs, rt)
                    end
                catch err
                end
            end
            u = v
            found_bisection_zero = false
        end

        if i < n
            if iszero(fs[i+1])
                found_bisection_zero = true # kinda
                push!(zs, pts[i+1])
            elseif  sfs[i] * sfs[i+1] < 0
                found_bisection_zero = true
                rt = bisection(f, x, pts[i+1])
                push!(zs, rt)
            end
        end
    end

    sort!(zs)
end

# the algorithm first scans for zeros using the naive approach, then
# splits (a,b) by these zeros. This struct holds the subintervals
struct Interval{T}
a::T
b::T
depth::Int
end

Base.show(io::IO, alpha::Interval) = print(io, "($(alpha.a), $(alpha.b))")

# check if f(a) is non zero using tolerances max(atol, eps()), rtol
function _non_zero(fa, a::T, atol, rtol) where {T}
    abs(fa) >= max(atol, abs(a)*rtol*oneunit(fa)/oneunit(a), oneunit(fa)*eps(T))
end

# After splitting by zeros we have intervals (zm, zn) this is used to shrink
# to (zm+, zn-) where both are non-zeros, as defined above
function find_non_zero(f, a::T, barrier, xatol, xrtol, atol, rtol) where {T}
    nan = (0*a)/(0*a) # try to get typed NaN
    xtol = max(xatol, abs(a) * xrtol, oneunit(a) * eps(T))
    sgn = barrier > a ? 1 : -1
    ctr = 0
    x = a + 2^ctr*sgn*xtol
    while !_non_zero(f(x), x, atol, rtol)
        ctr += 1
        x += 2^ctr*sgn*xtol
        ((sgn > 0 && x > barrier) || (sgn < 0 && x < barrier)) && return nan
        ctr > 100 && return nan
    end

    x
end

# split a < z1 < z2 < ... < zn < b into intervals (a+,z1-), (z1+, z2-), ...
# where possible; push! onto ints
function make_intervals!(ints, f, a, b, zs, depth, xatol, xrtol, atol, rtol)
    pts = vcat(a, zs, b)

    for (u,v) in zip(pts[1:end-1], pts[2:end])
        ur = find_non_zero(f, u, v, xatol, xrtol, atol, rtol)
        isnan(ur) && continue

        vl = find_non_zero(f, v, u, xatol, xrtol, atol, rtol)
        isnan(vl) && continue

        push!(ints, Interval(ur, vl, depth))
    end
end

# adjust what we mean by x1 ~ x2 for purposes of adding a new zero
function approx_close(z1, z2, xatol, xrtol)
    tol = max(sqrt(xatol), max(abs(z1), abs(z2)) * sqrt(xrtol))
    abs(z1 - z2) < tol
end

# is proposed not near xs? (false and we add proposed)
function not_near(proposed, xs, xatol, xrtol)
    n = length(xs)
    n <= 1 && return true
    ind = n + 1
    for (i,rt) in enumerate(xs)
        if proposed < rt
            ind = i
            break
        end
    end
    if ind > 1 # xs[ind-1] <= propose < xs[ind]
        rt = xs[ind-1]
        approx_close(proposed, rt, xatol, xrtol) && return false
    end
    if ind <= n # value to right
        rt = xs[ind]
        approx_close(proposed, rt, xatol, xrtol) && return false
    end
    return true
end

"""
    find_zeros(f, a, b; [no_pts=12, k=8, naive=false, xatol, xrtol, atol, rtol])

Search for zeros of `f` in the interval `[a,b]`.

# Examples

```jldoctest find_zeros
julia> using Roots

julia> find_zeros(x -> exp(x) - x^4, -5, 20)        # a few well-spaced zeros
3-element Vector{Float64}:
 -0.8155534188089606
  1.4296118247255556
  8.613169456441398

julia> find_zeros(x -> sin(x^2) + cos(x)^2, 0, 2pi)  # many zeros
12-element Vector{Float64}:
 1.78518032659534
 2.391345462376604
 3.2852368649448853
 3.3625557095737544
 4.016412952618305
 4.325091924521049
 4.68952781386834
 5.00494459113514
 5.35145266881871
 5.552319796014526
 5.974560835055425
 6.039177477770888

julia> find_zeros(x -> cos(x) + cos(2x), 0, 4pi)    # mix of simple, non-simple zeros
6-element Vector{Float64}:
  1.0471975511965976
  3.141592653589943
  5.235987755982988
  7.330382858376184
  9.424777960769228
 11.519173063162574

julia> f(x) = (x-0.5) * (x-0.5001) * (x-1)          # nearby zeros
f (generic function with 1 method)

julia> find_zeros(f, 0, 2)
3-element Vector{Float64}:
 0.5
 0.5001
 1.0

julia> f(x) = (x-0.5) * (x-0.5001) * (x-4) * (x-4.001) * (x-4.2)
f (generic function with 1 method)

julia> find_zeros(f, 0, 10)
3-element Vector{Float64}:
 0.5
 0.5001
 4.2

julia> f(x) = (x-0.5)^2 * (x-0.5001)^3 * (x-4) * (x-4.001) * (x-4.2)^2  # hard to identify
f (generic function with 1 method)

julia> find_zeros(f, 0, 10, no_pts=21)                # too hard for default
5-element Vector{Float64}:
 0.49999999999999994
 0.5001
 4.0
 4.001
 4.200000000000001
```

!!! note
    There are two cases where the number of zeros may be underreported:
    * if the initial interval, [a,b], is too wide
    * if there are zeros  that are very nearby

----

The basic algorithm checks for zeros among the endpoints, and then
divides the interval `(a,b)` into `no_pts-1` subintervals and then
proceeds to look for zeros through bisection or a derivative-free
method.  As checking for a bracketing interval is relatively cheap and
bisection is guaranteed to converge, each interval has `k` pairs of
intermediate points checked for a bracket.

If any zeros are found, the algorithm uses these to partition `(a,b)`
into subintervals. Each subinterval is shrunk so that the endpoints
are not zeros and the process is repeated on the subinterval. If the
initial interval is too large, then the naive scanning for zeros may
be fruitless and no zeros will be reported. If there are nearby zeros,
the shrinking of the interval may jump over them, though as seen in
the examples, nearby roots can be identified correctly, though for
really nearby points, or very flat functions, it may help to increase
`no_pts`.

The tolerances are used to shrink the intervals, but not to find zeros
within a search. For searches, bisection is guaranteed to converge
with no specified tolerance. For the derivative free search, a
modification of the `Order0` method is used, which at worst case
compares `|f(x)| <= 8*eps(x)` to identify a zero. The algorithm might
identify more than one value for a zero, due to floating point
approximations. If a potential pair of zeros satisfy
`isapprox(a,b,atol=sqrt(xatol), rtol=sqrt(xrtol))` then they are
consolidated.

The algorithm can make many function calls. When zeros are found in an
interval, the naive search is carried out on each subinterval. To cut
down on function calls, though with some increased chance of missing
some zeros, the adaptive nature can be skipped with the argument
`naive=true` or the number of points stepped down.

The algorithm is derived from one in a
[PR](https://github.com/JuliaMath/Roots.jl/pull/113) by @djsegal.

!!! Note
The `IntervalRootFinding` package provides a rigorous alternative to this heuristic one.
That package use interval arithmetic, so can compute bounds on the size of the image of
an interval under a function `f`. If this image includes `0`, then it can look for the zero.
Bisection, on the other hand, only will look for a zero if the two endpoints have different signs,
a much more rigid condition for a potential zero.

For example, this function (due to `@truculentmath`) is particularly tricky, as it is positive at every floating point number, but has two zeros (the vertical asymptote at `15//11` is negative between adjacent floating point values):

```jldoctest interval_root_finding
julia> using IntervalArithmetic, IntervalRootFinding, Roots

julia> g(x) = x^2 + 1 +log(abs( 11*x-15 ))/99
g (generic function with 1 method)

julia> find_zeros(g, -3, 3)
Float64[]

julia> IntervalRootFinding.roots(g, -3..3, IntervalRootFinding.Bisection)
1-element Vector{Root{Interval{Float64}}}:
 Root([1.36363, 1.36364], :unknown)
```

A less extreme usage might be the following, where `unique` indicates Bisection could be useful and indeed `find_zeros` will identify these values:

```jldoctest interval_root_finding
julia> g(x) = exp(x) - x^5
g (generic function with 1 method)

julia> rts = IntervalRootFinding.roots(g, -20..20)
2-element Vector{Root{Interval{Float64}}}:
 Root([12.7132, 12.7133], :unique)
 Root([1.29585, 1.29586], :unique)

julia> find_zeros(g, -20, 20)
2-element Vector{Float64}:
  1.2958555090953687
 12.713206788867632
```

"""
function find_zeros(f, a, b; no_pts = 12, k=8,
                    naive = false,
                    kwargs...
                    )

    a0, b0 = promote(float(a), float(b))
    a0 = isinf(a0) ? nextfloat(a0) : a0
    b0 = isinf(b0) ? prevfloat(b0) : b0

    # set tolerances if not specified
    fa0, fb0 = f(a0), f(b0)
    d = Dict(kwargs)
    T, S = eltype(a0), eltype(fa0)
    xatol::T = get(d, :xatol, eps(one(T))^(4//5) * oneunit(T))
    xrtol = get(d, :xrtol, eps(one(T)) * one(T))
    atol::S  = get(d, :atol,  eps(float(S)) * oneunit(S))
    rtol  = get(d, :rtol,  eps(float(S)) * one(S))

    zs = T[]  # collect zeros

    # check endpoints for exact zeros, then narrow
    abs(fa0) * oneunit(T) / oneunit(S) <= 8 * eps(a0) && push!(zs, a0)
    abs(fb0) * oneunit(T) / oneunit(S) <= 8 * eps(b0) && push!(zs, b0)
    a0 = find_non_zero(f, a0, b0, xatol, xrtol, atol, rtol)
    b0 = find_non_zero(f, b0, a0, xatol, xrtol, atol, rtol)

    _fz!(zs, f, a0, b0, no_pts,k)  # initial zeros

    ints = Interval{T}[] # collect subintervals
    !naive && !isempty(zs) &&  make_intervals!(ints, f, a0, b0, zs, 1, xatol, xrtol, atol, rtol)

    nzs = T[]
    cnt = 0

    while !naive && !isempty(ints)
        cnt += 1
        i = pop!(ints)
        # this is fussy. Ideally, we would want to explore intervals
        # with fewer points (each interval is already probed for k=4
        # bisections) but how many fewer?  We already had ~ (i.b-i.a) *
        # no_pts / (b-a) points. Would want an increased density so n
        # > (i.b - i.a) / (b - a) * no_pts but this doesn't perform as
        # well as we might expect

        # sub_no_pts = ceil(Int, (i.b - i.a) / (b-a) * no_pts * 2^(i.depth))
        #sub_no_pts <= 2 && continue  # stop on depth, always divide if roots
        #sub_no_pts = max(3, floor(Int, no_pts  / (2.0)^(i.depth)))
        sub_no_pts = floor(Int, no_pts  / (2.0)^(i.depth))

        empty!(nzs)
        if sub_no_pts >= 2
            _fz!(nzs, f, i.a, i.b, sub_no_pts, k)
        end

        if !isempty(nzs)
            azs = filter(rt->not_near(rt, zs, xatol, xrtol), nzs) # trim out nearby roots
            length(azs) == 0 && continue
            append!(zs, azs)
            sort!(zs)
            i.depth > 4 && continue
            make_intervals!(ints, f, i.a, i.b, azs, i.depth+1, xatol, xrtol, atol, rtol)
        end
    end

    length(zs) <= 1 && return zs
    sort!(zs)

    # may identify same zero with nearby values along the way
    # this trims out with a relaxed tolerance on how close
    # nearby roots can be. Default is epsilon^(2/5) ≈ 5e-7
    inds = Int[1]
    z1 = first(zs)
    for i in 2:length(zs)
        z2 = zs[i]
        if !approx_close(z1, z2, xatol, xrtol)
            push!(inds, i)
        end
        z1=z2
    end

    return zs[inds]

end


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

"""
    RegularGridDiscretizer(dims)

Discretize spatial region into regular grid of given
dimensions `dims`.
"""
struct RegularGridDiscretizer{N} <: AbstractDiscretizer
  dims::Dims{N}
end

discretize(r::Rectangle, d::RegularGridDiscretizer) =
  RegularGrid(extrema(r)..., dims=d.dims)


@fenicsclass Geometry

#functions necessary for creating meshes from geometrical objects.
#2D objects below
Circle(centre,radius) = Geometry(mshr.Circle(centre,radius))
Rectangle(corner1,corner2)=Geometry(mshr.Rectangle(corner1,corner2))
Ellipse(centre,horizontal_semi_axis,vertical_semi_axis,fragments)=Geometry(mshr.Ellipse(centre,horizontal_semi_axis,vertical_semi_axis,fragments))

#3d objects below
Box(corner1,corner2) = Geometry(mshr.Box(corner1,corner2))
Cone(top,bottom,bottom_radius,slices::Int)=Geometry(mshr.Cone(top,bottom,bottom_radius,slices))
Sphere(centre,radius) = Geometry(mshr.Sphere(centre,radius))

#we generate the mesh of some geometrical objects
generate_mesh(geom_object::Geometry,size::Int)=Mesh(mshr.generate_mesh(geom_object.pyobject,size))
export Circle,Rectangle,Ellipse,Box,Cone,Sphere,generate_mesh

#operator overloading for Geometry types so we can create composite shapes
+(geom_object1::Geometry, geom_object2::Geometry) = Geometry(geom_object1.pyobject.__add__(geom_object2.pyobject))
-(geom_object1::Geometry, geom_object2::Geometry) = Geometry(geom_object1.pyobject.__sub__(geom_object2.pyobject))
*(geom_object1::Geometry, geom_object2::Geometry) = Geometry(geom_object1.pyobject.__mul__(geom_object2.pyobject))

set_subdomain(object::Geometry,degree,domain2) = fenicspycall(object, :set_subdomain,degree,domain2.pyobject)
"""
Extrude2D

Extrudes a 2D geometry to 3D
"""
Extrude2D(object::Geometry,thickness) = Geometry(mshr.Extrude2D(object.pyobject,thickness))
export  set_subdomain, Extrude2D

CSGUnion(geom_object1::Geometry, geom_object2::Geometry) = Geometry(mshr.CSGUnion(geom_object1.pyobject,geom_object2.pyobject))
CSGIntersection(geom_object1::Geometry, geom_object2::Geometry) = Geometry(mshr.CSGIntersection(geom_object1.pyobject,geom_object2.pyobject))
CSGDifference(geom_object1::Geometry, geom_object2::Geometry) = Geometry(mshr.CSGDifference(geom_object1.pyobject,geom_object2.pyobject))

CSGScaling(geom_object1::Geometry, scale) = Geometry(mshr.CSGScaling(geom_object1.pyobject,scale))

CSGRotation(geom_object1::Geometry, angle) = Geometry(mshr.CSGRotation(geom_object1.pyobject,angle))
CSGRotation(geom_object1::Geometry, point, angle) = Geometry(mshr.CSGRotation(geom_object1.pyobject,point, angle))
CSGRotation(geom_object1::Geometry, point_axis, point_center, angle) = Geometry(mshr.CSGRotation(geom_object1.pyobject, point_axis, point_center, angle))

CSGTranslation(geom_object1::Geometry, translation_point) = Geometry(mshr.CSGTranslation(geom_object1.pyobject, translation_point))

export CSGUnion, CSGIntersection, CSGDifference, CSGScaling, CSGRotation, CSGTranslation


# Load Julia packages (libraries) needed.

using StatisticalRethinking
using StructuralCausalModels

ProjDir = @__DIR__
cd(ProjDir) #do

# ### snippet 5.1

df1 = DataFrame!(CSV.File(rel_path("..", "data", "WaffleDivorce.csv"), delim=';'));

df = DataFrame()
df[!, :A] = df1[:, :MedianAgeMarriage]
df[!, :M] = df1[:, :Marriage]
df[!, :D] = df1[:, :Divorce]

d = OrderedDict(
  :M => [:A],
  :D => [:A]
);

dag = DAG("waffles", d, df=df);
show(dag)

fname = ProjDir * "/AMD_3.dot"
to_graphviz(dag, fname)
Sys.isapple() && run(`open -a GraphViz.app $(fname)`)

display(dag.s); println()

shipley_test_dag_3 = shipley_test(dag)
if !isnothing(shipley_test_dag_3)
  display(shipley_test_dag_3)
end
println()

f = [:M]; s = [:D]; sel = vcat(f, s)
c = [:A]

e = d_separation(dag, f, s)
println("d_separation($(dag.name), $f, $s) = $e")

e = d_separation(dag, f, s, c=c)
println("d_separation($(dag.name), $f, $s, c=$c) = $e\n")

bs = basis_set(dag)
display(bs)

adjustmentsets = adjustment_sets(dag, f[1], s[1])
println("Adjustment sets:")
adjustmentsets |> display

#end


using Base.Test
using ReverseDiffSource

@printf("*********************\njulia version = %s, dir %s\n", Base.VERSION_STRING, pwd())

typealias Vec Vector{Float64}
typealias Mat Matrix{Float64}
typealias SymMat Symmetric{Float64,Matrix{Float64}}

AAt(A::Mat) = Symmetric(A*A')
AtA(A::Mat) = Symmetric(A'*A)
sumsq(x::Vec) = norm(x)^2

# Make matrix from diagonal and strict lower triangle,
# e.g. D = [d11 d22 d33 d44]
#      LT = [L21 L31 L32 L41 L42 L43]
# Outputting
#  [d11   0   0   0]
#  [L21 d12   0   0] # row r: Ls starting at sum_i=1^r
#  [L31 L32 d33   0]
#  [L41 L42 L43 d44]
function ltri_unpack(D::Vec, LT::Vec)
  d=length(D)
  make_row(r::Int, L) = hcat(reshape([ L[i] for i=1:r-1 ],1,r-1), D[r], zeros(1,d-r))
  row_start(r::Int) = div((r-1)*(r-2),2)
  inds(r) = row_start(r)+(1:r-1)
  vcat([ make_row(r, LT[inds(r)]) for r=1:d ]...)
end

ltri_unpack(D, L) = ltri_unpack([Float64(d) for d in D], [Float64(l) for l in L])

@test [11 0 0 0; 21 22 0 0 ; 31 32 33 0 ; 41 42 43 44] == ltri_unpack([11 22 33 44], [21 31 32 41 42 43])

LL = ltri_unpack([1 2 3 4], [21 31 32 41 42 43])
@printf("An example lower triangle made from diag and LT=\n%s\n", LL)

function ltri_pack{M<:AbstractMatrix}(L::M)
  d=size(L,1)

  make_row(r::Int, L) = hcat(reshape([ L[i] for i=1:r-1 ],1,r-1), D[r], zeros(1,d-r))
  row_start(r::Int) = (r-1)*(r-2)/2
  diag(L), hcat([L[r,1:r-1] for r=1:d ]...)
end

ltri_pack{T}(L::LowerTriangular{T, Matrix{T}}) = ltri_pack(full(L)) ## Until  packed storage is implemented for those

@test ltri_unpack(ltri_pack(LL)...) == LL

@printf("pak=%s\n", ltri_pack(LL))

# Conventional GMM.
# This doesn't even use inverse covariance, because as soon as you
# start down that route, you may as well go for lpGMM below.
type GMM
  n::Int           # number of Gaussians
  d::Int           # dimension of Gaussian
  alphas::Vec      # weights: n, require sum(alphas)==1
  mus::Array{Vec}  # means: n, each dx1
  sigmas::Array{SymMat}  # covariances: n, each dxd symmetric positive definite
end

function log_likelihood(g::GMM, x::Vec)
  total = 0
  for k=1:g.n
    mean = g.mus[k]
    weight =  g.alphas[k]
    Σ = g.sigmas[k]
    mahalanobis = dot(mean - x, inv(Σ) * (mean - x))
    @printf("m=%s, ", det(inv(full(Σ))))
    total += weight / sqrt(det(2pi*full(Σ))) * exp(-0.5*mahalanobis)
  end
  @printf("w=%f\n", total)
  log(total)
end

n=3
d=2
alphas=rand(n); alphas /= sum(alphas);
mus=[randn(d) for k=1:n]
sigmas=[AAt(randn(d,d)) for k=1:n]
test_gmm = GMM(n,d,alphas,mus,sigmas)
@printf("An example gmm = %s\n", test_gmm)

x = randn(d) # Test point

ll0 = log_likelihood(test_gmm, x)
@printf("ll0=%f\n", ll0)

##########################################################################
# Log-parametrized GMM.
# Weights are strictly positive, covariances are parameterized by their inverse
# square roots (lower triangular).
type lpGMM
  n::Int           # number of Gaussians
  d::Int           # dimension of Gaussian
  alphas::Vec      # log weights: n
  mus::Array{Vec}  # means: n, each dx1
  LDs::Array{Vec}  # square-root-inverse-covariances, log(diagonal): n, each d x 1
  LTs::Array{Vec}  # square-root-inverse-covariances, lower triangle: n, each d*(d-1)/2 x 1
end

# Convert simple GMM to lpGMM
function lpGMM(g::GMM)
  LTs = Array{Vec}(g.n)
  LDs = Array{Vec}(g.n)
  for k=1:g.n
    L = inv(chol(g.sigmas[k].data, Val{:L}))
    D, T = ltri_pack(L)
    LDs[k], LTs[k] = vec(log(D)), vec(T)
  end
  lpGMM(g.n,g.d,log(g.alphas),g.mus,LDs,LTs)
end

# Convert log-parameterized-GMM to simple GMM UnivariateGMM
function GMM(l::lpGMM)
  alphas::Vec = exp(l.alphas)/sum(exp(l.alphas))
  mus::Array{Vec} = l.mus
  Ls = [ltri_unpack(exp(l.LDs[i]), l.LTs[i]) for i=1:l.n]
  sigmas::Array{SymMat} = map(A->inv(Symmetric(A'*A)), Ls)
  GMM(l.n,l.d,alphas,mus,sigmas)
end

g = lpGMM(test_gmm)

#@printf("gmm=%s\n**\n", GMM(g))

########################################

const halflog2π = log(2π)/2

function log_likelihood(g::lpGMM, x::Vec)
  total = 0
  weight_normalizer = sum(exp(g.alphas))
  for k=1:g.n
    InvLowerTriangle = ltri_unpack(exp(g.LDs[k]), g.LTs[k])
    mean = g.mus[k]
    weight =  exp(g.alphas[k])/weight_normalizer # Weights parametrized as logs
    mahalanobis = sumsq(InvLowerTriangle * (mean - x))
    @printf("m=%s, ", det(InvLowerTriangle'*InvLowerTriangle))
    total += weight * det(InvLowerTriangle) * exp(-0.5*mahalanobis)
  end
  @printf("tot=%f, ", total)
  log(total) - halflog2π*g.d
end

ll1 = log_likelihood(g, x)
@printf("ll0=%f, ll1=%f, rat=%f\n", ll0, ll1, ll0/ll1)
@test_approx_eq_eps ll0 ll1 1e-12

########################################
## A better implementation, with some factorizing and logsumexp:

# logsumexp()
function logsumexp_both(x::Array{Float64,1})
  A = maximum(x);
  ema = exp(x-A);
  sema = sum(ema);
  l = log(sema) + A;
  Jacobian = ema/sema;
  return (l, Jacobian);
end
logsumexp(x::Array{Float64,1}) = logsumexp_both(x)[1]
r = rand(5);
@test_approx_eq_eps logsumexp(r) log(sum(exp(r))) 1.0e-8

@deriv_rule logsumexp(x::AbstractArray)  x  logsumexp_both(x)[2].*ds

function logsumexp_1st(x::Array{Float64,1})
  A = maximum(x);
  ema = exp(x-A);
  sema = sum(ema);
  log(sema) + A;
end
rdiff(logsumexp_1st, (zeros(3),)).code

mahal(L::Mat, mu::Vec, x::Vec) = sumsq(L*(x-mu))

function log_likelihood_2(g::lpGMM, x::Vec)

  d_mahal = [0.5*mahal(ltri_unpack(exp(g.LDs[i]), g.LTs[i]), g.mus[i], x) for i in 1:g.n]

  determinants = [sum(g.LDs[i]) for i in 1:g.n]

  logsumexp(g.alphas + determinants - d_mahal) - logsumexp(g.alphas) - halflog2π*g.d
end

ll2 = log_likelihood_2(g, x)
@printf("ll0=%f, ll2=%f, rat=%f\n", ll0, ll2, ll0/ll2)
@test_approx_eq_eps ll0 ll2 1e-12

# Matrix version, in case needed
function log_likelihood_mat(alphas::Vec, mus::Mat, LDs::Mat, LTs::Mat, x::Vec)
  n = length(alphas)
  d = size(mus,1)
  L(i::Int) = ltri_unpack(exp(LDs[:,i]), LTs[:,i])
  d_mahal = [0.5*mahal(L(i), mus[:,i], x) for i in 1:n]

  determinants = [sum(LDs[:,i]) for i in 1:n]

  logsumexp(alphas + determinants - d_mahal) - logsumexp(alphas) - halflog2π*d
end

ll3 = log_likelihood_mat(g.alphas, hcat(g.mus...), hcat(g.LDs...), hcat(g.LTs...), x)
@printf("ll0=%f, ll3=%f, rat=%f\n", ll0, ll3, ll0/ll3)

using Stats
using Distributions

function draw(g::GMM)
  rng = Distributions.Categorical(exp(alphas)/sum(exp(alphas)))
  k = rand(rng)
  rng = Distributions.MultivariateNormal(mus[k], make_L(exp(g.LDs[k]), g.LTs[k]))
  x::Vec = rand(rng)
end

#   using Gadfly
#   x = [draw(g) for k in 1:100]
#   Gadfly.plot(sin, 0,1)

###################
# Define the derivative of sumsq
rdiff(sumsq, (zeros(3),))

@deriv_rule sumsq(x)   x   2x

sumsq2(x::Number,y::Number) = x^2 + y^2
sumsq2_rdiff = rdiff(sumsq2, (0.,0.))

#d1_sumsq2(x,y) = sumsq2_rdiff(x,y)[2]
#d2_sumsq2(x,y) = sumsq2_rdiff(x,y)[3]
#@deriv_rule sumsq2(x,y)   x   d1_sumsq2(x,y)
#@deriv_rule sumsq2(x,y)   y   d2_sumsq2(x,y)

rosenbrock(x::Float32,y::Float32) = sumsq2(1 - x, 10*(x - y^2))
rosenbrock_diff = rdiff(rosenbrock, (0.,0.))
rosenbrock_diff(1.,2.)

rosenbrockN_ex = :((1 - x[1])^2 + 100*sumsq(x[2:end] - x[1:end-1].^2))
rdiff(rosenbrockN_ex, x=zeros(4))

mahal_ex= :(mahal(ltri_unpack(x[1:3], x[4:6]), x[7:9], x[10:12]))

@eval mahal_new(x) = $mahal_ex
mahal_new(rand(12))

f=rdiff(mahal, (zeros(3,3),zeros(3),zeros(3),), )
f.
f(rand(3,3), rand(3), rand(3))

# This one not so good: unmanaged type int
lla(alphas::Vec) = log_likelihood_mat(alphas, hcat(g.mus...), hcat(g.LDs...), hcat(g.LTs...), x)
dx = rdiff(lla, alphas=g.alphas)
dx([-1,-248,-240,-1,-2])

Pkg.dir("ReverseDiffSource")
#           Query(

# abstract Visitor

# type myvisitor <: Visitor
#   index::Int64

#   myvisitor() = new(0)
# end

# v = myvisitor()

# function visit(v::myvisitor, obj::Number)
#   ++v.index
#   println(v)
# end

# function visit(v::Visitor, obj::Any)
#   fields = names(obj)
#   if length(fields) == 0
#     println("Cannot handle a [", typeof(obj), "]")
#   else
#     for n in names(obj)
#       visit(v, obj.(n))
#     end
#   end
# end

# visit(v, gmm_test)

# methods(Array)


import Statistics, DataFrames, StatsBase

"""
    standardize(X; robust::Bool=false)

Standardize (scale and reduce, Z-score) X so that the variables are expressed in terms of standard deviation (*i.e.*, mean = 0, SD = 1).

# Arguments
- `X`: Array or DataFrame.
- `robust::Bool`: If true, the standardization will be based on `median` and `mad` instead of `mean` and `sd` (default).

!!! note

    **Ideas / help required:**
    - Deal with missing values (See [#4](https://github.com/neuropsychology/Psycho.jl/issues/4))

# Examples
```jldoctest
standardize([1, 2, 3])

# output

3-element Array{Float64,1}:
 -1.0
  0.0
  1.0
```
"""
function standardize end

"""
    standardize!(X; robust::Bool=false)

In-place version of [`standardize`](@ref).
"""
function standardize! end

function standardize(X::Array{<:Number}; robust::Bool=false)

    if robust == false
        z = (X .- Statistics.mean(X, dims=1)) ./ Statistics.std(X, dims=1)
    else
        z = (X .- Statistics.median(X, dims=1)) ./ mapslices(x -> StatsBase.mad(x, normalize=true), X, dims=1)
    end
    return z
end

function standardize(X::Vector{<:Union{Missing, Number}}; robust::Bool=false)

    complete = collect(skipmissing(X))

    if robust == false
        z = (X .- Statistics.mean(complete)) ./ Statistics.std(complete)
    else
        z = (X .- Statistics.median(complete)) ./ StatsBase.mad(complete, normalize=true)
    end
    return z
end

function standardize!(X::Array{<:Union{Missing, Number}}; robust=false)
    X[:] = standardize(X, robust=robust)
end

function standardize!(X::DataFrames.DataFrame; robust=false)
    for (colname, col) in DataFrames.eachcol(X)
        if eltype(col) <: Union{Missing, Number}
            if length(unique(col)) > 2
                X[colname] = standardize(col)
            end
        end
    end
end

function standardize(X::DataFrames.DataFrame; robust=false)
    df = copy(X)
    standardize!(df; robust=robust)
    return df
end


# TODO: use StatePrior here, too.

# a gibbs sampler to compare.

# general algorithm:
# * pull out data point x_i
# * compute p(z_i = k | x_i) \propto p(x_i | z_i=k) p(z_i=k)
# * sample z_i from this.

"""
    GibbsCRP

Non-parametric clustering via Gibbs sampling with a Chinese Restaurant Process.

# Constructors

    GibbsCRP(prior, α, x)

Initialize all data assigned to one component.

# Fields

* `prior::NormalInverseChisq`: Prior for each cluster
* `logα::Float64`: (log) pseudo observations reserved for a new cluster
* `components::Vector{Component}`: Clusters
* `empties::Vector{Int}`: Keeps track of which indices are empty and can be
   recycled
* `assignments::Vector{Int}`: Current cluster assignments for each data point
* `data::Vector{Float64}`: Data points to cluster

"""
mutable struct GibbsCRP
    prior::NormalInverseChisq
    logα::Float64
    components::Vector{Component}
    empties::Vector{Int}
    assignments::Vector{Int}
    data::Vector{Float64}
end

GibbsCRP(prior::NormalInverseChisq, α::Float64, x::Vector{Float64}) =
    GibbsCRP(prior,
             log(α),
             [Component(prior, suffstats(Normal, x)), Component(prior)],
             [2],
             ones(Int, length(x)),
             x)

function StatsBase.sample!(gc::GibbsCRP, i::Int)
    x = gc.data[i]
    old_comp = sub(gc.components[gc.assignments[i]], x)
    gc.components[gc.assignments[i]] = old_comp

    # recycle if empty
    isempty(old_comp) && push!(gc.empties, gc.assignments[i])

    log_probs =
        [isempty(comp) ?
           -Inf :
           logpdf(comp, x) + log(nobs(comp))
         for comp
         in gc.components]

    log_probs[gc.empties[1]] = logpdf(Component(gc.prior), x) + gc.logα

    new_k = sample(Weights(exp.(log_probs)))
    gc.components[new_k] = add(gc.components[new_k], x)
    gc.assignments[i] = new_k

    # clean up empties
    if new_k == gc.empties[1]
        popfirst!(gc.empties)
        if isempty(gc.empties)
            push!(gc.components, Component(gc.prior))
            push!(gc.empties, length(gc.components))
        end
    end

    return gc
end

function StatsBase.sample!(gc::GibbsCRP)
    isempty(gc.data) &&
        throw(ArgumentError("can't sample from Gibbs sampler with empty .data field"))
    for i in 1:length(gc.data)
        sample!(gc, i)
    end
    return gc
end

mutable struct GibbsCRPSamples
    g::GibbsCRP
    assignments::Matrix{Int}
    components::Vector{Vector{Component}}
end

function fit!(gc::GibbsCRP, progress=true;
              samples=5000, burnin=samples÷10)
    assignments = Matrix{Int}(length(gc.data), samples)
    components = Vector{Vector{Component}}()
    for _ in 1:burnin
        sample!(gc)
    end

    @showprogress (progress ? 1 : Inf) "Gibbs sampling" for i in 1:samples
        sample!(gc)
        assignments[:,i] .= gc.assignments
        push!(components, copy(gc.components))
    end
    GibbsCRPSamples(gc, assignments, components)
end

assignments(gf::GibbsCRPSamples) = gf.assignments
ncomponents_dist(gf::GibbsCRPSamples) =
    fit(Categorical, mapslices(x -> length(unique(x)), gf.assignments, 1))

marginal_log_posterior(g::GibbsCRP) = marginal_log_posterior(g.components, g.logα)


function JuliaBench(operationMode)

	allFunctions = [MatrixGeneration, MatrixAddition, MatrixMultiplication, MatrixQuadraticForm, MatrixReductions, ElementWiseOperations, MatrixExp, MatrixSqrt, Svd, Eig, CholDec, MatInv, LinearSystem, LeastSquares, CalcDistanceMatrix, KMeans];

	allFunctionsString = ["Matrix Generation", "Matrix Addition", "Matrix Multiplication", "Matrix Quadratic Form", "Matrix Reductions", "Element Wise Operations", "Matrix Exponential", "Matrix Square Root", "Singular Value Decomposition", "Eigen Decomposition","Cholesky Decomposition", "Matrix Inversion", "Linear System Solution", "Linear Least Squares", "Squared Distance Matrix", "K-Means"];

	if (operationMode == 1) # partial benchmark
		vMatrixSize =  dropdims(readdlm(joinpath("Inputs","vMatrixSizePartial.csv"), ',',Int64), dims=1);
		numIterations = dropdims(readdlm(joinpath("Inputs","numIterationsPartial.csv"), ',',Int64), dims=1);

	elseif (operationMode == 2) # full benchmark
		vMatrixSize = dropdims(readdlm(joinpath("Inputs","vMatrixSizeFull.csv"), ',',Int64), dims=1);
		numIterations =  dropdims(readdlm(joinpath("Inputs","numIterationsFull.csv"), ',',Int64), dims=1);

	elseif (operationMode == 0) # Test benchmark
		vMatrixSize = 2;
		numIterations =  1;

	end

	numIterations = numIterations[1]; # It is 1x1 Array -> Scalar

	mRunTime = zeros(length(vMatrixSize), length(allFunctions), numIterations);
	tRunTime= Array{Any}(undef,length(allFunctions)+1,length(vMatrixSize)+1)# a table containing all the information
	tRunTime[1,1]="FunctionName\\MatrixSize";

	for ii = 1:length(vMatrixSize)
		matrixSize = vMatrixSize[ii];
		mX = randn(matrixSize, matrixSize);
		mY = randn(matrixSize, matrixSize);
		println("Matrix Size - $matrixSize");

		jj=1;
		for fun in allFunctions
			println("Processing $(allFunctionsString[jj]) - MatrixSize= $matrixSize");

			for kk = 1:numIterations;

				benchIJK =@benchmark $fun($matrixSize, $mX, $mY)
				# t =@benchmarkable $fun($matrixSize, $mX, $mY);
				# tune!(t)
				# run(t)

				mRunTime[ii, jj, kk]=median(benchIJK).time/1e3;
				# println("$(mRunTime[ii, jj, kk])")

			end
			tRunTime[jj+1,1]="$(allFunctionsString[jj])";
			tRunTime[1,ii+1]="$matrixSize";
			tRunTime[jj+1,ii+1]=mean(mRunTime[ii, jj,:]);
			jj+=1;
		end

	end

	return tRunTime, mRunTime;
end

function MatrixGeneration( matrixSize, mX, mY )

	mA = randn(matrixSize, matrixSize);
	mB = rand(matrixSize, matrixSize);

	return mA;
end

function MatrixAddition( matrixSize, mX, mY )

	scalarA = rand();
	scalarB = rand();

	mA = (scalarA .* mX) .+ (scalarB .* mY);

	return mA;
end

function MatrixMultiplication( matrixSize, mX, mY )

	scalarA = rand();
	scalarB = rand();

	mA = (scalarA .+ mX) * (scalarB .+ mY);
	# mA = BLAS.gemm!('N', 'N', scalarA, mX, collect(I), scalarB, mY)

	return mA;
end

function MatrixQuadraticForm( matrixSize, mX, mY )

	vX = randn(matrixSize);
	vB = randn(matrixSize);
	scalarC = rand();

	mA = (transpose(mX * vX) * (mX * vX)) .+ (transpose(vB) * vX) .+ scalarC;

	return mA;
end

function MatrixReductions( matrixSize, mX, mY )

	mA = sum(mX, dims=1) .+ minimum(mY, dims=2); #Broadcasting

	return mA;
end

function ElementWiseOperations( matrixSize, mX, mY )

	mA = rand(matrixSize, matrixSize);
	mB = 3 .+ rand(matrixSize, matrixSize);
	mC = rand(matrixSize, matrixSize);

	mD = abs.(mA) .+ sin.(mB);
	mE = exp.(-(mA .^ 2));
	mF = (-mB .+ sqrt.((mB .^ 2) .- (4 .* mA .* mC))) ./ (2 .* mA);

	mA = mD .+ mE .+ mF;

	return mA;
end

function MatrixExp( matrixSize, mX, mY )

	mA = exp(mX);

	return mA;
end

function MatrixSqrt( matrixSize, mX, mY )

	mY = transpose(mX) * mX;

	mA = sqrt(mY);

	return mA;
end

function Svd( matrixSize, mX, mY )

	F = svd(mX, full = false); # F is SVD object
	mU, mS, mV = F;

	return mA=0;
end

function Eig( matrixSize, mX, mY )

	F  = eigen(mX); # F is eigen object
	mD, mV = F;

	return mA=0;
end

function CholDec( matrixSize, mX, mY )

	mY = transpose(mX) * mX;

	mA = cholesky(mY);

	return mA;
end

function MatInv( matrixSize, mX, mY )

	mY = transpose(mX) * mX;

	mA = inv(mY);
	mB = pinv(mX);

	mA = mA .+ mB;

	return mA;
end

function LinearSystem( matrixSize, mX, mY )

	mB = randn(matrixSize, matrixSize);
	vB = randn(matrixSize);

	vA = mX \ vB;
	mA = mX \ mB;

	mA = mA .+ vA;

	return mA;
end

function LeastSquares( matrixSize, mX, mY )

	mB = randn(matrixSize, matrixSize);
	vB = randn(matrixSize);

	mXT=transpose(mX);
	vA = ( mXT * mX) \ ( mXT * vB);
	mA = ( mXT * mX) \ ( mXT * mB);

	mA = mA .+ vA;

	return mA;
end

function CalcDistanceMatrix( matrixSize, mX, mY )

	mY = randn(matrixSize, matrixSize);

	mA = transpose( sum(mX .^ 2, dims=1) ) .- (2 .* transpose(mX) * mY) .+ sum(mY .^ 2, dims=1);

	return mA;
end

function KMeans( matrixSize, mX, mY )

	# Assuming Samples are slong Columns (Rows are features)
	numClusters  = Int64( max( round(matrixSize / 100), 1 ) ); # % max between 1 and round(...)
	numIterations = 10;

	# http://stackoverflow.com/questions/36047516/julia-generating-unique-random-integer-array
	mA = mX[:, randperm(matrixSize)[1:numClusters]]; #<! Cluster Centroids

	for ii = 1:numIterations
		vMinDist, mClusterId = findmin( transpose(sum(mA .^ 2, dims=1)) .- (2 .* transpose(mA)* mX), dims=1); #<! Is there a `~` equivalent in Julia?
		vClusterId = LinearIndices( dropdims(mClusterId, dims=1) ); # to be able to access it later

		for jj = 1:numClusters
			mA[:, jj] = mean( mX[:, vClusterId .== jj ], dims=2 );
		end
	end

	mA = mA[:, 1] .+ transpose(mA[:, end]);

	return mA;
end


using DrWatson
@quickactivate :FractalDimension # uses DynamicalSystems, PyPlot
include(srcdir("style.jl"))

N = Int(1e5)

qH = 1
qC = 2

system = :henonheiles_chaotic
data = :embed_system
embedding_ds = 2:7
Cmethod = "standard"

# Calculate values for H, C
eHs, eCs, Hs, Cs = [Vector{Float64}[] for i in 1:4]
for d in embedding_ds
    params = @strdict N d qH qC data system
    if Cmethod ≠ "standard"
        params["Cmethod"] = Cmethod
    end

    params["z"] = -1

    # This is the main call that calculates everything
    output, s = produce_or_load(
        datadir("main"), params, make_C_H;
        prefix = string(data), suffix = "jld2", force = false,
        ignores = ["data"], tag = false,
    )
    @unpack eH, eC, H, C = output
    push!(eHs, eH); push!(Hs, H); push!(eCs, eC); push!(Cs, C)
end

# Do the actual plot
labels = ["\$d=$(d)\$" for d in embedding_ds]
legendtitle = "delay embedding (Hénon-Heiles, Δ ≈ 3)"

fig, axs = mainplot(
    Hs, Cs, eHs, eCs, labels, legendtitle;
    qH, qC, tol = 0.25,
)

wsave(plotsdir("paper", "embedding_$(system)"), fig)


@doc raw"""
    Vector2

An object `v` of the type [`Vector2`](@ref) is a vector in the two-dimensional
vector space ``ℝ^2``.

Geometrically spoken, a vector is the difference of two points ``p_2, p_1`` and
and denotes the direction and the distance from ``p_1`` to ``p_2``.

CGAL defines a symbolic constant [`NULL_VECTOR`](@ref). We will explicitly
state where you can pass this constant as an argument instead of a vector
initialized with zeros.
"""
Vector2

"""
    Vector2(a::Point2, b::Point2)

Introduces the vector `b - a`.
"""
Vector2(a::Point2, b::Point2)

"""
    Vector2(s::Segment2)

Introduces the vector `target(s) - source(s)`.
"""
Vector2(s::Segment2)

"""
    Vector2(r::Ray2)

Introduces the vector having the same direction as `r`.
"""
Vector2(r::Ray2)

"""
    Vector2(l::Line2)

Introduces the vector having the same direction as `t`.
"""
Vector2(l::Line2)

"""
    Vector2(::NullVector = NULL_VECTOR)

Introduces a null vector `v`.
"""
Vector2(::NullVector = NULL_VECTOR)

"""
    Vector2(x::Real, y::Real)

Introduces a vector `v` initialized to ``(x, y)``.
"""
Vector2(x::Real, y::Real) = Vector2(convert(FT, x), convert(FT, y))

@doc raw"""
    Vector2(hx::Real, hy::Real, hw::Real = 1)

Introduces a vector `v` initialized to `(hx/hw,hy/hw)`.

!!! info "Precondition"

    `hw` ``≠`` `0`
"""
Vector2(hx::Real, hy::Real, hw::Real) =
    Vector2(convert(RT, hx), convert(RT, hy), convert(RT, hw))

"""
    hx(p::Vector2)

Returns the homogeneous ``x`` coordinate.
"""
hx(p::Vector2)

"""
    hy(p::Vector2)

Returns the homogeneous ``y`` coordinate.
"""
hy(p::Vector2)

"""
    hw(p::Vector2)

Returns the homogenizing coordinate.
"""
hw(p::Vector2)

"""
    x(p::Vector2)

Returns the Cartesian ``x`` coordinate, that is [`hx()`](@ref)/[`hw()`](@ref).
"""
x(p::Vector2)

"""
    y(p::Vector2)

Returns the Cartesian ``y`` coordinate, that is [`hy()`](@ref)/[`hw()`](@ref).
"""
y(p::Vector2)

@doc raw"""
    homogeneous(v::Vector2, i::Integer)

Returns the `i`ᵗʰ homogeneous coordinate of `v`.

!!! info "Precondition"

    ``0 ≤ i ≤ 2``.
"""
homogeneous(v::Vector2, i::Integer)

@doc raw"""
    cartesian(v::Vector2, i::Integer)

Returns the `i`ᵗʰ Cartesian coordinate of `v`.

!!! info "Precondition"

    ``0 ≤ i ≤ 1``.
"""
cartesian(v::Vector2, i::Integer)

"""
    dimension(v::Vector2)

Returns the dimension (the constant 2).
"""
dimension(v::Vector2)

"""
    direction(v::Vector2)

Returns the direction which passes through `v`.
"""
direction(v::Vector2)

"""
    transform(v::Vector2, t::AffTransformation2)

Returns the vector obtained by applying `t` on `v`.
"""
transform(v::Vector2, t::AffTransformation2)

"""
    perpendicular(v::Vector2, o::Orientation)

Returns the vector perpendicular to `v` in clockwise or counterclockwise
orientation.
"""
perpendicular(v::Vector2, o::Orientation)

"""
    squared_length(v::Vector2)

Returns the squared length of `v`.
"""
squared_length(v::Vector2)

"""
    ==(v::Vector2, w::Vector2)

Test for equality: two vectors are equal, iff their ``x`` and ``y`` coordinates
are equal.

You can compare a vector with the `NULL_VECTOR`.
"""
==(v::Vector2, w::Vector2)

"""
    +(v::Vector2, w::Vector2)

Addition.
"""
+(v::Vector2, w::Vector2)

"""
    -(v::Vector2, w::Vector2)

Subtraction.
"""
-(v::Vector2, w::Vector2)

"""
    -(v::Vector2)

Returns the opposite vector.
"""
-(v::Vector2)

"""
    *(v::Vector2, w::Vector2)

Returns the scalar product (= inner product) of the two vectors.
"""
*(v::Vector2, w::Vector2)

"""
    /(v::Vector2, s::Real)

Division by a scalar.
"""
/(v::Vector2, s::Real)

@cxxdereference Base.:/(v::Vector2, s::Real) = v / convert(RT, s)


using ADI
using HCIToolbox
using Test
using Statistics
using LinearAlgebra
using StableRNGs

rng = StableRNG(8799)

@testset "Median" begin include("median.jl") end
@testset "PCA" begin include("pca.jl") end
@testset "TPCA" begin include("tpca.jl") end
@testset "NMF" begin include("nmf.jl") end
@testset "GreeDS" begin include("greeds.jl") end

@testset "Metrics" begin include("metrics.jl") end

@testset "SDI" begin include("sdi.jl") end


type FloatVectorRepresentation <: Representation
  length::Int
  min::Float
  max::Float
  range::FloatRange

  FloatVectorRepresentation(length::Int, min::Float, max::Float) =
    new(length, min, max, min:eps():max)
end

"""
Float vectors are implemented as a fixed-length vector of floating point
numbers.

**Parameters:**
* `length::Int`, the length of this vector. Defaults to 80, if no value is
supplied.
* `min::Float`, the minimum value that a component of this vector may assume.
Defaults to 0.0, if no value is supplied.
* `max::Float`, the maximum value that a component of this vector may assume.
Defaults to 1.0, if no value is supplied.
"""
float_vector{S <: AbstractString}(args::Dict{S, Any}) =
  float_vector( get(args, "length", 80),
                get(args, "min", 0.0),
                get(args, "max", 1.0))
float_vector() = float_vector(80, 0.0, 1.0)
float_vector(length::Int) = float_vector(length, 0.0, 1.0)
float_vector(length::Int, min::Float, max::Float) =
  FloatVectorRepresentation(length, min, max)

minimum_value(r::FloatVectorRepresentation) = r.min
maximum_value(r::FloatVectorRepresentation) = r.max
codon_type(::FloatVectorRepresentation) = Float
chromosome(r::FloatVectorRepresentation) = Vector{Float}
describe(i::Vector{Float}) = "($(join(i, ", ")))"

"""
Generates a float vector at random using its representation.
"""
generate(r::FloatVectorRepresentation) = Base.rand(r.range, r.length)


using PyPlot, Seismic

dt = 0.002
fn = 1/(2*dt)

# Ricker
w1 = Ricker(dt=dt)
nw = length(w1)
nc = floor(Int, nw/2)
t1 = dt*collect(-nc:1:nc)
nf = 8*nextpow2(nw)
df = 1/(nf*dt)
f1 = df*collect(0:1:nf-1)
wpad = cat(1,w1,zeros(nf-nw))
W1 = abs(fft(wpad))
W1 = W1/maximum(W1)

# Ormsby
w2 = Ormsby(dt=dt, f=[5.0, 10.0, 30.0, 55.0])
nw = length(w2)
t2 = dt*collect(0:1:nw-1)
nf = 8*nextpow2(nw)
df = 1/(nf*dt)
f2 = df*collect(0:1:nf-1)
wpad = cat(1,w2,zeros(nf-nw))
W2 = abs(fft(wpad))
W2 = W2/maximum(W2)

# Berlage
w3 = Berlage(dt=dt)
nw = length(w3)
t3 = dt*collect(0:1:nw-1)
nf = 8*nextpow2(nw)
df = 1/(nf*dt)
f3 = df*collect(0:1:nf-1)
wpad = cat(1,w3,zeros(nf-nw))
W3 = abs(fft(wpad))
W3 = W3/maximum(W3)

figure("Suite of wavelets in time")

subplot(221)
plot(t1, w1)
axis("tight")
xlabel("Time (s)")
title("Ricker wavelet")

subplot(222)
plot(t2, w2)
axis("tight")
xlabel("Time (s)")
title("Ormsby wavelet")

subplot(223)
plot(t3, w3)
axis("tight")
xlabel("Time (s)")
title("Berlage wavelet")

tight_layout()

figure("Suite of amplitude spectrums")

subplot(221)
plot(f1, W1)
xlim([0, fn])
xlabel("Frequency (Hz)")
title("Ricker spectrum")

subplot(222)
plot(f2, W2)
xlim([0, fn])
xlabel("Frequency (Hz)")
title("Ormsby spectrum")

subplot(223)
plot(f3, W3)
xlim([0, fn])
xlabel("Frequency (Hz)")
title("Berlage spectrum")

tight_layout()


using EarthDataLab
using Test
import DataFrames: DataFrame,aggregate
import Dates: year
import Statistics: cor, mean

@testset "Dataframe representation" begin
c=Cube()

d = getCubeData(c,variable=["air_temperature_2m","gross_primary_productivity"],longitude=(10,11),latitude=(50,51),
                time=(Date("2002-01-01"),Date("2008-12-31")))

dmem=readcubedata(d)

function docor(xout,xin)
    #Inside this function, xin is now a data frame
    @test isa(xin,DataFrame)
    xout[1]=cor(xin[:air_temperature_2m],xin[:gross_primary_productivity])
end
indims = InDims(TimeAxis,VariableAxis,artype = AsDataFrame())
outdims = OutDims()
o = mapCube(docor,dmem,indims=indims,outdims=outdims)

@test all(isapprox.(o.data,[cor(dmem.data[i,j,:,1],dmem.data[i,j,:,2]) for i=1:4, j=1:4]))

function annMean(xout,xin)
    #xin is now a DataFrame where time is added as the third column
    #We derive the year and add it to the dataframe
    xin[!,:year] = year.(xin[!,:Time])
    #Now we do the annual aggregation, note that we have to exclude the time column, because we can't aggregate here
    x2 = aggregate(xin[!,[1,2,4]],:year,a->mean(skipmissing(a)))
    #We copy the results to our output
    xout[:,1] = x2[!,2]
    xout[:,2] = x2[!,3]
end
indims = InDims("time","var",artype=AsDataFrame(true))
outdims = OutDims(RangeAxis("Year",2002:2008),"var")

o = mapCube(annMean,dmem,indims=indims,outdims=outdims)

@test all(isapprox.(o.data[1,1,:,:],mean(dmem.data[:,:,1:46,1],dims=3)))
end


# ------------------------------------------------------------------
# Copyright (c) 2018, Júlio Hoffimann Mendes <juliohm@stanford.edu>
# Licensed under the ISC License. See LICENCE in the project root.
# ------------------------------------------------------------------

"""
    fit_mle(gev, bm)

Fit generalized extreme value distribution `gev` to block maxima
`bm` with constrained maximum likelihood estimation.
"""
function fit_mle(::Type{GeneralizedExtremeValue}, bm::BlockMaxima)
  # retrieve maxima values
  x = collect(bm)
  n = length(x)

  # MLE for case ξ ≠ 0
  mle = Model(with_optimizer(Ipopt.Optimizer, print_level=0, sb="yes"))
  @variable(mle, μ, start=0.0)
  @variable(mle, σ, start=1.0)
  @variable(mle, ξ, start=0.1)
  @NLobjective(mle, Max,
    -n*log(σ)
    -(1 + 1/ξ)*sum(log(1 + ξ*(x[i]-μ)/σ) for i in 1:n)
    -sum((1 + ξ*(x[i]-μ)/σ)^(-1/ξ) for i in 1:n)
  )
  @NLconstraint(mle, [i=1:n], 1 + ξ*(x[i]-μ)/σ ≥ 1e-6)
  @NLconstraint(mle, abs(ξ) ≥ 1e-6)
  @constraint(mle, σ ≥ 1e-6)

  # MLE for case ξ = 0
  mle₀ = Model(with_optimizer(Ipopt.Optimizer, print_level=0, sb="yes"))
  @variable(mle₀, μ₀, start=0.0)
  @variable(mle₀, σ₀, start=1.0)
  @NLobjective(mle₀, Max,
    -n*log(σ₀)
    -sum((x[i]-μ₀)/σ₀ for i in 1:n)
    -sum(exp(-(x[i]-μ₀)/σ₀) for i in 1:n)
  )
  @constraint(mle₀, σ₀ ≥ 1e-6)

  # attempt to solve both cases
  optimize!(mle)
  optimize!(mle₀)

  # retrieve solver status
  status  = termination_status(mle)
  status₀ = termination_status(mle₀)

  # acceptable statuses
  OK = (MOI.OPTIMAL, MOI.LOCALLY_SOLVED)

  if status ∈ OK && status₀ ∈ OK
    # choose the maximum amongst the two
    if objective_value(mle) > objective_value(mle₀)
      GeneralizedExtremeValue(value(μ), value(σ), value(ξ))
    else
      GeneralizedExtremeValue(value(μ₀), value(σ₀), 0.)
    end
  elseif status ∈ OK
    GeneralizedExtremeValue(value(μ), value(σ), value(ξ))
  elseif status₀ ∈ OK
    GeneralizedExtremeValue(value(μ₀), value(σ₀), 0.)
  else
    error("could not fit distribution to maxima")
  end
end

"""
    fit_mle(gp, pm)

Fit generalized Pareto distribution `gp` to peak over threshold
maxima `pm` with constrained maximum likelihood estimation.
"""
function fit_mle(::Type{GeneralizedPareto}, pm::PeakOverThreshold)
  # retrieve maxima values
  x = collect(pm)
  n = length(x)

  # compute excess
  y = x .- pm.u

  # MLE for case ξ ≠ 0
  mle = Model(with_optimizer(Ipopt.Optimizer, print_level=0, sb="yes"))
  @variable(mle, σ, start=1.0)
  @variable(mle, ξ, start=0.1)
  @NLobjective(mle, Max,
    -n*log(σ)
    -(1 + 1/ξ)*sum(log(1 + ξ*y[i]/σ) for i in 1:n)
  )
  @NLconstraint(mle, [i=1:n], 1 + ξ*y[i]/σ ≥ 1e-6)
  @NLconstraint(mle, abs(ξ) ≥ 1e-6)
  @constraint(mle, σ ≥ 1e-6)

  # MLE for case ξ = 0
  mle₀ = Model(with_optimizer(Ipopt.Optimizer, print_level=0, sb="yes"))
  @variable(mle₀, σ₀, start=1.0)
  @NLobjective(mle₀, Max,
    -n*log(σ₀)
    -sum(y[i] for i in 1:n)/σ₀
  )
  @constraint(mle₀, σ₀ ≥ 1e-6)

  # attempt to solve both cases
  optimize!(mle)
  optimize!(mle₀)

  # retrieve solver status
  status  = termination_status(mle)
  status₀ = termination_status(mle₀)

  # acceptable statuses
  OK = (MOI.OPTIMAL, MOI.LOCALLY_SOLVED)

  if status ∈ OK && status₀ ∈ OK
    # choose the maximum amongst the two
    if objective_value(mle) > objective_value(mle₀)
      GeneralizedPareto(0.0, value(σ), value(ξ))
    else
      GeneralizedPareto(0.0, value(σ₀), 0.0)
    end
  elseif status ∈ OK
    GeneralizedPareto(0.0, value(σ), value(ξ))
  elseif status₀ ∈ OK
    GeneralizedPareto(0.0, value(σ₀), 0.0)
  else
    error("could not fit distribution to maxima")
  end
end


using Distributions
import Distributions: logpdf, pdf, sample

using ConjugatePriors
import ConjugatePriors: NormalInverseGamma, NormalWishart, NormalInverseWishart, posterior

include("transform.jl")

import Base.LinAlg: Cholesky
import Base.Random: rand
import Base: getindex, mean, scale

export NormalInverseGamma, NormalWishart, NormalInverseWishart, posterior

flatten(x :: Float64, d :: Distribution) = ( [link(d,x)], a -> invlink(d,a[1]) )
flatten(xs :: Array{Float64}, d :: Distribution) = ( link(d,xs), a -> invlink(d,a) )

function flatten(xs :: Matrix{Float64}, d :: Distribution)
  ys = link(d,xs)
  shape = size(ys)
  rebuild(zs) = invlink(d, reshape(zs, shape))
  return (ys[:], rebuild)
end

function flatten(x :: Tuple{Real,Real}, d :: NormalInverseGamma)
  y = link(d,x)
  z = [y[1],y[2]]
  return (z, a -> invlink(d, (a[1],a[2])))
end


using Measures
include("helper_functions.jl")

begin
    λ1 = 0.4
    λ2 = 1.4
    α = 1
    μ(α) = 1 / (1.5 * (α))
    ν(α) = 1 / (0.7 * (α))
    parameters = (λ1, λ2, α, μ, ν)
end

ps, qs = old_equilibrium_probability(parameters...; returnall=true)

plim, qlim = old_equilibrium_probability(parameters...; returnall=false)

ps_plot = plot(ylabel = "Probability", xlabel = "", size = (600, 250), legend = :topright, lw = 1.5)
plot!(ps_plot, ps, label = "Best response", lc = :blue)
hline!([plim], lc = :red, ls = :dash, label = "Optimal probability")

qs_plot = plot(qs, label = "Opponent best response", size = (600, 250), legend = :bottomright, lc = :orange, lw = 1.5, ylabel = "Probability")
hline!([qlim], lc = :red, ls = :dash, label = "Optimal probability")

plot(ps_plot, qs_plot, xlabel = "Number of Iterations", margin = 2.2mm)
# savefig("outputs/convergence_plot.pdf")


module TestDiffMap
    using ManifoldLearning
    using Test

    @testset "DiffMap" begin

	d = 2
	t = 1
	X, L = swiss_roll()
	Y = transform(DiffMap, X; d=d, t=t)

	@test outdim(Y) == d
	@test size(projection(Y), 2) == size(X, 2)
    end
end


using Test
using SumOfSquares
using DynamicPolynomials

function motzkin_test(optimizer, config::MOIT.Config)
    @polyvar x y
    # Motzkin polynomial
    p = x^4*y^2 + x^2*y^4 + 1 - 3*x^2*y^2

    model = _model(optimizer)
    # We want to write `p ≥ 0` instead of `p in SOSCone()` so we need to set
    # the polymodule to `SumOfSquares` to tell `PolyJuMP` to interpret
    # `p ≥ 0` as a Sum-of-Squares constraint.
    PolyJuMP.setpolymodule!(model, SumOfSquares)
    con_ref = @constraint(model, p ≥ 0)

    JuMP.optimize!(model)
    @test JuMP.termination_status(model) == MOI.INFEASIBLE
    @test JuMP.dual_status(model) == MOI.INFEASIBILITY_CERTIFICATE

    q = (x^2 + y^2) * p
    delete(model, con_ref)
    @constraint(model, q >= 0)

    JuMP.optimize!(model)
    @test JuMP.termination_status(model) == config.optimal_status
    @test JuMP.primal_status(model) == MOI.FEASIBLE_POINT
end
sd_tests["motzkin"] = motzkin_test


function mw_marg_lf(x::Vector{Float64}, cov::Array{Float64,2})
# This is a scaled version of the density in King and Hillier (1985) Equation (2) @
# -- Note x is ncov x n ... compute n marginal likelihood values @
# Compute Marginal Likelihood value for x with covariance matrix cov
  x = x./sqrt(sum(x.^2, 1))'
  m = size(x, 1)
  d = det(cov)
  # (gauss code) x = chol(invpd(cov))
  c = chol(Hermitian(inv(x'x)))
  print(x)
  #x = c*x
  tmp = sum(x, 1)
  tmp = d*(tmp.^m)
  mlf = sqrt(1.0./tmp)[:]
  return mlf
end
function mw_marg_lf(x::Array{Float64,2}, cov::Array{Float64,2})
  nsim = size(x, 2)
  println(nsim)
  mlf_vec = Vector{Float64}(nsim)
  for isim ∈ 1:nsim
    mlf_vec[isim] = mw_marg_lf(x[:, isim], cov)[1]
  end
return mlf_vec
end


using Nemo, FastArithmetic, BenchmarkTools

const k=ResidueRing(ZZ,5) # ,u=FiniteField(5,1,"u")
const K,t=PolynomialRing(k,"t")

Base.size(k::Nemo.GenResRing{Nemo.fmpz}) = Int(modulus(k))
Base.size{T<:Nemo.FinField}(k::T) = Int(order(k))

function create(n::Int64,k::Nemo.Ring)
  A=Array(k,n)
  c=size(k)-1
  for i in 1:n
    A[i]=k(rand(0:c))
  end
  return A
end

function createPol(n::Int64,K::Nemo.Ring)
  k=base_ring(K)
  return K(create(n,k))
end

function createIrrPol(n::Int64,K::Nemo.Ring)
  P=createPol(n,K)
  while degree(P)!=n || length(factor(P))!=1
    P=createPol(n,K)
  end
  return P
end

function create2d(m::Int64,n::Int64,k::Nemo.Ring)
  A=Array(k,n,m)
  c=size(k)-1
  for i in 1:n, j in 1:m
    A[i,j]=k(rand(0:c))
  end
  return A
end


# A basic example of using the BayesianMixtures package.

using BayesianMixtures
B = BayesianMixtures

# Simulate some data
x = randn(500)

# Specify model, data, and MCMC options
n_total = 1000  # total number of MCMC sweeps to run
options = B.options("Normal","MFM",x,n_total, x_full=x_full)  # MFM model with univariate Normal components

# Run MCMC sampler
result = B.run_sampler(options)

# Display the posterior on k (number of components)
println("Posterior on k:")
for (k,pk) in enumerate(B.k_posterior(result; upto=10))
    @printf "%d: %.3g\n" k pk
end


# ALS for matrices without initial TT
function TTm_ALS(matrix::Matrix,middlesizes::Matrix{Int64},rnks::Vector{Int64})
    tensor  = reshape(matrix,(middlesizes[1,:]...,middlesizes[2,:]...));
    d       = length(rnks)-1
    permind = Tuple([ (i-1)*d+j for j in 1:d for i in 1:2 ]);
    tensor  = permutedims(tensor,permind);
    resind  = Tuple([prod(col) for col in eachcol(middlesizes)]);
    tensor  = reshape(tensor,resind)
    tt      = TT_ALS(tensor,rnks);
    return MPT( [reshape(tt[i],(rnks[i], middlesizes[:,i]..., rnks[i+1])) for i = 1:order(tt)],d );
end

# ALS for matrices with initial TT
function TTm_ALS(matrix::Matrix,ttm0::MPT{4})
    d           = order(ttm0)
    middlesizes = [[size(ttm0)[i][1] for i = 1:d]'; [size(ttm0)[i][2] for i = 1:d]'];
    ressz       = ([size(ttm0)[i][1] for i = 1:d]'..., [size(ttm0)[i][2] for i = 1:d]'... )
    tensor      = reshape(matrix,ressz);
    permind     = Tuple([ (i-1)*d+j for j in 1:d for i in 1:2 ]);
    tensor      = permutedims(tensor,permind);
    resind      = Tuple([prod(col) for col in eachcol(middlesizes)]);
    tensor      = reshape(tensor,resind)
    tt0         = mpo2mps(ttm0)
    tt          = TT_ALS(tensor,tt0);
    rnks        = rank(tt);
    return MPT( [reshape(tt[i],(rnks[i][1], middlesizes[i][1], middlesizes[i][1], rnks[i][2])) for i = 1:d],d );
end

# ALS for symmetric matrices
function TTm_ALS(symmat::Symmetric{Float64, Matrix{Float64}},middlesizes::Matrix{Int64},rnks::Vector{Int64})
    d     = length(rnks)-1
    cores = Vector{Array{Float64,4}}(undef,d);
    R     = Matrix{Float64}
    for i = 1:d
        cores[i] = zeros(rnks[i],middlesizes[1,i],middlesizes[2,i],rnks[i+1])
        for rleft = 1:rnks[i]
            for rright = 1:rnks[i+1]
                tmp = randn(middlesizes[1,i],middlesizes[2,i])
                cores[i][rleft,:,:,rright] = tmp*tmp'
            end
        end
        tmp      = reshape(cores[i],(rnks[i]*middlesizes[1,i]*middlesizes[2,i],rnks[i+1]))
        F        = qr(tmp)
        cores[i] = reshape(Matrix(F.Q),(rnks[i],middlesizes[1,i],middlesizes[2,i],rnks[i+1]))
        if i>1
            cores[i] = nmodeproduct(R,cores[i],1)
        end
        R = F.R
    end
    ttm0 = MPT(cores,d)
    return TTm_ALS(Matrix(symmat),ttm0)
end

# ALS for diagonal matrices
function TTm_ALS(diag::Vector,middlesizes::Matrix,rnks::Vector)
    d      = length(rnks)-1
    tensor = reshape(diag,Tuple(middlesizes))
    tt     = TT_ALS(tensor,rnks)
    cores  = Vector{Array{Float64,4}}(undef,N)
    for i = 1:N
        cores[i] = zeros(rnks[i],middlesizes[i],middlesizes[i],rnks[i+1])
        for ri = 1:rnks[i]
            for rii = 1:rnks[i+1]
                cores[i][ri,:,:,rii] = Matrix(Diagonal(tt[i][ri,:,rii]))
            end
        end
    end
    return MPT(cores,d)
end


################################################################################
#
#   COLLECTION OF FUNCTIONS FOR
#   BRILLOUIN ZONES
#
#   FILE CONTAINS
#   - mathematical helper functions for planes, lines and points
#   - construction of BZ parts (corners, faces, etc.)
#   - construction of BZ objects from reciprocal unitcells
#
################################################################################

################################################################################
#
#   MATHEMATICAL HELPER FUNCTIONS
#
################################################################################

# function to find out if two vectors are parallel
function parallelVectors(
            v1 :: Vector{<:Real},
            v2 :: Vector{<:Real}
        ) :: Bool

    # check if they can be scaled into one another (i.e. the rank of the matrix)
    r = rank([v1'; v2'])
    if r == 1
        return true
    elseif r == 2
        return false
    else
        error("Got " * string(r) * " from calculating the rank of two vectors:\n" *
        "v1 = " * string(v1) * ", and v2 = " * string(v2))
    end
end

# Find out if a point point_p is on a plane given by (plane_p,plane_n)
function pointOnPlane(
            point_p :: Vector{<:Real},
            plane_p :: Vector{<:Real},
            plane_n :: Vector{<:Real},
            error_allowed :: Real = 1e-15
        ) :: Bool

    # check if the plane normal is perpendicular to the connecting vector between plane basis point and point
    return dot(plane_p .- point_p, plane_n) < error_allowed
end

# Find out if two planes (p1,n2) and (p2,n2) are identical
function identicalPlanes(
            p1 :: Vector{<:Real},
            n1 :: Vector{<:Real},
            p2 :: Vector{<:Real},
            n2 :: Vector{<:Real},
            error_allowed :: Real = 1e-15
        ) :: Bool

    # check if p2 is on plane 1 and the normals are parallel
    return pointOnPlane(p2, p1,n1, error_allowed) && parallelVectors(n1,n2)
end

# Find out if three planes (pi,ni) have one intersecting point
function haveSingleIntersection(
            p1 :: Vector{<:Real},
            n1 :: Vector{<:Real},
            p2 :: Vector{<:Real},
            n2 :: Vector{<:Real},
            p3 :: Vector{<:Real},
            n3 :: Vector{<:Real},
            error_allowed :: Real = 1e-15
        ) :: Bool

    # calculate the determinant
    determinant = dot(n3, cross(n1, n2))
    # check if they intersect in a point
    if abs(determinant) < error_allowed
        # there will not be a single point of intersection
        return false
    else
        # they will intersect in one point
        return true
    end
end

# Get the intersecting point of three planes
function getSingleIntersection(
            p1 :: Vector{<:Real},
            n1 :: Vector{<:Real},
            p2 :: Vector{<:Real},
            n2 :: Vector{<:Real},
            p3 :: Vector{<:Real},
            n3 :: Vector{<:Real},
            error_allowed :: Real = 1e-15
        ) :: Vector{Float64}

    # calculate the main determinant
    determinant = dot(n3, cross(n1, n2))

    # calculate the other relevant determinants
    D = Float64[-dot(n1,p1), -dot(n2,p2), -dot(n3,p3)]
    det_x = det([
        D[1] n1[2] n1[3]
        D[2] n2[2] n2[3]
        D[3] n3[2] n3[3]
    ])
    det_y = det([
        n1[1] D[1] n1[3]
        n2[1] D[2] n2[3]
        n3[1] D[3] n3[3]
    ])
    det_z = det([
        n1[1] n1[2] D[1]
        n2[1] n2[2] D[2]
        n3[1] n3[2] D[3]
    ])
    # get the interesection point
    return [det_x, det_y, det_z] ./ determinant
end

################################################################################
#
#   COMPUTING BRILLOUIN ZONE OBJECTS
#   (GENERAL FUNCTIONS)
#
################################################################################

# obtain the Brillouin zone for a given reciprocal unitcell (L=2,N=2)
function getBrillouinZone(
            ::Type{BZ},
            reciprocal_unitcell :: RU
            ;
            max_r :: Real = 1.001
        ) :: BZ where {L,P<:AbstractReciprocalPoint{2},B<:AbstractBond{L,2},RU<:AbstractReciprocalUnitcell{P,B},BZ<:AbstractBrillouinZone{RU}}

    ##########
    # STEP 1 - Compute the reciprocal lattice points that are needed for the calculation
    ##########

    # build a test lattice
    reciprocal_lattice = getLatticePeriodic(reciprocal_unitcell, 2)

    # find the minimal distance to a the Gamma point
    max_distance = maximum([dot(k,k) for k in point.(sites(reciprocal_lattice))])
    min_distance = minimum([dot(k,k) for k in point.(sites(reciprocal_lattice))])

    # reconstruct the reciprocal lattice
    reciprocal_lattice = getLatticeInSphere(reciprocal_unitcell, sqrt(max_distance)*max_r)

    # list of mid points
    mid_points = Vector{Float64}[
        k.*0.5  for k in point.(sites(reciprocal_lattice)) if sum(k.*k) > 1e-5
    ]

    # list of normals
    normals = Vector{Float64}[
        [k[2], -k[1]] ./ (k[1]*k[1] + k[2]*k[2]) for k in mid_points
    ]

    ##########
    # STEP 2 - Compute the corners of the BZ
    ##########

    ############
    # STEP 2.1 - Compute insersection points
    ############

    # list of intersection points
    intersections = Vector{Float64}[]

    # find the intersections
    for i in 1:length(mid_points)
    for j in 1:length(mid_points)
        # continue if same index
        if i==j
            continue
        end
        # get data for point 1
        p1 = mid_points[i]
        d1 = normals[i]
        # get data for point 2
        p2 = mid_points[j]
        d2 = normals[j]
        # check the interesection point
        delta_1 =  d1[2]*p1[1] - d1[1]*p1[2]
        delta_2 =  d2[2]*p2[1] - d2[1]*p2[2]
        delta   = -d1[2]*d2[1] + d2[2]*d1[1]
        # push to the intersections (if not devided by zero)
        if abs(delta) > 1e-8
            push!(intersections,[
                (d1[1]*delta_2 - d2[1]*delta_1) / delta,
                (d1[2]*delta_2 - d2[2]*delta_1) / delta
            ])
        end
    end
    end

    ############
    # STEP 2.2 - Find closest intersection points --> points of BZ
    ############

    # list of distances
    intersection_distances = Float64[
        sum(isec.*isec) for isec in intersections
    ]

    # find the minimal intersection distance
    min_distance = minimum(intersection_distances)

    # list of points which are only the minimal intersection distance away (=corners)
    corners_BZ = Array{Float64, 1}[]

    # check all points
    for i in 1:length(intersections)
        # check the relative deviation to the calculated minimal distance
        if intersection_distances[i] < min_distance * (1.0 + 1e-8)
            # check if not inserted already
            inserted_already = false
            for p in corners_BZ
                if sum((p.-intersections[i]).*(p.-intersections[i])) < 1e-8
                    inserted_already = true
                end
            end
            # found a new point
            if !inserted_already
                push!(corners_BZ, intersections[i])
            end
        end
    end

    ##########
    # STEP 3 - Compute the faces of the BZ (based on corners)
    ##########

    # list of points if they are contained in the loop
    in_loop = Bool[false for p in corners_BZ]

    # the list of the loop (to bring the points in correct order)
    loop = Int64[]

    # start with the first point
    push!(loop, 1)
    in_loop[1] = true

    # look for next point (which is closest to the given point in all remaing ones)
    while length(loop) < length(corners_BZ)
        # check for the closest
        closest_distance = sum(corners_BZ[1].*corners_BZ[1]) * 1000
        closest_point = -1
        # check all other points
        for i in 1:length(corners_BZ)
            # if already in loop, continue
            if in_loop[i]
                continue
            end
            # compare to the nearest one
            if sum((corners_BZ[i].-corners_BZ[loop[end]]).*(corners_BZ[i].-corners_BZ[loop[end]])) < closest_distance
                closest_point    = i
                closest_distance = sum((corners_BZ[i].-corners_BZ[loop[end]]).*(corners_BZ[i].-corners_BZ[loop[end]]))
            end
        end
        # insert the closest point as the next in line
        push!(loop, closest_point)
        in_loop[closest_point] = true
    end

    # push the starting point into the loop to wrap up and make a closed line
    #push!(loop, loop[1])

    # list of edges (only contains this one loop)
    faces_BZ = Vector{Int64}[loop]

    ##########
    # STEP 4 - Construct a new object and return it
    ##########

    # return the finished BZ
    return newBrillouinZone(
        BZ,
        reciprocal_unitcell,
        corners_BZ,
        faces_BZ
    )
end

# obtain the Brillouin zone for a given reciprocal unitcell (L=3,N=3)
function getBrillouinZone(
            ::Type{BZ},
            reciprocal_unitcell :: RU
            ;
            max_r :: Real = 1.001
        ) :: BZ where {L,P<:AbstractReciprocalPoint{3},B<:AbstractBond{L,3},RU<:AbstractReciprocalUnitcell{P,B},BZ<:AbstractBrillouinZone{RU}}

    ##########
    # STEP 1 - Compute the reciprocal lattice points that are needed for the calculation
    ##########

    # build a test lattice
    reciprocal_lattice = getLatticePeriodic(reciprocal_unitcell, 2)

    # find the minimal distance to a the Gamma point
    max_distance = maximum([dot(k,k) for k in point.(sites(reciprocal_lattice))])
    min_distance = minimum([dot(k,k) for k in point.(sites(reciprocal_lattice))])

    # reconstruct the reciprocal lattice
    reciprocal_lattice = getLatticeInSphere(reciprocal_unitcell, sqrt(max_distance)*max_r)

    # k points
    k_points = point.(sites(reciprocal_lattice))

    # list of mid points
    mid_points = Vector{Float64}[
        k.*0.5  for k in point.(sites(reciprocal_lattice)) if sum(k.*k) > 1e-5
    ]

    # list of normals of planes
    normals = [m ./ dot(m,m) for m in mid_points]

    ##########
    # STEP 2 - Find all intersection points of 3 planes and check that they are closest to the Gamma point
    ##########

    # list of points which are only the minimal intersection distance to Gamma point
    points = Array{Float64, 1}[]

    # temp point
    point_tmp = zeros(Float64, 3)

    # lists for all planes that aided in the construction
    planes_n = Array{Float64,1}[]
    planes_p = Array{Float64,1}[]

    # find the intersections
    for i in 1:length(mid_points)
        # get data for plane 1
        p1 = mid_points[i]
        n1 = normals[i]
        # get a second plane
        for j in i+1:length(mid_points)
            # get data for plane 2
            p2 = mid_points[j]
            n2 = normals[j]
            # continue if they are parallel
            if dot(cross(n1,n2), cross(n1,n2)) < 1e-8
                continue
            end
            # get a third plane
            for l in j+1:length(mid_points)
                # get data for plane 3
                p3 = mid_points[l]
                n3 = normals[l]
                # calculate the determinant
                determinant = dot(n3, cross(n1, n2))
                # check if they intersect in a point
                if abs(determinant) < 1e-8
                    # there will not be one point of intersection
                    continue
                end
                # calculate the other relevant determinants
                #A = Float64[n1[1], n2[1], n3[1]]
                #B = Float64[n1[2], n2[2], n3[2]]
                #C = Float64[n1[3], n2[3], n3[3]]
                D = Float64[-dot(n1,p1), -dot(n2,p2), -dot(n3,p3)]
                #det_x = dot(D, cross(B, C))
                #det_y = dot(A, cross(D, C))
                #det_z = dot(A, cross(B, D))
                det_x = det([
                    D[1] n1[2] n1[3]
                    D[2] n2[2] n2[3]
                    D[3] n3[2] n3[3]
                ])
                det_y = det([
                    n1[1] D[1] n1[3]
                    n2[1] D[2] n2[3]
                    n3[1] D[3] n3[3]
                ])
                det_z = det([
                    n1[1] n1[2] D[1]
                    n2[1] n2[2] D[2]
                    n3[1] n3[2] D[3]
                ])
                # get the interesection point
                point_tmp[1] = det_x / determinant
                point_tmp[2] = det_y / determinant
                point_tmp[3] = det_z / determinant
                # check if the distance to any other k point in the lattice is shorter than to the Gamma point
                shorter_distance_found = false
                distance_to_Gamma = dot(point_tmp, point_tmp)
                # check all points
                for k in k_points
                    if dot(point_tmp .- k, point_tmp .- k)*(1 + 1e-8) < distance_to_Gamma
                        shorter_distance_found = true
                        break
                    end
                end
                # insert point into list if no shorter distance found
                if !shorter_distance_found
                    # check if not inserted already
                    inserted_already = false
                    # check all points
                    for p in points
                        if dot(p.-point_tmp, p.-point_tmp) < 1e-6
                            inserted_already = true
                            break
                        end
                    end
                    # if not inserted already, insert now
                    if !inserted_already
                        # insert the point into the list
                        push!(points, copy(point_tmp))
                        # see if the planes that cross here, have to be inserted as well

                        # check plane 1
                        insert_p1 = true
                        # look through all planes that are inserted already
                        for index in 1:length(planes_p)
                            # check if the point lies in the plane
                            if abs(dot(p1 .- planes_p[index], planes_n[index])) < 1e-6
                                # point lies in this plane, check the direction of the normal
                                if dot(cross(planes_n[index], n1), cross(planes_n[index], n1)) < 1e-6
                                    # normals also match in direction
                                    insert_p1 = false
                                    break
                                end
                            end
                        end
                        # if not found, insert into lists
                        if insert_p1
                            push!(planes_p, p1)
                            push!(planes_n, n1)
                        end

                        # check plane 2
                        insert_p2 = true
                        # look through all planes that are inserted already
                        for index in 1:length(planes_p)
                            # check if the point lies in the plane
                            if abs(dot(p2 .- planes_p[index], planes_n[index])) < 1e-6
                                # point lies in this plane, check the direction of the normal
                                if dot(cross(planes_n[index], n2), cross(planes_n[index], n2)) < 1e-6
                                    # normals also match in direction
                                    insert_p2 = false
                                    break
                                end
                            end
                        end
                        # if not found, insert into lists
                        if insert_p2
                            push!(planes_p, p2)
                            push!(planes_n, n2)
                        end

                        # check plane 3
                        insert_p3 = true
                        # look through all planes that are inserted already
                        for index in 1:length(planes_p)
                            # check if the point lies in the plane
                            if abs(dot(p3 .- planes_p[index], planes_n[index])) < 1e-6
                                # point lies in this plane, check the direction of the normal
                                if dot(cross(planes_n[index], n3), cross(planes_n[index], n3)) < 1e-6
                                    # normals also match in direction
                                    insert_p3 = false
                                    break
                                end
                            end
                        end
                        # if not found, insert into lists
                        if insert_p3
                            push!(planes_p, p3)
                            push!(planes_n, n3)
                        end

                    end
                end
            end
        end
    end

    ##########
    # STEP 5 - Connect the points in planes of the BZ --> edges of BZ
    ##########

    # the list of all the loops (i.e. all edge loops of the BZ)
    edges        = Array{Int64,1}[]
    # the list of all the loops but sorted without double point for the start+end
    edges_sorted = Array{Int64,1}[]

    # find all loops of all points by constructing planes

    # iterate over all planes
    for i in 1:length(planes_n)

        # find the subset of points that lie within the plane
        in_plane = Bool[abs(dot(p .- planes_p[i], planes_n[i])) < 1e-6 for p in points]

        # use the 2D loop finding within the plane
        point_indices = collect(1:length(points))[in_plane]

        # check if there are enough points to form a loop
        if length(point_indices) < 3
            # not enough points found
            continue
        end

        # list of points if they are contained in the loop
        in_loop = [false for i in point_indices]

        # the list of the loop (to bring the points in correct order)
        # with relative indices (have to be translated back)
        loop = Int64[]

        # start with the first (relative) point
        push!(loop, 1)
        in_loop[1] = true

        # look for next point (which is closest to the given point in all remaing ones)
        while length(loop) < length(point_indices)
            # check for the closest
            closest_distance = 1e32
            closest_point = -1
            # check all other points
            for i in 1:length(point_indices)
                # if already in loop, continue
                if in_loop[i]
                    continue
                end
                # compare to the nearest one
                if dot(points[point_indices[i]].-points[point_indices[loop[end]]], points[point_indices[i]].-points[point_indices[loop[end]]]) < closest_distance
                    closest_point    = i
                    closest_distance = dot(points[point_indices[i]].-points[point_indices[loop[end]]], points[point_indices[i]].-points[point_indices[loop[end]]])
                end
            end
            # security check
            if closest_point == -1
                println()
                println(closest_distance)
                println(points[in_plane])
            end
            # insert the closest point as the next in line
            push!(loop, closest_point)
            in_loop[closest_point] = true
        end

        # get back the real indices and reconstruct loop
        loop = point_indices[loop]
        # get the sorted
        loop_sorted = sort(loop)
        # push the starting point into the loop
        #push!(loop, loop[1])

        # loop not found yet
        loop_found = false
        # check if the loop already exists
        for l in edges_sorted
            if l==loop_sorted
                loop_found = true
                break
            end
        end
        # if the loop is not found, insert the loop into the edges list
        if !loop_found
            push!(edges, loop)
        end

    end

    ##########
    # STEP 4 - Construct a new object and return it
    ##########

    # return the finished BZ
    return newBrillouinZone(
        BZ,
        reciprocal_unitcell,
        points,
        edges
    )
end

################################################################################
#
#   COMPUTING BRILLOUIN ZONE OBJECTS
#   (Wrapper)
#
################################################################################

# wrapper for all reciprocal unitcells
function getBrillouinZone(
            reciprocal_unitcell :: RU
            ;
            kwargs...
        ) :: BrillouinZone{RU} where {L,N,D,P<:AbstractReciprocalPoint{D},B<:AbstractBond{L,N},RU<:AbstractReciprocalUnitcell{P,B}}

    # call the respective function
    return getBrillouinZone(BrillouinZone{RU}, reciprocal_unitcell; kwargs...)
end

# wrapper for all real space unitcells
function getBrillouinZone(
            unitcell :: U
            ;
            kwargs...
        ) :: BrillouinZone{ReciprocalUnitcell{ReciprocalPoint{D},B}} where {L,D,S<:AbstractSite{L,D},B,U<:AbstractUnitcell{S,B}}

    # call the respective function
    return getBrillouinZone(BrillouinZone{ReciprocalUnitcell{ReciprocalPoint{D},B}}, getReciprocalUnitcell(unitcell); kwargs...)
end

# Export the function
export getBrillouinZone

################################################################################
#
#   FUNCTIONS FOR HIGH SYMMETRY MESH
#
################################################################################

# contructor for symmetric mesh of given brillouin zone
function getPointsHighSymmetryMesh(
            bz              :: BZ,
            line_resolution :: Integer
            ;
            include_bz_corners :: Bool           = true,
            include_face_centers :: Bool         = true,
            include_bz_edges :: Bool             = true,  # edges of faces (i.e. lines connecting the corners)
            include_lines_in_faces :: Bool       = true,
            include_lines_to_face_center :: Bool = true,
            include_lines_to_gamma :: Bool       = true,
            include_random_face_points :: Bool   = true,
            face_resolution :: Integer = 100
        ) :: Vector{Vector{Float64}} where {D,B, P<:AbstractReciprocalPoint{D}, R<:AbstractReciprocalUnitcell{P,B}, BZ<:AbstractBrillouinZone{R}}

    # get the gamma point
    gamma = point(gammaPoint(reciprocalUnitcell(bz)))
    # the list of points that form the mesh
    mesh = Vector{Float64}[]

    # iterate over all faces of the BZ
    for f in faces(bz)

        # obtain the relevant corners
        k_corners = unique([corners(bz)[i] for i in f])
        # points that connect to gamma
        points_to_gamma = Vector{Float64}[]

        # INCLUDE CORNER POINTS IN MESH
        if include_bz_corners
            for c in k_corners
                # put the corner into the mesh
                push!(mesh, c)
                # put the corner into the list of points that connect to the gamma point
                push!(points_to_gamma, c)
            end
        end

        # INCLUDE FACE CENTERS IN MESH
        if include_face_centers
            # push the center into the mesh
            push!(mesh, getPointCenter(k_corners))
            # push the center into the list of connecting points
            push!(points_to_gamma, getPointCenter(k_corners))
        end

        # INCLUDE EDGES IN MESH
        if include_bz_edges
            # the edge that wraps in the point list
            line = getPointsOnLine(k_corners[end], k_corners[1], line_resolution)
            for p in line
                # put the point into the mesh
                push!(mesh, p)
            end
            # put the start, end and center of line into the gamma point connecting point list
            push!(points_to_gamma, k_corners[end])
            push!(points_to_gamma, k_corners[1])
            push!(points_to_gamma, getPointCenter(k_corners[1], k_corners[end]))
            # all other edges
            for i in 1:length(k_corners)-1
                # construct line
                line = getPointsOnLine(k_corners[i], k_corners[i+1], line_resolution)
                for p in line
                    # put the point into the mesh
                    push!(mesh, p)
                end
                # put the start, end and center of line into the gamma point connecting point list
                push!(points_to_gamma, k_corners[i])
                push!(points_to_gamma, k_corners[i+1])
                push!(points_to_gamma, getPointCenter(k_corners[i], k_corners[i+1]))
            end
        end

        # INCLUDE LINES IN THE FACE
        if include_lines_in_faces
            # iterate over all combination of points
            for i in 1:length(k_corners)
            for j in i+1:length(k_corners)
                # build the line from i to j
                line = getPointsOnLine(k_corners[i], k_corners[j], line_resolution)
                for p in line
                    push!(mesh, p)
                end
                # put the start, end and center of line into the gamma point connecting point list
                push!(points_to_gamma, k_corners[i])
                push!(points_to_gamma, k_corners[j])
                push!(points_to_gamma, getPointCenter(k_corners[i], k_corners[j]))
            end
            end
        end

        # INCLUDE LINES TO THE FACE CENTER
        if include_lines_to_face_center
            # build the face center
            center = getPointCenter(k_corners)
            # iterate over all points
            for i in 1:length(k_corners)
                # build the line from i to j
                line = getPointsOnLine(k_corners[i], center, line_resolution)
                for p in line
                    push!(mesh, p)
                end
            end
        end

        # INCLUDE LINES CONNECTING TO GAMMA POINT
        if include_lines_to_gamma
            # get a unique list of connecting points
            unique!(points_to_gamma)
            # build lines from these points to the gamma point (if they are not the gamma point)
            for ptg in points_to_gamma
                if dot(ptg.-gamma, ptg.-gamma) < 1e-8
                    continue
                end
                line = getPointsOnLine(ptg, gamma, line_resolution)
                for p in line
                    push!(mesh, p)
                end
            end
        end

        # INCLUDE FACE POINTS RANDOMLY SAMPLED
        if include_random_face_points
            points = getPointsRandomInConvexHull(face_resolution, k_corners...)
            for p in points
                push!(mesh, p)
            end
        end

    end

    # remove redundant points
    unique!(mesh)

    # return the mesh
    return mesh
end

export getPointsHighSymmetryMesh


function nonsym_ising_ham(;J = -1,spin = 1//2,lambda = 0.5,longit=0.0)
    (sx,sy,sz)=nonsym_spintensors(spin);
    id = one(sx);

    hamdat = Array{Union{Missing,typeof(sx)},3}(missing,1,3,3)
    hamdat[1,1,1] = id;
    hamdat[1,end,end] = id;
    hamdat[1,1,2] = J*sz;
    hamdat[1,2,end] = sz;
    hamdat[1,1,end] = lambda*sx+longit*sz;

    ham = MPOHamiltonian(hamdat);

    return ham
end

function nonsym_ising_mpo(;beta = log(1+sqrt(2))/2)
    t = [exp(beta) exp(-beta); exp(-beta) exp(beta)];

    r = eigen(t);
    nt = r.vectors*sqrt(Diagonal(r.values))*r.vectors;

    O = zeros(2,2,2,2);
    O[1,1,1,1]=1; O[2,2,2,2]=1;

    @tensor toret[-1 -2;-3 -4] := O[1 2;3 4]*nt[-1;1]*nt[-2;2]*nt[-3;3]*nt[-4;4];

    torett = TensorMap(complex(toret),ℂ^2*ℂ^2,ℂ^2*ℂ^2);
    return InfiniteMPO(torett);
end

function z2_ising_mpo(; beta = log(1+sqrt(2))/2)
  x = cosh(beta)
  y = sinh(beta)

  sec = ℤ₂Space(0 => 1, 1=> 1)
  mpo =  TensorMap(zeros, ComplexF64, sec * sec, sec * sec)
  blocks(mpo)[Irrep[ℤ₂](0)] = [2x^2 2x*y; 2x*y 2y^2]
  blocks(mpo)[Irrep[ℤ₂](1)] = [2x*y 2x*y; 2x*y 2x*y]

  return InfiniteMPO(mpo);
end