text:
export window1d

function window1d(x::Node, dims, ksize, padding, stride, dilation)
    Node(window1d, (x,dims,ksize,padding,stride,dilation))
end

function window1d(x::Var, dims, ksize, padding, stride, dilation)
    ydata = window1d(x.data, dims, ksize, padding, stride, dilation)
    Var(ydata, ∇window1d!, (x,dims,ksize,padding,stride,dilation))
end

function window1d(x::Matrix{T}, dims::Vector{Int}, ksize::Int, padding::Int, stride::Int, dilation::Int) where T
    outdims = map(dims) do d
        k = (ksize - 1) * dilation + 1
        (d + 2padding - k) ÷ stride + 1
    end
    cumdim = 0
    y = similar(x, ksize*size(x,1), sum(outdims))
    fill!(y, 0)
    yi = 1
    for n = 1:length(dims)
        ndims = dims[n]
        i = cumdim - padding + 1
        for d = 1:outdims[n]
            for j = i:dilation:i+(ksize-1)*dilation
                if cumdim < j <= cumdim+ndims
                    copyto!(y, yi, x, (j-1)*size(x,1)+1, size(x,1))
                end
                yi += size(x, 1)
            end
            i += stride
        end
        cumdim += ndims
    end
    y
end

@generated function window1d(x::CuMatrix{T}, dims::Vector{Int}, ksize, padding, stride, dilation) where T
    Ct = cstring(T)
    k = Kernel("""
    __global__ void window1d($Ct *y, $Ct *x, int *cumsizeY, int *cumsizeX,
        int batchsize, int m, int n, int ksize, int padding, int stride) {

        int idx = blockIdx.x * blockDim.x + threadIdx.x;
        int batchI = idx / m;
        if (batchI >= batchsize) return;
        idx -= m * batchI;
        int sizeY = cumsizeY[batchI+1] - cumsizeY[batchI];
        if (idx >= sizeY) return;

        int nj = idx / n;
        int ni = idx - nj * n;
        int kj = nj / ksize;
        int ki = nj - kj * ksize;
        int xj = cumsizeX[batchI] - padding + ki + kj*stride;
        int xi = ni + xj * n;
        int yi = idx + cumsizeY[batchI];
        if (xj >= cumsizeX[batchI] && xj < cumsizeX[batchI+1]) y[yi] = x[xi];
        else y[yi] = 0;
    }""")
    quote
        ydims = Array{Int}(undef, length(dims))
        for i = 1:length(dims)
            d = dims[i]
            k = (ksize - 1) * dilation + 1
            ydims[i] = (d + 2padding - k) ÷ stride + 1
        end
        y = similar(x, ksize*size(x,1), sum(ydims))
        cumsize_y = Array{Cint}(undef, length(dims)+1)
        cumsize_x = similar(cumsize_y)
        cumsize_y[1] = cumsize_x[1] = 0
        for i = 2:length(cumsize_y)
            cumsize_y[i] = cumsize_y[i-1] + size(y,1)*ydims[i-1]
            cumsize_x[i] = cumsize_x[i-1] + dims[i-1]
        end
        cumsize_y = CuArray(cumsize_y)
        cumsize_x = CuArray(cumsize_x)

        m = maximum(ydims) * size(y,1)
        gdims, bdims = cudims(m*length(ydims))
        $k(gdims, bdims, pointer(y), pointer(x), pointer(cumsize_y), pointer(cumsize_x),
            length(dims), m, size(x,1), ksize, padding, stride)
        y
    end
end

function ∇window1d!(y::Var, x::Var, dims, args...)
    isnothing(x.grad) || ∇window1d!(y.grad, x.grad, dims, args...)
end

function ∇window1d!(gy::Matrix{T}, gx::Matrix{T}, dims::Vector{Int}, ksize::Int, padding::Int, stride::Int, dilation::Int) where T
    outdims = map(dims) do d
        k = (ksize - 1) * dilation + 1
        (d + 2padding - k) ÷ stride + 1
    end
    cumdim = 0
    yi = 1
    for n = 1:length(dims)
        ndims = dims[n]
        i = cumdim - padding + 1
        for d = 1:outdims[n]
            for j = i:dilation:i+(ksize-1)*dilation
                if cumdim < j <= cumdim+ndims
                    addto!(gx, (j-1)*size(gx,1)+1, gy, yi, size(gx,1))
                end
                yi += size(gx, 1)
            end
            i += stride
        end
        cumdim += ndims
    end
end

@generated function ∇window1d!(gy::CuMatrix{T}, gx::CuMatrix{T}, dims::Vector{Int}, ksize, padding, stride, dilation) where T
    Ct = cstring(T)
    k = Kernel("""
    __global__ void window1d_grad($Ct *gy, $Ct *gx, int *cumsizeY, int *cumsizeX,
        int batchsize, int m, int n, int ksize, int padding, int stride) {

        int idx = blockIdx.x * blockDim.x + threadIdx.x;
        int batchI = idx / m;
        if (batchI >= batchsize) return;
        idx -= m * batchI;
        int sizeY = cumsizeY[batchI+1] - cumsizeY[batchI];
        if (idx >= sizeY) return;

        int nj = idx / n;
        int ni = idx - nj * n;
        int kj = nj / ksize;
        int ki = nj - kj * ksize;
        int xj = cumsizeX[batchI] - padding + ki + kj*stride;
        int xi = ni + xj * n;
        int yi = idx + cumsizeY[batchI];
        if (xj >= cumsizeX[batchI] && xj < cumsizeX[batchI+1]) atomicAdd(&gx[xi], gy[yi]);
    }""")
    quote
        ydims = Array{Int}(undef, length(dims))
        for i = 1:length(dims)
            d = dims[i]
            k = (ksize - 1) * dilation + 1
            ydims[i] = (d + 2padding - k) ÷ stride + 1
        end
        cumsize_y = Array{Cint}(undef, length(dims)+1)
        cumsize_x = similar(cumsize_y)
        cumsize_y[1] = cumsize_x[1] = 0
        for i = 2:length(cumsize_y)
            cumsize_y[i] = cumsize_y[i-1] + size(gy,1)*ydims[i-1]
            cumsize_x[i] = cumsize_x[i-1] + dims[i-1]
        end
        cumsize_y = CuArray(cumsize_y)
        cumsize_x = CuArray(cumsize_x)

        m = maximum(ydims) * size(gy,1)
        gdims, bdims = cudims(m*length(ydims))
        $k(gdims, bdims, pointer(gy), pointer(gx), pointer(cumsize_y), pointer(cumsize_x),
            length(dims), m, size(gx,1), ksize, padding, stride)
    end
end


text:
module Sandboxes

using MLStyle: @match
using ReplMaker: LineEdit, enter_mode!, initrepl

"""
    baremod()

Return an anonymous baremodule
"""
baremod() = @eval baremodule $(gensym(:Sandboxes_mod)) end

# :ccall is here because it's not exported, but is callable. It's a special-cased not-function.
# I don't think there are any others (`outer` used to be like this, but is now syntax).
const core_name_remaps = Dict(n => gensym(:Sandboxes_remaps) for n in [:ccall, names(Core)...])
const inverse_core_name_remaps = Dict(reverse.(collect(core_name_remaps)))

"""
    sandboxify(ex:Expr)

Return an expression that, when evaluated in a baremodule, cannot:

- access identifiers in `Core`
- import other modules
- define new modules (this would allow escape by `eval` or `.Core`)
- define macros (this allows escape by returning `:(:Core)`)

To punch holes in the sandbox, interpolate values into the expression like this:

```
sandboxify( :( nand(a, b) = \$(~)(\$(&)(a, b)) ) )
```

The above example will return an expression for a function that will internally
call `~` and `&` as defined in the scope you created the expression in, rather
than as defined in the sandbox.
"""
function sandboxify(ex)
    filterer(ex) = @match ex begin
        # Map names that are in Core to generated symbols so that they cannot be accessed.
        s::Symbol => get(core_name_remaps, s, s)

        # This is to prevent `baremodule X end; X.Core`.
        # QuoteNode remapping will unfortunately make it impossible to access a
        # property with a reserved name on a struct you've imported into the
        # sandbox and will make code that deals with literal QuoteNodes pretty
        # weird and error prone.
        # Maybe this can be done with less breakage.
        # QuoteNode(s) => QuoteNode(get(core_name_remaps, s, s))

        # For now let's ban all modules so we don't need to rewrite QuoteNodes
        Expr(:module, _...) => error("You cannot define modules in this sandbox.")

        # Not allowed to import stuff
        Expr(:using, _...) ||
        Expr(:import, _...) => error("Imports are not permitted in this sandbox: $ex")

        # Can't use a macro to escape the sandbox
        #
        # At the moment, I can't work out how to make macro definition inside the
        # sandbox safe (filtering the globalrefs in the macroexpand of the
        # expression, maybe, but then that interferes with macros you've
        # deliberately imported into the sandbox)
        #
        # Possibly I could put a special form inside the macro definitions that I
        # then use as a signal to turn on sandboxing, but that seems like a bunch
        # of work.
        #
        # So anyway, for now you can't define macros inside the sandbox at all.
        Expr(:macro, _...) => error("Macro definitions are not permitted in this sandbox: $ex")

        # This allows trivial escape via `eval` and `Base`
        Expr(:module, true, _...) => error("You cannot define non-bare modules in this sandbox.")

        # Recurse
        Expr(head, args...) => Expr(head, filterer.(args)...)

        # Fallback
        x => x
    end
    filterer(ex)
end

"""
    sandboxed_eval_expr(m::Module, ex::Expr)

Return an expression that, when evaluated, will safely evaluate some transformation of `ex` in `m`.
"""
function sandboxed_eval_expr(m::Module, ex)
    quote
        try
            @eval $m $(sandboxify(ex))
        catch e
            if e isa UndefVarError && haskey($inverse_core_name_remaps, e.var)
                rethrow(UndefVarError($inverse_core_name_remaps[e.var]))
            else
                rethrow()
            end
        end
    end
end

# REPL inside sandbox

"""
    repl_in(m::Module)

Define a REPL that eval's sandboxed code inside the baremodule `m`.
"""
function repl_in(m::Module)
    function valid_julia(s)
        input = String(take!(copy(LineEdit.buffer(s))))
        ex = Meta.parse(input)
        !(ex isa Expr && ex.head == :incomplete)
    end

    function sandbox_parser(s::String)
        ex = Meta.parse(s)
        sandboxed_eval_expr(m, ex)
    end

    sandbox_mode = initrepl(sandbox_parser;
                            prompt_text="Sandbox> ",
                            prompt_color = :yellow,
                            startup_text = false,
                            mode_name = :sandbox,
                            valid_input_checker = valid_julia)
end

"""
    enter_repl(m::Module)

Start a sandboxed REPL mode inside baremodule `m`.
"""
function enter_repl(m::Module)
    enter_mode!(repl_in(m))
end

end

ext: jl

max_stars_repo_name: cmcaine/Sandboxes

max_issues_repo_name: cmcaine/Sandboxes

max_forks_repo_name: cmcaine/Sandboxes


text:
"""
    issorted(g)

Return `true` if the timestamps of an evolving graph `g` is sorted and `false` otherwise.
"""
function issorted(g::EvolvingGraph)
    return issorted(g.timestamps)
end

"""
    sort_timestamps!(g)

Sort an evolving graph `g` according to the order of its timestamps.

```jldoctest
julia> using EvolvingGraphs

julia> g = EvolvingGraph()
Directed EvolvingGraph 0 nodes, 0 static edges, 0 timestamps

julia> add_bunch_of_edges!(g, [(1,2,2001),(3,4,2008),(5,6,2007),(2,1,2003)])
Directed EvolvingGraph 6 nodes, 4 static edges, 4 timestamps

julia> edges(g)
4-element Array{EvolvingGraphs.WeightedTimeEdge{EvolvingGraphs.Node{Int64},Int64,Float64},1}:
 Node(1)-1.0->Node(2) at time 2001
 Node(3)-1.0->Node(4) at time 2008
 Node(5)-1.0->Node(6) at time 2007
 Node(2)-1.0->Node(1) at time 2003

julia> sort_timestamps!(g)
Directed EvolvingGraph 6 nodes, 4 static edges, 4 timestamps

julia> edges(g)
4-element Array{EvolvingGraphs.WeightedTimeEdge{EvolvingGraphs.Node{Int64},Int64,Float64},1}:
 Node(1)-1.0->Node(2) at time 2001
 Node(2)-1.0->Node(1) at time 2003
 Node(5)-1.0->Node(6) at time 2007
 Node(3)-1.0->Node(4) at time 2008
```
"""
function sort_timestamps!(g::EvolvingGraph)
    p = sortperm(g.timestamps)
    g.edges = g.edges[p]
    g.timestamps = g.timestamps[p]
    g
end

"""
    sort_timestamps(g)

Sort evolving graph `g` according to timestamps, return a new sorted evolving graph.
"""
sort_timestamps(g::EvolvingGraph) = sort_timestamps!(deepcopy(g))

"""
    slice_timestamps!(g, t_min, t_max)

Slice an (sorted) evolving graph `g` between timestamp `t_min` and `t_max.`
"""
function slice_timestamps!{V,E,T}(g::EvolvingGraph{V,E,T}, t_min::T, t_max::T)
    issorted(g) || sorttime!(g)
    a = findfirst(g.timestamps, t_min)
    b = findlast(g.timestamps, t_max)
    g.edges = g.edges[a:b]
    g.timestamps =  g.timestamps[a:b]
    g
end

"""
    slice_timestamps(g, t_min, t_max)

Slice an (sorted) evolving graph `g` between timestamp `t_min` and `t_max.`, return a new evolving graph.

```jldoctest
julia> using EvolvingGraphs

julia> g = EvolvingGraph()
Directed EvolvingGraph 0 nodes, 0 static edges, 0 timestamps

julia> add_bunch_of_edges!(g, [(1,2,1), (2,3,1), (1,4,2), (2,3,3), (3,4,5)])
Directed EvolvingGraph 4 nodes, 5 static edges, 4 timestamps

julia> g2 = slice_timestamps(g, 2,3)
Directed EvolvingGraph 4 nodes, 2 static edges, 2 timestamps

julia> edges(g)
5-element Array{EvolvingGraphs.WeightedTimeEdge{EvolvingGraphs.Node{Int64},Int64,Float64},1}:
 Node(1)-1.0->Node(2) at time 1
 Node(2)-1.0->Node(3) at time 1
 Node(1)-1.0->Node(4) at time 2
 Node(2)-1.0->Node(3) at time 3
 Node(3)-1.0->Node(4) at time 5

julia> edges(g2)
2-element Array{EvolvingGraphs.WeightedTimeEdge{EvolvingGraphs.Node{Int64},Int64,Float64},1}:
 Node(1)-1.0->Node(4) at time 2
 Node(2)-1.0->Node(3) at time 3
```

"""
function slice_timestamps{V,E,T}(g::EvolvingGraph{V,E,T}, t_min::T, t_max::T)
    slice_timestamps!(deepcopy(g), t_min, t_max)
end


text: println("\a")

max_stars_repo_name: saltedcoffii/beep

max_issues_repo_name: saltedcoffii/beep

max_forks_repo_name: saltedcoffii/beep


text:
using JAC.AngularMomentum: clebschgordan, wigner3j, wigner6j, wigner9j

