romeo = unwrap # access unwrap function via alias romeo
romeo! = unwrap!

function mask_from_voxelquality(qmap, threshold)
    qmap_bin = qmap .> threshold # NaN defaults to false (0)
    max_hole_size = length(qmap) / 20
    qmap_bin = .!imfill(.!qmap_bin, (1, max_hole_size)) # fills all holes up to max_hole_size (uses 6 connectivity as default)
    return gaussiansmooth3d(qmap_bin; padding=true) .> 0.8 # hardcoded final threshold
end

function ROMEO.calculateweights(phase::AbstractArray{T,4}; TEs, template=2, p2ref=1, keyargs...) where T
    args = Dict{Symbol, Any}(keyargs)
    args[:phase2] = phase[:,:,:,p2ref]
    args[:TEs] = TEs[[template, p2ref]]
    if haskey(args, :mag)
        args[:mag] = args[:mag][:,:,:,template]
    end
    return ROMEO.calculateweights(view(phase,:,:,:,template); args...)
end

# calculates B0 in [Hz]
function calculateB0_unwrapped(unwrapped_phase, mag, TEs)
    dims = 4
    TEs = to_dim(TEs, 4)
    return (1000 / 2π) * sum(unwrapped_phase .* mag .* mag .* TEs; dims) ./ sum(mag .* mag .* TEs.^2; dims) |> I -> dropdims(I; dims)
end

"""
Calculates a quality for each voxel. Takes the same inputs as ROMEO
"""
function romeovoxelquality(phase; keyargs...)
    weights = ROMEO.calculateweights(phase; type=Float32, rescale=x->x, keyargs...)
    return dropdims(sum(weights; dims=1); dims=1) ./ 3 # [0;1]
end


using FastConv
using Test

A = zeros(5,5)
A[3,3] = 1
k = rand(3,3)
Ac = convn(A,k)
@test Ac[3:5,3:5] == k && all(a->a==0, Ac[:,[1,2,6,7]]) && all(a->a==0, Ac[[1,2,6,7], :])

A = zeros(5,6,7)
A[3,4,5] = 1
k = rand(3,3,3)
Ac = convn(A,k)
@test Ac[3:5,4:6,5:7] == k &&
    all(a->a==0, Ac[[1,2,6,7], :, :]) &&
    all(a->a==0, Ac[:, [1,2,3,7,8], :]) &&
    all(a->a==0, Ac[:, :, [1,2,3,4,8,9]])

A = randn(5,5)
k = ones(1,1)
Ac = convn(A,k)
@test Ac == A

A = randn(5,5)
k = zeros(3,3)
Ac = convn(A,k)
@test Ac == zeros(7,7)

A = randn(5,5)
k1 = rand(3,3)
k2 = rand(3,3)
A1_2 = convn(convn(A,k1), k2)
A_12 = fastconv(A, fastconv(k1, k2))
@test A1_2 \approx A_12


