module DataEnvelopmentAnalysis

    """
        DataEnvelopmentAnalysis
    A Julia package for efficiency and productivity measurement using Data Envelopment Analysis (DEA).
    [DataEnvelopmentAnalysis repository](https://github.com/javierbarbero/DataEnvelopmentAnalysis.jl).
    """

    using JuMP
    using GLPK
    using Ipopt
    using SparseArrays
    using LinearAlgebra

    using Printf: @sprintf
    using Statistics: std
    using StatsBase: CoefTable

    import StatsBase: nobs, mean

    export
    # optimizer
    DEAOptimizer,
    newdeamodel,

    # Types
    AbstractDEAModel,
    DEAModel,
    AbstractDEAPeers, AbstractDEAPeersDMU,
    DEAPeers, DEAPeersDMU,
    AbstractTechnicalDEAModel, AbstractRadialDEAModel,
    TechnicalDEAModel,
    RadialDEAModel, AdditiveDEAModel, DirectionalDEAModel, GeneralizedDFDEAModel,
    RussellDEAModel, EnhancedRussellGraphDEAModel, ModifiedDDFDEAModel,
    AbstractHolderDEAModel, HolderL1DEAModel, HolderL2DEAModel, HolderLInfDEAModel,
    ReverseDDFDEAModel,
    AbstractEconomicDEAModel,
    AbstractCostDEAModel, AbstractRevenueDEAModel, AbstractProfitDEAModel, AbstractProfitabilityDEAModel,
    CostDEAModel, RevenueDEAModel, ProfitDEAModel, ProfitabilityDEAModel,
    AbstractProductivityDEAModel,
    MalmquistDEAModel,
    # All models
    nobs, ninputs, noutputs,
    # Peers
    peers, peersmatrix, ispeer,
    # Technical models
    dea, deaadd, deaaddweights, deaddf, deagdf,
    dearussell, deaerg, deamddf, deaholder, dearddf,
    efficiency, slacks,
    # Economic models
    deamincost, deamaxrevenue, deamaxprofit,
    deacost, dearevenue, deaprofit, deaprofitability,
    normfactor, ismonetary,
    # Common technical and economic models
    targets,
    # Productivity models
    malmquist,
    nperiods, prodchange

    # Include code of functions
    include("optimizer.jl")

    include("model.jl")
    include("peers.jl")

    include("technical.jl")
    include("dea.jl")
    include("deaadd.jl")
    include("deaddf.jl")
    include("deagdf.jl")
    include("dearussell.jl")
    include("deaerg.jl")
    include("deamddf.jl")
    include("deaholder.jl")
    include("dearddf.jl")

    include("economic.jl")
    include("econoptim.jl")
    include("deacost.jl")
    include("dearevenue.jl")
    include("deaprofit.jl")
    include("deaprofitability.jl")

    include("productivity.jl")
    include("malmquist.jl")

    function __init__()

        nothing
    end

end # module


using Test
using Plots
using OnlineStats: CovMatrix
using LyapunovExponents
using LyapunovExponents: VecVariance, with_input_noise, with_output_noise
using LyapunovExponents.TestTools: @test_nothrow

@time @testset "Smoke test CLV" begin
    @test_nothrow begin
        prob = LyapunovExponents.lorenz_63(t_attr=3).prob :: LEProblem
        solver = CLVSolver(prob)
        solve!(solver)
    end
    @testset "Recording CLV" begin
        prob = LyapunovExponents.lorenz_63(t_attr=3).prob :: LEProblem
        solver = CLVSolver(
            prob;
            record = [:G, :C, :D, :x])
        solve!(solver)
        @test isdefined(solver.sol, :G_history)
        @test isdefined(solver.sol, :C_history)
        @test isdefined(solver.sol, :D_history)
        @test isdefined(solver.sol, :x_history)
        @test unique(size.(solver.sol.G_history)) == [(3, 3)]
        @test unique(size.(solver.sol.C_history)) == [(3, 3)]
        @test unique(size.(solver.sol.D_history)) == [(3,)]
        @test unique(size.(solver.sol.x_history)) == [(3,)]
    end
end

@time @testset "Smoke test demo" begin
    @test begin
        demo = solve!(LyapunovExponents.lorenz_63(t_attr=3))
        plot(demo, show = false)
        true
    end
    @test begin
        demo = solve!(LyapunovExponents.lorenz_63(t_attr=3, dim_lyap=1))
        plot(demo, show = false)
        true
    end
    @testset "main_stat = $ST" for ST in [VecVariance, CovMatrix]
        demo = LyapunovExponents.lorenz_63(t_attr=3)
        @test_nothrow solve!(demo; main_stat = ST)
        sol = demo.solver.sol
        @test isa(sol.main_stat, ST)
        @test_nothrow plot(demo, show = false)
    end
end

@time @testset "Smoke test RODE" begin
    ode = LyapunovExponents.lorenz_63().prob.phase_prob
    rode = with_input_noise(ode, 1)
    @test_nothrow begin
        t_attr = 3
        prob = LEProblem(rode; t_attr=t_attr)
        sol = solve(prob; record=true, integrator_options=[:dt => 1e-3])
        plot(sol)
    end

    rode = with_output_noise(ode, 1)
    @test_nothrow begin
        t_attr = 3
        prob = LEProblem(rode; t_attr=t_attr)
        sol = solve(prob; record=true, integrator_options=[:dt => 1e-3])
        plot(sol)
    end
end


#####
##### Tests for basic manifold functionality
#####

types = [Vector{Float64},
         SizedVector{3, Float64},
         MVector{3, Float64},
         Vector{Float32},
         SizedVector{3, Float32},
         MVector{3, Float32}]

@testset "Euclidean" begin
    M = Euclidean(3)
    base_points = [randn(3) for _=1:3]
    t = rand()

    for T in types
        @testset "Type $T" begin
            x₀ = convert(T, base_points[1])
            v₀ = convert(T, base_points[2])
            p₀ = convert(T, base_points[3])

            x, p = cogeodesic_flow(M, x₀, p₀, t)
            @test x ≈ x₀ + t * p₀
            @test p ≈ p₀

            x, v = geodesic_flow(M, x₀, v₀, t)
            @test x ≈ x₀ + t * v₀
            @test v ≈ v₀
        end
    end
end

@testset "Sphere" begin
    M = Sphere(3)

    base_points = [normalize(randn(3))]
    Π = I - base_points[1] * base_points[1]'
    push!(base_points, normalize(Π * randn(3)))
    push!(base_points, normalize(Π * randn(3)))

    for T in types
        @testset "Type $T" begin
            x₀ = convert(T, base_points[1])
            v₀ = convert(T, base_points[2])
            p₀ = convert(T, base_points[3])
            t = eltype(T)(π)

            x, p = cogeodesic_flow(M, x₀, p₀, t)
            @test x ≈ -x₀
            @test p ≈ -p₀

            x, v = geodesic_flow(M, x₀, v₀, t)
            @test x ≈ -x₀
            @test v ≈ -v₀
        end
    end
end


################################################################################
#
#   METHODS FOR CONSTRUCTION OF FERMI SURFACES (2D & 3D)
#
#   STRUCTURE OF THE FILE
#
#   1) CALCULATION OF FERMI SURFACE
#
#   2) PLOTTING OF FERMI SURFACE
#       - plotting from points
#       - plotting from unitcell
#
################################################################################

# Function for calculating Fermi surface in 2D (not exported)
function getFermiSurface2D(
            unitcell::Unitcell,
            N_points::Int64;
            fermi_energy::Float64=0.0,
            enforce_hermitian::Bool=false,
            epsilon::Float64=1e-10,
            epsilon_k::Float64=1e-10,
            bounds_lower::Array{Float64,1}=-2*pi.*ones(2),
            bounds_upper::Array{Float64,1}=2*pi.*ones(2),
            refold_to_first_BZ::Bool=true
        )

    # list of k values that contribute to the BZ
    k_values = zeros(Float64, N_points, 2)

    # funtion of energy
    function energy(kvector)
        eigenvalues = eigvals(getInteractionMatrixKSpace(unitcell, kvector, enforce_hermitian=enforce_hermitian)) .- fermi_energy
        eigenvalues = eigenvalues .* eigenvalues
        return minimum(eigenvalues)
    end

    # the current search index
    index = 1
    # search until there are enough points
    while index <= N_points

        # find a suitable starting point for the Newton algorithm
        k = Float64[rand(), rand()]
        for j in 1:length(k)
            k[j] = k[j] * bounds_lower[j] + (1-k[j]) * bounds_upper[j]
        end

        # start with the initial energy
        e0 = energy(k)
        # iterate i over 100 newton steps (maximum)
        for i in 1:100
            # check if the energy is already converged
            if e0 < epsilon
                # save the k vector
                k_values[index,:] = k
                # increment the index
                index = index+1
                # break the newton loop
                break
            end
            # the current energy
            H_0 = e0
            # the gradient of the energy
            H_eps =
            [
                energy(k .+ [epsilon_k, 0]),
                energy(k .+ [0, epsilon_k])
            ]
            dH = (H_eps - H_0) ./ epsilon_k
            # absolute value of the gradient
            dHdH = dot(dH, dH)
            # break the Newton loop if the gradient diverges or is flattened too much
            if abs(dHdH) < 1e-20 || abs(dHdH) > 1e20
                break
            end
            # increment k
            dk = dH * (H_0 / dHdH)
            k -= dk
            # calculate a new energy
            e0 = energy(k)
        end
    end

    # if refolding is enabled, refold all points
    if refold_to_first_BZ
        # get the lattice vectors
        a1 = unitcell.lattice_vectors[1]
        a2 = unitcell.lattice_vectors[2]
        # get the reciprocal lattice vectors
        b1 = [a2[2], -a2[1]]
        b2 = [a1[2], -a1[1]]
        # normalize the vectors
        b1 .*= 2*pi/sum(b1.*a1)
        b2 .*= 2*pi/sum(b2.*a2)

        # iterate over all k points
        for index in 1:N_points
            # assumend to be refolded to begin with
            refolded = true
            # iterate while still refolding possible
            while refolded
                # new iteration, not refolded
                refolded = false
                # try all possible refoldings
                for r in [b1, b2, -b1, -b2, b1+b2, -b2-b1, b1-b2, -b2+b1]
                    if sum((k_values[index,:].-r).*(k_values[index,:].-r)) < sum((k_values[index,:]).*(k_values[index,:]))
                        k_values[index,:] .-= r
                        refolded=true
                    end
                end
            end
        end
    end

    # return the list
    return k_values

end

# Function for calculating Fermi surface in 3D (not exported)
function getFermiSurface3D(
            unitcell::Unitcell,
            N_points::Int64;
            fermi_energy::Float64=0.0,
            enforce_hermitian::Bool=false,
            epsilon::Float64=1e-10,
            epsilon_k::Float64=1e-10,
            bounds_lower::Array{Float64,1}=-2*pi.*ones(3),
            bounds_upper::Array{Float64,1}=2*pi.*ones(3),
            refold_to_first_BZ::Bool=true
        )

    # list of k values that contribute to the BZ
    k_values = zeros(Float64, N_points, 3)

    # funtion of energy
    function energy(kvector)
        eigenvalues = eigvals(getInteractionMatrixKSpace(unitcell, kvector, enforce_hermitian=enforce_hermitian)) .- fermi_energy
        eigenvalues = eigenvalues .* eigenvalues
        return minimum(eigenvalues)
    end

    # the current search index
    index = 1
    # search until there are enough points
    while index <= N_points

        # find a suitable starting point for the Newton algorithm
        k = Float64[rand(), rand(), rand()]
        for j in 1:length(k)
            k[j] = k[j] * bounds_lower[j] + (1-k[j]) * bounds_upper[j]
        end

        # start with the initial energy
        e0 = energy(k)
        # iterate i over 100 newton steps (maximum)
        for i in 1:100
            # check if the energy is already converged
            if e0 < epsilon
                # save the k vector
                k_values[index,:] = k
                # increment the index
                index = index+1
                # break the newton loop
                break
            end
            # the current energy
            H_0 = e0
            # the gradient of the energy
            H_eps =
            [
                energy(k .+ [epsilon_k, 0, 0]),
                energy(k .+ [0, epsilon_k, 0]),
                energy(k .+ [0, 0, epsilon_k])
            ]
            dH = (H_eps - H_0) ./ epsilon_k
            # absolute value of the gradient
            dHdH = dot(dH, dH)
            # break the Newton loop if the gradient diverges or is flattened too much
            if abs(dHdH) < 1e-20 || abs(dHdH) > 1e20
                break
            end
            # increment k
            dk = dH * (H_0 / dHdH)
            k -= dk
            # calculate a new energy
            e0 = energy(k)
        end
    end

    # if refolding is enabled, refold all points
    if refold_to_first_BZ
        # get the lattice vectors
        a1 = unitcell.lattice_vectors[1]
        a2 = unitcell.lattice_vectors[2]
        a3 = unitcell.lattice_vectors[3]
        # get the reciprocal lattice vectors
        b1 = cross(a2, a3)
        b2 = cross(a3, a1)
        b3 = cross(a1, a2)
        # normalize the vectors
        b1 .*= 2*pi/sum(b1.*a1)
        b2 .*= 2*pi/sum(b2.*a2)
        b3 .*= 2*pi/sum(b3.*a3)
        # put together refolding
        refoldings = Array{Float64,1}[]
        for i in [-1,0,1]
        for j in [-1,0,1]
        for l in [-1,0,1]
            refolding = i*b1 + j*b2 + l*b3
            if sum(abs.(refolding))>1e-8
                push!(refoldings, refolding)
            end
        end
        end
        end
        # iterate over all k points
        for index in 1:N_points
            # assumend to be refolded to begin with
            refolded = true
            # iterate while still refolding possible
            while refolded
                # new iteration, not refolded
                refolded = false
                # try all possible refoldings
                for r in refoldings
                    if sum((k_values[index,:].-r).*(k_values[index,:].-r)) < sum((k_values[index,:]).*(k_values[index,:]))
                        k_values[index,:] .-= r
                        refolded=true
                    end
                end
            end
        end
    end

    # return the list
    return k_values

end

