text:
import Zygote

function linearIndex(nz, iz, ix)
    return (ix-1)*nz + iz
end

@NonlinearEquations.equations exclude=(nz, nx) function wave2diter(c, f, u, uim1, uim2, nz, nx, nt, dz, dx, dt) # 1st order
    NonlinearEquations.setnumequations(nz*nx)

    ## TOP ROW #########################
    # top row (x, not corners)
    iz = 1
    for ix=2:nx-1
        v1 = c[linearIndex(nz, 1, ix)] * dt/dz
        v2 = c[linearIndex(nz, 1, ix)] * dt/dx
        diff1 = v1*(-(uim1[linearIndex(nz, 2, ix)] - uim1[linearIndex(nz, 1, ix)] - uim2[linearIndex(nz, 2, ix)] - uim2[linearIndex(nz, 1, ix)]))
        diff2 = 0.5*v2^2*(uim1[linearIndex(nz, 1, ix-1)] -2*uim1[linearIndex(nz, 1, ix)] + uim1[linearIndex(nz, 1, ix+1)])
        NonlinearEquations.addterm((ix-1)*nz+iz, u[linearIndex(nz, 1, ix)] - 2*uim1[linearIndex(nz, 1, ix)] -diff1 - diff2)

    end

    ## SECOND FROM TOP ROW #########################
    # second to top row (x, not corners)
    iz = 2
    for ix=3:nx-2
        NonlinearEquations.addterm((ix-1)*nz + iz, -1*(-uim1[linearIndex(nz, 2, ix-2)] +16*uim1[linearIndex(nz, 2, ix-1)] + -30*uim1[linearIndex(nz, 2, ix)] +16*uim1[linearIndex(nz, 2, ix+1)] - uim1[linearIndex(nz, 2, ix+2)])/dx^2/12*c[linearIndex(nz, 2, ix)]^2)
    end
    # second from top row (z)
    for ix=1:nx
        NonlinearEquations.addterm((ix-1)*nz + iz, -1*(uim1[linearIndex(nz, 1, ix)] - 2*uim1[linearIndex(nz, 2, ix)] + uim1[linearIndex(nz, 3, ix)])/dz^2*c[linearIndex(nz, 2, ix)]^2)
    end

    ## BOTTOM ROW #########################

    for ix=2:nx-1
        v1 = c[linearIndex(nz, nz, ix)] * dt/dz
        v2 = c[linearIndex(nz, nz, ix)] * dt/dx
        diff1 = v1*(-(uim1[linearIndex(nz, nz, ix)] - uim1[linearIndex(nz, nz-1, ix)] - uim2[linearIndex(nz, nz, ix)] - uim2[linearIndex(nz, nz-1, ix)]))
        diff2 = 0.5*v2^2*(uim1[linearIndex(nz, nz, ix-1)] -2*uim1[linearIndex(nz, nz, ix)] + uim1[linearIndex(nz, nz, ix+1)])
        NonlinearEquations.addterm((ix-1)*nz+iz, u[linearIndex(nz, nz, ix)] - 2*uim1[linearIndex(nz, nz, ix)] -diff1 - diff2)

    end

    ## SECOND FROM BOTTOM ROW #########################
    # second to bottom row (x, not corners)
    iz = nz - 1
    for ix=3:nx-2
        NonlinearEquations.addterm((ix-1)*nz + iz, -1*(-uim1[linearIndex(nz, nz-1, ix-2)] +16*uim1[linearIndex(nz, nz-1, ix-1)] + -30*uim1[linearIndex(nz, nz-1, ix)] +16*uim1[linearIndex(nz, nz-1, ix+1)] - uim1[linearIndex(nz, nz-1, ix+2)])/dx^2/12*c[linearIndex(nz, nz-1, ix)]^2)
    end
    # second from bottom row (z)
    for ix=1:nx
        NonlinearEquations.addterm((ix-1)*nz + iz, -1*(uim1[linearIndex(nz, nz, ix)] - 2*uim1[linearIndex(nz, nz-1, ix)] + uim1[linearIndex(nz, nz-2, ix)])/dz^2*c[linearIndex(nz, nz-1, ix)]^2)
    end

    # LEFT COLUMN ####################
    ix = 1
    for iz=2:nz-1
        v1 = c[linearIndex(nz, iz, 1)] * dt/dz
        v2 = c[linearIndex(nz, iz, 1)] * dt/dx
        diff1 = 0.5*v1^2*(uim1[linearIndex(nz, iz-1, 1)] -2*uim1[linearIndex(nz, iz, 1)] + uim1[linearIndex(nz, iz+1, 1)])
        diff2 = v2*((uim1[linearIndex(nz, iz, 2)] - uim1[linearIndex(nz, iz, 1)] - uim2[linearIndex(nz, iz, 2)] - uim2[linearIndex(nz, iz, 1)]))
        NonlinearEquations.addterm((ix-1)*nz+iz, u[linearIndex(nz, iz, 1)] - 2*uim1[linearIndex(nz, iz, 1)] -diff1 - diff2)

    end

    # SECOND FROM LEFT COLUMN ####################
    # second to left edge (x)
    ix = 2
    for iz=1:nz
        NonlinearEquations.addterm((ix-1)*nz + iz, -1*(uim1[linearIndex(nz, iz, 1)] - 2*uim1[linearIndex(nz, iz, 2)] + uim1[linearIndex(nz, iz, 3)])/dx^2*c[linearIndex(nz, iz, 2)]^2)
    end

    # second to left edge (z, not corners)
    for iz=3:nz-2
        NonlinearEquations.addterm((ix-1)*nz + iz, -1*(-uim1[linearIndex(nz, iz-2, 2)] +16*uim1[linearIndex(nz, iz-1, 2)] + -30*uim1[linearIndex(nz, iz, 2)] +16*uim1[linearIndex(nz, iz+1, 2)] - uim1[linearIndex(nz, iz+2, 2)])/dz^2/12*c[linearIndex(nz, iz, 2)]^2)
    end

    # RIGHT COLUMN ####################
    ix = nx
    for iz=2:nz-1
        v1 = c[linearIndex(nz, iz, nx)] * dt/dz
        v2 = c[linearIndex(nz, iz, nx)] * dt/dx
        diff1 = 0.5*v1^2*(uim1[linearIndex(nz, iz-1, nx)] -2*uim1[linearIndex(nz, iz, nx)] + uim1[linearIndex(nz, iz+1, nx)])
        diff2 = v2*(-(uim1[linearIndex(nz, iz, nx)] - uim1[linearIndex(nz, iz, nx-1)] - uim2[linearIndex(nz, iz, nx)] - uim2[linearIndex(nz, iz, nx-1)]))
        NonlinearEquations.addterm((ix-1)*nz+iz, u[linearIndex(nz, iz, nx)] - 2*uim1[linearIndex(nz, iz, nx)] -diff1 - diff2)

    end

    # SECOND FROM RIGHT COLUMN ####################
    # second to right edge (x)
    ix = nx - 1
    for iz=1:nz
        NonlinearEquations.addterm((ix-1)*nz + iz, -1*(uim1[linearIndex(nz, iz, nx)] - 2*uim1[linearIndex(nz, iz, nx-1)] + uim1[linearIndex(nz, iz, nx-2)])/dx^2*c[linearIndex(nz, iz, nx-1)]^2)
    end

    # second to right edge (z, not corners)
    for iz=3:nz-2
        NonlinearEquations.addterm((ix-1)*nz + iz, -1*(-uim1[linearIndex(nz, iz-2, nx-1)] +16*uim1[linearIndex(nz, iz-1, nx-1)] + -30*uim1[linearIndex(nz, iz, nx-1)] +16*uim1[linearIndex(nz, iz+1, nx-1)] - uim1[linearIndex(nz, iz+2, nx-1)])/dz^2/12*c[linearIndex(nz, iz, nx-1)]^2)
    end

    # CENTER ####################
    for iz=3:nz-2
        for ix=3:nx-2
            NonlinearEquations.addterm((ix-1)*nz + iz, -1*(-uim1[linearIndex(nz, iz-2, ix)] +16*uim1[linearIndex(nz, iz-1, ix)] + -30*uim1[linearIndex(nz, iz, ix)] +16*uim1[linearIndex(nz, iz+1, ix)] - uim1[linearIndex(nz, iz+2, ix)])/(dz^2*12)*c[linearIndex(nz, iz, ix)]^2) # z
            NonlinearEquations.addterm((ix-1)*nz + iz, -1*(-uim1[linearIndex(nz, iz, ix-2)] +16*uim1[linearIndex(nz, iz, ix-1)] + -30*uim1[linearIndex(nz, iz, ix)] +16*uim1[linearIndex(nz, iz, ix+1)] - uim1[linearIndex(nz, iz, ix+2)])/(dx^2*12)*c[linearIndex(nz, iz, ix)]^2) # x
        end
    end

    ################################################################################
    # time derivative and source ####################
    for iz=1:nz
        for ix=1:nx
            # forcing term
            NonlinearEquations.addterm((ix-1)*nz + iz, f[linearIndex(nz, iz, ix)])

            # 1st order time derivative
            NonlinearEquations.addterm((ix-1)*nz + iz, (u[linearIndex(nz, iz, ix)] - 2*uim1[linearIndex(nz, iz, ix)] + uim2[linearIndex(nz, iz, ix)])/dt^2 )
        end
    end
end