@testset "@racahsum" begin
    # Clebsch–Gordan coefficients
    @test @racahsum(clebschgordan(5/2, 3/2, 0, 0, 5/2, M), M) ≈ clebschgordan(5/2, 3/2, 0, 0, 5/2, 3/2)
    @test @racahsum(clebschgordan(5/2, 3/2, 0, 0, 5/2,-M), M) ≈ clebschgordan(5/2, 3/2, 0, 0, 5/2, 3/2)
    @test @racahsum(clebschgordan(5/2, 3/2, 0, 0, J, M), J, M) ≈ clebschgordan(5/2, 3/2, 0, 0, 5/2, 3/2)
    @test @racahsum(clebschgordan(5/2, 3/2, 0, 0, J,-M), J, M) ≈ clebschgordan(5/2, 3/2, 0, 0, 5/2, 3/2)
    @test @racahsum(clebschgordan(5/2, 3/2, 0, 0, J, M), M, J) ≈ clebschgordan(5/2, 3/2, 0, 0, 5/2, 3/2)
    @test @racahsum(clebschgordan(5/2, 3/2, 0, 0, J,-M), M, J) ≈ clebschgordan(5/2, 3/2, 0, 0, 5/2, 3/2)
    @test @racahsum(clebschgordan(j₁, m₁, 0, 0, 5/2, 3/2), m₁, j₁) ≈ clebschgordan(5/2, 3/2, 0, 0, 5/2, 3/2)
    @test @racahsum(clebschgordan(j₁,-m₁, 0, 0, 5/2, 3/2), m₁, j₁) ≈ clebschgordan(5/2, 3/2, 0, 0, 5/2, 3/2)
    @test @racahsum(clebschgordan(5/2, 3/2, j₂, m₂, 5/2, 3/2), j₂, m₂) ≈
        sum(clebschgordan(5/2, 3/2, j₂, 0, 5/2, 3/2) for j₂=0:5)
    @test @racahsum(clebschgordan(5/2, 3/2, j₂,-m₂, 5/2, 3/2), j₂, m₂) ≈
        sum(clebschgordan(5/2, 3/2, j₂, 0, 5/2, 3/2) for j₂=0:5)
    @test @racahsum(clebschgordan(2, 0, 3/2, 1/2, J, 1/2), J) ≈
        sum(clebschgordan(2, 0, 3/2, 1/2, J, 1/2) for J = HalfInt(1/2):HalfInt(7/2))
    @test @racahsum(clebschgordan(2, 0, 3/2, 1/2, J, M), J, M) ≈
        sum(clebschgordan(2, 0, 3/2, 1/2, J, 1/2) for J = HalfInt(1/2):HalfInt(7/2))
    @test @racahsum(clebschgordan(2, 0, 3/2, 1/2, J,-M), J, M) ≈
        sum(clebschgordan(2, 0, 3/2, 1/2, J, 1/2) for J = HalfInt(1/2):HalfInt(7/2))
    @test @racahsum(clebschgordan(2, 0, 3/2, 1/2, J, M), J=HalfInt(3/2):HalfInt(9/2), M) ≈
        sum(clebschgordan(2, 0, 3/2, 1/2, J, 1/2) for J = HalfInt(3/2):HalfInt(7/2))
    @test @racahsum(clebschgordan(2, 0, 3/2, 1/2, J,-M), J=HalfInt(3/2):HalfInt(9/2), M) ≈
        sum(clebschgordan(2, 0, 3/2, 1/2, J, 1/2) for J = HalfInt(3/2):HalfInt(7/2))
    @test @racahsum(clebschgordan(2, 0, 3/2, 1/2, J, M), J=HalfInt(1):HalfInt(4), M) ≈ 0 atol=1e-15
    @test @racahsum(clebschgordan(2, 0, 3/2, 1/2, J,-M), J=HalfInt(1):HalfInt(4), M) ≈ 0 atol=1e-15
    @test @racahsum((-1)^Integer(5/2-m) * clebschgordan(5/2, m, 5/2, -m, 0, 0), m) ≈ √6
    @test @racahsum((-1)^Integer(5/2-m) * clebschgordan(5/2, m, 5/2, -m, 2, 0), m) ≈ 0 atol=1e-15
    # Orthogonality relations for Clebsch–Gordan coefficients
    @test @racahsum(clebschgordan(2, 2, 3/2, -3/2, J, M)*clebschgordan(2, 2, 3/2, -3/2, J, M), J, M) ≈ 1
    @test @racahsum(clebschgordan(2, 1, 3/2, -1/2, J, M)*clebschgordan(2, 1, 3/2, -1/2, J, M), J, M) ≈ 1
    @test @racahsum(clebschgordan(2, 2, 3/2, -3/2, J, M)*clebschgordan(2, 1, 3/2, -1/2, J, M), J, M) ≈ 0 atol=1e-15
    @test @racahsum(clebschgordan(2, 2, 3/2, -3/2, J,-M)*clebschgordan(2, 2, 3/2, -3/2, J,-M), J, M) ≈ 1
    @test @racahsum(clebschgordan(2, 1, 3/2, -1/2, J,-M)*clebschgordan(2, 1, 3/2, -1/2, J,-M), J, M) ≈ 1
    @test @racahsum(clebschgordan(2, 2, 3/2, -3/2, J,-M)*clebschgordan(2, 1, 3/2, -1/2, J,-M), J, M) ≈ 0 atol=1e-15
    @test @racahsum(clebschgordan(5/2, m₁, 3/2, m₂, 2, 0)*clebschgordan(5/2, m₁, 3/2, m₂, 2, 0), m₁, m₂) ≈ 1
    @test @racahsum(clebschgordan(5/2, m₁, 3/2, m₂, 3, 0)*clebschgordan(5/2, m₁, 3/2, m₂, 3, 0), m₁, m₂) ≈ 1
    @test @racahsum(clebschgordan(5/2, m₁, 3/2, m₂, 3, 1)*clebschgordan(5/2, m₁, 3/2, m₂, 3, 1), m₁, m₂) ≈ 1
    @test @racahsum(clebschgordan(5/2, m₁, 3/2, m₂, 2, 0)*clebschgordan(5/2, m₁, 3/2, m₂, 3, 0), m₁, m₂) ≈ 0 atol=1e-15
    @test @racahsum(clebschgordan(5/2, m₁, 3/2, m₂, 3, 0)*clebschgordan(5/2, m₁, 3/2, m₂, 3, 1), m₁, m₂) ≈ 0 atol=1e-15
    @test @racahsum(clebschgordan(5/2, m₁, 3/2, m₂, 2, 0)*clebschgordan(5/2, m₁, 3/2, m₂, 3, 1), m₁, m₂) ≈ 0 atol=1e-15
    @test @racahsum(clebschgordan(5/2,-m₁, 3/2, m₂, 2, 0)*clebschgordan(5/2,-m₁, 3/2, m₂, 2, 0), m₁, m₂) ≈ 1
    @test @racahsum(clebschgordan(5/2, m₁, 3/2,-m₂, 3, 0)*clebschgordan(5/2, m₁, 3/2,-m₂, 3, 0), m₁, m₂) ≈ 1
    @test @racahsum(clebschgordan(5/2,-m₁, 3/2,-m₂, 3, 1)*clebschgordan(5/2,-m₁, 3/2,-m₂, 3, 1), m₁, m₂) ≈ 1
    @test @racahsum(clebschgordan(5/2,-m₁, 3/2, m₂, 2, 0)*clebschgordan(5/2,-m₁, 3/2, m₂, 3, 0), m₁, m₂) ≈ 0 atol=1e-15
    @test @racahsum(clebschgordan(5/2, m₁, 3/2,-m₂, 3, 0)*clebschgordan(5/2, m₁, 3/2,-m₂, 3, 1), m₁, m₂) ≈ 0 atol=1e-15
    @test @racahsum(clebschgordan(5/2,-m₁, 3/2,-m₂, 2, 0)*clebschgordan(5/2,-m₁, 3/2,-m₂, 3, 1), m₁, m₂) ≈ 0 atol=1e-15
    # 3-j symbol
    @test @racahsum(wigner3j(2, j, 4, 0, 0, 0), j) ≈ sum(wigner3j(2, j, 4, 0, 0, 0) for j=2:2:6)
    @test @racahsum(wigner3j(2, j, 4, m, 0, m′), j, m, m′) ≈ sum(wigner3j(2, j, 4, m, 0, -m) for j=2:6, m=-2:2)
    @test @racahsum(wigner3j(2, j, 4,-m, 0, m′), j, m, m′) ≈ sum(wigner3j(2, j, 4, m, 0, -m) for j=2:6, m=-2:2)
    @test @racahsum(wigner3j(2, j, 4, m, 0,-m′), j, m, m′) ≈ sum(wigner3j(2, j, 4, m, 0, -m) for j=2:6, m=-2:2)
    @test @racahsum(wigner3j(2, j, 4,-m, 0,-m′), j, m, m′) ≈ sum(wigner3j(2, j, 4, m, 0, -m) for j=2:6, m=-2:2)
    @test @racahsum((-1)^Integer(3/2-m)*wigner3j(3/2, 3/2, 0,  m, -m, 0), m) ≈ 2
    @test @racahsum((-1)^Integer(3/2+m)*wigner3j(3/2, 3/2, 0, -m,  m, 0), m) ≈ 2
    @test @racahsum((-1)^Integer(3/2-m)*wigner3j(3/2, 3/2, 1,  m, -m, 0), m) ≈ 0 atol=1e-15
    @test @racahsum((-1)^Integer(3/2+m)*wigner3j(3/2, 3/2, 1, -m,  m, 0), m) ≈ 0 atol=1e-15
    @test @racahsum((-1)^Integer(4-m)*wigner3j(4, 4, 0,  m, -m, 0), m) ≈ 3
    @test @racahsum((-1)^Integer(4+m)*wigner3j(4, 4, 0, -m,  m, 0), m) ≈ 3
    @test @racahsum((-1)^Integer(4-m)*wigner3j(4, 4, 1,  m, -m, 0), m) ≈ 0 atol=1e-15
    @test @racahsum((-1)^Integer(4+m)*wigner3j(4, 4, 1, -m,  m, 0), m) ≈ 0 atol=1e-15
    # Orthogonality relations for 3-j symbols
    @test 9 * @racahsum(wigner3j(2, 6, 4,  m₁,  m₂, 0)*wigner3j(2, 6, 4,  m₁,  m₂, 0), m₁, m₂) ≈ 1
    @test 9 * @racahsum(wigner3j(2, 6, 4, -m₁,  m₂, 0)*wigner3j(2, 6, 4, -m₁,  m₂, 0), m₁, m₂) ≈ 1
    @test 9 * @racahsum(wigner3j(2, 6, 4,  m₁, -m₂, 0)*wigner3j(2, 6, 4,  m₁, -m₂, 0), m₁, m₂) ≈ 1
    @test 9 * @racahsum(wigner3j(2, 6, 4, -m₁, -m₂, 0)*wigner3j(2, 6, 4, -m₁, -m₂, 0), m₁, m₂) ≈ 1
    @test 9 * @racahsum(wigner3j(2, 6, 4,  m₁,  m₂, 0)*wigner3j(2, 6, 5,  m₁,  m₂, 0), m₁, m₂) ≈ 0 atol=1e-15
    @test 9 * @racahsum(wigner3j(2, 6, 4, -m₁,  m₂, 0)*wigner3j(2, 6, 5, -m₁,  m₂, 0), m₁, m₂) ≈ 0 atol=1e-15
    @test 9 * @racahsum(wigner3j(2, 6, 4,  m₁, -m₂, 0)*wigner3j(2, 6, 5,  m₁, -m₂, 0), m₁, m₂) ≈ 0 atol=1e-15
    @test 9 * @racahsum(wigner3j(2, 6, 4, -m₁, -m₂, 0)*wigner3j(2, 6, 5, -m₁, -m₂, 0), m₁, m₂) ≈ 0 atol=1e-15
    @test 9 * @racahsum(wigner3j(2, 6, 4,  m₁,  m₂, 0)*wigner3j(2, 6, 4,  m₁,  m₂, 1), m₁, m₂) ≈ 0 atol=1e-15
    @test 9 * @racahsum(wigner3j(2, 6, 4, -m₁,  m₂, 0)*wigner3j(2, 6, 4, -m₁,  m₂, 1), m₁, m₂) ≈ 0 atol=1e-15
    @test 9 * @racahsum(wigner3j(2, 6, 4,  m₁, -m₂, 0)*wigner3j(2, 6, 4,  m₁, -m₂, 1), m₁, m₂) ≈ 0 atol=1e-15
    @test 9 * @racahsum(wigner3j(2, 6, 4, -m₁, -m₂, 0)*wigner3j(2, 6, 4, -m₁, -m₂, 1), m₁, m₂) ≈ 0 atol=1e-15
    @test 6 * @racahsum(wigner3j(3/2, 1, 5/2, m₁, m₂, -1/2)*wigner3j(3/2, 1, 5/2, m₁, m₂, -1/2), m₁, m₂) ≈ 1
    @test 4 * @racahsum(wigner3j(3/2, 1, 3/2, m₁, m₂, -1/2)*wigner3j(3/2, 1, 3/2, m₁, m₂, -1/2), m₁, m₂) ≈ 1
    @test 6 * @racahsum(wigner3j(3/2, 1, 5/2, m₁, m₂, -1/2)*wigner3j(3/2, 1, 3/2, m₁, m₂, -1/2), m₁, m₂) ≈ 0 atol=1e-15
    @test 6 * @racahsum(wigner3j(3/2, 1, 5/2, m₁, m₂, -1/2)*wigner3j(3/2, 1, 5/2, m₁, m₂,  1/2), m₁, m₂) ≈ 0 atol=1e-15
    @test 4 * @racahsum(wigner3j(3/2, 1, 3/2, m₁, m₂, -1/2)*wigner3j(3/2, 1, 3/2, m₁, m₂, -3/2), m₁, m₂) ≈ 0 atol=1e-15
    @test @racahsum((2j₃+1)*wigner3j(2, 6, j₃, 0, 0, m₃)*wigner3j(2, 6, j₃, 0, 0, m₃), j₃, m₃) ≈ 1
    @test @racahsum((2j₃+1)*wigner3j(2, 6, j₃, 0, 0, m₃)*wigner3j(2, 6, j₃, 0, 1, m₃), j₃, m₃) ≈ 0 atol=1e-15
    @test @racahsum((2j₃+1)*wigner3j(2, 6, j₃, 0, 0, m₃)*wigner3j(2, 6, j₃, 1, 0, m₃), j₃, m₃) ≈ 0 atol=1e-15
    @test @racahsum((2j₃+1)*wigner3j(3/2, 1, j₃, 1/2, 0, m₃)*wigner3j(3/2, 1, j₃,  1/2, 0, m₃), j₃, m₃) ≈ 1
    @test @racahsum((2j₃+1)*wigner3j(3/2, 1, j₃, 1/2, 0, m₃)*wigner3j(3/2, 1, j₃,  1/2, 1, m₃), j₃, m₃) ≈ 0 atol=1e-15
    @test @racahsum((2j₃+1)*wigner3j(3/2, 1, j₃, 1/2, 0, m₃)*wigner3j(3/2, 1, j₃, -1/2, 0, m₃), j₃, m₃) ≈ 0 atol=1e-15
    # 6-j symbol as sum over 3-j symbols
    @test @racahsum((-1)^Integer(18-m₁-m₂-m₃-m₄-m₅-m₆) * wigner3j(3, 3, 3, -m₁, -m₂, -m₃) *
                    wigner3j(3, 3, 3, m₁, -m₅, m₆) * wigner3j(3, 3, 3, m₄, m₂, -m₆) *
                    wigner3j(3, 3, 3, -m₄, m₅, m₃), m₁, m₂, m₃, m₄, m₅, m₆) ≈ wigner6j(3, 3, 3, 3, 3, 3)
    @test @racahsum((-1)^Integer(30-m₁-m₂-m₃-m₄-m₅-m₆) * wigner3j(5, 5, 5, -m₁, -m₂, -m₃) *
                    wigner3j(5, 5, 5, m₁, -m₅, m₆) * wigner3j(5, 5, 5, m₄, m₂, -m₆) *
                    wigner3j(5, 5, 5, -m₄, m₅, m₃), m₁, m₂, m₃, m₄, m₅, m₆) ≈ wigner6j(5, 5, 5, 5, 5, 5)
    @test @racahsum((-1)^Integer(9-m₁-m₂-m₃-m₄-m₅-m₆) * wigner3j(1, 3/2, 5/2, -m₁, -m₂, -m₃) *
                    wigner3j(1, 5/2, 3/2, m₁, -m₅, m₆) * wigner3j(0, 3/2, 3/2, m₄, m₂, -m₆) *
                    wigner3j(0, 5/2, 5/2, -m₄, m₅, m₃), m₁, m₂, m₃, m₄, m₅, m₆) ≈ wigner6j(1, 3/2, 5/2, 0, 5/2, 3/2)
    # 6-j symbol
    @test @racahsum(wigner6j(1,1,1,1,1,j), j) ≈ wigner6j(1,1,1,1,1,0) + wigner6j(1,1,1,1,1,1) + wigner6j(1,1,1,1,1,2)
    @test @racahsum(j*wigner6j(1,1,1,1,1,j), j) ≈ wigner6j(1,1,1,1,1,1) + 2wigner6j(1,1,1,1,1,2)
    @test @racahsum(wigner6j(1,3/2,j,0,5/2,3/2), j) ≈ wigner6j(1,3/2,5/2,0,5/2,3/2)
    # Orthogonality relation for 6-j symbols
    @test @racahsum((2j₃+1)*wigner6j(1,3/2,j₃,0,5/2,3/2)*wigner6j(1,3/2,j₃,0,5/2,3/2), j₃) ≈ 1/4
    @test @racahsum((2j₃+1)*wigner6j(1,3/2,j₃,0,5/2,3/2)*wigner6j(1,3/2,j₃,0,5/2,5/2), j₃) ≈ 0 atol=1e-15
    @test @racahsum((2j₃+1)*wigner6j(1,1,j₃,1,1,1)*wigner6j(1,1,j₃,1,1,0), j₃) ≈ 0 atol=1e-15
    @test @racahsum((2j₃+1)*wigner6j(1,1,j₃,1,1,1)*wigner6j(1,1,j₃,1,1,1), j₃) ≈ 1/3
    @test @racahsum((2j₃+1)*wigner6j(1,1,j₃,1,1,1)*wigner6j(1,1,j₃,1,1,2), j₃) ≈ 0 atol=1e-15
    # 9-j symbol as sum over 6-j symbols
    @test @racahsum((-1)^2j*(2j+1)*wigner6j(1,1,1,1,0,j)*wigner6j(1,1,1,1,j,1)*wigner6j(1,1,0,j,1,1), j) ≈
        wigner9j(1,1,1,1,1,1,1,1,0)
    @test @racahsum((-1)^2j*(2j+1)*wigner6j(1,5/2,3/2,3/2,0,j)*wigner6j(3/2,0,3/2,5/2,j,5/2)*
                    wigner6j(5/2,5/2,0,j,1,3/2), j) ≈ wigner9j(1,3/2,5/2,5/2,0,5/2,3/2,3/2,0)
    @test @racahsum((-1)^2j*(2j+1)*wigner6j(1,5/2,3/2,3/2,0,j)*wigner6j(3/2,0,3/2,5/2,j,5/2)*
                    wigner6j(5/2,5/2,0,j,1,3/2), j) ≈ wigner9j(1,3/2,5/2,5/2,0,5/2,3/2,3/2,0)
    @test @racahsum((-1)^2j*(2j+1)*wigner6j(1/2,3,5/2,5/2,1,j)*wigner6j(1/2,3,5/2,3,j,2)*
                    wigner6j(1,2,1,j,1/2,1/2), j) ≈ wigner9j(1/2,1/2,1,3,3,2,5/2,5/2,1)
    # 9-j symbol
    @test @racahsum(wigner9j(1,1,1,1,1,1,1,1,j), j) ≈ sum(wigner9j(1,1,1,1,1,1,1,1,j) for j=0:2)
    @test @racahsum(j^2*wigner9j(1,1,1,1,1,1,1,1,j), j) ≈ wigner9j(1,1,1,1,1,1,1,1,1) + 4wigner9j(1,1,1,1,1,1,1,1,2)
    @test @racahsum(wigner9j(1,3/2,5/2,5/2,j,5/2,3/2,3/2,0), j) ≈ sum(wigner9j(1,3/2,5/2,5/2,j,5/2,3/2,3/2,0) for j=0:3)
    @test @racahsum(wigner9j(1/2,1/2,1,j,3,j,5/2,5/2,1), j) ≈ wigner9j(1/2,1/2,1,2,3,2,5/2,5/2,1)
    # Orthogonality relation for 9-j symbols
    @test @racahsum((2j₇+1)*(2j₈+1)*wigner9j(1/2,1/2,1,3,3,2,j₇,j₈,1)*wigner9j(1/2,1/2,1,3,3,2,j₇,j₈,1), j₇, j₈) ≈ 1/15
    @test @racahsum((2j₇+1)*(2j₈+1)*wigner9j(1/2,1/2,0,3,3,1,j₇,j₈,1)*wigner9j(1/2,1/2,0,3,3,1,j₇,j₈,1), j₇, j₈) ≈ 1/3
    @test @racahsum((2j₇+1)*(2j₈+1)*wigner9j(1/2,1/2,0,3,3,1,j₇,j₈,1)*wigner9j(1/2,1/2,1,3,3,2,j₇,j₈,1), j₇, j₈) ≈ 0 atol=1e-15
    # If summation bounds cannot be determined and were not specified explicitly, sum over 0:1/2:30
    @test @racahsum(wigner6j(j,0,j,0,j,0), j) ≈ sum(wigner6j(j,0,j,0,j,0) for j=0:HalfInt(1/2):30)
    @test @racahsum(wigner6j(j,0,j,0,j,0), j=0:10) ≈ sum(wigner6j(j,0,j,0,j,0) for j=0:10)