# obtain Fermi surface
"""
    getFermiSurface(
                unitcell::Unitcell,
                N_points::Int64
             [; fermi_energy::Float64=0.0,
                enforce_hermitian::Bool=false,
                epsilon::Float64=1e-10,
                epsilon_k::Float64=1e-10,
                bounds_lower::Array{Float64,1}=-2*pi.*ones(4),
                bounds_upper::Array{Float64,1}=2*pi.*ones(4),
                refold_to_first_BZ::Bool=true ]
            )

calculates `N_points` points which belong to the Fermi surface of the tight-binding model given by
the `Unitcell` object `unitcell`. The Fermi energy can be adjusted by using `fermi_energy=...`.

The points are located using Newton's method with random starting positions. The procedure can be
modified by using some of the optional keywords:
- `enforce_hermitian` determines if the matrix used for energy calculation is made hermitian by construction
- `epsilon` energy threshold for values to be considered close to the Fermi energy
- `epsilon_k` small k increment that is used for calculation of derivatives
- `bounds_upper` / `bounds_lower` the bounds between which the random starting location for the Newton method is searched

After the search has been completed, the values will optionally be refolded into the first brillouin zone by choosing `refold_to_first_BZ=true`.

# Examples

```julia-repl
julia> fermi_surace = getFermiSurface(unitcell, 100)    # 2D unitcell
100x2 Array{Float64}
...

julia> fermi_surace = getFermiSurface(unitcell, 100)    # 3D unitcell
100x3 Array{Float64}
...

julia> fermi_surace = getFermiSurface(unitcell, 100, fermi_energy=1.0)
100x2 Array{Float64}
...

julia> fermi_surace = getFermiSurface(unitcell, 100, refold_to_first_BZ=false)
100x2 Array{Float64}
...

```
"""
function getFermiSurface(
            unitcell::Unitcell,
            N_points::Int64;
            fermi_energy::Float64=0.0,
            enforce_hermitian::Bool=false,
            epsilon::Float64=1e-10,
            epsilon_k::Float64=1e-10,
            bounds_lower::Array{Float64,1}=-2*pi.*ones(3),
            bounds_upper::Array{Float64,1}=2*pi.*ones(3),
            refold_to_first_BZ::Bool=true
        )
    # check if the unitcell can be put in either method
    if length(unitcell.basis[1]) != length(unitcell.lattice_vectors)
        println("Unitcell has not the same number of lattice vectors as dimensions")
        return zeros(1,1)
    end
    # check which function to pass to
    if length(unitcell.lattice_vectors) == 2
        # 2D case
        return getFermiSurface2D(
                unitcell,
                N_points,
                fermi_energy=fermi_energy,
                enforce_hermitian=enforce_hermitian,
                epsilon=epsilon,
                epsilon_k=epsilon_k,
                bounds_lower=bounds_lower,
                bounds_upper=bounds_upper,
                refold_to_first_BZ=refold_to_first_BZ
            )
    else length(unitcell.lattice_vectors) == 3
        # 3D case
        return getFermiSurface3D(
                unitcell,
                N_points,
                fermi_energy=fermi_energy,
                enforce_hermitian=enforce_hermitian,
                epsilon=epsilon,
                epsilon_k=epsilon_k,
                bounds_lower=bounds_lower,
                bounds_upper=bounds_upper,
                refold_to_first_BZ=refold_to_first_BZ
            )
    end

end
export getFermiSurface

# plot the Fermi surface (2D)
function plotFermiSurface2D(
            k_values::Array{Float64, 2};
            brillouin_zone::BrillouinZone=BrillouinZone(),
            plot_title::String="",
            plot_color="b",
            figsize::Tuple=(6,6),
            showPlot::Bool=true,
            save_filename::String="NONE"
        )

    ###########################
    #   INITIAL SETTINGS
    ###########################

    # configure plot environment
    rc("font", family="serif")

    # create a new figure
    fig = figure(figsize=figsize)

    ###########################
    #   PLOT FERMI SURFACE
    ###########################

    # if brillouin zone not empty, plot it as well
    if length(brillouin_zone.points) > 0
        plotBrillouinZone(brillouin_zone, new_figure=false)
    end

    # scatter the points
    scatter(k_values[:,1], k_values[:,2], color=plot_color)

    ###########################
    #   CONFIGURE AXIS & TITLE
    ###########################

    # set the title
    if plot_title == "AUTO"
        # set the title to an automatically generated title
        title("Fermi surface")
    elseif plot_title == ""
        # do nothing title related
    else
        # set the title to the given title
        title(plot_title)
    end

    # equal axis
    gca()[:set_aspect]("equal")

    ###########################
    #   FINISH THE PLOT
    ###########################

    # tighten the layout
    tight_layout()

    # save the plot
    if save_filename != "NONE"
        # make sure the directory exists
        if contains(save_filename, "/")
    		# get the containing folder
    		folder = save_filename[1:findlast(save_filename, '/')]
    		# build the path to that folder
    		mkpath(folder)
    	end
        # save the plot
        savefig(save_filename)
    end

    # maybe show the plot
    if showPlot
        show()
    end

    # return the figure object
    return fig

end
function plotFermiSurface2D(
            unitcell::Unitcell,
            N_points::Int64;
            fermi_energy::Float64=0.0,
            enforce_hermitian::Bool=false,
            epsilon::Float64=1e-10,
            epsilon_k::Float64=1e-10,
            bounds_lower::Array{Float64,1}=-2*pi.*ones(4),
            bounds_upper::Array{Float64,1}=2*pi.*ones(4),
            refold_to_first_BZ::Bool=true,
            brillouin_zone::BrillouinZone=BrillouinZone(),
            plot_title::String="",
            plot_color="b",
            figsize::Tuple=(6,6),
            showPlot::Bool=true,
            save_filename::String="NONE"
        )

    # check if BZ should be calculated
    if length(brillouin_zone.points) == 0
        brillouin_zone = createBrillouinZone(unitcell)
    end

    # calculate the Fermi surface first
    fermi_surface = getFermiSurface2D(
            unitcell,
            N_points,
            fermi_energy=fermi_energy,
            enforce_hermitian=enforce_hermitian,
            epsilon=epsilon,
            epsilon_k=epsilon_k,
            bounds_lower=bounds_lower,
            bounds_upper=bounds_upper,
            refold_to_first_BZ=refold_to_first_BZ
        )

    # plot the fermi surface
    plotFermiSurface2D(
            fermi_surface,
            brillouin_zone=brillouin_zone,
            plot_title=plot_title,
            plot_color=plot_color,
            figsize=figsize,
            showPlot=showPlot,
            save_filename=save_filename
        )

end

# plot the Fermi surface (3D)
function plotFermiSurface3D(
            k_values::Array{Float64, 2};
            brillouin_zone::BrillouinZone=BrillouinZone(),
            plot_title::String="",
            plot_color="b",
            figsize::Tuple=(6,6),
            showPlot::Bool=true,
            save_filename::String="NONE"
        )

    ###########################
    #   INITIAL SETTINGS
    ###########################

    # configure plot environment
    rc("font", family="serif")

    # create a new figure
    fig = figure(figsize=figsize)

    ###########################
    #   PLOT FERMI SURFACE
    ###########################

    # if brillouin zone not empty, plot it as well
    if length(brillouin_zone.points) > 0
        plotBrillouinZone(brillouin_zone, new_figure=false)
    end

    # scatter the points
    scatter3D(k_values[:,1], k_values[:,2], k_values[:,3], color=plot_color)

    ###########################
    #   CONFIGURE AXIS & TITLE
    ###########################

    # set the title
    if plot_title == "AUTO"
        # set the title to an automatically generated title
        title("Fermi surface")
    elseif plot_title == ""
        # do nothing title related
    else
        # set the title to the given title
        title(plot_title)
    end

    # equal axis
    gca()[:set_aspect]("equal")

    ###########################
    #   FINISH THE PLOT
    ###########################

    # tighten the layout
    tight_layout()

    # save the plot
    if save_filename != "NONE"
        # make sure the directory exists
        if contains(save_filename, "/")
    		# get the containing folder
    		folder = save_filename[1:findlast(save_filename, '/')]
    		# build the path to that folder
    		mkpath(folder)
    	end
        # save the plot
        savefig(save_filename)
    end

    # maybe show the plot
    if showPlot
        show()
    end

    # return the figure object
    return fig

end
function plotFermiSurface3D(
            unitcell::Unitcell,
            N_points::Int64;
            fermi_energy::Float64=0.0,
            enforce_hermitian::Bool=false,
            epsilon::Float64=1e-10,
            epsilon_k::Float64=1e-10,
            bounds_lower::Array{Float64,1}=-2*pi.*ones(3),
            bounds_upper::Array{Float64,1}=2*pi.*ones(3),
            refold_to_first_BZ::Bool=true,
            brillouin_zone::BrillouinZone=BrillouinZone(),
            plot_title::String="",
            plot_color="b",
            figsize::Tuple=(6,6),
            showPlot::Bool=true,
            save_filename::String="NONE"
        )

    # check if BZ should be calculated
    if length(brillouin_zone.points) == 0
        brillouin_zone = createBrillouinZone(unitcell)
    end

    # calculate the Fermi surface first
    fermi_surface = getFermiSurface3D(
            unitcell,
            N_points,
            fermi_energy=fermi_energy,
            enforce_hermitian=enforce_hermitian,
            epsilon=epsilon,
            epsilon_k=epsilon_k,
            bounds_lower=bounds_lower,
            bounds_upper=bounds_upper,
            refold_to_first_BZ=refold_to_first_BZ
        )

    # plot the fermi surface
    plotFermiSurface3D(
            fermi_surface,
            brillouin_zone=brillouin_zone,
            plot_title=plot_title,
            plot_color=plot_color,
            figsize=figsize,
            showPlot=showPlot,
            save_filename=save_filename
        )

end

"""
    plotFermiSurface(
                unitcell::Unitcell,
                N_points::Int64
             [; fermi_energy::Float64=0.0,
                enforce_hermitian::Bool=false,
                epsilon::Float64=1e-10,
                epsilon_k::Float64=1e-10,
                bounds_lower::Array{Float64,1}=-2*pi.*ones(4),
                bounds_upper::Array{Float64,1}=2*pi.*ones(4),
                refold_to_first_BZ::Bool=true,
                brillouin_zone::BrillouinZone=BrillouinZone(),
                plot_title::String="",
                plot_color="b",
                figsize::Tuple=(6,6),
                showPlot::Bool=true,
                save_filename::String="NONE" ]
            )

calculates `N_points` points which belong to the Fermi surface of the tight-binding model given by
the `Unitcell` object `unitcell` and plots them using `PyPlot`. The Fermi energy can be adjusted by using `fermi_energy=...`.

The points are located using Newton's method with random starting positions. The procedure can be
modified by using some of the optional keywords:
- `enforce_hermitian` determines if the matrix used for energy calculation is made hermitian by construction
- `epsilon` energy threshold for values to be considered close to the Fermi energy
- `epsilon_k` small k increment that is used for calculation of derivatives
- `bounds_upper` / `bounds_lower` the bounds between which the random starting location for the Newton method is searched

After the search has been completed, the values will optionally be refolded into the first brillouin zone by choosing `refold_to_first_BZ=true`.

For plotting, several more options can be used.

# Examples

```julia-repl
julia> plotFermiSurface(unitcell, 100)
PyPlot.Figure(...)

julia> plotFermiSurface(unitcell, 100, fermi_energy=1.0, showPlot=false)
PyPlot.Figure(...)

```
"""
function plotFermiSurface(
            k_values::Array{Float64, 2};
            brillouin_zone::BrillouinZone=BrillouinZone(),
            plot_title::String="",
            plot_color="b",
            figsize::Tuple=(6,6),
            showPlot::Bool=true,
            save_filename::String="NONE"
        )

    # check which function to pass to
    if size(k_values,2) == 2
        # 2D case
        return plotFermiSurface2D(
                k_values,
                brillouin_zone=brillouin_zone,
                plot_title=plot_title,
                plot_color=plot_color,
                figsize=figsize,
                showPlot=showPlot,
                save_filename=save_filename
            )
    else size(k_values,2) == 3
        # 3D case
        return plotFermiSurface3D(
                k_values,
                brillouin_zone=brillouin_zone,
                plot_title=plot_title,
                plot_color=plot_color,
                figsize=figsize,
                showPlot=showPlot,
                save_filename=save_filename
            )
    end

end
function plotFermiSurface(
            unitcell::Unitcell,
            N_points::Int64;
            fermi_energy::Float64=0.0,
            enforce_hermitian::Bool=false,
            epsilon::Float64=1e-10,
            epsilon_k::Float64=1e-10,
            bounds_lower::Array{Float64,1}=-2*pi.*ones(4),
            bounds_upper::Array{Float64,1}=2*pi.*ones(4),
            refold_to_first_BZ::Bool=true,
            brillouin_zone::BrillouinZone=BrillouinZone(),
            plot_title::String="",
            plot_color="b",
            figsize::Tuple=(6,6),
            showPlot::Bool=true,
            save_filename::String="NONE"
        )

    # check if the unitcell can be put in either method
    if length(unitcell.basis[1]) != length(unitcell.lattice_vectors)
        println("Unitcell has not the same number of lattice vectors as dimensions")
        return zeros(1,1)
    end
    # check if BZ should be calculated
    if length(brillouin_zone.points) == 0
        brillouin_zone = createBrillouinZone(unitcell)
    end
    # check which function to pass to
    if length(unitcell.lattice_vectors) == 2
        # 2D case
        return plotFermiSurface2D(
                unitcell,
                N_points,
                fermi_energy=fermi_energy,
                enforce_hermitian=enforce_hermitian,
                epsilon=epsilon,
                epsilon_k=epsilon_k,
                bounds_lower=bounds_lower,
                bounds_upper=bounds_upper,
                refold_to_first_BZ=refold_to_first_BZ,
                brillouin_zone=brillouin_zone,
                plot_title=plot_title,
                plot_color=plot_color,
                figsize=figsize,
                showPlot=showPlot,
                save_filename=save_filename
            )
    else length(unitcell.lattice_vectors) == 3
        # 3D case
        return plotFermiSurface3D(
                unitcell,
                N_points,
                fermi_energy=fermi_energy,
                enforce_hermitian=enforce_hermitian,
                epsilon=epsilon,
                epsilon_k=epsilon_k,
                bounds_lower=bounds_lower,
                bounds_upper=bounds_upper,
                refold_to_first_BZ=refold_to_first_BZ,
                brillouin_zone=brillouin_zone,
                plot_title=plot_title,
                plot_color=plot_color,
                figsize=figsize,
                showPlot=showPlot,
                save_filename=save_filename
            )
    end

end
export plotFermiSurface


module AppendedTriangulationsTests