"""
QQQQ
"""
function test_fit_example1(;sim_runs=10^2)
    Λ₄, λ45, λ54, Λ₅ = 15.0, 5.0, 7.0, 16.0
    μ41, μ42, μ43, μ51, μ52, μ53 = 4.0, 3.0, 3.0, 1.0, 7.0, 1.0
    T_example = [-Λ₄ λ45; λ54 -Λ₅]
    T0_example = [μ41 μ42 μ43; μ51 μ52 μ53]
    initial_dist = [0.5,0.5]

    maph = MAPHDist(initial_dist',T_example, T0_example)

    @show maph

    absorbtion_probs_of_true_dist = absorption_probs(maph)
    mean_of_true_dist = mean(maph)
    scvs_of_true_dist = scv(maph)

    times = zeros(sim_runs)
    absorbing_states = zeros(Int,sim_runs)

    @showprogress "Simulating data" for i in 1:sim_runs
        time, state = rand(maph, full_trace = false)
        times[i] = time
        absorbing_states[i] = state
    end

    estimated_dist = fit_maph(times, absorbing_states, 4, max_iter = 0)

    absorbtion_probs_of_estimated_dist = absorption_probs(estimated_dist)
    mean_of_estimated_dist = mean(estimated_dist)
    scvs_of_estimated_dist = scv(estimated_dist)

    @show absorbtion_probs_of_true_dist,absorbtion_probs_of_estimated_dist
    @show mean_of_estimated_dist, mean_of_true_dist
    @show scvs_of_estimated_dist, scvs_of_true_dist

    return true
end


include("AlgAssOrd/AlgAssOrd.jl")
include("AlgAssOrd/Ideal.jl")
include("AlgAssOrd/FracIdeal.jl")
include("AlgAssOrd/Elem.jl")
include("AlgAssOrd/PicardGroup.jl")
include("AlgAssOrd/ResidueRingMultGrp.jl")
include("AlgAssOrd/LocallyFreeClassGroup.jl")


using H3
using H3.Lib
using H3.Lib: H3Index, GeoCoord, Geofence, GeoPolygon
import H3.API.geoToH3
using Geodesy
using DataStructures

using Plots
using HexGeometry

# The goal here is to search a plausible sized grid for the nearest n point.
# The MRMS data has a grid 3500 x 7000 points which are roughly 1km apart
# Here, we will use a test grid that is evenly spaced in latitude from 20-50 degrees
# and longitude from -55 to -120 degrees

struct Location
    lat::Float64
    lon::Float64
end

"""
grid(res=(7000,3500), r=1000, maxResolution=11, minResolution=6)

Creates a grid of points with a specified number of points in east-west and
north-south directions (defaul 7000ew x 3500ns). Points will have a specified
spacing (default 1000m = 1km).

In addition to locating points on a grid, an index will be created
so that points can be found using different scales of H3Index.
"""
function grid(;res=(7000,3500), r=1000.0, maxResolution=11, minResolution=6)
    center = (40.0,-95)
    # planar projection tangent to roughly the center of continental US
    tx = ENUfromLLA(LLA(center...,0), wgs84)
    ix = inv(tx)
    grid = Vector{Location}()
    index = DefaultDict{H3Index, Vector{Int32}}(Vector{Int32})
    x0, y0 = res ./ 2
    id = 0
    grid = Vector{Location}(undef, reduce(*, res .+ 1))
    for x = -x0:x0, y = -y0:y0
        id += 1
        px = ix(ENU(x*r, y*r, 0))
        grid[id] = Location(px.lat, px.lon)
        for resolution in minResolution:maxResolution
            hex = geoToH3(px.lat, px.lon, resolution)
            push!(index[hex], id)
        end
    end
    return grid, index
end

function nearest(query, grid, index, range)
    tx = ENUfromLLA(LLA(query...), wgs84)
    for res in max(range...):-1:min(range...)
        q = geoToH3(query..., res)
        r = [grid[z] for u in map(x -> i[x], H3.API.kRing(q, 1)) for z in u]
        if length(r) > 1
            p = map(x -> tx(LLA(x.lat, x.lon, 0.0)), r)
            distance = map(x->sum(x.^2), p)
            best = findfirst(distance .== reduce(min, distance))
            return [r[best], sqrt(distance[best])]
        end
    end
    return nothing
end


using QuasiCopula, LinearAlgebra, Random, GLM
using DataFrames, Statistics
using BenchmarkTools, Test

BLAS.set_num_threads(1)
Threads.nthreads()

p_fixed = 3    # number of fixed effects, including intercept
m = 1    # number of variance components
# true parameter values
Random.seed!(12345)
# try next
βtrue = rand(Uniform(-2, 2), p_fixed)
θtrue = [0.5]
rtrue = 10.0

d = NegativeBinomial()
link = LogLink()
D = typeof(d)
Link = typeof(link)
T = Float64

samplesize = 10000
ni = 25

gcs = Vector{NBCopulaVCObs{T, D, Link}}(undef, samplesize)

Γ = θtrue[1] * ones(ni, ni)

# for reproducibility I will simulate all the design matrices here
Random.seed!(12345)
X_samplesize = [randn(ni, p_fixed - 1) for i in 1:samplesize]

for i in 1:samplesize
    X = [ones(ni) X_samplesize[i]]
    η = X * βtrue
    μ = exp.(η)
    p = rtrue ./ (μ .+ rtrue)
    vecd = Vector{DiscreteUnivariateDistribution}(undef, ni)
    vecd = [NegativeBinomial(rtrue, p[i]) for i in 1:ni]
    nonmixed_multivariate_dist = NonMixedMultivariateDistribution(vecd, Γ)
    # simuate single vector y
    y = Vector{Float64}(undef, ni)
    res = Vector{Float64}(undef, ni)
    rand(nonmixed_multivariate_dist, y, res)
    V = [ones(ni, ni)]
    gcs[i] = NBCopulaVCObs(y, X, V, d, link)
end

# form VarLmmModel
gcm = NBCopulaVCModel(gcs)
# precompile
println("precompiling NB VCM fit")
gcm2 = deepcopy(gcm);
QuasiCopula.fit!(gcm2, maxBlockIter = 1);

fittime = @elapsed QuasiCopula.fit!(gcm)
@show fittime
@show gcm.β
@show gcm.θ
@show gcm.r
@show gcm.∇β
@show gcm.∇θ
@show gcm.∇r

loglikelihood!(gcm, true, true)
vcov!(gcm)
@show QuasiCopula.confint(gcm)
# mse and time under our model
mseβ, mser, mseθ = MSE(gcm, βtrue, rtrue, θtrue)
@show mseβ
@show mser
@show mseθ

@test mseβ < 0.01
@test mseθ < 0.01
@test mser < 0.1


using ProximalOperators
using ProximalOperators: Zero

ProximalOperators.Conjugate(f::Zero) = IndZero()
ProximalOperators.Conjugate(f::IndZero) = Zero()
ProximalOperators.Conjugate(f::SqrNormL2) = SqrNormL2(1.0 / f.lambda)

# TODO: Add other useful functions and calculus rules such as translation


using RandomWalks

# using all of the following just so Revise knows about them
using RandomWalks.Actors
using RandomWalks.Lattices
using RandomWalks.LatticeVars
using RandomWalks.LatticeWalks
using RandomWalks.Points
using RandomWalks.WalksBase
using RandomWalks.Walks

using Distributions, BenchmarkTools
using EmpiricalCDFs
using DataFrames

struct WalkParams
    dimension::Int
    ntrials::Int
    nsteps::Int
    lambda::Float64
    alpha::Float64
end

WalkParams() = WalkParams(0,0,0,0.0,0.0)

function Base.print(io::IO, wp::WalkParams)
    println(io, "dimension = ", wp.dimension)
    println(io, "ntrials = ", wp.ntrials)
    println(io, "lambda = ", wp.lambda)
    print(io, "alpha = ", wp.alpha)
#    println(io, "nsteps = ", wp.nsteps)
end

function getcdfs(; dimension = 2, ntrials = 10^3, nsteps = 10^3, lambda = 1e2)
    alpha = 1/2
    wparams = WalkParams(
        dimension,
        ntrials,
        nsteps,
        lambda,
        alpha
        )

    lattice_rates = LatticeVar{dimension}(Exponential(lambda))
    waiting_times = LatticeVarParam{dimension}(Pareto(alpha))
    quenched_lattice = Lattice(waiting_times, lattice_rates)

    walk = MortalWalk(WalkB{dimension}())
    lattice_walk = LatticeWalk(quenched_lattice, walk)

    log_store_range = 1:0.5:12.0

    storing_actor = storing_nsteps_num_sites_visited_actor(logrange(log_store_range))
#    storing_actor = storing_num_sites_visited_actor(logrange(log_store_range))
#    storing_actor2 = storing_num_sites_visited_actor(logrange(log_store_range))

    step_limit_actor = StepLimitActor(nsteps)
    theactors = ActorSet(step_limit_actor, storing_actor)
    walk_plan = WalkPlan(lattice_walk, theactors)
    ecdfs_actor = ECDFsActor(storing_actor);
    max_step_counter = CountActor(step_limit_actor)
    trial_actors = ActorSet(max_step_counter, ecdfs_actor)
    @time trial!(walk_plan, SampleLoopActor(ProgressIter(ntrials), trial_actors));
    check_counter(max_step_counter, "Hit step limit")
    tcdfs = get_ecdf_times(ecdfs_actor);
    ptcdfs = prune(tcdfs)
    return (ptcdfs, ntrials, trial_actors, wparams)
end

function cdfrep(ptcdfs, ntrials)
    cdfs = get_cdfs(ptcdfs)
    times = get_times(ptcdfs)
    nsites = cdfs[:,2]
    nsteps = cdfs[:,1]
    DataFrame([times, length.(nsteps), length.(nsteps) ./ ntrials,  median.(nsteps), extrema.(nsites),
               median.(nsites), mean.(nsites), std.(nsites)],
              [:time, :ncounts, :survp, :median_steps, :minmax_sites, :median_nsites, :mean_nsites, :std_nsites])
end

function vstrials(; dimension = 2, nsteps = 10^2, lambda = 1e3)
    for nexp in 1:7
        for fac in (1, 3, 6)
            (tcdfs, ntrials, trial_actors, wparams) = getcdfs(;ntrials = fac * 10^nexp,
                               nsteps=nsteps, lambda=lambda, dimension = dimension)
            df = cdfrep(tcdfs, ntrials)
            println(wparams)
            println(df)
            println()
        end
    end
end

nothing


using PyPlot
using Jagot.plotting
plot_style("ggplot")

using PyCall
Cycler = pyimport("cycler")
plt.rc("axes", prop_cycle=plt.rcParams["axes.prop_cycle"])

using SphericalBesselFunctions

include(joinpath(@__DIR__, "../test/reference.jl"))

function simple_example()
    nx = 1000
    x = range(-12, stop=12, length=nx)

    nℓ = 10
    j, j′, y, y′ = bessels(x, nℓ)

    cfigure("simple example",figsize=(7,6)) do
        csubplot(221,nox=true) do
            plot(x, j)
            margins(0, 0.1)
            ylim(-0.5, 0.5)
            ylabel(L"j_n(x)")
        end
        csubplot(222,nox=true) do
            plot(x, y)
            margins(0, 0.1)
            ylim(-0.5, 0.5)
            ylabel(L"y_n(x)")
            axes_labels_opposite(:y)
        end
        csubplot(223) do
            plot(x, j′)
            margins(0, 0.1)
            ylim(-0.5, 0.5)
            xlabel(L"x")
            ylabel(L"j'_n(x)")
        end
        csubplot(224) do
            plot(x, y′)
            margins(0, 0.1)
            ylim(-0.5, 0.5)
            xlabel(L"x")
            ylabel(L"y'_n(x)")
            axes_labels_opposite(:y)
        end
    end
    savefig("docs/src/figures/simple-example.svg")
end

function accuracy()
    nx = 1001
    x = 10 .^ range(-1, stop=4, length=nx)
    nℓ = 105
    j, j′, y, y′ = @time bessels(x, nℓ)

    sel = 1:min(nℓ,104)

    nℓref = min(nℓ,length(sel))
    jref = zeros(nx, nℓref)
    j′ref = zeros(nx, nℓref)
    yref = zeros(nx, nℓref)
    y′ref = zeros(nx, nℓref)

    @time for (i,x) in enumerate(x)
        bessels_ref!(view(jref, i, :),
                     view(j′ref, i, :),
                     view(yref, i, :),
                     view(y′ref, i, :),
                     x)
    end

    j = j[:,sel]
    j′ = j′[:,sel]
    y = y[:,sel]
    y′ = y′[:,sel]
    jref = jref[:,sel]
    j′ref = j′ref[:,sel]
    yref = yref[:,sel]
    y′ref = y′ref[:,sel]

    cfigure("accuracy", figsize=(9,10)) do
        m,n = 4,3
        ks = reshape(1:m*n, n, m)'

        csubplot(m,n,ks[1,1],nox=true) do
            semilogx(x, j, rasterized=true)
            ylim(-0.5,0.5)
            ylabel(L"j_n(x)")
            title("SphericalBesselFunctions.jl")
        end
        csubplot(m,n,ks[2,1],nox=true) do
            semilogx(x, j′, rasterized=true)
            ylim(-0.5,0.5)
            ylabel(L"j_n'(x)")
        end
        csubplot(m,n,ks[3,1],nox=true) do
            semilogx(x, y, rasterized=true)
            ylim(-0.5,0.5)
            ylabel(L"y_n(x)")
        end
        csubplot(m,n,ks[4,1]) do
            semilogx(x, y′, rasterized=true)
            ylim(-0.5,0.5)
            xlabel(L"x")
            ylabel(L"y_n'(x)")
        end

        csubplot(m,n,ks[1,2],nox=true,noy=true) do
            semilogx(x, jref, rasterized=true)
            ylim(-0.5,0.5)
            title("SpecialFunctions.jl")
        end
        csubplot(m,n,ks[2,2],nox=true,noy=true) do
            semilogx(x, j′ref, rasterized=true)
            ylim(-0.5,0.5)
        end
        csubplot(m,n,ks[3,2],nox=true,noy=true) do
            semilogx(x, yref, rasterized=true)
            ylim(-0.5,0.5)
        end
        csubplot(m,n,ks[4,2],noy=true) do
            semilogx(x, y′ref, rasterized=true)
            ylim(-0.5,0.5)
            xlabel(L"x")
        end

        csubplot(m,n,ks[1,3],nox=true) do
            jerror = abs.(j-jref)
            loglog(x, jerror, rasterized=true)
            ylim(max(1e-19,minimum(jerror)), max(maximum(jerror),1e-4))
            title("Difference")
            axes_labels_opposite(:y)
        end
        csubplot(m,n,ks[2,3],nox=true) do
            j′error = abs.(j′-j′ref)
            loglog(x, j′error, rasterized=true)
            ylim(max(1e-19,minimum(j′error)), max(maximum(j′error),1e-4))
            axes_labels_opposite(:y)
        end
        csubplot(m,n,ks[3,3],nox=true) do
            yerror = abs.(y-yref)
            loglog(x, yerror, rasterized=true)
            ylim(max(1e-19,minimum(yerror)), max(maximum(yerror),1e-4))
            axes_labels_opposite(:y)
        end
        csubplot(m,n,ks[4,3]) do
            loglog(x, abs.(y′-y′ref), rasterized=true)
            xlabel(L"x")
            axes_labels_opposite(:y)
        end
    end
    savefig("docs/src/figures/accuracy.svg")
end

macro echo(expr)
    println(expr)
    :(@time $expr)
end

@info "Documentation plots"
mkpath("docs/src/figures")
@echo simple_example()
@echo accuracy()


# Type
"""
	(type) Ray{T<:AnyComplex} in the complex plane
Each `Ray` type is parameterized according to the common type of its complex input arguments.
"""
struct Ray{T<:AnyComplex} <: AbstractCurve
	base::T
	angle::AbstractFloat
	reverse::Bool
	function Ray{T}(a,d,rev=false) where T<:AnyComplex
		new(a,mod2pi(d),rev)
	end
end

# Untyped constructor
"""
	Ray(a,θ,reverse=false)
Construct the ray starting at `a` and extending to infinity at the angle `θ`. If `reverse` is true, the ray is considered to extend from infinity to `a` at angle `-θ`.
"""
function Ray(a::Number,θ::Real,rev=false)
	a = complex(float(a))
	Ray{typeof(a)}(a,float(θ),rev)
end

# Complex type converters
for ctype in [:Spherical,:Polar,:Complex]
	@eval begin
		function $ctype(R::Ray{T}) where T<:AnyComplex
			Ray($ctype(R.base),R.angle,R.reverse)
		end
	end
end

# Required methods
arclength(R::Ray) = Inf
function point(R::Ray{T},t::Real) where T
	if R.reverse
		t = 1-t
	end
	# avoid NaNs
	if t==0
		R.base
	elseif t==1
		T(Inf)
	else
		R.base + t/(1-t)*exp(complex(0,R.angle))
	end
end
(C::Ray)(t::Real) = point(C,t)

"""
	arg(R::Ray,z)
Find the parameter argument `t` such that `R(t)==z` is true.

This gives undefined results if `z` is not actually on the ray.
"""
function arg(R::Ray,z::Number)
	if isinf(z)
		t = 1
	else
		δ = abs(z - R.base)
		t = δ / (1+δ)
	end
	return R.reverse ? 1-t : t
end

function unittangent(R::Ray{T},t::Real=0) where T <: AnyComplex
	τ = T( exp(complex(0,R.angle)) )
	R.reverse ? -τ : τ
end

function tangent(R::Ray{T},t::Real) where T <: AnyComplex
	τ = exp(complex(0,R.angle))
	R.reverse ? T(-τ/t^2) : T(τ/(1-t)^2)
end

+(R::Ray,z::Number) = Ray(R.base+z,R.angle,R.reverse)
-(R::Ray) = Ray(-R.base,mod2pi(R.angle+pi),R.reverse)
*(R::Ray,z::Number) = Ray(z*R.base,mod2pi(R.angle+sign(z)),R.reverse)

"""
	inv(R)
Invert the ray `R` through the origin. In general the inverse is an `Arc`.
"""
function inv(R::Ray)
	w = 1 ./ point(R,[0,0.5,1])
	Arc(w...)
end

# Other methods
isfinite(::Ray) = false
conj(R::Ray) = Ray(conj(R.base),-R.angle,R.reverse)
reverse(R::Ray) = Ray(R.base,R.angle,!R.reverse)

"""
	isapprox(R1::Ray,R2::Ray; tol=<default>)
	R1 ≈ R2
Determine if `R1` and `R2` represent the same segment, irrespective of the type or values of its parameters. Identity is determined by agreement within `tol`, which is interpreted as the weaker of absolute and relative differences.
"""
function isapprox(R1::Ray,R2::Ray;tol=DEFAULT[:tol])
	return isapprox(R1.base,R2.base,rtol=tol,atol=tol) && (abs(mod2pi(R1.angle-R2.angle)) < tol)
end

"""
	isleft(z,R::Ray)
Determine whether the number `z` lies "to the left" of ray `R`. This means that the angle it makes with `tangent(R)` is in the interval (0,π).

Note that `isleft` and `isright` are *not* logical opposites; a point on the (extended) ray should give `false` in both cases.
"""
function isleft(z::Number,R::Ray)
	a,b = point(R,[0.2,0.8])  # accounts for reversal
	(real(b)-real(a)) * (imag(z)-imag(a)) > (real(z)-real(a)) * (imag(b)-imag(a))
end

"""
	isright(z,R::Ray)
Determine whether the number `z` lies "to the right" of ray `R`. This means that the angle it makes with `tangent(R)` is in the interval (-π,0).

Note that `isleft` and `isright` are *not* logical opposites; a point on the (extended) ray should give `false` in both cases.
"""
function isright(z::Number,R::Ray)
	a,b = point(R,[0.2,0.8])  # accounts for reversal
	(real(b)-real(a)) * (imag(z)-imag(a)) < (real(z)-real(a)) * (imag(b)-imag(a))
end

"""
	dist(z,R::Ray)
Compute the distance from number `z` to the ray `R`.
"""
dist(z::Number,R::Ray) = abs(z - closest(z,R))

"""
	closest(z,R::Ray)

Find the point on ray `R` that lies closest to `z`.
"""
function closest(z::Number,R::Ray)
	# translate and rotate to positive Re axis
	s = exp(complex(0,R.angle))
	ζ = (z-R.base)/s
	R.base + max(real(ζ),0)*s
end

sign(R::Ray) = unittangent(R)

# Display methods
function show(io::IO,R::Ray{T}) where {T}
	print(IOContext(io,:compact=>true),"Ray(",R.base,",",R.angle,",",R.reverse,")")
end

function show(io::IO,::MIME"text/plain",R::Ray{T}) where {T}
	if R.reverse
		print(io,"Ray from ∞ to ",R.base," at angle ",R.angle)
	else
		print(io,"Ray from ",R.base," to ∞ at angle ",R.angle)
	end
end

# in the plane, use two (finite) points for plotting
plotdata(R::Ray{T}) where T<:Union{Complex,Polar} = R.reverse ? [R(0.3),R.base] : [R.base,R(0.7)]


count_nonzeros(IJ, mn) = [count(isequal(ij), IJ) for ij in 1:mn]

"""
    nonzero_minors(N, overlap) -> (ks,ls)

Find all (distinct) minor determinants of order `N` of the orbital
`overlap` matrix that do not vanish, i.e. all non-vanishing minors are
guaranteed to be present, but not all of the returned minors are
guaranteed to be non-zero. Vanishing minors returned arise when the
overlap matrix is rank deficient, which is unlikely to happen when
computing energy expressions, but must still be guarded against. This
is most easily checked by actually calculating the [`cofactor`](@ref),
which is most likely desired anyway.
"""
function nonzero_minors(N::Integer, overlap::AbstractSparseMatrix)
    ks = Vector{Vector{Int}}()
    ls = Vector{Vector{Int}}()

    (I,J,V) = findnz(overlap)
    m,n = size(overlap)
    m == n || throw(ArgumentError("Overlap matrix must be square"))

    # First deal with trivial special cases.
    N > m && return ks, ls
    if N == m
        push!(ks, collect(1:m))
        push!(ls, collect(1:m))
        return ks, ls
    end

    # Number of non-zero matrix elements for each row/column.
    nzrowcount = count_nonzeros(I, m)
    nzcolcount = count_nonzeros(J, n)
    # Vanishing rows/columns.
    zrows = findall(iszero, nzrowcount)
    zcols = findall(iszero, nzcolcount)
    # Minimum order of Γ⁽ᵖ⁾
    prmin = length(zrows)
    pcmin = length(zcols)
    pmin = max(prmin, pcmin)

    # Then deal with next set of trivial special cases.
    N < pmin && return ks, ls
    if N == prmin && N == pcmin
        push!(ks, zrows)
        push!(ls, zcols)
        return ks, ls
    end

    # The general case.
    Nr = N - prmin
    Nc = N - pcmin
    # Find all rows supporting a specific column j.
    I′ = [I[findall(isequal(j), J)] for j in 1:n]
    # Find all columns supporting a specific row i.
    J′ = [J[findall(isequal(i), I)] for i in 1:m]

    for k in combinations(unique(I),Nr)
        # Find all columns supported by any of the rows in the
        # combination k.
        cols = unique(vcat([J′[i] for i in k]...))
        # For each combination k of rows, we can at most strike
        # out Nc columns.
        colbudget = Nc
        colsmuststrike = Vector{Int}()

        for j in cols
            # Test if column j is solely supported by the row
            # combination k.
            colsupport = setdiff(I′[j], k)
            if isempty(colsupport)
                colbudget -= 1
                colbudget < 0 && break
                push!(colsmuststrike, j)
            end
        end
        colbudget < 0 && continue

        Ncm = length(colsmuststrike)

        # The columns we must strike out fill the column budget.
        if Ncm == Nc
            push!(ks, sort(vcat(k,zrows)))
            push!(ls, sort(vcat(colsmuststrike,zcols)))
            continue
        end

        # Find all combinations of the remaining columns.
        colchoices = combinations(setdiff(unique(J), colsmuststrike), Nc-Ncm)
        for l in colchoices
            # Find all rows that would be affected if column
            # combination l is stricken out.
            colsupport = unique(vcat([I′[j] for j in l]...))
            supported = true
            for i in colsupport
                # If the row i supporting one of the columns in the
                # column combination l is already a candidate to be
                # stricken out, it does not matter if its support
                # vanishes.
                i ∈ k && continue
                # Otherwise, if its support vanishes, then l is an
                # unviable column combination.
                rowsupport = setdiff(J′[i], l)
                if isempty(rowsupport)
                    supported = false
                    break
                end
            end
            if supported
                push!(ks, sort(vcat(k,zrows)))
                push!(ls, sort(vcat(l,colsmuststrike,zcols)))
            end
        end
    end

    ks, ls
end


"""
Optimization function to calculate warping for elastic regression

    regression_warp(beta, timet, q, y, alpha)
    :param beta: regression function
    :param timet: vector describing time samples
    :param q: vector describing srsf
    :param y: response value
    :param alpha: intercept

    Returns
    :return gamma_new: new gamma
"""
function regression_warp(beta::Vector, timet::Vector, q::Vector, y::Float64,
                         alpha::Float64)
    gam_M = optimum_reparam(beta, timet, q, method="DP2");
    qM = warp_q_gamma(timet, q, gam_M);
    y_M = trapz(timet, qM.*beta);

    gam_m = optimum_reparam(-1.*beta, timet, q, method="DP2");
    qm = warp_q_gamma(timet, q, gam_m);
    y_m = trapz(timet, qm.*beta);

    if y > alpha+y_M
        gamma_new = gam_M;
    elseif y < alpha + y_m
        gamma_new = gam_m;
    else
        gamma_new = zero_crossing(y-alpha, q, beta, timet, y_M, y_m,
                                  gam_M, gam_m);
    end

    return gamma_new

end

"""
Calculate warping for logistic regression

    logistic_warp(beta, timet, q, y)
    :param beta: regression function
    :param timet: time samples
    :param q: srsf
    :param y: response

    Returns
    :return gamma: new gamma
"""
function logistic_warp(beta::Vector, timet::Vector, q::Array, y)
    if y == 1
        gamma = optimum_reparam(beta, timet, q, method="DP2");
    elseif y == -1
        gamma = optimum_reparam(-1.*beta, timet, q, method="DP2");
    end

    return gamma

end

"""
Calculate logistic optimization function

    logit_optm(x::Vector, grad::Vector, Phi, y)
    :param x: samples
    :param grad: gradient
    :param Phi: coefficient matrix
    :param y: response
"""
function logit_optm(x::Vector, grad::Vector, Phi, y)
    if length(grad) > 0
        logit_gradient!(x, grad, Phi, y);
    end
    out = logit_loss(x, Phi, y);
    return out
end

"""
Logistic function

    phi(t)

    Returns
    :return out: phi(t)
"""
function phi(t)
    idx = t .> 0;
    out = Array(Float64, size(t));
    if sum(idx) > 0
        out[idx] = 1./(1+exp(-1*t[idx]));
    end
    exp_t = exp(t[~idx]);
    out[~idx] = exp_t ./ (1 + exp_t);
    return out
end

"""
Calculate logistic loss function

    logit_loss(b, X, y)
    :param b: coefficients
    :param X: matrix
    :param y: response

    Returns
    :return out: loss function
"""
function logit_loss(b, X, y)
    z = X * b;
    yz = y .* z;
    idx = yz .> 0;
    out = zeros(yz);
    out[idx] = log(1+exp(-1.*yz[idx]));
    out[~idx] = (-yz[~idx] + log(1+exp(yz[~idx])));
    out = sum(out);
    return out
end

"""
Calculate gradient of logistic optimization in place

    logit_gradient!(b, grad, X, y)
    :param b: coefficients
    :param grad: gradient
    :param X: matrix
    :param y: response
"""
function logit_gradient!(b, grad, X, y)
    z = X * b;
    z = phi(y.*z);
    z0 = (z-1) .* y;
    grad[:] = X.' * z0;
end

"""
Calculate Hessian of logistic optimization

    logit_hessian(s, b, X, y)
    :param s:
    :param b: coefficients
    :param X: matrix
    :param y: response
"""
function logit_hessian(s, b, X, y)
    z = X*b;
    z = phi(y.*z);
    d = z*(1-z);
    wa = d.*(X * s);
    Hs = X.' * wa;
    return Hs
end

"""
Calculate m-logistic warping using gradient method

    mlogit_warp_grad(alpha, beta, timet, q, y; max_itr=8000, tol=1e-10,
                     delt=0.008, display=0)
    :param alpha: intercept
    :param beta: regression function
    :param timet: vector describing time samples
    :param q: srsf
    :param y: response
    :param max_itr: maximum number of iterations
    :param tol: stopping tolerance
    :param delt: gradient step size
    :param display: display optimization iterations
"""
function mlogit_warp_grad(alpha, beta, timet, q, y; max_itr=8000,
                          tol=1e-10, delt=0.008, display=0)
    m1 = length(timet);
    m = size(beta,2);
    q /= norm(q);
    alpha /= norm(alpha);
    for i in 1:m
        beta[:, i] = beta[:, i] / norm(beta[:, i]);
    end
    gam1 = collect(linspace(0, 1, m1));
    gamout = zeros(m1);
    beta1 = reshape(beta, m1*m, 1);

    ccall((:mlogit_warp_grad, libfdasrsf), Void,
          (Ptr{Int32}, Ptr{Int32}, Ptr{Float64}, Ptr{Float64}, Ptr{Float64},
          Ptr{Float64}, Ptr{Float64}, Ptr{Int32}, Ptr{Int32}, Ptr{Float64},
          Ptr{Float64},  Ptr{Int32}, Ptr{Float64}),
          &m1, &m, alpha, beta1, timet, gam1, q, y, &max_itr, &tol, &delt,
          &display, gamout)

    return gamout
end

"""
Calculate warping for m-logistic elastic regression

    mlogit_optim(x, grad, Phi, Y)
    :param x: sample
    :param grad: gradient
    :param Phi: matrix
    :param Y: response matrix
"""
function mlogit_optm(x::Vector, grad::Vector, Phi, Y)
    if length(grad) > 0
        mlogit_gradient!(x, grad, Phi, Y);
    end
    out = mlogit_loss(x, Phi, Y);
    return out
end

"""
Calculate loss for m-logistic regression

    mlogit_loss(b, X, Y)
    :param b: coefficients
    :param X: matrix
    :param Y: response matrix
"""
function mlogit_loss(b, X, Y)
    N, m = size(Y);  # n_samples, n_classes
    M = size(X,2);   # n_features
    B = reshape(b, M, m);
    Yhat = X * B;
    Yhat -= repmat(minimum(Yhat, 2),1,size(Yhat,2));
    Yhat = exp(-1.*Yhat);
    # l1-normalize
    Yhat = Yhat./repmat(sum(Yhat, 2),1,size(Yhat,2));

    Yhat = Yhat .* Y;
    nll = sum(log(sum(Yhat,2)));
    nll /= -N;

    return nll
end

"""
Calculate m-logistic elastic regression loss function gradient in place

    mlogit_gradient!(b, grad, X, Y)
    :param b: coefficients
    :param grad: gradient
    :param X: matrix
    :param Y: response matrix
"""
function mlogit_gradient!(b, grad, X, Y)
    N, m = size(Y);  # n_samples, n_classes
    M = size(X,2); # n_features
    B = reshape(b, M, m);
    Yhat = X * B;
    Yhat -= repmat(minimum(Yhat, 2),1,size(Yhat,2));
    Yhat = exp(-1.*Yhat);
    # l1-normalize
    Yhat = Yhat./repmat(sum(Yhat, 2),1,size(Yhat,2));

    _Yhat = Yhat .* Y;
    _Yhat = _Yhat./repmat(sum(_Yhat,2),1,size(_Yhat,2));
    Yhat -= _Yhat;
    grad1 = X.' * Yhat;
    grad1 /= -N;
    grad[:] = reshape(grad1, M*m, 1);
end


##
# Ultraspherical
##

struct UltrasphericalWeight{T,Λ} <: AbstractJacobiWeight{T}
    λ::Λ
end

UltrasphericalWeight{T}(λ) where T = UltrasphericalWeight{T,typeof(λ)}(λ)
UltrasphericalWeight(λ) = UltrasphericalWeight{float(typeof(λ)),typeof(λ)}(λ)

==(a::UltrasphericalWeight, b::UltrasphericalWeight) = a.λ == b.λ

function getindex(w::UltrasphericalWeight, x::Number)
    x ∈ axes(w,1) || throw(BoundsError())
    (1-x^2)^(w.λ-one(w.λ)/2)
end

sum(w::UltrasphericalWeight{T}) where T = sqrt(convert(T,π))*exp(loggamma(one(T)/2 + w.λ)-loggamma(1+w.λ))

struct Ultraspherical{T,Λ} <: AbstractJacobi{T}
    λ::Λ
end
Ultraspherical{T}(λ::Λ) where {T,Λ} = Ultraspherical{T,Λ}(λ)
Ultraspherical(λ::Λ) where Λ = Ultraspherical{Float64,Λ}(λ)
Ultraspherical(P::Legendre{T}) where T = Ultraspherical(one(T)/2)
function Ultraspherical(P::Jacobi{T}) where T
    P.a == P.b || throw(ArgumentError("$P is not ultraspherical"))
    Ultraspherical(P.a+one(T)/2)
end

Ultraspherical(::ChebyshevU{T}) where T = Ultraspherical{T}(1)

const WeightedUltraspherical{T} = WeightedBasis{T,<:UltrasphericalWeight,<:Ultraspherical}

WeightedUltraspherical(λ) = UltrasphericalWeight(λ) .* Ultraspherical(λ)
WeightedUltraspherical{T}(λ) where T = UltrasphericalWeight{T}(λ) .* Ultraspherical{T}(λ)
orthogonalityweight(C::Ultraspherical) = UltrasphericalWeight(C.λ)

ultrasphericalc(n::Integer, λ, z::Number) = Base.unsafe_getindex(Ultraspherical{promote_type(typeof(λ),typeof(z))}(λ), z, n+1)
ultraspherical(λ, d::AbstractInterval{T}) where T = Ultraspherical{float(promote_type(eltype(λ),T))}(λ)[affine(d,ChebyshevInterval{T}()), :]

==(a::Ultraspherical, b::Ultraspherical) = a.λ == b.λ
==(::Ultraspherical, ::ChebyshevT) = false
==(::ChebyshevT, ::Ultraspherical) = false
==(C::Ultraspherical, ::ChebyshevU) = isone(C.λ)
==(::ChebyshevU, C::Ultraspherical) = isone(C.λ)
==(P::Ultraspherical, Q::Jacobi) = isone(2P.λ) && Jacobi(P) == Q
==(P::Jacobi, Q::Ultraspherical) = isone(2Q.λ) && P == Jacobi(Q)
==(P::Ultraspherical, Q::Legendre) = isone(2P.λ)
==(P::Legendre, Q::Ultraspherical) = isone(2Q.λ)

###
# interrelationships
###

Jacobi(C::Ultraspherical{T}) where T = Jacobi(C.λ-one(T)/2,C.λ-one(T)/2)

########
# Jacobi Matrix
########

function jacobimatrix(P::Ultraspherical{T}) where T
    λ = P.λ
    Tridiagonal((one(T):∞) ./ (2 .*((zero(T):∞) .+ λ)),
                Zeros{T}(∞),
                ((2λ):∞) ./ (2 .*((one(T):∞) .+ λ)))
end

# These return vectors A[k], B[k], C[k] are from DLMF. Cause of MikaelSlevinsky we need an extra entry in C ... for now.
function recurrencecoefficients(C::Ultraspherical{T}) where T
    λ = convert(T,C.λ)
    n = 0:∞
    (2(n .+ λ) ./ (n .+ 1), Zeros{typeof(λ)}(∞), (n .+ (2λ-1)) ./ (n .+ 1))
end

##########
# Derivatives
##########

# Ultraspherical(1)\(D*Chebyshev())
@simplify *(D::Derivative{<:Any,<:ChebyshevInterval}, S::ChebyshevU) = D * Ultraspherical(S)

# Ultraspherical(1/2)\(D*Legendre())
@simplify function *(D::Derivative{<:Any,<:ChebyshevInterval}, S::Legendre)
    T = promote_type(eltype(D),eltype(S))
    A = _BandedMatrix(Ones{T}(1,∞), ℵ₀, -1,1)
    ApplyQuasiMatrix(*, Ultraspherical{T}(convert(T,3)/2), A)
end

# Ultraspherical(λ+1)\(D*Ultraspherical(λ))
@simplify function *(D::Derivative{<:Any,<:ChebyshevInterval}, S::Ultraspherical)
    T = promote_type(eltype(D),eltype(S))
    A = _BandedMatrix(Fill(2convert(T,S.λ),1,∞), ℵ₀, -1,1)
    ApplyQuasiMatrix(*, Ultraspherical{T}(S.λ+1), A)
end

# Ultraspherical(λ-1)\ (D*wUltraspherical(λ))
@simplify function *(D::Derivative{<:Any,<:AbstractInterval}, WS::Weighted{<:Any,<:Ultraspherical})
    S = WS.P
    λ = S.λ
    T = promote_type(eltype(D),eltype(WS))
    if λ == 1
        A = _BandedMatrix((-(one(T):∞))', ℵ₀, 1,-1)
        ApplyQuasiMatrix(*, Weighted(ChebyshevT{T}()), A)
    else
        n = (0:∞)
        A = _BandedMatrix((-one(T)/(2*(λ-1)) * ((n.+1) .* (n .+ (2λ-1))))', ℵ₀, 1,-1)
        ApplyQuasiMatrix(*, Weighted(Ultraspherical{T}(λ-1)), A)
    end
end

# Ultraspherical(λ-1)\ (D*w*Ultraspherical(λ))
@simplify function *(D::Derivative{<:Any,<:AbstractInterval}, WS::WeightedUltraspherical)
    w,S = WS.args
    λ = S.λ
    T = eltype(WS)
    if iszero(w.λ)
        D*S
    elseif isorthogonalityweighted(WS) # weights match
        D * Weighted(S)
    else
        error("Not implemented")
    end
end

##########
# Conversion
##########

\(A::Ultraspherical, B::Legendre) = A\Ultraspherical(B)
\(A::Legendre, B::Ultraspherical) = Ultraspherical(A)\B

function \(A::Ultraspherical, B::Jacobi)
    Ã = Jacobi(A)
    Diagonal(Ã[1,:]./A[1,:]) * (Ã\B)
end
function \(A::Jacobi, B::Ultraspherical)
    B̃ = Jacobi(B)
    (A\B̃)*Diagonal(B[1,:]./B̃[1,:])
end

function \(U::Ultraspherical{<:Any,<:Integer}, C::ChebyshevT)
    T = promote_type(eltype(U), eltype(C))
    (U\Ultraspherical{T}(1)) * (ChebyshevU{T}()\C)
end

function \(U::Ultraspherical{<:Any,<:Integer}, C::ChebyshevU)
    T = promote_type(eltype(U), eltype(C))
    U\Ultraspherical(C)
end

\(T::Chebyshev, C::Ultraspherical) = inv(C \ T)

function \(C2::Ultraspherical{<:Any,<:Integer}, C1::Ultraspherical{<:Any,<:Integer})
    T = promote_type(eltype(C2), eltype(C1))
    λ_Int = C1.λ
    λ = convert(T,λ_Int)
    if C2.λ == λ_Int+1
        _BandedMatrix( Vcat(-(λ ./ ((0:∞) .+ λ))', Zeros{T}(1,∞), (λ ./ ((0:∞) .+ λ))'), ℵ₀, 0, 2)
    elseif C2.λ == λ_Int
        Eye{T}(∞)
    elseif C2.λ > λ_Int
        (C2 \ Ultraspherical(λ_Int+1)) * (Ultraspherical(λ_Int+1)\C1)
    else
        error("Not implemented")
    end
end

function \(C2::Ultraspherical, C1::Ultraspherical)
    T = promote_type(eltype(C2), eltype(C1))
    λ_Int = C1.λ
    λ = convert(T,λ_Int)
    if C2.λ == λ+1
        _BandedMatrix( Vcat(-(λ ./ ((0:∞) .+ λ))', Zeros{T}(1,∞), (λ ./ ((0:∞) .+ λ))'), ℵ₀, 0, 2)
    elseif C2.λ == λ_Int
        Eye{T}(∞)
    elseif isinteger(C2.λ-λ_Int) && C2.λ > λ_Int
        Cm = Ultraspherical{T}(λ_Int+1)
        (C2 \ Cm) * (Cm \ C1)
    elseif isinteger(C2.λ-λ_Int)
        inv(C1 \ C2)
    else
        error("Not implemented")
    end
end

function \(w_A::WeightedUltraspherical, w_B::WeightedUltraspherical)
    wA,A = w_A.args
    wB,B = w_B.args
    T = promote_type(eltype(w_A),eltype(w_B))

    if wA == wB
        A \ B
    elseif B.λ == A.λ+1 && wB.λ == wA.λ+1 # Lower
        λ = convert(T,A.λ)
        _BandedMatrix(Vcat(((2λ:∞) .* ((2λ+1):∞) ./ (4λ .* (λ+1:∞)))',
                            Zeros{T}(1,∞),
                            (-(1:∞) .* (2:∞) ./ (4λ .* (λ+1:∞)))'), ℵ₀, 2,0)
    else
        error("not implemented for $A and $wB")
    end
end

\(A::Legendre, wB::WeightedUltraspherical) = Ultraspherical(A) \ wB

function \(A::Ultraspherical, w_B::WeightedUltraspherical)
    (UltrasphericalWeight(zero(A.λ)) .* A) \ w_B
end


module SparsityOperators

import Base: length, iterate, \
using LinearAlgebra
import LinearAlgebra.BLAS: gemv, gemv!
import LinearAlgebra: BlasFloat, normalize!, norm, rmul!, lmul!
using SparseArrays
using Random
using CUDA

using Reexport
@reexport using Reexport
@reexport using LinearOperators
@reexport using FFTW
@reexport using Wavelets

const Trafo = Union{AbstractMatrix, AbstractLinearOperator, Nothing}
const FuncOrNothing = Union{Function, Nothing}

include("FFTOp.jl")
include("DCTOp.jl")
include("DSTOp.jl")
include("WaveletOp.jl")
include("GradientOp.jl")
include("SamplingOp.jl")
include("WeightingOp.jl")
include("NormalOp.jl")

export linearOperator, linearOperatorList

linearOperator(op::Nothing,shape,T::Type=ComplexF32) = nothing

"""
  returns a list of currently implemented `LinearOperator`s
"""
function linearOperatorList()
  return ["DCT-II", "DCT-IV", "FFT", "DST", "Wavelet", "Gradient"]
end

"""
    linearOperator(op::AbstractString, shape)

returns the `LinearOperator` with name `op`.

# valid names
* `"FFT"`
* `"DCT-II"`
* `"DCT-IV"`
* `"DST"`
* `"Wavelet"`
* `"Gradient"`
"""
function linearOperator(op::AbstractString, shape, T::Type=ComplexF32)
  shape_ = tuple(shape...)
  if op == "FFT"
    trafo = FFTOp(T, shape_, false) #FFTOperator(shape)
  elseif op == "DCT-II"
    shape_ = tuple(shape[shape .!= 1]...)
    trafo = DCTOp(T, shape_, 2)
  elseif op == "DCT-IV"
    shape_ = tuple(shape[shape .!= 1]...)
    trafo = DCTOp(T, shape_, 4)
  elseif op == "DST"
    trafo = DSTOp(T, shape_)
  elseif op == "Wavelet"
    trafo = WaveletOp(T,shape_)
  elseif op=="Gradient"
    trafo = GradientOp(T,shape_)
  else
    error("Unknown transformation")
  end
  trafo
end

end # module


module NNlib

using Requires

export σ, sigmoid, relu, leakyrelu, elu, swish, selu, softplus, softsign, logσ, logsigmoid,
  softmax, logsoftmax, maxpool, meanpool

include("numeric.jl")
include("activation.jl")
include("softmax.jl")
include("logsoftmax.jl")
include("linalg.jl")
include("conv.jl")
include("cubroadcast.jl")

end # module


"""
    R = JenksClassification(n,X)

Jenks Natural Breaks Optimzation algorithm. Default parameters

    JenksClassification(n::Int,                             # number of classes
                        X::Array{T,1};                      # data
                        errornorm::Int=1,                   # L1, or L2 errornorm
                        maxiter::Int=200,                   # number of iterations
                        flux::Real=0.1,                     # flux [fraction] of values from one class to another
                        fluxadjust::Real=1.03,              # change in flux depending on previous flux direction
                        fluxadjust_bothways::Bool=true,     # lower flux or also increase flux (=bothways)
                        inimode::String="maxentropy",       # initialisation method
                        feedback::Bool=true                 # feedback on progress and convergence

# Examples

```jldoc
julia> R = JenksClassification(5,rand(10000));
100%: ARE=0.05014468, finished in 0.01s.

julia> R.breaks
6-element Array{Int64,1}:
     1
  1972
  3993
  6041
  8002
 10001
```
"""
function JenksClassification(n::Int,X::Array{T,1};
                            errornorm::Int=1,
                            maxiter::Int=200,
                            flux::Real=0.1,
                            fluxadjust::Real=1.03,
                            fluxadjust_bothways::Bool=true,
                            inimode::String="maxentropy",
                            feedback::Bool=true) where {T<:AbstractFloat}

    sort!(X)
    ndata = length(X)

    # initialise a JenksResult struct
    JR = JenksResult(n,X,errornorm=errornorm,maxiter=maxiter,flux=flux)
    JR.ndata = ndata

    # Initialisation
    breaks = JR.breaks                      # breaks is a view on the data in JR.breaks
    breaks[:] = BreaksInit(n,X,inimode)     # lower bounds of every class as index of X
    SDAM = SqDeviation(X)                   # squared deviation from array mean (this is constant)
    DEVs = Array{Float64,1}(undef,n)        # Deviations from class centre for each class (lin/sq)
    DEVs_pre = Array{Float64,1}(undef,n)    # Deviations from previous iteration

    # Rename deviation function (linear vs squared)
    if errornorm == 1
        Deviation = LinDeviation
    elseif errornorm == 2
        Deviation = SqDeviation
    else
        throw(error("Errornorm '$errornorm' not defined."))
    end

    if errornorm == 1
        JR.AREhistory[1] = ARE(X,breaks)
    else
        JR.GVFhistory[1] = GVF(X,breaks)
    end

    # automatic flux adjustment
    fluxes = JR.fluxes

    t0 = time()
    for iter in 1:maxiter

        DEVs_pre .= DEVs

        # Get the lin/squared deviations for all classes
        for iclass in 1:n
            DEVs[iclass] = Deviation(ClassValues(X,breaks,iclass))
        end

        # shift class bounds
        for iclass in 1:n-1

            a,b,c = breaks[iclass:iclass+2]         # for class sizes

            # compare deviations between two adjacent classes
            if DEVs[iclass] < DEVs[iclass+1]

                # adjust fluxes based on history
                if iter > 1
                    if DEVs_pre[iclass] >= DEVs_pre[iclass+1]    # flux was previously in opposite direction
                        fluxes[iclass] /= fluxadjust             # decrease flux
                    elseif fluxadjust_bothways                          # flux was previously in same direction
                        fluxes[iclass] *= fluxadjust             # increase flux
                    end
                end

                # class size factor: Bigger classes give away more data points
                # shift data points to the class with smaller DEV
                cs_factor = Int(round((b-a)*fluxes[iclass]))
                newbreak = breaks[iclass+1]+cs_factor
                breaks[iclass+1] = min(newbreak,breaks[iclass+2]-2)
            else

                # adjust fluxes based on history
                if iter > 1
                    if DEVs_pre[iclass] <= DEVs_pre[iclass+1]    # flux was previously in opposite direction
                        fluxes[iclass] /= fluxadjust
                    elseif fluxadjust_bothways
                        fluxes[iclass] *= fluxadjust
                    end
                end

                cs_factor = Int(round((c-b)*fluxes[iclass]))
                newbreak = breaks[iclass+1]-cs_factor
                breaks[iclass+1] = max(newbreak,breaks[iclass]+2)
            end
        end

        # calculate global errors and store them
        if errornorm == 1
            AREnow = sum(DEVs)/ndata
            JR.AREhistory[iter+1] = AREnow
        else
            GVFnow = 1-sum(DEVs)/SDAM
            JR.GVFhistory[iter+1] = GVFnow
        end

        # provide average rounding error / goodness of variance feedback
        if feedback
            percent = Int(round(iter/maxiter*100))

            if errornorm == 1
                AREstring = @sprintf "%.8f" AREnow
                print("\r\u1b[K")
                print("$percent%: ARE=$AREstring")
            else
                GVFstring = @sprintf "%.8f" GVFnow
                print("\r\u1b[K")
                print("$percent%: GVF=$GVFstring")
            end
        end
    end
    JR.dt = time() - t0
    feedback ? println(@sprintf ", finished in %.2fs." JR.dt) : nothing

    # Store both errornorms in JR struct
    if errornorm == 1
        JR.ARE = sum(DEVs)/ndata
        JR.GVF = GVF(X,breaks)
    else
        JR.GVF = 1-sum(DEVs)/SDAM
        JR.ARE = ARE(X,breaks)
    end

    # convert breaks to centres and class sizes
    Breaks2Centres!(JR,X)
    Breaks2ClassSize!(JR)
    Breaks2Bounds!(JR,X)

    return JR
end

"""Returns the data point values for a given class i, provided the data X and
the breaks. Assumes the data array X to be sorted."""
function ClassValues(X::Array{T,1},breaks::Array{Int,1},i::Int) where {T<:AbstractFloat}
    return X[breaks[i]:breaks[i+1]-1]
end


@cache struct RichardsonEulerCache{uType,rateType,tabType, arrayType} <: OrdinaryDiffEqMutableCache
  u::uType
  uprev::uType
  tmp::uType
  k::rateType
  fsalfirst::rateType
  tab::tabType
  T::arrayType
end

struct RichardsonEulerConstantCache <: OrdinaryDiffEqConstantCache
  m::Int #Generalise type of m
  # Can add different order of n_{i} here
  function RichardsonEulerConstantCache()
    m = 4
    new(m)
  end
end

function alg_cache(alg::RichardsonEuler,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{true}})
  tmp = similar(u)
  k = zero(rate_prototype)
  fsalfirst = zero(rate_prototype)
  tab = RichardsonEulerConstantCache()
  T = fill(zeros(eltype(u), size(u)), (tab.m, tab.m))
  RichardsonEulerCache(u,uprev,tmp,k,fsalfirst,tab,T)
end

function alg_cache(alg::RichardsonEuler,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{false}})
  RichardsonEulerConstantCache()
end


using SparseArrays, MatrixNetworks, Graphs, SimpleWeightedGraphs
using SimpleValueGraphs
# using IndexedGraphs
include("../src/IndexedGraphs.jl")
using .IndexedGraphs
using Random
using LinearAlgebra

# build directed graph
N = 10_000
W = sprand(N, N, 5/N)
# remove self loops
for i in 1:size(W,2); W[i,i] = 0; end
dropzeros!(W)
W = W + W'
w = sparse(UpperTriangular(W)).nzval    # unique wights for Graph
wD = nonzeros(permutedims(W))           # duplicate weights for DiGraph

# pick source at random
rng = MersenneTwister(0)
s = rand(rng, 1:size(W,2))

g = IndexedGraph(W)
gD = IndexedDiGraph(W)
g_MN = MatrixNetwork(W)
g_Graphs = SimpleGraph(W)
g_SWG = SimpleWeightedGraph(W)
g_SVG = ValGraph(g_Graphs,
    edgeval_types=(Float64,), edgeval_init=(s, d) -> (W[s,d],))

ds_IG = dijkstra_shortest_paths(g, s, w)
d_IG = ds_IG.dists; p_IG  = ds_IG.parents

ds_IDG = dijkstra_shortest_paths(gD, s, wD)
d_IDG = ds_IDG.dists; p_IDG  = ds_IDG.parents

d_MN, p_MN = MatrixNetworks.dijkstra(g_MN, s)

ds_Graphs = dijkstra_shortest_paths(g_Graphs, s, W)
d_Graphs = ds_Graphs.dists; p_Graphs = ds_Graphs.parents

ds_SWG = dijkstra_shortest_paths(g_SWG, s)
d_SWG = ds_SWG.dists; p_SWG = ds_SWG.parents

ds_SVG = SimpleValueGraphs.Experimental.dijkstra_shortest_paths(g_SVG, s)
d_SVG = ds_SVG.dists; p_SVG = ds_SVG.parents

# check that results are correct
@assert d_IG == d_MN == d_Graphs == d_SWG == d_SVG == d_IDG
@assert p_IG == p_MN == p_Graphs == p_SWG == p_SVG == p_IDG

### BENCHMARK
using BenchmarkTools
println("IndexedDiGraph:")
@btime dijkstra_shortest_paths($gD, $s, $wD)
println("IndexedGraph:")
@btime dijkstra_shortest_paths($g, $s, $w)
println("MatrixNetwork:")
@btime MatrixNetworks.dijkstra($g_MN, $s)
println("SimpleGraph")
@btime dijkstra_shortest_paths($g_Graphs, $s, $W)
println("SimpleWeightedGraph:")
@btime dijkstra_shortest_paths($g_SWG, $s);
println("ValGraph (SimpleValueGraphs.Experimental):")
@btime SimpleValueGraphs.Experimental.dijkstra_shortest_paths($g_SVG, $s);


# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/Materials.jl/blob/master/LICENSE

using FEMMaterials, Tensors, Test

analysis, problem, element, bc_elements, ip = get_one_element_material_analysis(:Chaboche)
update!(element, "E", 200.0e3)
update!(element, "nu", 0.3)
update!(element, "R0", 100.0)
update!(element, "Kn", 10.0)
update!(element, "nn", 20.0)
update!(element, "C1", 0.0)
update!(element, "D1", 100.0)
update!(element, "C2", 0.0)
update!(element, "D2", 1000.0)
update!(element, "Q", 0.0)
update!(element, "b", 0.1)

times = [0.0]
loads = [0.0]
dt = 0.5
E = 200.0e3
nu = 0.3
G = 0.5*E/(1+nu)
syield = 100.0
# vm = sqrt(3)*G*ga | ea = ga
ea = 2*syield/(sqrt(3)*G)
# Go to elastic border
push!(times, times[end]+dt)
push!(loads, loads[end] + ea*dt)
 # Proceed to plastic flow
push!(times, times[end]+dt)
push!(loads, loads[end] + ea*dt)
 # Reverse direction
push!(times, times[end]+dt)
push!(loads, loads[end] - ea*dt)
 # Continue and pass yield criterion
push!(times, times[end]+dt)
push!(loads, loads[end] - 2*ea*dt)
loading = ShearStrainLoading(times, loads)
update_bc_elements!(bc_elements, loading)

analysis.properties.t1 = maximum(times)
run!(analysis)
stresses = zeros(length(times), 6)
for (i,t) in zip(1:length(times), times)
    stresses[i,:] = tovoigt(ip("stress", t))
    @test isapprox(stresses[i,[1,2,3,4,6]], zeros(5); atol=1e-6)
end
s31 = stresses[:,5]
eeqs = [ip("cumeq", t) for t in times]

@test isapprox(s31[2], syield/sqrt(3.0))
@test isapprox(s31[3]*sqrt(3.0), syield + 10.0*((eeqs[3]-eeqs[2])/dt)^(1.0/20.0); rtol=1e-2)
@test isapprox(s31[4], s31[3]-G*ea*dt)
@test isapprox(s31[5]*sqrt(3.0), -(syield + 10.0*((eeqs[5]-eeqs[4])/dt)^(1.0/20.0)); rtol=1e-2)


using OpenQuantumTools

# Define a function to construct an eigendecomposition routine for the
# Hamiltonian. The routine should have the signature: (H, t, lvl) -> (w, v).
# The argument of this function is the cache used by the Hamiltonian object.
function build_user_eigen(u_cache)
    EIGS = function(H, t, lvl)
        println("I am the user defined eigendecomposition routine.")
        w, v = eigen(Hermitian(H(t)))
        w[1:lvl], v[:, 1:lvl]
    end
end

H = DenseHamiltonian([(s)->1-s, (s)->s], [σx, σz], EIGS=build_user_eigen)

eigen_decomp(H, 0.5, lvl=2)

function build_user_eigen(u_cache)
    # note that to keep the unit consistent, the unit of any value inside the routine should be 1/h
    w, v = eigen(Hermitian(2*π*(σx+σz)))
    EIGS = function(H, t, lvl)
        w[1:lvl], v[:, 1:lvl]
    end
end

H = DenseHamiltonian([(s)->1.0], [σx+σz], EIGS=build_user_eigen)

print(eigen_decomp(H, 0.5, lvl=2))
print(eigen_decomp(H, 0.0, lvl=2))

Hd = standard_driver(4, sp=true);
Hp = two_local_term(rand(3), [[1,2],[2,3],[3,4]], 4, sp=true)
H = SparseHamiltonian([(s)->1-s, (s)->s], [Hd, Hp], unit=:ħ)

# the default eigen_decomposition using the dense matrices algorithm
w, v = eigen_decomp(H, 0.1, lvl=4)

using Arpack

function build_user_eigen(u_cache)
    function (H, t, lvl)
        hmat = H(t)
        println("Using sparse matrix algorithm")
        # define all the Arpack routine parameters here
        eigs(hmat, nev = lvl, which=:SR, tol=0.0, maxiter=300)
    end
end

Hd = standard_driver(4, sp=true);
Hp = two_local_term(rand(3), [[1,2],[2,3],[3,4]], 4, sp=true)
H = SparseHamiltonian([(s)->1-s, (s)->s], [Hd, Hp], unit=:ħ, EIGS =build_user_eigen)

eigen_decomp(H, 0.1, lvl=4)


# using Distributions

immutable GaussianMixtureModelCommonCovarDiagonal <: Distribution
    mu::Array{Float64,2}
    probs::Vector{Float64}
    covar::Vector{Float64}
    aliastable::Distributions.AliasTable
    function GaussianMixtureModelCommonCovarDiagonal(m::Array{Float64,2}, p::Vector{Float64}, ic::Vector{Float64})
        if size(m,2) != length(p)
            error("means and probs must have the same number of elements")
        end
if( size(m,1) != size(ic,1) )
    error("means and covar matrix not compatible sizes: ",size(m)," vs ",size(ic) )
end
        sump = 0.0
        for i in 1:length(p)
            if p[i] < 0.0
                error("MixtureModel: probabilities must be non-negative")
            end
            sump += p[i]
        end
        table = Distributions.AliasTable(p)
        new(m, p ./ sump, ic, table)
    end
end

function mean(d::GaussianMixtureModelCommonCovarDiagonal)
    np = size(d.mu,2)
    m = zeros(np)
    for i in 1:length(d.probs)
        m += vec(d.mu[:,i]) .* d.probs[i]
    end
    return m
end

function pdf(d::GaussianMixtureModelCommonCovarDiagonal, x::Any)
    p = 0.0
    for i in 1:length(d.probs)
        p += Distributions.pdf(Distributions.MvNormal(sqrt(d.covar)), x .- vec(d.mu[:,i]) ) * d.probs[i]
    end
    return p
end

function rand(d::GaussianMixtureModelCommonCovarDiagonal)
    i = Distributions.rand(d.aliastable)
    return  Distributions.rand(Distributions.MvNormal(vec(d.mu[:,i]),sqrt(d.covar)))
end


using BasisFunctions: hasderivative, hasantiderivative, resize

P = 80

@testset "$(rpad("Bernstein polynomials",P))" begin

    ξ = [0.0, 1.0]
    b = Bernstein(ξ)

    BasisFunctions.Test.test_generic_dict_interface(Bernstein(10))

    @testset "$(rpad("Basic functionality",P))" begin
        @test nodes(b) == ξ
        @test nnodes(b) == 2
        @test degree(b) == 1
        @test support(b) == BernsteinInterval{eltype(ξ)}()
        @test nodes(b) == nodes(similar(b, eltype(ξ), 2))
        @test nodes(b) == nodes(similar(b, eltype(ξ), ξ))
        @test nodes(resize(b, 2)) == ξ
        @test nodes(resize(b, 5)) == [0.0, 0.25, 0.50, 0.75, 1.0]
    end

    @testset "$(rpad("Basis evaluation",P))" begin
        @test b[1](0.0) == 1.0
        @test b[1](0.5) == 0.5
        @test b[1](1.0) == 0.0

        @test b[2](0.0) == 0.0
        @test b[2](0.5) == 0.5
        @test b[2](1.0) == 1.0
    end

    @testset "$(rpad("Basis derivative",P))" begin
        @test eval_element_derivative(b, 1, 0.0) == -1.0
        @test eval_element_derivative(b, 1, 0.5) == -1.0
        @test eval_element_derivative(b, 1, 1.0) == -1.0

        @test eval_element_derivative(b, 2, 0.0) == +1.0
        @test eval_element_derivative(b, 2, 0.5) == +1.0
        @test eval_element_derivative(b, 2, 1.0) == +1.0
    end

end


# This file is a part of EuclidianNormalizingFlows.jl, licensed under the MIT License (MIT).

std_normal_logpdf(x::Real) = -(abs2(x) + log2π)/2

function mvnormal_negll_trafo(trafo::Function, X::AbstractMatrix{<:Real})
    nsamples = size(X, 2) # normalize by number of samples to be independent of batch size:
    Y, ladj = with_logabsdet_jacobian(trafo, X)
    #ref_ll = sum(sum(std_normal_logpdf.(Y), dims = 1) .+ ladj) / nsamples
    # Faster:
    ll = (sum(std_normal_logpdf.(Y)) + sum(ladj)) / nsamples
    #@assert ref_ll ≈ ll
    return -ll
end

function mvnormal_negll_trafograd(trafo::Function, X::AbstractMatrix{<:Real})
    negll, back = Zygote.pullback(mvnormal_negll_trafo, trafo, X)
    d_trafo = back(one(eltype(X)))[1]
    return negll, d_trafo
end

function optimize_whitening(
    smpls::VectorOfSimilarVectors{<:Real}, initial_trafo::Function, optimizer;
    nbatches::Integer = 100, nepochs::Integer = 100,
    optstate = Optimisers.state(optimizer, deepcopy(initial_trafo)),
    negll_history = Vector{Float64}()
)
    batchsize = round(Int, length(smpls) / nbatches)
    batches = collect(Iterators.partition(smpls, batchsize))
    trafo = deepcopy(initial_trafo)
    state = deepcopy(optstate)
    negll_hist = Vector{Float64}()
    for i in 1:nepochs
        for batch in batches
            X = flatview(batch)
            negll, d_trafo = mvnormal_negll_trafograd(trafo, X)
            state, trafo = Optimisers.update(optimizer, state, trafo, d_trafo)
            push!(negll_hist, negll)
        end
    end
    (result = trafo, optimizer_state = state, negll_history = vcat(negll_history, negll_hist))
end
export optimize_trafo


export misfit, reducedMisfit
"""
function misfit(param,M,w;doDerivative,pwr)

Compute the following quantities

misfit:             1/pwr * || R*(M*w - d) ||_pwr^pwr
residual:           M*w - d
weighted residual:  R*(M*w - d)

and optionally the gradient and Hessian of the misfit with respect to w.

Here, R is the inverse square root of the covariance matrix Gamma.

If M is a 3D array, all values are computed separately for each binary matrix
stored in M.

Inputs:

	1. param  - StrainReconParam
	2. M      - binary matrix or a 3D binary array
	3. w      - vector of weights

Keyword arguments:

    4. doDerivative   - flag for computing derivatives (default = false)
    5. pwr            - exponent in the likelihood function (default = param.pwr)

Output:
	1. Jc             - misfit
	2. res            - residual
	3. wres           - weighted residual
	4. dJ             - gradient of misfit w.r.t. w (if doDerivative = true)
	5. H              - Hessian of misfit w.r.t. w (if doDerivative = true)
"""
function misfit(param::StrainReconParam,M::BitArray{2},w::Vector{Float64}; doDerivative::Bool=false,pwr::Float64=param.pwr)
	d    = param.d
	R    = param.R

	res  = M*w-d
	wres = R*res
	Jc   = 1/pwr * sum(abs.(wres).^pwr)

	if doDerivative
		if pwr != 2.0
		    error("derivative for pwr!=2 not implemented");
		end
		dJ = M'*(R'*wres)
		H  = M'*((R'*R)*M)
		return Jc, res, wres, dJ, H
	else
		return Jc, res, wres
	end
