module NetSci02

using LightGraphs
using Distributions
using Loess

include("utils.jl")
include("tests.jl")

"""
    samplepercolations(measure::Function, graph::Graph, node_distributions::Function,
                       repetitions::Integer, stepsize = 0.1)

Sample `repetitions` node removals on `graph`, over deletion rates `0.0:stepsize:1.0`, calculating
`measure` each time.  The probability of deleting a nodes is given by `node_distributions(φ)`
for each φ (so `node_distributions` must have type Float -> Vector{Bernoulli}).
"""
function samplepercolations(measure::Function,
                            graph::Graph,
                            node_distributions::Function,
                            repetitions::Integer,
                            stepsize = 0.05,
                            result_type = typeof(measure(Graph(1, 0))))
    const vertices = nv(graph)
    const range = 0.0:stepsize:1.0
    curves = SharedArray(result_type, repetitions, length(range))

    pmap(1:repetitions) do r
        for (i, φ) in enumerate(range)
            victim = copy(graph)
            distributions = node_distributions(φ)

            for v = vertices:-1:1
                if rand(distributions[v]) == 1
                    rem_vertex!(victim, v)
                end
            end

            curves[r, i] = measure(victim)
        end
    end

    return curves
end

"""
    informed_probabilities(graph::Graph, weights::Vector{Float64}, infos::Vector{Symbol})

Construct a paramtrized vector of distributions over nodes which integrates knowledge about local
clustering, degree, and second-order neighbourhood, taking into account `weights`.  `infos` can
be used to choose which information to include.
"""
function informed_probabilities(graph::Graph,
                                weights::Vector{Float64},
                                infos::Vector{Symbol})::Function
    const maxdegrees = nv(graph) * (nv(graph) + 1) / 2

    # wherever meaningful, there is a normalized (*_n) and an unnormalized variant
    const information_functions = Dict(
        :degree => g -> degree(g),
        :degree_n => g -> degree(g) / maxdegrees,
        :inv_degree => g -> 1 ./ (degree(g) + eps()),
        :inv_degree_n => g -> (1 ./ (degree(g) + eps())) / maxdegrees,
        :local_clustering => g -> local_clustering_coefficient(g),
        :neighborhood2 => g -> map(v -> length(neighborhood(g, v, 2)), vertices(g)),
        :neighborhood2_n => g -> map(v -> length(neighborhood(g, v, 2)), vertices(g)) / maxdegrees,
        :neighborhood3 => g -> map(v -> length(neighborhood(g, v, 3)), vertices(g)),
        :neighborhood3_n => g -> map(v -> length(neighborhood(g, v, 3)), vertices(g)) / maxdegrees
    )

    information = Array(Float64, nv(graph), length(infos))
    for (i, info) in enumerate(infos)
        information[:, i] = information_functions[info](graph)
    end

    return makedistributions(information * weights)
end

"""
    uninformed_probabilities(graph::Graph)

Like `informed_probabilities`, but just use the same distribution based on `φ` for every node.
"""
function uninformed_probabilities(graph::Graph)
    return φ -> Bernoulli.(fill(φ, nv(graph)))
end

"""
    makedistributions(logits::Vector{Float64})::Function

Convert a vector of weight values for every node (`logits`) into a parametrized Bernoulli
distribution for each node, such that the expected fraction of nodes after using this for
percolation depends on the parameter (but the individual probabilities may differ).
"""
function makedistributions(logits::AbstractArray{Float64, 1})::Function
    probabilities = logistic.(logits)
    return function (φ)
        scaled_probabilities = probabilities .* (φ / mean(probabilities))
        Bernoulli.(clamp(scaled_probabilities, 0.0, 1.0)) # because rounding errors
    end
end

"""
    smoothcurves(curves::AbstractArray{Float64, 2})

Return the smoothed average of the given `curves` (using Loess regression).
"""
function smoothcurves{T<:Number}(curves::AbstractArray{T, 2}, use_loess = true)
    average_curve = squeeze(mapslices(mean, curves, 1), 1)

    if use_loess
        stepsize = 1 / (size(curves)[end] - 1)
        range = 0.0:stepsize:1.0
        return predict(loess(range, average_curve), range)
    else
        return average_curve
    end
end

"""
    normalizecurves(curves::AbstractArray{Float64, 2})

Normalize all curves to the range [0, 1] by their maximum.
"""
normalizecurves{T<:Number}(curves::AbstractArray{T, 2}) = curves ./ mapslices(maximum, curves, 2)

"""
    percolation_area{T<:Number}(curves::AbstractArray{T, 2}, smooth = true)

Estimate the area under the sampled percolation curves by averaging and calculating the area.
If `smooth` is true, uses Loess regression and numerical integration, otherwise simply sum up and
weight by the stepsize.
"""
function percolation_area{T<:Number}(curves::AbstractArray{T, 2}, smooth = true)
    normalized_curves = normalizecurves(curves)
    samples = size(curves)[1]
    stepsize = 1 / (size(curves)[end] - 1)

    if smooth
        # numerical integration over a local regression curve
        average_curve = squeeze(mapslices(mean, normalized_curves, 1), 1)
        smoothed_curve = loess(0.0:stepsize:1.0, average_curve)
        return quadgk(x -> predict(smoothed_curve, x), 0.0, 1.0)[1]
    else
        # just take the average of everything...
        return sum(normalized_curves) * stepsize / samples
    end
end

"""
    max_component_size(g::Graph)

Calculate the size of the larges connected component of `g`.
"""
max_component_size(g::Graph) = reduce(max, 0, length.(connected_components(g)))

"""
    percolation_loss(graph::Graph, infos::Vector{Symbol}, samples, stepsize)

Construct a loss function which samples the expected area of the percolation curve of `graph`,
using the parameter as weightings for local information.
"""
function percolation_loss(graph::Graph, infos::Vector{Symbol}, samples = 10, stepsize = 0.05,
                          measure = max_component_size)
    function loss(parameter)
        dist = informed_probabilities(graph, parameter, infos)
        curves = samplepercolations(measure, graph, dist, samples, stepsize)
        return percolation_area(curves)
    end
end

####################################################################################################
# TRAINING AND TESTING
####################################################################################################

function log_annealing(;kwargs...)
    args = Dict(kwargs)
    # push!(error_trace, args[:f_new])

    if args[:steps] % 100 == 0
        println(args[:steps], " steps, T = ", args[:T])
        println("\tParameters: ", args[:x_new])
        println("\tError: ", args[:f_new])
    end
end

function trainandtest(name::String, train_graph::Graph, test_graph::Graph,
                      results_file::String, variant::String;
                      mode = "a", parameters...)
    for (s, v) in parameters
        @eval $s = $v
    end

    println("Training on ", name, " using ", variant, "...")
    trained_param = annealing(percolation_loss(train_graph, infos, samples, stepsize),
                              fill(0.0, length(infos)),
                              init_temp, temp_decay, exploration_size,
                              (s, v) -> s > iterations;
                              debug_callback = log_annealing)
    println("Learned parameters (", name, "): ", trained_param, "\n")

    baseline_dist = uninformed_probabilities(test_graph)
    optimized_dist = informed_probabilities(test_graph, trained_param, infos)

    baseline_samples = samplepercolations(max_component_size, test_graph,
                                          baseline_dist, samples, stepsize)
    optimized_samples = samplepercolations(max_component_size, test_graph,
                                           optimized_dist, samples, stepsize)

    baseline_curve = normalizecurves(baseline_samples)
    optimized_curve = normalizecurves(optimized_samples)

    range = 0.0:stepsize:1.0
    open(results_file, mode) do f
        for (x, y, y2) in zip(range,
                              smoothcurves(baseline_curve, false),
                              smoothcurves(baseline_curve, true))
            write(f, "$name $variant baseline $x $y $y2\n")
        end

        for (x, y, y2) in zip(range,
                          smoothcurves(optimized_curve, false),
                          smoothcurves(optimized_curve, true))
            write(f, "$name $variant optimized $x $y $y2\n")
        end
    end
end

function main()
    fb1 = readnetwork("../data/facebook1.txt")
    fb2 = readnetwork("../data/facebook2.txt")
    austin = readnetwork("../data/austin.txt")
    philadelphia = readnetwork("../data/philadelphia.txt")
    random1 = Graph(2000, 5000)
    random2 = Graph(2000, 5000)

    examples = ["facebook" => (fb1, fb2),
                "traffic" => (austin, philadelphia),
                "random" => (random1, random2)]

    info_variants = Dict("d" => [:degree],
                         "i" => [:inv_degree],
                         "l" => [:local_clustering],
                         "n" => [:neighborhood2],
                         "dl" => [:degree, :local_clustering],
                         "dln" => [:degree, :local_clustering, :neighborhood2])

    info_variants_n = Dict("i" => [:inv_degree_n],
                           "d" => [:degree_n],
                           "l" => [:local_clustering],
                           "n" => [:neighborhood2_n],
                           "dl" => [:degree_n, :local_clustering],
                           "dln" => [:degree_n, :local_clustering, :neighborhood2_n])

    training_parameters = Dict(
        :stepsize => 0.05,
        :samples => 10,
        :init_temp => 10.0,
        :temp_decay => 0.9,
        :exploration_size => 0.05,
        :iterations => 400)

    for (variant, infos) in info_variants
        for (name, (train, test)) in examples
            trainandtest(name, train, test,
                         "../evaluation/results1.txt", variant;
                         mode = "a", infos = infos, training_parameters...)
        end
    end

    # this turns out to be worse...
    # for (variant, infos) in info_variants_n
    #     for (name, (train, test)) in examples
    #         trainandtest(name, train, test,
    #                      "../evaluation/results2.txt", variant;
    #                      mode = "a", infos = infos, training_parameters...)
    #     end
    # end
end

end


@testset "Midpoints" begin
    @test_throws Exception Midpoints{Int}(2)
    @test eltype(@inferred(Midpoints(2))) == Float64

    for T in Floats
        @test eltype(@inferred(Midpoints{T}(3))) == T
        @test typeof(@inferred(Midpoints{T}(3)[1])) == T

        @test Midpoints{T}(0) ≈ T[]
        @test Midpoints{T}(1) ≈ T[0]  atol = eps(real(T))
        @test Midpoints{T}(2) ≈ T[-0.5,0.5]
        @test Midpoints{T}(3) ≈ T[-T(2)/3,0,T(2)/3]
        @test Midpoints{T}(4) ≈ T[-0.75,-0.25,0.25,0.75]
    end
end


#%%

using Plots
gr()
# %%
xs=-pi:0.1:pi
anim = @animate for i=1:100
    plot(xs,sin.(i*xs));
end
gif(anim, "anim_fps15.gif", fps = 15)
gif(anim, "anim_fps30.gif", fps = 30)

# %%
@gif for i=1:100
    xs=-pi:0.1:pi
    plot(xs,sin.(i*xs));
end every 10


function Point2ID(X)
    Y = 0
    for ii = 1:length(X)
        Y += X[ii]*2^(ii-1);
    end
    Y
end

function GenNeighbor(X)
    N = length(X)
    Y = zeros(N)
    for ii = 1:N
        temp = copy(X)
        temp[ii] = 1-temp[ii]
        Y[ii] = Point2ID(temp) + 1
    end
    sort!(Y)
    Y
end

function GenCuts(X)
    N = length(X)
    Y = zeros(N)
    for ii = 1:N
        Y[ii] = ii + X[ii]*N
    end
    sort!(Y)
end

function Int2Bin(X,N)
    Y = zeros(N)
    for ii = N:-1:1
        Y[ii] = floor(X/2^(ii-1))
        X = X - Y[ii]*2^(ii-1)
    end
    Y
end
include("FindIntersectingPoint.jl")
include("Vert.jl")
include("PTVer.jl")


using Statistics
using Elefridge
using PyPlot
using StatsBase

N = 10000
m = 10000
q = 7

meanerror = fill(0f0,3,q,N)
linerror = fill(0f0,3,q,N)
logerror = fill(0f0,3,q,N)

for i in 1:N
    # sample various distributions
    for (s,D) = enumerate([ rand(Float32,m),
                            4*rand(Float32,m),
                            rand(Float32,m).+1,
                            randn(Float32,m),
                            4*randn(Float32,m),
                            randn(Float32,m).+1,
                            exp.(randn(Float32,m))])

        Dround = round(D,7)
        Dgroom = groom(D,7)
        Dshave = shave(D,7)

        # mean error
        meanerror[1,s,i] = abs(mean(D - Dround))
        meanerror[2,s,i] = abs(mean(D - Dgroom))
        meanerror[3,s,i] = abs(mean(D - Dshave))

        # max lin error
        linerror[1,s,i] = maximum(abs.(D - Dround))
        linerror[2,s,i] = maximum(abs.(D - Dgroom))
        linerror[3,s,i] = maximum(abs.(D - Dshave))

        # max log error
        logerror[1,s,i] = maximum(abs.(log10.(D./Dround)))
        logerror[2,s,i] = maximum(abs.(log10.(D./Dgroom)))
        logerror[3,s,i] = maximum(abs.(log10.(D./Dshave)))
    end
 end

## fit histogram
ion()
pygui(true)
fig,(ax1,ax2,ax3) = subplots(1,3,figsize=(10,4),sharey=true)

lwmm = 1
lw10 = 3
lw25 = 8

distr = [   "Uniform(0,1)","Uniform(0,4)","Uniform(1,2)",
            "Normal(0,1)","Normal(0,4)","Normal(1,1)",
            "Lognormal(0,1)"]

for (ax,error) in zip([ax1,ax2,ax3],[meanerror,linerror,logerror])
    for i in 1:q
        for (s,e) in enumerate([error[1,i,:],error[2,i,:],error[3,i,:]])
            y = i-s/4
            c = "C"*string(s-1)
            ax.semilogx([minimum(e),maximum(e)],[y,y],c,lw=lwmm)
            p10 = percentile(e,10)
            p25 = percentile(e,25)
            p50 = percentile(e,50)
            p75 = percentile(e,75)
            p90 = percentile(e,90)
            ax.semilogx([p10,p90],[y,y],c,lw=lw10)
            ax.semilogx([p25,p75],[y,y],c,lw=lw25)
            ax.scatter(p50,y,100,c,edgecolor="k",zorder=5)
            ax.axhline(i,color="k",lw=0.5)
        end
    end
end