end


text:
# # Temperature and metallic systems
#
# In this example we consider the modeling of a magnesium lattice
# as a simple example for a metallic system.
# For our treatment we will use the PBE exchange-correlation functional.
# First we import required packages and setup the lattice.
# Again notice that DFTK uses the convention that lattice vectors are
# specified column by column.

using DFTK
using Plots
using Unitful
using UnitfulAtomic

a = 3.01794  # bohr
b = 5.22722  # bohr
c = 9.77362  # bohr
lattice = [[-a -a  0]; [-b  b  0]; [0   0 -c]]
Mg = ElementPsp(:Mg, psp=load_psp("hgh/pbe/Mg-q2"))
atoms = [Mg => [[2/3, 1/3, 1/4], [1/3, 2/3, 3/4]]];

# Next we build the PBE model and discretize it.
# Since magnesium is a metal we apply a small smearing
# temperature to ease convergence using the Fermi-Dirac
# smearing scheme. Note that both the `Ecut` is too small
# as well as the minimal ``k``-point spacing
# `kspacing` far too large to give a converged result.
# These have been selected to obtain a fast execution time.
# By default `PlaneWaveBasis` chooses a `kspacing`
# of `2π * 0.022` inverse Bohrs, which is much more reasonable.

kspacing = 0.945 / u"angstrom"        # Minimal spacing of k-points,
##                                      in units of wavevectors (inverse Bohrs)
Ecut = 5                              # Kinetic energy cutoff in Hartree
temperature = 0.01                    # Smearing temperature in Hartree
smearing = DFTK.Smearing.FermiDirac() # Smearing method
##                                      also supported: Gaussian,
##                                      MarzariVanderbilt,
##                                      and MethfesselPaxton(order)

model = model_DFT(lattice, atoms, [:gga_x_pbe, :gga_c_pbe];
                  temperature=temperature,
                  smearing=smearing)
kgrid = kgrid_from_minimal_spacing(lattice, kspacing)
basis = PlaneWaveBasis(model; Ecut, kgrid);

# Finally we run the SCF. Two magnesium atoms in
# our pseudopotential model result in four valence electrons being explicitly
# treated. Nevertheless this SCF will solve for eight bands by default
# in order to capture partial occupations beyond the Fermi level due to
# the employed smearing scheme. In this example we use a damping of `0.8`.
# The default `LdosMixing` should be suitable to converge metallic systems
# like the one we model here. For the sake of demonstration we still switch to
# Kerker mixing here.
scfres = self_consistent_field(basis, damping=0.8, mixing=KerkerMixing());
#-
scfres.occupation[1]
#-
scfres.energies

# The fact that magnesium is a metal is confirmed
# by plotting the density of states around the Fermi level.
plot_dos(scfres)


text:
using Materials
using Base.Test

@testset "Materialize" begin
    @test Materials.materialize(Materials) === nothing
end

@testset "Material" begin
    @test isfile(material(Materials, "P53"))
    @material Materials P53
    @test P53 isa String
    @test isfile(P53)

    @material Materials primes
    @test primes isa Vector{Int}
    @test length(primes) == 10_000
end


text:
"""
    GlmResp

The response vector and various derived vectors in a generalized linear model.
"""
struct GlmResp{V<:FPVector,D<:UnivariateDistribution,L<:Link} <: ModResp
    "`y`: response vector"
    y::V
    d::D
    "`devresid`: the squared deviance residuals"
    devresid::V
    "`eta`: the linear predictor"
    eta::V
    "`mu`: mean response"
    mu::V
    "`offset:` offset added to `Xβ` to form `eta`.  Can be of length 0"
    offset::V
    "`wts:` prior case weights.  Can be of length 0."
    wts::V
    "`wrkwt`: working case weights for the Iteratively Reweighted Least Squares (IRLS) algorithm"
    wrkwt::V
    "`wrkresid`: working residuals for IRLS"
    wrkresid::V
end

function GlmResp(y::V, d::D, l::L, η::V, μ::V, off::V, wts::V) where {V<:FPVector, D, L}
    n  = length(y)
    nη = length(η)
    nμ = length(μ)
    lw = length(wts)
    lo = length(off)

    # Check y values
    checky(y, d)

    # Lengths of y, η, and η all need to be n
    if !(nη == nμ == n)
        throw(DimensionMismatch("lengths of η, μ, and y ($nη, $nμ, $n) are not equal"))
    end

    # Lengths of wts and off can be either n or 0
    if lw != 0 && lw != n
        throw(DimensionMismatch("wts must have length $n or length 0 but was $lw"))
    end
    if lo != 0 && lo != n
        throw(DimensionMismatch("offset must have length $n or length 0 but was $lo"))
    end

    return GlmResp{V,D,L}(y, d, similar(y), η, μ, off, wts, similar(y), similar(y))
end

function GlmResp(y::V, d::D, l::L, off::V, wts::V) where {V<:FPVector,D,L}
    η   = similar(y)
    μ   = similar(y)
    r   = GlmResp(y, d, l, η, μ, off, wts)
    initialeta!(r.eta, d, l, y, wts, off)
    updateμ!(r, r.eta)
    return r
end

deviance(r::GlmResp) = sum(r.devresid)

"""
    cancancel(r::GlmResp{V,D,L})

Returns `true` if dμ/dη for link `L` is the variance function for distribution `D`

When `L` is the canonical link for `D` the derivative of the inverse link is a multiple
of the variance function for `D`.  If they are the same a numerator and denominator term in
the expression for the working weights will cancel.
"""
cancancel(::GlmResp) = false
cancancel(::GlmResp{V,D,LogitLink}) where {V,D<:Union{Bernoulli,Binomial}} = true
cancancel(::GlmResp{V,D,NegativeBinomialLink}) where {V,D<:NegativeBinomial} = true
cancancel(::GlmResp{V,D,IdentityLink}) where {V,D<:Normal} = true
cancancel(::GlmResp{V,D,LogLink}) where {V,D<:Poisson} = true

"""
    updateμ!{T<:FPVector}(r::GlmResp{T}, linPr::T)

Update the mean, working weights and working residuals, in `r` given a value of
the linear predictor, `linPr`.
"""
function updateμ! end

function updateμ!(r::GlmResp{T}, linPr::T) where T<:FPVector
    isempty(r.offset) ? copyto!(r.eta, linPr) : broadcast!(+, r.eta, linPr, r.offset)
    updateμ!(r)
    if !isempty(r.wts)
        map!(*, r.devresid, r.devresid, r.wts)
        map!(*, r.wrkwt, r.wrkwt, r.wts)
    end
    r
end

function updateμ!(r::GlmResp{V,D,L}) where {V<:FPVector,D,L}
    y, η, μ, wrkres, wrkwt, dres = r.y, r.eta, r.mu, r.wrkresid, r.wrkwt, r.devresid

    @inbounds for i in eachindex(y, η, μ, wrkres, wrkwt, dres)
        μi, dμdη = inverselink(L(), η[i])
        μ[i] = μi
        yi = y[i]
        wrkres[i] = (yi - μi) / dμdη
        wrkwt[i] = cancancel(r) ? dμdη : abs2(dμdη) / glmvar(r.d, μi)
        dres[i] = devresid(r.d, yi, μi)
    end
end

function updateμ!(r::GlmResp{V,D,L}) where {V<:FPVector,D<:Union{Bernoulli,Binomial},L<:Link01}
    y, η, μ, wrkres, wrkwt, dres = r.y, r.eta, r.mu, r.wrkresid, r.wrkwt, r.devresid

    @inbounds for i in eachindex(y, η, μ, wrkres, wrkwt, dres)
        μi, dμdη, μomμ = inverselink(L(), η[i])
        μ[i] = μi
        yi = y[i]
        wrkres[i] = (yi - μi) / dμdη
        wrkwt[i] = cancancel(r) ? dμdη : abs2(dμdη) / μomμ
        dres[i] = devresid(r.d, yi, μi)
    end
end

function updateμ!(r::GlmResp{V,D,L}) where {V<:FPVector,D<:NegativeBinomial,L<:NegativeBinomialLink}
    y, η, μ, wrkres, wrkwt, dres = r.y, r.eta, r.mu, r.wrkresid, r.wrkwt, r.devresid

    @inbounds for i in eachindex(y, η, μ, wrkres, wrkwt, dres)
        θ = r.d.r # the shape parameter of the negative binomial distribution
        μi, dμdη, μomμ = inverselink(L(θ), η[i])
        μ[i] = μi
        yi = y[i]
        wrkres[i] = (yi - μi) / dμdη
        wrkwt[i] = dμdη
        dres[i] = devresid(r.d, yi, μi)
    end
end

"""
    wrkresp(r::GlmResp)

The working response, `r.eta + r.wrkresid - r.offset`.
"""
wrkresp(r::GlmResp) = wrkresp!(similar(r.eta), r)

"""
    wrkresp!{T<:FPVector}(v::T, r::GlmResp{T})

Overwrite `v` with the working response of `r`
"""
function wrkresp!(v::T, r::GlmResp{T}) where T<:FPVector
    broadcast!(+, v, r.eta, r.wrkresid)
    isempty(r.offset) ? v : broadcast!(-, v, v, r.offset)
end

abstract type AbstractGLM <: LinPredModel end

mutable struct GeneralizedLinearModel{G<:GlmResp,L<:LinPred} <: AbstractGLM
    rr::G
    pp::L
    fit::Bool