end

function misfit(param::StrainReconParam,M::BitArray{3},w::Vector{Float64}; doDerivative::Bool=false,pwr::Float64=param.pwr)
	if doDerivative
		error("derivative for multiple binary matrices not implemented")
	end

	d    = param.d
	n    = param.n
	nM   = size(M,3)
	m    = length(d)

	# allocate space
	Jc   = zeros(nM)
	res  = zeros(m,nM)
	wres = zeros(m,nM)

	# compute misfit and residuals for each M
	for k=1:nM
		Jc[k],res[:,k],wres[:,k] = misfit(param,M[:,:,k],w,pwr=pwr)
	end
	return Jc, res, wres
end

"""
function reducedMisfit(param,w;nsol,doDerivative,pwr)

Compute reduced misfit

min_M 1/pwr * || R*(M*w - d) ||_pwr^pwr

where the minimum is taken over all binary matrices M in Omega_M.

If nsol>1, then nsol matrices, corresponding to nsol smallest misfit values, are
considered and the output is similar to misfit with 3D input array.

Input:

	1. param  - StrainReconParam
	2. w      - vector of weights

Keyword arguments:

	3. nsol           - consider nsol binary matrices (default = 1)
    4. doDerivative   - flag for computing derivatives (default = false)
	5. p              - exponent in the likelihood function (default = param.pwr)

Output:
	1. Jc             - misfit
	2. res            - residual
	3. wres           - weighted residual
	4. M              - the minimizing matrix (or a 3D array if nsol>1)
	5. dJ             - gradient of misfit w.r.t. w (if doDerivative = true)
	6. H              - Hessian of misfit w.r.t. w (if doDerivative = true)

"""
function reducedMisfit(param::StrainReconParam,w::Vector{Float64};nsol::Int=1,doDerivative::Bool=false,pwr::Float64=param.pwr)
	M = findM(param,w,nsol=nsol,pwr=pwr)
	if !doDerivative
		Jc, res, wres = misfit(param,M,w,pwr=pwr)
		return Jc, res, wres, M
	else
		Jc, res, wres, dJ, H = misfit(param,M,w,doDerivative=doDerivative,pwr=pwr)
		return Jc, res, wres, M, dJ, H
	end