using Test
using Gridap.ReferenceFEs
using Gridap.Arrays
using Gridap.Geometry
using Gridap.Fields
using Gridap.Integration
using LinearAlgebra: ⋅
using Gridap.CellData
using Gridap.FESpaces
using Gridap.ReferenceFEs
using FillArrays

domain = (0,1,0,1)
partition = (10,10)
model = CartesianDiscreteModel(domain,partition)

ncells = num_cells(model)
nin = ceil(Int,2*ncells/3)
cell_to_mask = fill(false,ncells)
cell_to_mask[1:nin] .= true

grid = get_grid(model)

Ω_in = RestrictedTriangulation(grid,cell_to_mask)
Ω_out = RestrictedTriangulation(grid,.! cell_to_mask)
Ω = lazy_append(Ω_out,Ω_in)
test_triangulation(Ω)

order = 1
degree = 2*order
quad_in = CellQuadrature(Ω_in,degree)
quad_out = CellQuadrature(Ω_out,degree)
quad = CellQuadrature(Ω,degree)
@test isa(quad.trian,AppendedTriangulation)
@test isa(quad.cell_quad,AppendedArray)
@test isa(quad.cell_point,AppendedArray)
@test isa(quad.cell_weight,AppendedArray)

V = TestFESpace(model,ReferenceFE(lagrangian,Float64,order),conformity=:H1)

v(x) = x[1]+x[2]

vh = interpolate(v,V)

e = v - vh
el2 = sqrt(sum(integrate(e*e,quad)))
@test el2 < 1.0e-8

x = get_cell_points(quad)

dv = get_cell_shapefuns(V)
du = get_cell_shapefuns_trial(V)
@test isa(dv(x),AppendedArray)
@test isa(∇(dv)(x),AppendedArray)
@test isa(du(x),AppendedArray)
@test isa(∇(du)(x),AppendedArray)

cellmat = integrate( ∇(dv)⋅∇(du), quad )

@test isa(cellmat,AppendedArray)
@test isa(cellmat.a,Fill)
@test isa(cellmat.b,Fill)

end # module


function jt03()
#
#  test polynomial suggested by Jenkins and Traub
#
    z = 1 ./ exp10.(1:8)
    p = poly(z)
    z = [z ones(8)]
    p, PolyZeros(z)
end


"""
    Calculates ECI Position of SV

    $SIGNATURES
    ´sat_state´: satellite state, combining decoded data, code- and carrierphase

    This function calculates the position of the SV in ECI coorinates.
    The implementation follows IS-GPS-200K
"""
function sat_position_ECI(sat_state::SatelliteState)
    is_sat_healthy_and_decodable(sat_state.decoder_state) || throw(BadData("SV not decoded properly"))

    t = calc_corrected_time(sat_state)
    tₖ = t - sat_state.decoder_state.data.t_oe
    tₖ = check_crossover(tₖ)

    # Excentric and true anomaly
    tₛᵥ = calc_uncorrected_time(sat_state)
    dtₛᵥ = code_phase_offset(sat_state.decoder_state, tₛᵥ)
    Eₖ = calc_eccentric_anomaly(sat_state.decoder_state, t, dtₛᵥ)

    e  = sat_state.decoder_state.data.e
    vₖ =  2 * atan(sqrt((1 + e) / (1 - e)) * tan(Eₖ / 2))
    # Argument of latitude
    Φₖ = vₖ + sat_state.decoder_state.data.ω

    # Second harmonic perturbations
    δuₖ = sat_state.decoder_state.data.C_us * sin(2 * Φₖ) + sat_state.decoder_state.data.C_uc * cos(2 * Φₖ);
    δrₖ = sat_state.decoder_state.data.C_rs * sin(2 * Φₖ) + sat_state.decoder_state.data.C_rc * cos(2 * Φₖ);
    δiₖ = sat_state.decoder_state.data.C_is * sin(2 * Φₖ) + sat_state.decoder_state.data.C_ic * cos(2 * Φₖ);

    # Corrected orbit parameters
    uₖ = Φₖ + δuₖ;
    rₖ = sat_state.decoder_state.data.sqrt_A^2 * (1 - e * cos(Eₖ)) + δrₖ;
    iₖ = sat_state.decoder_state.data.i_0 + δiₖ + sat_state.decoder_state.data.IDOT * tₖ;

    # Positions in orbital plane
    xₖs = rₖ * cos(uₖ);
    yₖs = rₖ * sin(uₖ);

    # Corrected longitude of ascending node
    Ωₖ = sat_state.decoder_state.data.Ω_0 + sat_state.decoder_state.data.Ω_dot * tₖ;

    # ECI coordinates
    sinΩₖ = sin(Ωₖ);
    cosΩₖ = cos(Ωₖ);
    cosiₖ = cos(iₖ);
    siniₖ = sin(iₖ);

    xₖ = xₖs * cosΩₖ - yₖs * cosiₖ * sinΩₖ;
    yₖ = xₖs * sinΩₖ + yₖs * cosiₖ * cosΩₖ;
    zₖ = yₖs * siniₖ;

    return [xₖ, yₖ, zₖ]
end

"""
    Calculates ECEF Position by converting ECI position

    $SIGNATURES
    ´sat_state´: satellite state, combining decoded data, code- and carrierphase

    This function calculates the position of the SV in ECI coorinates and
    converts this position into ECEF coordinates.
    The implementation follows IS-GPS-200K
"""
function sat_position_ECI_2_ECEF(sat_state::SatelliteState)
    pos_ECI = sat_position_ECI(sat_state)

    tₖ = calc_corrected_time(sat_state)
    tₖ = tₖ - sat_state.decoder_state.data.t_oe
    tₖ = check_crossover(tₖ)
    θ = sat_state.decoder_state.constants.Ω_dot_e * (tₖ + sat_state.decoder_state.data.t_oe)
    cosθ  = cos(θ)
    sinθ  = sin(θ)
    R = [cosθ sinθ; -sinθ cosθ]
    pos_ECI[1:2] = R * pos_ECI[1:2]
    return pos_ECI
end

"""
    Calculates ECEF Position of SV

    $SIGNATURES
    ´sat_state´: satellite state, combining decoded data, code- and carrierphase

    This function calculates the position of the SV in ECEF coorinates.
    The implementation follows IS-GPS-200K Table 20-IV.
"""
function sat_position_ECEF(sat_state::SatelliteState)

    is_sat_healthy_and_decodable(sat_state.decoder_state) || throw(BadData("SV not decoded properly"))
    F = sat_state.decoder_state.constants.F
    e = sat_state.decoder_state.data.e
    μ = sat_state.decoder_state.constants.μ
    Ω_dot_e = sat_state.decoder_state.constants.Ω_dot_e

    A = sat_state.decoder_state.data.sqrt_A^2
    n₀ = sqrt(μ / (A^3))
    n = n₀ + sat_state.decoder_state.data.Δn

    t = calc_corrected_time(sat_state)
    tₖ = t - sat_state.decoder_state.data.t_oe
    tₖ = check_crossover(tₖ)

    # Excentric and true anomaly
    t_sv = calc_uncorrected_time(sat_state)
    dt_sv = code_phase_offset(sat_state.decoder_state, t_sv)
    E = calc_eccentric_anomaly(sat_state.decoder_state, t, dt_sv)

    vₖ = 2 * atan(sqrt((1 + e) / (1 - e)) * tan(E / 2))

    Φₖ = vₖ + sat_state.decoder_state.data.ω

    δuₖ = sat_state.decoder_state.data.C_us * sin(2 * Φₖ) + sat_state.decoder_state.data.C_uc * cos(2 * Φₖ)
    δrₖ = sat_state.decoder_state.data.C_rs * sin(2 * Φₖ) + sat_state.decoder_state.data.C_rc * cos(2 * Φₖ)
    δiₖ = sat_state.decoder_state.data.C_is * sin(2 * Φₖ) + sat_state.decoder_state.data.C_ic * cos(2 * Φₖ)

    uₖ = Φₖ + δuₖ
    rₖ = A * (1 - e * cos(E)) + δrₖ
    iₖ = sat_state.decoder_state.data.i_0 + δiₖ + sat_state.decoder_state.data.IDOT * tₖ

    xₖs = rₖ * cos(uₖ)
    yₖs = rₖ * sin(uₖ)

    Ωₖ = sat_state.decoder_state.data.Ω_0 +
        (sat_state.decoder_state.data.Ω_dot - Ω_dot_e) * tₖ -
        sat_state.decoder_state.data.t_oe * Ω_dot_e

    xₖ = xₖs * cos(Ωₖ) - yₖs * cos(iₖ) * sin(Ωₖ)
    yₖ = xₖs * sin(Ωₖ) + yₖs * cos(iₖ) * cos(Ωₖ)
    zₖ = yₖs * sin(iₖ)

    position = [xₖ, yₖ, zₖ]
    return position
end

"""
    Checks if satellite data is correctly tranmitted
    $SIGNATURES
    ´decoder_state´: Decoder

    Checks if satellite data contains the needed information and
    if errors during decoding occured
"""
function is_sat_healthy_and_decodable(decoder_state::GNSSDecoderState)
    decoder_state.data.svhealth == "000000" ? decoder_state.subframes_decoded[1:3] == [1,1,1] : false
end


include("spherical_helper.jl")

"""
    power_spectrum_1d(::AtmosGCMConfigType, var_grid, z, lat, lon, weight)

Calculates the 1D (zonal) power spectra using the fourier transform at each latitude and level from a 3D velocity field.
The snapshots of these spectra should be averaged to obtain a time-average.
The input velocities must be interpolated to a Gaussian grid.

# References
- [Wiin1967](@cite)
- [Koshyk2001](@cite)

# Arguments
- var_grid: variable (typically u or v) on a Gausian (lon, lat, z) grid to be transformed
- z: vertical coordinate (height or pressure)
- lat: latitude
- lon: longitude
- mass_weight: for mass-weighted calculations
"""
function power_spectrum_1d(::AtmosGCMConfigType, var_grid, z, lat, lon, weight)
    num_lev = length(z)
    num_lat = length(lat)
    num_lon = length(lon)
    num_fourier = Int64(num_lon)

    # get number of positive Fourier coefficients incl. 0
    if mod(num_lon, 2) == 0 # even
        num_pfourier = div(num_lon, 2)
    else # odd
        num_pfourier = div(num_lon, 2) + 1
    end

    zon_spectrum = zeros(Float64, num_pfourier, num_lat, num_lev)
    freqs = zeros(Float64, num_pfourier, num_lat, num_lev)

    for k in 1:num_lev
        for j in 1:num_lat
            # compute fft frequencies for each latitude
            x = lon ./ 180.0 .* π
            dx = (lon[2] - lon[1]) ./ 180.0 .* π

            freqs_ = fftfreq(num_fourier, 1.0 / dx) # 0,+ve freq,-ve freqs (lowest to highest)
            freqs[:, j, k] = freqs_[1:num_pfourier] .* 2.0 .* π

            # compute the fourier coefficients for all latitudes
            fourier = fft(var_grid[:, j, k]) # e.g. vcos_grid, ucos_grid
            fourier = (fourier / num_fourier)

            # convert to energy spectra
            zon_spectrum[1, j, k] =
                zon_spectrum[1, j, k] +
                weight[k] * fourier[1] .* conj(fourier[1])

            for m in 2:num_pfourier
                zon_spectrum[m, j, k] =
                    zon_spectrum[m, j, k] +
                    2.0 * weight[k] * fourier[m] * conj(fourier[m]) # factor 2 for neg freq contribution
            end
        end
    end
    return zon_spectrum, freqs
end

"""
    power_spectrum_2d(::AtmosGCMConfigType, var_grid, mass_weight)

- transform variable on grid to the 2d spectral space using fft on latitude circles
(as for the 1D spectrum) and Legendre polynomials for meridians, and calculate spectra

# Arguments
- var_grid: variable (typically u or v) on a Gausian (lon, lat, z) grid to be transformed
- mass_weight: weight for mass-weighted calculations

# References
- [Baer1972](@cite)
"""
function power_spectrum_2d(::AtmosGCMConfigType, var_grid, mass_weight)
    #  initialize spherical mesh variables
    nθ, nd = (size(var_grid, 2), size(var_grid, 3))
    mesh = SpectralSphericalMesh(nθ, nd)
    var_spectrum = mesh.var_spectrum
    var_spherical = mesh.var_spherical

    sinθ, wts = compute_gaussian!(mesh.nθ) # latitude weights using Gaussian quadrature, to orthogonalize Legendre polynomials upon summation
    mesh.qnm =
        compute_legendre!(mesh.num_fourier, mesh.num_spherical, sinθ, mesh.nθ) #  normalized associated Legendre polynomials

    for k in 1:(mesh.nd)
        # apply Gaussian quadrature weights
        for i in 1:(mesh.nθ)
            mesh.qwg[:, :, i] .= mesh.qnm[:, :, i] * wts[i] * mass_weight[k]
        end

        # Transform variable using spherical harmonics
        var_spherical[:, :, k, :] =
            trans_grid_to_spherical!(mesh, var_grid[:, :, k]) # var_spherical[m,n,k,sinθ]

        # Calculate energy spectra
        var_spectrum[:, :, k] =
            2.0 .* sum(var_spherical[:, :, k, :], dims = 3) .*
            conj(sum(var_spherical[:, :, k, :], dims = 3))  # var_spectrum[m,n,k] # factor 2 to account for negative Fourier frequencies
        var_spectrum[1, :, k] = var_spectrum[1, :, k] ./ 2.0 # m=0
    end
    return var_spectrum, mesh.wave_numbers, var_spherical, mesh
end

# TODO: enable mass weighted and vertically integrated calculations


# Gaussian drift proposal
@gen function gaussian_drift(tr::Trace, params=Dict(:β => (0.3, 0.01, Inf)))
    for (addr, distargs) in params
        if distargs isa Tuple && length(distargs) == 3
            sigma, lb, ub = distargs
            {addr} ~ trunc_normal(tr[addr], sigma, lb, ub)()
        else
            sigma = distargs
            {addr} ~ normal(tr[addr], sigma)()
        end
    end
end

# Guess beta from the data from initial infection counts
@gen function init_beta(obs::ChoiceMap, t::Int)
    params = create_params()
    cities = 1:length(params[:Ns])
    infected_0 = sum(params[:Is])
    infected_t = sum(obs[:step => t => :obs => :infected => c] for c in cities)
    β_hat = (infected_t/infected_0) ^ (1/t) - 1
    β ~ trunc_normal(β_hat, 0.1, 0, Inf)()
end