end

function coeftable(mm::AbstractGLM)
    cc = coef(mm)
    se = stderror(mm)
    zz = cc ./ se
    CoefTable(hcat(cc,se,zz,2.0 * ccdf.(Ref(Normal()), abs.(zz))),
              ["Estimate","Std.Error","z value", "Pr(>|z|)"],
              ["x$i" for i = 1:size(mm.pp.X, 2)], 4)
end

function confint(obj::AbstractGLM, level::Real)
    hcat(coef(obj),coef(obj)) + stderror(obj)*quantile(Normal(),(1. -level)/2.)*[1. -1.]
end
confint(obj::AbstractGLM) = confint(obj, 0.95)

deviance(m::AbstractGLM) = deviance(m.rr)

function loglikelihood(m::AbstractGLM)
    r   = m.rr
    wts = r.wts
    y   = r.y
    mu  = r.mu
    d   = r.d
    ll  = zero(eltype(mu))
    if length(wts) == length(y)
        ϕ = deviance(m)/sum(wts)
        @inbounds for i in eachindex(y, mu, wts)
            ll += loglik_obs(d, y[i], mu[i], wts[i], ϕ)
        end
    else
        ϕ = deviance(m)/length(y)
        @inbounds for i in eachindex(y, mu)
            ll += loglik_obs(d, y[i], mu[i], 1, ϕ)
        end
    end
    ll
end

dof(x::GeneralizedLinearModel) = dispersion_parameter(x.rr.d) ? length(coef(x)) + 1 : length(coef(x))

function _fit!(m::AbstractGLM, verbose::Bool, maxIter::Integer, minStepFac::Real,
              convTol::Real, start)

    # Return early if model has the fit flag set
    m.fit && return m

    # Check arguments
    maxIter >= 1       || throw(ArgumentError("maxIter must be positive"))
    0 < minStepFac < 1 || throw(ArgumentError("minStepFac must be in (0, 1)"))

    # Extract fields and set convergence flag
    cvg, p, r = false, m.pp, m.rr
    lp = r.mu

    # Initialize β, μ, and compute deviance
    if start == nothing || isempty(start)
        # Compute beta update based on default response value
        # if no starting values have been passed
        delbeta!(p, wrkresp(r), r.wrkwt)
        linpred!(lp, p)
        updateμ!(r, lp)
        installbeta!(p)
    else
        # otherwise copy starting values for β
        copy!(p.beta0, start)
        fill!(p.delbeta, 0)
        linpred!(lp, p, 0)
        updateμ!(r, lp)
    end
    devold = deviance(m)

    for i = 1:maxIter
        f = 1.0 # line search factor
        local dev

        # Compute the change to β, update μ and compute deviance
        try
            delbeta!(p, r.wrkresid, r.wrkwt)
            linpred!(lp, p)
            updateμ!(r, lp)
            dev = deviance(m)
        catch e
            isa(e, DomainError) ? (dev = Inf) : rethrow(e)
        end

        # Line search
        ## If the deviance isn't declining then half the step size
        ## The convTol*dev term is to avoid failure when deviance
        ## is unchanged except for rouding errors.
        while dev > devold + convTol*dev
            f /= 2
            f > minStepFac || error("step-halving failed at beta0 = $(p.beta0)")
            try
                updateμ!(r, linpred(p, f))
                dev = deviance(m)
            catch e
                isa(e, DomainError) ? (dev = Inf) : rethrow(e)
            end
        end
        installbeta!(p, f)

        # Test for convergence
        crit = (devold - dev)/dev
        verbose && println("$i: $dev, $crit")
        if crit < convTol || dev == 0
            cvg = true
            break
        end
        @assert isfinite(crit)
        devold = dev
    end
    cvg || throw(ConvergenceException(maxIter))
    m.fit = true
    m
end

StatsBase.fit!(m::AbstractGLM; verbose::Bool=false, maxIter::Integer=30,
              minStepFac::Real=0.001, convTol::Real=1.e-6, start=nothing) =
    _fit!(m, verbose, maxIter, minStepFac, convTol, start)

function StatsBase.fit!(m::AbstractGLM, y; wts=nothing, offset=nothing, dofit::Bool=true,
                        verbose::Bool=false, maxIter::Integer=30, minStepFac::Real=0.001, convTol::Real=1.e-6,
                        start=nothing)
    r = m.rr
    V = typeof(r.y)
    r.y = copy!(r.y, y)
    isa(wts, Nothing) || copy!(r.wts, wts)
    isa(offset, Nothing) || copy!(r.offset, offset)
    initialeta!(r.eta, r.d, r.l, r.y, r.wts, r.offset)
    updateμ!(r, r.eta)
    fill!(m.pp.beta0, 0)
    m.fit = false
    if dofit
        _fit!(m, verbose, maxIter, minStepFac, convTol, start)
    else
        m
    end
end

function fit(::Type{M},
    X::Union{Matrix{T},SparseMatrixCSC{T}},
    y::V,
    d::UnivariateDistribution,
    l::Link = canonicallink(d);
    dofit::Bool = true,
    wts::V      = similar(y, 0),
    offset::V   = similar(y, 0),
    fitargs...) where {M<:AbstractGLM,T<:FP,V<:FPVector}

    # Check that X and y have the same number of observations
    if size(X, 1) != size(y, 1)
        throw(DimensionMismatch("number of rows in X and y must match"))
    end

    rr = GlmResp(y, d, l, offset, wts)
    res = M(rr, cholpred(X), false)
    return dofit ? fit!(res; fitargs...) : res
end

fit(::Type{M},
X::Union{Matrix,SparseMatrixCSC},
y::AbstractVector,
d::UnivariateDistribution,
l::Link=canonicallink(d); kwargs...) where {M<:AbstractGLM} =
    fit(M, float(X), float(y), d, l; kwargs...)

glm(F, D, args...; kwargs...) = fit(GeneralizedLinearModel, F, D, args...; kwargs...)

GLM.Link(mm::AbstractGLM) = mm.l
GLM.Link(r::GlmResp{T,D,L}) where {T,D,L} = L()
GLM.Link(r::GlmResp{T,D,L}) where {T,D<:NegativeBinomial,L<:NegativeBinomialLink} = L(r.d.r)
GLM.Link(m::GeneralizedLinearModel) = Link(m.rr)

Distributions.Distribution(r::GlmResp{T,D,L}) where {T,D,L} = D
Distributions.Distribution(m::GeneralizedLinearModel) = Distribution(m.rr)

"""
    dispersion(m::AbstractGLM, sqr::Bool=false)

Return the estimated dispersion (or scale) parameter for a model's distribution,
generally written σ² for linear models and ϕ for generalized linear models.
It is, by definition, equal to 1 for the Bernoulli, Binomial, and Poisson families.

If `sqr` is `true`, the squared dispersion parameter is returned.
"""
function dispersion(m::AbstractGLM, sqr::Bool=false)
    r = m.rr
    if dispersion_parameter(r.d)
        wrkwt, wrkresid = r.wrkwt, r.wrkresid
        s = sum(i -> wrkwt[i] * abs2(wrkresid[i]), eachindex(wrkwt, wrkresid)) / dof_residual(m)
        sqr ? s : sqrt(s)
    else
        one(eltype(r.mu))
    end
end

"""
    predict(mm::AbstractGLM, newX::AbstractMatrix; offset::FPVector=Vector{eltype(newX)}(0))

Form the predicted response of model `mm` from covariate values `newX` and, optionally,
an offset.
"""
function predict(mm::AbstractGLM, newX::AbstractMatrix;
                 offset::FPVector=eltype(newX)[])
    eta = newX * coef(mm)
    if !isempty(mm.rr.offset)
        length(offset) == size(newX, 1) ||
            throw(ArgumentError("fit with offset, so `offset` kw arg must be an offset of length `size(newX, 1)`"))
        broadcast!(+, eta, eta, offset)
    else
        length(offset) > 0 && throw(ArgumentError("fit without offset, so value of `offset` kw arg does not make sense"))
    end
    mu = [linkinv(Link(mm), x) for x in eta]
end

# A helper function to choose default values for eta
function initialeta!(eta::AbstractVector,
                    dist::UnivariateDistribution,
                    link::Link,
                    y::AbstractVector,
                    wts::AbstractVector,
                    off::AbstractVector)

    n  = length(y)
    lw = length(wts)
    lo = length(off)

    if lw == n
        @inbounds @simd for i = eachindex(y, eta, wts)
            μ      = mustart(dist, y[i], wts[i])
            eta[i] = linkfun(link, μ)
        end
    elseif lw == 0
        @inbounds @simd for i = eachindex(y, eta)
            μ      = mustart(dist, y[i], 1)
            eta[i] = linkfun(link, μ)
        end
    else
        throw(ArgumentError("length of wts must be either $n or 0 but was $lw"))
    end

    if lo == n
        @inbounds @simd for i = eachindex(eta, off)
            eta[i] -= off[i]
        end
    elseif lo != 0
        throw(ArgumentError("length of off must be either $n or 0 but was $lo"))
    end

    return eta
end

# Helper function to check that the values of y are in the allowed domain
function checky(y, d::Distribution)
    if any(x -> !insupport(d, x), y)
        throw(ArgumentError("y must be in the support of D"))
    end
    return nothing
end
function checky(y, d::Binomial)
    for yy in y
        0 ≤ yy ≤ 1 || throw(ArgumentError("$yy in y is not in [0,1]"))
    end
    return nothing
end


text:
using Test
using DiffEqSensitivity
using OrdinaryDiffEq
import DifferentiableBackwardEuler
import NonlinearEquations
import Random
import SparseArrays
import Zygote

import NonlinearEquations
import SparseArrays
@NonlinearEquations.equations function gw1d(h, Qs)
	NonlinearEquations.setnumequations(length(h))
	NonlinearEquations.addterm(1, -2 * h[1] + h[2] + Qs[1])
	for i = 2:length(h) - 1
		NonlinearEquations.addterm(i, h[i - 1] - 2 * h[i] + h[i + 1] + Qs[i])
	end
	NonlinearEquations.addterm(length(h), -2 * h[end] + h[end - 1] + Qs[end])
end

function solveforh(Qs)
	res(h) = gw1d_residuals(h, Qs)
	jac(h) = gw1d_h(h, Qs)
	h0 = zeros(N)
	h = jac(h0) \ -res(h0)
	return h
end

N = 25
Qs = 8 * ones(N) / N ^ 2
h = solveforh(Qs)

#compute the gradient for steady state
Random.seed!(1)
obsnode = rand(1:length(h))
g(h, Qs) = h[obsnode]
function g_h(h, p)
	retval = zeros(length(h))
	retval[obsnode] = 1.0
	return retval
end
g_Qs(h, Qs) = zeros(length(Qs))
f_h(h, Qs) = gw1d_h(h, Qs)
f_Qs(h, Qs) = gw1d_Qs(h, Qs)
grad = NonlinearEquations.gradient(h, Qs, g_h, g_Qs, f_h, f_Qs)
t1 = @elapsed grad = NonlinearEquations.gradient(h, Qs, g_h, g_Qs, f_h, f_Qs)

fdgrad = similar(grad)
sortedgradindices = sort(1:length(grad); by=i->abs(grad[i]), rev=true)
t2 = @elapsed for i = 1:length(fdgrad)
	global fdgrad
	global Qs
	global h
	dk = 1e-8
	h0 = h
	Qs0 = copy(Qs)
	thisQs = copy(Qs)
	thisQs[i] += dk
	thish = solveforh(thisQs)
	fdgrad[i] = (thish[obsnode] - h0[obsnode]) / dk
end
@test isapprox(grad, fdgrad; rtol=1e-4)#test the steady state gradients are working

u0 = zeros(N)
p = Qs
tspan = [0.0, 1e4]
f(u, p, t) = gw1d_residuals(u, p)
f_u(u, p, t) = gw1d_h(u, p)
f_u(storage, u, p, t) = gw1d_h!(storage, u, p)
f_p(u, p, t) = gw1d_Qs(u, p)
f_p(storage, u, p, t) = gw1d_Qs!(storage, u, p)
odef = ODEFunction(f; jac=f_u, jac_prototype=f_u(u0, p, 0.0), paramjac=f_p)
prob = ODEProblem(odef, u0, tspan, p)
soln_diffeq = solve(prob, ImplicitEuler(); u0=prob.u0, p=prob.p, abstol=1e-8, reltol=1e-8)
soln_dbe = DifferentiableBackwardEuler.steps(u0, f, f_u, f_p, (args...)->zeros(length(u0)), p, soln_diffeq.t; ftol=1e-12)
@test isapprox(soln_diffeq[:, :], soln_dbe)#test DifferentiableBackwardEuler is working for the forward solve
loss(p) = DifferentiableBackwardEuler.steps(u0, f, f_u, f_p, (args...)->zeros(length(u0)), p, soln_diffeq.t; ftol=1e-12)[obsnode, end]
dgdp_zygote = Zygote.gradient(loss, p)[1]
@test isapprox(dgdp_zygote, grad)#test that the steady state derivative matches the equivalent derivative obtained with the transient solution


text:
############################################################
## joAbstractFosterLinearOperator - extra functions

# elements(jo)
elements(A::joLooseLinearFunction{DDT,RDT}) where {DDT,RDT} = joHelpers_etc.elements_helper(A)
elements(A::joLooseMatrix{DDT,RDT}) where {DDT,RDT} = A*jo_speye(DDT,A.n)

# hasinverse(jo)
hasinverse(A::joAbstractFosterLinearOperator{DDT,RDT}) where {DDT,RDT} = !isnull(A.iop)

# issquare(jo)
issquare(A::joAbstractFosterLinearOperator{DDT,RDT}) where {DDT,RDT} = (A.m == A.n)

# istall(jo)
istall(A::joAbstractFosterLinearOperator{DDT,RDT}) where {DDT,RDT} = (A.m > A.n)

# iswide(jo)
iswide(A::joAbstractFosterLinearOperator{DDT,RDT}) where {DDT,RDT} = (A.m < A.n)

# iscomplex(jo)
iscomplex(A::joAbstractFosterLinearOperator{DDT,RDT}) where {DDT,RDT} = !(DDT<:Real && RDT<:Real)

# islinear(jo)

# isadjoint(jo)


text:
#== # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Description
#
#   Functions related to the associated Legendre functions.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# References
#
#   [1] Holmes, S. A. and W. E. Featherstone, 2002. A unified approach to the
#       Clenshaw summation and the recursive computation of very high degree and
#       order normalised associated Legendre functions. Journal of Geodesy,
#       76(5), pp. 279-299.
#
#       For more info.: http://mitgcm.org/~mlosch/geoidcookbook/node11.html
#
#   [2] Vallado, D. A (2013). Fundamentals of Astrodynamics and Applications.
#       Microcosm Press, Hawthorn, CA, USA.
#
#   [3] Schmidt, A (1917). Erdmagnetismus, Enzykl. Math. Wiss., 6, pp. 265–396.
#
#   [4] Winch, D. E., Ivers, D. J., Turner, J. P. R., Stening R. J (2005).
#       Geomagnetism and Schmidt quasi-normalization. Geophysical Journal
#       International, 160(2), pp. 487-504.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # ==#

export denormalize_legendre, legendre!, legendre

export legendre_fully_normalized!, legendre_fully_normalized
export legendre_schmidt_quasi_normalized!, legendre_schmidt_quasi_normalized
export legendre_conventional!, legendre_conventional