end


###################################################################
# Positive Poisson GLM
###################################################################

###################################################################
# Distributions.jl extensions
###################################################################
"""
    PositivePoisson(λ)

A *PositivePoisson distribution* (aka zero-truncated Poisson, ZTP) descibes the number of
independent events occurring within a unit time interval, given the average rate of occurrence `λ`
and, importantly, given that the number is not zero.

``P(X = k) = \\frac{λ^k}{k!(1-e^{-λ})} e^{-λ}, \\quad \\text{ for } k = 1,2,\\ldots.``

```julia
PositivePoisson()        # PositivePoisson distribution with rate parameter 1
PositivePoisson(lambda)       # PositivePoisson distribution with rate parameter lambda

params(d)        # Get the parameters, i.e. (λ,)
mean(d)          # Get the mean arrival rate, i.e. λ
```

External links:

* [PositivePoisson distribution on Wikipedia](https://en.wikipedia.org/wiki/Zero-truncated_Poisson_distribution)

"""
struct PositivePoisson <: DiscreteUnivariateDistribution
    λ::Float64

    PositivePoisson(λ::Real) = new(λ)
    PositivePoisson() = new(1.0)
end

import Base.minimum, Base.maximum
Distributions.@distr_support PositivePoisson 1 (d.λ == 0.0 ? 0 : Inf)