"""
Variational inference for Bayesian DimensionalityReduction
"""
module DimensionalityReduction

using Distributions
#using ProgressMeter

export DRModel
export sample_data, VI

####################
## Types
struct DRModel
    D::Int
    M::Int
    sigma2_y::Float64
    m_W::Array{Float64, 2} # MxD
    Sigma_W::Array{Float64, 3} # MxMxD
    m_mu::Array{Float64, 1} # D
    Sigma_mu::Array{Float64, 2} # DxD
end

####################
## functions
function sqsum(mat::Array{Float64}, idx::Int)
    return squeeze(sum(mat, idx), idx)
end

"""
Sample data given hyperparameters.
"""
function sample_data(N::Int, model::DRModel)
    D = model.D
    M = model.M
    W = zeros(M, D)
    mu = zeros(D)
    for d in 1 : D
        W[:,d] = rand(MvNormal(model.m_W[:,d], model.Sigma_W[:,:,d]))
    end
    mu = rand(MvNormal(model.m_mu, model.Sigma_mu))

    Y = zeros(D, N)
    X = randn(M, N)
    for n in 1 : N
        Y[:,n] = rand(MvNormal(W'*X[:,n] + mu, model.sigma2_y*eye(D)))
    end
    return Y, X, W, mu
end

function init(Y::Array{Float64, 2}, prior::DRModel)
    M = prior.M
    D, N = size(Y)
    X = randn(M, N)
    XX = zeros(M, M, N)
    for n in 1 : N
        XX[:,:,n] = X[:,n]*X[:,n]' + eye(M)
    end
    return X, XX
end

function update_W(Y::Array{Float64, 2}, prior::DRModel, posterior::DRModel,
                  X::Array{Float64, 2}, XX::Array{Float64, 3})
    D = prior.D
    M = prior.M
    N = size(Y, 2)
    m_W = zeros(M, D)
    Sigma_W = zeros(M, M, D)
    mu = posterior.m_mu
    for d in 1 : D
        Sigma_W[:,:,d] = inv(inv(prior.sigma2_y)*sqsum(XX, 3) + inv(prior.Sigma_W[:,:,d]))
        m_W[:,d] = Sigma_W[:,:,d]*(inv(prior.sigma2_y)*X*(Y[[d],:] - mu[d]*ones(1, N))'
                                   + inv(prior.Sigma_W[:,:,d])*prior.m_W[:,d])
    end
    return DRModel(D, M, prior.sigma2_y, m_W, Sigma_W, posterior.m_mu, posterior.Sigma_mu)
end

function update_mu(Y::Array{Float64, 2}, prior::DRModel, posterior::DRModel,
                   X::Array{Float64, 2}, XX::Array{Float64, 3})
    N = size(Y, 2)
    D = prior.D
    M = prior.M
    W = posterior.m_W
    Sigma_mu = inv(N*inv(prior.sigma2_y)*eye(D) + inv(prior.Sigma_mu))
    m_mu = Sigma_mu*(inv(prior.sigma2_y)*sqsum(Y - W'*X, 2) + inv(prior.Sigma_mu)*prior.m_mu)
    return DRModel(D, M, prior.sigma2_y, posterior.m_W, posterior.Sigma_W, m_mu, Sigma_mu)
end

function update_X(Y::Array{Float64, 2}, posterior::DRModel)
    D, N = size(Y)
    M = posterior.M

    W = posterior.m_W
    WW = zeros(M, M, D)
    for d in 1 : D
        WW[:,:,d] = W[:,d]*W[:,d]' + posterior.Sigma_W[:,:,d]
    end
    mu = posterior.m_mu
    X = zeros(M, N)
    XX = zeros(M, M, N)
    for n in 1 : N
        Sigma = inv(inv(posterior.sigma2_y)*sqsum(WW, 3) + eye(M))
        X[:,n] = inv(posterior.sigma2_y)*Sigma*W*(Y[:,n] - mu)
        XX[:,:,n] = X[:,n] * X[:,n]' + Sigma
    end
    return X, XX
end

function interpolate(mask::BitArray{2}, X::Array{Float64, 2}, posterior::DRModel)
    Y_est = posterior.m_W'*X + repmat(posterior.m_mu, 1, size(X, 2))
    return return Y_est[mask]
end

"""
Compute variational posterior distributions.
"""
function VI(Y::Array{Float64, 2}, prior::DRModel, max_iter::Int)
    X, XX = init(Y, prior)
    mask = isnan.(Y)
    sum_nan = sum(mask)
    posterior = deepcopy(prior)

    #progress = Progress(max_iter)
    for iter in 1 : max_iter
        # progress
        #next!(progress)

        # Interpolate
        if sum_nan > 0
            Y[mask] = interpolate(mask, X, posterior)
        end

        # M-step
        posterior = update_W(Y, prior, posterior, X, XX)
        posterior = update_mu(Y, prior, posterior, X, XX)

        # E-step
        X, XX = update_X(Y, posterior)
    end

    return posterior, X
end

end


@doc doc"""
    Hyperbolic{N} <: Manifold

The hyperbolic space $\mathbb H^n$ represented by $n+1$-Tuples, i.e. in by
vectors in $\mathbb R^{n+1}$ using the Minkowsi metric, i.e.

```math
\mathbb H^n = \Bigl\{x\in\mathbb R^{n+1}
\ \Big|\ \langle x,x \rangle_{\mathrm{M}}= -x_{n+1}^2
+ \displaystyle\sum_{k=1}^n x_k^2 = -1, x_{n+1} > 0\Bigr\},
```

where $\langle\cdot,\cdot\rangle_{\mathrm{M}}$ denotes the [`minkowski_dot`](@ref)
is Minkowski inner product. The tangent space $T_x\mathbb H^n$ is given by

````math
T_x\mathbb H^n \coloneqq \bigl\{
v \in \mathbb R^{n+1} \ \bigl |\ \langle x,v\rangle_{\mathrm{M}} = 0
\bigr\}.
````
The Minkowski inner product inntroduces the [`MinkowskiMetric`](@ref), which is
a Riemannian metric on the tangent bundle $T\mathbb H^n$.

# Constructor

    Hyperbolic(n)

Generate the $\mathbb H^{n}\subset \mathbb R^{n+1}$
"""
struct Hyperbolic{N} <: Manifold end

Hyperbolic(n::Int) = Hyperbolic{n}()

@doc doc"""
    MinkowskiMetric <: LorentzMetric

The Minkowski metric is a [`LorentzMetric`](@ref) with, i.e.

````math
\langle a,b\rangle_{\mathrm{M}} = -a_{n+1}b_{n+1} +
\displaystyle\sum_{k=1}^n a_kb_k.
````
It is also the default metric e.g. for the [`Hyperbolic`](@ref) space.

!!! note
    While the `MinkowskiMetric` itself is not positive definite in the whole embedded space,
    it is positive definite when restricted to a tangent space $T_x\mathcal M$,
    $x\in\mathcal M$, of the [`Hyperbolic`](@ref) space $\mathcal M$.
"""
struct MinkowskiMetric <: LorentzMetric end

"""
    check_manifold_point(M::Hyperbolic, x; kwargs...)

Check whether `x` is a valid point on the [`Hyperbolic`](@ref) `M`, i.e. is a vector with
[`minkowski_dot`](@ref) -1. The tolerance for the last test can be set using the `kwargs...`.
"""
function check_manifold_point(M::Hyperbolic, x; kwargs...)
    if size(x) != representation_size(M)
        return DomainError(
            size(x),
            "The point $(x) does not lie on $(M), since its size is not $(representation_size(M)).",
        )
    end
    if !isapprox(minkowski_dot(x, x), -1.0; kwargs...)
        return DomainError(
            minkowski_dot(x, x),
            "The point $(x) does not lie on $(M) since its Minkowski inner product is not -1.",
        )
    end
    return nothing
end

"""
    check_tangent_vector(M::Hyperbolic, x, v; kwargs... )

Check whether `v` is a tangent vector to `x` on the [`Hyperbolic`](@ref) `M`, i.e.
after [`check_manifold_point`](@ref)`(M,x)`, `v` has to be of same dimension as `x`
and orthogonal to `x` with respect to [`minkowski_dot`](@ref).
The tolerance for the last test can be set using the `kwargs...`.
"""
function check_tangent_vector(M::Hyperbolic, x, v; kwargs...)
    perr = check_manifold_point(M, x)
    perr === nothing || return perr
    if size(v) != representation_size(M)
        return DomainError(
            size(v),
            "The vector $(v) is not a tangent to a point on $M since its size does not match $(representation_size(M)).",
        )
    end
    if !isapprox(minkowski_dot(x, v), 0.0; kwargs...)
        return DomainError(
            abs(minkowski_dot(x, v)),
            "The vector $(v) is not a tangent vector to $(x) on $(M), since it is not orthogonal (with respect to the Minkowski inner product) in the embedding.",
        )
    end
    return nothing
end

@doc doc"""
    distance(M::Hyperbolic, x, y)

Compute the distance on the [`Hyperbolic`](@ref) `M`, which reads

````math
d_{\mathbb H^n}(x,y) = \operatorname{acosh}( - \langle x, y \rangle_{\mathrm{M}}),
````

where $\langle\cdot,\cdot\rangle_{\mathrm{M}}$ denotes the [`minkowski_dot`](@ref).
"""
distance(M::Hyperbolic, x, y) = acosh(max(-minkowski_dot(x, y), 1.0))

@doc doc"""
    exp(M::Hyperbolic, x, v)

Compute the exponential map on the [`Hyperbolic`](@ref) space $\mathbb H^n$ eminating
from `x` towards `v`, which is optionally scaled by `t`. The formula reads

````math
\exp_x v = \cosh(\sqrt{\langle v,v\rangle_{\mathrm{M}}})x
+ \sinh(\sqrt{\langle v,v\rangle_{\mathrm{M}}})\frac{v}{\sqrt{\langle v,v\rangle_{\mathrm{M}}}},
````

where $\langle\cdot,\cdot\rangle_{\mathrm{M}}$ denotes the [`minkowski_dot`](@ref).
"""
exp(::Hyperbolic, ::Any...)

function exp!(M::Hyperbolic, y, x, v)
    vn = sqrt(max(minkowski_dot(v, v), 0.0))
    vn < eps(eltype(x)) && return copyto!(y, x)
    return copyto!(y, cosh(vn) * x + sinh(vn) / vn * v)
end

flat!(M::Hyperbolic, v::CoTFVector, x, w::TFVector) = copyto!(v, w)

@doc doc"""
    injectivity_radius(M::Hyperbolic[, x])

Return the injectivity radius on the [`Hyperbolic`](@ref), which is always $\infty$.
"""
injectivity_radius(H::Hyperbolic, args...) = Inf

@doc doc"""
    inner(M::Hyperbolic, x, v, w)

Compute the Riemannian inner product for two tangent vectors `v` and `w`
from $T_x\mathbb H^n$ of the [`Hyperbolic`](@ref) space $\mathbb H^n$ given by
$\langle w, v \rangle_{\mathrm{M}}$ the [`minkowski_dot`](@ref) Minkowski
inner product on $\mathbb R^{n+1}$.
"""
@inline inner(M::Hyperbolic, x, w, v) = minkowski_dot(w, v)

is_default_metric(::Hyperbolic, ::MinkowskiMetric) = Val(true)

@doc doc"""
    log(M::Hyperbolic, x, y)

Compute the logarithmic map on the [`Hyperbolic`](@ref) space $\mathbb H^n$, the tangent
vector representing the [`geodesic`](@ref) starting from `x`
reaches `y` after time 1 on the [`Hyperbolic`](@ref) space `M`.
The formula reads for $x\neq y$

```math
\log_x y = d_{\mathbb H^n}(x,y)
\frac{y-\langle x,y\rangle_{\mathrm{M}} x}{\lVert y-\langle x,y\rangle_{\mathrm{M}} x \rVert_2}
```
and is zero otherwise.
"""
log(::Hyperbolic, ::Any...)

function log!(M::Hyperbolic, v, x, y)
    scp = minkowski_dot(x, y)
    w = y + scp * x
    wn = sqrt(max(scp .^ 2 - 1, 0.0))
    wn < eps(eltype(x)) && return zero_tangent_vector!(M, v, x)
    v .= acosh(max(1.0, -scp)) / wn .* w
    return v
end

@doc doc"""
    minkowski_dot(a,b)
Compute the Minkowski inner product of two Vectors `a` and `b` of same length
`n+1`, i.e.

````math
\langle a,b\rangle_{\mathrm{M}} = -a_{n+1}b_{n+1} + \displaystyle\sum_{k=1}^n a_kb_k.
````
"""
function minkowski_dot(a::AbstractVector, b::AbstractVector)
    return -a[end] * b[end] + sum(a[1:end-1] .* b[1:end-1])
end

@doc doc"""
    manifold_dimension(H::Hyperbolic)

Return the dimension of the hyperbolic space manifold $\mathbb H^n$, i.e. $n$.
"""
manifold_dimension(::Hyperbolic{N}) where {N} = N

"""
    mean(
        M::Hyperbolic,
        x::AbstractVector,
        [w::AbstractWeights,]
        method = CyclicProximalPointEstimationMethod();
        kwargs...,
    )

Compute the Riemannian [`mean`](@ref mean(M::Manifold, args...)) of `x` on the
[`Hyperbolic`](@ref) space using [`CyclicProximalPointEstimation`](@ref).
"""
mean(::Hyperbolic, ::Any...)

function mean!(M::Hyperbolic, y, x::AbstractVector, w::AbstractVector; kwargs...)
    return mean!(M, y, x, w, CyclicProximalPointEstimation(); kwargs...)
end

@doc doc"""
    project_tangent(M::Hyperbolic, x, v)

Perform an orthogonal projection with respect to the Minkowski inner product of `v` onto
the tangent space at `x` of the [`Hyperbolic`](@ref) space `M`.

The formula reads
````math
w = v + \langle x,v\rangle_{\mathrm{M}} x,
````
where $\langle \cdot, \cdot \rangle_{\mathrm{M}}$ denotes the Minkowski inner
product in the embedding, see [`minkowski_dot`](@ref).
"""
project_tangent(::Hyperbolic, ::Any...)

project_tangent!(::Hyperbolic, w, x, v) = (w .= v .+ minkowski_dot(x, v) .* x)

@doc doc"""
    representation_size(M::Hyperbolic)

Return the representation size on the [`Hyperbolic`](@ref), i.e. for the `n`-diomensional
hyperbolic manifold the dimention of the embedding, i.e. `n+1`.
"""
@generated representation_size(::Hyperbolic{N}) where {N} = (N + 1,)

sharp!(M::Hyperbolic, v::TFVector, x, w::CoTFVector) = copyto!(v, w)

@doc doc"""
    vector_transport_to(M::Hyperbolic, x, v, y, ::ParallelTransport)

Compute the paralllel transport of the `v` from the tangent space at `x` on the
[`Hyperbolic`](@ref) space $\mathbb H^n$ to the tangent at `y` along the [`geodesic`](@ref)
connecting `x` and `y`. The formula reads

````math
P_{y\gets x}(v) = v - \frac{\langle \log_xy,v\rangle_x}{d^2_{\mathbb H^n}(x,y)}
\bigl(\log_xy + \log_yx \bigr).
````
"""
vector_transport_to(::Hyperbolic, ::Any, ::Any, ::Any, ::ParallelTransport)

function vector_transport_to!(M::Hyperbolic, vto, x, v, y, ::ParallelTransport)
    w = log(M, x, y)
    wn = norm(M, x, w)
    wn < eps(eltype(x + y)) && return copyto!(vto, v)
    return copyto!(vto, v - (inner(M, x, w, v) * (w + log(M, y, x)) / wn^2))
end

@doc doc"""
    zero_tangent_vector(M::Hyperbolic, x)

Return the zero vector from the tangent space at `x` of the [`Hyperbolic`](@ref) `M`.
"""
zero_tangent_vector(::HybridArray, ::Any...)

zero_tangent_vector!(M::Hyperbolic, v, x) = fill!(v, 0)


# Compare with tabulated values,  classical hypergeometric formulas

@testset "Bessel" begin
    T = Float64

    for α in (3 / 2, 1 / 2, 1, 2)
        P = Bessel{α,T}
        β = 2

        p0, p1, p2, p3, p4 = basis.(P, 0:4)

        # Krall and Frink
        for x in range(0, 1, length=5)
            @test p0(x) ≈ 1
            @test p1(x) ≈ 1 + α * (x / β)
            @test p2(x) ≈ 1 + 2(α + 1) * (x / β) + (α + 1) * (α + 2) * (x / β)^2
            @test p3(x) ≈
                  1 +
                  3(α + 2) * (x / β) +
                  3(α + 2) * (α + 3) * (x / β)^2 +
                  (α + 2) * (α + 3) * (α + 4) * (x / β)^3
            @test p4(x) ≈
                  1 +
                  4(α + 3) * (x / β) +
                  6(α + 3) * (α + 4) * (x / β)^2 +
                  4(α + 3) * (α + 4) * (α + 5) * (x / β)^3 +
                  (α + 3) * (α + 4) * (α + 5) * (α + 6) * (x / β)^4
        end
    end

    x = variable(Polynomial)
    for α in (1 / 4, 1 / 2, 3 / 4, 1, 2)
        P = Bessel{α}
        for n in 0:5
            p = basis(P, n)
            @test p(x) ≈ SP.classical_hypergeometric(P, n, x)
        end
    end
end

@testset "ChebyshevT" begin
    p0 = Chebyshev([1])  # 1
    p1 = Chebyshev([0, 1]) # x
    p2 = Chebyshev([0, 0, 1]) #
    p3 = Chebyshev([0, 0, 0, 1]) #
    p4 = Chebyshev([0, 0, 0, 0, 1]) #

    for x in range(-1, 1, length=5)
        @test p0(x) ≈ 1
        @test p1(x) ≈ x
        @test p2(x) ≈ 2x^2 - 1
        @test p3(x) ≈ 4x^3 - 3x
        @test p4(x) ≈ 8x^4 - 8x^2 + 1
    end
end

@testset "ChebyshevU" begin
    p0 = ChebyshevU([1])  # 1
    p1 = ChebyshevU([0, 1]) # x
    p2 = ChebyshevU([0, 0, 1]) #
    p3 = ChebyshevU([0, 0, 0, 1]) #
    p4 = ChebyshevU([0, 0, 0, 0, 1]) #

    for x in range(-1, 1, length=5)
        @test p0(x) ≈ 1
        @test p1(x) ≈ 2x
        @test p2(x) ≈ 4x^2 - 1
        @test p3(x) ≈ 8x^3 - 4x
        @test p4(x) ≈ 16x^4 - 12x^2 + 1
    end
end

@testset "Hermite" begin
    p0 = Hermite([1])  # 1
    p1 = Hermite([0, 1]) # x
    p2 = Hermite([0, 0, 1]) #
    p3 = Hermite([0, 0, 0, 1]) #
    p4 = Hermite([0, 0, 0, 0, 1]) #

    for x in range(-1, 1, length=5)
        @test p0(x) ≈ 1
        @test p1(x) ≈ 2x
        @test p2(x) ≈ 4x^2 - 2
        @test p3(x) ≈ 8x^3 - 12x
        @test p4(x) ≈ 16x^4 - 48x^2 + 12
    end

    x = variable(LaurentPolynomial)
    P = Hermite
    for n in 2:5
        p = basis(P, n)
        @test p(x) ≈ SP.classical_hypergeometric(P, n, x)
    end
end

@testset "ChebyshevHermite" begin
    p0 = ChebyshevHermite([1])  # 1
    p1 = ChebyshevHermite([0, 1]) # x
    p2 = ChebyshevHermite([0, 0, 1]) #
    p3 = ChebyshevHermite([0, 0, 0, 1]) #
    p4 = ChebyshevHermite([0, 0, 0, 0, 1]) #

    for x in range(-1, 1, length=5)
        @test p0(x) ≈ 1
        @test p1(x) ≈ x
        @test p2(x) ≈ x^2 - 1
        @test p3(x) ≈ x^3 - 3x
        @test p4(x) ≈ x^4 - 6x^2 + 3
    end

    x = variable(LaurentPolynomial)
    P = ChebyshevHermite
    for n in 2:5
        p = basis(P, n)
        @test p(x) ≈ SP.classical_hypergeometric(P, n, x) atol = 1e-8
    end
end

@testset "Legendre" begin
    p0 = Legendre([1])  # 1
    p1 = Legendre([0, 1]) # x
    p2 = Legendre([0, 0, 1]) #
    p3 = Legendre([0, 0, 0, 1]) #
    p4 = Legendre([0, 0, 0, 0, 1]) #

    for x in range(-1, stop=1, length=5)
        @test p0(x) ≈ 1
        @test p1(x) ≈ x
        @test p2(x) ≈ 1 / 2 * (3x^2 - 1)
        @test p3(x) ≈ 1 / 2 * (5x^3 - 3x)
        @test p4(x) ≈ 1 / 8 * (35x^4 - 30x^2 + 3)
    end
end

@testset "Shifted Legendre" begin
    T = Float64
    P = ShiftedLegendre{T}

    p0 = P([1])  # 1
    p1 = P([0, 1]) # x
    p2 = P([0, 0, 1]) #
    p3 = P([0, 0, 0, 1]) #
    p4 = P([0, 0, 0, 0, 1]) #

    for x in range(0, stop=1, length=5)
        @test p0(x) ≈ 1
        @test p1(x) ≈ 2x - 1
        @test p2(x) ≈ 6x^2 - 6x + 1
        @test p3(x) ≈ 20x^3 - 30x^2 + 12x - 1
        @test p4(x) ≈ 70x^4 - 140x^3 + 90x^2 - 20x + 1
    end
end

@testset "Laguerre" begin
    P = Laguerre{0}
    p0 = P([1])  # 1
    p1 = P([0, 1]) # x
    p2 = P([0, 0, 1]) #
    p3 = P([0, 0, 0, 1]) #
    p4 = P([0, 0, 0, 0, 1]) #

    for x in range(0, stop=5, length=5)
        @test p0(x) ≈ 1
        @test p1(x) ≈ -x + 1
        @test p2(x) ≈ 1 / 2 * (x^2 - 4x + 2)
        @test p3(x) ≈ 1 / 6 * (-x^3 + 9x^2 - 18x + 6)
        @test p4(x) ≈ 1 / 24 * (x^4 - 16x^3 + 72x^2 - 96x + 24)
    end

    x = variable(Polynomial)
    for α in (0, 1 / 4, 1 / 2, 3 / 4)
        P = Laguerre{α}
        for n in 2:5
            p = basis(P, n)
            @test p(x) ≈ SP.classical_hypergeometric(P, n, x)
        end
    end
end

@testset "Gegenbauer" begin
    for α in (1 / 4, 1 / 3, 1 / 2, 1, 2)
        P = Gegenbauer{α,Float64}
        p0 = P([1])  # 1
        p1 = P([0, 1]) # x
        p2 = P([0, 0, 1]) #
        p3 = P([0, 0, 0, 1]) #
        p4 = P([0, 0, 0, 0, 1]) #

        for x in range(-1, stop=1, length=5)
            @test p0(x) ≈ 1
            @test p1(x) ≈ 2α * x
            @test p2(x) ≈ -α + 2α * (1 + α) * x^2
            @test p3(x) ≈ -2α * (1 + α) * x + 4 / 3 * α * (1 + α) * (2 + α) * x^3
        end
    end

    x = variable(Polynomial)
    for α in (1 / 4, 1 / 2, 3 / 4, 1)
        P = Gegenbauer{α}
        for n in 2:5
            p = basis(P, n)
            @test p(x) ≈ SP.classical_hypergeometric(P, n, x)
        end
    end
end

@testset "Jacobi" begin
    T = Float64
    x = variable()

    # ChebyshevT = J(-1/2, -1/2)
    for (alpha_beta, Q) in (
        ((-1 / 2, -1 / 2), Chebyshev{T}),
        ((1 / 2, 1 / 2), ChebyshevU{T}),
        ((0, 0), Legendre{T}),
    )
        P = Jacobi{alpha_beta...,T}
        for i in 2:6
            p = basis(P, i)
            q = basis(Q, i)
            @test SP.monic(p)(x) ≈ SP.monic(q)(x)
        end
    end

    x = variable(Polynomial)
    for (α, β) in ((1 / 2, 1 / 2), (-1 / 2, 1 / 2), (1 / 2, -1 / 2), (-1 / 2, -1 / 2))
        P = Jacobi{α,β}
        for n in 2:5
            p = basis(P, n)
            @test p(x) ≈ SP.classical_hypergeometric(P, n, x)
        end
    end

    ## We have `jacobi_eval` for α and β outside [-1, ∞)
    α, β = 1 / 2, -1 / 2
    for n in 2:5
        x0 = 1 / 2
        @test SP.jacobi_eval(α, β, n, x0) ≈ basis(Jacobi{α,β}, n)(x0)
    end
end

## CDOP

@testset "Charlier" begin
    x = variable()
    for μ in (1 / 4, 1 / 2, 1, 2)
        P = Charlier{μ}
        @test basis(P, 0)(x) ≈ one(x)
        @test basis(P, 1)(x) ≈ 1 - x / μ
        @test basis(P, 2)(x) ≈ (x^2 + μ^2 - x * (1 + 2μ)) / μ^2
    end

    for μ in (1 / 4, 1 / 2, 1, 2)
        P = Charlier{μ}
        for i in 0:5
            @test basis(P, i)(x) ≈ SP.classical_hypergeometric(P, i, x)
        end
    end
end

@testset "Hahn" begin
    x = variable()
    for α in (1 / 4, 1 / 2)
        for β in (1 / 4, 1 / 2)
            𝐍 = 5

            P = Hahn{α,β,𝐍}
            for i in 0:𝐍
                @test basis(P, i)(x) ≈ SP.classical_hypergeometric(P, i, x)
            end
        end
    end
end

@testset "Krawchouk" begin
    x = variable()
    for p in (1 / 4, 1 / 2, 1, 2)
        for N in (3, 6)
            P = Krawchouk{p,N}
            @test basis(P, 0)(x) ≈ one(x)
            @test basis(P, 1)(x) ≈ -N * p + x
            @test basis(P, 2)(x) ≈ 1 / 2 * (N^2 * p^2 + x * (2p + x - 1) - N * p * (p + 2x))
        end
    end

    for p in (1 / 4, 1 / 2, 1, 2)
        for N in (3, 6)
            P = Krawchouk{p,N}
            for i in 0:N
                @test basis(P, i)(x) ≈ SP.classical_hypergeometric(P, i, x)
            end
        end
    end
end

@testset "Meixner" begin
    x = variable()
    for γ in (1 / 4, 1 / 2, 3 / 4)
        for μ in (1 / 4, 1 / 2, 3 / 4)
            P = Meixner{γ,μ}
            for i in 0:4
                @test basis(P, i)(x) ≈ SP.classical_hypergeometric(P, i, x)
            end
        end
    end
end


module MDSim

export init, simulate!

using Statistics: mean

mutable struct MDSystem{DT<:AbstractFloat}
    l::DT
    h::DT
    U::DT
    K::DT
    T::DT
    P::DT
    N::Integer
    f::Matrix{DT}
    r::Matrix{DT}
    v::Matrix{DT}
    function MDSystem(N::Integer, T₀::DT, h::DT, l::DT) where {DT<:AbstractFloat}
        r = rand(DT, 2, N) * l
        v = rand(DT, 2, N) .- 0.5
        v *= √(T₀ ÷ mean(v .^ 2))
        f = Matrix{DT}(undef, 2, N)
        return new{DT}(l, h, DT(NaN), DT(NaN), DT(NaN), DT(NaN), N, f, r, v)
    end
end

function init(; N::Integer, T₀::DT, h::DT, l::DT) where {DT<:AbstractFloat}
    sys = MDSystem(N, T₀, h, l)
    alignleft!(sys)
    update_forces!(sys)
    update_phase!(sys)
    check_boundary!(sys)
    update_potential!(sys)
    update_kinetic!(sys)
    update_temperature!(sys)
    update_pressure!(sys)
    return sys
end

function simulate!(sys::MDSystem)
    update_phase!(sys)
    check_boundary!(sys)
    update_potential!(sys)
    update_kinetic!(sys)
    update_temperature!(sys)
    update_pressure!(sys)
end

function alignleft!(sys::MDSystem)    #Thanks to my dearest friend for his beautiful method in setting the initial conditions SLH Hero
    xcount = ceil(Integer, (2 * sys.N)^(1 / 2) / 2)
    ycount = 2 * xcount

    xs = collect(range(0.1 * sys.l / 2, 0.9 * sys.l / 2, length = xcount))
    sys.r[1, :] .= repeat(xs, outer = ycount)[1:sys.N]

    ys = collect(range(0.1 * sys.l, 0.9 * sys.l, length = ycount))
    sys.r[2, :] .= repeat(ys, inner = xcount, outer = 1)[1:sys.N]
end

function update_forces!(sys::MDSystem)
    for i ∈ 1:sys.N
        sys.f[:, i] *= 0
        for j ∈ setdiff(1:sys.N, i)
            rmirror = sys.r[:, j]
            if sys.r[1, i] - rmirror[1] > sys.l / 2
                rmirror[1] += sys.l
            elseif sys.r[1, i] - rmirror[1] < -sys.l / 2
                rmirror[1] -= sys.l
            end
            if sys.r[2, i] - rmirror[2] > sys.l / 2
                rmirror[2] += sys.l
            elseif sys.r[2, i] - rmirror[2] < -sys.l / 2
                rmirror[2] -= sys.l
            end
            Δr = √sum((sys.r[:, i] - rmirror) .^ 2)
            sys.f[:, i] += 48 * (((Δr)^-14) - 0.5 * ((Δr)^-8)) * (sys.r[:, i] - rmirror)
        end
    end
end

function update_force!(sys::MDSystem, i::Integer)
    sys.f[:, i] *= 0
    for j ∈ setdiff(1:sys.N, i)
        rmirror = sys.r[:, j]
        if sys.r[1, i] - rmirror[1] > sys.l / 2
            rmirror[1] += sys.l
        elseif sys.r[1, i] - rmirror[1] < -sys.l / 2
            rmirror[1] -= sys.l
        end
        if sys.r[2, i] - rmirror[2] > sys.l / 2
            rmirror[2] += sys.l
        elseif sys.r[2, i] - rmirror[2] < -sys.l / 2
            rmirror[2] -= sys.l
        end
        Δr = √sum((sys.r[:, i] - rmirror) .^ 2)
        sys.f[:, i] += 48 * (((Δr)^-14) - 0.5 * ((Δr)^-8)) * (sys.r[:, i] - rmirror)
    end
end

"""
important note:
    update_forces!() will update the force of the whole system,
    but update_force! only updates the selected particle.
"""

function update_phase!(sys::MDSystem)
    for i ∈ 1:sys.N
        f₁ = sys.f[:, i]
        sys.r[:, i] += sys.h * sys.v[:, i] + 0.5 * sys.h * sys.h * f₁
        update_force!(sys, i)
        f₂ = sys.f[:, i]
        sys.v[:, i] += sys.h * (f₁ + f₂) / 2
    end
end

function check_boundary!(sys::MDSystem)
    sys.r = (sys.r .+ sys.l) .% sys.l
end

function update_potential!(sys::MDSystem)
    Σrterm = 0.0

    for i ∈ 1:(sys.N-1)
        for j ∈ i+1:sys.N
            rmirror = sys.r[:, j]
            if sys.r[1, i] - rmirror[1] > sys.l / 2
                rmirror[1] += sys.l
            elseif sys.r[1, i] - rmirror[1] < -sys.l / 2
                rmirror[1] -= sys.l
            end
            if sys.r[2, i] - rmirror[2] > sys.l / 2
                rmirror[2] += sys.l
            elseif sys.r[2, i] - rmirror[2] < -sys.l / 2
                rmirror[2] -= sys.l
            end
            Δr = √sum((sys.r[:, i] - rmirror) .^ 2)
            Σrterm += ((Δr)^-12) - ((Δr)^-6)
        end
    end

    U = 4 * Σrterm

    sys.U = U
end

function update_kinetic!(sys::MDSystem)
    sys.K = sum(sys.v .^ 2) / 2
end

function update_temperature!(sys::MDSystem)
    sys.T = sum(sys.v .^ 2) / (2 * sys.N)
end

function update_pressure!(sys::MDSystem)
    Σv² = sum(sys.v .^ 2)
    Σrterm = 0.0

    for i ∈ 1:(sys.N-1)
        for j ∈ i+1:sys.N
            rmirror = sys.r[:, j]
            if sys.r[1, i] - rmirror[1] > sys.l / 2
                rmirror[1] += sys.l
            elseif sys.r[1, i] - rmirror[1] < -sys.l / 2
                rmirror[1] -= sys.l
            end
            if sys.r[2, i] - rmirror[2] > sys.l / 2
                rmirror[2] += sys.l
            elseif sys.r[2, i] - rmirror[2] < -sys.l / 2
                rmirror[2] -= sys.l
            end
            Δr = √sum((sys.r[:, i] - rmirror) .^ 2)
            Σrterm += ((Δr)^-12) - 0.5 * ((Δr)^-6)
        end
    end
    sys.P = ((48 * Σrterm) + Σv²) / (2 * sys.l * sys.l)
end

end


using Fatou
using SyntaxTree
using LaTeXStrings
using Test

Fatou.Reduce.load_package(:rlfi)
@test newton(:(z^3-1)) |> typeof <: Fatou.Define
#@test orbit(:(1*x),x->x, [-1.7 2π -1.2],17,3,147) == nothing
#@test orbit(:(1*x),x->x, [-1.7 2π -1.2],0,1) == nothing
@test basin(juliafill(:(z^3-1)),1) |> typeof == LaTeXStrings.LaTeXString
@test basin(newton(:(z^3-1)),1) |> typeof == LaTeXStrings.LaTeXString
@test newton(:(z^3-1)) |> fatou |> typeof <: Fatou.FilledSet
@test mandelbrot(:(z^2+c)) |> fatou |> typeof <: Fatou.FilledSet
@test juliafill(:(z^2-0.06+0.67im)) |> fatou |> typeof <: Fatou.FilledSet
#@test (η = Fatou.newton_raphson(:(z^3-1),1); ξ = genfun(η,[:z,:c]); ζ = 2.1; [ξ(ζ,0), [Fatou.recomp(η,ζ,k) for k ∈ 2:3]...]) |> typeof <: Array
# fix for case k ∈ 2:4+


"""
    UnitSimplexOracle(right_side)

Represents the scaled unit simplex:
```
C = {x ∈ R^n_+, ∑x ≤ right_side}
```
"""
struct UnitSimplexOracle{T} <: LinearMinimizationOracle
    right_side::T
end

UnitSimplexOracle{T}() where {T} = UnitSimplexOracle{T}(one(T))

UnitSimplexOracle(rhs::Integer) = UnitSimplexOracle{Rational{BigInt}}(rhs)

"""
LMO for scaled unit simplex:
`∑ x_i = τ`
Returns either vector of zeros or vector with one active value equal to RHS if
there exists an improving direction.
"""
function compute_extreme_point(lmo::UnitSimplexOracle{T}, direction; v = nothing, kwargs...) where {T}
    idx = argmin(direction)
    if direction[idx] < 0
        return ScaledHotVector(lmo.right_side, idx, length(direction))
    end
    return ScaledHotVector(zero(T), idx, length(direction))
end

function convert_mathopt(
    lmo::UnitSimplexOracle{T},
    optimizer::OT;
    dimension::Integer,
    kwargs...,
) where {T,OT}
    MOI.empty!(optimizer)
    τ = lmo.right_side
    n = dimension
    (x, _) = MOI.add_constrained_variables(optimizer, [MOI.Interval(0.0, 1.0) for _ in 1:n])
    MOI.add_constraint(
        optimizer,
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(ones(n), x), 0.0),
        MOI.LessThan(τ),
    )
    return MathOptLMO(optimizer)
end

"""
Dual costs for a given primal solution to form a primal dual pair
for scaled unit simplex.
Returns two vectors. The first one is the dual costs associated with the constraints
and the second is the reduced costs for the variables.
"""
function compute_dual_solution(::UnitSimplexOracle{T}, direction, primalSolution) where {T}
    idx = argmax(primalSolution)
    critical = min(direction[idx], 0)
    lambda = [critical]
    mu = direction .- lambda
    return lambda, mu
end

"""
    ProbabilitySimplexOracle(right_side)

Represents the scaled probability simplex:
```
C = {x ∈ R^n_+, ∑x = right_side}
```
"""
struct ProbabilitySimplexOracle{T} <: LinearMinimizationOracle
    right_side::T
end

ProbabilitySimplexOracle{T}() where {T} = ProbabilitySimplexOracle{T}(one(T))

ProbabilitySimplexOracle(rhs::Integer) = ProbabilitySimplexOracle{Float64}(rhs)

"""
LMO for scaled probability simplex.
Returns a vector with one active value equal to RHS in the
most improving (or least degrading) direction.
"""
function compute_extreme_point(lmo::ProbabilitySimplexOracle{T}, direction; v = nothing, kwargs...) where {T}
    idx = argmin(direction)
    return ScaledHotVector(lmo.right_side, idx, length(direction))
end

function convert_mathopt(
    lmo::ProbabilitySimplexOracle{T},
    optimizer::OT;
    dimension::Integer,
    kwargs...,
) where {T,OT}
    MOI.empty!(optimizer)
    τ = lmo.right_side
    n = dimension
    (x, _) = MOI.add_constrained_variables(optimizer, [MOI.Interval(0.0, 1.0) for _ in 1:n])
    MOI.add_constraint(
        optimizer,
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(ones(n), x), 0.0),
        MOI.EqualTo(τ),
    )
    return MathOptLMO(optimizer)
end

"""
Dual costs for a given primal solution to form a primal dual pair
for scaled probability simplex.
Returns two vectors. The first one is the dual costs associated with the constraints
and the second is the reduced costs for the variables.
"""
function compute_dual_solution(
    ::ProbabilitySimplexOracle{T},
    direction,
    primal_solution;
    kwargs...,
) where {T}
    idx = argmax(primal_solution)
    lambda = [direction[idx]]
    mu = direction .- lambda
    return lambda, mu
end


using DEBTrait, Test
using Random
Random.seed!(123)

######## benchmark K_D against Tang et al. (2019)
D_S     = 1.4e-9*ones(1)
N_p     = 3000*ones(1,1)
r_c     = 1e-6
V_cell  = (4/3)*pi*r_c^3*ones(1)
ρ_p     = DEBTrait.transporter_density(V_cell, N_p)
K_D_0   = DEBTrait.specific_reference_affinity(ρ_p, V_cell, D_S)
@test K_D_0[1] ≈ 5.7969874109e-05 atol=1e-6
#
pct_sand = 70.0
pct_clay = 30.0
s_sat    = 0.10101010101010102
N_cell   = 10.0*ones(1)
Soil = BaseSoil(pct_sand, pct_clay, s_sat, N_cell)
K_D     = DEBTrait.affinity_constant(ρ_p, V_cell, D_S, Soil.pct_sand, Soil.pct_clay, Soil.s_sat, Soil.N_cell)
@test K_D[1] ≈ 1.447412099595345 atol=1e-6
#
Soil.s_sat = 0.8080808080808082
K_D     = DEBTrait.affinity_constant(ρ_p, V_cell, D_S, Soil.pct_sand, Soil.pct_clay, Soil.s_sat, Soil.N_cell)
@test K_D[1] ≈ 9.390492906747054e-05 atol=1e-6
#
N_SB    = DEBTrait.monomer_uptake_sites(V_cell, ρ_p)

########
n_microbes      = 1
n_substrates    = 1
Isolate         = BaseGenome(1e6*ones(n_microbes), 10*ones(n_microbes), 1.0*ones(n_microbes), 1.0*ones(n_substrates,n_microbes))
Monomer         = BaseMonomer(1e-10*ones(n_substrates), 6*ones(n_substrates), 3*ones(n_substrates), 0.3*ones(n_substrates))
Soil            = BaseSoil(0.0, 0.0, 1.0, 1*ones(n_microbes))
MetabolismC     = PMetabolismC{BaseGenome, Array{Float64,1}}(Isolate)
MetabolismC.y_EM = ones(n_microbes)
MetabolismC.y_EX = ones(n_microbes)
MetabolismC.α    = zeros(n_microbes)
AssimilationC   = PAssimilation(Isolate, Monomer, Soil, MetabolismC)
@test size(AssimilationC.ρ_p) == (n_substrates, n_microbes)
@test size(AssimilationC.N_SB) == (n_substrates, n_microbes)
@test size(AssimilationC.K_D) == (n_substrates, n_microbes)
@test AssimilationC.K_D_0 ≈ AssimilationC.K_D atol = 1e-4
Soil.s_sat = 0.05
AssimilationC   = PAssimilation(Isolate, Monomer, Soil, MetabolismC)
@test AssimilationC.K_D_0[1] < AssimilationC.K_D[1]

n_microbes      = rand(1:10)
n_substrates    = rand(1:10)
Isolate         = BaseGenome(1e6*ones(n_microbes), 10*ones(n_microbes), 1.0*ones(n_microbes), 1.0*ones(n_substrates,n_microbes))
Monomer         = BaseMonomer(1e-10*ones(n_substrates), 6*ones(n_substrates), 3*ones(n_substrates), 0.3*ones(n_substrates))
Soil            = BaseSoil(0.0, 0.0, 1.0, 10*ones(n_microbes))
MetabolismC     = PMetabolismC{BaseGenome, Array{Float64,1}}(Isolate)
MetabolismC.y_EM = ones(n_microbes)
MetabolismC.y_EX = ones(n_microbes)
MetabolismC.α    = zeros(n_microbes)
AssimilationC   = PAssimilation(Isolate, Monomer, Soil, MetabolismC)
@test size(AssimilationC.ρ_p) == (n_substrates, n_microbes)
@test size(AssimilationC.N_SB) == (n_substrates, n_microbes)
@test size(AssimilationC.K_D) == (n_substrates, n_microbes)


using DataFrames
using CSV
using HorseML.Clustering
using HorseML.Clustering: fit!

@testset "Clustering" begin
    sy, st = rand(), rand()
    y, t = rand(20), rand(20)
    my1, my2, my3 = fill!(Array{Float64}(undef, 1, 1), sy), y[:, :], rand(20, 2)

    #Load the data
    df = CSV.read("clustering.csv", DataFrame)
    K = 3
    x = Array(df)

    #Test for Kmeans
    model = Kmeans(K)
    @test_nowarn fit!(model, x)
    @test_nowarn model(x)

    μ = model.μ
    @test_throws DomainError dm(sy, st, μ)
    @test typeof(dm(y, t, μ)) == Float64
    @test typeof(dm(y, t, μ, reduction = "sum")) == Float64
    @test size(dm(y, t, μ, reduction = "none")) == (20, )
    @test_throws ArgumentError dm(y, t, μ, reduction="NaN")
    @test_throws DomainError dm(my1, st, μ)
    @test dm(my2, st, μ) == dm(vec(my2), st, μ)
    @test_throws DimensionMismatch dm(my3, st, μ)
    @test dm(my2, t, μ) == dm(y, t, μ)
    @test_throws DimensionMismatch dm(my3, t, μ)

    #Test for GMM(Gaussian Mixture Model)
    model = GMM(K)
    @test_nowarn fit!(model, x)
    @test_nowarn model(x)

    π, μ, Σ = model.π, model.μ, model.Σ
    @test nlh(x, π, μ, Σ) ≈ 215.78213657665862
end


using Jchemo

n = 100 ; p = 10
X = rand(n, p)
y = rand([3, 4, 10], n)
#y = rand(["a", "b", "c"], n) ;
Xtrain = X[1:70, :] ; ytrain = y[1:70]
Xtest = X[71:n, :] ; ytest = y[71:n]

tab(ytrain)
tab(ytest)

############ PLSR-DA

nlv = 0:3
pars = mpars(nlv = nlv)
gridscore(Xtrain, ytrain, Xtest, ytest;
    score = err, fun = plsrda, pars = pars)

# Faster
gridscorelv(Xtrain, ytrain, Xtest, ytest;
    score = err, fun = plsrda, nlv = nlv)

############ PLS-LDA/QDA

nlv = 1:3
pars = mpars(nlv = nlv)
gridscore(Xtrain, ytrain, Xtest, ytest;
    score = err, fun = plslda, pars = pars)

# Faster
gridscorelv(Xtrain, ytrain, Xtest, ytest;
    score = err, fun = plslda, nlv = nlv)

gridscorelv(Xtrain, ytrain, Xtest, ytest;
    score = err, fun = plsqda, nlv = nlv)

############ RR-DA

lb = [.001; .01; .1]
pars = mpars(lb = lb)
gridscore(Xtrain, ytrain, Xtest, ytest;
    score = err, fun = rrda, pars = pars)

# Faster
gridscorelb(Xtrain, ytrain, Xtest, ytest;
    score = err, fun = rrda, lb = lb)

############ KRR-DA

lb = [.001; .01; .1]
pars = mpars(lb = lb, gamma = [1; 10])
#pars = mpars(lb = lb, kern = ["kpol";], degree = 1:2)
gridscore(Xtrain, ytrain, Xtest, ytest;
    score = err, fun = krrda, pars = pars)

# Faster
pars = mpars(gamma = [1; 10])
gridscorelb(Xtrain, ytrain, Xtest, ytest;
    score = err, fun = krrda, lb = lb, pars = pars)

############ kNN-DA

nlvdis = 5 ; metric = ["mahal";]
h = [1; 3] ; k = [20; 10]
pars = mpars(nlvdis = nlvdis, metric = metric, h = h, k = k)
gridscore(Xtrain, ytrain, Xtest, ytest;
    score = msep, fun = knnda, pars = pars, verbose = true)

############ kNN-LWPLSR-DA

nlvdis = 5 ; metric = ["mahal";]
h = [1; 3] ; k = [20; 10]
nlv = 1:2
pars = mpars(nlv = nlv, nlvdis = nlvdis, metric = metric, h = h, k = k)
gridscore(Xtrain, ytrain, Xtest, ytest;
    score = msep, fun = lwplsrda, pars = pars, verbose = true)

# Faster
pars = mpars(nlvdis = nlvdis, metric = metric, h = h, k = k)
gridscorelv(Xtrain, ytrain, Xtest, ytest;
    score = msep, fun = lwplsrda, nlv = nlv, pars = pars, verbose = false)

############ kNN-LWPLS-LDA

nlvdis = 5 ; metric = ["mahal";]
h = [1; 3] ; k = [20; 10]
nlv = 1:2
pars = mpars(nlv = nlv, nlvdis = nlvdis, metric = metric, h = h, k = k)
gridscore(Xtrain, ytrain, Xtest, ytest;
    score = msep, fun = lwplslda, pars = pars, verbose = true)

# Faster
pars = mpars(nlvdis = nlvdis, metric = metric, h = h, k = k)
gridscorelv(Xtrain, ytrain, Xtest, ytest;
    score = msep, fun = lwplslda, nlv = nlv, pars = pars, verbose = false)

########### PLSR-DA-AGG

# Here there is no sense to use gridscorelv

pars = mpars(nlv = ["1:2"; "1:3"])
gridscore(Xtrain, ytrain, Xtest, ytest;
    score = msep, fun = plsrda_agg, pars = pars)

############ kNN-LWPLSR-DA-AGG

# Here there is no sense to use gridscorelv

nlvdis = 5 ; metric = ["mahal";] ;
h = [1.; 3.] ; k = [20; 10]
nlv = ["1:2"; "2:5"] ;
pars = mpars(nlv = nlv, nlvdis = nlvdis, metric = metric, h = h, k = k) ;
gridscore(Xtrain, ytrain, Xtest, ytest;
    score = msep, fun = lwplsrda_agg, pars = pars, verbose = true)


type SDF <: SDFDynamicSystem
    omg::Float64
    zta::Float64
    u0::Float64
    v0::Float64

    SDF(omg::Float64, zta::Float64) = new(omg, zta, 0.0, 0.0)
end

function SDF(m::Float64, c::Float64, k::Float64)
    omg = sqrt(k/m)
    zta = c/(2.0*m*omg)
    SDF(omg, zta)
end

function systemEquation(s::SDF, p::Function)
    function dydt(y::AbstractArray{Float64,1}, t::Float64)
        [y[2]; p(t)-2.0*s.zta*s.omg*y[2]-s.omg*s.omg*y[1]]
    end
    dydt
end

function response(s::SDF, p::Function, t::AbstractArray{Float64,1})
    y0 = [s.u0; s.v0]

    f = systemEquation(s, p)

    y = ode4(f, y0, t)

    u = y[:,1]
    v = y[:,2]
    fdz = 2.0*s.zta*s.omg*v
    a = map(p,t)-fdz-s.omg*s.omg*u

    [u v a fdz zeros(length(t))]
end


module RobustAdaptiveMetropolisSampler

using LinearAlgebra, Random, Distributions, PDMats, ProgressMeter

export RAM_sample

# The following methods cover different ways to pass in a co-variance matrix
function RAM_sample(logtarget, x0::AbstractVector{<:Number}, M0::AbstractMatrix{<:Real}, n::Int; kwargs...)
    return RAM_sample(logtarget, x0, PDMat(M0), n; kwargs...)
end

function RAM_sample(logtarget, x0::AbstractVector{<:Number}, M0::AbstractVector{<:Real}, n::Int; kwargs...)
    return RAM_sample(logtarget, x0, PDiagMat(abs2.(M0)), n; kwargs...)
end

function RAM_sample(logtarget, x0::AbstractVector{<:Number}, M0::Diagonal{<:Real}, n::Int; kwargs...)
    return RAM_sample(logtarget, x0, PDiagMat(diag(M0)), n; kwargs...)
end

function RAM_sample(logtarget, x0::AbstractVector{<:Number}, M0::Real, n::Int; kwargs...)
    return RAM_sample(logtarget, x0, ScalMat(length(x0), abs2(M0)), n; kwargs...)
end

# Actual sampling code

function RAM_sample(logtarget, x0::AbstractVector{<:Number}, M0::AbstractPDMat, n::Int; opt_α=0.234, γ=2/3, q=Normal(), show_progress::Bool=true)
    length(x0) == size(M0, 1) || error("Covariance matrix M0 must match size of x0.")
    n > 0 || error("n must be larger than 0.")
    0 < opt_α < 1 || error("opt_α must be between 0 and 1.")
    0.5 < γ <= 1 || error("γ must be between 0.5 and 1.")

    x = copy(x0)
    s = cholesky(M0)

    d = length(x0)

    u = zeros(d)
    y = zeros(d)

    log_probability_x = logtarget(x0)

    output_chain = Matrix{Float64}(undef, n, d)

    stats_accepted_values = 0

    progress_meter = Progress(n, show_progress ? 0.1 : Inf)

    # This is a pre-allocated vector used in the loop
    scaled_proposal_vector = Vector{Float64}(undef, d)

    for i in 1:n
        # Step R1
        rand!(q, u)

        y[:] .= x .+ mul!(scaled_proposal_vector, s.L, u)

        # Step R2
        log_probability_y = logtarget(y)

        α = min(1, exp(log_probability_y - log_probability_x))

        if α > rand()
            stats_accepted_values += 1
            x, y = y, x
            log_probability_x = log_probability_y
        end

        # Step R3

        # This is taken from the second paragraph of section 5
        η = min(1, d * i^-γ)

        # Compute the new covariance matrix
        M = s.L * (I + η * (α-opt_α) * (u * u') / norm(u)^2 ) * s.L'

        # The paper has a proof that M is symmetric, so we declare that fact
        # to work around numerical rounding errors
        s = cholesky(Symmetric(M))

        output_chain[i, :] .= x

        next!(progress_meter; showvalues = [(:acceptance_rate,stats_accepted_values/i)])
    end

    return (chain=output_chain, acceptance_rate=stats_accepted_values/n, M=s.L * s.L')
end

end


#=
Common parameters for estimating the probability of for a diffusion with X(0) = x₀
satisfying X(T) ∈ B for the Muller potential.
=#

using LinearAlgebra
using Random
using BasicMD
using ForwardDiff
using Optim
using TestLandscapes: Muller

#  set up potential
V = x-> Muller(x)
# commented out due to threading issue with ForwardDiff
# cfg = ForwardDiff.GradientConfig(V, zeros(Float64,2));
∇V! = (gradV, X) -> ForwardDiff.gradient!(gradV, V, X);

# Identify minima for defining x₀ and target set B
mins=[optimize(V, x) for x in [[-0.5, 1.5],[0., 0.5],[0.5, 0.]]];
x₀ = copy(mins[1].minimizer);   # initial point
xt = copy(mins[3].minimizer);   # point defining target set B

T = 10;     # terminal time
Δt = 1e-4;  # time step
nΔt = Int(T/Δt);
β = 0.1;    # inverse temperature

# define observable as the indicator function on B = {x∣ |x-xt|≤r}
r = 0.5;
f = X-> Int(norm(X.-xt)≤r);

sampler = MALA(V,∇V!, β, Δt);


module ConditionalMean

using Base: indices1, tail
export condmean, condvar, condstd
export nanmean, nanstd

"""
`M = condmean(A, cond, region)` calculates the mean value of A along the dimensions listed in `region`, ignoring values when cond is not true, or values equal to cond if cond is a number.
`M = condmean(f, A, cond, region)` calculates the mean of f(A).
"""
#condmean{T<:Real}(A::AbstractArray{T}, cond::Function, region=1) = _condmean(A, T, region)
# vectorizing mean(A[:]) is trivial, so keep default region definition of 1
condmean{T<:Number}(             A::AbstractArray{T}, cond::Function=(x->true), region=1; Km1::Bool=false) = _condmean(identity, A, T, cond      , region, Km1)
condmean{T<:Number}(             A::AbstractArray{T}, cond::Number            , region=1; Km1::Bool=false) = _condmean(identity, A, T, x->x!=cond, region, Km1)
condmean{T<:Number}(f::Function, A::AbstractArray{T}, cond::Function=(x->true), region=1; Km1::Bool=false) = _condmean(f       , A, T, cond      , region, Km1)
condmean{T<:Number}(f::Function, A::AbstractArray{T}, cond::Number            , region=1; Km1::Bool=false) = _condmean(f       , A, T, x->x!=cond, region, Km1)

# redirect _condmean to _condmeanoffset for Array{Float32} to improve precision
#_condmean{T<:Float32}(f::Function, A::AbstractArray, ::Type{T}, cond, region, Km1::Bool=false) = _condmeanoffset(f, A, T, cond, region, Km1)

# # general condition cond generalizes to all Number types, but behavior maybe unpredictable for some
# function _condmean{T<:Number}(f::Function, A::AbstractArray, ::Type{T}, cond, region, Km1::Bool=false)
#     inds = Base.reduced_indices(A, region)
#     K = similar(Array{Int}, inds)
#     S = similar(array{eltype(A)}, inds)
#     fill!(K,0)
#     fill!(S, zero(eltype(A)))
#     condsum!(S, K, A, f, cond)
#     if Km1
#        S./(K-1)
#     else
#        S./K
#     end
# end

# add offset to improve precision for floating point numerics, behavior unpredictable for other types
function _condmeanoffset{T<:Number}(f::Function, A::AbstractArray, ::Type{T}, cond, region, Km1::Bool=false)
    inds = Base.reduced_indices(A, region)
    K = similar(Array{Int}, inds)
    S = similar(array{eltype(A)}, inds)
    P = similar(array{eltype(A)}, inds)
    fill!(K,0)
    fill!(S, zero(eltype(A)))
    fill!(P, zero(eltype(A)))
    condsumoffset!(S, P, K, A, f, cond)
    if Km1
        (S + P.*K) ./ (K-1)
    else
        S./K + P
    end
end

# always redirect _condmean to _condmeanoffset to improve precision
_condmean=_condmeanoffset

using Base: check_reducedims, reducedim1, safe_tail
using Base.Broadcast: newindex

"""
    condsum!(S, K, A, f, cond)

Compute the sum `S` and number of contributing values `K` for
reductions of the array `A` over dimensions. `S` and `K` must have
identical indices, and either match `A` or have singleton-dimensions for
the dimensions that are being summed over. Only values x with cond(x)==true
are included in the tallies of `S` and `K`.

Note that the mean is just S./K.
"""
function condsum!{T,N}(S, K, A::AbstractArray{T,N}, f, cond)
    check_reducedims(S, A)
    isempty(A) && return S, K
    indices(S) == indices(K) || throw(DimensionMismatch("S and K must have identical indices"))

    indsAt, indsSt = safe_tail(indices(A)), safe_tail(indices(S))
    keep, Idefault = _newindexer(indsAt, indsSt)
    if reducedim1(S, A)
        # keep the accumulators as a local variable when reducing along the first dimension
        i1 = first(indices1(S))
        @inbounds for IA in CartesianRange(indsAt)
            IS = newindex(IA, keep, Idefault)
            s, k = S[i1,IS], K[i1,IS]
            for i in indices(A, 1)
                tmp = A[i, IA]
                if cond(tmp)
                    s += f(tmp)
                    k += 1
                end
            end
            S[i1,IS], K[i1,IS] = s, k
        end
    else
        @inbounds for IA in CartesianRange(indsAt)
            IS = newindex(IA, keep, Idefault)
            for i in indices(A, 1)
                tmp = A[i, IA]
                if cond(tmp)
                    S[i, IS] += f(tmp)
                    K[i, IS] += 1
                end
            end
        end
    end
    S, K
end
if VERSION < v"0.6.0-dev.693"
    _newindexer(shape, inds) = Base.Broadcast.newindexer(shape, inds)
else
    _newindexer(shape, inds) = Base.Broadcast.shapeindexer(shape, inds)
end

# add offset to improve precision for floating point numerics, behavior unpredictable for other types
"""
    condsumoffset!(S, P, K, A, f, cond)

Compute the adjusted sum `S` and number of contributing values `K` for
reductions of `f` mapped on array `A` over dimensions.
`S`, `P`, and `K` must have identical indices, and either match `A`
or have singleton-dimensions for the dimensions that are being summed
over. Only values of A with cond(A[i])==true are included in the tallies
of `S` and `K`. `P` is the first good value in each A set of `A` values
to be summed.

By subtracting off the first good value `P` accumulated to each sum,
condsumoffset! improves precision in case the values of A have a nearly
constant offset. The improved precision is beneficial if `K`×`P` is so large
that interesting variations of A are truncated in the sum S. However, precision
of this implementation will be worse than condsum! if the first good value is an
outlier farther from zero than the mean of the values of `A` summed.

The value mean computed in `condmeanprecise` is S./K+P.

See also `condsum!` and `condmeanoffset`.
"""
function condsumoffset!{T,N}(S, P, K, A::AbstractArray{T,N}, f, cond)
  #loops over dimensions of A _then_ determines dimensions of reduced S
    check_reducedims(S, A)
    isempty(A) && return S, K
    indices(S) == indices(K) || throw(DimensionMismatch("S and K must have identical indices"))

    indsAt, indsSt = safe_tail(indices(A)), safe_tail(indices(S))
    keep, Idefault = _newindexer(indsAt, indsSt)
    if reducedim1(S, A)
        # keep the accumulators as a local variable when reducing along the first dimension
        i1 = first(indices1(S))
        @inbounds for IA in CartesianRange(indsAt)
            IS = newindex(IA, keep, Idefault)
            s, p, k = S[i1,IS], P[i1,IS], K[i1,IS]
            #isfirst = true
            for i in indices(A, 1)
                tmp = A[i, IA]
                if cond(tmp)
                    if k==0 #isfirst
                        p = f(tmp)
                        #isfirst = false
                    else
                        s += f(tmp)-p
                    end
                    k += 1
                end
            end
            S[i1,IS], P[i1,IS], K[i1,IS] = s, p, k
        end
    else
        @inbounds for IA in CartesianRange(indsAt)
            IS = newindex(IA, keep, Idefault)
            #isfirst = true
            for i in indices(A, 1)
                tmp = A[i, IA]
                if cond(tmp)
                    if K[i, IS]==0 #isfirst
                        P[i, IS] = f(tmp)
                        #isfirst = false
                    else
                        S[i, IS] += f(tmp)-P[i, IS]
                    end
                    K[i, IS] += 1
                end
            end
        end
    end
    S, P, K
end
#=
if testing for k==0 is slow, then an alternative that works in many cases for
nearly constant offsets across the whole of A would be to just subtract off the
first good value in all of A. This would be good for surface pressure, but not
good for a significant gradient of A, like atmospheric pressure from the surface
to the stratosphere, where p varies by an order of magnitude.
Another solution, that I don't think will be faster, is to make a boolean array
`isfirst` like `K` but private to condsumoffset that stores whether the offset
has been recorded in `P`.
=#

# NOT TESTED
# precise, but requires two passes of condmean through A
# instead could write a version of condsum that accumulates multiple outputs of f, e.g. f(x)=[x, x*x]
"""
`condvar(A, cond, region; m=knownmean, corrected=true)` gives the unbiased
estimate of the variance of A (normalized by N-1) along dims in region.
`corrected=false` gives the variance normalized by N.
"""
condvar(A,cond,region; m=condmean(A,cond,region), corrected::Bool=true) =       condmean(x->abs(x)*abs(x), broadcast(-,A,m), cond, region, Km1=corrected)
condstd(A,cond,region; m=condmean(A,cond,region), corrected::Bool=true) = sqrt.(condmean(x->abs(x)*abs(x), broadcast(-,A,m), cond, region, Km1=corrected))
# corrected=true normalizes variance by K-1

"nanmean(X,r) mean of X long dimension(s) r, ignoring NaNs."
nanmean(X,r)=condmean(X,!(isnan),r)
"nanstd(X,r) standard deviation of X along dimension(s) r, ignoring NaNs."
nanstd(X,r)=sqrt.(condmean(x->abs(x)*abs(x),X.-nanmean(X,r),!(isnan),r))
end # module


using Test
using ConstraintSolver
const CS = ConstraintSolver

import MathOptInterface
const MOI = MathOptInterface

import ForwardDiff
using JuMP
using LinearAlgebra: dot

@testset "StableSet" begin
    @testset "Weighted stable set" begin
        matrix = [
            0 1 1 0
            1 0 1 0
            1 1 0 1
            0 0 1 0
        ]
        model = CS.Optimizer(logging = [], keep_logs = true)
        x = [MOI.add_constrained_variable(model, MOI.ZeroOne()) for _ in 1:4]
        for i in 1:4, j in 1:4
            if matrix[i, j] == 1 && i < j
                (z, _) = MOI.add_constrained_variable(model, MOI.GreaterThan(0.0))
                MOI.add_constraint(model, z, MOI.Integer())
                MOI.add_constraint(model, z, MOI.LessThan(1.0))
                f = MOI.ScalarAffineFunction(
                    [
                        MOI.ScalarAffineTerm(1.0, x[i][1]),
                        MOI.ScalarAffineTerm(1.0, x[j][1]),
                        MOI.ScalarAffineTerm(1.0, z),
                    ],
                    0.0,
                )
                MOI.add_constraint(model, f, MOI.EqualTo(1.0))
            end
        end
        weights = [0.2, 0.1, 0.2, 0.1]
        terms = [MOI.ScalarAffineTerm(weights[i], x[i][1]) for i in eachindex(x)]
        objective = MOI.ScalarAffineFunction(terms, 0.0)
        MOI.set(model, MOI.ObjectiveFunction{typeof(objective)}(), objective)
        MOI.set(model, MOI.ObjectiveSense(), MOI.MAX_SENSE)
        MOI.optimize!(model)
        var_x = [x[i][1] for i in 1:4]
        CS.save_logs(CS.get_inner_model(model), "stable_set.json", :x => var_x)
        rm("stable_set.json")

        @test MOI.get(model, MOI.VariablePrimal(), x[4][1]) == 1
        @test MOI.get(model, MOI.VariablePrimal(), x[1][1]) == 1
        @test MOI.get(model, MOI.ObjectiveValue()) ≈ 0.3
    end

    @testset "Weighted stable set JuMP" begin
        matrix = [
            0 1 1 0
            1 0 1 0
            1 1 0 1
            0 0 1 0
        ]
        m = Model(CSJuMPTestOptimizer())
        x = @variable(m, x[1:4], Bin)
        for i in 1:4, j in (i + 1):4
            if matrix[i, j] == 1
                zcomp = @variable(m)
                JuMP.set_binary(zcomp)
                @constraint(m, x[i] + x[j] + zcomp == 1)
            end
        end
        w = [0.2, 0.1, 0.2, 0.1]
        @objective(m, Max, dot(w, x))
        optimize!(m)
        @test JuMP.value(x[4]) ≈ 1
        @test JuMP.value(x[1]) ≈ 1
        @test JuMP.value(x[2]) ≈ 0
        @test JuMP.value(x[3]) ≈ 0
        @test JuMP.objective_value(m) ≈ 0.3
    end

    @testset "Bigger stable set using <= " begin
        matrix = [
            0 1 0 1 0 0 1 0 0 1
            1 0 1 0 1 0 0 0 0 0
            0 1 0 1 0 0 0 0 1 0
            1 0 1 0 1 1 0 0 0 0
            0 1 0 1 0 1 1 1 0 0
            0 0 0 1 1 0 1 0 0 0
            1 0 0 0 1 1 0 1 0 0
            0 0 0 0 1 0 1 0 1 0
            0 0 1 0 0 0 0 1 0 1
            1 0 0 0 0 0 0 0 1 0
        ]
        n = 10
        m = Model(CSJuMPTestOptimizer())
        @variable(m, x[1:n], Bin)
        nconstraints = 0
        for i in 1:n, j in (i + 1):n
            if matrix[i, j] == 1
                @constraint(m, x[i] + x[j] <= 1)
                nconstraints += 1
            end
        end
        w = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
        @objective(m, Max, dot(w, x))
        optimize!(m)
        com = CS.get_inner_model(m)
        @test is_solved(com)
        @test com.info.n_constraint_types.inequality == length(com.constraints)
        @test com.info.n_constraint_types.inequality == nconstraints

        @test MOI.get(m, MOI.ObjectiveValue()) ≈ 2.7
    end

    @testset "Bigger stable set using with Cbc" begin
        matrix = [
            0 1 0 1 0 0 1 0 0 1
            1 0 1 0 1 0 0 0 0 0
            0 1 0 1 0 0 0 0 1 0
            1 0 1 0 1 1 0 0 0 0
            0 1 0 1 0 1 1 1 0 0
            0 0 0 1 1 0 1 0 0 0
            1 0 0 0 1 1 0 1 0 0
            0 0 0 0 1 0 1 0 1 0
            0 0 1 0 0 0 0 1 0 1
            1 0 0 0 0 0 0 0 1 0
        ]
        n = 10
        cbc_optimizer = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
        m = Model(optimizer_with_attributes(
            CS.Optimizer,
            "lp_optimizer" => cbc_optimizer,
            "logging" => [],
        ))
        @variable(m, x[1:n], Bin)
        nconstraints = 0
        for i in 1:n, j in (i + 1):n
            if matrix[i, j] == 1
                @constraint(m, x[i] + x[j] <= 1)
                nconstraints += 1
            end
        end
        w = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
        @objective(m, Max, dot(w, x))
        optimize!(m)
        com = CS.get_inner_model(m)
        @test com.info.n_constraint_types.inequality == length(com.constraints)
        @test com.info.n_constraint_types.inequality == nconstraints

        @test MOI.get(m, MOI.ObjectiveValue()) ≈ 2.7
    end

    function weighted_stable_set(w)
        matrix = [
            0 1 1 0
            1 0 1 0
            1 1 0 1
            0 0 1 0
        ]
        model = CS.Optimizer(solution_type = Real, logging = [])
        x = [MOI.add_constrained_variable(model, MOI.ZeroOne()) for _ in 1:4]
        for i in 1:4, j in 1:4
            if matrix[i, j] == 1 && i < j
                (z, _) = MOI.add_constrained_variable(model, MOI.GreaterThan(0.0))
                MOI.add_constraint(model, z, MOI.Integer())
                MOI.add_constraint(model, z, MOI.LessThan(1.0))
                f = MOI.ScalarAffineFunction(
                    [
                        MOI.ScalarAffineTerm(1.0, x[i][1]),
                        MOI.ScalarAffineTerm(1.0, x[j][1]),
                        MOI.ScalarAffineTerm(1.0, z),
                    ],
                    0.0,
                )
                MOI.add_constraint(model, f, MOI.EqualTo(1.0))
            end
        end
        # sum(w[i] * x[i] for i in V) - stable_set == 0
        terms = [MOI.ScalarAffineTerm(w[i], x[i][1]) for i in eachindex(x)]
        objective = MOI.ScalarAffineFunction(terms, zero(eltype(w)))
        MOI.set(model, MOI.ObjectiveFunction{typeof(objective)}(), objective)
        MOI.set(model, MOI.ObjectiveSense(), MOI.MAX_SENSE)
        MOI.optimize!(model)
        return MOI.get(model, MOI.ObjectiveValue())
    end

    # type piracy for the greater good
    # see https://github.com/JuliaDiff/ForwardDiff.jl/issues/318
    function Base.typemin(::Type{ForwardDiff.Dual{T,V,N}}) where {T,V,N}
        ForwardDiff.Dual{T,V,N}(typemin(V))
    end

    @testset "Differentiating stable set" begin
        weights = [0.2, 0.1, 0.2, 0.1]
        # ∇w = Zygote.gradient(weighted_stable_set, weights)
        ∇w = ForwardDiff.gradient(weighted_stable_set, weights)
        @test ∇w[1] ≈ 1
        @test ∇w[4] ≈ 1
        @test ∇w[2] ≈ ∇w[3] ≈ 0
    end
end


"""Polynomial regression interface

Taken from https://juliastats.github.io/StatsModels.jl/stable/internals/#extending-1
"""

module Polynomials

using StatsModels
using StatsBase

poly(x, n) = x.^n

const POLY_CONTEXT = Any

struct PolyTerm{T,D} <: AbstractTerm
    term::T
    deg::D
end

Base.show(io::IO, p::PolyTerm) = print(io, "poly($(p.term), $(p.deg))")

poly(t::Symbol, d::Int) = PolyTerm(term(t), term(d))

function StatsModels.apply_schema(t::FunctionTerm{typeof(poly)},
                                  sch::StatsModels.Schema,
                                  Mod::Type{<:POLY_CONTEXT})
    apply_schema(PolyTerm(t.args_parsed...), sch, Mod)
end

function StatsModels.apply_schema(t::PolyTerm,
                                  sch::StatsModels.Schema,
                                  Mod::Type{<:POLY_CONTEXT})
    term = apply_schema(t.term, sch, Mod)
    isa(term, ContinuousTerm) ||
        throw(ArgumentError("PolyTerm only works with continuous terms (got $term)"))
    isa(t.deg, ConstantTerm) ||
        throw(ArgumentError("PolyTerm degree must be a number (got $t.deg)"))
    PolyTerm(term, t.deg.n)
end

function StatsModels.modelcols(p::PolyTerm, d::NamedTuple)
    col = modelcols(p.term, d)
    reduce(hcat, [col.^n for n in 1:p.deg])
end

StatsModels.terms(p::PolyTerm) = terms(p.term)
StatsModels.termvars(p::PolyTerm) = StatsModels.termvars(p.term)
StatsModels.width(p::PolyTerm) = p.deg

StatsBase.coefnames(p::PolyTerm) = coefnames(p.term) .* "^" .* string.(1:p.deg)

export poly
end