"""
    legendre!([N,] P::AbstractMatrix, ϕ::Number, ph_term::Bool = false, n_max::Integer = -1, m_max::Integer = -1)

Compute the associated Legendre function `P_n,m[cos(ϕ)]`. The maximum degree and
order that will be computed are given by the parameters `n_max` and `m_max`. If
they are negative, then the dimensions of matrix `P` will be used.

The result will be stored at matrix `P`.

The optional parameter `N` can be used to select the normalization. The
following values are valid:

* `Val(:full)`: Compute the fully normalized associated Legendre function (see
  `legendre_fully_normalized!`).
* `Val(:schmidt)`: Compute the Schmidt quasi-normalized associated Legendre
  function (see `legendre_schmidt_quasi_normalized!`).
* `Val(:conv)`: Compute the conventional associated Legendre function (see
  `legendre_conventional!`).

If `N` is omitted, then the full normalization will be used.

If `ph_term` is set to `true`, then the Condon-Shortley phase term `(-1)ᵐ` will
be included. If `ph_term` is not present, then it defaults to `false`.

"""
legendre!(P::AbstractMatrix, ϕ::Number, ph_term::Bool = false,
          n_max::Integer = -1, m_max::Integer = -1) =
    legendre_fully_normalized!(P, float(ϕ), ph_term, n_max, m_max)

legendre!(::Val{:full}, P::AbstractMatrix, ϕ::Number, ph_term::Bool = false,
          n_max::Integer = -1, m_max::Integer = -1) =
    legendre_fully_normalized!(P, float(ϕ), ph_term, n_max, m_max)

legendre!(::Val{:schmidt}, P::AbstractMatrix, ϕ::Number, ph_term::Bool = false,
          n_max::Integer = -1, m_max::Integer = -1) =
    legendre_schmidt_quasi_normalized!(P, float(ϕ), ph_term, n_max, m_max)

legendre!(::Val{:conv}, P::AbstractMatrix, ϕ::Number, ph_term::Bool = false,
          n_max::Integer = -1, m_max::Integer = -1) =
    legendre_conventional!(P, float(ϕ), ph_term, n_max, m_max)

"""
    legendre([N,] ϕ::Number, n_max::Integer, m_max::Integer = -1, ph_term::Bool = false)

Compute the associated Legendre function `P_n,m[cos(ϕ)]`. The maximum degree
that will be computed is `n_max` and the maximum order is `m_max`. Notice that
if `m_max` is higher than `n_max` or negative, than it is set to `n_max`.

The optional parameter `N` can be used to select the normalization. The
following values are valid:

* `Val(:full)`: Compute the fully normalized associated Legendre function (see
  `legendre_fully_normalized`).
* `Val(:schmidt)`: Compute the Schmidt quasi-normalized associated Legendre
  function (see `legendre_schmidt_quasi_normalized`).
* `Val(:conv)`: Compute the conventional associated Legendre function (see
  `legendre_conventional`).

If `N` is omitted, then the full normalization will be used (`Val(:full)`).

If `ph_term` is set to `true`, then the Condon-Shortley phase term `(-1)ᵐ` will
be included. If `ph_term` is not present, then it defaults to `false`.

# Returns

A matrix with the Legendre associated functions `P_n,m[cos(ϕ)]`.

"""
legendre(ϕ::Number, n_max::Integer, m_max::Integer = -1, ph_term::Bool = false) =
    legendre_fully_normalized(float(ϕ), n_max, m_max, ph_term)

legendre(::Val{:full}, ϕ::Number, n_max::Integer, m_max::Integer = -1,
         ph_term::Bool = false) =
    legendre_fully_normalized(float(ϕ), n_max, m_max, ph_term)

legendre(::Val{:schmidt}, ϕ::Number, n_max::Integer, m_max::Integer = -1,
         ph_term::Bool = false) =
    legendre_schmidt_quasi_normalized(float(ϕ), n_max, m_max, ph_term)

legendre(::Val{:conv}, ϕ::Number, n_max::Integer, m_max::Integer = -1,
         ph_term::Bool = false) =
    legendre_conventional(float(ϕ), n_max, m_max, ph_term)

################################################################################
#                Fully Normalized Associated Legendre Functions
################################################################################

"""
    legendre_fully_normalized!(P::AbstractMatrix, ϕ::Number, ph_term::Bool = false, n_max::Integer = -1, m_max::Integer = -1)

Compute the fully normalized associated Legendre function `P_n,m[cos(ϕ)]`. The
maximum degree and order that will be computed are given by the parameters
`n_max` and `m_max`. If they are negative, then the dimensions of matrix `P`
will be used:

    maximum degree -> number of rows
    maximum order  -> number of columns

The result will be stored at matrix `P`.

If `ph_term` is set to `true`, then the Condon-Shortley phase term `(-1)ᵐ` will
be included. If `ph_term` is not present, then it defaults to `false`.

# Remarks

This algorithm was based on [1]. Our definition of fully normalized associated
Legendre function can be seen in [2, p. 546]. The conversion is obtained by:

                 _                     -
                |  (n-m)! . k . (2n+1)  |      k = 1 if m  = 0
    K_n,m = sqrt| --------------------- |,     k = 2 if m != 0
                |         (n+m)!        |
                 -                     -
    _
    P_n,m = P_n,m * K_n,m,

          _
    where P_n,m is the fully normalized Legendre associated function.

"""
function legendre_fully_normalized!(P::AbstractMatrix, ϕ::Number,
                                    ph_term::Bool = false, n_max::Integer = -1,
                                    m_max::Integer = -1)

    # Obtain the maximum degree and order that must be computed.
    n_max, m_max = _get_degree_and_order(P, n_max, m_max)

    # Auxiliary variables to improve code performance.
    c = cos(ϕ)
    s = sqrt(1-c^2)

    s_fact = !ph_term ? +s : -s

    @inbounds for n = 0:n_max
        # Starting values.
        if n == 0
            P[0+1,0+1] = 1
            continue
        elseif n == 1
            P[1+1,0+1] = +sqrt(3)*c

            if m_max > 0
                P[1+1,1+1] = +sqrt(3)*s_fact
            end

            continue
        end

        aux_n = (2n-1)*(2n+1)

        for m = 0:n

            if n == m
                P[n+1,n+1] = s_fact*sqrt( (2n+1)/(2n) )*P[n-1+1,n-1+1]
            else
                aux_nm = (n-m)*(n+m)
                a_nm   = sqrt( aux_n / aux_nm )
                b_nm   = sqrt( ( (2n+1)*(n+m-1)*(n-m-1) ) / ( aux_nm*(2n-3) ) )

                # We assume that the matrix is not initialized. Hence, we must
                # not access elements on the upper triangle.
                if m != n-1
                    P[n+1,m+1] = a_nm*c*P[n-1+1,m+1] - b_nm*P[n-2+1,m+1]
                else
                    P[n+1,m+1] = a_nm*c*P[n-1+1,m+1]
                end
            end

            # Check if the maximum desired order has been reached.
            m == m_max && break
        end
    end

    nothing
end

"""
    legendre_fully_normalized(ϕ::T, n_max::Integer, m_max::Integer = -1, ph_term::Bool = false) where T<:AbstractFloat

Compute the fully normalized associated Legendre function `P_n,m[cos(ϕ)]`. The
maximum degree that will be computed is `n_max` and the maximum order is
`m_max`. Notice that if `m_max` is higher than `n_max` or negative, than it is
set to `n_max`.

If `ph_term` is set to `true`, then the Condon-Shortley phase term `(-1)ᵐ` will
be included. If `ph_term` is not present, then it defaults to `false`.

# Returns

A matrix with the Legendre associated functions `P_n,m[cos(ϕ)]`.

# Remarks

This algorithm was based on [1]. Our definition of fully normalized associated
Legendre function can be seen in [2, p. 546]. The conversion is obtained by:

                 _                     -
                |  (n-m)! . k . (2n+1)  |      k = 1 if m  = 0
    K_n,m = sqrt| --------------------- |,     k = 2 if m != 0
                |         (n+m)!        |
                 -                     -
    _
    P_n,m = P_n,m * K_n,m,

          _
    where P_n,m is the fully normalized Legendre associated function.

"""
function legendre_fully_normalized(ϕ::T, n_max::Integer, m_max::Integer = -1,
                                   ph_term::Bool = false) where T<:AbstractFloat

    (n_max < 0) && throw(ArgumentError("n_max must be positive."))

    ( (m_max < 0) || (m_max > n_max) ) && (m_max = n_max)

    P = zeros(T, n_max+1, m_max+1)
    legendre_fully_normalized!(P, ϕ, ph_term)
    return P
end

################################################################################
#            Schmidt Quasi-Normalized Associated Legendre Functions
################################################################################

"""
    legendre_schmidt_quasi_normalized!(P::AbstractMatrix, ϕ::Number, ph_term::Bool = false, n_max::Integer = -1, m_max::Integer = -1)

Compute the Schmidt quasi-normalized associated Legendre function
`P_n,m[cos(ϕ)]` [3,4]. The maximum degree and order that will be computed are
given by the parameters `n_max` and `m_max`. If they are negative, then the
dimensions of matrix `P` will be used:

    maximum degree -> number of rows
    maximum order  -> number of columns

The result will be stored at matrix `P`.

If `ph_term` is set to `true`, then the Condon-Shortley phase term `(-1)ᵐ` will
be included. If `ph_term` is not present, then it defaults to `false`.

# Remarks

This algorithm was based on [3,4]. The conversion is obtained by:

                 _           -
                |     (n-m)!  |    k = 1 if m  = 0
    K_n,m = sqrt| k. -------- |,   k = 2 if m != 0
                |     (n+m)!  |
                 -           -

    =
    P_n,m = P_n,m * K_n,m,

          =
    where P_n,m is the quasi-normalized normalized Legendre associated function.

"""
function legendre_schmidt_quasi_normalized!(P::AbstractMatrix, ϕ::Number,
                                            ph_term::Bool = false,
                                            n_max::Integer = -1,
                                            m_max::Integer = -1)

    # Obtain the maximum degree and order that must be computed.
    n_max, m_max = _get_degree_and_order(P, n_max, m_max)

    # Auxiliary variables to improve code performance.
    c = cos(ϕ)
    s = sqrt(1-c^2)

    s_fact = !ph_term ? +s : -s

    @inbounds for n = 0:n_max
        # Starting values.
        if n == 0
            P[0+1,0+1] = 1
            continue

        elseif n == 1
            P[1+1,0+1] = +c

            if m_max > 0
                P[1+1,1+1] = +s_fact
            end

            continue
        end

        aux_n = 2n-1 # -> sqrt( (2n-1)*(2n-1) )

        for m = 0:n

            if m == n
                P[n+1,n+1] = s_fact*sqrt( aux_n/(2n) )*P[n-1+1,n-1+1]
            else
                aux_nm = sqrt( (n-m)*(n+m) )
                a_nm   = aux_n / aux_nm
                b_nm   = sqrt( (n+m-1)*(n-m-1) ) / aux_nm

                # We assume that the matrix is not initialized. Hence, we must not
                # access elements on the upper triangle.
                if m != n-1
                    P[n+1,m+1] = a_nm*c*P[n-1+1,m+1] - b_nm*P[n-2+1,m+1]
                else
                    P[n+1,m+1] = a_nm*c*P[n-1+1,m+1]
                end
            end

            # Check if the maximum desired order has been reached.
            m == m_max && break
        end
    end

    nothing
end

"""
    legendre_schmidt_quasi_normalized(ϕ::T, n_max::Integer, m_max::Integer = -1, ph_term::Bool = false) where T<:AbstractFloat

Compute the Schmidt quasi-normalized associated Legendre function
`P_n,m[cos(ϕ)]`. The maximum degree that will be computed is `n_max` and the
maximum order is `m_max`. Notice that if `m_max` is higher than `n_max` or
negative, than it is set to `n_max`.

If `ph_term` is set to `true`, then the Condon-Shortley phase term `(-1)ᵐ` will
be included. If `ph_term` is not present, then it defaults to `false`.

# Returns

A matrix with the Legendre associated functions `P_n,m[cos(ϕ)]`.

# Remarks

This algorithm was based on [3,4]. The conversion is obtained by:

                 _           -
                |     (n-m)!  |    k = 1 if m  = 0
    K_n,m = sqrt| k. -------- |,   k = 2 if m != 0
                |     (n+m)!  |
                 -           -

    =
    P_n,m = P_n,m * K_n,m,

          =
    where P_n,m is the quasi-normalized normalized Legendre associated function.

"""
function legendre_schmidt_quasi_normalized(ϕ::T, n_max::Integer,
                                           m_max::Integer = -1,
                                           ph_term::Bool = false) where T<:AbstractFloat

    (n_max < 0) && throw(ArgumentError("n_max must be positive."))

    ( (m_max < 0) || (m_max > n_max) ) && (m_max = n_max)

    P = zeros(T, n_max+1, m_max+1)
    legendre_schmidt_quasi_normalized!(P, ϕ, ph_term)
    return P
end

################################################################################
#                  Conventional Associated Legendre Function
################################################################################

"""
    legendre_conventional!(P::AbstractMatrix, ϕ::Number, ph_term::Bool = false, n_max::Integer = -1, m_max::Integer = -1)

Compute the conventional associated Legendre function `P_n,m[cos(ϕ)]`. The
maximum degree and order that will be computed are given by the parameters
`n_max` and `m_max`. If they are negative, then the dimensions of matrix `P`
will be used:

    maximum degree -> number of rows
    maximum order  -> number of columns

The result will be stored at matrix `P`.

If `ph_term` is set to `true`, then the Condon-Shortley phase term `(-1)ᵐ` will
be included. If `ph_term` is not present, then it defaults to `false`.

"""
function legendre_conventional!(P::AbstractMatrix, ϕ::Number,
                                ph_term::Bool = false, n_max::Integer = -1,
                                m_max::Integer = -1)

    # Obtain the maximum degree and order that must be computed.
    n_max, m_max = _get_degree_and_order(P, n_max, m_max)

    # Auxiliary variables to improve code performance.
    c = cos(ϕ)
    s = sqrt(1-c^2)

    s_fact = !ph_term ? +s : -s

    @inbounds for n = 0:n_max
        # Starting values.
        if n == 0
            P[0+1,0+1] = 1
            continue
        elseif n == 1
            P[1+1,0+1] = +c

            if m_max > 0
                P[1+1,1+1] = +s_fact
            end

            continue
        end

        aux_n = 2n-1 # -> sqrt( (2n-1)*(2n-1) )

        for m = 0:n

            if n == m
                P[n+1,n+1] = s_fact*aux_n*P[n-1+1,n-1+1]
            else
                aux_nm = n-m # -> sqrt( (n-m)*(n-m) )
                a_nm   = aux_n / aux_nm
                b_nm   = (n+m-1) / aux_nm # -> sqrt( (n+m-1)*(n+m-1) ) / aux_nm

                # We assume that the matrix is not initialized. Hence, we must
                # not access elements on the upper triangle.
                if m != n-1
                    P[n+1,m+1] = a_nm*c*P[n-1+1,m+1] - b_nm*P[n-2+1,m+1]
                else
                    P[n+1,m+1] = a_nm*c*P[n-1+1,m+1]
                end
            end

            # Check if the maximum desired order has been reached.
            m == m_max && break
        end
    end

    nothing
end