## WAVE PROPAGATION ####################
function uoneiter(ui, uim1, uim2, c, f, nz, nx, nt, dz, dx, dt) # 1st order, one iteration of wave propagation
    return ui - dt^2 * wave2diter_residuals(c, f, ui, uim1, uim2, nz, nx, nt, dz, dx, dt)
end

# multiple iterations
function getuIters(cval, f, nz, nx, nt, dz, dx, dt)
    u = zeros(nz*nx, nt + 2)
    for i=3:nt + 2
        u[:, i] = uoneiter(zeros(size(u, 1)), u[:, i - 1], u[:, i - 2], cval, f[:, i - 2], nz, nx, nt, dz, dx, dt)
        #if mod(i, 10) == 0
        #   @printf("Iteration %i\n", i)
        #end
    end
    return u
end

function ChainRulesCore.rrule(::typeof(getuIters), c, f, nz, nx, nt, dz, dx, dt)
    u = getuIters(c, f, nz, nx, nt, dz, dx, dt)
    return u, delta->begin
        ds = []
        #last iteration
        _, pb = Zygote.pullback(uoneiter, zeros(size(u, 1)), u[:, size(u, 2) - 1], u[:, size(u, 2) - 2], c, f[:, end], nz, nx, nt, dz, dx, dt)
        d1, d2, d3, d4, d5 = pb(delta[:, size(u, 2)])
        push!(ds, (d1, d2, d3))
        dc = d4
        df = hcat(d5)
        #second to last iteration
        _, pb = Zygote.pullback(uoneiter, zeros(size(u, 1)), u[:, size(u, 2) - 2], u[:, size(u, 2) - 3], c, f[:, end - 1], nz, nx, nt, dz, dx, dt)
        d1, d2, d3, d4, d5 = pb(delta[:, size(u, 2) - 1] + ds[1][2])
        push!(ds, (d1, d2, d3))
        dc += d4
        df = hcat(d5, df)
        #now do the general iteration
        for i = size(u, 2) - 2:-1:3
            _, pb = Zygote.pullback(uoneiter, zeros(size(u, 1)), u[:, i - 1], u[:, i - 2], c, f[:, i - 2], nz, nx, nt, dz, dx, dt)
            d1, d2, d3, d4, d5 = pb(delta[:, i] + ds[2][2] + ds[1][3])
            dc += d4
            df = hcat(d5, df)
            push!(ds, (d1, d2, d3))
            ds = ds[2:end]
        end
        return (ChainRulesCore.NO_FIELDS,#getuIters
                dc,
                df,
                ChainRulesCore.NO_FIELDS,#nz
                ChainRulesCore.NO_FIELDS,#nx
                ChainRulesCore.NO_FIELDS,#nt
                ChainRulesCore.NO_FIELDS,#dz
                ChainRulesCore.NO_FIELDS,#dx
                ChainRulesCore.NO_FIELDS)#dt
    end
end