ax1.set_yticks([1,2,3,4,5,6,7].-0.5)
ax1.yaxis.set_tick_params(length=0)
ax2.yaxis.set_tick_params(length=0)
ax3.yaxis.set_tick_params(length=0)
ax1.set_yticklabels(distr)
ax1.set_ylim(-1.75,7)

ax1.set_xlim(1e-6,1e-2)
ax2.set_xlim(1e-3,1e0)
ax3.set_xlim(1e-3,1e0)

e = meanerror[2,1,:]
p10 = percentile(e,10)
p25 = percentile(e,25)
p50 = percentile(e,50)
p75 = percentile(e,75)
p90 = percentile(e,90)
maxx = maximum(e)
ax1.text(p10,0.1,"10%",rotation=90,ha="center",va="top",fontsize=9)
ax1.text(p25,0.1,"25%",rotation=90,ha="center",va="top",fontsize=9)
ax1.text(p50,0.1,"median",rotation=90,ha="center",va="top",fontsize=9)
ax1.text(p75,0.1,"75%",rotation=90,ha="center",va="top",fontsize=9)
ax1.text(p90,0.1,"90%",rotation=90,ha="center",va="top",fontsize=9)
ax1.text(maxx,0.1,"max",rotation=90,ha="center",va="top",fontsize=9)

ax2.plot(0,0,"C0",lw=8,label="Round-to-nearest")
ax2.plot(0,0,"C1",lw=8,label="Bit-grooming")
ax2.plot(0,0,"C2",lw=8,label="Bit-shaving")
ax2.legend(loc=4)

ax1.set_xlabel("absolute mean error")
ax2.set_xlabel("maximum absolute error")
ax3.set_xlabel("maximum decimal error")

ax1.set_title("Mean error",loc="left")
ax2.set_title("Absolute error",loc="left")
ax3.set_title("Decimal error",loc="left")

ax1.set_title("a",loc="right",fontweight="bold")
ax2.set_title("b",loc="right",fontweight="bold")
ax3.set_title("c",loc="right",fontweight="bold")

tight_layout()


module Williamson2ShallowWaterIMEXTests

using FillArrays
using Test
using WriteVTK
using Gridap
using GridapGeosciences

include("Williamson2InitialConditions.jl")

# Solves the steady state Williamson2 test case for the shallow water equations on a sphere
# of physical radius 6371220m. Involves a modified coriolis term that exactly balances
# the potential gradient term to achieve a steady state
# reference:
# D. L. Williamson, J. B. Drake, J. J.HackRüdiger Jakob, P. N.Swarztrauber, (1992)
# J Comp. Phys. 102 211-224

l2_err_u = [0.009874787793428196, 0.0028512787026009232, 0.0007426759094111235]
l2_err_h = [0.005610917625840937, 0.0014555948928746177, 0.0003683350733739   ]

order  = 1
degree = 4

for i in 1:3
  n      = 2*2^i
  nstep  = 5*n
  Uc     = sqrt(g*H₀)
  dx     = 2.0*π*rₑ/(4*n)
  dt     = 0.05*dx/Uc
  println("timestep: ", dt)   # gravity wave time step

  model = CubedSphereDiscreteModel(n; radius=rₑ)
  hf, uf = shallow_water_imex_time_stepper(model, order, degree,
                                      h₀, u₀, f₀, g,
                                      dt, 0.0, nstep;
                                      write_solution=false,
                                      write_solution_freq=5,
                                      write_diagnostics=true,
                                      write_diagnostics_freq=1,
                                      dump_diagnostics_on_screen=true)

  Ω     = Triangulation(model)
  dΩ    = Measure(Ω, degree)
  hc    = CellField(h₀, Ω)
  e     = h₀-hf
  err_h = sqrt(sum(∫(e⋅e)*dΩ))/sqrt(sum(∫(hc⋅hc)*dΩ))
  uc    = CellField(u₀, Ω)
  e     = u₀-uf
  err_u = sqrt(sum(∫(e⋅e)*dΩ))/sqrt(sum(∫(uc⋅uc)*dΩ))
  println("n=", n, ",\terr_u: ", err_u, ",\terr_h: ", err_h)

  @test abs(err_u - l2_err_u[i]) < 10.0^-12
  @test abs(err_h - l2_err_h[i]) < 10.0^-12
end

end


"""
Create normal budget distribution for each product with mean given by
the number of resources in the product.
"""
function normal_budgets_per_resource(P::AbstractArray{<:Product}, resource_budget_μ ::Float64, σ::Float64)
    P_sizes = Float64[sum(p) for p in P]
    budget_means = P_sizes .* resource_budget_μ
    [Normal(μ, σ) for μ in budget_means]
end


module LinearKernelModels

using OffsetArrays, LinearAlgebra

export solve_for_kernel, compute_Cd, compute_r, compute_kerr

"""
    C, d = compute_Cd(S::AbstractVecOrMat, r::AbstractVector, trange::AbstractUnitRange)
    C, d = compute_Cd(S::AbstractVecOrMat, r::AbstractVector, nt::Integer)

Given one or more stimuli, where `S[t, i]` is the strength of the
`i`th stimulus at time `t`, and a response vs. time `r[t]`, calculate
`C` and `d` that can be used to solve for a kernel-matrix `k`, where
`k` convolved with `S` best reconstructs `r`. In particular,

    k = reshape(C \\ d, (nt, nstim))

`nt` (an integer) is the number of timepoints in the kernel.
If the kernel is allowed to look into `r`'s future, instead supply `trange = tpast:tfuture`
where `tpast` might typically be negative and `tfuture` typically positive.
(The meaning is the offset relative to the current moment `t=0`.)

If you have temporal gaps in your measured response, splice the blocks
together into a single vector, separating the blocks by stretches of
NaNs of length `nt` or larger. All entries in `S` should be
finite and represent the actual value of the stimulus at these
times. If you have periods of uncertainty about the stimulus history,
then for any time points within `nt` after these unknown
blocks, replace the corresponding entries in `r` with NaN. You can
then set the unknown entries in `S` to an arbitrary value (e.g.,
zero).

See also [`compute_r`](@ref), [`compute_kerrs`](@ref).
"""
function compute_Cd(S::AbstractVecOrMat, r::AbstractVector, trange::AbstractUnitRange)
    T, nstim = size(S, 1), size(S, 2)
    axes(S, 1) == 1:T || error("S must start indexing at t=1, got $(axes(S))")
    axes(r) == (1:T,) || error("r must agree with S's indexing, got $(axes(r)) versus $(axes(S))")
    kernelaxes = (trange, oftype(trange, Base.OneTo(nstim)))
    L = prod(length.(kernelaxes))
    C = zeros(L, L)
    LI = LinearIndices(OffsetArrays.IdentityUnitRange.(kernelaxes))
    w = (!).(isnan.(r))
    for i1 = 1:nstim, i2 = i1:nstim
        for τ1 in trange, τ2 in trange
            s = 0.0
            for t = max(1, 1-τ1, 1-τ2):min(T, T-τ1, T-τ2)
                s += w[t] * S[t+τ1, i1] * S[t+τ2, i2]
            end
            ind1, ind2 = LI[τ1, i1], LI[τ2, i2]
            C[ind1, ind2] = C[ind2, ind1] = s
        end
    end
    d = zeros(L)
    for i = 1:nstim
        for τ in trange
            s = 0.0
            for t = max(1, 1-τ):min(T, T-τ)
                p = r[t]*S[t+τ, i]
                s += ifelse(w[t], p, zero(p))
            end
            d[LI[τ, i]] = s
        end
    end
    return C, d
end
compute_Cd(S::AbstractVecOrMat, r::AbstractVector, nt::Integer) =
    compute_Cd(S, r, -(nt-1):0)

"""
    r = compute_r(S::AbstractVecOrMat, k::AbstractVecOrMat)

Given one or more stimuli, where `S[t, i]` is the strength of the
`i`th stimulus at time `t`, and a kernel `k[τ, i]` for each
stimulus, compute the summed-convolution between `S` and `k` to
predict the response `r`.
"""
function compute_r(S, k)
    T, nstim = size(S, 1), size(S, 2)
    trange = axes(k, 1)
    @assert(axes(k, 2) == 1:nstim)
    r = zeros(T)
    for i = 1:nstim
        for τ in trange
            kτi = k[τ,i]
            for t = max(1, 1-τ):min(T, T-τ)
                r[t] += S[t+τ,i]*kτi
            end
        end
    end
    return r
end