### Parameters
Distributions.params(d::PositivePoisson) = (d.λ,)
Distributions.rate(d::PositivePoisson) = d.λ
### Statistics
Distributions.mean(d::PositivePoisson) = d.λ / (1.0-exp(-d.λ))

function logfactorialapprox(n::Integer)
  x=n+1.0::Float64
  (x-0.5)*log(x) - x + 0.5*log(2π) + 1.0/(12.0*x)
end

function logfactorial(n::Integer)
  if n<2
    return 0.0
  end
  x=0.0::Float64
  for i=2:n
    x+=log(i)
  end
  x
end

"""
Evaluates log(exp(x)-1.0) even when exp(x) blows up and the answer is simply x
(solution is related to more general logsumexp trick)
"""
function logexpm1(x::T) where {T}
  x + log(one(T)-exp(-x))
end

# Note that the factorial will blow up with moderately large x
Distributions.pdf(d::PositivePoisson, x::Int) = d.λ^x / ((exp(d.λ)-1.0) * factorial(x))

# calculating the log factorial is too expensive, we use an approximation
# given here: http://www.johndcook.com/blog/2010/08/16/how-to-compute-log-factorial/
Distributions.logpdf(d::PositivePoisson, x::Int) = x*log(d.λ) - logexpm1(d.λ) - logfactorialapprox(x)