"""
    legendre_conventional(ϕ::T, n_max::Integer, m_max::Integer = -1, ph_term::Bool = false) where T<:AbstractFloat

Compute the conventional associated Legendre function `P_n,m[cos(ϕ)]`. The
maximum degree that will be computed is `n_max` and the maximum order is
`m_max`. Notice that if `m_max` is higher than `n_max` or negative, than it is
set to `n_max`.

If `ph_term` is set to `true`, then the Condon-Shortley phase term `(-1)ᵐ` will
be included. If `ph_term` is not present, then it defaults to `false`.

# Returns

A matrix with the Legendre associated functions `P_n,m[cos(ϕ)]`.

"""
function legendre_conventional(ϕ::T, n_max::Integer, m_max::Integer = -1,
                               ph_term::Bool = false) where T<:AbstractFloat

    (n_max < 0) && throw(ArgumentError("n_max must be positive."))

    ( (m_max < 0) || (m_max > n_max) ) && (m_max = n_max)

    P = zeros(T, n_max+1, m_max+1)
    legendre_conventional!(P, ϕ, ph_term)
    return P
end

################################################################################
#                                   Private
################################################################################

"""
    _get_degree_and_order(P, n_max, m_max)

Return the maximum degree and order to compute the Legendre associated functions
given the matrix `P` and the configuration values `n_max` and `m_max`.

"""
@inline function _get_degree_and_order(P, n_max, m_max)
    # Get the size of the matrix.
    (rows, cols) = size(P)

    # If the order or degree is less than 0, then the user wants to use all the
    # available memory.
    n_max < 0 && (n_max = rows - 1)
    m_max < 0 && (m_max = cols <= rows ? cols - 1 : n_max)

    # Make sure that the degree and order fits the matrix.
    n_max > rows - 1 && (n_max = rows - 1)
    ( (m_max > cols - 1) || (m_max > n_max) ) && (m_max = min(cols-1, n_max))

    return n_max, m_max
end


text:
using DataFrames
import CSV
using EcologicalNetworks
using StatsPlots
using LinearAlgebra
using Statistics

# Functions
include("lib.jl")

# Read clover
include("read_clover.jl")

# Get array parameters
virus_to_predict = parse(Int64, get(ENV, "SLURM_ARRAY_TASK_ID", "1"))
imputation_rank = 12

# Get them matrices
A = Array(clover.edges)

# Pools
filled = findall(A)
empty = findall(.!(A))

dfs = [DataFrame(value = Bool[], host = String[], virus = String[], host_pos = Int64[], virus_pos = Int64[], initial = Float64[], updated = Float64[]) for i in 1:Threads.nthreads()]

QUER = [Float64.(A) for i in 1:Threads.nthreads()]
TMPL = [n1(Q) for Q in QUER]

@time Threads.@threads for i in 1:size(clover, 2)
    Qu = QUER[Threads.threadid()]
    try
        pos = CartesianIndex(virus_to_predict, i)
        result = impute(Qu, TMPL[Threads.threadid()], pos; rank=12)
        push!(dfs[Threads.threadid()],
            (false, hosts[i], viruses[virus_to_predict], i, virus_to_predict, result.initial, result.updated)
        )
    catch
        @warn "Unable to do prediction for $(viruses[virus_to_predict]) and host $(hosts[i])"
        continue
    end
    CSV.write("svd-prediction-thread-$(Threads.threadid())-virus-$(virus_to_predict).csv", dfs[Threads.threadid()])
end

max_stars_repo_name: viralemergence/trefle

max_issues_repo_name: viralemergence/trefle

max_forks_repo_name: viralemergence/trefle


text:
using Test, Random, AutoGrad, Knet
include(AutoGrad.dir("test","gradcheck.jl"))


text:
"""
  xrayshoot(A::Array{Float64,2} , b::Array{Float64,1} , x0::Array{Float64,2} , v::Array{Float64,2})

returns: ( lambda_min , lambda_max , x_intersection_min , x_intersection_max )

# xrayshoot considers the system of linear inequalities A*x<b, and assumes
that all columns of x0 lie inside the solution set.
Then this function returns the most negative / most postitive lambdas
such that x0_i + lambda_i * v_i <= A*b holds.

# TEST:
 A1=[eye(2);-eye(2)]; b1=[1;1;1;1];
   lambda_min , lambda_max , x_intersection_min , x_intersection_max = XRayShoot(A1,b1,zeros(2,4),[[1;1],[1;0],[-1;1],[0.1;0.2]]);
   lambda_min , lambda_max , x_intersection_min , x_intersection_max = XRayShoot(A1,b1,zeros(2,1),[[1;1],[1;0],[-1;1],[0.1;0.2]]);
"""
function xrayshoot( A::Array{Float64,2} , b::Array{Float64,1} , x0::Array{Float64,2} , v::Array{Float64,2} )
    # if only one point, then replicate it once for every supplied
    # direction
    if ( (size(x0,2)==1) && (size(v,2)>1) )
        x0 = repmat(x0,1,size(v,2))
    end

    ## 1. compute intersections with constraints:
    AZ = A*v
    S::Array{Float64,2} = broadcast(-,b,A*x0)./AZ
    smax = copy(S)
    smax[AZ.<=0] = Inf
    smax = minimum(smax,1)
    lambda_max = smax;

    smin = S
    smin[AZ.>=0] = -Inf
    #print(smin)
    smin = maximum(smin,1)
    lambda_min = smin

    ## also return the intersecting points:
    #if(nargout >= 3),
    x_intersection_min = x0 + broadcast( * , lambda_min , v )
    x_intersection_max = x0 + broadcast( * , lambda_max , v )

    return lambda_min , lambda_max , x_intersection_min , x_intersection_max
end

mutable struct HRSamplingTrackingData
    X::Vector{Vector{Float64}}
    X_fdens::Vector{Float64}
    X_fcost::Vector{Float64}
    X_viable::Vector{Bool}
    X_opcode::Vector{Int64}

    fevals_expand::Int64
    fevals_bisect::Int64
    fevals_linesearch::Int64
    fevals_fdens::Int64
end
function HRSamplingTrackingData()
    return HRSamplingTrackingData( Vector{Vector{Float64}}() , Vector{Float64}() , Vector{Float64}() , Vector{Bool}() , Vector{Int64}() , 0 ,0 ,0 , 0)
end
function Base.hcat(a::HRSamplingTrackingData,b::HRSamplingTrackingData)
    return HRSamplingTrackingData( [a.X;b.X] , [a.X_fdens;b.X_fdens] , [a.X_fcost;b.X_fcost] , [a.X_viable;b.X_viable] , [a.X_opcode;b.X_opcode] , a.fevals_expand+b.fevals_expand , a.fevals_bisect+b.fevals_bisect , a.fevals_linesearch+b.fevals_linesearch , a.fevals_fdens+b.fevals_fdens )
end

struct HRStepConfig
    bisect_max_it::Int
    bisect_rel_tol::Float64
    bisect_abs_tol::Float64
    linesearch_max_tries::Int
    linedensity_fevals::Int

    track_chords::Int
    track_bisection::Int # 1: viable samples, 2: all samples
    track_linesearch::Int  # 1: viable samples, 2: all samples
end
function HRStepConfig( ; bisect_max_it::Int=Int64(-1) , bisect_rel_tol=NaN , bisect_abs_tol=Inf , linesearch_max_tries::Int=Int64(-1) , linedensity_fevals::Int=Int64(100) ,  track_chords::Int=Int64(0) , track_bisection::Int=Int64(0) , track_linesearch::Int=Int64(0) )
    HRStepConfig(bisect_max_it,bisect_rel_tol,bisect_abs_tol,linesearch_max_tries,linedensity_fevals,track_chords,track_bisection,track_linesearch)
end

function push_hrsample!(data::HRSamplingTrackingData,x::Vector{Float64},xfdens::Float64,xfcost::Float64,xviable::Bool,opcode::Int64)
    push!(data.X,x)
    push!(data.X_fdens,xfdens)
    push!(data.X_fcost,xfcost)
    push!(data.X_viable,xviable)
    push!(data.X_opcode,opcode)
end

"""
  bisection( f , threshold , a , b , max_it , max_rel_tol )

bisection performs bisection of the function of f on the line
from a to b such that the point is viable (i.e. f value is below threshold) and close to a. We assume that
f(b) is viable, f(a) most likely is not.

  f : function which returns as a first return value a
      boolean which describes if the point is viable.

  a,b : specify a line segment. b is the point which is viable, and we try
        to find the point closest to a which is still viable
  max_it : maximum iterations for moving the viable point close to a.
           NOTE: we may require more iterations than max_it if we have to
           move close to b to find a viable point.

# Output:
- x : the discovered boundary
- num_fcalls  : number of function evaluations
- x_viable    : tested points that are viable (without x).
- x_nonviable : tested points that are non-viable.

# TEST:
function
x_result1 = XBisection( x1 -> x1>7 , 5 , 10 , 10 );
x_result2 = XBisection( x1 -> x1>7 , 5 , 1e6 , 6 );
x_result3 = XBisection( x1 -> x1>(1e6-4) , 5 , 1e6 , 5 );

Thomas Liphardt, May 2018
"""
function bisection( f , threshold::Float64 , a::Array{Float64,1} , b::Array{Float64,1} , tdata::HRSamplingTrackingData ; bisect_max_it::Int=-1 , bisect_max_rel_tol::Float64=NaN , bisect_max_abs_tol=Inf , bisect_track_samples::Int64=Int64(0))
    #----- default parameters: --------------------------------
    if(bisect_max_it<0); bisect_max_it = 16; end
    if(isnan(bisect_max_rel_tol)); bisect_max_rel_tol = 1e-3; end
    #----- end default parameters: ----------------------------

    dim = size(a,1)

    #if(nargin<4), max_it      = 10; end
    #if(nargin<4), max_it      = 8; end
    #if(nargin<5), max_rel_tol = 1e-4; end

    n_it::Int64 = 1

    # We parametrize the line via x: (x=0) is a, (x=1) is b
    f_mix = (pa::Array{Float64,1},pb::Array{Float64,1},px::Float64) -> pa + (pb-pa)*px
    xa::Float64 = 0; xb::Float64 = 1

    #x_viable    = zeros(dim,0)
    #x_nonviable = zeros(dim,0)
    found_viable::Bool = false
    viable_fval::Float64   = NaN
    x_best_viable::Float64      = NaN
    x_fval_best_viable::Float64 = NaN

    x::Float64           = NaN

    norm_ab = norm(b-a)

    while( n_it<bisect_max_it && ( (xb-xa)>bisect_max_rel_tol  || (xb-xa)*norm_ab>bisect_max_abs_tol ) )
        #fprintf('\n%10.8f - %10.8f    --> x = %10.8f ',xa,xb,(xa+xb)/2);

        # bisect
        x = (xa+xb)/2.

        # check viability:
        tdata.fevals_bisect += 1
        viable_fval = f( f_mix(a,b,x) )
        viable = viable_fval < threshold

        if(viable)
            found_viable = true
            #x_viable(:,end+1) = f_mix(a,b,x)
            #x_viable = hcat( f_mix(a,b,x) , f_mix(a,b,x) )
            # if viable, go further left:
            x_best_viable = x
            x_fval_best_viable = viable_fval
            xb = x
            if(bisect_track_samples>=1); push_hrsample!(tdata,f_mix(a,b,x_best_viable),NaN,viable_fval,true,65); end
        else
            #x_nonviable(:,end+1) = f_mix(a,b,x)
            #x_nonviable = hcat( x_nonviable , f_mix(a,b,x) )
            # if not viable, go further right:
            xa = x
            if(bisect_track_samples>=2); push_hrsample!(tdata,f_mix(a,b,x),NaN,viable_fval,false,66); end
        end

        n_it = n_it+1
        if(n_it>1e3)
            if( ~f(b) )
                error("\n[ERROR] : XBisection : Point B was non-viable!\n")
            else
                @printf("[Problem] : XBisection : Bisection did not find other point..\nReturn original b..\n")
            end
        end
    end

    # if the last actual point is not viable, then we take the xb point.
    # xb is always viable, that's an assertion.
    x_result::Array{Float64,1} = b
    if(!found_viable)
    else
        x_result = f_mix(a,b,x_best_viable)
    end

    # remove the last point that we added to x_viable:
    #x_viable(:,end) = []
    #x_viable = x_viable(:,setdiff(1:size(x_viable,2),size(x_viable,2)))
    num_fcalls = n_it

    return (found_viable,x_result,x_fval_best_viable,num_fcalls) # :: Tuple{Bool,Vector{Float64},Float64,Integer64}

end

"""
  hr_step_unif!(A::Array{Float64,2} , b::Array{Float64,1} , x::Array{Float64,2} , nsteps::Int64)
hr steps for each column of x will be performed in place
random numbers are generated by the Julia GLOBAL_RNG
"""
function hr_step_unif!(A::Array{Float64,2} , b::Array{Float64,1} , x::Array{Float64,2} , nsteps::Int64)
    (d,n) = size(x)
    for zi=1:nsteps
        rv = randn(d,n)
        #rv = broadcast( / , rv , sum(rv.^2,1) )
        ru = rand(1, n)
        (lambda_min , lambda_max , x_intersection_min , x_intersection_max) = xrayshoot( A , b , x , rv )
        x[:,:] = x_intersection_min + broadcast( * , x_intersection_max-x_intersection_min , ru )
    end
    #return xi
end

function sample_unif_hr( G::Array{Float64,2} , h::Array{Float64,1} , x::Array{Float64,1} , nsamples::Int64 , nsteps::Int64 )
    x2 = repmat(x,1,nsamples)
    hr_step_unif!(G , h , x2 , nsteps)
    return x2
end