"""
    k = solve_for_kernel(S, r, isnz; rtol=sqrt(eps()))

Solve for a kernel, imposing constraints on entries that are allowed to be non-zero.
`S` is the stimulus and `r` the response vector. (See [`compute_Cd`](@ref) for more information.)
`isnz` must have the axes of the kernel, and be true in every slot where the solved kernel
may be non-zero. If you have no need for constraints, supply `trues(axs)` where `axs`
encodes the axes of the kernel.

# Example
```
julia> using LinearKernelModels, OffsetArrays

julia> S = rand(100,3);

julia> ktrue = OffsetArray(rand(5, 3), -4:0, 1:3)
OffsetArray(::Array{Float64,2}, -4:0, 1:3) with eltype Float64 with indices -4:0×1:3:
 0.155762  0.877222  0.457633
 0.88788   0.355702  0.0268748
 0.276609  0.914133  0.230344
 0.718295  0.158363  0.0402396
 0.195513  0.783832  0.5007

julia> ktrue[-4, 1] = 0
0

julia> isnz = ktrue .!= 0
OffsetArray(::BitArray{2}, -4:0, 1:3) with eltype Bool with indices -4:0×1:3:
 false  true  true
  true  true  true
  true  true  true
  true  true  true
  true  true  true

julia> r = compute_r(S, ktrue);

julia> k = solve_for_kernel(S, r, isnz)
OffsetArray(::Array{Float64,2}, -4:0, 1:3) with eltype Float64 with indices -4:0×1:3:
 6.99441e-15  0.877222  0.457633
 0.88788      0.355702  0.0268748
 0.276609     0.914133  0.230344
 0.718295     0.158363  0.0402396
 0.195513     0.783832  0.5007
```
"""
function solve_for_kernel(S::AbstractVecOrMat, r::AbstractVector, isnz::AbstractVecOrMat{Bool}; rtol=sqrt(eps()))
    trange = axes(isnz, 1)
    C, d = compute_Cd(S, r, trange)
    nconstrained = prod(length.(axes(isnz))) - sum(isnz)
    Q = zeros(size(C,1), nconstrained)
    col = 0
    for (i, isunconstrained) in enumerate(isnz)
        if !isunconstrained
            col += 1
            Q[i,col] = 1
        end
    end
    A = [C Q; Q' zeros(nconstrained, nconstrained)]
    dcat = [d; zeros(nconstrained)]
    F = svd(A)
    S = F.S
    Smax = maximum(S)
    flag = S .< rtol*Smax
    S[flag] .= Inf
    soln = F \ dcat
    reshape(soln[1:length(d)], axes(isnz))
end

"""
    kerr = compute_kerr(k, C, S, r)

Compute the estimated standard error in the parameters `k` of the
kernels, where `C` is determined from [`compute_Cd`](@ref), `S` is the
stimulus, and `r` the response.
"""
function compute_kerr(k, C, S, r)
    rbar = compute_r(S, k)
    N = 0
    Σ = 0.0
    for t = 1:length(r)
        rt = r[t]
        if !isnan(rt)
            Σ += (rt - rbar[t])^2
            N += 1
        end
    end
    lk = prod(length.(axes(k)))
    σ2 = Σ/(N-lk)  # replace with length(k)
    return reshape(sqrt(σ2) * sqrt.(diag(inv(C))), axes(k))
end

end # module


using LinearAlgebra
using Plots
using Random
using Ripserer
pgfplotsx()

function circ_rand(n, r, center)
    map(rand(n)) do φ
        (r * sin(2π * φ), r * cos(2π * φ)) .+ center
    end
end

function circ_equi(n, r, center)
    map(range(0, 1, length=n+1)[1:end-1]) do φ
        (r * sin(2π * φ), r * cos(2π * φ)) .+ center
    end
end

function annulus(n, r1, r2, center)
    points = Tuple{Float64,Float64}[]
    while length(points) < n
        v = 2r1 * (rand(2) .- 0.5)
        if r1 ≥ norm(v) > r2
            push!(points, (v[1], v[2]) .+ center)
        end
    end
    return points
end

Random.seed!(1337)
data = [
    circ_rand(20, 1, (0, 2));
    circ_equi(12, 1.5, (0, -2));
    annulus(40, 2, 1, (2 * √3, 0));
]

res_co = ripserer(data; reps=true)[2][end-2:end]
res_ho = ripserer(data; alg=:involuted)[2][end-2:end]

function plot_representatives(data, ints; kwargs...)
    plt = scatter(data; legend=false, aspect_ratio=1, markersize=2, kwargs...)
    foreach(ints) do int
        plot!(plt, int, data; threshold_strict=death(int))
    end
    return plt
end

plot(
    plot_representatives(data, res_co; title="Representative Cocycles"),
    plot_representatives(data, res_ho; title="Representative Cycles"),
    size=(600, 300),
)
savefig(joinpath(@__DIR__, "figures/cycle-cocycle-comparison.pdf"))


"""
    Nlplex{T <: Real} = Nilplexification{T}

A nilplex number is a nilplexification of the real numbers.
"""
const Nlplex{T <: Real} = Nilplexification{T}

function Nlplex(a::T, b::T) where T <: Real
    Nlplex{T}(a, b)
end

function Nlplex(a::T) where T <: Real
    Nlplex{T}(a, zero(T))
end

function Nlplex(a::Real, b::Real)
    Nlplex(promote(a, b)...)
end

"""
    unreal(z::Nlplex)

The unreal part of a nilplex number is a 1-vector.
"""
function unreal(z::Nlplex{T}) where T <: Real
    T[z.r]
end

function show(io::IO, z::Nlplex)
    print(io, "[1: ")
    print(io, z.l)
    print(io, ", A: ")
    print(io, z.r)
    print(io, "]")
end

function random(::Type{Nlplex{T}}) where T <: Real
    Nlplex{T}(rand(T), rand(T))
end

function random(::Type{Nlplex{BigInt}})
    Nlplex{BigInt}(randomBigInt(), randomBigInt())
end

function random(::Type{Nlplex{BigFloat}})
    Nlplex{BigFloat}(randomBigFloat(), randomBigFloat())
end

function random(::Type{Nlplex{Rational{BigInt}}})
    Nlplex{Rational{BigInt}}(randomBigRational(), randomBigRational())
end

function abs2(z::Nlplex)
    abs2(z.l)
end

function barA(z::Nlplex)
    Nlplex(z.l, -(z.r))
end

function iszerodivisor(z::Nlplex)
    iszero(z.l)
end

function inv(z::Nlplex)
    if iszerodivisor(z)
        error(ZeroDivisorInverse)
    end

    barA(z) / abs2(z)
end


module _FuzzEn
export FuzzEn
using Statistics: mean, std
using LinearAlgebra: UpperTriangular, I, inv
    """
        Fuzz, Ps1, Ps2 = FuzzEn(Sig)

    Returns the fuzzy entropy estimates `Fuzz` and the average fuzzy distances
    (`m`:Ps1, `m+1`:Ps2) for `m` = [1,2] estimated from the data sequence `Sig`
    using the default parameters: embedding dimension = 2, time delay = 1,
    fuzzy function (`Fx`) = "default", fuzzy function parameters (`r`) = [0.2, 2],
    logarithm = natural

        Fuzz, Ps1, Ps2 = FuzzEn(Sig::AbstractArray{T,1} where T<:Real; m::Int=2, tau::Int=1, r::Union{Real,Tuple{Real,Real}}=(.2,2), Fx::String="default", Logx::Real=exp(1))

    Returns the fuzzy entropy estimates `Fuzz` for dimensions = [1,...,`m`]
    estimated for the data sequence `Sig` using the specified keyword arguments:

    # Arguments:
    `m`     - Embedding Dimension, a positive integer  [default: 2]\n
    `tau`   - Time Delay, a positive integer        [default: 1]\n
    `Fx`    - Fuzzy function name, one of the following:
              {`"sigmoid", "modsampen", "default", "gudermannian", "linear"`}\n
    `r`     - Fuzzy function parameters, a 1 element scalar or a 2 element
              tuple of positive values. The `r` parameters for each fuzzy
              function are defined as follows:      [default: [.2 2]]\n
                sigmoid:      r(1) = divisor of the exponential argument
                              r(2) = value subtracted from argument (pre-division)
                modsampen:    r(1) = divisor of the exponential argument
                              r(2) = value subtracted from argument (pre-division)
                default:      r(1) = divisor of the exponential argument
                              r(2) = argument exponent (pre-division)
                gudermannian: r  = a scalar whose value is the numerator of
                                    argument to gudermannian function:
                                    GD(x) = atan(tanh(r/x))
                linear:       r  = an integer value. When r = 0, the
                                    argument of the exponential function is
                                    normalised between [0 1]. When r = 1,
                                    the minimuum value of the exponential
                                    argument is set to 0.   \n
    `Logx`  - Logarithm base, a positive scalar  [default: natural]

    ## For further information on keyword arguments, see the EntropyHub guide.

    # See also `SampEn`, `ApEn`, `PermEn`, `DispEn`, `XFuzzEn`, `FuzzEn2D`, `MSEn`

    # References:
        [1] Weiting Chen, et al.
            "Characterization of surface EMG signal based on fuzzy entropy."
            IEEE Transactions on neural systems and rehabilitation engineering
            15.2 (2007): 266-272.

        [2] Hong-Bo Xie, Wei-Xing He, and Hui Liu
            "Measuring time series regularity using nonlinear
            similarity-based sample entropy."
            Physics Letters A
            372.48 (2008): 7140-7146.

    """
    function FuzzEn(Sig::AbstractArray{T,1} where T<:Real; m::Int=2, tau::Int=1,
         r::Union{Real,Tuple{Real,Real}}=(.2,2), Fx::String="default", Logx::Real=exp(1))

    N = size(Sig)[1]
    (N>10) ? nothing : error("Sig:   must be a numeric vector")
    (m > 0) ? nothing : error("m:     must be an integer > 0")
    (tau > 0) ? nothing : error("tau:   must be an integer > 0")
    (minimum(r)>=0 && length(r)<=2) ?
        nothing : error("r:  must be a scalar or 2 element tuple of positive values")
    (lowercase(Fx) in ["default","sigmoid","modsampen","gudermannian","linear"]) ?
        nothing : error("Fx:    must be one of the following strings -
        'default', 'sigmoid', 'modsampen', 'gudermannian', 'linear'")
    (Logx>0) ? nothing : error("Logx:   must be a positive number > 0")

    if length(r) == 2 && lowercase(Fx)=="linear"
        r = 0;
        print("Multiple values for r entered. Default value (0) used.\n")
    elseif length(r) == 2 && lowercase(Fx)=="gudermannian"
        r = r[1]
        print("Multiple values for r entered. First value used.\n")
    end

    m += 1
    Fun = getfield(_FuzzEn,Symbol(lowercase(Fx)))
    Sx = zeros(N,m)
    for k = 1:m
        Sx[1:N-(k-1)*tau,k] = Sig[1+(k-1)*tau:end]
    end
    Ps1 = zeros(m)
    Ps2 = zeros(m-1)
    Ps1[1] = .5

    for k = 2:m
        N1 = N - k*tau;    N2 = N - (k-1)*tau;
        T2 = Sx[1:N2,1:k] .- mean(Sx[1:N2,1:k],dims=2)
        d2 = zeros(N2-1,N2-1)

        for p = 1:N2-1
            Mu2 = maximum(abs.(repeat(transpose(T2[p,:]),N2-p) - T2[p+1:end,:]),dims=2)
            d2[p,p:end] = Fun(Mu2,r)
        end
        d1 = d2[1:N1-1,1:N1-1]
        Ps1[k] = sum(d1)/(N1*(N1-1))
        Ps2[k-1] = sum(d2)/(N2*(N2-1))
    end

    Fuzz = (log.(Logx, Ps1[1:end-1]) - log.(Logx,Ps2))

    return Fuzz, Ps1, Ps2
    end

    function sigmoid(x,r)
        if length(r) == 1
            error("When Fx = 'Sigmoid', r must be a two-element tuple.")
        end
        y = inv.(1 .+ exp.((x.-r[2])/r[1]))
        return y
    end

    function modsampen(x,r)
        if length(r) == 1
            error("When Fx = 'Modsampen', r must be a two-element tuple.")
        end
        y = inv.(1 .+ exp.((x.-r[2])/r[1]))
        return y
    end

    function default(x,r)
        if length(r) == 1
            error("When Fx = 'Default', r must be a two-element tuple.")
        end
        y = exp.(-(x.^r[2])/r[1])
        return y
    end

    function gudermannian(x,r)
        if r <= 0
            error("When Fx = 'Gudermannian', r must be a scalar > 0.")
        end
        y = atan.(tanh.(r[1]./x))
        y ./= maximum(y)
        return y
    end

    function linear(x,r)
        if r == 0 && length(x)>1
            y = exp.(-(x .- minimum(x))/(maximum(x)-minimum(x)))
        elseif r == 1
            y = exp.(-(x .- minimum(x)))
        elseif r == 0 && length(x)==1
            y = [0]
        else
            error("When Fx = 'Linear', r must be 0 or 1.")
        end
        return y
    end

end
"""
Copyright 2021 Matthew W. Flood, EntropyHub

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

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

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

For Terms of Use see https://github.com/MattWillFlood/EntropyHub
"""


# Import packages
using Test
using SafeTestsets
using ModelCheck

using Distributions
using Random
using DataFrames
using StatsBase
using GLM
using Gadfly
using Statistics
using Compose
using Colors
using SafeTestsets

# Import simulated data and Poisson general linear model to test functions
include("test_simulated_model.jl")

##############################################
# - Run unit tests
##############################################

@testset "ModelCheck Tests" begin

# Test #1: Test that plot_qq returns an object of type 'Plot' when using Poisson GLM
@test   typeof(plot_qqplot(model = modPoisson,
                           df = df,
                           y = "y")) == Plot

# Test #2: Test that plot_residual_fitted returns an object of type 'Plot' when using Poisson GLM
@test   typeof(plot_residual_fitted(model = modPoisson,
                                    df = df,
                                    y = "y")) == Plot

# Test #3: Test that plot_scale_location returns an object of type 'Plot' when using Poisson GLM
@test   typeof(plot_scale_location(model = modPoisson,
                                   df = df,
                                   y = "y")) == Plot

# Test #4: Test that plot_model_check returns an object of type 'Context' when using Poisson GLM
@test   typeof(plot_model_check(model = modPoisson,
                                df = df,
                                y = "y")) == Context

# Test #5: Test that plot_model_check returns an object of type 'Vector{Float64}' when using Poisson GLM
@test   typeof(calc_deviance_resids(model = modPoisson,
                                    df = df,
                                    y = "y")) == Vector{Float64}

end


using IterativeSolvers
using Base.Test

@testset "Orthogonalization" begin

srand(1234321)
n = 10
m = 3

@testset "Eltype $T" for T = (Complex64, Float64)

    # Create an orthonormal matrix V
    V, = qr(rand(T, n, m))

    # And a random vector to be orth. to V.
    w_original = rand(T, n)

    # Test whether w is w_original orthonormalized w.r.t. V,
    # given the projection h = V' * h and the norm of V * V' * h
    is_orthonormalized = (w, h, nrm) -> begin
        # Normality
        @test norm(w) ≈ one(real(T))

        # Orthogonality
        @test norm(V'w) ≈ zero(real(T)) atol = 10eps(real(T))

        # Denormalizing and adding the components in V should give back the original
        @test nrm * w + V * h ≈ w_original
    end

    # Assuming V is a matrix
    @testset "Using $method" for method = (DGKS, ClassicalGramSchmidt, ModifiedGramSchmidt)

        # Projection size
        h = zeros(T, m)

        # Orthogonalize w in-place
        w = copy(w_original)
        nrm = orthogonalize_and_normalize!(V, w, h, method)

        is_orthonormalized(w, h, nrm)
    end

    # Assuming V is a vector.
    @testset "ModifiedGramSchmidt with vectors" begin
        V_vec = [V[:, i] for i = 1 : m]

        # Projection size
        h = zeros(T, m)

        # Orthogonalize w in-place
        w = copy(w_original)
        nrm = orthogonalize_and_normalize!(V_vec, w, h, ModifiedGramSchmidt)

        is_orthonormalized(w, h, nrm)
    end
end
end


#=
see description in native.jl

allows objective to minimize frobenius norm or nuclear norm of residual matrix,
plus regularization; similar to
https://arxiv.org/ftp/arxiv/papers/1405/1405.1207.pdf
=#

using SparseArrays

struct MatrixRegressionJuMP{T <: Real} <: ExampleInstanceJuMP{T}
    Y::Matrix{T}
    X::Matrix{T}
    nuc_obj::Bool # use nuclear norm loss, else squared loss
    lam_fro::Real # penalty on Frobenius norm
    lam_nuc::Real # penalty on nuclear norm
    lam_las::Real # penalty on l1 norm
    lam_glr::Real # penalty on penalty on row group l1 norm
    lam_glc::Real # penalty on penalty on column group l1 norm
end

function MatrixRegressionJuMP{Float64}(
    n::Int,
    m::Int,
    p::Int,
    args...;
    A_max_rank::Int = div(m, 2) + 1,
    A_sparsity::Real = max(0.2, inv(sqrt(m * p))),
    Y_noise::Real = 0.01,
    )
    @assert p >= m
    @assert 1 <= A_max_rank <= m
    @assert 0 < A_sparsity <= 1
    A_left = sprandn(p, A_max_rank, A_sparsity)
    A_right = sprandn(A_max_rank, m, A_sparsity)
    A = 10 * A_left * A_right
    X = randn(n, p)
    Y = X * A + Y_noise * randn(n, m)
    return MatrixRegressionJuMP{Float64}(Y, X, args...)
end

function MatrixRegressionJuMP{Float64}(n::Int, m::Int)
    @assert n >= m >= 1
    X = randn(n, m)
    Y = randn(n, m)
    return MatrixRegressionJuMP{Float64}(Y, X, true, 0.1, 0, 0, 0, 0)
end

function build(inst::MatrixRegressionJuMP{T}) where {T <: Float64}
    (Y, X) = (inst.Y, inst.X)
    @assert min(inst.lam_fro, inst.lam_nuc, inst.lam_las,
        inst.lam_glr, inst.lam_glc) >= 0
    (data_n, data_m) = size(Y)
    data_p = size(X, 2)
    @assert size(X, 1) == data_n
    @assert data_p >= data_m
    @assert !inst.nuc_obj || data_m <= data_n
    model = JuMP.Model()
    JuMP.@variable(model, A[1:data_p, 1:data_m])
    JuMP.@variable(model, loss)
    loss_mat = Y - X * A

    if inst.nuc_obj
        JuMP.@constraint(model, vcat(loss, vec(loss_mat')) in
            MOI.NormNuclearCone(data_m, data_n))
    else
        if data_n > data_p
            # dimension reduction via QR
            F = qr(X, ColumnNorm())
            loss_mat = (F.Q' * Y)[1:data_p, :] - F.R[1:data_p, 1:data_p] *
                F.P' * A
        end
        JuMP.@constraint(model, vcat(loss, 1, vec(loss_mat) / sqrt(data_n)) in
            JuMP.RotatedSecondOrderCone())
    end

    obj = one(T) * loss
    if !iszero(inst.lam_fro)
        JuMP.@variable(model, t_fro)
        JuMP.@constraint(model, vcat(t_fro, inst.lam_fro * vec(A)) in
            JuMP.SecondOrderCone())
        obj += t_fro
    end
    if !iszero(inst.lam_nuc)
        JuMP.@variable(model, t_nuc)
        JuMP.@constraint(model, vcat(t_nuc, inst.lam_nuc * vec(A)) in
            MOI.NormNuclearCone(data_p, data_m))
        obj += t_nuc
    end
    if !iszero(inst.lam_las)
        JuMP.@variable(model, t_las)
        JuMP.@constraint(model, vcat(t_las, inst.lam_las * vec(A)) in
            MOI.NormOneCone(data_p * data_m + 1))
        obj += t_las
    end
    if !iszero(inst.lam_glr)
        JuMP.@variable(model, t_glr[1:data_p])
        JuMP.@constraint(model, [i = 1:data_p], vcat(t_glr[i],
            inst.lam_glr * A[i, :]) in JuMP.SecondOrderCone())
        obj += sum(t_glr)
    end
    if !iszero(inst.lam_glc)
        JuMP.@variable(model, t_glc[1:data_m])
        JuMP.@constraint(model, [i = 1:data_m], vcat(t_glc[i],
            inst.lam_glc * A[:, i]) in JuMP.SecondOrderCone())
        obj += sum(t_glc)
    end
    JuMP.@objective(model, Min, obj)

    model.ext[:A_var] = A # save for use in tests

    return model
end

function test_extra(inst::MatrixRegressionJuMP{T}, model::JuMP.Model) where T
    stat = JuMP.termination_status(model)
    @test stat == MOI.OPTIMAL
    (stat == MOI.OPTIMAL) || return

    # check objective value is correct
    tol = eps(T)^0.2
    (Y, X) = (inst.Y, inst.X)
    A_opt = JuMP.value.(model.ext[:A_var])
    loss_mat = Y - X * A_opt
    if inst.nuc_obj
        loss = sum(svdvals(loss_mat))
    else
        loss = sum(abs2, loss_mat) / (2 * size(Y, 1))
    end
    obj_result = loss +
        inst.lam_fro * norm(A_opt, 2) +
        inst.lam_nuc * sum(svdvals(A_opt)) +
        inst.lam_las * norm(A_opt, 1) +
        inst.lam_glr * sum(norm, eachrow(A_opt)) +
        inst.lam_glc * sum(norm, eachcol(A_opt))
    @test JuMP.objective_value(model) ≈ obj_result atol=tol rtol=tol
    return
end


using MonteCarloMarkovKernels
using Random
import Statistics: mean, var
using Test

seed!(12345)

function makeAR1Kernel(c::Float64, φ::Float64, σ::Float64)
  function P(x::Float64)
    return c + φ * x + σ * randn()
  end
  return P
end

const c = 2.32
const φ = 0.98
const σ = 2.3
P = makeAR1Kernel(c, φ, σ)

chain = simulateChain(P, 0.0, 1024*1024)

@test mean(chain) ≈ c/(1-φ) rtol=0.1
@test var(chain) ≈ σ^2/(1-φ^2) rtol=0.1

@test estimateBM(chain) ≈ σ^2/(1-φ^2) * (1+φ)/(1-φ) rtol=0.5
@test estimateSV(chain, :ModifiedBartlett) ≈
  σ^2/(1-φ^2) * (1+φ)/(1-φ) rtol=0.5


@doc raw"""
    AbstractGroupOperation

Abstract type for smooth binary operations $∘$ on elements of a Lie group $\mathcal{G}$:
```math
∘ : \mathcal{G} × \mathcal{G} → \mathcal{G}
```
An operation can be either defined for a specific group manifold over
number system `𝔽` or in general, by defining for an operation `Op` the following methods:

    identity_element!(::AbstractDecoratorManifold, q, q)
    inv!(::AbstractDecoratorManifold, q, p)
    _compose!(::AbstractDecoratorManifold, x, p, q)

Note that a manifold is connected with an operation by wrapping it with a decorator,
[`AbstractDecoratorManifold`](https://juliamanifolds.github.io/ManifoldsBase.jl/stable/decorator.html#ManifoldsBase.AbstractDecoratorManifold) using the [`IsGroupManifold`](@ref) to specify the operation.
For a concrete case the concrete wrapper [`GroupManifold`](@ref) can be used.
"""
abstract type AbstractGroupOperation end

"""
    IsGroupManifold{O<:AbstractGroupOperation} <: AbstractTrait

A trait to declare an [`AbstractManifold`](https://juliamanifolds.github.io/ManifoldsBase.jl/stable/types.html#ManifoldsBase.AbstractManifold)  as a manifold with group structure
with operation of type `O`.

Using this trait you can turn a manifold that you implement _implictly_ into a Lie group.
If you wish to decorate an existing manifold with one (or different) [`AbstractGroupAction`](@ref)s,
see [`GroupManifold`](@ref).

# Constructor

    IsGroupManifold(op)
"""
struct IsGroupManifold{O<:AbstractGroupOperation} <: AbstractTrait
    op::O
end

"""
    AbstractInvarianceTrait <: AbstractTrait

A common supertype for anz [`AbstractTrait`](https://juliamanifolds.github.io/ManifoldsBase.jl/stable/decorator.html#ManifoldsBase.AbstractTrait) related to metric invariance
"""
abstract type AbstractInvarianceTrait <: AbstractTrait end

"""
    HasLeftInvariantMetric <: AbstractInvarianceTrait

Specify that a certain the metric of a [`GroupManifold`](@ref) is a left-invariant metric
"""
struct HasLeftInvariantMetric <: AbstractInvarianceTrait end

direction(::HasLeftInvariantMetric) = LeftAction()
direction(::Type{HasLeftInvariantMetric}) = LeftAction()

"""
    HasRightInvariantMetric <: AbstractInvarianceTrait

Specify that a certain the metric of a [`GroupManifold`](@ref) is a right-invariant metric
"""
struct HasRightInvariantMetric <: AbstractInvarianceTrait end

direction(::HasRightInvariantMetric) = RightAction()
direction(::Type{HasRightInvariantMetric}) = RightAction()

"""
    HasBiinvariantMetric <: AbstractInvarianceTrait

Specify that a certain the metric of a [`GroupManifold`](@ref) is a bi-invariant metric
"""
struct HasBiinvariantMetric <: AbstractInvarianceTrait end
function parent_trait(::HasBiinvariantMetric)
    return ManifoldsBase.TraitList(HasLeftInvariantMetric(), HasRightInvariantMetric())
end

"""
    is_group_manifold(G::GroupManifold)
    is_group_manifoldd(G::AbstractManifold, o::AbstractGroupOperation)

returns whether an [`AbstractDecoratorManifold`](https://juliamanifolds.github.io/ManifoldsBase.jl/stable/decorator.html#ManifoldsBase.AbstractDecoratorManifold) is a group manifold with
[`AbstractGroupOperation`](@ref) `o`.
For a [`GroupManifold`](@ref) `G` this checks whether the right operations is stored within `G`.
"""
is_group_manifold(::AbstractManifold, ::AbstractGroupOperation) = false

@trait_function is_group_manifold(M::AbstractDecoratorManifold, op::AbstractGroupOperation)
function is_group_manifold(
    ::TraitList{<:IsGroupManifold{<:O}},
    ::AbstractDecoratorManifold,
    ::O,
) where {O<:AbstractGroupOperation}
    return true
end
@trait_function is_group_manifold(M::AbstractDecoratorManifold)
is_group_manifold(::AbstractManifold) = false
function is_group_manifold(
    t::TraitList{<:IsGroupManifold{<:AbstractGroupOperation}},
    M::AbstractDecoratorManifold,
)
    return is_group_manifold(M, t.head.op)
end

"""
    ActionDirection

Direction of action on a manifold, either [`LeftAction`](@ref) or [`RightAction`](@ref).
"""
abstract type ActionDirection end

"""
    LeftAction()

Left action of a group on a manifold.
"""
struct LeftAction <: ActionDirection end

"""
    RightAction()

Right action of a group on a manifold.
"""
struct RightAction <: ActionDirection end

"""
    switch_direction(::ActionDirection)

Returns a [`RightAction`](@ref) when given a [`LeftAction`](@ref) and vice versa.
"""
switch_direction(::ActionDirection)
switch_direction(::LeftAction) = RightAction()
switch_direction(::RightAction) = LeftAction()

@doc raw"""
    Identity{O<:AbstractGroupOperation}

Represent the group identity element ``e ∈ \mathcal{G}`` on a Lie group ``\mathcal G``
with [`AbstractGroupOperation`](@ref) of type `O`.

Similar to the philosophy that points are agnostic of their group at hand, the identity
does not store the group `g` it belongs to. However it depends on the type of the [`AbstractGroupOperation`](@ref) used.

See also [`identity_element`](@ref) on how to obtain the corresponding [`AbstractManifoldPoint`](https://juliamanifolds.github.io/ManifoldsBase.jl/stable/types.html#ManifoldsBase.AbstractManifoldPoint) or array representation.

# Constructors

    Identity(G::AbstractDecoratorManifold{𝔽})
    Identity(o::O)
    Identity(::Type{O})

create the identity of the corresponding subtype `O<:`[`AbstractGroupOperation`](@ref)
"""
struct Identity{O<:AbstractGroupOperation} end

@trait_function Identity(M::AbstractDecoratorManifold)
function Identity(
    ::TraitList{<:IsGroupManifold{O}},
    ::AbstractDecoratorManifold,
) where {O<:AbstractGroupOperation}
    return Identity{O}()
end
Identity(::O) where {O<:AbstractGroupOperation} = Identity(O)
Identity(::Type{O}) where {O<:AbstractGroupOperation} = Identity{O}()

# To ensure allocate_result_type works in general if idenitty apears in the tuple
number_eltype(::Identity) = Bool

@doc raw"""
    identity_element(G)

Return a point representation of the [`Identity`](@ref) on the [`IsGroupManifold`](@ref) `G`.
By default this representation is the default array or number representation.
It should return the corresponding default representation of ``e`` as a point on `G` if
points are not represented by arrays.
"""
identity_element(G::AbstractDecoratorManifold)
@trait_function identity_element(G::AbstractDecoratorManifold)
function identity_element(::TraitList{<:IsGroupManifold}, G::AbstractDecoratorManifold)
    q = allocate_result(G, identity_element)
    return identity_element!(G, q)
end

@trait_function identity_element!(G::AbstractDecoratorManifold, p)

function allocate_result(G::AbstractDecoratorManifold, ::typeof(identity_element))
    return zeros(representation_size(G)...)
end

@doc raw"""
    identity_element(G::AbstractDecoratorManifold, p)

Return a point representation of the [`Identity`](@ref) on the [`IsGroupManifold`](@ref) `G`,
where `p` indicates the type to represent the identity.
"""
identity_element(G::AbstractDecoratorManifold, p)
@trait_function identity_element(G::AbstractDecoratorManifold, p)
function identity_element(::TraitList{<:IsGroupManifold}, G::AbstractDecoratorManifold, p)
    q = allocate_result(G, identity_element, p)
    return identity_element!(G, q)
end
function check_size(
    ::TraitList{<:IsGroupManifold{<:O}},
    M::AbstractDecoratorManifold,
    ::Identity{<:O},
) where {O<:AbstractGroupOperation}
    return nothing
end
function check_size(::EmptyTrait, M::AbstractDecoratorManifold, e::Identity)
    return DomainError(0, "$M seems to not be a group manifold with $e.")
end

@doc raw"""
    is_identity(G::AbstractDecoratorManifold, q; kwargs)

Check whether `q` is the identity on the [`IsGroupManifold`](@ref) `G`, i.e. it is either
the [`Identity`](@ref)`{O}` with the corresponding [`AbstractGroupOperation`](@ref) `O`, or
(approximately) the correct point representation.
"""
is_identity(G::AbstractDecoratorManifold, q)
@trait_function is_identity(G::AbstractDecoratorManifold, q; kwargs...)
function is_identity(
    ::TraitList{<:IsGroupManifold},
    G::AbstractDecoratorManifold,
    q;
    kwargs...,
)
    return isapprox(G, identity_element(G), q; kwargs...)
end
function is_identity(
    ::TraitList{<:IsGroupManifold{O}},
    G::AbstractDecoratorManifold,
    ::Identity{O};
    kwargs...,
) where {O<:AbstractGroupOperation}
    return true
end
function is_identity(
    ::TraitList{<:IsGroupManifold},
    ::AbstractDecoratorManifold,
    ::Identity;
    kwargs...,
)
    return false
end

@inline function isapprox(
    ::TraitList{<:IsGroupManifold{O}},
    G::AbstractDecoratorManifold,
    p::Identity{O},
    q;
    kwargs...,
) where {O<:AbstractGroupOperation}
    return is_identity(G, q; kwargs...)
end
@inline function isapprox(
    ::TraitList{<:IsGroupManifold{O}},
    G::AbstractDecoratorManifold,
    p,
    q::Identity{O};
    kwargs...,
) where {O<:AbstractGroupOperation}
    return is_identity(G, p; kwargs...)
end
function isapprox(
    ::TraitList{<:IsGroupManifold{O}},
    G::AbstractDecoratorManifold,
    p::Identity{O},
    q::Identity{O};
    kwargs...,
) where {O<:AbstractGroupOperation}
    return true
end
function isapprox(
    ::TraitList{<:IsGroupManifold{O}},
    G::AbstractDecoratorManifold,
    p::Identity{O},
    q::Identity;
    kwargs...,
) where {O<:AbstractGroupOperation}
    return false
end
function isapprox(
    ::TraitList{<:IsGroupManifold{O}},
    G::AbstractDecoratorManifold,
    p::Identity,
    q::Identity{O};
    kwargs...,
) where {O<:AbstractGroupOperation}
    return false
end

@inline function isapprox(
    ::TraitList{IsGroupManifold{O}},
    G::AbstractDecoratorManifold,
    p::Identity{O},
    X,
    Y;
    kwargs...,
) where {𝔽,O<:AbstractGroupOperation}
    return isapprox(G, identity_element(G), X, Y; kwargs...)
end
function Base.isapprox(
    ::TraitList{<:IsGroupManifold},
    ::AbstractDecoratorManifold,
    ::Identity,
    ::Identity;
    kwargs...,
)
    return false
end

function Base.show(io::IO, ::Identity{O}) where {O<:AbstractGroupOperation}
    return print(io, "Identity($O)")
end

function check_point(
    ::TraitList{<:IsGroupManifold{O1}},
    G::AbstractDecoratorManifold,
    e::Identity{O2};
    kwargs...,
) where {O1<:AbstractGroupOperation,O2<:AbstractGroupOperation}
    return DomainError(
        e,
        "The Identity $e does not lie on $G, since its the identity with respect to $O2 and not $O1.",
    )
end

function is_point(
    ::TraitList{<:IsGroupManifold},
    G::AbstractDecoratorManifold,
    e::Identity,
    te=false;
    kwargs...,
)
    ie = is_identity(G, e; kwargs...)
    (te && !ie) && throw(DomainError(e, "The provided identity is not a point on $G."))
    return ie
end

function is_vector(
    t::TraitList{<:IsGroupManifold},
    G::AbstractDecoratorManifold,
    e::Identity,
    X,
    te=false,
    cbp=true;
    kwargs...,
)
    if cbp
        ie = is_point(G, e; kwargs...)
        (!te) && return ie
    end
    return is_vector(next_trait(t), G, identity_element(G), X, te, false; kwargs...)
end

@doc raw"""
    adjoint_action(G::AbstractDecoratorManifold, p, X)

Adjoint action of the element `p` of the Lie group `G` on the element `X`
of the corresponding Lie algebra.

It is defined as the differential of the group authomorphism ``Ψ_p(q) = pqp⁻¹`` at
the identity of `G`.

The formula reads
````math
\operatorname{Ad}_p(X) = dΨ_p(e)[X]
````
where $e$ is the identity element of `G`.

Note that the adjoint representation of a Lie group isn't generally faithful.
Notably the adjoint representation of SO(2) is trivial.
"""
adjoint_action(G::AbstractDecoratorManifold, p, X)
@trait_function adjoint_action(G::AbstractDecoratorManifold, p, Xₑ)
function adjoint_action(::TraitList{<:IsGroupManifold}, G::AbstractDecoratorManifold, p, Xₑ)
    Xₚ = translate_diff(G, p, Identity(G), Xₑ, LeftAction())
    Y = inverse_translate_diff(G, p, p, Xₚ, RightAction())
    return Y
end

@trait_function adjoint_action!(G::AbstractDecoratorManifold, Y, p, Xₑ)
function adjoint_action!(
    ::TraitList{<:IsGroupManifold},
    G::AbstractDecoratorManifold,
    Y,
    p,
    Xₑ,
)
    Xₚ = translate_diff(G, p, Identity(G), Xₑ, LeftAction())
    inverse_translate_diff!(G, Y, p, p, Xₚ, RightAction())
    return Y
end

@doc raw"""
    inv(G::AbstractDecoratorManifold, p)

Inverse $p^{-1} ∈ \mathcal{G}$ of an element $p ∈ \mathcal{G}$, such that
$p \circ p^{-1} = p^{-1} \circ p = e ∈ \mathcal{G}$, where $e$ is the [`Identity`](@ref)
element of $\mathcal{G}$.
"""
inv(::AbstractDecoratorManifold, ::Any...)
@trait_function Base.inv(G::AbstractDecoratorManifold, p)
function Base.inv(::TraitList{<:IsGroupManifold}, G::AbstractDecoratorManifold, p)
    q = allocate_result(G, inv, p)
    return inv!(G, q, p)
end

function Base.inv(
    ::TraitList{IsGroupManifold{O}},
    ::AbstractDecoratorManifold,
    e::Identity{O},
) where {O<:AbstractGroupOperation}
    return e
end

@trait_function inv!(G::AbstractDecoratorManifold, q, p)
function inv!(::TraitList{<:IsGroupManifold}, G::AbstractDecoratorManifold, q, p)
    return inv!(G.manifold, q, p)
end

function inv!(
    ::TraitList{IsGroupManifold{O}},
    G::AbstractDecoratorManifold,
    q,
    ::Identity{O},
) where {O<:AbstractGroupOperation}
    return identity_element!(G, q)
end

function Base.copyto!(
    ::TraitList{IsGroupManifold{O}},
    ::AbstractDecoratorManifold,
    e::Identity{O},
    ::Identity{O},
) where {O<:AbstractGroupOperation}
    return e
end
function Base.copyto!(
    ::TraitList{IsGroupManifold{O}},
    G::AbstractDecoratorManifold,
    p,
    ::Identity{O},
) where {O<:AbstractGroupOperation}
    return identity_element!(G, p)
end

@doc raw"""
    compose(G::AbstractDecoratorManifold, p, q)

Compose elements ``p,q ∈ \mathcal{G}`` using the group operation ``p \circ q``.

For implementing composition on a new group manifold, please overload `_compose`
instead so that methods with [`Identity`](@ref) arguments are not ambiguous.
"""
compose(::AbstractDecoratorManifold, ::Any...)

@trait_function compose(G::AbstractDecoratorManifold, p, q)
function compose(::TraitList{<:IsGroupManifold}, G::AbstractDecoratorManifold, p, q)
    return _compose(G, p, q)
end
function compose(
    ::AbstractDecoratorManifold,
    ::Identity{O},
    p,
) where {O<:AbstractGroupOperation}
    return p
end
function compose(
    ::AbstractDecoratorManifold,
    p,
    ::Identity{O},
) where {O<:AbstractGroupOperation}
    return p
end
function compose(
    ::AbstractDecoratorManifold,
    e::Identity{O},
    ::Identity{O},
) where {O<:AbstractGroupOperation}
    return e
end

function _compose(G::AbstractDecoratorManifold, p, q)
    x = allocate_result(G, compose, p, q)
    return _compose!(G, x, p, q)
end

@trait_function compose!(M::AbstractDecoratorManifold, x, p, q)

function compose!(::TraitList{<:IsGroupManifold}, G::AbstractDecoratorManifold, x, q, p)
    return _compose!(G, x, q, p)
end
function compose!(
    G::AbstractDecoratorManifold,
    q,
    p,
    ::Identity{O},
) where {O<:AbstractGroupOperation}
    return copyto!(G, q, p)
end
function compose!(
    G::AbstractDecoratorManifold,
    q,
    ::Identity{O},
    p,
) where {O<:AbstractGroupOperation}
    return copyto!(G, q, p)
end
function compose!(
    G::AbstractDecoratorManifold,
    q,
    ::Identity{O},
    e::Identity{O},
) where {O<:AbstractGroupOperation}
    return identity_element!(G, q)
end
function compose!(
    ::AbstractDecoratorManifold,
    e::Identity{O},
    ::Identity{O},
    ::Identity{O},
) where {O<:AbstractGroupOperation}
    return e
end

transpose(e::Identity) = e

@doc raw"""
    hat(M::AbstractDecoratorManifold{𝔽,O}, ::Identity{O}, Xⁱ) where {𝔽,O<:AbstractGroupOperation}

Given a basis $e_i$ on the tangent space at a the [`Identity`}(@ref) and tangent
component vector ``X^i``, compute the equivalent vector representation
``X=X^i e_i**, where Einstein summation notation is used:

````math
∧ : X^i ↦ X^i e_i
````

For array manifolds, this converts a vector representation of the tangent
vector to an array representation. The [`vee`](@ref) map is the `hat` map's
inverse.
"""
function hat(
    M::AbstractDecoratorManifold,
    ::Identity{O},
    X,
) where {O<:AbstractGroupOperation}
    return get_vector_lie(M, X, VeeOrthogonalBasis())
end
function hat!(
    M::AbstractDecoratorManifold,
    Y,
    ::Identity{O},
    X,
) where {O<:AbstractGroupOperation}
    return get_vector_lie!(M, Y, X, VeeOrthogonalBasis())
end
function hat(M::AbstractManifold, e::Identity, ::Any)
    return throw(ErrorException("On $M there exsists no identity $e"))
end
function hat!(M::AbstractManifold, ::Any, e::Identity, ::Any)
    return throw(ErrorException("On $M there exsists no identity $e"))
end

@doc raw"""
    vee(M::AbstractManifold, p, X)

Given a basis $e_i$ on the tangent space at a point `p` and tangent
vector `X`, compute the vector components $X^i$, such that $X = X^i e_i$, where
Einstein summation notation is used:

````math
\vee : X^i e_i ↦ X^i
````

For array manifolds, this converts an array representation of the tangent
vector to a vector representation. The [`hat`](@ref) map is the `vee` map's
inverse.
"""
function vee(
    M::AbstractDecoratorManifold,
    ::Identity{O},
    X,
) where {O<:AbstractGroupOperation}
    return get_coordinates_lie(M, X, VeeOrthogonalBasis())
end
function vee!(
    M::AbstractDecoratorManifold,
    Y,
    ::Identity{O},
    X,
) where {O<:AbstractGroupOperation}
    return get_coordinates_lie!(M, Y, X, VeeOrthogonalBasis())
end
function vee(M::AbstractManifold, e::Identity, ::Any)
    return throw(ErrorException("On $M there exsists no identity $e"))
end
function vee!(M::AbstractManifold, ::Any, e::Identity, ::Any)
    return throw(ErrorException("On $M there exsists no identity $e"))
end

"""
    lie_bracket(G::AbstractDecoratorManifold, X, Y)

Lie bracket between elements `X` and `Y` of the Lie algebra corresponding to
the Lie group `G`, cf. [`IsGroupManifold`](@ref).

This can be used to compute the adjoint representation of a Lie algebra.
Note that this representation isn't generally faithful. Notably the adjoint
representation of 𝔰𝔬(2) is trivial.
"""
lie_bracket(G::AbstractDecoratorManifold, X, Y)
@trait_function lie_bracket(M::AbstractDecoratorManifold, X, Y)

@trait_function lie_bracket!(M::AbstractDecoratorManifold, Z, X, Y)

_action_order(p, q, ::LeftAction) = (p, q)
_action_order(p, q, ::RightAction) = (q, p)

@doc raw"""
    translate(G::AbstractDecoratorManifold, p, q, conv::ActionDirection=LeftAction()])

Translate group element $q$ by $p$ with the translation $τ_p$ with the specified
`conv`ention, either left ($L_p$) or right ($R_p$), defined as
```math
\begin{aligned}
L_p &: q ↦ p \circ q\\
R_p &: q ↦ q \circ p.
\end{aligned}
```
"""
translate(::AbstractDecoratorManifold, ::Any...)
@trait_function translate(
    G::AbstractDecoratorManifold,
    p,
    q,
    conv::ActionDirection=LeftAction(),
)
function translate(
    ::TraitList{<:IsGroupManifold},
    G::AbstractDecoratorManifold,
    p,
    q,
    conv::ActionDirection,
)
    return compose(G, _action_order(p, q, conv)...)
end

@trait_function translate!(
    G::AbstractDecoratorManifold,
    X,
    p,
    q,
    conv::ActionDirection=LeftAction(),
)
function translate!(
    ::TraitList{<:IsGroupManifold},
    G::AbstractDecoratorManifold,
    X,
    p,
    q,
    conv::ActionDirection,
)
    return compose!(G, X, _action_order(p, q, conv)...)
end

@doc raw"""
    inverse_translate(G::AbstractDecoratorManifold, p, q, conv::ActionDirection=LeftAction())

Inverse translate group element $q$ by $p$ with the inverse translation $τ_p^{-1}$ with the
specified `conv`ention, either left ($L_p^{-1}$) or right ($R_p^{-1}$), defined as
```math
\begin{aligned}
L_p^{-1} &: q ↦ p^{-1} \circ q\\
R_p^{-1} &: q ↦ q \circ p^{-1}.
\end{aligned}
```
"""
inverse_translate(::AbstractDecoratorManifold, ::Any...)
@trait_function inverse_translate(
    G::AbstractDecoratorManifold,
    p,
    q,
    conv::ActionDirection=LeftAction(),
)
function inverse_translate(
    ::TraitList{<:IsGroupManifold},
    G::AbstractDecoratorManifold,
    p,
    q,
    conv::ActionDirection,
)
    return translate(G, inv(G, p), q, conv)
end

@trait_function inverse_translate!(
    G::AbstractDecoratorManifold,
    X,
    p,
    q,
    conv::ActionDirection=LeftAction(),
)
function inverse_translate!(
    ::TraitList{<:IsGroupManifold},
    G::AbstractDecoratorManifold,
    X,
    p,
    q,
    conv::ActionDirection,
)
    return translate!(G, X, inv(G, p), q, conv)
end

@doc raw"""
    translate_diff(G::AbstractDecoratorManifold, p, q, X, conv::ActionDirection=LeftAction())

For group elements $p, q ∈ \mathcal{G}$ and tangent vector $X ∈ T_q \mathcal{G}$, compute
the action of the differential of the translation $τ_p$ by $p$ on $X$, with the specified
left or right `conv`ention. The differential transports vectors:
```math
(\mathrm{d}τ_p)_q : T_q \mathcal{G} → T_{τ_p q} \mathcal{G}\\
```
"""
translate_diff(::AbstractDecoratorManifold, ::Any...)
@trait_function translate_diff(
    G::AbstractDecoratorManifold,
    p,
    q,
    X,
    conv::ActionDirection=LeftAction(),
)
function translate_diff(
    ::TraitList{<:IsGroupManifold},
    G::AbstractDecoratorManifold,
    p,
    q,
    X,
    conv::ActionDirection,
)
    Y = allocate_result(G, translate_diff, X, p, q)
    translate_diff!(G, Y, p, q, X, conv)
    return Y
end
@trait_function translate_diff!(
    G::AbstractDecoratorManifold,
    Y,
    p,
    q,
    X,
    conv::ActionDirection=LeftAction(),
)

@doc raw"""
    inverse_translate_diff(G::AbstractDecoratorManifold, p, q, X, conv::ActionDirection=LeftAction())

For group elements $p, q ∈ \mathcal{G}$ and tangent vector $X ∈ T_q \mathcal{G}$, compute
the action on $X$ of the differential of the inverse translation $τ_p$ by $p$, with the
specified left or right `conv`ention. The differential transports vectors:
```math
(\mathrm{d}τ_p^{-1})_q : T_q \mathcal{G} → T_{τ_p^{-1} q} \mathcal{G}\\
```
"""
inverse_translate_diff(::AbstractDecoratorManifold, ::Any...)
@trait_function inverse_translate_diff(
    G::AbstractDecoratorManifold,
    p,
    q,
    X,
    conv::ActionDirection=LeftAction(),
)
function inverse_translate_diff(
    ::TraitList{<:IsGroupManifold},
    G::AbstractDecoratorManifold,
    p,
    q,
    X,
    conv::ActionDirection,
)
    return translate_diff(G, inv(G, p), q, X, conv)
end

@trait_function inverse_translate_diff!(
    G::AbstractDecoratorManifold,
    Y,
    p,
    q,
    X,
    conv::ActionDirection=LeftAction(),
)
function inverse_translate_diff!(
    ::TraitList{<:IsGroupManifold},
    G::AbstractDecoratorManifold,
    Y,
    p,
    q,
    X,
    conv::ActionDirection,
)
    return translate_diff!(G, Y, inv(G, p), q, X, conv)
end

@doc raw"""
    exp_lie(G, X)
    exp_lie!(G, q, X)

Compute the group exponential of the Lie algebra element `X`. It is equivalent to the
exponential map defined by the [`CartanSchoutenMinus`](@ref) connection.

Given an element $X ∈ 𝔤 = T_e \mathcal{G}$, where $e$ is the [`Identity`](@ref) element of
the group $\mathcal{G}$, and $𝔤$ is its Lie algebra, the group exponential is the map

````math
\exp : 𝔤 → \mathcal{G},
````
such that for $t,s ∈ ℝ$, $γ(t) = \exp (t X)$ defines a one-parameter subgroup with the
following properties:

````math
\begin{aligned}
γ(t) &= γ(-t)^{-1}\\
γ(t + s) &= γ(t) \circ γ(s) = γ(s) \circ γ(t)\\
γ(0) &= e\\
\lim_{t → 0} \frac{d}{dt} γ(t) &= X.
\end{aligned}
````

!!! note
    In general, the group exponential map is distinct from the Riemannian exponential map
    [`exp`](@ref).

For example for the [`MultiplicationOperation`](@ref) and either `Number` or `AbstractMatrix`
the Lie exponential is the numeric/matrix exponential.

````math
\exp X = \operatorname{Exp} X = \sum_{n=0}^∞ \frac{1}{n!} X^n.
````

Since this function also depends on the group operation, make sure to implement
the corresponding trait version `exp_lie(::TraitList{<:IsGroupManifold}, G, X)`.
"""
exp_lie(G::AbstractManifold, X)
@trait_function exp_lie(M::AbstractDecoratorManifold, X)
function exp_lie(::TraitList{<:IsGroupManifold}, G::AbstractDecoratorManifold, X)
    q = allocate_result(G, exp_lie, X)
    return exp_lie!(G, q, X)
end

@trait_function exp_lie!(M::AbstractDecoratorManifold, q, X)

@doc raw"""
    log_lie(G, q)
    log_lie!(G, X, q)

Compute the Lie group logarithm of the Lie group element `q`. It is equivalent to the
logarithmic map defined by the [`CartanSchoutenMinus`](@ref) connection.

Given an element $q ∈ \mathcal{G}$, compute the right inverse of the group exponential map
[`exp_lie`](@ref), that is, the element $\log q = X ∈ 𝔤 = T_e \mathcal{G}$, such that
$q = \exp X$

!!! note
    In general, the group logarithm map is distinct from the Riemannian logarithm map
    [`log`](@ref).

    For matrix Lie groups this is equal to the (matrix) logarithm:

````math
\log q = \operatorname{Log} q = \sum_{n=1}^∞ \frac{(-1)^{n+1}}{n} (q - e)^n,
````

where $e$ here is the [`Identity`](@ref) element, that is, $1$ for numeric $q$ or the
identity matrix $I_m$ for matrix $q ∈ ℝ^{m × m}$.

Since this function also depends on the group operation, make sure to implement
either
* `_log_lie(G, q)` and `_log_lie!(G, X, q)` for the points not being the [`Identity`](@ref)
* the trait version `log_lie(::TraitList{<:IsGroupManifold}, G, e)`, `log_lie(::TraitList{<:IsGroupManifold}, G, X, e)` for own implementations of the identity case.
"""
log_lie(::AbstractDecoratorManifold, q)
@trait_function log_lie(G::AbstractDecoratorManifold, q)
function log_lie(::TraitList{<:IsGroupManifold}, G::AbstractDecoratorManifold, q)
    return _log_lie(G, q)
end
function log_lie(
    ::TraitList{<:IsGroupManifold{O}},
    G::AbstractDecoratorManifold,
    ::Identity{O},
) where {O<:AbstractGroupOperation}
    return zero_vector(G, identity_element(G))
end
# though identity was taken care of – as usual restart decorator dispatch
function _log_lie(G::AbstractDecoratorManifold, q)
    X = allocate_result(G, log_lie, q)
    return log_lie!(G, X, q)
end

@trait_function log_lie!(G::AbstractDecoratorManifold, X, q)
function log_lie!(::TraitList{<:IsGroupManifold}, G::AbstractDecoratorManifold, X, q)
    return _log_lie!(G, X, q)
end
function log_lie!(
    ::TraitList{<:IsGroupManifold{O}},
    G::AbstractDecoratorManifold,
    X,
    ::Identity{O},
) where {O<:AbstractGroupOperation}
    return zero_vector!(G, X, identity_element(G))
end

"""
    GroupExponentialRetraction{D<:ActionDirection} <: AbstractRetractionMethod

Retraction using the group exponential [`exp_lie`](@ref) "translated" to any point on the
manifold.

For more details, see
[`retract`](@ref retract(::GroupManifold, p, X, ::GroupExponentialRetraction)).

# Constructor

    GroupExponentialRetraction(conv::ActionDirection = LeftAction())
"""
struct GroupExponentialRetraction{D<:ActionDirection} <: AbstractRetractionMethod end

function GroupExponentialRetraction(conv::ActionDirection=LeftAction())
    return GroupExponentialRetraction{typeof(conv)}()
end

"""
    GroupLogarithmicInverseRetraction{D<:ActionDirection} <: AbstractInverseRetractionMethod

Retraction using the group logarithm [`log_lie`](@ref) "translated" to any point on the
manifold.

For more details, see
[`inverse_retract`](@ref inverse_retract(::GroupManifold, p, q ::GroupLogarithmicInverseRetraction)).

# Constructor

    GroupLogarithmicInverseRetraction(conv::ActionDirection = LeftAction())
"""
struct GroupLogarithmicInverseRetraction{D<:ActionDirection} <:
       AbstractInverseRetractionMethod end

function GroupLogarithmicInverseRetraction(conv::ActionDirection=LeftAction())
    return GroupLogarithmicInverseRetraction{typeof(conv)}()
end

direction(::GroupExponentialRetraction{D}) where {D} = D()
direction(::GroupLogarithmicInverseRetraction{D}) where {D} = D()

@doc raw"""
    retract(
        G::AbstractDecoratorManifold,
        p,
        X,
        method::GroupExponentialRetraction{<:ActionDirection},
    )

Compute the retraction using the group exponential [`exp_lie`](@ref) "translated" to any
point on the manifold.
With a group translation ([`translate`](@ref)) $τ_p$ in a specified direction, the
retraction is

````math
\operatorname{retr}_p = τ_p \circ \exp \circ (\mathrm{d}τ_p^{-1})_p,
````

where $\exp$ is the group exponential ([`exp_lie`](@ref)), and $(\mathrm{d}τ_p^{-1})_p$ is
the action of the differential of inverse translation $τ_p^{-1}$ evaluated at $p$ (see
[`inverse_translate_diff`](@ref)).
"""
function retract(
    ::TraitList{<:IsGroupManifold},
    G::AbstractDecoratorManifold,
    p,
    X,
    method::GroupExponentialRetraction,
)
    conv = direction(method)
    Xₑ = inverse_translate_diff(G, p, p, X, conv)
    pinvq = exp_lie(G, Xₑ)
    q = translate(G, p, pinvq, conv)
    return q
end

function retract!(
    ::TraitList{<:IsGroupManifold},
    G::AbstractDecoratorManifold,
    q,
    p,
    X,
    method::GroupExponentialRetraction,
)
    conv = direction(method)
    Xₑ = inverse_translate_diff(G, p, p, X, conv)
    pinvq = exp_lie(G, Xₑ)
    return translate!(G, q, p, pinvq, conv)
end

@doc raw"""
    inverse_retract(
        G::AbstractDecoratorManifold,
        p,
        X,
        method::GroupLogarithmicInverseRetraction{<:ActionDirection},
    )

Compute the inverse retraction using the group logarithm [`log_lie`](@ref) "translated"
to any point on the manifold.
With a group translation ([`translate`](@ref)) $τ_p$ in a specified direction, the
retraction is

````math
\operatorname{retr}_p^{-1} = (\mathrm{d}τ_p)_e \circ \log \circ τ_p^{-1},
````

where $\log$ is the group logarithm ([`log_lie`](@ref)), and $(\mathrm{d}τ_p)_e$ is the
action of the differential of translation $τ_p$ evaluated at the identity element $e$
(see [`translate_diff`](@ref)).
"""
function inverse_retract(
    ::TraitList{<:IsGroupManifold},
    G::AbstractDecoratorManifold,
    p,
    q,
    method::GroupLogarithmicInverseRetraction,
)
    conv = direction(method)
    pinvq = inverse_translate(G, p, q, conv)
    Xₑ = log_lie(G, pinvq)
    return translate_diff(G, p, Identity(G), Xₑ, conv)
end

function inverse_retract!(
    ::TraitList{<:IsGroupManifold},
    G::AbstractDecoratorManifold,
    X,
    p,
    q,
    method::GroupLogarithmicInverseRetraction,
)
    conv = direction(method)
    pinvq = inverse_translate(G, p, q, conv)
    Xₑ = log_lie(G, pinvq)
    return translate_diff!(G, X, p, Identity(G), Xₑ, conv)
end

@doc raw"""
    get_vector_lie(G::AbstractDecoratorManifold, a, B::AbstractBasis)

Reconstruct a tangent vector from the Lie algebra of `G` from cooordinates `a` of a basis `B`.
This is similar to calling [`get_vector`](@ref) at the `p=`[`Identity`](@ref)`(G)`.
"""
function get_vector_lie(G::AbstractManifold, X, B::AbstractBasis)
    return get_vector(base_manifold(G), identity_element(G), X, B)
end
function get_vector_lie!(G::AbstractManifold, Y, X, B::AbstractBasis)
    return get_vector!(base_manifold(G), Y, identity_element(G), X, B)
end

@doc raw"""
    get_coordinates_lie(G::AbstractManifold, X, B::AbstractBasis)

Get the coordinates of an element `X` from the Lie algebra og `G` with respect to a basis `B`.
This is similar to calling [`get_coordinates`](@ref) at the `p=`[`Identity`](@ref)`(G)`.
"""
function get_coordinates_lie(G::AbstractManifold, X, B::AbstractBasis)
    return get_coordinates(base_manifold(G), identity_element(G), X, B)
end
function get_coordinates_lie!(G::AbstractManifold, a, X, B::AbstractBasis)
    return get_coordinates!(base_manifold(G), a, identity_element(G), X, B)
end


module GlobalSensitivity

using Statistics, RecursiveArrayTools, LinearAlgebra
using QuasiMonteCarlo, ForwardDiff
using Parameters: @unpack, @with_kw
using FFTW, Distributions

abstract type GSAMethod end

include("morris_sensitivity.jl")
include("sobol_sensitivity.jl")
include("regression_sensitivity.jl")
include("DGSM_sensitivity.jl")
include("eFAST_sensitivity.jl")

export Sobol, Morris, gsa,
       SensitivityAlg, RegressionGSA, DGSM, eFAST

end # module


function exp_rand()
    # cdf of Possion+(log(2)), since P(X=0) = 1/2
    # for positive Possion, q[k] = sum(log(2)^k / k!)
    q = [
    0.6931471805599453,
    0.9333736875190459,
    0.9888777961838675,
    0.9984959252914960,
    0.9998292811061389,
    0.9999833164100727,
    0.9999985691438767,
    0.9999998906925558,
    0.9999999924734159,
    0.9999999995283275,
    0.9999999999728814,
    0.9999999999985598,
    0.9999999999999289,
    0.9999999999999968,
    0.9999999999999999,
    1.0000000000000000
    ]
    a = 0.0
    u = rand()
    while (u <= 0) | (u >= 1)
        u = rand()
    end
    while true
        u += u;
        if (u > 1.0)
            break
        end
        a += q[1]  # μ = q[1], and the times of loop is M, so a = μM
        # TODO: but why the times of loop follows a geometric (1/2) ??
    end
    u -= 1.0 # an independent uniform variate
    if u <= q[1]
        return a + u
    end
    i = 1 # one-based
    ustar = rand()
    umin = ustar
    while true
        ustar = rand()
        if (umin > ustar)
            umin = ustar
        end
        i += 1
        if u <= q[i]
            return a + umin * q[1]
        end
    end
end


include("graph.jl")

function contact_freq_to_distance_matrix(contact_file::AbstractString; remove_inf=true, scale=1, exponent=1/3)
    # turns a csv file of contact frequences into a distance matrix
    contact_data = readcsv(contact_file, Float64)
    # TODO: use the graphical approach... distances are 1/freq, use shortest
    # paths algorithm to get distances that don't exist

    # arbitrarily scale distances by some scale
    distances = scale./(contact_data).^exponent
    dist = distances
    # removing zero points
    if remove_inf
        dist_index = 1
        while dist_index <= size(dist)[1]
            if mapreduce(x->(isinf(x) || x==0), &, dist[dist_index, 1:end-1])
                dist = [dist[1:dist_index-1, :], dist[dist_index+1:end, :]]
                dist = [dist[:, 1:dist_index-1] dist[:,dist_index+1:end]]
            else
                dist_index += 1
            end
        end
    else

    end
    distances = dist
    n, n2 = size(distances)
    g = init_graph()
    for i in range(1, n)
        add_node!(g, i)
    end
    for i in range(1, n-1)
        for i2 in (i+1):n
            if distances[i, i2] != Inf
                add_edge!(g, i, i2, distances[i, i2])
            else
                #distances[i, i2] = 1.0
                #add_edge!(g, i, i2, distances[i, i2])
            end
        end
    end
    # done with constructing graph, now to find the distances
    dist = all_shortest_paths(g)
    # TODO: what if there's an "Inf" in the distance? replace by 1?
    # this is an incredibly hacky way of doing things, and it might mess
    # everything up entirely.
    # ... and it does seem to mess up everything entirely.
    # how else would this be done?
    #dist[dist.==Inf] = 1
    return dist
end

function mds_metric(dist::Array{Float64, 2}; dims = 3)
    # 'dist' is a matrix of distances (should be symmetric)
    # metric MDS - assumes that the data is metric
    # TODO: do some crazy matrix operations
    N, N = size(dist)
    d = zeros(N)
    b = 1/N^2*sum([sum([dist[j,k]^2 for k in j:N]) for j in range(1,N)])
    for i in range(1, N)
        d[i] = 1/N * sum([dist[i, j]^2 for j in range(1, N)]) - b
    end
    M = zeros(N, N)
    for i in range(1, N)
        for j in range(1, N)
            M[i,j] = 1/2*(d[i] + d[j] - dist[i,j]^2)
        end
    end
    lambda, E = eig(M)
    # largest 3 eigenvalues
    v = E[:, end-dims+1:end] .* sqrt(lambda[end-dims+1:end])'
    return fliplr(v[:,1:dims]) #[v[:,3] v[:,2] v[:,1]]
end

function test()
    m = "../test_data/combined_You_Pro_merged_chr18.csv"
    dist = contact_freq_to_distance_matrix(m)
    return dist, mds_metric(dist)
end


struct DFNGammaStrategy <: SVDD.InitializationStrategyGamma
    min_γ
    max_γ
end

DFNGammaStrategy() = DFNGammaStrategy(0.0, 100.0)

function dfn(nn, gamma)
    near = mean(exp(-gamma*nn_dist[2]) for nn_dist in nn[2])
    far = mean(exp(-gamma*nn_dist[end]) for nn_dist in nn[2])
    return 2*near - 2*far
end

function _dfn(data)
    kdtree = KDTree(data)
    nn = knn(kdtree, data, size(data, 2), true)
    return ((gamma) -> dfn(nn, gamma))
end

function SVDD.calculate_gamma(model, strategy::DFNGammaStrategy)
    f(gamma) = - _dfn(model.data)(gamma)
    res = Optim.optimize(f, strategy.min_γ, strategy.max_γ, abs_tol=0.1)
    return res.minimizer
end


using Yota
using Zygote
using Zygote: @adjoint
# include("../src/core.jl")
using BenchmarkTools
using Distributions
import NNlib

logistic(x) = 1 / (1 + exp(-x))
softplus(x) = log(exp(x) + 1)

# Yota diff rules
@diffrule logistic(x::Number) x (logistic(x) * (1 - logistic(x)) * dy)
@diffrule softplus(x::Number) x logistic(x) * dy

# Zygote diff rules
@adjoint logistic(x) = logistic(x), Δ -> ((logistic(x) * (1 - logistic(x)) * Δ),)
@adjoint softplus(x) = softplus(x), Δ -> ((logistic(x) * Δ),)

function perf_test(f, args...)
    println("Compiling derivatives using Yota")
    @time Yota.grad(f, args...)
    r1 = @benchmark Yota.grad($f, $(args)...)
    show(stdout, MIME{Symbol("text/plain")}(), r1)
    println("\n")

    println("Compiling derivatives using Zygote")
    @time Zygote.gradient(f, args...)
    r2 = @benchmark Zygote.gradient($f, $(args)...)
    show(stdout, MIME{Symbol("text/plain")}(), r2)
    println("\n")
end

# VAE benchmark

function xavier_init(dim_in, dim_out; c=1)
    low = -c * sqrt(6.0 / (dim_in + dim_out))
    high = c * sqrt(6.0 / (dim_in + dim_out))
    return rand(Uniform(low, high), dim_in, dim_out)
end

# see examples/vae.jl for a more convenient API
function encode(We1, be1, We2, be2, We3, be3, We4, be4, x)
    he1 = softplus.(We1 * x .+ be1)
    he2 = softplus.(We2 * he1 .+ be2)
    mu = We3 * he2 .+ be3
    log_sigma2 = We4 * he2 .+ be4
    return mu, log_sigma2
end

function decode(Wd1, bd1, Wd2, bd2, Wd3, bd3, z)
    hd1 = softplus.(Wd1 * z .+ bd1)
    hd2 = softplus.(Wd2 * hd1 .+ bd2)
    x_rec = logistic.(Wd3 * hd2 .+ bd3)
    return x_rec
end

function vae_cost(We1, be1, We2, be2, We3, be3, We4, be4,
                  Wd1, bd1, Wd2, bd2, Wd3, bd3, eps, x)
    mu, log_sigma2 = encode(We1, be1, We2, be2, We3, be3, We4, be4, x)
    z = mu .+ sqrt.(exp.(log_sigma2)) .* eps
    x_rec = decode(Wd1, bd1, Wd2, bd2, Wd3, bd3, z)
    # loss
    rec_loss = -sum(x .* log.(1e-10 .+ x_rec) .+ (1 .- x) .* log.(1e-10 + 1.0 .- x_rec); dims=1)
    # KLD = -0.5 .* sum(1 .+ log_sigma2 .- mu .^ 2.0f0 - exp.(log_sigma2); dims=1)
    # cost = mean(rec_loss .+ KLD)
    cost = mean(rec_loss)
end

function benchmark_vae()
    n_inp, n_he1, n_he2, n_z, n_hd1, n_hd2, n_out = 784, 500, 500, 20, 500, 500, 784
    We1, be1, We2, be2, We3, be3, We4, be4 = (xavier_init(n_he1, n_inp),
                                              zeros(n_he1),
                                              xavier_init(n_he2, n_he1),
                                              zeros(n_he2),
                                              xavier_init(n_z, n_he2),
                                              zeros(n_z),
                                              xavier_init(n_z, n_he2),
                                              zeros(n_z))
    Wd1, bd1, Wd2, bd2, Wd3, bd3 = (xavier_init(n_hd1, n_z),
                                    zeros(n_hd1),
                                    xavier_init(n_hd2, n_hd1),
                                    zeros(n_hd2),
                                    xavier_init(n_out, n_hd2),
                                    zeros(n_out))
    eps = (rand(Normal(0, 1), size(We3, 1), 100))
    x = rand(784, 100)
    perf_test(vae_cost, We1, be1, We2, be2, We3, be3, We4, be4, Wd1, bd1, Wd2, bd2, Wd3, bd3, eps, x)
end


import Base.norm

function perfTest(N=100000)
  start = time()
  x = Variable(N)
  c = ones(N, 1)
  p = Problem(:minimize, dot(c,x), [c<=x, x<=2*c])
  println(time() - start)
  solve!(p)
  println(time() - start)
end

function test()
  start = time()
  N = 200
  x = Variable(1)
  y = Variable(N, N)
  c = ones(N, 1)
  p = Problem(:minimize, c' * (y + x) * c, [x >= 3, 2*y>=0, y <= x])
  println(time() - start)
  solve!(p)
  println(p.optval)
  println(time() - start)
end

# Assuming speed of calling randn, sprandn, etc is negligible for now...
# This is untrue. We'll soon generate data files.
# Also we'll need lucky data for this to be feasible.
# test2(10, 15, 10); sometimes gets feasible, bounded problems
function test2(m=1000, n=1000, q=1000)
  x = Variable(n)
  A = randn(m,n)
  b = randn(m)
  C = randn(q,n)
  d = C * randn(n)
  c = randn(n)
  p = Problem(:minimize, c' * x, [A*x <=b, C*x == d])
  solve!(p)
  println(p.optval)
  return p
end

# test2sparse(10,10,5,0.4); sometimes gets bounded, feasible problems
function test2sparse(m=1000, n=1000, q=1000, p=0.1)
  x = Variable(n)
  A = sprandn(m,n,p)
  b = randn(m)
  C = sprandn(q,n,p)
  d = C * randn(n)
  c = randn(n)
  p = Problem(:minimize, c' * x, [A*x <=b, C*x ==d])
  solve!(p)
  println(p.optval)
  return p
end

# http://see.stanford.edu/materials/lsocoee364a/hw4sol.pdf
# close to optimal with 10, 10
function test3(m=1000,n=1000)
  t = Variable(1)
  A = randn(m,n)
  x = randn(n); # TODO: The more interesting problem is to make x also a variable
  # but then we don't have the solution analytically
  # we can compare it to minimize norm(A*x-b, 'inf') in Matlab with variable x
  b = randn(m)
  p = Problem(:minimize, t, [A*x - b <= t*ones(m,1), -t*ones(m,1) <= A*x - b])
  solve!(p)
  println(p.optval)
  soln = Base.norm(A*x - b, Inf)
  println("Solution and deviation from solution")
  println(soln)
  println(p.optval - soln)
  return p
end

# http://see.stanford.edu/materials/lsocoee364a/hw4sol.pdf
# we get close to optimal with 10, 10
function test4(m=1000, n=1000)
  s = Variable(m)
  x = randn(n); # TODO: This is also more interesting with x = Variable(n)
  # We can compare it to minimize norm(A*x-b, 1) in Matlab with variable x
  A = randn(m,n)
  b = randn(m)
  p = Problem(:minimize, sum(s), [A*x - b <= s, A*x - b >= -s])
  solve!(p)
  return p
end


#Ensemble Kalman Sampler: specific structures and function definitions

"""
    Sampler{FT<:AbstractFloat,IT<:Int} <: Process

An ensemble Kalman Sampler process.
"""
struct Sampler{FT <: AbstractFloat} <: Process
    ""
    prior_mean::Vector{FT}
    ""
    prior_cov::Array{FT, 2}
end

function update_ensemble!(ekp::EnsembleKalmanProcess{FT, IT, Sampler{FT}}, g_in::Array{FT, 2}) where {FT, IT}

    #catch works when g_in non-square
    if !(size(g_in)[2] == ekp.N_ens)
        throw(DimensionMismatch("ensemble size in EnsembleKalmanProcess and g_in do not match, try transposing or check ensemble size"))
    end

    # u: N_ens × N_par
    # g: N_ens × N_obs
    u_old = get_u_final(ekp)
    u_old = permutedims(u_old, (2, 1))
    u = u_old
    g = permutedims(g_in, (2, 1))

    # u_mean: N_par × 1
    u_mean = mean(u', dims = 2)
    # g_mean: N_obs × 1
    g_mean = mean(g', dims = 2)
    # g_cov: N_obs × N_obs
    g_cov = cov(g, corrected = false)
    # u_cov: N_par × N_par
    u_cov = cov(u, corrected = false)

    # Building tmp matrices for EKS update:
    E = g' .- g_mean
    R = g' .- ekp.obs_mean
    # D: N_ens × N_ens
    D = (1 / ekp.N_ens) * (E' * (ekp.obs_noise_cov \ R))

    Δt = 1 / (norm(D) + 1e-8)

    noise = MvNormal(u_cov)

    implicit =
        (1 * Matrix(I, size(u)[2], size(u)[2]) + Δt * (ekp.process.prior_cov' \ u_cov')') \
        (u' .- Δt * (u' .- u_mean) * D .+ Δt * u_cov * (ekp.process.prior_cov \ ekp.process.prior_mean))

    u = implicit' + sqrt(2 * Δt) * rand(noise, ekp.N_ens)'

    # store new parameters (and model outputs)
    push!(ekp.u, DataContainer(u, data_are_columns = false))
    push!(ekp.g, DataContainer(g, data_are_columns = false))
    # u_old is N_ens × N_par, g is N_ens × N_obs,
    # but stored in data container with N_ens as the 2nd dim

    compute_error!(ekp)

end


using JuMP
using Clp
using Test
using JuMPLPUtils
import Random
Random.seed!(1234)
A = rand(40, 10)
c = rand(10)
b = rand(40)
m = Model()
J = 1:10
I = 1:40
@variable(m, x[J] >= 0)

#minimize the cost
z = @objective(m, Min, sum(c[j]*x[j] for j in J))
#Satisfy the constraints
@constraint(m, constraint[i in I], sum(A[i, j]*x[j] for j in J) >= b[i])
setsolver(m, ClpSolver(LogLevel = 0))
solve(m)

N = [1,3,6,7,8,9]
B = [2,4,5,10]
x_val =getvalue(x).innerArray

@test sum(x_val[N]) <= 1e-6
@test abs(sum(x_val[B]) -3.179) <= 1e-3
@test abs(getslack(constraint[10])-.775) <= 1e-3

s = A*x_val - b
@test sum(abs(s[i] -getslack(constraint[i])) for i in I) <= 1e-10

c_int, b_int = getoptranges(m)
z_val = getvalue(z)
pi = getdual(constraint)
for j in J
    dc = .8*max(c_int[j,1],-1.0)
    c[j] += dc
    @objective(m, Min, sum(c[j]*x[j] for j in J))
    solve(m)
    @test sum(abs(x_val[j] - getvalue(x)[j]) for j in J) <= 1e-10
    c[j] -= dc
    dc = .8*min(c_int[j,2],1.0)
    c[j] += dc
    @objective(m, Min, sum(c[j]*x[j] for j in J))
    solve(m)
    @test sum(abs(x_val[j] - getvalue(x)[j]) for j in J) <= 1e-10
    c[j] -= dc
end
@objective(m, Min, sum(c[j]*x[j] for j in J))

for i in I
    db = .8*max(b_int[i,1],-1.0)
    JuMP.setRHS(constraint[i], b[i] + db)
    solve(m)
    @test sum(getvalue(x).innerArray[N]) <= 1e-10
    @test abs(z_val - getvalue(z) + db*pi[i]) <= 1e-10
    db= .8*min(b_int[i,2], 1.0)
    JuMP.setRHS(constraint[i], b[i] + db)
    solve(m)
    @test sum(getvalue(x).innerArray[N]) <= 1e-10
    @test abs(z_val - getvalue(z) + db*pi[i]) <= 1e-10
    JuMP.setRHS(constraint[i], b[i])
end

m = Model()
@variable(m, x[J] >= 0)
#minimize the cost
z = @objective(m, Max, sum(c[j]*x[j] for j in J))
#Satisfy the constraints
@constraint(m, constraint[i in I], sum(A[i, j]*x[j] for j in J) <= b[i])
setsolver(m, ClpSolver(LogLevel = 0))
solve(m)

N = [1,2,4,5,6,7,8,9,10]
B = [3]
x_val =getvalue(x).innerArray

@test sum(x_val[N]) <= 1e-6
@test abs(sum(x_val[B]) - .063688) <= 1e-3
@test abs(getslack(constraint[10])-.5704) <= 1e-3

s = b- A*x_val
@test sum(abs(s[i] -getslack(constraint[i])) for i in I) <= 1e-10

c_int, b_int = getoptranges(m)
z_val = getvalue(z)
pi = getdual(constraint)
for j in J
    dc = .99*max(c_int[j,1],-1.0)
    c[j] += dc
    @objective(m, Max, sum(c[j]*x[j] for j in J))
    solve(m)
    @test sum(abs(x_val[j] - getvalue(x)[j]) for j in J) <= 1e-10
    c[j] -= dc
    dc = .99*min(c_int[j,2],1.0)
    c[j] += dc
    @objective(m, Max, sum(c[j]*x[j] for j in J))
    solve(m)
    @test sum(abs(x_val[j] - getvalue(x)[j]) for j in J) <= 1e-10
    c[j] -= dc
end
@objective(m, Max, sum(c[j]*x[j] for j in J))

for i in I
    db = .99*max(b_int[i,1],-1.0)
    JuMP.setRHS(constraint[i], b[i] + db)
    solve(m)
    @test sum(getvalue(x).innerArray[N]) <= 1e-10
    @test abs(z_val - getvalue(z) + db*pi[i]) <= 1e-10
    db= .99*min(b_int[i,2], 1.0)
    JuMP.setRHS(constraint[i], b[i] + db)
    solve(m)
    @test sum(getvalue(x).innerArray[N]) <= 1e-10
    @test abs(z_val - getvalue(z) + db*pi[i]) <= 1e-10
    JuMP.setRHS(constraint[i], b[i])
end


"""
    SeisMP(in,operator,param,iter,lp,imax)

Robust Matching Pursuit algorithm via lp-norm inner product

# Arguments
- `in`: Noisy data
- `operator`: Any operator passed the dot test
- `param`: parameters for the operator
- `iter`: Max iterations for the matching pursuit
- `lp`: lp-norm space inner product
- `imax`: Max interation for the lp-norm inner product

"""
function SeisMP(in,operator,param,iter,lp,imax)
    Madj=operator(in,param,-1,lp,imax);
    x=zeros(size(Madj));
    dr=zeros(size(in));
    r=in
    for i=1:iter
        Madj=operator(r,param,-1,lp,imax);
        ind=findmax(abs.(Madj))[2];
        M=zeros(size(Madj));
        M[ind[1],ind[2]]=Madj[ind[1],ind[2]];
        g=operator(M,param,1,lp,imax);
        alpha=LP_norm(vec(g),vec(r),lp,imax,1);
        dr=dr+alpha*g
        x=x+alpha*M;
        r=r-alpha*g;
    end
    return dr
end


"""
    tridiagonalize(A, d::Integer=0)

The symmetric real or Hermitian input matrix `A` is transformed to a real
`SymTridiagonal` matrix.
For general matrices only `d` superdiagonals are processed.
Works efficiently for `BandedMatrices`.
"""
function tridiagonalize(A::Union{Symmetric{<:Real},Hermitian}, d::Integer=0)
    m, n = size(A)
    m == n || throw(DimensionMismatch("matrix is not square: dimensions are $((m, n))"))
    dmax = bandwidth(A) + 1
    if d == 0
        d = min(n, dmax)
    end

    0 < d <= dmax || throw(ArgumentError("number of diagonals $d not in 1:$dmax"))

    B = copybands(A, d)
    td = _tridiag_algorithm!(B)
    dv = [real(B[1,i]) for i in 1:n]
    ev = d > 1 ? [-abs(B[2,i]) for i = 2:n] : zeros(real(eltype(A)), n-1)
    SymTridiagonal(dv, ev)
end

# B[1:q,1:n] has the bands of a (2q-1)-banded hermitian nxn-matrix.
# superdiagonal i valid in B[i,i+1:n]. B is completely overwritten.
# The result is in B[1:2,1:nvens2]
function _tridiag_algorithm!(B)
    d, n = size(B)
    0 < d <= n || throw(ArgumentError("number of diagonals $d not in 1:$n"))

    @inbounds for bm = d-1:-1:2
        for k = 1:n-bm
            kp = k
            apiv = B[bm+1,bm+kp]
            iszero(apiv) && continue
            for i = bm+k-1:bm:n-1
                b = B[ i-kp+1,i]
                c, s, r = LinearAlgebra.givensAlgorithm(b, apiv)
                u, v = B[1,i], B[1,i+1]
                upx = (u + v) / 2
                B[1,i] = (u - v) / 2
                B[i-kp+1,i] = r
                for j = kp+1:i
                    u = B[i-j+1,i]
                    v = B[i-j+2,i+1]
                    B[i-j+1,i], B[i-j+2,i+1] = u * c + v * s, -u * s' + v * c
                end
                B[1,i+1] = -(B[1,i])'
                ip = i + bm
                for j = i+1:min(ip, n)
                    u = B[j-i+1,j]
                    v = B[j-i,j]
                    B[j-i+1,j], B[j-i,j] = u * c + v * s', -u * s + v * c
                end
                w = real(B[1,i+1])
                B[1,i] = upx - w
                B[1,i+1] = upx + w
                if ip < n
                    v = B[ip-i+1,ip+1]
                    apiv, B[ip-i+1,ip+1] = v * s', v * c
                end
                kp = i
            end
        end
    end
    B
end

# generalization of method for symmetric BandedMatrices
bandwidth(A::AbstractMatrix) = min(size(A)...) - 1

function copybands(A::AbstractMatrix{T}, d::Integer) where T
    n = min(size(A)...)
    d = min(d, n)
    B = Matrix{T}(undef, d, n)
    for i = 1:d
        B[i,1:i-1] .= zero(T)
        for j = i:n
            B[i,j] = A[j-i+1,j]
        end
    end
    B
end


module Clustering

using LinearAlgebra
using Statistics

export GMM, Kmeans, Xmeans, DBSCAN, HDBSCAN

include("utils.jl")
include("GMM.jl")
include("Kmeans.jl")
include("DBSCAN.jl")
include("HDBSCAN.jl")
include("Xmeans.jl")

end


export PoreData

@doc raw"""
`PoreData` is a collection of physical parameters for coupled geomechanics and flow simulation

- `M`: Biot modulus
- `b`: Biot coefficient
- `ρb`: Bulk density
- `ρf`: Fluid density
- `kp`: Permeability
- `E`: Young modulus
- `ν`: Poisson ratio
- `μ`: Fluid viscosity
- `Pi`: Initial pressure
- `Bf`: formation volume, $B_f=\frac{\rho_{f,0}}{\rho_f}$
- `g`: Gravity acceleration
"""
@with_kw mutable struct PoreData
    M::Float64 = 0.1
    b::Float64 = 1.0
    ρb::Float64 = 2400.
    ρf::Float64 = 1000.
    kp::Float64 = 50e6 * 9.869233e-13 # Darcy
    E::Float64 = 350e6
    ν::Float64 = 0.0
    μ::Float64 = 8.9e-4 # viscosity of water
    Pi::Float64 = 2.125e6
    Bf::Float64 = 1.0
    g::Float64 = 9.807
end


#=
Thermodynamic relationships
=#

####### BASIC PERFECT GAS EQUATIONS OF STATE #######

Temperature(ρ::Density,p::Pressure;gas::PerfectGas=DefaultPerfectGas) = Temperature(p/(ρ*GasConstant(gas)))
Temperature(p::Pressure,ρ::Density;gas::PerfectGas=DefaultPerfectGas) = Temperature(ρ,p,gas=gas)

Density(p::Pressure,T::Temperature;gas::PerfectGas=DefaultPerfectGas) = Density(p/(GasConstant(gas)*T))
Density(T::Temperature,p::Pressure;gas::PerfectGas=DefaultPerfectGas) = Density(p,T,gas=gas)

Pressure(ρ::Density,T::Temperature;gas::PerfectGas=DefaultPerfectGas) = Pressure(ρ*GasConstant(gas)*T)
Pressure(T::Temperature,ρ::Density;gas::PerfectGas=DefaultPerfectGas) = Pressure(ρ,T,gas=gas)

StagnationTemperature(ρ0::StagnationDensity,p0::StagnationPressure;gas::PerfectGas=DefaultPerfectGas) = StagnationTemperature(p0/(ρ0*GasConstant(gas)))
StagnationTemperature(p0::StagnationPressure,ρ0::StagnationDensity;gas::PerfectGas=DefaultPerfectGas) = StagnationTemperature(ρ0,p0,gas=gas)

StagnationDensity(p0::StagnationPressure,T0::StagnationTemperature;gas::PerfectGas=DefaultPerfectGas) = StagnationDensity(p0/(GasConstant(gas)*T0))
StagnationDensity(T0::StagnationTemperature,p0::StagnationPressure;gas::PerfectGas=DefaultPerfectGas) = StagnationDensity(p0,T0,gas=gas)

StagnationPressure(ρ0::StagnationDensity,T0::StagnationTemperature;gas::PerfectGas=DefaultPerfectGas) = StagnationPressure(ρ0*GasConstant(gas)*T0)
StagnationPressure(T0::StagnationTemperature,ρ0::StagnationDensity;gas::PerfectGas=DefaultPerfectGas) = StagnationPressure(ρ0,T0,gas=gas)

"""
    SoundSpeed(T::Temperature[,γ::SpecificHeatRatio=1.4][,R::GasConstant=287])

Compute the speed of sound of a perfect gas, based on the absolute temperature `T`. The ratio of specific heats `γ` and `R`
can also be supplied, but default to the values for air at standard conditions
(1.4 and 287 J/kg.K, respectively).
"""
SoundSpeed(T::Temperature;gas::PerfectGas=DefaultPerfectGas) = SoundSpeed(sqrt(SpecificHeatRatio(gas)*GasConstant(gas)*T))

Enthalpy(T::Temperature;gas::PerfectGas=DefaultPerfectGas) = Enthalpy(SpecificHeatPressure(gas)*T)
Temperature(h::Enthalpy;gas::PerfectGas=DefaultPerfectGas) = Temperature(h/SpecificHeatPressure(gas))

StagnationEnthalpy(T0::StagnationTemperature;gas::PerfectGas=DefaultPerfectGas) = StagnationEnthalpy(SpecificHeatPressure(gas)*T0)
StagnationTemperature(h0::StagnationEnthalpy;gas::PerfectGas=DefaultPerfectGas) = StagnationTemperature(h0/SpecificHeatPressure(gas))

InternalEnergy(T::Temperature;gas::PerfectGas=DefaultPerfectGas) = InternalEnergy(SpecificHeatVolume(gas)*T)
Temperature(e::InternalEnergy;gas::PerfectGas=DefaultPerfectGas) = Temperature(e/SpecificHeatVolume(gas))

StagnationInternalEnergy(T0::StagnationTemperature;gas::PerfectGas=DefaultPerfectGas) = StagnationInternalEnergy(SpecificHeatVolume(gas)*T0)
StagnationTemperature(e0::StagnationInternalEnergy;gas::PerfectGas=DefaultPerfectGas) = StagnationTemperature(e0/SpecificHeatVolume(gas))

StagnationEnthalpy(h::Enthalpy,u::Velocity) = StagnationEnthalpy(value(h)+0.5*value(u)^2)
Enthalpy(h0::StagnationEnthalpy,u::Velocity) = Enthalpy(value(h0)-0.5*value(u)^2)

# Mach number - velocity relations
Velocity(a::SoundSpeed,M::MachNumber) = Velocity(M*a)
MachNumber(u::Velocity,a::SoundSpeed) = MachNumber(u/a)
SoundSpeed(u::Velocity,M::MachNumber) = SoundSpeed(u/M)

# mass flow rate
MassFlowRate(ρ::Density,u::Velocity,A::Area) = MassFlowRate(ρ*u*A)


@testset "test_pid_design" begin
@test stabregionPID("exp(-sqrt(s))", exp10.(range(-3, stop=3, length=50)), kd=1)[2][1] ≈ -1.022356911142034

P = tf(1,[1,1])^4
gangoffourplot(P,tf(1))
gangoffourplot(ss(P),ss(1))
ωp = 0.8
kp,ki,C = loopshapingPI(P,ωp,phasemargin=60, doplot=true)
@test kp ≈ 0.82274734724854
@test ki ≈ -0.45472580052281375
end


@testset "Reducing ODE mod p" begin

    ode = @ODEmodel(
        x'(t) = 1//2 * x(t) - 5 * y(t),
        y'(t) = (3//5 * x(t) + 17 * y(t)) // (y(t) - 1),
        z(t) = y(t)
    )

    ode_red = reduce_ode_mod_p(ode, 17)

    x = str_to_var("x", ode_red.poly_ring)
    y = str_to_var("y", ode_red.poly_ring)
    @test ode_red.x_equations[x] == 9 * x + 12 * y
    @test ode_red.x_equations[y] == 4 * x // (y + 16)

end


function init_2(t,
                fft_u0, fft_v0,
                A1, A2,
                matr, conjmatr,
                sigma, llambda,
                ktau, epsilon, ntau)

    champu, champv = ftau(t, fft_u0, fft_v0, A1, A2,
                          matr, conjmatr, sigma, llambda)

    champu_fft = fft(champu, 1)
    champv_fft = fft(champv, 1)

    champu_fft[1, :] .= 0.0
    champv_fft[1, :] .= 0.0

    champu = ifft(champu_fft ./ (1im * ktau), 1)
    champv = ifft(champv_fft ./ (1im * ktau), 1)

    fft_ubar = fft_u0 .- epsilon * transpose(champu[1, :])
    fft_vbar = fft_v0 .- epsilon * transpose(champv[1, :])

    println("fft_ubar:",size(fft_ubar))

    C1u, C1v = C1(t, fft_ubar, fft_vbar, A1, A2, matr,
                  conjmatr, sigma, llambda, ktau, epsilon, ntau)

    fft_ug = fft_u0 .- transpose(C1u[1, :])
    fft_vg = fft_v0 .- transpose(C1v[1, :])

    fft_ubar, fft_vbar, fft_ug, fft_vg

end


# specialized likelihood for MCMC using net
function LL(θ, m, S, model, info)
    k = size(m,1)
    ms = zeros(S, k)
    Threads.@threads for s = 1:S
        ms[s,:] = Float64.(model(transform(WileE_model(θ[:])', info)'))
    end
    mbar = mean(ms,dims=1)[:]
    Σ = cov(ms)
    x = (m .- mbar)[:]
    lnL = try
        lnL =-0.5*log(det(Σ)) - 0.5*x'*inv(Σ)*x # for Bayesian
    catch
        lnL = -Inf
    end
    return lnL
end

# version without net
function LL(θ, m, S)
    k = size(m,1)
    ms = zeros(S, k)
    Threads.@threads for s = 1:S
        ms[s,:] = WileE_model(θ[:])
    end
    mbar = mean(ms,dims=1)[:]
    Σ = cov(ms)
    x = (m .- mbar)[:]
    lnL = try
        lnL =-0.5*log(det(Σ)) - 0.5*x'*inv(Σ)*x # for Bayesian
    catch
        lnL = -Inf
    end
    return lnL
end