# for comparison we define the exact logpdf for the positive poisson
logpdf_exact(d::PositivePoisson, x::Integer) = x*log(d.λ) - log(exp(d.λ)-1.0) - logfactorial(x)

# similar calculation for the regular poisson
logpdf_approx(d::Poisson, x::Integer) = x*log(d.λ) - d.λ - logfactorialapprox(x)

###################################################################
# GLM.jl extensions
###################################################################
struct LogProductLogLink <: Link end

# denoting by λ=e^η the intensity of the untruncated poisson
LAMBERTW_USE_NAN=true
λfn(μ) = μ + lambertw(-exp(-μ)*μ)

GLM.linkfun(::LogProductLogLink, μ) = log(λfn(μ))

# exp(η) / (1.0 - exp(-exp(η))) = 1.0 / (exp(-η) - exp(-η)exp(-exp(η)))
#                               = 1.0 / (exp(-η) - exp(-exp(η)-η))
function GLM.linkinv(::LogProductLogLink, η::T) where {T}
  λ = exp(η)
  λ / (one(T) - exp(-λ))
end

function GLM.mueta(::LogProductLogLink, η::T) where {T}
  λ = exp(η)
  expmλ = exp(-λ)
  μ = λ / (one(T) - expmλ)
  μ * (one(T) - μ*expmλ)