## PULLBACK FUNCTIONS ####################
function make_wave2diter_pullback(args...)
    function wave2diter_pullback(delta)
		retval = (ChainRulesCore.NO_FIELDS,#function
                @ChainRulesCore.thunk(wave2diter_c(args...)' * delta),#c
				@ChainRulesCore.thunk(wave2diter_f(args...)' * delta),#f
				@ChainRulesCore.thunk(wave2diter_u(args...)' * delta),#u
				@ChainRulesCore.thunk(wave2diter_uim1(args...)' * delta),#uim1
				@ChainRulesCore.thunk(wave2diter_uim2(args...)' * delta),#uim2
                ChainRulesCore.NO_FIELDS,#nz
                ChainRulesCore.NO_FIELDS,#nx
                ChainRulesCore.NO_FIELDS,#nt
                ChainRulesCore.NO_FIELDS,#dz
                ChainRulesCore.NO_FIELDS,#dx
                ChainRulesCore.NO_FIELDS)#dt
        return retval
    end
    return wave2diter_pullback
end

function ChainRulesCore.rrule(::typeof(wave2diter_residuals), c, f, u, uim1, uim2, nz, nx, nt, dz, dx, dt)
    args = (c, f, u, uim1, uim2, nz, nx, nt, dz, dx, dt)
    residuals = wave2diter_residuals(args...)
    pullback = make_wave2diter_pullback(args...)
    return residuals, pullback
end


text:
"""
    function r2star_ll(
        mag::AbstractArray{<:AbstractFloat, N > 1},
        TEs::NTuple{NT > 1, Real},
        mask::Union{Nothing, AbstractArray{Bool}} = nothing
    ) -> typeof(similar(mag, size(mag)[1:N-1]))

Log-linear fit.

### Arguments
- `mag::AbstractArray{<:AbstractFloat, N > 1}`: multi-echo magnitude
- `TEs::NTuple{NT > 1, Real}`: echo times
- `mask::Union{Nothing, AbstractArray{Bool}} = nothing`: binary mask of region of interest

### Returns
- `typeof(similar(mag, size(mag)[1:N-1]))`: R2* map (1 / units of TEs)
"""
function r2star_ll(
    mag::AbstractArray{T, N},
    TEs::NTuple{NT, Real},
    mask::Union{Nothing, AbstractArray{Bool}} = nothing
) where {T<:AbstractFloat, N, NT}
    N > 1 || throw(ArgumentError("array must contain echoes in last dimension"))
    NT > 1 || throw(ArgumentError("data must be multi-echo"))

    size(mag, N) == NT || throw(DimensionMismatch())
    mask === nothing || length(mask) == length(mag) ÷ NT || throw(DimensionMismatch())

    r2s = tzero(mag, size(mag)[1:N-1])

    vmag = reshape(mag, :, NT)
    vr2s = vec(r2s)

    A = qr(Matrix{T}([-[TEs...] ones(NT)]))

    if mask === nothing
        b = tcopy(log, transpose(vmag))

        x̂ = ldiv!(A, b)

        @inbounds @batch for I in eachindex(vr2s)
            vr2s[I] = x̂[1,I]
        end

    else
        vmask = vec(mask)

        i = 0
        R = Vector{Int}(undef, sum(vmask))
        @inbounds for I in eachindex(vmask)
            if vmask[I]
                R[i += 1] = I
            end
        end

        b = similar(mag, (NT, length(R)))
        @inbounds for t in Base.OneTo(NT)
            bt = @view(b[t,:])
            mt = @view(vmag[:,t])
            @batch for I in eachindex(R)
                bt[I] = log(mt[R[I]])
            end
        end

        x̂ = ldiv!(A, b)

        @inbounds @batch for I in eachindex(R)
            vr2s[R[I]] = x̂[1,I]
        end
    end

    return r2s
end

"""
    function r2star_arlo(
        mag::AbstractArray{<:AbstractFloat, N > 1},
        TEs::NTuple{NT > 1, Real},
        mask::Union{Nothing, AbstractArray{Bool}} = nothing
    ) -> typeof(similar(mag, size(mag)[1:N-1]))

Auto-Regression on Linear Operations (ARLO) [1].

### Arguments
- `mag::AbstractArray{<:AbstractFloat, N > 1}`: multi-echo magnitude
- `TEs::NTuple{NT > 1, Real}`: echo times
- `mask::Union{Nothing, AbstractArray{Bool}} = nothing`: binary mask of region of interest

### Returns
- `typeof(similar(mag, size(mag)[1:N-1]))`: R2* map (1 / units of TEs)

### References
[1] Pei M, Nguyen TD, Thimmappa ND, Salustri C, Dong F, Cooper MA, Li J,
    Prince MR, Wang Y. Algorithm for fast monoexponential fitting based on
    auto‐regression on linear operations (ARLO) of data.
    Magnetic resonance in medicine. 2015 Feb;73(2):843-50.
"""
function r2star_arlo(
    mag::AbstractArray{T, N},
    TEs::NTuple{NT, Real},
    mask::Union{Nothing, AbstractArray{Bool}} = nothing
) where {T<:AbstractFloat, N, NT}
    N > 1 || throw(ArgumentError("array must contain echoes in last dimension"))
    NT > 2 || throw(ArgumentError("ARLO requires at least 3 echoes"))

    size(mag, N) == NT || throw(DimensionMismatch())
    mask === nothing || length(mask) == length(mag) ÷ NT || throw(DimensionMismatch())

    all((≈)(TEs[2]-TEs[1]), TEs[2:end].-TEs[1:end-1]) ||
        throw(DomainError("ARLO requires equidistant echoes"))

    r2s = tzero(mag, size(mag)[1:N-1])

    vmag = reshape(mag, :, NT)
    vr2s = vec(r2s)

    _zeroT = zero(T)
    _twoT  = convert(T, 2)
    _fourT = convert(T, 4)

    α = convert(T, 3 / (TEs[2]-TEs[1]))

    if mask === nothing
        @inbounds @batch for I in eachindex(vr2s)
            m0 = vmag[I,1]
            m1 = vmag[I,2]
            m2 = vmag[I,3]

            δ = m0 - m2
            s = m0 + muladd(_fourT, m1, m2)
            a = muladd(_twoT, m1, m0)

            num = δ * a
            den = s * a

            for t in 2:NT-2
                m0 = m1
                m1 = m2
                m2 = vmag[I,t+2]

                δ = m0 - m2
                s = m0 + muladd(_fourT, m1, m2)
                a = muladd(_twoT, m1, m0)

                num = muladd(δ, a, num)
                den = muladd(s, a, den)
            end

            vr2s[I] = iszero(den) ? _zeroT : α * num * inv(den)
        end

    else
        vmask = vec(mask)

        @inbounds @batch for I in eachindex(vr2s)
            if vmask[I]
                m0 = vmag[I,1]
                m1 = vmag[I,2]
                m2 = vmag[I,3]

                δ = m0 - m2
                s = m0 + muladd(_fourT, m1, m2)
                a = muladd(_twoT, m1, m0)

                num = δ * a
                den = s * a

                for t in 2:NT-2
                    m0 = m1
                    m1 = m2
                    m2 = vmag[I,t+2]

                    δ = m0 - m2
                    s = m0 + muladd(_fourT, m1, m2)
                    a = muladd(_twoT, m1, m0)

                    num = muladd(δ, a, num)
                    den = muladd(s, a, den)
                end

                vr2s[I] = iszero(den) ? _zeroT : α * num * inv(den)
            end
        end
    end

    return r2s
end

"""
    function r2star_crsi(
        mag::AbstractArray{<:AbstractFloat, N > 1},
        TEs::NTuple{NT > 1, Real},
        mask::Union{Nothing, AbstractArray{Bool}} = nothing;
        M::Integer = 3,
        sigma::Union{Nothing, Real} = nothing,
        Rsz::NTuple{N-1, Integer} = size(mag)[1:N-1] .÷ 20,
    ) -> typeof(similar(mag, size(mag)[1:N-1]))

Calculation of Relaxivities by Signal Integration (CRSI) [1].

### Arguments
- `mag::AbstractArray{<:AbstractFloat, N > 1}`: multi-echo magnitude
- `TEs::NTuple{NT > 1, Real}`: echo times
- `mask::Union{Nothing, AbstractArray{Bool}} = nothing`: binary mask of region of interest

### Keywords
- `M::Integer = 3`: interpolation factor
- `sigma::Union{Nothing, Real} = nothing`: noise
- `Rsz::NTuple{N-1, Integer} = size(mag)[1:N-1] .÷ 20`:
    - `sigma isa Real`: unused
    - `sigma isa Nothing`: size of kernels used to calculate the noise from the
        background signal of the magnitude.

### Returns
- `typeof(similar(mag, size(mag)[1:N-1]))`: R2* map (1 / units of TEs)

### References
[1] Song R, Loeffler RB, Holtrop JL, McCarville MB, Hankins JS, Hillenbrand CM.
    Fast quantitative parameter maps without fitting: Integration yields
    accurate mono‐exponential signal decay rates.
    Magnetic resonance in medicine. 2018 Jun;79(6):2978-85.
"""
function r2star_crsi(
    mag::AbstractArray{T, N},
    TEs::NTuple{NT, Real},
    mask::Union{Nothing, AbstractArray{Bool}} = nothing;
    M::Integer = 3,
    sigma::Union{Nothing, Real} = nothing,
    Rsz::NTuple{NR, Integer} = size(mag)[1:N-1] .÷ 20,
) where {T<:AbstractFloat, N, NT, NR}
    N > 1 || throw(ArgumentError("array must contain echoes in last dimension"))
    NT > 1 || throw(ArgumentError("data must be multi-echo"))

    size(mag, N) == NT || throw(DimensionMismatch())
    mask === nothing || length(mask) == length(mag) ÷ NT || throw(DimensionMismatch())
    sigma === nothing || NR == N-1 || throw(DimensionMismatch())
    M > 0 || throw(ArgumentError("interpolation factor M must be greater than 0"))

    P = tcopy(x -> x*x, mag)
    r2s = tzero(mag, size(mag)[1:N-1])

    vP = reshape(P, :, NT)
    vr2s = vec(r2s)

    if sigma !== nothing
        σ2 = convert(T, 2*sigma*sigma)
    else
        σ2 = _noise_crsi(P, Rsz, mask)
    end

    _zeroT = zero(T)
    α = convert(T, 1//2)
    β = convert(T, -α * σ2 * (TEs[end] - TEs[1]))

    τ  = SVector{NT-1, T}((TEs[2:end] .- TEs[1:end-1]) ./ (M+1))
    γ0 = SVector{M+1, T}([(2*M - 2*m + 1) / (2*M + 2) for m in 0:M]...)
    γ1 = SVector{M+1, T}([(2*m + 1) / (2*M + 2) for m in 0:M]...)

    if mask === nothing
        @inbounds @batch per=thread for I in eachindex(vr2s)
            den = β
            for t in Base.OneTo(NT-1)
                p0 = vP[I,t]
                p1 = vP[I,t+1]
                p = pow(p0, γ0[1]) * pow(p1, γ1[1])
                for m in 2:M+1
                    p = muladd(pow(p0, γ0[m]), pow(p1, γ1[m]), p)
                end
                den = muladd(τ[t], p, den)
            end

            vr2s[I] = iszero(den) ? _zeroT : α * (vP[I,1] - vP[I,end]) * inv(den)
        end

    else
        vmask = vec(mask)

        @inbounds @batch per=thread for I in eachindex(vr2s)
            if vmask[I]
                den = β
                for t in Base.OneTo(NT-1)
                    p0 = vP[I,t]
                    p1 = vP[I,t+1]
                    p = pow(p0, γ0[1]) * pow(p1, γ1[1])
                    for m in 2:M+1
                        p = muladd(pow(p0, γ0[m]), pow(p1, γ1[m]), p)
                    end
                    den = muladd(τ[t], p, den)
                end

                vr2s[I] = iszero(den) ? _zeroT : α * (vP[I,1] - vP[I,end]) * inv(den)
            end
        end
    end

    return r2s
end

function _noise_crsi(
    P::AbstractArray{T, N},
    Rsz::NTuple{M, Integer},
    mask::Union{Nothing, AbstractArray{Bool}} = nothing
) where {T<:AbstractFloat, N, M}
    M == N-1 || throw(DimensionMismatch())

    sz = size(P)
    rsz = ntuple(n -> min(Rsz[n], (sz[n]-2)÷2), Val(N-1))

    lb = rsz .+ 2
    ub = sz[1:M] .- 1 .- rsz

    outer = CartesianIndices(ntuple(n -> 2:sz[n]-1, Val(N-1)))
    inner = CartesianIndices(ntuple(n -> 2+rsz[n]:sz[n]-1-rsz[n], Val(N-1)))

    n = 0
    σ2 = zero(T)

    if mask === nothing
        @inbounds for t in axes(P, N)
            for I in EdgeIterator(outer, inner)
                if all(n -> I[n] < lb[n] || I[n] > ub[n], 1:M)
                    σ2 += P[I,t]
                    n += 1
                end
            end
        end

    else
        @inbounds for t in axes(P, N)
            for I in EdgeIterator(outer, inner)
                if all(n -> I[n] < lb[n] || I[n] > ub[n], 1:M) && !mask[I]
                    σ2 += P[I,t]
                    n += 1
                end
            end
        end
    end

    return n == 0 ? zero(T) : σ2 / n
end

"""
    function r2star_numart2s(
        mag::AbstractArray{<:AbstractFloat, N > 1},
        TEs::NTuple{NT > 1, Real},
        mask::Union{Nothing, AbstractArray{Bool}} = nothing
    ) -> typeof(similar(mag, size(mag)[1:N-1]))

Numerical Algorithm for Real-time T2* mapping (NumART2*) [1].

### Arguments
- `mag::AbstractArray{<:AbstractFloat, N > 1}`: multi-echo magnitude
- `TEs::NTuple{NT > 1, Real}`: echo times
- `mask::Union{Nothing, AbstractArray{Bool}} = nothing`: binary mask of region of interest

### Returns
- `typeof(similar(mag, size(mag)[1:N-1]))`: R2* map (1 / units of TEs)

### References
[1] Hagberg GE, Indovina I, Sanes JN, Posse S. Real‐time quantification of T2*
    changes using multiecho planar imaging and numerical methods.
    Magnetic Resonance in Medicine: An Official Journal of the International
    Society for Magnetic Resonance in Medicine. 2002 Nov;48(5):877-82.
"""
function r2star_numart2s(
    mag::AbstractArray{T, N},
    TEs::NTuple{NT, Real},
    mask::Union{Nothing, AbstractArray{Bool}} = nothing
) where {T<:AbstractFloat, N, NT}
    N > 1 || throw(ArgumentError("array must contain echoes in last dimension"))
    NT > 1 || throw(ArgumentError("data must be multi-echo"))

    size(mag, N) == NT || throw(DimensionMismatch())
    mask === nothing || length(mask) == length(mag) ÷ NT || throw(DimensionMismatch())

    r2s = tzero(mag, size(mag)[1:N-1])

    vmag = reshape(mag, :, NT)
    vr2s = vec(r2s)

    _zeroT = zero(T)
    α = convert(T, 2*(NT - 1) / (TEs[end] - TEs[1]))

    if mask === nothing
        @inbounds @batch for I in eachindex(vr2s)
            den = vmag[I,1]
            for t in 2:NT-1
                den += vmag[I,t] + vmag[I,t]
            end
            den += vmag[I,NT]
            vr2s[I] = iszero(den) ? _zeroT : α * (vmag[I,1] - vmag[I,NT]) * inv(den)
        end

    else
        vmask = vec(mask)

        @inbounds @batch for I in eachindex(vr2s)
            if vmask[I]
                den = vmag[I,1]
                for t in 2:NT-1
                    den += vmag[I,t] + vmag[I,t]
                end
                den += vmag[I,NT]
                vr2s[I] = iszero(den) ? _zeroT : α * (vmag[I,1] - vmag[I,NT]) * inv(den)
            end
        end
    end

    return r2s
end

ext: jl


text:
macro add_properties(T, props...)
    e = quote end
    for prop in props
        f = quote $(esc(prop))(a::$(esc(T))) = true end
        append!(e.args, f.args)
    end
    e
end

tolerance(x) = tolerance(typeof(x))
tolerance(::Type{T}) where {T} = tolerance(prectype(T))
tolerance(T::Type{<:AbstractFloat}) = sqrt(eps(T))

default_threshold(x) = default_threshold(typeof(x))
default_threshold(T::Type{<:AbstractFloat}) = 100eps(T)
default_threshold(::Type{T}) where {T} = default_threshold(prectype(T))

linspace(a, b, n=100) = range(a, stop=b, length=n)

# delinearize_coefficients!(dest::BlockVector, src::AbstractVector) =
#     dest[:] .= src[:]
#
# linearize_coefficients!(dest::AbstractVector, src::BlockVector) =
#     dest[:] .= src[:]

delinearize_coefficients!(dest::AbstractArray{T,N}, src::AbstractVector{T}) where {T,N} =
    dest[:] .= src[:]

linearize_coefficients!(dest::AbstractVector{T}, src::AbstractArray{T,N}) where {T,N} =
    dest[:] .= src[:]

components(bv::BlockVector) = [view(bv, Block(i)) for i in 1:blocklength(bv)]

function BlockArrays.BlockVector(arrays::AbstractVector{T}...) where {T}
    A = BlockArray{T}(undef_blocks, [length(array) for array in arrays])
    for (i,a) in enumerate(arrays)
        A[Block(i)] = a
        # setblock!(A, a, i)
    end
    A
end

function BlockArrays.BlockMatrix(arrays::AbstractMatrix{T}...) where {T}
    A = BlockArray{T}(undef_blocks, [length(array) for array in arrays])
    for (i,a) in enumerate(arrays)
        A[Block(i)] = a
        # setblock!(A, a, i)
    end
    A
end

"Return true if the set is indexable and has elements whose type is a subtype of T."
indexable_list(set, ::Type{T}) where {T} = typeof(set[1]) <: T

indexable_list(set::Array{S}, ::Type{T}) where {S,T} = S <: T
indexable_list(set::NTuple{N,S}, ::Type{T}) where {N,S,T}= S <: T

#An efficient way to access elements of a Tuple type using index j
@generated function tuple_index(::Type{T}, j) where {T <: Tuple}
    :($T.parameters[j])
end

@generated function tuple_length(::Type{T}) where {T <: Tuple}
    :($length(T.parameters))
end

function insert_at(a::Vector, idx, b)
    if idx > 1
        [a[1:idx-1]..., b..., a[idx:end]...]
    else
        [b..., a...]
    end
end

isdyadic(n::Int) = n == 1<<round(Int, log2(n))

# This is a candidate for a better implementation. How does one generate a
# unit vector in a tuple?
dimension_tuple(::Val{N}, dim::Int) where N = ntuple(k -> ((k==dim) ? 1 : 0), Val(N))
dimension_tuple(N::Int, dim::Int) = ntuple(k -> ((k==dim) ? 1 : 0), N)

function matrix_by_mul(A::AbstractMatrix{T}) where T
    Z = zeros(T, size(A))
    matrix_by_mul!(Z,A)
    Z
end

function matrix_by_mul!(Z::Matrix{T}, A::AbstractMatrix{T}) where T
    e = zeros(T,size(Z,2))
    for i in 1:size(Z,2)
        e[i] = 1
        mul!(view(Z, :, i), A, e)
        e[i] = 0
    end
end

# TODO: move these definitions to DomainSets
###############
# Subeltype
###############

"Return the type of the elements of `x`."
subeltype(x) = subeltype(typeof(x))
subeltype(::Type{T}) where {T} = eltype(eltype(T))

###############
# Dimension
###############

dimension(x) = dimension(typeof(x))
dimension(::Type{T}) where {T <: Number} = 1
dimension(::Type{SVector{N,T}}) where {N,T} = N
dimension(::Type{<:NTuple{N,Any}}) where {N} = N
dimension(::Type{CartesianIndex{N}}) where {N} = N
dimension(::Type{T}) where {T} = 1

const Domain1d{T <: Number} = Domain{T}
const Domain2d{T} = EuclideanDomain{2,T}
const Domain3d{T} = EuclideanDomain{3,T}
const Domain4d{T} = EuclideanDomain{4,T}
export Domain1d, Domain2d, Domain3d, Domain4d

iscompatible(map1::AbstractMap, map2::AbstractMap) = map1==map2
iscompatible(map1::AffineMap, map2::AffineMap) = (map1.A ≈ map2.A) && (map1.b ≈ map2.b)
iscompatible(domain1::Domain, domain2::Domain) = domain1 == domain2


text:
####################################################
#### Quantecon state-action representation #########
####################################################

# this function sets up the state-action pairs and feeds them into the QuantEcon solver
function _solve(p::DDP, method::Type{All},  mTransition, mReward, disp::Bool)

    # ONLY PROGRAMMED FOR SINGLECHOICEVAR SO FAR!

    # state action pair representation
    nStateVars = length(p.tStateVectors)
    tNodes = length.(p.tStateVectors)

    nChoiceVars = length(p.tChoiceVectors)
    nChoices = prod(length.(p.tChoiceVectors))
    # nChoices = length(p.tChoiceVectors[1])

    # tAll = (tNodes..., nChoices)
    # mSA = gridmake(p.tStateVectors..., p.vChoiceVector)

    # mSA_coord = gridmake(1:prod(tNodes), 1:nChoices)
    mSA_coord = gridmake(1:nChoices,1:prod(tNodes))
    # mS_coord = mSA_coord[:,1:1]
    # mA_coord = mSA_coord[:,2:2]
    mS_coord = mSA_coord[:,2:2]
    mA_coord = mSA_coord[:,1:1]

    # rewardmatrix gives choices x states, want states x choices
    # R = mReward' # is choices x states
    R = mReward
    R_sa = R[:]

    # get transition matrix, which elements are allowed?
    Q = mTransition

    # is state-action representation, since the choices don't need to be equal
    # to the states, it is possible that some choices are not admissible
    # this will be the case where probabilities of the transition matrix are not
    # between 0 and 1
    # mark all rows where probabilities are within 0 and 1
    admissible = all((Q.>=0.) .& (Q.<=1.), dims=2)[:]
    # would be more effcient to do this on g rather than Q, huge matrix

    # delete state-action pairs that are not admissible
    mS_coord = mS_coord[admissible]
    mA_coord = mA_coord[admissible]
    R_sa = R_sa[admissible]
    # Q = Q[find(admissible),:]
    Q = Q[(LinearIndices(admissible))[findall(admissible)],:]

    # R_sa = reward
    # Q = transition probability array
    # mS_coord: Action Indices
    # mA_coord: Action Index Pointers
    ddp = DiscreteDP(R_sa, Q, p.β, mS_coord, mA_coord)

    results = solve(ddp, VFI)

    if disp
        println(string("number of iterations = ", results.num_iter))
    end

    # results.sigma is a vector of indices of optimal choices
    # mPolFun = p.tChoiceVectors[1][results.sigma]
    vChoiceVector = gridmake(p.tChoiceVectors...)
    mPolFun = vChoiceVector[results.sigma, :]

    # meshPolFun = reshape(mPolFun, tuple(tNodes...))
    meshValFun = reshape(results.v, tuple(tNodes...))

    # return createsolution(p, meshValFun, (meshPolFun,))
    meshPolFun = Vector{Array{Float64}}()
    for k = 1 : nChoiceVars
        push!(meshPolFun, reshape(mPolFun[:,k], tuple(tNodes...)))
    end

    return meshValFun, Tuple(meshPolFun)
end

# function solve(p::TwoChoiceVar, method::Type{SA},  mTransition, mReward, disp::Bool)
#
#     # state action pair representation
#     nStateVars = length(p.tStateVectors)
#     tNodes = length.(p.tStateVectors)
#
#     nChoiceVars = length(p.tChoiceVectors)
#     nChoices = prod(length.(p.tChoiceVectors))
#
#     # tAll = (tNodes..., nChoices)
#     # mSA = gridmake(p.tStateVectors..., p.vChoiceVector)
#
#     mSA_coord = gridmake(1:prod(tNodes), 1:nChoices)
#     mS_coord = mSA_coord[:,1:1]
#     mA_coord = mSA_coord[:,2:2]
#
#     # reward usually gives choices x states, want states x choices
#     R = mReward' # is choices x states
#     # R_sa = addDim(R[:]) # can convert into state-action repr
#     R_sa = R[:]
#
#
#     # get transition matrix, which elements are allowed?
#     Q = transitionmatrix(p, method)
#
#     # mark all rows where probabilities are not within 0 and 1
#     admissible = all((Q.>=0.) .& (Q.<=1.), dims=2)[:]
#     # would be more effcient to do this on g than Q, huge matrix
#
#     # delete s-a pairs that are not admissible
#     mS_coord = mS_coord[admissible]
#     mA_coord = mA_coord[admissible]
#     R_sa = R_sa[admissible]
#     Q = Q[(LinearIndices(admissible))[findall(admissible)], :]
#
#
#     # R_sa = reward
#     # Q = transition probability array
#     # mS_coord: Action Indices
#     # mA_coord: Action Index Pointers
#     ddp = DiscreteDP(R_sa, Q, p.params.β, mS_coord, mA_coord)
#
#     results = solve(ddp, VFI)
#
#     if disp
#         println(string("number of iterations = ", results.num_iter))
#     end
#
#     # size of results.sigma is (nK x nz) x 1
#     # values are the indeces of choices
#     # @show size(results.sigma)
#     # @show results.sigma[1:10]
#
#     vChoiceVector = gridmake(p.tChoiceVectors...)
#     mPolFun = vChoiceVector[results.sigma, :]
#
#     meshPolFun1 = reshape(mPolFun[:,1], tuple(tNodes...))
#     meshPolFun2 = reshape(mPolFun[:,2], tuple(tNodes...))
#     meshValFun = reshape(results.v, tuple(tNodes...))
#
#     return createsolution(p, meshValFun, (meshPolFun1, meshPolFun2))
#
# end


text:
using DiffMesh
using Test
using Pkg, SafeTestsets

function activate_examples_env()
    Pkg.activate("../examples")
    Pkg.develop(PackageSpec(path=dirname(@__DIR__)))
    Pkg.instantiate()
end

@testset "DiffMesh.jl" begin
    activate_examples_env()
    @time @safetestset "Tridiagonal" begin include("tdg.jl") end
    @time @safetestset "Examples" begin include("examples.jl") end
end


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

mutable struct UpdateVirtualNetworkRuleWithAccountParameters <: SwaggerModel
    name::Any # spec type: Union{ Nothing, String } # spec name: name
    properties::Any # spec type: Union{ Nothing, UpdateVirtualNetworkRuleProperties } # spec name: properties

    function UpdateVirtualNetworkRuleWithAccountParameters(;name=nothing, properties=nothing)
        o = new()
        validate_property(UpdateVirtualNetworkRuleWithAccountParameters, Symbol("name"), name)
        setfield!(o, Symbol("name"), name)
        validate_property(UpdateVirtualNetworkRuleWithAccountParameters, Symbol("properties"), properties)
        setfield!(o, Symbol("properties"), properties)
        o
    end
end # type UpdateVirtualNetworkRuleWithAccountParameters

const _property_map_UpdateVirtualNetworkRuleWithAccountParameters = Dict{Symbol,Symbol}(Symbol("name")=>Symbol("name"), Symbol("properties")=>Symbol("properties"))
const _property_types_UpdateVirtualNetworkRuleWithAccountParameters = Dict{Symbol,String}(Symbol("name")=>"String", Symbol("properties")=>"UpdateVirtualNetworkRuleProperties")
Base.propertynames(::Type{ UpdateVirtualNetworkRuleWithAccountParameters }) = collect(keys(_property_map_UpdateVirtualNetworkRuleWithAccountParameters))
Swagger.property_type(::Type{ UpdateVirtualNetworkRuleWithAccountParameters }, name::Symbol) = Union{Nothing,eval(Base.Meta.parse(_property_types_UpdateVirtualNetworkRuleWithAccountParameters[name]))}
Swagger.field_name(::Type{ UpdateVirtualNetworkRuleWithAccountParameters }, property_name::Symbol) =  _property_map_UpdateVirtualNetworkRuleWithAccountParameters[property_name]

function check_required(o::UpdateVirtualNetworkRuleWithAccountParameters)
    (getproperty(o, Symbol("name")) === nothing) && (return false)
    true
end

function validate_property(::Type{ UpdateVirtualNetworkRuleWithAccountParameters }, name::Symbol, val)
end


text:
"""
    module Geometry

Module defining basic geometrical concepts.
"""
module Geometry

import Base: ==, length, in, iterate, getindex, lastindex, isempty, eltype, keys
import Base: union, setdiff, intersect, issubset

using StaticArrays
using LinearAlgebra
using OrderedCollections
using AbstractTrees
using RecipesBase
using Printf
using Statistics:median

using WavePropBase
using WavePropBase.Utils

WavePropBase.@import_interface

include("point.jl")
include("referenceshapes.jl")
include("entities.jl")
include("domain.jl")

export
    # abstract types
    AbstractReferenceShape,
    AbstractEntity,
    AbstractParametricBody,
    # types
    ElementaryEntity,
    ParametricEntity,
    ParametricCurve,
    Domain,
    ReferencePoint,
    ReferenceLine,
    ReferenceTriangle,
    ReferenceTetrahedron,
    ReferenceSquare,
    ReferenceHyperCube,
    Point2D,
    Point3D,
    # functions
    clear_entities!,
    entities,
    tag,
    key,
    assertequaldim,
    boundary,
    skeleton,
    internal_boundary,
    external_boundary,
    low_corner,
    high_corner,
    measure,
    vertices,
    line,
    new_tag,
    global_add_entity!,
    index_range,
    parent,
    container,
    loc2glob,
    points,
    # global variables
    TAGS,
    ENTITIES

end # module

max_stars_repo_name: WaveProp/WavePropBase

max_issues_repo_name: WaveProp/WavePropBase

max_forks_repo_name: WaveProp/WavePropBase


text:
module Day7

using Advent: save_input

struct Bag
    color::AbstractString
    quant::Int
end

function parse_rules(filename)
    rules = Dict{String, Vector{Bag}}()
    for line in eachline(filename)
        color, rest = split(line, r" bags contain ")
        rawrules = eachmatch(r"(\d+) (\w+ \w+) bag", rest)
        rules[color] = [
            Bag.(color, parse(Int, quant))
            for (quant, color) in Iterators.map(x -> x.captures, rawrules)
        ]
    end
    rules
end

function resolve(rules, color)
    color == "shiny gold" || any(cc -> resolve(rules, cc.color), rules[color])
end

function solve_a(filename)
    rules = parse_rules(filename)
    count(keys(rules)) do color
        color != "shiny gold" && resolve(rules, color)
    end
end

function resolve2(rules, color)
    if length(rules[color]) > 0
        1 + sum(cc -> cc.quant * resolve2(rules, cc.color), rules[color])
    else
        1
    end
end

function solve_b(filename)
    rules = parse_rules(filename)
    resolve2(rules, "shiny gold") - 1
end

solve(filename = save_input(2020, 7)) = solve_a(filename), solve_b(filename)

end

max_stars_repo_name: AnAverageHuman/competitive

max_issues_repo_name: AnAverageHuman/competitive

max_forks_repo_name: AnAverageHuman/competitive


text:
function log_all_plots(results :: Dict{String, Pair}, games :: Dict{String, Vector{Game}}, league_name :: String)
    teams = [team for team ∈ keys(games)]
    gr()
    for team ∈ teams
        plot(map(x -> x.date, games[team]), results[team].first, ribbon=2 .* results[team].second, fillalpha=0.15, legend=false, title="$(team) Form", xlabel="Matches", ylabel="Form Index", size=(1000,500))
        savefig("form/$(league_name)/$(team)_form.png")
    end

    plt = plot(size = (2000, 1000), legend=:outertopright, title="Form Comparison $(league_name)", xlabel="Matches", ylabel="Form Index")
    for team ∈ teams
        plot!(plt, map(x -> x.date, games[team]), results[team].first, ribbon=2 .* results[team].second, fillalpha=0.15, labels=team)
    end
    savefig("form/$(league_name)/all_form.png")

    plotly()
    plt = plot(size = (1500, 750), legend=:outertopright, title="Form Comparison $(league_name)", xlabel="Matches", ylabel="Form Index")
    for team ∈ teams
        plot!(plt, map(x -> x.date, games[team]), results[team].first, ribbon=2 .* results[team].second, fillalpha=0.15, labels=team)
    end
    display(plt)
    savefig(plt, "form/$(league_name)/all_form.html")
    savefig(plt, "docs/all_form_$(replace(league_name, " " => "")).html")
end

max_stars_repo_name: shalinkpatel/football-form

max_issues_repo_name: shalinkpatel/football-form

max_forks_repo_name: shalinkpatel/football-form


text:
using ModelingToolkit, MethodOfLines, LinearAlgebra, OrdinaryDiffEq
using DomainSets

@testset "Inviscid Burgers equation, 1D, u(0, x) ~ x" begin
    @parameters x t
    @variables u(..)
    Dx = Differential(x)
    Dt = Differential(t)
    x_min = 0.0
    x_max = 1.0
    t_min = 0.0
    t_max = 6.0

    analytic_u(t, x) = x / (t + 1)

    eq = Dt(u(t, x)) ~ -u(t, x) * Dx(u(t, x))

    bcs = [u(0, x) ~ x,
        u(t, x_min) ~ analytic_u(t, x_min),
        u(t, x_max) ~ analytic_u(t, x_max)]

    domains = [t ∈ Interval(t_min, t_max),
        x ∈ Interval(x_min, x_max)]

    dx = 0.05

    @named pdesys = PDESystem(eq, bcs, domains, [t, x], [u(t, x)])

    disc = MOLFiniteDifference([x => dx], t, upwind_order=1)

    prob = discretize(pdesys, disc)

    sol = solve(prob, Tsit5())

    grid = get_discrete(pdesys, disc)
    x_disc = grid[x]
    solu = [map(d -> sol[d][i], grid[u(t, x)]) for i in 1:length(sol[t])]

    for (i, t) in enumerate(sol.t[1:end])
        u_analytic = analytic_u.([t], x_disc)
        u_disc = solu[i]
        @test all(isapprox.(u_analytic, u_disc, atol=1e-3))
    end
end

@testset "Inviscid Burgers equation, 1D, u(0, x) ~ x, Non-Uniform" begin
    @parameters x t
    @variables u(..)
    Dx = Differential(x)
    Dt = Differential(t)
    x_min = 0.0
    x_max = 1.0
    t_min = 0.0
    t_max = 1.0

    analytic_u(t_, x_) = x_ / (t_ + 1.0)

    eq = Dt(u(t, x)) ~ -u(t, x) * Dx(u(t, x))

    bcs = [u(0, x) ~ x,
        u(t, x_min) ~ analytic_u(t, x_min),
        u(t, x_max) ~ analytic_u(t, x_max)]

    domains = [t ∈ Interval(t_min, t_max),
        x ∈ Interval(x_min, x_max)]

    @named pdesys = PDESystem(eq, bcs, domains, [t, x], [u(t, x)])

    dx = 0:0.05:1
    dx = collect(dx)
    dx[2:end-1] .= dx[2:end-1] .+ rand([0.001, -0.001], length(dx[2:end-1]))

    disc = MOLFiniteDifference([x => dx], t, upwind_order=1)

    prob = discretize(pdesys, disc)

    sol = solve(prob, Tsit5())

    grid = get_discrete(pdesys, disc)
    x_disc = grid[x]
    x_disc = getfield.(x_disc, [:val])

    solu = [map(d -> sol[d][i], grid[u(t, x)]) for i in 1:length(sol[t])]

    for (i, t_disc) in enumerate(sol.t[1:end])
        u_analytic = analytic_u.([t_disc], x_disc)
        u_disc = solu[i]
        @test all(isapprox.(u_analytic, u_disc, atol=1 * 10^(-2.5)))
    end
end

# Exact solutions from: https://www.sciencedirect.com/science/article/pii/S0898122110003883
@testset "Test 01: Burger's Equation 2D" begin
    @parameters x y t
    @variables u(..) v(..)
    Dt = Differential(t)
    Dx = Differential(x)
    Dy = Differential(y)
    Dxx = Differential(x)^2
    Dyy = Differential(y)^2

    R = 80.0

    x_min = y_min = t_min = 0.0
    x_max = y_max = 1.0
    t_max = 1.0

    u_exact(x, y, t) = 3 / 4 - 1 / (4 * (1 + exp(R * (-t - 4x + 4y) / 32)))
    v_exact(x, y, t) = 3 / 4 + 1 / (4 * (1 + exp(R * (-t - 4x + 4y) / 32)))

    eq = [Dt(u(x, y, t)) + u(x, y, t) * Dx(u(x, y, t)) + v(x, y, t) * Dy(u(x, y, t)) ~ 1 / R * (Dxx(u(x, y, t)) + Dyy(u(x, y, t))),
        Dt(v(x, y, t)) + u(x, y, t) * Dx(v(x, y, t)) + v(x, y, t) * Dy(v(x, y, t)) ~ 1 / R * (Dxx(v(x, y, t)) + Dyy(v(x, y, t)))]

    domains = [x ∈ Interval(x_min, x_max),
        y ∈ Interval(y_min, y_max),
        t ∈ Interval(t_min, t_max)]

    bcs = [u(x, y, 0) ~ u_exact(x, y, 0),
        u(0, y, t) ~ u_exact(0, y, t),
        u(x, 0, t) ~ u_exact(x, 0, t),
        u(1, y, t) ~ u_exact(1, y, t),
        u(x, 1, t) ~ u_exact(x, 1, t), v(x, y, 0) ~ v_exact(x, y, 0),
        v(0, y, t) ~ v_exact(0, y, t),
        v(x, 0, t) ~ v_exact(x, 0, t),
        v(1, y, t) ~ v_exact(1, y, t),
        v(x, 1, t) ~ v_exact(x, 1, t)]

    @named pdesys = PDESystem(eq, bcs, domains, [t, x, y], [u(x, y, t), v(x, y, t)])

    # Method of lines discretization
    dx = 0.1
    dy = 0.1

    # Try 4th approx order
    order = 4

    discretization = MOLFiniteDifference([x => dx, y => dy], t, approx_order=order)

    # Convert the PDE problem into an ODE problem
    prob = discretize(pdesys, discretization)

    sol = solve(prob, Tsit5())

    Nx = floor(Int64, (x_max - x_min) / dx) + 1
    Ny = floor(Int64, (y_max - y_min) / dy) + 1

    # anim = @animate for k in 1:length(t)
        #        solu′ = reshape([sol[u[(i-1)*Ny+j]][k] for i in 1:Nx for j in 1:Ny],(Nx,Ny))
        #        solv′ = reshape([sol[v[(i-1)*Ny+j]][k] for i in 1:Nx for j in 1:Ny],(Nx,Ny))
        #        heatmap(solu′)
        # end
        # gif(anim, "plots/Burgers2Dsol.gif", fps = 5)
    grid = get_discrete(pdesys, discretization)

    solu′ = map(d -> sol[d][end], grid[u(x, y, t)])
    solv′ = map(d -> sol[d][end], grid[v(x, y, t)])

    t = sol[t]
    r_space_x = grid[x]
    r_space_y = grid[y]

    asfu = reshape([u_exact(t_max, r_space_x[i], r_space_y[j]) for j in 1:Ny for i in 1:Nx], (Nx, Ny))
    asfv = reshape([v_exact(t_max, r_space_x[i], r_space_y[j]) for j in 1:Ny for i in 1:Nx], (Nx, Ny))

    asfu[1, 1] = asfu[1, end] = asfu[end, 1] = asfu[end, end] = 0.0
    asfv[1, 1] = asfv[1, end] = asfv[end, 1] = asfv[end, end] = 0.0

    # anim = @animate for T in t
    #        asfu = reshape([u_exact(T,r_space_x[i],r_space_y[j]) for j in 1:Ny for i in 1:Nx],(Nx,Ny))
    #        asfv = reshape([v_exact(T,r_space_x[i],r_space_y[j]) for j in 1:Ny for i in 1:Nx],(Nx,Ny))

    #        heatmap(asfu)
    # end
    # gif(anim, "plots/Burgers2Dexact.gif", fps = 5)

    #    mu = max(asfu...)
    #    mv = max(asfv...)
    @test_broken asfu ≈ solu′ atol = 0.2
    @test_broken asfv ≈ solv′ atol = 0.2
end


text:
"""
    is_leap_year(year)

Return `true` if `year` is a leap year in the gregorian calendar.

"""
function is_leap_year(year)

end


text:
using Cairo
using Test

# CairoSurface
#     ptr::Ptr{Nothing}
#     width::Float64
#     height::Float64
#     data::Matrix{T}

c = CairoARGBSurface(160, 90)
@test c isa CairoSurface{UInt32}

ctx = CairoContext(c)
@test ctx isa CairoContext


text:
"""
    get_coefficients(α,θ,m)

α should be a vector of shapes of length n, θ should be a matrix of θ of size (n,d) for a MultivariateGammaConvolution with d marginals.

This function produce the coefficients of the multivariate (tensorised) expensions of the density. It is type stable and quite optimized.
it is based on some precomputations that are done in a certain precision.

"""
function get_coefficients(α, θ, m; P = nothing)
    # α must be an array with size (n,)
    # θ must be an array with size (n,d)
    # max_p must be a Tuple of integers with size (d,)
    T = Base.promote_eltype(α,θ,[1.0]) # At least float.
    α = T.(α)
    θ = T.(θ)

    # Trim down null αs values:
    are_pos = findall(α .> eps(T))
    if ndims(θ) == 1
        θ = θ[are_pos]
    else
        θ = θ[are_pos,:] # produces a matrix when θ is only a vector..
    end
    α = α[are_pos]
    if isnothing(P)
        P = get_precomp(T,sum(m))::PreComp{T}
    end
    coefs = zeros(T,m)
    build_coefficients!(coefs,α,θ,T(1.0),m,P)
    return coefs
end

function build_coefficients!(coefs,α,θ,cst1,m,P)

    # This is the default method, for higher dimensions.

    # Allocates ressources
    κ = deepcopy(coefs)
    μ = deepcopy(coefs)
    n = size(θ)[1]
    d = length(m)
    I = CartesianIndices(coefs)
    S = θ ./ (cst1 .+ sum(θ,dims=2))
    S_pow = [s^k for k in (0:Base.maximum(m)), s in S]

    # Starting the algorithm: there is an edge case for the Oth cumulant, 0th moment and 0th coef:
    to_log = cst1 .- sum(S,dims=2)
    to_log .= ifelse.(to_log .< 0, to_log .+ eps(Base.eltype(α)), to_log)
    κ[1] = sum(α .* log.(to_log))
    μ[1] = exp(κ[1])
    coefs[1] = μ[1]

    @inbounds for k in I[2:length(I)]
        # Indices and organisation
        k_arr = Tuple(k)
        degree = findfirst(k_arr .!= 1)
        sk = sum(k_arr)

        # Main Computations:
        for i in 1:n
            rez = α[i]
            for j in 1:d
                rez *= S_pow[k[j],i,j]
            end
            κ[k] += rez
        end
        κ[k] *= P.FACTS[sk-d]

        for j in CartesianIndices(k)
            rez_coefs = μ[j]
            if j[degree] < k[degree]
                rez_mu = μ[j] * κ[k - j + I[1]]
                for i in 1:d
                    rez_mu *= P.BINS[j[i],k[i]-Int(i==degree)]
                    rez_coefs *= P.LAGUERRE[j[i], k[i]]
                end
                μ[k] += rez_mu
            else
                for i in 1:d
                    rez_coefs *= P.LAGUERRE[j[i], k[i]]
                end
            end
            coefs[k] += rez_coefs
        end
    end
    coefs .*= sqrt(2*cst1)^d
end

function build_coefficients!(coefs::AbstractArray{T,1},α::AbstractArray{T,1},θ::AbstractArray{T,1},cst1::T,m::NTuple{1, T2},P::PreComp{T}) where {T2 <: Int,T <: Real}

    #kappa = Array{Base.eltype(α), 1}(undef, m[1])
    kappa = zeros(T,m[1])
    mu = zeros(T, m[1])

    @. θ /= 1 + θ

    # Edge case for the Oth cumulant, 0th moment and 0th coef:
    to_log = cst1 .- θ
    to_log .= ifelse.(to_log .< 0, to_log .+ eps(Base.eltype(α)), to_log) # correcting rounding errors.
    kappa[1] = sum(α .* log.(to_log))
    coefs[1] = mu[1] = exp(kappa[1]) # this is the only place kappa[1] is used.
    #coefs[1] = mu[1] = prod(to_log .^ α)
    @inbounds for k in 2:m[1]
        # Main Computations:
        α .*= θ
        kappa[k] = sum(α)
        @inbounds for j in 1:k-1
            mu[k] += mu[j] * kappa[k-j+1]
        end
        mu[k] /= k-1
        @views coefs[k] = P.LAGUERRE2[1:k, k]'mu[1:k]
    end
end

function build_coefficients!(coefs::AbstractArray{T,1},α::AbstractArray{T,1},θ::AbstractArray{T,2},cst1::T,m::NTuple{2, T2},P::PreComp{T}) where {T2 <: Int,T <: Real}
    # Allocates ressources
    κ = Array{Base.eltype(α), 2}(undef, m)
    μ = zeros(Base.eltype(α), m)

    #n = size(θ)[1]
    d = length(m)
    I = CartesianIndices(coefs)
    θ ./= (cst1 .+ sum(θ,dims=2))
    S_pow = [s^k for k in (0:Base.maximum(m)), s in θ]

    # reduce the number of multiplications :
    for i in 1:size(α,1)
        S_pow[:,i,:] .*= sqrt(α[i])
    end

    # Starting the algorithm: there is an edge case for the Oth cumulant, 0th moment and 0th coef:
    to_log = cst1 .- sum(θ,dims=2)
    to_log .= ifelse.(to_log .< 0, to_log .+ eps(Base.eltype(α)), to_log)
    κ[1] = sum(α .* log.(to_log)) # this log fails ifsum(S,dims=2) is greater than 1, which occurs because of rounding errors. We add an eps to avoid that.
    μ[1] = exp(κ[1])
    coefs[1] = μ[1]

    @inbounds for k in I[2:length(I)]
        # Indices and organisation
        degree = k[1]!=1 ? 1 : 2
        deg1 = Int(1==degree)
        deg2 = Int(2==degree)

        κ[k] = sum(S_pow[k[1],:,1] .* S_pow[k[2],:,2]) * P.FACTS[sum(Tuple(k))-d]
        for j in CartesianIndices(k)
            if j[degree] < k[degree]
                μ[k] += μ[j] * κ[k - j + I[1]] *  P.BINS[j[1],k[1]-deg1] * P.BINS[j[2],k[2]-deg2]
            end
            coefs[k] += μ[j] * P.LAGUERRE[j[1], k[1]] * P.LAGUERRE[j[2], k[2]]
        end
    end
    coefs .*= sqrt(2*cst1)^d
end

"""
    laguerre_L_2x(x,p)

Compute univariate laguerre polynomials Lₚ(2x).
"""
function laguerre_L_2x(x,p)
    P = get_precomp(Base.eltype(x),p+1)
    sum(P.LAGUERRE[1:(p+1),p+1] .* x .^ (0:p))
end

"""
    laguerre_phi(x,p)

Computes the function ϕₚ(x) = √2 e⁻ˣ Lₚ(2x). This functions form an orthonormal basis of R₊
"""
function laguerre_phi(x,p)
    exp(-x) * sqrt(2) * laguerre_L_2x(x,p)
end

"""
    laguerre_density(x,coefs)

Given some laguerre coefficients, computes the correpsonding function at the point x.
"""
function laguerre_density(x,coefs)
    rez = zero(Base.eltype(coefs))
    for p in CartesianIndices(coefs)
        rez += coefs[p] * prod(laguerre_phi.(x,Tuple(p) .-1))
    end
    return max(rez,0)
end

"""
    laguerre_phi_several_pts(x,max_p)

Computes the laguerre_phi for each row of x and each p in CartesianIndex(maxp)
This is a lot more efficient than broadcasting the laguerre_phi function,
but this is ugly in the sense that we re-wrote some of the code. This machanisme could be re-written more efficiently.
"""
function laguerre_phi_several_pts(x::AbstractArray{T},max_p) where {T <: Real}
    P = get_precomp(T,Base.maximum(max_p)+1)::PreComp{T}
    d,n = size(x)
    rez = ones(T,(n,max_p...))
    MP = Base.maximum(max_p)

    # println("Computing exponentials...")
    exponentials = dropdims(exp.(-sum(x,dims=1)),dims=1)

    laguerre_L = zeros(T,(d,MP,n))
    # println("Computing powers...")
    powers = x[:,na,:] .^ (0:(MP-1))[na,:,na]
    # println("Computing laguerre_L")
    Threads.@threads for p in 1:MP
        laguerre_L[:,p:p,:] = dropdims(sum(P.LAGUERRE[1:p,p:p][na,:,:,na] .* powers[:,1:p,na,:],dims=2),dims=2)
        print(p,"\n")
    end

    # println("Computing L(2x)")
    Threads.@threads for p in CartesianIndices(max_p)
        for i in 1:d
            rez[[CartesianIndex((i,Tuple(p)...)) for i in 1:n]] .*= laguerre_L[i,p[i],:]
        end
        print(p,"\n")
    end
    return rez .* sqrt(T(2))^d .* exponentials
end

"""
    empirical_coefs(x,maxp)

Efficiently Compute the empirical laguerre coefficients of the density of the random vector x, given as a Matrix with n row (number of samples) and d columns (number of dimensions)
"""
function empirical_coefs(x,max_p)
    T = Base.promote_eltype(x,[1.0])
    x = T.(x)
    y = laguerre_phi_several_pts(x,max_p)
    # println("Taking the means...")
    return dropdims(sum(y,dims=1)/size(y,1),dims=1)
end

function empirical_coefs(x,max_p::NTuple{2,M}) where M<:Int
    T = Base.promote_eltype(x,[1.0])
    x = T.(x)
    P = get_precomp(T,Base.maximum(max_p)+1)::PreComp{T}
    d,n = size(x)
    rez = ones(T,max_p)
    MP = Base.maximum(max_p)

    exponentials = dropdims(exp.(-sum(x,dims=1)),dims=1)

    laguerre_L = zeros(T,(d,MP,n))
    powers = x[:,na,:] .^ (0:(MP-1))[na,:,na]
    Threads.@threads for p in 1:MP
        laguerre_L[:,p:p,:] = dropdims(sum(P.LAGUERRE[1:p,p:p][na,:,:,na] .* powers[:,1:p,na,:],dims=2),dims=2)
    end

    Threads.@threads for p in CartesianIndices(max_p)
        rez[p] = sum(laguerre_L[1,p[1],:] .* laguerre_L[2,p[2],:] .* exponentials)
    end
    rez = (rez .* 2) ./ n
    return rez
end

"""
    L2Objective(par,emp_coefs)

A L2 distance to be minimized between laguerre coefficients of MultivariateGammaConvolutions and empirical laguerre coefficients.
This distance is some kind of very high degree polynomial * exponentials, so minimizing it is very hard.
We recomand using some global minimisation routine, we found the Particle Swarm optimizers from Optim.jl particularly efficient.
"""
function L2Objective(par, emp_coefs)
    m = size(emp_coefs)
    d = length(m)
    n = Int((size(par)[1])/(d+1))
    α = par[1:n] .^2 #make them positives
    rates = reshape(par[(n+1):(d+1)*n],(n,d)) .^ 2 # make them positives
    coefs = get_coefficients(α,rates,m)
    return sum((coefs .- emp_coefs) .^2)
end

function αθ_from_par(par)
    n = Int(length(par)/2)
    return par[1:n].^2, par[(n+1):2n].^2
end


text:
# Tests for angle SOCP constriants

using TrajOptSOCPs
using Test
using SparseArrays, LinearAlgebra

@testset "Simple Angle SOCP" begin
    sI = Diagonal([1; 0])
    tI = [0; 1]
    asocp = AL_simpleAngleCone(0.2, sI, tI)

    # The first and third should be satisified. The second should not be
    # satisfied.
    xTestArr = [[0.1; 2], [2; 2], [-0.4; 3]]
    ws = TrajOptSOCPs.whichSatisfied(asocp, xTestArr)

    @test ws[1]
    @test !ws[2]
    @test ws[3]

    gN1 = TrajOptSOCPs.getNormToProjVals(asocp, xTestArr[1])
    gN2 = TrajOptSOCPs.getNormToProjVals(asocp, xTestArr[2])

    @test gN1'gN1 == 0
    @test gN2'gN2 != 0

    # Now the calculus
    grad1 = TrajOptSOCPs.getGradC(asocp, xTestArr[1])
    hess1 = TrajOptSOCPs.getHessC(asocp, xTestArr[1])
    alhess1 = TrajOptSOCPs.getHessC_ALTerm(asocp, xTestArr[1])

    @test size(grad1, 1) == 2
    @test size(hess1) == (2, 2)
    @test size(alhess1) == (2, 2)

    grad2 = TrajOptSOCPs.getGradC(asocp, xTestArr[2])
    hess2 = TrajOptSOCPs.getHessC(asocp, xTestArr[2])
    alhess2 = TrajOptSOCPs.getHessC_ALTerm(asocp, xTestArr[2])

    @test size(grad2, 1) == 2
    @test size(hess2) == (2, 2)
    @test size(alhess2) == (2, 2)

end

@testset "Test Values of Angle SOCP" begin
    sI = Diagonal([1; 0])
    tI = [0; 1]

    alpha = TrajOptSOCPs.getAlpha(20, true)
    @test alpha == -tand(20)

    asocp = AL_simpleAngleCone(alpha, sI, tI)

    xTest1 = [1; -1] # Angle of 45 deg > 20 deg
    xTest2 = [0.18; -1] # Angle of ~10 deg < 20 deg
    xTest3 = [alpha; -1] # Exactly 20 deg

    @test !TrajOptSOCPs.satisfied(asocp, xTest1)
    @test TrajOptSOCPs.satisfied(asocp, xTest2)
    @test TrajOptSOCPs.satisfied(asocp, xTest3)

    # Test the constraint violation
    gN1 = TrajOptSOCPs.getNormToProjVals(asocp, xTest1)
    gN2 = TrajOptSOCPs.getNormToProjVals(asocp, xTest2)
    gN3 = TrajOptSOCPs.getNormToProjVals(asocp, xTest3)

    @test gN1'gN1 != 0.0
    @test gN2'gN2 == 0.0
    @test gN3'gN3 == 0.0

end

@testset "Test Calculus of Angle SOCP" begin
    sI = Diagonal([1; 0])
    tI = [0; 1]

    alpha = TrajOptSOCPs.getAlpha(20, true)

    asocp = AL_simpleAngleCone(alpha, sI, tI)

    xTest1 = [1; -1] # Angle of 45 deg > 20 deg
    xTest2 = [0.18; -1] # Angle of ~10 deg < 20 deg
    xTest3 = [alpha; -1] # Exactly 20 deg

    # Test the constraint violation
    g1 = TrajOptSOCPs.getGradC(asocp, xTest1)
    g2 = TrajOptSOCPs.getGradC(asocp, xTest2)
    g3 = TrajOptSOCPs.getGradC(asocp, xTest3)

    @test size(g1, 1) == 2
    @test size(g2, 1) == 2
    @test size(g3, 1) == 2

    @test g1 == [1; -alpha]
    @test g2 == [0.0; 0.0] #[1; -alpha]
    @test g3 == [0.0; 0.0] #[sign(alpha); -alpha]

    h1 = TrajOptSOCPs.getHessC(asocp, xTest1)
    h2 = TrajOptSOCPs.getHessC(asocp, xTest2)
    h3 = TrajOptSOCPs.getHessC(asocp, xTest3)

    @test h1 == h2
    @test h2 == h3
    @test h3 == spzeros(2, 2)

    ah1 = TrajOptSOCPs.getHessC_ALTerm(asocp, xTest1)
    ah2 = TrajOptSOCPs.getHessC_ALTerm(asocp, xTest2)
    ah3 = TrajOptSOCPs.getHessC_ALTerm(asocp, xTest3)

    @test ah1 == [1 -alpha; -alpha alpha^2]
    @test ah2 == ah3
    @test h3 == ah3

end

@testset "Test Calculus of Angle SOCP 3D" begin
    sI = Diagonal([1; 0; 1])
    tI = [0; 1; 0]

    alpha = TrajOptSOCPs.getAlpha(20, true)

    asocp = AL_simpleAngleCone(alpha, sI, tI)

    xTest1 = [1; -1; 1] # Angle of 45 deg > 20 deg
    xTest2 = [0.125; -1; 0.125] # Angle of ~10 deg < 20 deg
    xTest3 = [0; -1; alpha] # Exactly 20 deg
    xTest4 = [0.1; -0.2; 0.3] # Angle of ~58 deg > 20 deg

    @test !TrajOptSOCPs.satisfied(asocp, xTest4)

    # Test the constraint violation
    g1 = TrajOptSOCPs.getGradC(asocp, xTest1)
    g2 = TrajOptSOCPs.getGradC(asocp, xTest2)
    g3 = TrajOptSOCPs.getGradC(asocp, xTest3)
    g4 = TrajOptSOCPs.getGradC(asocp, xTest4)

    @test size(g1, 1) == 3
    @test size(g2, 1) == 3
    @test size(g3, 1) == 3
    @test size(g4, 1) == 3

    @test g1 == [1 / sqrt(2); -alpha; 1 / sqrt(2)]
    @test g2 == zeros(3) #[1 / sqrt(2); -alpha; 1 / sqrt(2)]
    @test g3 == zeros(3) #[0; -alpha; sign(alpha)]
    @test g4 == [0.1 / norm([0.1; 0.3]); -alpha; 0.3 / norm([0.1; 0.3])]

    h1 = TrajOptSOCPs.getHessC(asocp, xTest1)
    h2 = TrajOptSOCPs.getHessC(asocp, xTest2)
    h3 = TrajOptSOCPs.getHessC(asocp, xTest3)
    h4 = TrajOptSOCPs.getHessC(asocp, xTest4)

    @test h1 == h2
    @test h2 == h3
    @test h3 == h4
    @test h4 == spzeros(3, 3)

    ah1 = TrajOptSOCPs.getHessC_ALTerm(asocp, xTest1)
    ah2 = TrajOptSOCPs.getHessC_ALTerm(asocp, xTest2)
    ah3 = TrajOptSOCPs.getHessC_ALTerm(asocp, xTest3)
    ah4 = TrajOptSOCPs.getHessC_ALTerm(asocp, xTest4)

    sq2 = - alpha / sqrt(2)
    ah1_check = [1/2 sq2 1/2; sq2 alpha^2 sq2; 1/2 sq2 1/2]

    ns4 = norm([0.1; 0.3])
    sq4 = - alpha / norm([0.1; 0.3])

    ah4_check = [((0.1^2) / ns4^2) (0.1 * sq4) ((0.1 * 0.3) / ns4^2);
                (0.1 * sq4)         alpha^2         (0.3 * sq4);
                ((0.1 * 0.3) / ns4^2) (0.3 * sq4) ((0.3^2) / ns4^2)]

    @test round.(ah1, digits = 8) == round.(ah1_check, digits = 8)
    @test ah2 == ah3
    @test h3 == ah3
    @test round.(ah4, digits = 8) == round.(ah4_check, digits = 8)

end

@testset "Multiple Angle SOCP" begin
    sI = Diagonal([1; 0])
    tI = [0; 1]
    alpha = 0.2

    nDim = 2
    indicator = findnz(sparse(
                    [0; 0; 0; 0; 1; 0; 0; 0; 0; 0; 1; 0; 0; 0; 0; 0; 1; 0]))[1]

    aMany = TrajOptSOCPs.makeAL_Multiple_AngleCone(alpha, sI, tI,
                                                    nDim, indicator)

    # Check that it worked
    @test typeof(aMany) == AL_Multiple_AngleCone

    xTest = [1; 2; 3; 4; 5; 6; 6; 5; 4; 3; 2; 1; 1; 2; 3; 4; 5; 6]

    # The code needs to notice that there are three vectors to extract and then
    # get those and evaluate it.
    rawVec = TrajOptSOCPs.getRaw(aMany, xTest)

    @test size(rawVec, 1) == 3

    gN = TrajOptSOCPs.getNormToProjVals(aMany, xTest)

    @test size(gN, 1) == 3
    @test gN'gN != 0.0

    jacobMany = TrajOptSOCPs.getGradC(aMany, xTest)
    hessMany = TrajOptSOCPs.getHessC(aMany, xTest)
    ahessMany = TrajOptSOCPs.getHessC_ALTerm(aMany, xTest)

    xSize = size(xTest, 1)

    @test size(jacobMany) == (3, xSize)
    @test size(hessMany) == (xSize, xSize)
    @test size(ahessMany) == (xSize, xSize)

end

max_stars_repo_name: danineamati/TrajOptSOCPs


text:
# TODO: this file is pretty sparse...

using InteractiveUtils: subtypes

inferred_retval(f, sig) = code_typed(f, sig)[1].first.code[end].args[1]

@context CanOverdubTestCtx
ctx = CanOverdubTestCtx()
# test that the result value of `canoverdub` is inferred exactly for `T <: Core.Builtin`
for T in subtypes(Core.Builtin)
    if !(T <: typeof(Core._apply))
        @test !inferred_retval(canoverdub, (typeof(ctx), T))
        @test !inferred_retval(canoverdub, (typeof(ctx), typeof(Core._apply), T))
        @test !inferred_retval(canoverdub, (typeof(ctx), typeof(Core.invoke), T))
    end
end

@test canoverdub(ctx, hypot, 1, 2)
@test canoverdub(ctx, Core.invoke, hypot, Tuple{Int,Int}, 1, 2)
@test canoverdub(ctx, Core._apply, hypot, (1, 2))
@test canoverdub(ctx, Core._apply, Core.invoke, (hypot, Tuple{Int,Int}, 1, 2))