"""
  hr_step( f_density , f_cost , threshold::Float64 , G::Array{Float64,2} , h::Array{Float64,1} , x::Array{Float64,1} , x_candidate_a::Array{Float64,1}  , x_candidate_b::Array{Float64,1} , hrconf::HRStepConfig ,  tdata::HRSamplingTrackingData  )

# Returns: ( status , x_next , f_density(x_next) , f_cost(x_next) , n_fevals )
# status:  1: full_success , 2: success, but rbs_left failed, 3: success, but rbs_right failed
          -1: line_search failed , -2 both rbs failed
"""
function hr_step( f_density , f , threshold::Float64 , G::Array{Float64,2} , h::Array{Float64,1} , x::Array{Float64,1} , x_candidate_a::Array{Float64,1}  , x_candidate_b::Array{Float64,1} , hrconf::HRStepConfig ,  tdata::HRSamplingTrackingData  )
    #default parameter
    conf_linesearch_max_tries::Int64 = 20;
    conf_linedensity_fevals::Int64   = 100;
    if(hrconf.linesearch_max_tries>0); conf_linesearch_max_tries=hrconf.linesearch_max_tries; end
    if(hrconf.linedensity_fevals>0); conf_linedensity_fevals=hrconf.linedensity_fevals; end

    rbs_result_l = xrayshoot_and_bisect(f,threshold,G,h,x,x_candidate_a , hrconf , tdata )
    rbs_result_r = xrayshoot_and_bisect(f,threshold,G,h,x,x_candidate_b , hrconf , tdata )
    # ::(bool,x,f(x),num_fcalls)

    xa = rbs_result_l[2]
    xb = rbs_result_r[2]

    if(hrconf.track_chords>0)
        if(rbs_result_l[1]); push_hrsample!(tdata,xa,NaN,rbs_result_l[3],true,17);
        elseif(hrconf.track_chords>=2); push_hrsample!(tdata,xa,NaN,rbs_result_l[3],true,18) end
        if(rbs_result_r[1]); push_hrsample!(tdata,xb,NaN,rbs_result_r[3],true,17);
        elseif(hrconf.track_chords>=2); push_hrsample!(tdata,xb,NaN,rbs_result_r[3],true,18) end
    end

    if( !rbs_result_l[1] && !rbs_result_r[1] )
        print("Warning: bisection double fail..\n")
        # bisection double fail..
        return (-2,x,NaN,rbs_result_l[4]+rbs_result_r[4])
    end

    # line search
    ls_non_uniform = !isa(f_density,Void)
    if(!ls_non_uniform)
        # use linspace over range (we could also use random points..)
        ls_points = shuffle( linspace(1.e-8,1.-(1e-8),conf_linesearch_max_tries) )
        #ls_points = rand(linesearch_max_tries)
        ls_matrix = broadcast( + , xa , (xb-xa) * ls_points' )
    else
        # sample the 1d function:
        dist_border = 1.e-8
        fd_grid = linspace(0+dist_border,1-dist_border,conf_linedensity_fevals)
        fd_xx   = broadcast( + , xa , (xb-xa) * fd_grid' )
        fd_f    = NaN * ones(conf_linedensity_fevals)
        for zi in 1:conf_linedensity_fevals
            fd_f[zi] = f_density(fd_xx[:,zi])
        end
        tdata.fevals_fdens += conf_linedensity_fevals
        # use piecewise lin. function approx, and integrate
        fd_fi_pos = ( fd_grid[1:end-1]+fd_grid[2:end] ) ./ 2
        fd_fi     = ( fd_f[1:end-1]+fd_f[2:end] ) ./ 2
        fd_cumsum = cumsum(fd_fi)
        fd_cumsum = fd_cumsum ./ last(fd_cumsum)
        # draw conf_linesearch_max_tries from this distribution:
        r_cs_vals = rand(conf_linesearch_max_tries)

        ls_points = NaN * ones(conf_linesearch_max_tries)
        for zi=1:conf_linesearch_max_tries
            ls_points[zi] = fd_fi_pos[ findfirst( x -> x>r_cs_vals[zi] , fd_cumsum ) ]
        end
        ls_matrix = broadcast( + , xa , (xb-xa) * ls_points' )
    end
    x_next_fdens = NaN
    x_next_fcost = NaN
    x_next       = NaN * ones( length(x) )
    ls_success   = false
    zi           = 0
    for zi in 1:conf_linesearch_max_tries
        x_next_fcost = f(ls_matrix[:,zi])
        if( x_next_fcost < threshold )
            x_next       = ls_matrix[:,zi]
            if(ls_non_uniform); x_next_fdens = f_density(ls_matrix[:,zi]); end # note that we did not yet compute the exact fdens at this point.. thats why we compute it here

            ls_success = true
            break
        else
            if(hrconf.track_linesearch>=2); push_hrsample!(tdata,x_next,x_next_fdens,x_next_fcost,false,34) ;end
        end
    end
    tdata.fevals_linesearch += zi

    if(ls_success)
        success_code = 1
        if(!rbs_result_l[1]); success_code = 2;
        elseif(!rbs_result_r[1]); success_code = 3;
        end
        if(hrconf.track_linesearch>=1); push_hrsample!(tdata,x_next,x_next_fdens,x_next_fcost,true,33) ;end
        return (success_code,x_next,x_next_fdens,x_next_fcost,zi +rbs_result_l[4]+rbs_result_r[4])
    else
        return (-1,x,NaN,NaN,zi+rbs_result_l[4]+rbs_result_r[4])
    end
end

"""
xrayshoot_and_bisect(f , threshold , G::Array{Float64,2} , h::Array{Float64,1} , x::Array{Float64,1} , x_cand::Array{Float64,1} , hrconf::HRStepConfig , tdata::HRSamplingTrackingData )
shoots from x into direction of x_cand, then performs bisection to detect
point x_p with f(x_p) slightly below the threshold
"""
function xrayshoot_and_bisect(f , threshold , G::Array{Float64,2} , h::Array{Float64,1} , x::Array{Float64,1} , x_cand::Array{Float64,1} , hrconf::HRStepConfig , tdata::HRSamplingTrackingData )
    rv          = x_cand - x
    norm_rv     = norm(rv)
    rv_normed   = rv/norm_rv
    (lambda_min_a , lambda_max_a , x_intersection_min_a , x_intersection_max_a) = xrayshoot( G , h , x[:,:] , rv_normed[:,:] )

    lambda_max::Float64 = lambda_max_a[1]
    # check if candidate sample is outside of polytope
    if( lambda_max < norm_rv )
        x_cand = x + (1-1.e-6) * lambda_max * rv_normed
    end

    # check if candidate sample is viable, if so, then move further out
    fevals_expansion = 1
    xc_f        = f(x_cand)
    found_point = false
    reached_boundary = false # only set true if we reach it..
    if( xc_f < threshold )
        #print("expand..\n")
        # try to find viable point by taking points further away..
        xd  = x_cand-x
        nxd = norm(xd)
        multiplier = 1
        reached_boundary = false
        while( xc_f <= threshold && !reached_boundary )
            if( 2*multiplier*nxd < lambda_max)
                multiplier    =  2*multiplier
                x_cand = x + multiplier * nxd * xd
            else
                #print("reached max_expansion..\n")
                x_cand = x + (1-(1.e-6)) * lambda_max * (xd/nxd) # little bit inside
                reached_boundary = true
            end
            fevals_expansion += 1
            tdata.fevals_expand += 1
            xc_f = f(x_cand)
        end
    end
    if(xc_f <= threshold)
        #print("expansion yielded two viable points\n")
        if(!reached_boundary); print("Warning: max expansion did not reach boundary?! That's strange.."); end
        #return (false,NaN*ones(length(x)),NaN,fevals_expansion)
        # ok, then just accept these two points..
        return (true,x_cand,xc_f,fevals_expansion)
    end

    bisect_result = bisection( f , threshold ,  x_cand , x , tdata ; bisect_max_it=hrconf.bisect_max_it , bisect_max_rel_tol=hrconf.bisect_rel_tol , bisect_max_abs_tol=hrconf.bisect_abs_tol , bisect_track_samples = hrconf.track_bisection )
    # ::(bool,x,f(x),num_fcalls)
    return (bisect_result[1],bisect_result[2],bisect_result[3],bisect_result[4]+fevals_expansion)
end


text:
function save_sparse_matrix(fid, A::SparseMatrixCSC, name::String)

    n, m = size(A)
    Is, Js, Vs = findnz(A)

    fid["$(name)_Is"] = Int32[Is...]
    fid["$(name)_Js"] = Int32[Js...]
    fid["$(name)_Vs"] = Float64[Vs...]
    fid["$(name)_size"] = Int32[n, m]

    @info "Sparse Matrix '$(name)' saved to disk."
end

function load_sparse_matrix(fid, name::String)

    if isempty(fid["$(name)_Is"])
        Is = Int32[1]
        Js = Int32[1]
        Vs = Float64[0]
    else
        Is = fid["$(name)_Is"][:]
        Js = fid["$(name)_Js"][:]
        Vs = fid["$(name)_Vs"][:]
    end
    n, m = fid["$(name)_size"][:]

    @info "Loaded Sparse Matrix '$(name)'."

    return sparse(Is, Js, Vs, n, m)
end

function save_grid(fid, grd::Grid, name::String)

    # create group for grid
    create_group(fid, name)

    # save projection
    save_sparse_matrix(fid, grd.Proj, "$(name)/Proj")
    fid["$(name)/r0"] = grd.r0[:]
    fid["$(name)/r1"] = grd.r1[:]
    fid["$(name)/points"] = grd.points[:,:]
    fid["$(name)/spacing"] = Float64[grd.dx, grd.dy]
    fid["$(name)/_size"] = Int64[grd.Nk, grd.Nz, grd._Nx, grd._Ny, grd._Nbuffer]
    fid["$(name)/_NaNs"] = grd._nan_outside_boundaries[:,:]

    @info "Grid '$(name)' saved to disk."

end

function load_grid(fid, name::String)

    Proj = load_sparse_matrix(fid, "$(name)/Proj")
    r0 = fid["$(name)/r0"][:]
    r1 = fid["$(name)/r1"][:]
    points = fid["$(name)/points"][:,:]
    dx, dy = fid["$(name)/spacing"][:]
    Nk, Nz, Nx, Ny, Nbuffer = fid["$(name)/_size"][:]
    NaNs = fid["$(name)/_NaNs"][:,:]

    @info "Loaded Grid '$(name)'."

    return Grid(r0, r1, points, Proj, dx, dy, Nk, Nz, Nx, Ny, Nbuffer, NaNs)
end

function save_mirror(fid, M::Mirror, name::String)

    fid["$(name)/verts"] = M.verts[:,:]
    fid["$(name)/arms"] = M.arms[:,:,:]

end

function load_mirror(fid, name::String)

    verts = fid["$(name)/verts"][:,:]
    arms = fid["$(name)/arms"][:,:,:]

    return Mirror(verts, arms)
end

function save_ghost_data(fid, gd::GhostData, name::String)

    # create group
    create_group(fid, name)

    save_sparse_matrix(fid, gd.Proj, "$(name)/Proj")
    save_sparse_matrix(fid, gd.R, "$(name)/R")
    fid["$(name)/flip"] = gd.flip_factors[:,:]
end

function load_ghost_data(fid, name::String)

    Proj = load_sparse_matrix(fid, "$(name)/Proj")
    R = load_sparse_matrix(fid, "$(name)/R")
    flip_factors = fid["$(name)/flip"][:,:]

    return GhostData(Proj, R, flip_factors)
end

function boundary_setup(bdry_path::String, N1, N2, N3)

    count = 0

    # define the magnetic field
    psi, bx, by, bz = solovev_flux_function(0.1, 0.1, 0.3, 1.7, 10.0, downsep=[1, -0.561])
    Bp(r) = Float64[bx(r...), by(r...)]
    bp(r) = Bp(r) / norm(Bp(r))

    # values for flux surfaces
    psi_in = psi(1 + 0.28, 0)
    psi_out = psi(1 + 0.34, 0)
    psi_priv = psi(1, -(0.561 + 0.01))

    # find where the inner flux surface intersects y=0.1 on the outboard side
    bracket1 = [1.0, 0.1]
    bracket2 = [2.0, 0.1]
    r0 = rootfind_bisection(r->psi(r...), psi_in, bracket1, bracket2)

    # resolve inner boundary
    inner_chain_1, _ = trace_fieldline(r0, bp, r->psi(r...), 0.001) do r
        r[2] < 0
    end
    inner_chain_2, _ = trace_fieldline(inner_chain_1[:,end], bp, r->psi(r...), 0.001) do r
        r[2] > 0.1
    end
    skip = maximum([div(size(inner_chain_1)[2], N1), 1])
    inner_chain = hcat(inner_chain_1[:,1:skip:end-1], inner_chain_2[:,1:skip:end-1])
    Nin = size(inner_chain)[2]

    inner_range = 1:Nin
    count += Nin

    # find where outer flux surface intersects right diverter plate
    bracket1 = [1.0, -(0.561 + 0.04)]
    bracket2 = [2.0, -(0.561 + 0.04)]
    r0 = rootfind_bisection(r->psi(r...), psi_out, bracket1, bracket2)

    # resolve outer flux surface
    outer_chain, _ = trace_fieldline(r0, bp, r->psi(r...), 0.001) do r
        r[2] < -(0.561 + 0.04)
    end
    skip = maximum([div(size(outer_chain)[2], N2), 1])
    outer_chain = outer_chain[:,1:skip:end]
    Nout = size(outer_chain)[2]

    outer_range_1 = count+1:count+Nout
    count += Nout

    # find where outer flux surface intersects left diverter plate, add to outer_chain
    bracket1 = [0.0, -(0.561 + 0.04)]
    bracket2 = [1.0, -(0.561 + 0.04)]
    rl = rootfind_bisection(r->psi(r...), psi_out, bracket1, bracket2)
    outer_chain = hcat(outer_chain, rl)

    div_range_1 = count+1:count+1
    count += 1

    # find where private flux surface intersects left diverter plate
    bracket1 = [rl[1], -(0.561 + 0.04)]
    bracket2 = [1.0, -(0.561 + 0.04)]
    rl = rootfind_bisection(r->psi(r...), psi_priv, bracket1, bracket2)

    # trace private flux surface
    private_chain, _ = trace_fieldline(rl, r->-bp(r), r->psi(r...), 0.001) do r
        r[2] < -(0.561 + 0.04)
    end
    skip = maximum([div(size(private_chain)[2], N3), 1])
    private_chain = private_chain[:,1:skip:end-1]
    outer_chain = hcat(outer_chain, private_chain)
    Npriv = size(private_chain)[2]

    outer_range_2 = count+1:count+Npriv
    count += Npriv

    # find where private flux surface intersects right diverter plate
    bracket1 = [1.0, -(0.561 + 0.04)]
    bracket2 = [2.0, -(0.561 + 0.04)]
    rr = rootfind_bisection(r->psi(r...), psi_priv, bracket1, bracket2)
    outer_chain = hcat(outer_chain, rr)

    div_range_2 = count+1:count+1
    count += 1

    Nv = Nin + Nout + Npriv + 2

    inner_boundary = fill(false, Nv)
    inner_boundary[inner_range] .= true
    outer_boundary = fill(false, Nv)
    outer_boundary[outer_range_1] .= true
    outer_boundary[outer_range_2] .= true
    div_boundary = fill(false, Nv)
    div_boundary[div_range_1] .= true
    div_boundary[div_range_2] .= true

    M = Mirror([reverse(inner_chain, dims=2), outer_chain])

    fid = h5open(bdry_path, "w")
    save_mirror(fid, M, "Mirror")
    fid["inner_boundary"] = inner_boundary[:]
    fid["outer_boundary"] = outer_boundary[:]
    fid["div_boundary"] = div_boundary[:]
    close(fid)
end

function boundary_no_diverter(bdry_path::String, N1, N2, N3)

    count = 0

    # define the magnetic field
    psi, bx, by, bz = solovev_flux_function(0.1, 0.0, 0.3, 1.0, 10.0)
    Bp(r) = Float64[bx(r...), by(r...)]
    bp(r) = Bp(r) / norm(Bp(r))

    psi_in = psi(1 + 0.28, 0)
    psi_out = psi(1 + 0.34, 0)

    # find where the inner flux surface intersects y=0.1 on the outboard side
    bracket1 = [1.0, 0.1]
    bracket2 = [2.0, 0.1]
    r0 = rootfind_bisection(r->psi(r...), psi_in, bracket1, bracket2)

    # resolve inner boundary
    inner_chain_1, _ = trace_fieldline(r0, bp, r->psi(r...), 0.001) do r
        r[2] < 0
    end
    inner_chain_2, _ = trace_fieldline(inner_chain_1[:,end], bp, r->psi(r...), 0.001) do r
        r[2] > 0.1
    end
    skip = maximum([div(size(inner_chain_1)[2], N1), 1])
    inner_chain = hcat(inner_chain_1[:,1:skip:end-1], inner_chain_2[:,1:skip:end-1])
    Nin = size(inner_chain)[2]

    inner_range = 1:Nin
    count += Nin

    # find where the outer flux surface intersects y=0.1 on the outboard side
    bracket1 = [1.0, 0.1]
    bracket2 = [2.0, 0.1]
    r0 = rootfind_bisection(r->psi(r...), psi_out, bracket1, bracket2)

    # resolve outer boundary
    outer_chain_1, _ = trace_fieldline(r0, bp, r->psi(r...), 0.001) do r
        r[2] < 0
    end
    outer_chain_2, _ = trace_fieldline(outer_chain_1[:,end], bp, r->psi(r...), 0.001) do r
        r[2] > 0.1
    end
    skip = maximum([div(size(outer_chain_1)[2], N2), 1])
    outer_chain = hcat(outer_chain_1[:,1:skip:end-1], outer_chain_2[:,1:skip:end-1])
    Nout = size(outer_chain)[2]

    outer_range = count+1:count+Nout
    count += Nout

    Nv = Nin + Nout

    inner_boundary = fill(false, Nv)
    inner_boundary[inner_range] .= true
    outer_boundary = fill(false, Nv)
    outer_boundary[outer_range] .= true

    M = Mirror([reverse(inner_chain, dims=2), outer_chain])

    fid = h5open(bdry_path, "w")
    save_mirror(fid, M, "Mirror")
    fid["inner_boundary"] = inner_boundary[:]
    fid["outer_boundary"] = outer_boundary[:]
    close(fid)
end

function geometry_no_divertor(geo_path::String, bdry_path::String, Nx, Ny)

    fid = h5open(bdry_path, "r")
    M = load_mirror(fid, "Mirror")
    inner_boundary = fid["inner_boundary"][:]
    outer_boundary = fid["outer_boundary"][:]
    close(fid)

    grd = Grid(Nx, Ny, 1) do
        inside(x,y) = distance_to_mirror(x, y, M)[1] > -0.05
        r0 = Float64[0.5, -1]
        r1 = Float64[1.5, 1]
        inside, r0, r1
    end

    # compute penalization values as step functions
    pen = zeros(grd.Nk, 2)
    pen[:,1] = f_to_grid(grd) do x,y
        smoothstep(x, y, 0.02, 0.02, M[inner_boundary])
    end
    pen[:,2] = f_to_grid(grd) do x,y
        smoothstep(x, y, 0.02, 0.02, M[outer_boundary])
    end

    # K1 and K2 are operators for penalized forward-euler timesteps
    K1 = sparse(Diagonal(pen[:,1] .* pen[:,2]))
    K2 = 100 * (I - K1)

    gd_nmann = GhostData(M, grd)

    fid = h5open(geo_path, "w")
    save_grid(fid, grd, "Grid")
    save_sparse_matrix(fid, K1, "K1")
    save_sparse_matrix(fid, K2, "K2")
    save_ghost_data(fid, gd_nmann, "gd")
    close(fid)
end

function simple_geometry_setup(geo_path::String, bdry_path::String, Nx, Ny)

    fid = h5open(bdry_path, "r")
    M = load_mirror(fid, "Mirror")
    inner_boundary = fid["inner_boundary"][:]
    outer_boundary = fid["outer_boundary"][:]
    div_boundary = fid["div_boundary"][:]
    close(fid)

    grd = Grid(Nx, Ny, 1) do
        inside(x,y) = distance_to_mirror(x, y, M)[1] > -0.05
        r0 = Float64[0.5, -1]
        r1 = Float64[1.5, 1]
        inside, r0, r1
    end

    # compute penalization values as step functions
    pen = zeros(grd.Nk, 3)
    pen[:,1] = f_to_grid(grd) do x,y
        smoothstep(x, y, 0.02, 0.02, M[inner_boundary])
    end
    pen[:,2] = f_to_grid(grd) do x,y
        smoothstep(x, y, 0.02, 0.02, M[outer_boundary])
    end
    pen[:,3] = f_to_grid(grd) do x,y
        smoothstep(x, y, 0.02, 0.02, M[div_boundary])
    end

    # K1 and K2 are operators for penalized forward-euler timesteps
    K1 = sparse(Diagonal(pen[:,1] .* pen[:,2] .* pen[:,3]))
    K2 = 100 * (I - K1)

    gd_nmann = GhostData(M, grd)

    fid = h5open(geo_path, "w")
    save_grid(fid, grd, "Grid")
    save_sparse_matrix(fid, K1, "K1")
    save_sparse_matrix(fid, K2, "K2")
    save_ghost_data(fid, gd_nmann, "gd")
    close(fid)
end

function simple_physics_setup(w0, psi0, L_w, L_psi, init_path::String, geo_path::String)

    fid = h5open(geo_path, "r")
    grd = load_grid(fid, "Grid")
    gd = load_ghost_data(fid, "gd")
    close(fid)

    psi_func, bx, by, bz = solovev_flux_function(0.1, 0.1, 0.3, 1.7, 10.0, downsep=[1, -0.561])
    psi_mid = psi_func(0.685, 0)
    dp_w = abs(ForwardDiff.derivative(u -> psi_func(u,0), 0.65)) * L_w
    dp_psi = dp_w * L_psi / L_w

    w_vec = f_to_grid(grd) do x,y
        w0*tanh((psi_func(x,y) - psi_mid) / dp_w)
    end
    psi_vec = f_to_grid(grd) do x,y
        psi0*tanh((psi_func(x,y) - psi_mid) / dp_psi)
    end

    w = hcat(w_vec, w_vec, w_vec)
    psi = hcat(psi_vec, psi_vec, psi_vec)

    # compute derivatives
    Minv = stencil1d(3)
    Dx = derivative_matrix(1, 0, Minv, Minv, grd) * 0.3
    Dy = derivative_matrix(0, 1, Minv, Minv, grd) * 0.3
    Dxx = derivative_matrix(2, 0, Minv, Minv, grd) * (0.3)^2
    Dyy = derivative_matrix(0, 2, Minv, Minv, grd) * (0.3)^2

    # solve vorticity equation
    phi = solve_vorticity_eqn(zeros(grd.Nk), w[:,2], Dxx, Dyy, gd)

    fid = h5open(init_path, "w")
    save_sparse_matrix(fid, Dx, "Dx")
    save_sparse_matrix(fid, Dy, "Dy")
    save_sparse_matrix(fid, Dxx, "Dxx")
    save_sparse_matrix(fid, Dyy, "Dyy")
    fid["w"] = w[:,:]
    fid["phi"] = phi[:]
    fid["psi"] = psi[:,:]
    close(fid)
end

function simulate(Nt, k_w, k_psi, am, dt, output_path, init_path, geo_path)

    fid = h5open(geo_path, "r")
    grd = load_grid(fid, "Grid")
    K1 = load_sparse_matrix(fid, "K1")
    K2 = load_sparse_matrix(fid, "K2")
    gd = load_ghost_data(fid, "gd")
    close(fid)

    fid = h5open(init_path, "r")
    Dx = load_sparse_matrix(fid, "Dx")
    Dy = load_sparse_matrix(fid, "Dy")
    Dxx = load_sparse_matrix(fid, "Dxx")
    Dyy = load_sparse_matrix(fid, "Dyy")
    w = fid["w"][:,:]
    phi = fid["phi"][:]
    psi = fid["psi"][:,:]
    close(fid)

    fid = h5open(output_path, "w")
    save_grid(fid, grd, "Grid")
    create_dataset(fid, "w", datatype(Float64), dataspace(grd.Nk*grd.Nz,Nt), chunk=(grd.Nk*grd.Nz,1))
    create_dataset(fid, "phi", datatype(Float64), dataspace(grd.Nk*grd.Nz,Nt), chunk=(grd.Nk*grd.Nz,1))
    create_dataset(fid, "psi", datatype(Float64), dataspace(grd.Nk*grd.Nz,Nt), chunk=(grd.Nk*grd.Nz,1))
    fid["w"][:,1] = w[:,2]
    fid["phi"][:,1] = phi[:]
    fid["psi"][:,1] = psi[:,2]
    close(fid)

    Nskip = 1
    p = Progress(Nskip*(Nt-1), "Running simple model... ")
    for t=2:Nt
        for _=1:Nskip
            leapfrog!(w, psi, phi, Dx, Dxx, Dy, Dyy, K1, K2, gd, dt, am, k_psi, k_w)
            next!(p)
        end

        fid = h5open(output_path, "r+")
        fid["w"][:,t] = w[:,2]
        fid["phi"][:,t] = phi[:]
        fid["psi"][:,t] = psi[:,2]
        close(fid)
    end

end

#
#
#
# function simple_geometry_setup(path::String, Nx, Ny)
#
#     # define the magnetic field
#     psi, bx, by, bz = solovev_flux_function(0.1, 0.1, 0.3, 1.7, 10.0, downsep=[1, -0.561])
#
#     # values for flux surfaces
#     psi_in = psi(1 - 0.28, 0)
#     psi_out = psi(1 - 0.34, 0)
#     psi_priv = psi(1, -(0.561 + 0.01))
#
#     # define barriers for flux surfaces
#     inner_flux = Barrier() do
#         func(x,y) = (y > -0.561) ? psi(x,y) : abs(psi(x,y))
#         rmap(x,y) = trace_reflection(x, y, psi, psi_in, 0.001)
#         func, psi_in, rmap, 1
#     end
#
#     outer_flux = Barrier() do
#         func(x,y) = (norm([x-1,y+0.6]) > 0.05) ? psi(x,y) : psi_out + psi_priv - psi(x,y)
#         rmap(x,y) = trace_reflection(x, y, func, psi_out, 0.001)
#         func, psi_out, rmap, -1
#     end
#
#     # define barrier for divertor target plates
#     target = Barrier() do
#         func(x,y) = y
#         val = -(0.561 + 0.04)
#         rmap(x,y) = x, 2*val - y
#         func, val, rmap, 1
#     end
#
#     # compute step function widths
#     dr = 0.01
#     dp_in = abs(ForwardDiff.derivative(u -> psi(u,0), 0.7)) * dr
#     dp_out = abs(ForwardDiff.derivative(u -> psi(u,0), 0.6)) * dr
#     dp_priv = abs(ForwardDiff.derivative(u -> psi(1,u), -(0.561 + 0.05))) * dr
#
#     # create Grid
#     deltas = Float64[dp_in, 0.5*(dp_out + dp_priv), dr]
#     barrs = [inner_flux, outer_flux, target]
#     grd = Grid(Nx, Ny, 1) do
#         inside(x,y) = check_if_inside(x, y, 10*deltas, barrs)
#         r0 = Float64[0.5, -1]
#         r1 = Float64[1.5, 1]
#         inside, r0, r1
#     end
#
#     # compute penalization values as step functions
#     pen = zeros(grd.Nk, 3)
#     for i=1:3
#         pen[:,i] = f_to_grid(grd) do x,y
#             smoothstep(x, y, deltas[i], barrs[i])
#         end
#     end
#
#     # K1 and K2 are operators for penalized forward-euler timesteps
#     K1 = sparse(Diagonal(pen[:,1] .* pen[:,2] .* pen[:,3]))
#     K2 = 100 * (I - K1)
#
#     gc_nmann = GhostConditions(barrs, grd)
#     gc = swap_sign(gc_nmann, 1, 2, 3)
#
#     fid = h5open(path, "w")
#     save_grid(fid, grd, "Grid")
#     save_sparse_matrix(fid, K1, "K1")
#     save_sparse_matrix(fid, K2, "K2")
#     save_ghost_conditions(fid, gc_nmann, "gc")
#     close(fid)
# end
#
#
# function simple_physics_setup(w0, psi0, L_w, L_psi, init_path::String, geo_path::String)
#
#     fid = h5open(geo_path, "r")
#     grd = load_grid(fid, "Grid")
#     gc = load_ghost_conditions(fid, "gc")
#     close(fid)
#
#     psi_func, bx, by, bz = solovev_flux_function(0.1, 0.1, 0.3, 1.7, 10.0, downsep=[1, -0.561])
#     psi_mid = psi_func(0.65, 0)
#     dp_w = abs(ForwardDiff.derivative(u -> psi_func(u,0), 0.65)) * L_w
#     dp_psi = dp_w * L_psi / L_w
#
#     w_vec = f_to_grid(grd) do x,y
#         w0*tanh((psi_func(x,y) - psi_mid) / dp_w)
#     end
#     psi_vec = f_to_grid(grd) do x,y
#         psi0*tanh((psi_func(x,y) - psi_mid) / dp_psi)
#     end
#
#     w = hcat(w_vec, w_vec, w_vec)
#     psi = hcat(psi_vec, psi_vec, psi_vec)
#
#     # compute derivatives
#     Minv = stencil1d(3)
#     Dx = derivative_matrix(1, 0, Minv, Minv, grd) * 0.3
#     Dy = derivative_matrix(0, 1, Minv, Minv, grd) * 0.3
#     Dxx = derivative_matrix(2, 0, Minv, Minv, grd) * (0.3)^2
#     Dyy = derivative_matrix(0, 2, Minv, Minv, grd) * (0.3)^2
#
#     # solve vorticity equation
#     phi = solve_vorticity_eqn(zeros(grd.Nk), w[:,2], Dxx, Dyy, gc)
#
#     fid = h5open(init_path, "w")
#     save_sparse_matrix(fid, Dx, "Dx")
#     save_sparse_matrix(fid, Dy, "Dy")
#     save_sparse_matrix(fid, Dxx, "Dxx")
#     save_sparse_matrix(fid, Dyy, "Dyy")
#     fid["w"] = w[:,:]
#     fid["phi"] = phi[:]
#     fid["psi"] = psi[:,:]
#     close(fid)
# end


text:
"""
    struct Sparsity.Preorder{S <: Sparsity.Pattern, C <: SumOfSquares.Certificate.AbstractPreorderCertificate} <: SumOfSquares.Certificate.AbstractPreorderCertificate
        sparsity::S
        certificate::C
    end

Same certificate as `C` except that the Sum-of-Squares polynomials `σ_i`
are modelled as a sum of Sum-of-Squares polynomials with smaller basis
using the sparsity reduction `sparsity`.
"""
struct Preorder{S <: Pattern, C <: SumOfSquares.Certificate.AbstractPreorderCertificate} <: SumOfSquares.Certificate.AbstractPreorderCertificate
    sparsity::S
    certificate::C
end

SumOfSquares.Certificate.cone(certificate::Preorder) = SumOfSquares.Certificate.cone(certificate.certificate)

struct Domain{S, P, B}
    domain::S
    processed::P
    bases::Vector{Vector{B}}
end

function SumOfSquares.Certificate.preprocessed_domain(certificate::Preorder, domain::SemialgebraicSets.BasicSemialgebraicSet, p)
    basis, Preorder_bases = sparsity(p, domain, certificate.sparsity, certificate.certificate)
    return Domain(domain, SumOfSquares.Certificate.preprocessed_domain(certificate.certificate, domain, p), Preorder_bases)
end

function SumOfSquares.Certificate.preorder_indices(certificate::Preorder, domain::Domain)
    return SumOfSquares.Certificate.preorder_indices(certificate.certificate, domain.processed)
end

function SumOfSquares.Certificate.multiplier_basis(::Preorder, index::SumOfSquares.Certificate.PreorderIndex, domain::Domain)
    return domain.bases[index.value]
end
function SumOfSquares.Certificate.multiplier_basis_type(::Type{Preorder{S, C}}) where {S, C}
    return Vector{SumOfSquares.Certificate.multiplier_basis_type(C)}
end

function SumOfSquares.Certificate.generator(certificate::Preorder, index::SumOfSquares.Certificate.PreorderIndex, domain::Domain)
    return SumOfSquares.Certificate.generator(certificate.certificate, index, domain.processed)
end

SumOfSquares.Certificate.ideal_certificate(certificate::Preorder) = Ideal(certificate.sparsity, SumOfSquares.Certificate.ideal_certificate(certificate.certificate))
SumOfSquares.Certificate.ideal_certificate(::Type{<:Preorder{S, C}}) where {S, C} = Ideal{S, SumOfSquares.Certificate.ideal_certificate(C)}

SumOfSquares.matrix_cone_type(::Type{Preorder{S, C}}) where {S, C} = SumOfSquares.matrix_cone_type(C)

max_stars_repo_name: blegat/SumOfSquares

max_issues_repo_name: blegat/SumOfSquares

max_forks_repo_name: blegat/SumOfSquares