end

function GLM.inverselink(::LogProductLogLink, η::T) where {T}
    λ = exp(η)
    expmλ = exp(-λ)
    μ = λ / (one(T) - expmλ)
    dμdη = μ * (one(T) - μ*expmλ)
    μ, dμdη, dμdη # third arg is the glmvar which in this case equals the derivative
end

GLM.canonicallink(::PositivePoisson) = LogProductLogLink()
# GLM.glmvar(::PositivePoisson, ::LogProductLogLink, μ, η) = μ * (1.0 - μ*exp(-exp(η)))
# GLM.glmvar(::PositivePoisson, μ) = μ * (1.0 - μ*exp(-exp(log(λfn(μ)))))
GLM.glmvar(::PositivePoisson, μ::T) where {T} = μ * (one(T) - μ*exp(-λfn(μ)))
GLM.mustart(::PositivePoisson, y, wt) = y + oftype(y, 0.1)

function GLM.devresid(::PositivePoisson, y::T, μ::T) where {T}
  if y>one(T)
    if μ>one(T)
      λμ = λfn(μ)
      λy = λfn(y)
      2 * (y*(log(λy)-log(λμ)) - logexpm1(λy) + logexpm1(λμ))
    else
      typemax(T) # +∞
    end
  else
    if μ>one(T)
      2 * (λfn(μ) - log(μ))
    else
      zero(T)
    end
  end
end

GLM.dispersion_parameter(::PositivePoisson) = false

function GLM.loglik_obs(::PositivePoisson, y, μ, wt, ϕ)
  wt*logpdf(PositivePoisson(λfn(μ)), y)
end

# function GLM.dispersion(m::GLM.AbstractGLM, sqr::Bool=false)
#     wrkwts = m.rr.wrkwts
#     wrkresid = m.rr.wrkresid
#
#     if isa(m.rr.d, @compat Union{Binomial, Poisson, PositivePoisson})
#         return one(eltype(wrkwts))
#     end
#
#     s = zero(eltype(wrkwts))
#     @inbounds @simd for i = eachindex(wrkwts,wrkresid)
#         s += wrkwts[i]*abs2(wrkresid[i])
#     end
#     s /= df_residual(m)
#     sqr ? s : sqrt(s)
# end


using Zygote, Test, LinearAlgebra
zygotefun2(a, b) = sigmoid_fast(a + b) + LoopVectorization.relu(a * b) * tanh(a - b)
zygotefun3(a, b, c) = log(tanh_fast(a * b) + 1 + exp(b - c))
@testset "Zygote" begin
  @show @__LINE__

  for T ∈ (Float32, Float64)
    x = randn(T, 217)
    y = randn(T, 217)
    z = randn(T, 217)
    w = randn(T, 217)
    # Test 1-arg and `tanh_fast` vs `tanh`
    gref1 = gradient(x -> dot(w, map(tanh, x)), x)
    glv_1 = gradient(x -> dot(w, vmap(tanh_fast, x)), x)

    @test only(gref1) ≈ only(glv_1)

    # Test 2 arguments
    # Test 3 arguments
    gref2 = gradient(xyz -> dot(w, map(zygotefun2, xyz...)), (x, y))
    g_lv2 = gradient(xyz -> dot(w, vmap(zygotefun2, xyz...)), (x, y))

    @test all(map(≈, only(gref2), only(g_lv2)))

    gref3 = gradient(xyz -> dot(w, map(zygotefun3, xyz...)), (x, y, z))
    g_lv3 = gradient(xyz -> dot(w, vmap(zygotefun3, xyz...)), (x, y, z))

    @test all(map(≈, only(gref3), only(g_lv3)))
  end
end


## extended methods for computing covariance and scatter matrix

# auxiliary functions

function _symmetrize!(a::DenseMatrix)
    m, n = size(a)
    m == n || error("a must be a square matrix.")
    for j = 1:n
        @inbounds for i = j+1:n
            vl = a[i,j]
            vr = a[j,i]
            a[i,j] = a[j,i] = middle(vl, vr)
        end
    end
    return a
end

function _scalevars(x::DenseMatrix, s::DenseVector, dims::Int)
    dims == 1 ? Diagonal(s) * x :
    dims == 2 ? x * Diagonal(s) :
    error("dims should be either 1 or 2.")
end

## scatter matrix

_unscaled_covzm(x::DenseMatrix, dims::Colon)   = unscaled_covzm(x)
_unscaled_covzm(x::DenseMatrix, dims::Integer) = unscaled_covzm(x, dims)

_unscaled_covzm(x::DenseMatrix, wv::AbstractWeights, dims::Colon)   =
    _symmetrize!(unscaled_covzm(x, _scalevars(x, values(wv))))
_unscaled_covzm(x::DenseMatrix, wv::AbstractWeights, dims::Integer) =
    _symmetrize!(unscaled_covzm(x, _scalevars(x, values(wv), dims), dims))

"""
    scattermat(X, [wv::AbstractWeights]; mean=nothing, dims=1)

Compute the scatter matrix, which is an unnormalized covariance matrix.
A weighting vector `wv` can be specified to weight
the estimate.

# Arguments
* `mean=nothing`: a known mean value. `nothing` indicates that the mean is
  unknown, and the function will compute the mean. Specifying `mean=0` indicates
  that the data are centered and hence there's no need to subtract the mean.
* `dims=1`: the dimension along which the variables are organized.
  When `dims = 1`, the variables are considered columns with observations in rows;
  when `dims = 2`, variables are in rows with observations in columns.
"""
function scattermat end

"""
    cov(X, w::AbstractWeights, vardim=1; mean=nothing,  corrected=false)

Compute the weighted covariance matrix. Similar to `var` and `std` the biased covariance
matrix (`corrected=false`) is computed by multiplying `scattermat(X, w)` by
``\\frac{1}{\\sum{w}}`` to normalize. However, the unbiased covariance matrix
(`corrected=true`) is dependent on the type of weights used:
* `AnalyticWeights`: ``\\frac{1}{\\sum w - \\sum {w^2} / \\sum w}``
* `FrequencyWeights`: ``\\frac{1}{\\sum{w} - 1}``
* `ProbabilityWeights`: ``\\frac{n}{(n - 1) \\sum w}`` where ``n`` equals `count(!iszero, w)`
* `Weights`: `ArgumentError` (bias correction not supported)
"""
cov

"""
    mean_and_cov(x, [wv::AbstractWeights,] vardim=1; corrected=false) -> (mean, cov)

Return the mean and covariance matrix as a tuple. A weighting
vector `wv` can be specified. `vardim` that designates whether
the variables are columns in the matrix (`1`) or rows (`2`).
Finally, bias correction is applied to the covariance calculation if
`corrected=true`. See [`cov`](@ref) documentation for more details.
"""
function mean_and_cov end

scattermat(x::DenseMatrix; mean=nothing, dims::Int=1) =
    _scattermatm(x, mean, dims)
_scattermatm(x::DenseMatrix, ::Nothing, dims::Int) =
    _unscaled_covzm(x .- mean(x, dims=dims), dims)
_scattermatm(x::DenseMatrix, mean, dims::Int=1) =
    _unscaled_covzm(x .- mean, dims)

scattermat(x::DenseMatrix, wv::AbstractWeights; mean=nothing, dims::Int=1) =
    _scattermatm(x, wv, mean, dims)
_scattermatm(x::DenseMatrix, wv::AbstractWeights, ::Nothing, dims::Int) =
    _unscaled_covzm(x .- mean(x, wv, dims=dims), wv, dims)
_scattermatm(x::DenseMatrix, wv::AbstractWeights, mean, dims::Int) =
    _unscaled_covzm(x .- mean, wv, dims)

## weighted cov
covm(x::DenseMatrix, mean, w::AbstractWeights, dims::Int=1;
     corrected::DepBool=nothing) =
    rmul!(scattermat(x, w, mean=mean, dims=dims), varcorrection(w, depcheck(:covm, corrected)))

cov(x::DenseMatrix, w::AbstractWeights, dims::Int=1; corrected::DepBool=nothing) =
    covm(x, mean(x, w, dims=dims), w, dims; corrected=depcheck(:cov, corrected))

function corm(x::DenseMatrix, mean, w::AbstractWeights, vardim::Int=1)
    c = covm(x, mean, w, vardim; corrected=false)
    s = stdm(x, w, mean, vardim; corrected=false)
    cov2cor!(c, s)
end

"""
    cor(X, w::AbstractWeights, dims=1)

Compute the Pearson correlation matrix of `X` along the dimension
`dims` with a weighting `w` .
"""
cor(x::DenseMatrix, w::AbstractWeights, dims::Int=1) =
    corm(x, mean(x, w, dims=dims), w, dims)

function mean_and_cov(x::DenseMatrix, dims::Int=1; corrected::Bool=true)
    m = mean(x, dims=dims)
    return m, covm(x, m, dims, corrected=corrected)
end
function mean_and_cov(x::DenseMatrix, wv::AbstractWeights, dims::Int=1;
                      corrected::DepBool=nothing)
    m = mean(x, wv, dims=dims)
    return m, cov(x, wv, dims; corrected=depcheck(:mean_and_cov, corrected))
end

"""
    cov2cor(C, s)

Compute the correlation matrix from the covariance matrix `C` and a vector of standard
deviations `s`. Use `StatsBase.cov2cor!` for an in-place version.
"""
cov2cor(C::AbstractMatrix, s::AbstractArray) = cov2cor!(copy(C), s)

"""
    cor2cov(C, s)

Compute the covariance matrix from the correlation matrix `C` and a vector of standard
deviations `s`. Use `StatsBase.cor2cov!` for an in-place version.
"""
cor2cov(C::AbstractMatrix, s::AbstractArray) = cor2cov!(copy(C), s)

"""
    cor2cov!(C, s)

Converts the correlation matrix `C` to a covariance matrix in-place using a vector of
standard deviations `s`.
"""
function cor2cov!(C::AbstractMatrix, s::AbstractArray)
    n = length(s)
    size(C) == (n, n) || throw(DimensionMismatch("inconsistent dimensions"))
    for i in CartesianIndices(size(C))
        @inbounds C[i] *= s[i[1]] * s[i[2]]
    end
    return C
end


@testset "Pretest Tests" begin
    @testset "Levene" begin
        @testset "Multidimensional Array" begin
            @testset "1-way ANOVA" begin
                expected = [AnovaValue(  "Total", 48.24912, 19),
                            AnovaResult("Groups",  1.69792,  3, 0.56597333, 0.19452932, 0.89858, -0.137424287),
                            AnovaFactor( "Error", 46.5512,  16, 2.90945)]

                @testset "Replicate Vectors" begin
                    observations = Array{Vector{Float64}, 1}(undef, 4)
                    observations[1] = [ 60.8,  57.0,  65.0, 58.6,  61.7]
                    observations[2] = [ 68.7,  67.7,  74.9, 66.3,  69.8]
                    observations[3] = [102.6, 102.1, 100.2, 96.5, 100.4]
                    observations[4] = [ 87.9,  84.2,  83.1, 85.7,  90.3]

                    results = levene(observations)

                    @test all(expected .≈ results.effects)
                end

                @testset "Replicate First Dimension" begin
                    observations = [60.8 68.7 102.6 87.9;
                                    57.0 67.7 102.1 84.2;
                                    65.0 74.9 100.2 83.1;
                                    58.6 66.3  96.5 85.7;
                                    61.7 69.8 100.4 90.3]

                    results = levene(observations)

                    @test all(expected .≈ results.effects)
                end
            end

            @testset "2-way ANOVA" begin
                expected = [AnovaValue(  "Total", 135.172,   19),
                            AnovaResult("Groups",  89.24992,  3,  29.749973, 10.365375, 0.00049480716, 0.58416608),
                            AnovaFactor( "Error",  45.92208, 16,   2.87013)]

                @testset "Replicate Vectors" begin
                    observations = Array{Vector{Float64}, 2}(undef, 2, 2)
                    observations[1,1] = [16.5, 18.4, 12.7, 14.0, 12.8]
                    observations[1,2] = [14.5, 11.0, 10.8, 14.3, 10.0]
                    observations[2,1] = [39.1, 26.2, 21.3, 35.8, 40.2]
                    observations[2,2] = [32.0, 23.8, 28.8, 25.0, 29.3]

                    results = levene(observations)

                    @test all(expected .≈ results.effects)
                end

                @testset "Replicate First Dimension" begin
                    observations = cat(hcat([16.5, 18.4, 12.7, 14.0, 12.8], [39.1, 26.2, 21.3, 35.8, 40.2]),
                                        hcat([14.5, 11.0, 10.8, 14.3, 10.0], [32.0, 23.8, 28.8, 25.0, 29.3]), dims = 3)

                    results = levene(observations)

                    @test all(expected .≈ results.effects)
                end
            end

            @testset "2-way ANOVA without replicates" begin
                @testset "Fixed-effects (Model I)" begin
                    @testset "Replicate Vectors" begin
                        observations = Array{Vector{Float64}, 2}(undef, 3, 4)
                        observations[1,1] = [123]
                        observations[1,2] = [138]
                        observations[1,3] = [110]
                        observations[1,4] = [151]
                        observations[2,1] = [145]
                        observations[2,2] = [165]
                        observations[2,3] = [140]
                        observations[2,4] = [167]
                        observations[3,1] = [156]
                        observations[3,2] = [176]
                        observations[3,3] = [185]
                        observations[3,4] = [175]

                        @test_throws ErrorException levene(observations)
                    end
                end
            end

            @testset "3-way ANOVA" begin
                expected = [AnovaValue(  "Total", 0.50277778,  71),
                            AnovaResult("Groups", 0.077777778, 17, 0.0045751634, 0.58131488, 0.8912173, -0.109700816),
                            AnovaFactor( "Error", 0.425,       54, 0.0078703704)]

                @testset "Replicate Vectors" begin
                    observations = Array{Vector{Float64}, 3}(undef, 2, 3, 3)
                    observations[1,1,1] = [1.9, 1.8, 1.6, 1.4]
                    observations[1,1,2] = [2.1, 2.0, 1.8, 2.2]
                    observations[1,1,3] = [1.1, 1.2, 1.0, 1.4]
                    observations[1,2,1] = [2.3, 2.1, 2.0, 2.6]
                    observations[1,2,2] = [2.4, 2.6, 2.7, 2.3]
                    observations[1,2,3] = [2.0, 2.1, 1.9, 2.2]
                    observations[1,3,1] = [2.9, 2.8, 3.4, 3.2]
                    observations[1,3,2] = [3.6, 3.1, 3.4, 3.2]
                    observations[1,3,3] = [2.9, 2.8, 3.0, 3.1]
                    observations[2,1,1] = [1.8, 1.7, 1.4, 1.5]
                    observations[2,1,2] = [2.3, 2.0, 1.9, 1.7]
                    observations[2,1,3] = [1.4, 1.0, 1.3, 1.2]
                    observations[2,2,1] = [2.4, 2.7, 2.4, 2.6]
                    observations[2,2,2] = [2.0, 2.3, 2.1, 2.4]
                    observations[2,2,3] = [2.4, 2.6, 2.3, 2.2]
                    observations[2,3,1] = [3.0, 3.1, 3.0, 2.7]
                    observations[2,3,2] = [3.1, 3.0, 2.8, 3.2]
                    observations[2,3,3] = [3.2, 2.9, 2.8, 2.9]

                    results = levene(observations)

                    @test all(expected .≈ results.effects)
                end

                @testset "Replicate First Dimension" begin
                    observations = cat(cat(hcat([1.9, 1.8, 1.6, 1.4], [1.8, 1.7, 1.4, 1.5]),
                                            hcat([2.3, 2.1, 2.0, 2.6], [2.4, 2.7, 2.4, 2.6]),
                                            hcat([2.9, 2.8, 3.4, 3.2], [3.0, 3.1, 3.0, 2.7]), dims = 3),
                                        cat(hcat([2.1, 2.0, 1.8, 2.2], [2.3, 2.0, 1.9, 1.7]),
                                            hcat([2.4, 2.6, 2.7, 2.3], [2.0, 2.3, 2.1, 2.4]),
                                            hcat([3.6, 3.1, 3.4, 3.2], [3.1, 3.0, 2.8, 3.2]), dims = 3),
                                        cat(hcat([1.1, 1.2, 1.0, 1.4], [1.4, 1.0, 1.3, 1.2]),
                                            hcat([2.0, 2.1, 1.9, 2.2], [2.4, 2.6, 2.3, 2.2]),
                                            hcat([2.9, 2.8, 3.0, 3.1], [3.2, 2.9, 2.8, 2.9]), dims = 3), dims = 4)

                    results = levene(observations)

                    @test all(expected .≈ results.effects)
                end
            end
        end

        @testset "Vectors" begin
            observations = [1.9, 1.8, 1.6, 1.4, 2.1, 2.0, 1.8, 2.2, 1.1, 1.2, 1.0, 1.4,
                            2.3, 2.1, 2.0, 2.6, 2.4, 2.6, 2.7, 2.3, 2.0, 2.1, 1.9, 2.2,
                            2.9, 2.8, 3.4, 3.2, 3.6, 3.1, 3.4, 3.2, 2.9, 2.8, 3.0, 3.1,
                            1.8, 1.7, 1.4, 1.5, 2.3, 2.0, 1.9, 1.7, 1.4, 1.0, 1.3, 1.2,
                            2.4, 2.7, 2.4, 2.6, 2.0, 2.3, 2.1, 2.4, 2.4, 2.6, 2.3, 2.2,
                            3.0, 3.1, 3.0, 2.7, 3.1, 3.0, 2.8, 3.2, 3.2, 2.9, 2.8, 2.9]

            factorassignments = [[repeat([1], 36); repeat([2], 36)],
                                 repeat([repeat([1], 12); repeat([2], 12); repeat([3], 12)], 2),
                                 repeat([repeat([1], 4); repeat([2], 4); repeat([3], 4)], 6)]

            expected = [AnovaValue(  "Total", 0.50277778,  71),
                        AnovaResult("Groups", 0.077777778, 17, 0.0045751634, 0.58131488, 0.8912173, -0.109700816),
                        AnovaFactor( "Error", 0.425,       54, 0.0078703704)]

            @testset "3-way ANOVA ordered" begin
                results = levene(observations, factorassignments)
                @test all(expected .≈ results.effects)
            end

            @testset "3-way ANOVA shuffled" begin
                randorder = [65,  6, 31, 70, 45,  2, 41, 59,  3, 39, 40, 47, 61, 55, 13, 68, 67, 25,
                             60, 66, 44, 33, 42, 52, 56, 15, 49, 43, 69, 10, 51, 58, 19, 22,  8, 37,
                             71, 23, 46, 24, 54, 53, 16, 18, 11, 32, 29, 35, 63, 34, 64, 62, 28, 72,
                             30, 38, 17, 21,  9,  5, 50, 12,  7, 36, 48, 27, 14, 20,  4,  1, 57, 26]
                shuffledobservations = observations[randorder]
                shuffledfactorassignments = [fa[randorder] for fa ∈ factorassignments]

                results = levene(shuffledobservations, shuffledfactorassignments)

                @test all(expected .≈ results.effects)
            end

            @testset "3-way ANOVA tolerates nonconsecutive factor levels" begin
                nonconsecutivefactorassignments = [fa .* 2 for fa ∈ factorassignments]

                results = levene(observations, nonconsecutivefactorassignments)

                @test all(expected .≈ results.effects)
            end

            @testset "3-way ANOVA tolerates noninteger factor levels" begin
                factorassignmentstrings = [repeat(["1"], 36); repeat(["2"], 36)]
                stringsinfactorassignments = [[factorassignmentstrings]; factorassignments[2:3]]

                results = levene(observations, stringsinfactorassignments)

                @test all(expected .≈ results.effects)
            end
        end

        @testset "DataFrame" begin
            observations = [1.9, 1.8, 1.6, 1.4, 2.1, 2.0, 1.8, 2.2, 1.1, 1.2, 1.0, 1.4,
                            2.3, 2.1, 2.0, 2.6, 2.4, 2.6, 2.7, 2.3, 2.0, 2.1, 1.9, 2.2,
                            2.9, 2.8, 3.4, 3.2, 3.6, 3.1, 3.4, 3.2, 2.9, 2.8, 3.0, 3.1,
                            1.8, 1.7, 1.4, 1.5, 2.3, 2.0, 1.9, 1.7, 1.4, 1.0, 1.3, 1.2,
                            2.4, 2.7, 2.4, 2.6, 2.0, 2.3, 2.1, 2.4, 2.4, 2.6, 2.3, 2.2,
                            3.0, 3.1, 3.0, 2.7, 3.1, 3.0, 2.8, 3.2, 3.2, 2.9, 2.8, 2.9]

            factorassignments = [[repeat([1], 36); repeat([2], 36)],
                                 repeat([repeat([1], 12); repeat([2], 12); repeat([3], 12)], 2),
                                 repeat([repeat([1], 4); repeat([2], 4); repeat([3], 4)], 6)]

            expected = [AnovaValue(  "Total", 0.50277778,  71),
                        AnovaResult("Groups", 0.077777778, 17, 0.0045751634, 0.58131488, 0.8912173, -0.109700816),
                        AnovaFactor( "Error", 0.425,       54, 0.0078703704)]

            @testset "3-way ANOVA shuffled" begin
                df = DataFrame([[observations]; factorassignments], [:observations; :FactorC; :FactorB; :FactorA])

                results = levene(df, :observations, [:FactorC; :FactorB; :FactorA])

                @test all(expected .≈ results.effects)
            end
        end
    end
end