text:
######## THIS FILE IS THE SAME AS `haplotype_pair.jl`
######## except it solves the least squares objective using a
######## heuristic stepwise search

function haplopair_stepscreen!(
    X::AbstractMatrix,
    H::AbstractMatrix;
    inv_sqrt_allele_var::Union{AbstractVector, Nothing} = nothing,
    r::Int = 1,
    # preallocated vectors
    happair1::AbstractVector          = ones(Int32, size(X, 2)),            # length n
    happair2::AbstractVector          = ones(Int32, size(X, 2)),            # length n
    hapscore::AbstractVector          = Vector{Float32}(undef, size(X, 2)), # length n
    maxindx ::AbstractVector{Int32}   = Vector{Int32}(undef, r),            # length r
    maxgrad ::AbstractVector{Float32} = Vector{Float32}(undef, r),          # length r
    # preallocated matrices
    Xwork :: AbstractMatrix{Float32} = Matrix{Float32}(undef, size(X, 1), size(X, 2)), # p × n
    Hwork :: AbstractMatrix{Float32} = convert(Matrix{Float32}, H),                 # p × d
    M  :: AbstractMatrix{Float32} = Matrix{Float32}(undef, size(H, 2), size(H, 2)), # d × d
    Nt :: AbstractMatrix{Float32} = Matrix{Float32}(undef, size(H, 2), size(X, 2)), # d × n
    )

    p, n  = size(X)
    d     = size(H, 2)

    # global search
    if r ≥ d
        return haplopair!(Xw_aligned, Hw_aligned, happair1=happair1,
            happair2=happair2, hapscore=hapscore, Xwork=Xwork, Hwork=Hwork,
            M=M, N=Transpose(Nt))
    end

    # create views for matrices
    t6 = @elapsed begin
        if size(Xwork, 1) != size(H, 1) # last window
            Xwork = Matrix{Float32}(undef, p, n)
            Hwork = Matrix{Float32}(undef, p, d)
        end
        Mwork = view(M, 1:d, 1:d)
        Ntwork = view(Nt, 1:d, :)
        Hwork_view = view(Hwork, :, 1:d)
        copyto!(Hwork_view, H)
    end

    # initialize missing data
    t5 = @elapsed initXfloat!(Xwork, X)

    # assemble M (symmetric)
    t2 = @elapsed begin
        if !isnothing(inv_sqrt_allele_var)
            Hwork_view .*= inv_sqrt_allele_var # wᵢ = 1/√2p(1-p)
        end
        mul!(Mwork, Transpose(Hwork_view), Hwork_view)
        for j in 1:d, i in 1:(j - 1) # off-diagonal
            @inbounds Mwork[i, j] = 2Mwork[i, j] + Mwork[i, i] + Mwork[j, j]
        end
        for j in 1:d # diagonal
            @inbounds Mwork[j, j] *= 4
        end
        LinearAlgebra.copytri!(Mwork, 'U')

        # assemble N
        if !isnothing(inv_sqrt_allele_var)
            Hwork_view .*= inv_sqrt_allele_var # wᵢ = 1/2p(1-p)
        end
        mul!(Ntwork, Transpose(Hwork_view), Xwork)
        @simd for I in eachindex(Ntwork)
            @inbounds Ntwork[I] *= 2
        end
    end

    # computational routine
    t3 = @elapsed begin
        if r == 1
            haplopair_stepwise!(happair1, happair2, hapscore, Mwork, Ntwork)
        else
            haplopair_topr!(happair1, happair2, hapscore, Mwork, Ntwork, r,
                maxindx, maxgrad)
        end
    end

    # supplement the constant terms in objective
    t2 += @elapsed begin
        @inbounds for j in 1:n
            @simd for i in 1:p
                hapscore[j] += abs2(Xwork[i, j])
            end
        end
    end

    t1 = t4 = 0.0 # no haplotype rescreening or computing dist(X, H)

    return t1, t2, t3, t4, t5, t6
end

function haplopair_stepwise!(
    happair1 :: AbstractVector{<:Integer},
    happair2 :: AbstractVector{<:Integer},
    hapmin   :: AbstractVector{T},
    M        :: AbstractMatrix{T}, # d x d
    Nt       :: AbstractMatrix{T}, # d x n
    ) where T <: Real
    d, n = size(Nt)
    fill!(hapmin, typemax(T))
    @inbounds for k in 1:n
        # find the first haplotype
        gmax, i1 = Nt[1, k], 1
        for i in 2:d
            g = Nt[i, k]
            if g > gmax
                gmax = g
                i1   = i
            end
        end
        happair1[k] = i1
        # find the optimal second haplotype given i1 using LS criterion
        for i in 1:d
            score = M[i, i1] - Nt[i, k]
            if score < hapmin[k]
                hapmin[k] = score
                happair2[k] = i
            end
        end
        hapmin[k] -= Nt[i1, k]
    end
    return nothing
end

"""
    findtopr!(A, col, maxval, index)

Find the largest `length(index)` elements of column `A[:, col]`. `val` is
filled with the found largest `r` elements in sorted order and `idx` is filled
with their corresponding indices.
"""
@inline function findtopr!(
    A      :: AbstractMatrix{T},
    col    :: Integer,
    maxval :: AbstractVector{T},
    index  :: AbstractVector{<:Integer}
    ) where T <: Real
    fill!(maxval, typemin(T))
    @inbounds for row in 1:size(A, 1)
        a = A[row, col]
        k = searchsortedfirst(maxval, a)
        if k > 1
            popinsert!(maxval, k-1, a)
            popinsert!(index, k-1, row)
        end
    end
    nothing
end

"""
    popinsert!(v, k, vk)

Move elements in `v[2:k]` to `v[1:k-1]` and insert `vk` at position `k` of
vector `v`.
"""
@inline function popinsert!(v::AbstractVector, k::Integer, vk)
    @inbounds for i in 1:k-1
        v[i] = v[i+1]
    end
    v[k] = vk
    v
end

# same as haplopair_stepwise! but searches top `r` largest gradient
function haplopair_topr!(
    happair1 :: AbstractVector{<:Integer},
    happair2 :: AbstractVector{<:Integer},
    hapmin   :: AbstractVector{T},
    M        :: AbstractMatrix{T}, # d x d
    Nt       :: AbstractMatrix{T}, # d x n
    r        :: Integer = 5,
    maxindx  :: Vector{<:Integer} = Vector{Int32}(undef, r),
    maxgrad  :: Vector{T} = Vector{T}(undef, r)
    ) where T <: Real
    d, n = size(Nt)
    fill!(hapmin, typemax(T))
    chunks = div(d, 4)
    @inbounds for k in 1:n
        # find the top r haplotypes
        findtopr!(Nt, k, maxgrad, maxindx)
        # for each top haplotype, find the optimal second one
        for riter in 1:r
            i  = 1
            i1 = maxindx[riter]
            gmax = maxgrad[riter]
            for chunk in 1:chunks
                score = M[i, i1] - gmax - Nt[i, k]
                if score < hapmin[k]
                    hapmin[k], happair1[k], happair2[k] = score, i1, i
                end
                score = M[i+1, i1] - gmax - Nt[i+1, k]
                if score < hapmin[k]
                    hapmin[k], happair1[k], happair2[k] = score, i1, i+1
                end
                score = M[i+2, i1] - gmax - Nt[i+2, k]
                if score < hapmin[k]
                    hapmin[k], happair1[k], happair2[k] = score, i1, i+2
                end
                score = M[i+3, i1] - gmax - Nt[i+3, k]
                if score < hapmin[k]
                    hapmin[k], happair1[k], happair2[k] = score, i1, i+3
                end
                i += 4
            end
            # handle remaining terms
            while i ≤ d
                score = M[i, i1] - gmax - Nt[i, k]
                if score < hapmin[k]
                    hapmin[k], happair1[k], happair2[k] = score, i1, i
                end
                i += 1
            end
        end
    end
    return nothing
end

max_stars_repo_name: Hua-Zhou/MendelImpute


text:
function powerflow(
    polar::PolarForm,
    algo::AbstractNonLinearSolver;
    linear_solver=DirectSolver(),
)
    buffer = get(polar, PhysicalState())
    init_buffer!(polar, buffer)
    Jₓ = AutoDiff.Jacobian(polar, power_balance, State())
    return powerflow(polar, Jₓ, buffer, algo; linear_solver=linear_solver)
end

function powerflow(
    polar::PolarForm{T, IT, VT, MT},
    jacobian::AutoDiff.Jacobian,
    buffer::PolarNetworkState{IT,VT},
    algo::NewtonRaphson;
    linear_solver=DirectSolver(),
) where {T, IT, VT, MT}
    # Retrieve parameter and initial voltage guess
    Vm, Va, pbus, qbus = buffer.vmag, buffer.vang, buffer.pinj, buffer.qinj

    nbus = PS.get(polar.network, PS.NumberOfBuses())
    ngen = PS.get(polar.network, PS.NumberOfGenerators())
    npv = PS.get(polar.network, PS.NumberOfPVBuses())
    npq = PS.get(polar.network, PS.NumberOfPQBuses())
    n_states = get(polar, NumberOfState())
    nvbus = length(polar.network.vbus)
    ybus_re, ybus_im = get(polar.topology, PS.BusAdmittanceMatrix())

    ref = polar.indexing.index_ref
    pv = polar.indexing.index_pv
    pq = polar.indexing.index_pq

    # iteration variables
    iter = 0
    converged = false

    # indices
    j1 = 1
    j2 = npv
    j3 = j2 + 1
    j4 = j2 + npq
    j5 = j4 + 1
    j6 = j4 + npq

    # form residual function directly on target device
    F = buffer.balance
    dx = buffer.dx
    fill!(F, zero(T))
    fill!(dx, zero(T))

    # Evaluate residual function
    power_balance(polar, F, buffer)

    # check for convergence
    normF = norm(F, Inf)
    if algo.verbose >= VERBOSE_LEVEL_LOW
        @printf("Iteration %d. Residual norm: %g.\n", iter, normF)
    end
    if normF < algo.tol
        converged = true
    end

    linsol_iters = Int[]
    Vapv = view(Va, pv)
    Vapq = view(Va, pq)
    Vmpq = view(Vm, pq)
    dx12 = view(dx, j5:j6) # Vmqp
    dx34 = view(dx, j3:j4) # Vapq
    dx56 = view(dx, j1:j2) # Vapv

    @timeit TIMER "Newton" while ((!converged) && (iter < algo.maxiter))

        iter += 1

        @timeit TIMER "Jacobian" begin
            J = AutoDiff.jacobian!(polar, jacobian, buffer)
        end

        # Find descent direction
        if isa(linear_solver, LinearSolvers.AbstractIterativeLinearSolver)
            @timeit TIMER "Preconditioner" LinearSolvers.update!(linear_solver, J)
        end
        @timeit TIMER "Linear Solver" n_iters = LinearSolvers.ldiv!(linear_solver, dx, J, F)
        push!(linsol_iters, n_iters)

        # update voltage
        @timeit TIMER "Update voltage" begin
            # Sometimes it is better to move backward
            if (npv != 0)
                # Va[pv] .= Va[pv] .+ dx[j5:j6]
                Vapv .= Vapv .- dx56
            end
            if (npq != 0)
                # Vm[pq] .= Vm[pq] .+ dx[j1:j2]
                Vmpq .= Vmpq .- dx12
                # Va[pq] .= Va[pq] .+ dx[j3:j4]
                Vapq .= Vapq .- dx34
            end
        end

        fill!(F, zero(T))
        @timeit TIMER "Residual function" begin
            power_balance(polar, F, buffer)
        end

        @timeit TIMER "Norm" normF = xnorm(F)
        if algo.verbose >= VERBOSE_LEVEL_LOW
            @printf("Iteration %d. Residual norm: %g.\n", iter, normF)
        end

        if normF < algo.tol
            converged = true
        end
    end

    if algo.verbose >= VERBOSE_LEVEL_HIGH
        if converged
            @printf("N-R converged in %d iterations.\n", iter)
        else
            @printf("N-R did not converge.\n")
        end
    end

    # Timer outputs display
    if algo.verbose >= VERBOSE_LEVEL_MEDIUM
        show(TIMER)
        println("")
    end
    return ConvergenceStatus(converged, iter, normF, sum(linsol_iters))
end

ext: jl


text:
struct CloseOpen{L <: Union{Int,StaticInt}, U <: Union{Int,StaticInt}} <: AbstractUnitRange{Int}
    start::L
    upper::U
    @inline CloseOpen(s::StaticInt{L}, u::StaticInt{U}) where {L,U} = new{StaticInt{L},StaticInt{U}}(s, u)
    @inline CloseOpen(s::Integer, u::StaticInt{U}) where {U} = new{Int,StaticInt{U}}(s % Int, u)
    @inline CloseOpen(s::StaticInt{L}, u::Integer) where {L} = new{StaticInt{L},Int}(s, u % Int)
    @inline CloseOpen(s::Integer, u::Integer) = new{Int,Int}(s % Int, u % Int)
end
@inline CloseOpen(len::Integer) = CloseOpen(Zero(), len)

@inline Base.first(r::CloseOpen) = r.start
@inline Base.step(::CloseOpen) = One()
@inline Base.last(r::CloseOpen) = r.upper - One()
@inline Base.length(r::CloseOpen) = r.upper - r.start
@inline Base.length(r::CloseOpen{Zero}) = r.upper

@inline Base.iterate(r::CloseOpen) = (i = Int(first(r)); (i, i))
@inline Base.iterate(r::CloseOpen, i::Int) = (i += 1) ≥ r.upper ? nothing : (i, i)

ArrayInterface.known_first(::Type{<:CloseOpen{StaticInt{F}}}) where {F} = F
ArrayInterface.known_step(::Type{<:CloseOpen}) = 1
ArrayInterface.known_last(::Type{<:CloseOpen{<:Any,StaticInt{L}}}) where {L} = L - 1
ArrayInterface.known_length(::Type{CloseOpen{StaticInt{F},StaticInt{L}}}) where {F,L} = L - F

@inline canonicalize_range(r::OptionallyStaticUnitRange) = r
@inline canonicalize_range(r::CloseOpen) = r
@inline canonicalize_range(r::AbstractUnitRange) = maybestaticfirst(r):maybestaticlast(r)
@inline canonicalize_range(r::CartesianIndices) = CartesianIndices(map(canonicalize_range, r.indices))

Base.IteratorSize(::Type{<:CloseOpen}) = Base.HasShape{1}()
Base.IteratorEltype(::Type{<:CloseOpen}) = Base.HasEltype()
@inline Base.size(r::CloseOpen) = (length(r),)
Base.eltype(::CloseOpen) = Int


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

@recipe function f(events::DetectorHitEvents)

    edep_flat = collect(flatview(events.edep))
    edep_unit = unit(eltype(edep_flat))
    edep_unitless = Array(ustrip(edep_flat))

    layout := (2,2)
    unitformat --> :square

    pos_flat = flatview(collect(flatview(events.pos)))
    pos_x = pos_flat[1,:]
    pos_y = pos_flat[2,:]
    pos_z = pos_flat[3,:]

    @series begin
        subplot := 1
        seriestype := :histogram2d
        weights := edep_unitless
        aspect_ratio := 1.0

        bins --> 500

        #title := "Detector Hits, XY"
        xguide := "x"
        yguide := "y"

        (pos_x, pos_y)
    end

    @series begin
        subplot := 2
        seriestype := :histogram2d
        weights := edep_unitless
        aspect_ratio := 1.0

        bins --> 500

        #title := "Detector Hits, ZY"
        xguide := "z"
        yguide := "y"

        (pos_z, pos_y)
    end

    @series begin
        subplot := 3
        seriestype := :histogram2d
        weights := edep_unitless
        aspect_ratio := 1.0

        bins --> 500

        # title := "Detector Hits, XZ"
        xguide := "x"
        yguide := "z"

        (pos_x, pos_z)
    end

    @series begin
        subplot := 4
        seriestype := :stephist

        bins --> 1000
        yscale := :log10

        normalize := :density
        legend := false

        # title:="Energy Spectrum"
        xguide := "E"
        yguide := edep_unit == NoUnits ? "cts / E" : "cts / $edep_unit"

        sum.(events.edep)
    end
end

@recipe function f(wf::RDWaveform)

    seriestype = get(plotattributes, :seriestype, :line)

    if seriestype in (:line, :scatter)
        @series begin
            seriestype := seriestype
            xguide --> "t"
            yguide --> "sample value"
            unitformat --> :square
            wf.time, wf.value
        end

    #=
    elseif seriestype == :histogram2d
        n = length(linearindices(waveforms))
        h = pulse_hist(waveforms, ybinning)
        @series begin
            seriestype --> :bins2d
            x := h.edges[1]
            y := h.edges[2]
            z := Surface(h.weights)
            title --> "($n waveforms)"
            xlabel --> X_label
            ylabel --> Y_label
            ()
        end
    end
    =#
    else
        error("seriestype $seriestype not supported")
    end
end

@recipe function f(wfs::ArrayOfRDWaveforms)
    for wf in wfs
        @series begin
            wf
        end
    end
end
@recipe function f(wfs::Array{<:RDWaveform})
    @series begin
        ArrayOfRDWaveforms(wfs)
    end
end


text:
@testset "test eigen $T by syevd $interface" for T in [Float32, Float64, ComplexF32, ComplexF64], interface in ["CPU", "GPU"]

    # randomly generate a 2 by 2 matrix for testing
    size = 2
    matrixToTest = rand(T, size, size)
    matrixToTest .+= matrixToTest'

    # use the default Linear Algebra lib to calculate the right answer for testing
    right_answer = eigvals(matrixToTest)
    S = right_answer

    # to test the GPU interface, one should convert the matrix data to cuda
    if interface == "GPU"
        matrixToTest = CuArray(matrixToTest)
    end

    # initialize the MAGMA lib, serving as a necessary part before working
    magma_init()

    # call the basic (overloaded) wrapper gesvd! for syevd subroutines
    result = magma_syevd!('V','U',matrixToTest)

    # in the result, the wanted answer lies in the second position
    s = result[1]

    # finalize the MAGMA lib, serving as a necessary part after working
    magma_finalize()

    # if S is approximately equal to s, we defined then it's alright
    @test S ≈ s

end


text:
SensitivityAlg(args...;kwargs...) = @error("The SensitivtyAlg choice mechanism was completely overhauled. Please consult the local sensitivity documentation for more information")

abstract type AbstractForwardSensitivityAlgorithm{CS,AD,FDT} <: DiffEqBase.AbstractSensitivityAlgorithm{CS,AD,FDT} end
abstract type AbstractAdjointSensitivityAlgorithm{CS,AD,FDT} <: DiffEqBase.AbstractSensitivityAlgorithm{CS,AD,FDT} end

struct ForwardSensitivity{CS,AD,FDT} <: AbstractForwardSensitivityAlgorithm{CS,AD,FDT}
  autojacvec::Bool
end
Base.@pure function ForwardSensitivity(;
                                       chunk_size=0,autodiff=true,
                                       diff_type=Val{:central},
                                       autojacvec=autodiff)
  ForwardSensitivity{chunk_size,autodiff,diff_type}(autojacvec)
end

struct ForwardDiffSensitivity{CS,CTS} <: AbstractForwardSensitivityAlgorithm{CS,Nothing,Nothing}
end
Base.@pure function ForwardDiffSensitivity(;chunk_size=0,convert_tspan=true)
  ForwardDiffSensitivity{chunk_size,convert_tspan}()
end

struct BacksolveAdjoint{CS,AD,FDT} <: AbstractAdjointSensitivityAlgorithm{CS,AD,FDT}
  autojacvec::Bool
end
Base.@pure function BacksolveAdjoint(;chunk_size=0,autodiff=true,
                                      diff_type=Val{:central},
                                      autojacvec=autodiff)
  BacksolveAdjoint{chunk_size,autodiff,diff_type}(autojacvec)
end

struct InterpolatingAdjoint{CS,AD,FDT} <: AbstractAdjointSensitivityAlgorithm{CS,AD,FDT}
  autojacvec::Bool
  checkpointing::Union{Nothing,Bool}
end
Base.@pure function InterpolatingAdjoint(;chunk_size=0,autodiff=true,
                                         diff_type=Val{:central},
                                         autojacvec=autodiff,
                                         checkpointing=nothing)
  InterpolatingAdjoint{chunk_size,autodiff,diff_type}(autojacvec,checkpointing)
end

struct QuadratureAdjoint{CS,AD,FDT} <: AbstractAdjointSensitivityAlgorithm{CS,AD,FDT}
  autojacvec::Bool
end
Base.@pure function QuadratureAdjoint(;chunk_size=0,autodiff=true,
                                         diff_type=Val{:central},
                                         autojacvec=autodiff)
  QuadratureAdjoint{chunk_size,autodiff,diff_type}(autojacvec)
end

struct TrackerAdjoint <: AbstractAdjointSensitivityAlgorithm{nothing,true,nothing} end
struct ZygoteAdjoint <: AbstractAdjointSensitivityAlgorithm{nothing,true,nothing} end

@inline convert_tspan(::ForwardDiffSensitivity{CS,CTS}) where {CS,CTS} = CTS
@inline alg_autodiff(alg::DiffEqBase.AbstractSensitivityAlgorithm{CS,AD,FDT}) where {CS,AD,FDT} = AD
@inline get_chunksize(alg::DiffEqBase.AbstractSensitivityAlgorithm{CS,AD,FDT}) where {CS,AD,FDT} = CS
@inline diff_type(alg::DiffEqBase.AbstractSensitivityAlgorithm{CS,AD,FDT}) where {CS,AD,FDT} = FDT
@inline get_jacvec(alg::DiffEqBase.AbstractSensitivityAlgorithm{CS,AD,FDT}) where {CS,AD,FDT} = alg.autojacvec
@inline ischeckpointing(alg::DiffEqBase.AbstractSensitivityAlgorithm{CS,AD,FDT}) where {CS,AD,FDT} = alg.checkpointing


text:
##############################################################################
##
## Solve fokker-planck equation
##
##############################################################################
function solve_fp!(fd::FDExp, sol::SolutionExp; maxit::Int64 = 1500, tolFK::Float64 = 1e-7)

    #== sizes ==#
    b, a, z = fd.b, fd.a, fd.z
    bn  = length(b)
    an  = length(a)
    zn  = length(fd.z)
    abn = length(b)*length(a)

    ΔTab = fd.ΔTab
    ΔTagrid = fd.ΔTagrid
    ΔTbgrid = fd.ΔTbgrid

    #== Put density in 3dim matrix ==#
    # gvec_old = vcat(sol.g...)
    # gvec_out = copy(gvec_old)
    #
    # ΔTab = zeros(gvec_old)
    # bvec = zeros(gvec_old)
    # avec = zeros(gvec_old)
    # ijk = zero(Int)
    #
    # for zi = 1:zn, abi = 1:abn
    #     ijk += 1
    #     ΔTab[ijk] = ΔTagrid[afromab(abi)] * ΔTbgrid[bfromab(abi)]
    #     bvec[ijk] = b[bfromab(abi)]
    #     avec[ijk] = a[afromab(abi)]
    # end

    #== build LHS ==#
    fillB!(fd, sol)
    B = fd.B

    #== build RHS ==#
    λoff  = fd.λoff
    invΔᴷ = fd.invΔᴷ
    C = invΔᴷ*eye(zn) + λoff

    distance = 1.0
    gout = hcat(sol.g...)
    gold = copy(gout)

    it = 1
    while distance>tolFK && it<=maxit
        g̃ = gold * C
        updateg!(gout, g̃, B)

        ## Aggregating household decisions ##
        # copy!(gvec_old, gold)
        # copy!(gvec_out, gout)

        # Eaold = sum(avec .* gvec_out .* ΔTab)
        # Eaout = sum(avec .* gvec_out .* ΔTab)
        # Eb = sum(bvec .* gvec .* ΔTab)

        distance = chebyshev(vec(gout), vec(gold))
        # distance = abs(Eaold - Eaout)

        if distance < tolFK
            println("kfe solved : $(it) iterations")
            # println(out_text,"kfe solved : $(it) iterations")
            break
        else
            (it % 50 == 0) && @printf("  density iteration %d, distance %.4f \n", it, distance)
            (it % 50 == 0) && @printf("   sum %.4f --> %.4f\n", sum(gout'*ΔTab), sum(gold'*ΔTab) )
            # (it % 50 == 0) && @printf(out_text,"  density iteration %d, distance %.4f \n", it, distance)
            # (it % 50 == 0) && @printf(out_text,"   sum %.4f --> %.4f\n", sum(gout'*ΔTab), sum(gold'*ΔTab) )
        end

        copy!(gold, gout)
        it += 1
    end

    #== Update the converged density on solution ==#
    for zi=1:zn
        sol.g[zi] = gout[:,zi]
    end

    return Void
end

function updateg!(gout, g̃, B)

    zn = size(gout,2)
    for zi=1:zn
        #== Solve the system B(k) gⁿ⁺¹ (k) = C[] ==#
        gout[:,zi] = B[zi] \ g̃[:,zi]
    end
    return Void
end

function fillB!(fd::FDExp, sol::SolutionExp)

    abn = length(fd.a)*length(fd.b)
    zn  = length(fd.z)
    #= Construct a vector nb*na with the terms Δb̃[bi] * Δã[ai] =#
    ΔTab = fd.ΔTab

    #= D matrix =#
    D = ΔTab
    Dinv = 1./ΔTab
    # =========================================================================

    ### Allocate LHS/RHS matrices ###
    ## ........................................................................

    #== Back out information from fd ==#
    λdiag = fd.λdiag
    invΔᴷ = fd.invΔᴷ

    #== matrices to be filled ==#
    Au = sol.Au          # intensity matrix Au
    A  = sol.A           # use A from HJB as storage
    B  = fd.B            # storage matrix

    for zi = 1:zn
        Avals = nonzeros(A[zi])
        fill!(Avals, zero(Float64))

        Auvals = nonzeros(Au[zi])
        Arows = rowvals(Au[zi])
        for abi in 1:abn
            for k in nzrange(Au[zi], abi)
                row = Arows[k]
                current   = Dinv[row]*Auvals[k]*D[abi]
                Avals[k] = current
            end
        end
    end

    @inbounds for zi in 1:zn

        Avals = nonzeros(A[zi])
        Bvals = nonzeros(B[zi])

        #== erase Bvals ==#
        fill!(Bvals, zero(Float64))

        Brows = rowvals(B[zi])
        for abi in 1:abn
            for k in nzrange(B[zi], abi)
                # loop over elements in the column ij
                row      = Brows[k]
                Bvals[k] = -Avals[k]
                row == abi && (Bvals[k] += invΔᴷ - λdiag[zi])
                #NOTE:110 `λdiag` is entering only in this part of the code
            end
        end
    end
    # =========================================================================

    return Void
end

function solve_fp!(fd::FDImp; maxit::Int64 = 400, tolFK::Float64 = 1e-5)

    #== Back out information ==#
    invΔᴷ = fd.invΔᴷ
    Au = fd.Au          # use Au
    B = fd.B            # storage matrix

    ΔTbgrid = fd.ΔTbgrid
    ΔTagrid = fd.ΔTagrid

    abn = length(ΔTbgrid) * length(ΔTagrid)
    zn  = length(fd.z)

    #== Construct a vector nb*na with the terms Δb̃[bi] * Δã[ai] ==#
    ΔTab = zeros(fd.g⁰)
    ijk = zero(Int)
    for zi = 1:zn

        for abi = 1:abn
            ijk += 1
            ΔTab[ijk] = ΔTagrid[afromab(abi)] * ΔTbgrid[bfromab(abi)]
        end
    end
    D = ΔTab
    Dinv = 1./ΔTab

    Avals = nonzeros(Au)
    Arows = rowvals(Au)
    ijk = zero(Int)
    for zi in 1:zn, abi in 1:abn
        ijk += 1

        for k in nzrange(Au, ijk)
            row = Arows[k]
            current  = Dinv[row]*Avals[k]*D[ijk]
            Avals[k] = current
        end
    end

    Bvals = nonzeros(B); fill!(Bvals, zero(Float64))
    Brows = rowvals(B)
    ijk = zero(Int)
    @inbounds for zi in 1:zn, abi in 1:abn
        ijk += 1
        for k in nzrange(B, ijk)
            # loop over elements in the column ij
            row = Brows[k]
            Bvals[k] = - Avals[k]
            row == ijk && (Bvals[k] += invΔᴷ)
        end
    end

    gold = copy(fd.g⁰)
    gout = zeros(gold)
    it = 1
    distance = 1
    while distance>tolFK && it<=maxit
        g̃ = gold * invΔᴷ
        gout = B \ g̃

        distance = chebyshev(vec(gout), vec(gold))

        if it % 5 == 0
            @printf("Density iteration %d, distance %.4f \n", it, distance)

            #== Check if the summation is held constant across iterations ==#
            @printf("sum gold %.4f \n", sum(gold'*ΔTab))
            @printf("sum gout %.4f \n", sum(gout'*ΔTab))
        end
        copy!(gold, gout)
        it += 1
    end

    fd.g = gout

    return Void
end

# function solve_fp!(fd::TwoAssetsFD)
#
#     g⁰ = fd.g⁰
#     Δa = fd.Δa; Δb = fd.Δb
#
#     # == A' in the notes == #
#     A = fd.A
#
#     i_fix = 1
#     Aurows = rowvals(A)
#     Auvals = nonzeros(A)
#     @inbounds for ij in 1:size(A, 2)
#         for k in nzrange(A, ij)
#             Arows[k] == i_fix && ( Avals[k] = zero(Float64) )
#         end
#     end
#     A[i_fix,i_fix] = one(Float64)
#
#     # == solve system Ag = b == #
#     gg = A \ g⁰
#
#     total = 1/( Δb*Δa*sum(gg) )
#     fd.gg = total * gg
#
#     return Void
# end


text:
######################################################################
# diromni.jl
######################################################################

mutable struct DirOmni <: Sensor
    means::Vector{Float64}
    stds::Vector{Float64}

    function DirOmni(file::AbstractString)
        means = vec(readcsv(file)[:,2])
        stds = 2*ones(360)
        return new(means, stds)
    end
end

# returns a density
#function observe(m::SearchDomain, s::DirOmni, p::Pose)
function observe(theta::TargetTuple, s::DirOmni, p::Pose)
# determine the relative bearing
    rel_bearing = p[3] - true_bearing(p, theta)
    if rel_bearing < 0.0
        rel_bearing += 360.0
    end
    rel_int = round(Int, rel_bearing, RoundDown) + 1

    return s.means[rel_int] + s.stds[rel_int]*randn()
end

function O(s::DirOmni, theta::LocTuple, p::Pose, o)

    # calculate the relative int
    rel_bearing = p[3] - true_bearing(p, theta)
    if rel_bearing < 0.0
        rel_bearing += 360.0
    end
    rel_int = round(Int, rel_bearing, RoundDown) + 1

    # Calculate expected measurement
    o_diff = o - s.means[rel_int]

    # now look at probability
    d = Normal(0, s.stds[rel_int])
    #p = cdf(d, rel_end) - cdf(d, rel_start)
    p = pdf(d, o_diff)
    return p
end


text:
module BenchMixLookup
using BenchmarkTools
using SymSpellChecker
using SymSpellChecker: VerbosityALL, VerbosityCLOSEST

ASSETS_PATH = joinpath(@__DIR__, "..", "assets")

function load_noisy()
    map(readlines(joinpath(ASSETS_PATH, "noisy_query_en_1000.txt"))) do x
        String(strip(split(x)[1]))
    end
end

function mix(d, data, verbosity)
    sum(length(lookup(d, word, verbosity = verbosity)) for word in data)
end

dicts = [
    joinpath(ASSETS_PATH, "frequency_dictionary_en_30_000.txt"),
    joinpath(ASSETS_PATH, "frequency_dictionary_en_82_765.txt"),
    joinpath(ASSETS_PATH, "frequency_dictionary_en_500_000.txt"),
]

query1k = load_noisy()

suite = BenchmarkGroup()

for edit_distance in 1:3
    for prefix_length in 5:7
        d = SymSpell(dicts[2], max_dictionary_edit_distance = edit_distance, prefix_length = prefix_length)
        suite["mix_lookup_closest_$(edit_distance)_$(prefix_length)"] = @benchmarkable mix($d, $query1k, $VerbosityCLOSEST)
    end
end
end

BenchMixLookup.suite


text:
using SemiseparableMatrices, LazyArrays, MatrixFactorizations, Test

@testset "Inverse of banded" begin
    @testset "Bidiagonal" begin
        @testset "SemiseparableMatrix" begin
            n = 5
            L = Bidiagonal(randn(n), randn(n-1), :L)
            U = Bidiagonal(randn(n), randn(n-1), :U)
            Li = SemiseparableMatrix(ApplyArray(inv, L))
            Ui = SemiseparableMatrix(ApplyArray(inv, U))
            @test Li ≈ inv(L)
            @test Ui ≈ inv(U)
        end

        @testset "Derivation" begin
            n = 5
            L = Bidiagonal(ones(n), randn(n-1), :L)
            Li = ApplyArray(inv, L)
            ℓ = [1; cumprod(-L.ev)]
            @test tril(ℓ * inv.(ℓ)') ≈ Li

            L = Bidiagonal(randn(n), randn(n-1), :L)
            L̃ = Bidiagonal(ones(n), L.ev ./ L.dv[1:end-1], :L)
            # inv(L) == inv(L̃ * D) == inv(D) * inv(L̃)
            ℓ = [1; cumprod(-L.ev./L.dv[1:end-1])]
            @test tril(ℓ * inv.(ℓ)') ≈ inv(L̃)
            @test tril((inv.(L.dv) .* ℓ) * inv.(ℓ)') ≈ inv(L)

            n = 5
            U = Bidiagonal(ones(n), randn(n-1), :U)
            u = [1; cumprod(-U.ev)]
            @test triu(inv.(u) * u') ≈ inv(U)

            U = Bidiagonal(randn(n), randn(n-1), :U)
            u = [1; cumprod(-U.ev./U.dv[1:end-1])]
            @test triu(inv.(u) * (inv.(U.dv) .* u)') ≈ inv(U)
        end

        @testset "Tridiagonal inverse" begin
            n = 10
            A = Tridiagonal(randn(n-1), randn(n), randn(n-1)) + 10I
            U,L = ul(A, Val(false))
            Li = SemiseparableMatrix(ApplyArray(inv, L))
            Ui = SemiseparableMatrix(ApplyArray(inv, U))
            Ai = inv(A)
            @test Li * Ui ≈ Ai
            @test rank(triu(Ai)[1:5,5:end]) == 1
            a,b = Li.L.args; b = b'
            c,d = Ui.U.args; d = d'
            @test Ai[1,1] ≈ a[1]b[1]c[1]d[1]
            @test Ai[1,2] ≈ a[1]b[1]c[1]d[2]
            @test Ai[1,3] ≈ a[1]b[1]c[1]d[3]
            @test Ai[1,5] ≈ a[1]b[1]c[1]d[5]
            @test Ai[2,1] ≈ a[2]b[1]c[1]d[1]
            @test Ai[2,2] ≈ a[2]b[1]c[1]d[2]+a[2]b[2]c[2]d[2]
            @test Ai[2,3] ≈ a[2]b[1]c[1]d[3]+a[2]b[2]c[2]d[3]
            @test Ai[3,1] ≈ a[3]b[1]c[1]d[1]
            @test Ai[3,2] ≈ a[3]b[1]c[1]d[2]+a[3]b[2]c[2]d[2]
            @test Ai[3,3] ≈ a[3]b[1]c[1]d[3]+a[3]b[2]c[2]d[3]+a[3]b[3]c[3]d[3]
            @test Ai[3,4] ≈ a[3]b[1]c[1]d[4]+a[3]b[2]c[2]d[4]+a[3]b[3]c[3]d[4]

            @test Ai[1:3,3:5] ≈ [a[1]b[1]c[1]; a[2]*(b[1]c[1]+b[2]c[2]); a[3]*(b[1]c[1]+b[2]c[2]+b[3]c[3])]*d[3:5]'
            @test Ai[3:5,1:3] ≈ a[3:5] * [b[1]c[1]d[1] (b[1]c[1]+b[2]c[2])*d[2] (b[1]c[1]+b[2]c[2]+b[3]c[3])*d[3]]

            b̃ = cumsum(b .* c) .* d
            c̃ = a .* cumsum(b .* c)
            @test tril(a*b̃') ≈ tril(Ai)
            @test triu(c̃*d') ≈ triu(Ai)
        end
    end
end


text:
BlackBoxOptim.MethodNames
BlackBoxNames=Array{String}(undef,0)
push!(BlackBoxNames,"Adapt. Diff. Evol. (rand 1 bin)")
push!(BlackBoxNames,"Adapt. Diff. Evol. (rand 1 bin, radious limited)")
push!(BlackBoxNames,"Borg Multiobj. Evol. Alg.")
push!(BlackBoxNames,"Diff. Evol. (rand 1 bin)")
push!(BlackBoxNames,"Diff. Evol. (rand 1 bin, radious limited)")
push!(BlackBoxNames,"Diff. Evol. (rand 2 bin)")
push!(BlackBoxNames,"Diff. Evol. (rand 2 bin, radious limited)")
push!(BlackBoxNames,"Distance-weighted Exponential Nat. Evol. Strat.")
push!(BlackBoxNames,"Compass Coordinate Direct Search")
push!(BlackBoxNames,"Probabilistic Descent Direct search")
push!(BlackBoxNames,"Random Search")
push!(BlackBoxNames,"Resampling Inheritance Memetic Search")
push!(BlackBoxNames,"Resampling Memetic Search")
push!(BlackBoxNames,"Separable Natural Evolution Strategies")
push!(BlackBoxNames,"Simultaneous Perturbation Stochastic Approximation")
push!(BlackBoxNames,"Exponential Natural Evolution Strategies")

ibb2=[ibb 17 18 19 20]
history_plot=history_all[1:iter,ibb2[1,:],:] # func_evals,n_optimizers,length(funcs)
n_optimizers=length(ibb2)
history_plotINIT=copy(history_plot)
for i=1:n_optimizers
    for j=1:size(history_plot,3)
        history_plotINIT[1:3,i,j].=maximum(history_plotINIT[1:3,:,j])
        mi=minimum(history_plotINIT[:,:,j])
        ma=maximum(history_plotINIT[:,:,j])
        # println(mi," ",ma)
        history_plot[:,i,j]=(history_plotINIT[:,i,j].-mi)./(ma.-mi)
    end
end
history_plot=mean(history_plot, dims=3)[:,:,1]
mi=minimum(history_plot[end,:])
ma=maximum(history_plot[end,:])
sR=history_plot[end,length(ibb)+1]
ssR=@sprintf("mi= %.2E ma= %.2E ITS= %.2E",mi,ma,sR)
println(ssR)
# display(plot(history_plot,width=1,label=[string.(ibb) "oht"],legend=:topleft,title=ssR))
display(plot(history_plot[:,1:length(ibb)],width=1,label=BlackBoxNames[ibb],legend=:topright,title=ssR,titlefont = font("Times", 9),size = (sizex, sizey)))
display(plot!(history_plot[:,length(ibb)+2],width=1,label="Nelder-Mead",legend=:topright,title=ssR,titlefont = font("Times", 9),size = (sizex, sizey)))
display(plot!(history_plot[:,length(ibb)+3],width=1,label="Particle Swarm",legend=:topright,title=ssR,titlefont = font("Times", 9),size = (sizex, sizey)))
display(plot!(history_plot[:,length(ibb)+4],width=1,label="Simulated Annealing",legend=:topright,title=ssR,titlefont = font("Times", 9),size = (sizex, sizey)))
display(plot!(history_plot[:,length(ibb)+1],width=2,label="Proposed-BNO",legendfontsize=8,color=:black,size = (sizex, sizey)))
display(xlims!(0,sf1*evals1))
display(xticks!(0:1000:evals1))
xaxis!("Function Evaluations")
yaxis!("Normalized Responces")
for i=[1:length(ibb);length(ibb)+2;length(ibb)+3;length(ibb)+4]
    println(minimum(history_plot[:,i]))
end
println(minimum(history_plot[:,length(ibb)+1]))

history_plot=history_all[1:iter,ibb2[1,:],:] # func_evals,n_optimizers,length(funcs)
n_optimizers=length(ibb2)
history_plotINIT=copy(history_plot)
for i=1:n_optimizers
    for j=1:size(history_plot,3)
        history_plotINIT[1:3,i,j].=maximum(history_plotINIT[1:3,:,j])
        mi=minimum(history_plotINIT[:,:,j])
        ma=maximum(history_plotINIT[:,:,j])
        history_plot[:,i,j]=(history_plotINIT[:,i,j].-mi)./(ma.-mi)
    end
end
# history_plot=(mean(history_plot, dims=3)[:,:,1].^(1/10))
history_plot=mean(history_plot, dims=3)[:,:,1].^(1/10)
# mi=minimum(history_plot[end,:])
# ma=maximum(history_plot[end,:])
# sR=history_plot[end,length(ibb)+1]
# ssR=@sprintf("mi= %.2E ma= %.2E ITS= %.2E",mi,ma,sR)
# println(ssR)
# display(plot(history_plot,width=1,label=[string.(ibb) "oht"],legend=:topleft,title=ssR))
display(plot(history_plot[:,1:length(ibb)],width=1,label=BlackBoxNames[ibb],legend=:false,titlefont = font("Times", 9),size = (sizex, sizey)))
display(plot!(history_plot[:,length(ibb)+2],width=1,label="Nelder-Mead",legend=:false,titlefont = font("Times", 9),size = (sizex, sizey)))
display(plot!(history_plot[:,length(ibb)+3],width=1,label="Particle Swarm",legend=:false,titlefont = font("Times", 9),size = (sizex, sizey)))
display(plot!(history_plot[:,length(ibb)+4],width=1,label="Simulated Annealing",legend=:false,titlefont = font("Times", 9),size = (sizex, sizey)))
display(plot!(history_plot[:,length(ibb)+1],width=2,label="Proposed-BNO",legend=:false,color=:black,size = (sizex, sizey)))
# display(xlims!(0,sf1*evals1))
display(xticks!(0:1000:evals1))
xaxis!("Function Evaluations")
yaxis!("Normalized Responces")
# xlims()


text:
using Debug

type FactorAnalysis
  Ψ::Vector{Float64}
  F::Matrix{Float64}
  Vbar::Vector{Float64}
end

n_components(fa::FactorAnalysis) = size(fa.F, 2)

const SMALL = 1e-12

@debug function fit(::Type{FactorAnalysis}, V::Matrix{Float64}, n_components::Int, tol::Float64=1e-2, max_iter::Int=1000, Ψ::Vector{Float64} = ones(Float64, size(V,1)))
  Vbar = vec(mean(V, 2))
  σ2 = vec(var(V, 2, corrected=false)) # corrected = false used to get same results as scikit-learn
  D, N = size(V)
  X = centralize(V, Vbar)
  L_old = -Inf
  L = -Inf
  F = zeros(D, n_components)
  k = 0
  for k = 1:max_iter
    Ψ_sqrt = sqrt(Ψ) + SMALL
    X_tilde = 1./Ψ_sqrt .* X/(sqrt(N))
    U, Λ_tilde, _ = svd(X_tilde, thin=true)
    Λ = Λ_tilde.^2
    U_H = U[:,1:n_components]
    Λ_H = Λ[1:n_components]
    # Factor update
    @bp
    F = Ψ_sqrt .* U_H * diagm(sqrt(max(Λ_H .- 1, 0)))
    # Log likelihood
    L_old = L
    L = n_components + log(prod(2π*Ψ))
    for i = 1:n_components
      @inbounds L += log(Λ[i])
    end
    for i = n_components+1:D
      @inbounds L += Λ[i]
    end
    L *= -N/2
    # Noise update
    Ψ = vec(max(σ2 - sum(F.^2, 2), SMALL))

    if L - L_old < tol
      info("Factor analysis converged, log-likelihood = $L")
      break
    end
  end
  if k == max_iter
    info("Factor analysis did not converge in $max_iter iterations.")
  end
  FactorAnalysis(Ψ, F, Vbar), L
end

function transform(fa::FactorAnalysis, data::Matrix{Float64})
# FIXME
  X = centralize(data, fa.Vbar)
  F_Ψ = fa.F ./ fa.Ψ
  cov_z = inv(eye(n_components(fa)) .+ F_Ψ' * fa.Ψ)
  cov_z' * F_Ψ' * X
end


text:
# FeynmanAthermalAsymptoticComparison.jl - reproduce some 1950s numeric results

push!(LOAD_PATH,"../src/") # load module from local directory
using PolaronMobility

using Plots

# Here we collect all the data we need for plotting
αrange=0.2:0.1:10 # unstable below α=0.2
vs=[ feynmanvw(α)[1]  for α in αrange] # takes a while; computes all vw
ws=[ feynmanvw(α)[2]  for α in αrange]
# Urgh - can't seem to do this more elegantly; recalculating all the params just to strip the second component of the tuple

Es=map(F,vs,ws,αrange); # calculate polaron energy for these params

p=plot(xlim=αrange,ylim=(0,),
    xlabel="α parameter", ylabel="Reduced polaron units",
    title="Athermal (numeric, variational) v,w, and E for α" )

plot!(αrange,vs,label="v")
plot!(αrange,ws,label="w")
plot!(αrange,Es,label="E")
hline!([0.0],label="",color=:black)

# I like this plot!
# You can see how v=w=3 as α-->0.0
# w then tends down to 1.0 in a nice sigmoid form as α --> ~10+
# v increases monotonically

savefig("AthermalvwE.pdf")

# Feynman Stat Mech (1972), p. 240
smallαw(α)=3.0
P(α)=2/smallαw(α) * (sqrt(1+smallαw(α))-1) # Nb: error in Feynman Stat Mech: sqrt(1+W) , as in Feynman1955
smallαv(α)=3*(1 + 2*α*(1-P(α))/(3*smallαw(α)))

c=0.5772 # Euler Mascheroni constant

largeαw(α)=1.0
largeαv(α)=(4*α^2/9π) - (4*(log(2) + c/2)-1)

p=plot(xlim=αrange,ylim=(0,15), xlabel="α parameter", ylabel="v,w, Reduced polaron units", title="Asymtotic limits for v,w" )

plot!(αrange,α->smallαw(α),label="w, smallαw",style=:dash)
plot!(αrange,α->smallαv(α),label="v, smallαv",style=:solid)

plot!(αrange,α->largeαw(α),label="w, largeαw",style=:dash)
plot!(αrange,α->largeαv(α),label="v, largeαv",style=:solid)

plot!(αrange,vs,label="v, numeric",w=3,style=:solid)
plot!(αrange,ws,label="w, numeric",w=3,style=:dash)

savefig("AthermalvwAsymptotic.pdf")

smallαE(α)=-α-α^2/81
FeynmanStatMechlargeαE(α)=-α^2/(3π) - 3/2 *(2*log(2)+c)-3/4 # As in Feynman Stat Mech; with 2π corrected to 3π ...

Feynman1955largeαE(α)=-α^2/3π - 3*log(2) # as in Feynman I, 1955

p=plot(xlim=αrange,ylim=(-Inf,0), xlabel="α parameter", ylabel="Energy, Reduced polaron units", title="Asymtotic limits for E(α)" )

plot!(αrange,α->smallαE(α),label="smallαE")
plot!(αrange,α->FeynmanStatMechlargeαE(α),label="FeynmanStatMechlargeαE")
plot!(αrange,α->Feynman1955largeαE(α),label="Feynman1955largeαE")
plot!(αrange,Es,label="E (numeric integration)")

savefig("AthermalAsymptoticEnergy.pdf")


text:
#module ArraySlices

using Compat.view

import Base: length, size, eltype, getindex

export slices, columns, rows

# Type parameters
#
# F : type of SubArray
# D : indexed dimension of the array
# A : array type
#
immutable SliceIterator{F, D, A<:AbstractArray} <: AbstractVector{F}
    array::A
end

# ensure compatibility between versions
if VERSION < v"0.5-dev"
    _get_L(D, N) = D == 1 ? N : D
else
    _get_L(D, N) = D == 1 || D == N ? true : false
end

"""
    slices(array, dim)
Return a `SliceIterator` object to loop over the slices of `array` along
dimension `dim`.
"""
@generated function slices{T, N, D}(array::AbstractArray{T, N}, ::Type{Val{D}})
    # checks
    1 <= D <= N || error("invalid slice dimension")

    # construct incomplete type of slice, then fill
    # example: SubArray{Float64, 1, Array{Float64,2}, Tuple{Int64, Colon}, true}
    F = :(SubArray{$T, $(N-1), $array})

    # construct tuple of indices
    tupexpr = :(Tuple{})
    for i = 1:N
        push!(tupexpr.args, :Colon)
    end
    tupexpr.args[1+D] = :Int
    push!(F.args, tupexpr)

    # add L/LD parameter
    push!(F.args, _get_L(D, N))

    # build and return iterator
    :(SliceIterator{$F, $D, $array}(array))
end

# ~~~ Array interface ~~~
eltype{F}(s::SliceIterator{F}) = F
length{F, D}(s::SliceIterator{F, D}) = size(s.array, D)
size(s::SliceIterator) = (length(s), )

# build code that produces slices with the correct indexing
@generated function getindex{F, D}(s::SliceIterator{F, D}, i::Integer)
    # get ndims of parent array
    N = s.parameters[1].parameters[3].parameters[2]
    expr = :()
    expr.head = :call
    push!(expr.args, :view)
    push!(expr.args, :(getfield(s, :array)))
    # fill in with `Colon`s
    for i = 1:N
        push!(expr.args, Colon())
    end
    # then replace the indexed dimension
    expr.args[2 + D] = :i
    expr
end

# ~~~ Convenience functions for 2D arrays ~~~

"""
    columns(array)
Return a `SliceIterator` object to loop over the columns of `array`.
"""
columns(array::AbstractMatrix) = slices(array, Val{2})

"""
    rows(array)
Return a `SliceIterator` object to loop over the rows of `array`.
"""
rows(array::AbstractMatrix) = slices(array, Val{1})

#end


text:
@testset "Linear Spaces" begin
    @testset "AffineSubspace" begin
        A = AffineSubspace([1 0 3; 2 1 3], [5, -2])
        @test dim(A) == dim(intrinsic(A)) == dim(extrinsic(A)) == 1
        @test codim(A) == codim(intrinsic(A)) == codim(extrinsic(A)) == 2
        @test ambient_dim(A) == 3
        @test startswith(sprint(show, A), "1-dim. affine subspace")
        @test identity.(A) == A

        @test intrinsic(A) isa AffineIntrinsic
        @test startswith(sprint(show, intrinsic(A)), "AffineIntrinsic")
        @test extrinsic(A) isa AffineExtrinsic
        @test startswith(sprint(show, extrinsic(A)), "AffineExtrinsic")

        @test identity.(extrinsic(A)) == extrinsic(A)
        @test identity.(intrinsic(A)) == intrinsic(A)
        @test identity.(Intrinsic) == Intrinsic

        u = rand(1)
        x = A(u, Intrinsic)
        @test coord_change(A, Extrinsic, Intrinsic, x) ≈ u rtol = 1e-12
        @test coord_change(A, Intrinsic, Extrinsic, u) ≈ x rtol = 1e-12
        @test norm(A(x, Extrinsic)) ≈ 0 atol = 1e-14
        @test norm(AffineExtrinsic(intrinsic(A))(x)) ≈ 0.0 atol = 1e-12

        B = rand_affine_subspace(3; codim = 2)
        @test B != A
        copy!(B, A)
        @test extrinsic(B) == extrinsic(A)
        @test intrinsic(B) == intrinsic(A)
        @test B == A
        @test copy(A) == A
        @test copy(A) !== A

        C = rand_affine_subspace(3, dim = 1, real = true)
        @test geodesic_distance(A, C) > 0
        γ = geodesic(A, C)
        @test size(γ(1)) == size(intrinsic(A).Y)
    end

    @testset "AffineSubspaceHomotopy" begin
        @var x[1:4]
        f1 = rand_poly(x, 2)
        f2 = rand_poly(x, 2)
        F = System([f1, f2], x)

        A = rand_affine_subspace(4; codim = 2)
        # Commpute witness set
        L = A(x, Extrinsic)
        res = track.(total_degree(F ∩ L)...)
        W = solution.(res)

        B = rand_affine_subspace(4, codim = 2)

        graff_tracker = Tracker(
            AffineSubspaceHomotopy(F, A, B),
            options = (automatic_differentiation = 3,),
        )
        graff_result = track.(graff_tracker, W)
        @test all(is_success, graff_result)

        C = rand_affine_subspace(4, codim = 2)
        copy_A = copy(A)
        set_subspaces!(graff_tracker.homotopy, B, C)
        @test A == copy_A
        @test all(is_success, track.(graff_tracker, solution.(graff_result)))

        graff_path_tracker = EndgameTracker(AffineSubspaceHomotopy(F, A, B))
        graff_path_result = track.(graff_path_tracker, W)
        @test all(is_success, graff_path_result)

        @test all(isapprox.(
            solution.(graff_path_result),
            solution.(graff_result),
            rtol = 1e-12,
        ))
    end
end


text:
module CodeEnvironments

using SHA
using Base64
using Pkg

export
    CodeEnvironment,
    use_environment,
    environment_code

# A type which allows us to hide the details of Manifest.toml on a single
# source code line in jupyter, etc.
struct Base64Data
    data::String
end
Base.convert(::Type{String}, data::Base64Data) = String(base64decode(data.data))

"""
    CodeEnvironment(; project, manifest, name="")

"""
struct CodeEnvironment
    project::String
    manifest::String
    name::String
    fullname::String
    function CodeEnvironment(project::String, manifest::String, name::String)
        fullname = "CodeEnvironment_"              *
                   (isempty(name) ? "" : name*"_") *
                   bytes2hex(sha1(project*manifest))
        new(project, manifest, name, fullname)
    end
end

function CodeEnvironment(; project, manifest, name="", extras...)
    if !isempty(extras)
        # `extras` is present for forward compatibility: We may later want
        # extra info in the environment, but we don't yet know what it may be.
        @warn "CodeEnvironment() constructor: ignoring extra key value pairs" extras
    end
    project_str = convert(String, project)
    manifest_str = convert(String, manifest)
    CodeEnvironment(project_str, manifest_str, name)
end

function Base.show(io::IO, env::CodeEnvironment)
    manifest_nlines = sum(isequal('a'), env.manifest)
    print(io, env.fullname, "\n",
          "With project declaration and explicit dependencies (Project.toml):\n",
          env.project,
          "(+ detailed manifest of ", manifest_nlines, " lines)\n")
end

function _get_paths(env)
    envpath = joinpath(Pkg.envdir(), env.fullname)
    project_path = joinpath(envpath, "Project.toml")
    manifest_path = joinpath(envpath, "Manifest.toml")
    (envpath, project_path, manifest_path)
end

function check_environment(env)
    envpath, project_path, manifest_path = _get_paths(env)
    if !isdir(envpath)
        @info "Environment $(env.fullname) is not instantiated"
        return
    end

    project_ok  = env.project == read(project_path, String)
    manifest_ok = env.manifest == read(manifest_path, String)

    if !project_ok || !manifest_ok
        @error """
               Your code environment is inconsistent - it appears you have
               $(!project_ok ?
                 "added or removed packages from it manually." :
                 "updated the package versions in it manually.")

               We have loaded your installed version of this environment anyway
               but it is no longer a faithful copy of the reference environment!

               To fix this you have two options:
               * You may *update* the environment used by this script by
                 pasting the output of `environment_code()` over the top of the
                 code which calls `use_environment()`.
               * You may *revert* the local changes to your installed packages
                 with `use_environment(code_env, overwrite=true)`.
               """ reference_environment=env
    end
end

"""
    use_environment(env; overwrite=false)

Activate the environment `env`, instantiating it if necessary (installing,
building and precompiling all packages), or checking the existing environment
with the given name for consistency.
"""
function use_environment(env::CodeEnvironment; overwrite=false)
    envpath, project_path, manifest_path = _get_paths(env)
    newenv = !isdir(envpath)
    if overwrite || newenv
        if newenv
            mkpath(envpath)
            @info "Instantiating CodeEnvironment at $envpath"
        end
        write(project_path, env.project)
        write(manifest_path, env.manifest)
        Pkg.activate(envpath)
        Pkg.instantiate()
        Pkg.API.precompile()
    else
        @info "Activating existing environment at $envpath"
        check_environment(env)
        Pkg.activate(envpath)
    end
    nothing
end

"""
    environment_code([env::CodeEnvironment])

Generate julia code which may be pasted into a script or jupyter notebook in
order to set the environment.

If `env` is omitted, clone the current environment. Be careful doing this with
the default environment - you may end up with a lot of packages which your
current project doesn't need!

## Examples:

To get the current environment and put it in the system clipboard for pasting:

    clipboard(environment_code())

For use in jupyter, you may to print it to the screen instead if you are
running your jupyter kernel remotely:

    println(environment_code())
"""
function environment_code(env::CodeEnvironment)
    # To hide the environment guff in vim, use `set foldmethod=marker`
    namestr = isempty(env.name) ? "" : "name=$(repr(env.name)),"
    @info "Generating code for environment" env
    """
    using CodeEnvironments # {{{ Code environment setup
    code_env = CodeEnvironment($(namestr)project=CodeEnvironments.Base64Data($(repr(base64encode(env.project)))),manifest=CodeEnvironments.Base64Data($(repr(base64encode(env.manifest))))) # }}}
    use_environment(code_env)
    """
end

function environment_code()
    project_path = Pkg.Types.find_project_file() # FIXME: Call a public Pkg function for this?
    @info "Seeding CodeEnvironment with current project file" project_path
    environment_code(load_environment(project_path))
end

function load_environment(project_path)
    manifest_path = joinpath(dirname(Pkg.Types.find_project_file()), "Manifest.toml")
    CodeEnvironment(project=read(project_path, String), manifest=read(manifest_path, String))
end

end # module


text:
@auto_hash_equals struct LOCF <: Imputor
    vardim::Int
    context::AbstractContext
end

"""
    LOCF(; vardim=2, context=Context())

Last observation carried forward (LOCF) iterates forwards through the `data` and fills
missing data with the last existing observation. The current implementation is univariate,
so each variable in a table or matrix will be handled independently.

See also:
- [NOCB](@ref): Next Observation Carried Backward

WARNING: missing elements at the head of the array may not be imputed if there is no
existing observation to carry forward. As a result, this method does not guarantee
that all missing values will be imputed.

# Keyword Arguments
* `vardim=2::Int`: Specify the dimension for variables in matrix input data
* `context::AbstractContext`: A context which keeps track of missing data
  summary information

# Example
```jldoctest
julia> using Impute: LOCF, Context, impute

julia> M = [1.0 2.0 missing missing 5.0; 1.1 2.2 3.3 missing 5.5]
2×5 Array{Union{Missing, Float64},2}:
 1.0  2.0   missing  missing  5.0
 1.1  2.2  3.3       missing  5.5

julia> impute(M, LOCF(; vardim=1, context=Context(; limit=1.0)))
2×5 Array{Union{Missing, Float64},2}:
 1.0  2.0  2.0  2.0  5.0
 1.1  2.2  3.3  3.3  5.5
```
"""
LOCF(; vardim=2, context=Context()) = LOCF(vardim, context)

function impute!(data::AbstractVector, imp::LOCF)
    imp.context() do c
        start_idx = findfirst(c, data) + 1
        for i in start_idx:lastindex(data)
            if ismissing(c, data[i])
                data[i] = data[i-1]
            end
        end

        return data
    end
end


text:
#=
VecE3: a 3d euclidean vector
=#

struct VecE3{R<:Real} <: VecE{3,R}
    x::R
    y::R
    z::R
end
VecE3() = VecE3(0.0,0.0,0.0)
VecE3(x::Integer, y::Integer, z::Integer) = VecE3(float(x), float(y), float(z))
VecE3(t::Tuple) = VecE3(promote(t...)...)

Base.show(io::IO, a::VecE3) = @printf(io, "VecE3(%.3f, %.3f, %.3f)", a.x, a.y, a.z)

function dist(a::VecE3, b::VecE3)
    Δx = a.x-b.x
    Δy = a.y-b.y
    Δz = a.z-b.z
    sqrt(Δx*Δx + Δy*Δy + Δz*Δz)
end
function dist2(a::VecE3, b::VecE3)
    Δx = a.x-b.x
    Δy = a.y-b.y
    Δz = a.z-b.z
    Δx*Δx + Δy*Δy + Δz*Δz
end

proj(a::VecE3, b::VecE3, ::Type{Float64}) = (a.x*b.x + a.y*b.y + a.z*b.z) / sqrt(b.x*b.x + b.y*b.y + b.z*b.z) # dot(a,b) / |b|
function proj(a::VecE3, b::VecE3, ::Type{VecE3})
    # dot(a,b) / dot(b,b) ⋅ b
    s = (a.x*b.x + a.y*b.y + a.z*b.z) / (b.x*b.x + b.y*b.y + b.z*b.z)
    VecE3(s*b.x, s*b.y, s*b.z)
end

lerp(a::VecE3, b::VecE3, t::Real) = VecE3(a.x + (b.x-a.x)*t, a.y + (b.y-a.y)*t, a.z + (b.z-a.z)*t)

function rot(a::VecE3, axis::VecE3, θ::Real)

    #=
    rotate the point a: (x,y,z) about the axis: (u,v,w) by the angle θ following the
    right hand rule
    =#

    x,y,z = a.x, a.y, a.z
    u,v,w = axis.x, axis.y, axis.z

    u² = u*u
    v² = v*v
    w² = w*w

    m = u² + v² + w²

    c = cos(θ)
    s = sin(θ)
    mc = 1.0 - c

    rtms = sqrt(m)*s

    magc = m*c
    ux_vy_wz_mc = (u*x + v*y + w*z)*mc

    new_x = (u*ux_vy_wz_mc + x*magc + rtms*(-w*y + v*z)) / m
    new_y = (v*ux_vy_wz_mc + y*magc + rtms*( w*x - u*z)) / m
    new_z = (w*ux_vy_wz_mc + z*magc + rtms*(-v*x + u*y)) / m

    VecE3(new_x, new_y, new_z)
end
function rot_normalized(a::VecE3, axis::VecE3, θ::Real)

    #=
    rotate the point a: (x,y,z) about the axis: (u,v,w) by the angle θ following the
    right hand rule

    Here we assume axis is normalized
    =#

    x,y,z = a.x, a.y, a.z
    u,v,w = axis.x, axis.y, axis.z

    u² = u*u
    v² = v*v
    w² = w*w

    m = u² + v² + w²

    c = cos(θ)
    s = sin(θ)
    mc = 1.0 - c

    ux_vy_wz_mc = (u*x + v*y + w*z)*mc

    new_x = u*ux_vy_wz_mc + x*c + (-w*y + v*z)*s
    new_y = v*ux_vy_wz_mc + y*c + ( w*x - u*z)*s
    new_z = w*ux_vy_wz_mc + z*c + (-v*x + u*y)*s

    VecE3(new_x, new_y, new_z)
end


text:
using Test
using TFRecords

include("testparsebuild.jl")
include("testreadwrite.jl")


text:
"""
    consistency_tests(
        rng::AbstractRNG,
        f::OILMM,
        f_naive::GaussianProcessProbabilisticProgramme,
        x::AbstractVector,
        Y::ColVecs{<:Real},
    )

Verify that an OILMM `f` is self-consistent and consistent with the naive GP implementation
`f_naive`.
"""
function consistency_tests(
    rng::AbstractRNG,
    f::OILMM,
    f_naive::GaussianProcessProbabilisticProgramme;
    x_tr::AbstractVector,
    x_te::AbstractVector,
    x_naive_tr::AbstractVector,
    x_naive_te::AbstractVector,
    y_tr::AbstractVector,
    y_te::AbstractVector,
    σ²::Real,
)
    # Set up the finite-dimensional marginals.
    fx_tr = f(x_tr, σ²)
    fx_tr_naive = f_naive(x_naive_tr, σ²)

    # Check log p(Ytr) agrees with naive implementation.
    @test logpdf(fx_tr, y_tr) ≈ logpdf(fx_tr_naive, y_tr)

    # Check that prior marginals agree with naive implementation.
    @test mean(fx_tr) ≈ mean(fx_tr_naive)
    @test var(fx_tr) ≈ var(fx_tr_naive)

    # Check that noise-free prior marginals agree with the naive implementation.
    oilmm_marginals = denoised_marginals(f(x_tr))
    naive_marginals = marginals(f_naive(x_naive_tr))
    @test mean.(oilmm_marginals) ≈ mean.(naive_marginals)
    @test std.(oilmm_marginals) ≈ std.(naive_marginals)

    # Check that both versions of rand produce samples. This is not a correctness test.
    @test rand_latent(rng, f(x_tr)) isa Vector{<:Real}
    @test rand(rng, f(x_tr)) isa Vector{<:Real}

    # Ensure that log p(Yte | Ytr) is self-consistent.
    f_posterior = posterior(f(x_tr, σ²), y_tr)
    x_all = vcat(x_tr, x_te)

    y_tr_cv = reshape(y_tr, :, x_tr.out_dim)
    y_te_cv = reshape(y_te, :, x_te.out_dim)
    y_all = vec(vcat(y_tr_cv, y_te_cv))
    @test logpdf(f_posterior(x_te, σ²), y_te) ≈
        logpdf(f(x_all, σ²), y_all) - logpdf(f(x_tr, σ²), y_tr)

    # Construct the posterior naively using Stheno.
    f_posterior = posterior(f(x_tr, σ²), y_tr)
    f_posterior_naive = posterior(f_naive(x_naive_tr, σ²), y_tr)

    # Compare posterior marginals.
    @test mean(f_posterior(x_tr, σ²)) ≈ mean(f_posterior_naive(x_naive_tr, σ²))
    @test var(f_posterior(x_tr)) ≈ var(f_posterior_naive(x_naive_tr))

    @test mean(f_posterior(x_te)) ≈ mean(f_posterior_naive(x_naive_te))
    @test var(f_posterior(x_te)) ≈ var(f_posterior_naive(x_naive_te))

    # Check that the gradient w.r.t. the logpdf and be computed w.r.t. the observations.
    lml_zygote, pb = Zygote.pullback((f, x, y) -> logpdf(f(x, σ²), y), f, x_tr, y_tr)
    @test lml_zygote ≈ logpdf(f(x_tr, σ²), y_tr)

    Δ = randn(rng)
    cotangents_fd = FiniteDifferences.j′vp(
        central_fdm(5, 1),
        y -> logpdf(f(x_tr, σ²), y),
        Δ,
        y_tr,
    )
    cotangents_ad = pb(Δ)

    @test cotangents_fd[1] ≈ cotangents_ad[3]
end

function consistency_tests(
    rng::AbstractRNG,
    f::OILMM,
    f_naive::AbstractGP,
    x::AbstractVector,
    Y::ColVecs{<:Real},
)
    consistency_tests(rng, f, [f_naive], x, Y)
end

function FiniteDifferences.to_vec(X::ColVecs)
    x_vec, back = to_vec(X.X)
    function from_vec_ColVecs(x)
        return ColVecs(back(x))
    end
    return x_vec, from_vec_ColVecs
end


text:
# array.jl array of images
z = [ones(6,4), rand(6,4), rand(6,4), rand(6,4), rand(6,4)]
@isplot jim(z, yflip=false, title=L"test3 x^2_i")
@isplot jim(z, ncol=2)
@isplot jim(z, nrow=3)
@isplot jim(-3:2, -2:1, z) # zyflip

z = [reshape(1:12,3,4), ones(4,5)]
@test MIRTjim._maxgood(z) == 12
@test MIRTjim._mingood(z) == 1


text:
@testset "Number fields" begin
  include("NumField/Elem.jl")
end


text:
using Test

using SBA

include("test_indices.jl")
include("test_core.jl")
include("test_update.jl")
include("test_param_check.jl")


text:
######################################################################
# costs.jl
#
######################################################################

abstract type CostModel end

function get_cost(cm::CostModel, m::SearchDomain, x::Vehicle, f::AbstractFilter, a)
    error("get_cost not implemented for this cost model type.")
end
function get_cost(cm::CostModel, m::SearchDomain, x::Vehicle, f::AbstractFilter)
    error("get_cost not implemented for this cost model type.")
end

struct ConstantCost <: CostModel
    value::Float64
end
function get_cost(cc::ConstantCost,
                  m::SearchDomain,
                  x::Vehicle,
                  f::AbstractFilter,
                  a
                 )
    return cc.value
end
get_cost(cc::ConstantCost, m::SearchDomain, x::Vehicle, f::AbstractFilter) = cc.value

struct MoveCost <: CostModel
    speed::Float64

    MoveCost(speed::Real) = new(float(speed))
    MoveCost() = new(1.0)
end
function get_cost(mc::MoveCost, x::Vehicle, f::AbstractFilter, a)
    dx = a[1]
    dy = a[2]
    dist = sqrt(dx*dx + dy*dy)
    return (dist / mc.speed)
end

struct MoveCost2 <: CostModel
    speed::Float64

    MoveCost2(speed::Real) = new(float(speed))
    MoveCost2() = new(1.0)
end

# TODO: should I incorporate speed somehow?
function get_cost(mc::MoveCost2, x::Vehicle, f::AbstractFilter, a)
    dx = a[1]
    dy = a[2]
    return dx*dx + dy*dy
end

# Move and rotate
struct MoveAndRotateCost <: CostModel
    speed::Float64
    time_per_rotation::Float64
end
function get_cost(marc::MoveAndRotateCost, x::Vehicle, f::AbstractFilter, a)
    dx = a[1]
    dy = a[2]
    dist = sqrt(dx*dx + dy*dy)
    return (dist / marc.speed) + marc.time_per_rotation
end
get_cost(marc::MoveAndRotateCost, f::AbstractFilter) = marc.time_per_rotation

export EntropyCost
struct EntropyCost <: CostModel
    L::Float64
    n_cells::Int
end

function get_cost(ec::EntropyCost, m::SearchDomain, x::Vehicle, f::PF, a)
    get_cost(ec, m, x, f)
end
function get_cost(ec::EntropyCost, m::SearchDomain, x::Vehicle, f::PF)
    cheap_entropy2(f, ec.L, ec.n_cells)
    #c = covariance(f)
    #earr = eigvals(c)
    #return maximum(eigvals(c))
end

export MaxEigCost
struct MaxEigCost <: CostModel
    take_sqrt::Bool
end
MaxEigCost() = MaxEigCost(true)

function get_cost(mec::MaxEigCost, m::SearchDomain, x::Vehicle, f::PF, a)
    get_cost(mec, m, x, f)
end
function get_cost(mec::MaxEigCost, m::SearchDomain, x::Vehicle, f::PF)
    me = maximum(eigvals(covariance(f)))
    if mec.take_sqrt
        me = sqrt(me)
    end
    return me
end


text:
println("Fitting linear, time-invariant model")
n,λ = 150,2.48
M2 = RegularizedLTIModel(n,N1,λ)

lti = estfun(M2,p1,Q1)


text:
include("discrete.jl")

MAP = [
    "+---------+",
    "|R: | : :G|",
    "| : : : : |",
    "| : : : : |",
    "| | : | : |",
    "|Y| : |B: |",
    "+---------+",
]

mutable struct TaxiEnv <: DiscreteEnv
    desc
    locs
    discenv_obj::DiscreteEnvObj
end

function TaxiEnv()
    desc = Char.(vcat([Int.(collect(row))' for row in MAP]...))

    locs = [[1, 1], [1, 5], [5, 1], [5, 4]]

    num_states = 500
    num_rows = 5
    num_columns = 5
    max_row = num_rows
    max_col = num_columns
    initial_state_distrib = zeros(Float32, num_states)
    num_actions = 6
    P = Dict(state => Dict(action => [] for action=1:num_actions) for state=1:num_states)
    for row=1:num_rows
        for col=1:num_columns
            for pass_idx=1:length(locs)+1
                for dest_idx=1:length(locs)
                    state = encode(row, col, pass_idx, dest_idx)
                    if pass_idx ≤ 4 && pass_idx != dest_idx
                        initial_state_distrib[state] += 1
                    end
                    for action=1:num_actions
                        # defaults
                        new_row, new_col, new_pass_idx = row, col, pass_idx
                        reward = -1
                        done = false
                        taxi_loc = [row, col]

                        if action == 1
                            new_row = min(row + 1, max_row)
                        elseif action == 2
                            new_row = max(row - 1, 1)
                        end

                        if action == 3 && desc[row+1, 2*(col-1) + 3] == ':'
                            new_col = min(col + 1, max_col)
                        elseif action == 4 && desc[row+1, 2 * (col-1) + 1] == ':'
                            new_col = max(col - 1, 1)
                        elseif action == 5 # pickup
                            if pass_idx ≤ 4 && taxi_loc == locs[pass_idx]
                                new_pass_idx = 5
                            else # passenger not at location
                                reward = -10
                            end
                        elseif action == 6 # dropoff
                            if (taxi_loc == locs[dest_idx]) && pass_idx == 5
                                new_pass_idx = dest_idx
                                done = true
                                reward = 20
                            elseif taxi_loc ∈ locs && pass_idx == 4
                                new_pass_idx = findall(x -> x == taxi_loc, locs)[1]
                            else
                                reward = -10
                            end
                        end
                        new_state = encode(new_row, new_col, new_pass_idx, dest_idx)
                        push!(P[state][action], (1.0, new_state, reward, done))
                    end
                end
            end
        end
    end
    initial_state_distrib ./= sum(initial_state_distrib)
    discenv_obj = DiscreteEnvObj(num_states, num_actions, P, initial_state_distrib)
    TaxiEnv(desc, locs, discenv_obj)
end

function encode(taxi_row, taxi_col, pass_loc, dest_idx)
    i = taxi_row - 1
    i *= 5
    i += taxi_col - 1
    i *= 5
    i += pass_loc - 1
    i *= 4
    i += dest_idx
    return i
end

function decode(i)
    i -= 1
    out = []
    push!(out, i%4+1)
    i = i ÷ 4
    push!(out, i%5+1)
    i = i ÷ 5
    push!(out, i%5+1)
    i = i ÷ 5
    @assert 1 ≤ i+1 ≤ 5
    push!(out, i+1)
    reverse!(out)
    return out
end


text:
export sympCayleyTransform, inverseSympCayleyTransform, sympRound

# Symplectic Cayley Transform: Output a symplectic matrix when the input is a symmetric matrix
function sympCayleyTransform(M::AbstractMatrix)::AbstractMatrix
    n = isSquare(M) ÷ 2
    Ω = cat(fill([0 1; -1 0], n)...; dims=(1,2))
    return I - (Ω * M  + I / 2  )^-1
end
# Invese Symplectic Cayley Transform: Output a symmetric matrix when the input is a symplectic matrix
function inverseSympCayleyTransform(S::AbstractMatrix)::AbstractMatrix
    n = isSquare(S) ÷ 2
    Ω = cat(fill([0 1; -1 0], n)...; dims=(1,2))
    return Symmetric(-Ω*((I - S)^-1 - I/2))
end
# Rounding a Matrix using Cayley Transform
function sympRound(S::AbstractMatrix)::AbstractMatrix
    n = isSquare(S) ÷ 2
    Ω = cat(fill([0 1; -1 0], n)...; dims=(1,2))
    if det(S - I) < tolerance
        return -Ω * ( (sympCayleyTransform ∘ inverseSympCayleyTransform)(Ω * S) )
    else
        return (sympCayleyTransform ∘ inverseSympCayleyTransform)( S )
    end
end

# Rounding a Symplectic Matrix using Cayley Transform
sympRound(S::Symp)::Symp = sympRound(S.S)


text:
# ### `09/clip-02.jl`

# ### Snippet 9.2

using StatisticalRethinking, LinearAlgebra
#gr(size=(600, 600))

ProjDir = @__DIR__

# Number of samples

T = 1000

# Compute radial distance

rad_dist(x) = sqrt(sum(x .^ 2))

# Plot densities

p = density(xlabel="Radial distance from mode", ylabel="Density")

for d in [1, 10, 100, 1000]
  m = MvNormal(zeros(d), Diagonal(ones(d)))
  y = rand(m, T)
  rd = [rad_dist( y[:, i] ) for i in 1:T]
  density!(p, rd, lab="d=$d")
end

plot(p)

# End of `09/clip-02.jl


text:
module BuildingAddIR

# Translated from: https://github.com/tachukao/ocaml-mlir/blob/master/tests/ir.ml

using MLIR
using MLIR.IR

function populate_loop_body(ctx, loop_body, loc, func_body)
    iv = IR.get_arg(loop_body, 0)
    func_arg0 = IR.get_arg(func_body, 0)
    func_arg1 = IR.get_arg(func_body, 1)
    f32 = IR.parse_type(ctx, "f32")

    # (lhs)
    load_lhs_state = IR.create_operation_state("std.load", loc)
    load_lhs_operands = [func_arg0, iv]
    IR.push!(load_lhs_state, load_lhs_operands)
    IR.push!(load_lhs_state, f32)
    load_lhs = IR.Operation(load_lhs_state)
    IR.push!(loop_body, load_lhs)

    # (rhs)
    load_rhs_state = IR.create_operation_state("std.load", loc)
    load_rhs_operands = [func_arg1, iv]
    IR.push!(load_rhs_state, load_rhs_operands)
    IR.push!(load_rhs_state, f32)
    load_rhs = IR.Operation(load_rhs_state)
    IR.push!(loop_body, load_rhs)

    # (add)
    add_state = IR.create_operation_state("std.addf", loc)
    add_operands = [IR.get_result(load_lhs, 0), IR.get_result(load_rhs, 0)]
    IR.push!(add_state, add_operands)
    IR.push!(add_state, f32)
    add = IR.Operation(add_state)
    IR.push!(loop_body, add)

    # (store state)
    store_state = IR.create_operation_state("std.store", loc)
    store_operands = [IR.get_result(add, 0), func_arg0, iv]
    IR.push!(store_state, store_operands)
    store = IR.Operation(store_state)
    IR.push!(loop_body, store)

    # (yield state)
    yield_state = IR.create_operation_state("scf.yield", loc)
    yield = IR.Operation(yield_state)
    IR.push!(loop_body, yield)
end

function make_and_dump_add(ctx, loc)
    module_op = IR.Module(loc)
    module_body = IR.get_body(module_op)
    memref_type = IR.parse_type(ctx, "memref<?xf32>")
    func_body_arg_types = [memref_type, memref_type]
    func_region = IR.create_region()
    func_body = IR.create_block(func_body_arg_types)
    IR.push!(func_region, func_body)
    func_type_attr = IR.parse_attribute(ctx, "(memref<?xf32>, memref<?xf32>) -> ()")
    func_name_attr = IR.parse_attribute(ctx, "\"add\"")
    func_attrs = [IR.NamedAttribute(ctx, "type", func_type_attr), IR.NamedAttribute(ctx, "sym_name", func_name_attr)]
    func_state = IR.create_operation_state("func", loc)
    IR.push!(func_state, func_attrs)
    IR.push!(func_state, func_region)
    func = IR.Operation(func_state)
    IR.insert!(module_body, 0, func)
    index_type = IR.parse_type(ctx, "index")
    index_zero_literal = IR.parse_attribute(ctx, "0: index")
    index_zero_value_attr = IR.NamedAttribute(ctx, "value", index_zero_literal)
    const_zero_state = IR.OperationState("std.constant", loc)
    IR.push!(const_zero_state, index_type)
    IR.push!(const_zero_state, index_zero_value_attr)
    const_zero = IR.Operation(const_zero_state)
    IR.push!(func_body, const_zero)
    func_arg0 = IR.get_arg(func_body, 0)
    const_zero_value = IR.get_result(const_zero, 0)
    dim_operands = [func_arg0, const_zero_value]
    dim_state = IR.OperationState("std.dim", loc)
    IR.push!(dim_state, dim_operands)
    IR.push!(dim_state, index_type)
    dim = IR.Operation(dim_state)
    IR.push!(func_body, dim)
    loop_body_region = IR.create_region()
    loop_body = IR.Block(index_type)
    IR.push!(loop_body_region, loop_body)
    index_one_literal = IR.parse_attribute(ctx, "1: index")
    index_one_value_attr = IR.NamedAttribute(ctx, "value", index_one_literal)
    const_one_state = IR.OperationState("std.constant", loc)
    IR.push!(const_one_state, index_type)
    IR.push!(const_one_state, index_one_value_attr)
    const_one = IR.Operation(const_one_state)
    IR.push!(func_body, const_one)
    dim_value = IR.get_result(dim, 0)
    const_one_value = IR.get_result(const_one, 0)
    loop_operands = [const_zero_value, dim_value, const_one_value]
    loop_state = IR.OperationState("scf.for", loc)
    IR.push!(loop_state, loop_operands)
    IR.push!(loop_state, loop_body_region)
    loop = IR.Operation(loop_state)
    IR.push!(func_body, loop)
    populate_loop_body(ctx, loop_body, loc, func_body)
    ret_state = IR.OperationState("std.return", loc)
    ret = IR.Operation(ret_state)
    IR.push!(func_body, ret)
    modu = IR.get_operation(module_op)
    modu
end

construct_and_traverse_ir = () -> begin
    ctx = IR.create_context()
    loc = IR.create_unknown_location(ctx)
    MLIR.IR.register_standard_dialect!(ctx)
    modu = make_and_dump_add(ctx, loc)
    IR.verify(modu)
    IR.dump(modu)
end

modu = construct_and_traverse_ir()

end # module


text:
struct InstructionList
    instructions::Vector{Tuple{Symbol,Tuple{Symbol,Any,Any}}}
    var::Symbol
    n::Base.RefValue{Int}
end

function InstructionList(; var::Symbol = :_ι, n::Base.RefValue{Int} = Ref(0))
    InstructionList(Tuple{Symbol,Tuple{Symbol,Any,Any}}[], var, n)
end

function instruction_list(
    exprs::Vector{Expression};
    subexpressions::Dict{Expression,Expression} = Dict{Expression,Expression}(),
    perform_cse::Bool = true,
)
    v = InstructionList()
    PSE = Set{Symbol}()
    if perform_cse
        exprs, CSE = cse(exprs)
        merge!(CSE, subexpressions)
    else
        CSE = subexpressions
    end
    v, map(ex -> flat_expr!(v, ex, CSE, PSE), exprs)
end

Base.length(v::InstructionList) = v.n[]
function Base.show(io::IO, v::InstructionList)
    for (id, arg) in v.instructions
        println(io, :($id = $(Expr(:call, arg...))))
    end
end

function Base.push!(v::InstructionList, x)
    id = Symbol(v.var, v.n[] += 1)
    push!(v.instructions, (id, x))
    id
end
function Base.push!(v::InstructionList, x::Pair{Symbol,<:Tuple{Symbol,Any,Any}})
    push!(v.instructions, (first(x), last(x)))
    first(x)
end

function add!(list::InstructionList, @nospecialize(a), @nospecialize(b))
    if a === nothing
        b === nothing && return nothing
        return b
    else
        b === nothing && return a
        return push!(list, (:+, a, b))
    end
end
function sub!(list::InstructionList, @nospecialize(a), @nospecialize(b))
    if a === nothing
        b === nothing && return nothing
        return push!(list, (:-, b, nothing))
    else
        b === nothing && return a
        return push!(list, (:-, a, b))
    end
end

function mul!(list::InstructionList, @nospecialize(a), @nospecialize(b))
    (a === nothing || b === nothing) && return nothing
    if a == 1
        b == 1 && return a
        return b
    else
        b == 1 && return a
        return push!(list, (:*, a, b))
    end
end

function muladd!(
    list::InstructionList,
    @nospecialize(a),
    @nospecialize(b),
    @nospecialize(c),
)
    add!(list, c, mul!(list, a, b))
end

function div!(list::InstructionList, @nospecialize(a), @nospecialize(b))
    a === nothing && return nothing
    b == 1 && return a
    push!(list, (:/, a, b))
end

function pow!(list, @nospecialize(a), b::Int)
    @assert b ≥ 0
    b == 1 && return a
    b == 0 && return 1
    a == 1 && return a
    a === nothing && return nothing
    push!(list, (:^, a, b))
end

function flat_expr!(
    v,
    ex::Basic,
    # common sub expression
    CSE::Dict{Expression,Expression},
    # processed sub expression
    PSE::Set{Symbol} = Set{Symbol}(),
)
    t = class(ex)
    if t == :Symbol
        s = Symbol(to_string(ex))
        if haskey(CSE, ex) && s ∉ PSE
            push!(PSE, s)
            val = flat_expr!(v, CSE[ex], CSE, PSE)
            if val isa Symbol
                v.instructions[end] = (s, last(v.instructions[end]))
            else
                push!(v, s => val)
            end
        end
        return s
    elseif t == :Pow
        x, k = ModelKit.args(ex)
        return push!(v, (:^, flat_expr!(v, x, CSE, PSE), convert(Int, k)))
    elseif t == :Add || t == :Mul
        op = t == :Add ? :+ : :*
        vec = ModelKit.args(ex)
        op_arg = flat_expr!(v, vec[1], CSE, PSE)
        for i = 2:length(vec)
            op_arg = push!(v, (op, op_arg, flat_expr!(v, vec[i], CSE, PSE)))
        end
        return op_arg
    elseif t == :Div
        x, y = ModelKit.args(ex)
        return push!(v, (:/, flat_expr!(v, x, CSE, PSE), flat_expr!(v, y, CSE, PSE)))
    else
        return to_number(ex)
    end
end

struct DiffMap
    D::Dict{Tuple{Symbol,Int},Any}
end
DiffMap() = DiffMap(Dict{Tuple{Symbol,Int},Any}())
function Base.getindex(D::DiffMap, s::Symbol, i::Int)
    i == 0 && return s
    get(D.D, (s, i), nothing)
end
function Base.getindex(D::DiffMap, s, i)
    i == 0 && return s
    nothing
end
Base.setindex!(D::DiffMap, v::Any, s::Symbol, i::Int) = D.D[(s, i)] = v
Base.setindex!(D::DiffMap, v::Nothing, s::Symbol, i::Int) = nothing

function Base.diff(list::InstructionList, vars::Vector{Symbol}, f::Vector{Symbol})
    diff_map = DiffMap()
    for (i, v) in enumerate(vars)
        diff_map[v, i] = 1
    end
    dlist = diff!(list, length(vars), diff_map)

    J = Matrix{Union{Nothing,Symbol,Number}}(undef, length(f), length(vars))
    for (j, v) in enumerate(vars), (i, fi) in enumerate(f)
        J[i, j] = diff_map[fi, j]
    end
    dlist, J
end

"""
    diff!(list::InstructionList, N::Int, diff_map)

Returns an `InstructionList` list computing the derivatives of `N` variables.
The derivatives of the instructions are stored in `diff_map`.
"""
function diff!(list::InstructionList, N::Int, D::DiffMap)
    n = length(list)
    v = InstructionList(n = Ref(n))
    for (id, el) in list.instructions
        (op, arg1, arg2) = el
        if op == :^
            p2 = :NONE
            instr_added = false
            for ∂i = 1:N
                exp::Int = arg2
                arg1_∂i = D[arg1, ∂i]
                if arg1_∂i !== nothing
                    if p2 == :NONE
                        p1 = pow!(v, arg1, exp - 1)
                        p2 = mul!(v, exp, p1)
                    end
                    if !instr_added
                        if exp == 2
                            push!(v, id => (:^, arg1, 2))
                        else
                            push!(v, id => (:*, p1, arg1))
                        end
                        instr_added = true
                    end
                    D[id, ∂i] = mul!(v, p2, arg1_∂i)
                elseif p2 != :NONE && !instr_added
                    push!(v, id => el)
                    instr_added = true
                end

            end
            if !instr_added
                push!(v, id => el)
            end
        elseif op == :*
            for ∂i = 1:N
                ∂i == 1 && push!(v, id => el)
                a = mul!(v, D[arg1, ∂i], arg2)
                b = mul!(v, arg1, D[arg2, ∂i])
                D[id, ∂i] = add!(v, a, b)
            end
        elseif op == :/
            for ∂i = 1:N
                ∂i == 1 && push!(v, id => el)

                D[id, ∂i] = div!(
                    v,
                    sub!(v, mul!(v, D[arg1, ∂i], arg2), mul!(v, arg1, D[arg2, ∂i])),
                    pow!(v, arg2, 2),
                )
            end
        elseif op == :+
            for ∂i = 1:N
                ∂i == 1 && push!(v, id => el)
                D[id, ∂i] = add!(v, D[arg1, ∂i], D[arg2, ∂i])
            end
        end
    end
    v
end

## Higher order AD. This follows Chapter 13 of Griewank and Walther
"""
    univariate_diff!(list::InstructionList, K::Int, diff_map)

Returns an `InstructionList` list computing the first `K` derivatives.
The derivatives of the instructions are stored in `diff_map`.
"""
function univariate_diff!(list::InstructionList, K::Int, D::DiffMap)
    n = length(list)
    v = InstructionList(n = Ref(n))
    for (id, el) in list.instructions
        (op, arg1, arg2) = el
        if op == :^
            r = arg2
            arg1_1 = D[arg1, 1]
            if arg1_1 !== nothing && r == 2
                push!(v, id => el)
                for k = 1:K
                    w_k = nothing
                    # Compute ∑_{j=0}^k u_j U_{k - j} =
                    #         2(∑_{j=0}^div(k - 1,2) u_j U_{k - j}) +
                    #         iseven(k) * u_{k-1}^2
                    for j = 0:div(k - 1, 2)
                        w_k = muladd!(v, D[arg1, j], D[arg1, k-j], w_k)
                    end
                    w_k = mul!(v, 2, w_k)
                    if iseven(k)
                        w_k = add!(v, w_k, pow!(v, D[arg1, div(k, 2)], 2))
                    end
                    D[id, k] = w_k
                end
            elseif arg1_1 !== nothing
                w_0 = push!(v, id => (:^, arg1, r))
                u_0_inv = push!(v, (:inv_not_zero, arg1, nothing))
                # w_k is v_k in Griewank
                w = Symbol[]
                for k = 1:K
                    s = nothing
                    for j = 1:k
                        ũ_j = mul!(v, j, D[arg1, j])
                        w_kj = k == j ? w_0 : w[k-j]
                        s = muladd!(v, w_kj, ũ_j, s)
                    end
                    s = mul!(v, r, s)

                    t = nothing
                    for j = 1:k-1
                        u_kj = D[arg1, k-j]
                        w̃_j = mul!(v, j, w[j])
                        t = muladd!(v, u_kj, w̃_j, t)
                    end
                    w̃_k = mul!(v, u_0_inv, sub!(v, s, t))
                    w_k = div!(v, w̃_k, k)
                    push!(w, w_k)
                    D[id, k] = w_k
                end
            else
                push!(v, id => el)
            end
        elseif op == :*
            push!(v, id => el)
            for k = 1:K
                c_k = nothing
                for j = 0:k
                    c_k = muladd!(v, D[arg1, j], D[arg2, k-j], c_k)
                end
                # c_k === nothing && break
                D[id, k] = c_k
            end
        elseif op == :+
            push!(v, id => el)
            for k = 1:K
                D[id, k] = add!(v, D[arg1, k], D[arg2, k])
            end
        elseif op == :/
            push!(v, id => el)
            for k = 1:K
                s = nothing
                for j = 0:(k-1)
                    s = muladd!(v, D[id, j], D[arg2, k-j], s)
                end
                D[id, k] = div!(v, sub!(v, D[arg1, k], s), arg2)
            end
        end
    end
    v
end

@inline sqr(x) = x^2
@inline function sqr(z::Complex)
    x, y = reim(z)
    Complex((x + y) * (x - y), (x + x) * y)
end

"""
    inv_not_zero(x)

Invert x unless it is 0, then return 0.
"""
@inline inv_not_zero(x) = ifelse(iszero(x), x, inv(x))

function unroll_pow(var, n)
    n == 0 && return :(one($var))
    n == 1 && return var
    n == 2 && return :(sqr($var))
    n == 3 && return :($var * sqr($var))
    n == 4 && return :(sqr(sqr($var)))
    n == 5 && return :($var * sqr(sqr($var)))

    # base to expansion shows up which power it is needed to compute
    d = digits(n, base = 2)
    x = :x
    exprs = map(2:length(d)) do i
        :(local $(Symbol(x, 1 << (i - 1))) = sqr($(Symbol(x, 1 << (i - 2)))))
    end
    prods = Symbol[]
    for (i, di) in enumerate(d)
        if !iszero(di)
            push!(prods, Symbol(x, 1 << (i - 1)))
        end
    end
    if length(prods) > 1
        push!(exprs, Expr(:call, :*, prods...))
    end
    Expr(:let, :($(Symbol(x, 1)) = $var), Expr(:block, exprs...))
end

quick_div(a, b::Complex) = Base.FastMath.div_fast(a, b)
quick_div(a, b) = a / b

function to_expr(
    list::InstructionList,
    var_map = Dict{Symbol,Union{Expr,Symbol}}(),
    assignements = Dict{Symbol,Vector{Expr}}(),
)
    block = Expr(:block)
    exprs = block.args
    for (id, (op, arg1, arg2)) in list.instructions
        if op == :^
            x::Union{Symbol,Expr} = arg1
            if arg2 isa Int
                k::Int = arg2
                if k < 0
                    push!(exprs, :($id = inv($(unroll_pow(get(var_map, x, x), -k)))))
                else
                    push!(exprs, :($id = $(unroll_pow(get(var_map, x, x), k))))
                end
            else
                push!(exprs, :($id = $(get(var_map, x, x))^$arg2))
            end
        elseif op == :/
            a = get(var_map, arg1, arg1)
            b = get(var_map, arg2, arg2)
            push!(exprs, :($id = quick_div($a, $b)))
        else
            a = get(var_map, arg1, arg1)
            if arg2 !== nothing
                b = get(var_map, arg2, arg2)
                push!(exprs, :($id = $(Expr(:call, op, a, b))))
            else
                push!(exprs, :($id = $(Expr(:call, op, a))))
            end
        end

        if haskey(assignements, id)
            append!(exprs, assignements[id])
        end
    end
    block
end

function untuple(varname, tuplename, M)
    :($(Expr(:tuple, varname, (Symbol(varname, k) for k = 1:M-1)...)) = $tuplename)
end

function _impl_taylor_bivariate(K, M, N, op; inline = true)
    list = ModelKit.InstructionList()
    id = push!(list, (op, :x, :y))
    diff_map = ModelKit.DiffMap()
    for i = 1:M-1
        diff_map[:x, i] = Symbol(:x, i)
    end
    for j = 1:N-1
        diff_map[:y, j] = Symbol(:y, j)
    end
    dlist = ModelKit.univariate_diff!(list, K, diff_map)
    quote
        $(inline ? :(Base.@_inline_meta) : :())
        $(untuple(:x, :tx, M))
        $(untuple(:y, :ty, N))
        $(ModelKit.to_expr(dlist))
        $(Expr(:tuple, id, ((d = diff_map[id, k];
        d === nothing ? :(zero(x)) : d) for k = 1:K)...))
    end
end

@generated function taylor(
    ::Type{Val{op}},
    ::Type{Val{K}},
    tx::NTuple{M},
    ty::NTuple{N},
) where {op,K,M,N}
    _impl_taylor_bivariate(K, M, N, op)
end

make_ntuple(t) = t
make_ntuple(t::Number) = (t,)
make_ntuple(t::Tuple{A,B}) where {A,B} = convert(NTuple{2,promote_type(A, B)}, t)
make_ntuple(t::Tuple{A,B,C}) where {A,B,C} = convert(NTuple{3,promote_type(A, B, C)}, t)
make_ntuple(t::Tuple{A,B,C,D}) where {A,B,C,D} =
    convert(NTuple{4,promote_type(A, B, C, D)}, t)
make_ntuple(t::Tuple{A,B,C,D,E}) where {A,B,C,D,E} =
    convert(NTuple{5,promote_type(A, B, C, D, E)}, t)

taylor(op::Type{T1}, v::Type{T2}, tx, ty) where {T1,T2} =
    taylor(op::Type{T1}, v::Type{T2}, make_ntuple(tx), make_ntuple(ty))
function taylor(op::Type{T1}, v::Type{T2}, tx) where {T1,T2}
    taylor(op::Type{T1}, v::Type{T2}, make_ntuple(tx))
end

function _impl_taylor_pow(K, M)
    list = ModelKit.InstructionList()
    id = push!(list, (:^, :x, :r))
    diff_map = ModelKit.DiffMap()
    for i = 1:M-1
        diff_map[:x, i] = Symbol(:x, i)
    end
    dlist = ModelKit.univariate_diff!(list, K, diff_map)
    quote
        $(untuple(:x, :tx, M))
        $(ModelKit.to_expr(dlist))
        $(Expr(:tuple, id, ((d = diff_map[id, k];
        d === nothing ? :(zero(x)) : d) for k = 1:K)...))
    end
end

@generated function taylor(
    ::Type{Val{:^}},
    ::Type{Val{K}},
    tx::NTuple{M},
    r::Integer,
) where {K,M}
    _impl_taylor_pow(K, M)
end

function _impl_taylor_sqr(K, M)
    list = ModelKit.InstructionList()
    id = push!(list, (:^, :x, 2))
    diff_map = ModelKit.DiffMap()
    for i = 1:M-1
        diff_map[:x, i] = Symbol(:x, i)
    end
    dlist = ModelKit.univariate_diff!(list, K, diff_map)
    quote
        $(untuple(:x, :tx, M))
        $(ModelKit.to_expr(dlist))
        $(Expr(:tuple, id, ((d = diff_map[id, k];
        d === nothing ? :(zero(x)) : d) for k = 1:K)...))
    end
end

@generated function taylor(::Type{Val{:sqr}}, ::Type{Val{K}}, tx::NTuple{M}) where {K,M}
    _impl_taylor_sqr(K, M)
end


text:
import RayTracing:
    Track, Segment,
    DirectionType, Forward, Backward,
    universal_id, bc_fwd, bc_bwd, dir_next_track_fwd, dir_next_track_bwd

# TODO: Gridap solution object might be useful for interpolation at any point in space...
struct MoCSolution{T<:Real,P<:MoCProblem} <: TransportSolution
    prob::P

    keff::T
    residual::T
    iterations::Int

    # scalar flux
    φ::Vector{T}

    # reduced source
    q::Vector{T}

    # total source (integrated in energy given per region) used for convergence purposes
    Q::Vector{T}

    # angular flux at boundary
    boundary_ψ::Vector{T}
    start_boundary_ψ::Vector{T}
end

function MoCSolution{T}(prob::MoCProblem) where {T}
    @unpack nφ, nψ = prob
    NRegions = nregions(prob)
    φ = Vector{T}(undef, nφ)
    q = Vector{T}(undef, nφ)
    Q = Vector{T}(undef, NRegions)
    boundary_ψ = Vector{T}(undef, nψ)
    start_boundary_ψ = Vector{T}(undef, nψ)
    return MoCSolution(prob, one(T), zero(T), 0, φ, q, Q, boundary_ψ, start_boundary_ψ)
end

function show(io::IO, sol::MoCSolution)
    @unpack keff, residual, iterations = sol
    println(io, "  keff: ", keff)
    println(io, "  Residual: ", residual)
    print(io,   "  Iterations: ", iterations)
end

function (sol::MoCSolution)(i::Int, g::Int)
    NGroups = ngroups(sol.prob)
    g in 1:NGroups || throw(DomainError(g, "`g` is outside of domain."))
    return sol.φ[@region_index(i, g)]
end

function (sol::MoCSolution)(g::Int)
    NGroups = ngroups(sol.prob)
    g in 1:NGroups || throw(DomainError(g, "`g` is outside of domain."))
    return view(sol.φ, g:NGroups:lastindex(sol.φ))
end

function solve(
    prob::MoCProblem;
    max_iterations::Int=1000, max_residual::Real=1e-7, debug::Bool=false
)
    return _solve_eigenvalue_problem(prob, max_iterations, max_residual, debug)
end

function _solve_eigenvalue_problem(prob::MoCProblem, max_iter::Int, max_ϵ::Real, debug::Bool)

    T =  eltype(prob)
    sol = MoCSolution{T}(prob)

    optical_length!(prob)

    @set! sol.keff = one(T)
    set_uniform_φ!(sol, one(T))
    set_uniform_Q!(sol, zero(T))
    set_uniform_start_boundary_ψ!(sol, zero(T))
    update_boundary_ψ!(sol)

    debug && @info "MoC iterations start..."
    ϵ = Inf
    iter = 0
    while iter < max_iter

        normalize_fluxes!(sol, prob)
        compute_q!(sol, prob)
        compute_φ!(sol, prob)
        @set! sol.keff *= total_fission_source(sol, prob)
        ϵ = residual(sol, prob)

        debug && iszero(iter % 10) && @info "iteration $(iter)" sol.keff ϵ

        # we need at least three iterations:
        #  0. boundary conditions do not exists (unless everything is vaccum)
        #  1. we can compute a residual but it is not correct (previous iteration is fruit)
        #  2. now we can compute a correct residual

        if iter > 1 && isless(ϵ, max_ϵ)
            break
        end

        iter += 1
    end

    @set! sol.residual = ϵ
    @set! sol.iterations = iter

    return sol
end

function optical_length!(prob::MoCProblem)
    @unpack trackgenerator = prob
    @unpack tracks_by_uid = trackgenerator

    for track in tracks_by_uid
        _optical_length!(prob, track)
    end

    return nothing
end

function _optical_length!(prob::MoCProblem, track::Track)
    NGroups = ngroups(prob)

    for segment in track.segments
        @unpack ℓ, τ = segment

        resize!(τ, NGroups)

        xs = getxs(prob, segment.element)
        @unpack Σt = xs
        @inbounds for g in 1:NGroups
            τ[g] = Σt[g] * ℓ
        end
    end
end

@inline set_uniform_φ!(sol::MoCSolution, φ::Real) = fill!(sol.φ, φ)
@inline set_uniform_Q!(sol::MoCSolution, Q::Real) = fill!(sol.Q, Q)
@inline set_uniform_start_boundary_ψ!(sol::MoCSolution, ψ::Real) = fill!(sol.start_boundary_ψ, ψ)
@inline update_boundary_ψ!(sol::MoCSolution) = copy!(sol.boundary_ψ, sol.start_boundary_ψ)

function normalize_fluxes!(sol::MoCSolution, prob::MoCProblem)
    @unpack φ, boundary_ψ, start_boundary_ψ = sol

    # total fission source
    qft = total_fission_source(sol, prob)

    # λ is the normalization factor
    λ = 1 / qft
    φ .*= λ
    boundary_ψ .*= λ
    start_boundary_ψ .*= λ

    return nothing
end

function total_fission_source(sol::MoCSolution{T}, prob::MoCProblem) where {T}
    NGroups = ngroups(prob)
    NRegions = nregions(prob)
    @unpack φ = sol
    @unpack trackgenerator = prob
    @unpack volumes = trackgenerator

    qft = zero(T)
    @inbounds for i in 1:NRegions
        xs = getxs(prob, i)
        @unpack νΣf = xs
        fissionable = isfissionable(xs)

        if fissionable
            for g′ in 1:NGroups
                ig′ = @region_index(i, g′)
                qft += νΣf[g′] * φ[ig′] * volumes[i]
            end
            # qft += sum(
            #     νΣf[g′] * φ[@region_index(i, g′)] * volumes[i] for g′ in 1:NGroups
            # )
        end
    end

    return qft
end

function compute_q!(sol::MoCSolution{T}, prob::MoCProblem) where {T}
    NGroups = ngroups(prob)
    NRegions = nregions(prob)
    @unpack keff, φ, q = sol

    @inbounds for i in 1:NRegions
        xs = getxs(prob, i)
        @unpack χ, Σt, νΣf, Σs0 = xs
        fissionable = isfissionable(xs)

        for g in 1:NGroups
            ig = @region_index(i, g)
            qig = zero(T)
            for g′ in 1:NGroups
                ig′ = @region_index(i, g′)
                qig += Σs0[g′, g] * φ[ig′]
                if fissionable
                    qig += 1 / keff * χ[g] * νΣf[g′] * φ[ig′]
                end
            end
            # qig = sum(
            #     Σs0[g′, g] * φ[@region_index(i, g′)] +
            #     1 / keff * χ[g] * νΣf[g′] * φ[@region_index(i, g′)]
            #     for g′ in 1:NGroups
            # )

            qig /= (4π * Σt[g])
            q[ig] = qig
        end
    end

    return nothing
end

function compute_φ!(sol::MoCSolution{T}, prob::MoCProblem) where {T}
    @unpack trackgenerator = prob
    @unpack tracks_by_uid = trackgenerator

    set_uniform_φ!(sol, zero(T))
    update_boundary_ψ!(sol) # update entry ψ for all rays (including d, p and g dependence)

    for track in tracks_by_uid
        tally!(sol, prob, track, Forward)
        tally!(sol, prob, track, Backward)
    end

    add_q_to_φ!(sol, prob)

    return nothing
end

function tally!(sol::MoCSolution, prob::MoCProblem, track::Track, dir::DirectionType)
    NGroups = ngroups(prob)
    @unpack quadrature = prob

    n_polar_2 = npolar2(quadrature.polar)

    t = universal_id(track)
    d = Int32(dir)

    i = @angular_index(t, d, 1, 1)
    j = i + NGroups * n_polar_2 - 1
    boundary_ψ = @view sol.boundary_ψ[i:j] # boundary_ψ in for a given track in a given direction as function of (p, g)

    segments = dir == Forward ? track.segments : reverse!(track.segments)

    for segment in segments
        tally_φ!(sol, prob, track, segment, boundary_ψ)
    end

    set_start_boundary_ψ!(sol, prob, track, boundary_ψ, dir)

    if dir == Backward
        reverse!(track.segments)
    end

    return nothing
end

function tally_φ!(
    sol::MoCSolution,
    prob::MoCProblem,
    track::Track,
    segment::Segment,
    boundary_ψ::AbstractVector
)
    NGroups = ngroups(prob)
    @unpack φ, q = sol
    @unpack quadrature = prob
    @unpack polar, ω = quadrature
    @unpack sinθs = polar
    @unpack τ = segment

    i = segment.element
    a = track.azim_idx
    n_polar_2 = npolar2(polar)

    # TODO: is this the best possible loop order?
    @inbounds for g in 1:NGroups, p in 1:n_polar_2
        pg = @reduced_angular_index(p, g)
        ig = @region_index(i, g)
        Δψ = (boundary_ψ[pg] - q[ig]) * (1 - exp(-τ[g] / sinθs[p]))
        φ[ig] += 2 * ω[a, p] * Δψ # multiplied by 2 because we only loop in n_polar_2
        boundary_ψ[pg] -= Δψ
    end

    return nothing
end

function set_start_boundary_ψ!(
    sol::MoCSolution{T},
    prob::MoCProblem,
    current_track::Track,
    boundary_ψ::AbstractVector,
    dir::DirectionType
) where {T}

    NGroups = ngroups(prob)
    @unpack start_boundary_ψ = sol
    @unpack quadrature = prob

    n_polar_2 = npolar2(quadrature.polar)

    if dir == Forward
        next_track = current_track.next_track_fwd
        next_track_dir = dir_next_track_fwd(current_track)
        flag = !isequal(bc_fwd(current_track), Vaccum)
    elseif dir == Backward
        next_track = current_track.next_track_bwd
        next_track_dir = dir_next_track_bwd(current_track)
        flag = !isequal(bc_bwd(current_track), Vaccum)
    end

    t = universal_id(next_track)
    d = Int32(next_track_dir)

    # TODO: is this the best possible loop order?
    @inbounds for g in 1:NGroups, p in 1:n_polar_2
        tdpg = @angular_index(t, d, p, g)
        pg = @reduced_angular_index(p, g)
        start_boundary_ψ[tdpg] = flag ? boundary_ψ[pg] : zero(T)
    end

    return nothing
end

function add_q_to_φ!(sol::MoCSolution, prob::MoCProblem)
    NGroups = ngroups(prob)
    NRegions = nregions(prob)
    @unpack φ, q = sol
    @unpack trackgenerator = prob
    @unpack volumes = trackgenerator

    @inbounds for i in 1:NRegions
        xs = getxs(prob, i)
        @unpack Σt = xs

        for g in 1:NGroups
            ig = @region_index(i, g)
            φ[ig] /= (Σt[g] * volumes[i])
            φ[ig] += (4π * q[ig])
        end
    end

    return nothing
end

function residual(sol::MoCSolution{T}, prob::MoCProblem) where {T}
    NGroups = ngroups(prob)
    NRegions = nregions(prob)
    @unpack keff, φ, Q = sol

    ϵ = zero(T)
    @inbounds for i in 1:NRegions

        xs = getxs(prob, i)
        @unpack νΣf, Σs0 = xs
        fissionable = isfissionable(xs)

        old_qi = Q[i]
        new_qi = zero(T)

        # total fission source in each region (χ sum 1 when ∫ in g)
        if fissionable
            for g′ in 1:NGroups
                ig′ = @region_index(i, g′)
                νΣfg′ = νΣf[g′]
                new_qi += νΣfg′ * φ[ig′]
            end
            # new_qi = sum(νΣf[g′] * φ[ig′] for g′ in 1:NGroups)
        end

        new_qi /= keff

        # total scattering source
        for g in 1:NGroups, g′ in 1:NGroups
            ig′ = @region_index(i, g′)
            Σsgg′ = Σs0[g′, g]
            new_qi += Σsgg′ * φ[ig′]
        end

        if old_qi > 0
            ϵ += ((new_qi - old_qi) / old_qi)^2
        end

        Q[i] = new_qi
    end

    ϵ = sqrt(ϵ / NRegions)

    return ϵ
end


text:
module Functions

using ..TopOpt: dim, whichdevice, CPU, GPU, TopOpt, PENALTY_BEFORE_INTERPOLATION
using ..TopOptProblems, ..TrussTopOptProblems
using ..TopOptProblems: initialize_K, getdh
using ..FEA, ..CheqFilters
using ..Utilities, ForwardDiff, LinearAlgebra, Requires
using ..TrussTopOptProblems: getA, compute_local_axes

using Parameters: @unpack
using TimerOutputs, Ferrite, StaticArrays, StatsFuns
using SparseArrays, Statistics, ChainRulesCore, Zygote
using Nonconvex: Nonconvex
using Flux

export Volume,
    Compliance,
    Displacement,
    MeanCompliance,
    BlockCompliance,
    AbstractFunction,
    getfevals,
    getmaxfevals,
    maxedfevals,
    MicroVonMisesStress,
    MacroVonMisesStress,
    project,
    generate_scenarios,
    hutch_rand!,
    hadamard!,
    TrussStress,
    AssembleK,
    TrussElementKσ,
    ElementK,
    apply_boundary_with_zerodiag!,
    apply_boundary_with_meandiag!,
    NeuralNetwork,
    PredictFunction,
    TrainFunction

const to = TimerOutput()

abstract type AbstractFunction{T} <: Nonconvex.NonconvexCore.AbstractFunction end

include("function_utils.jl")
include("compliance.jl")
include("displacement.jl")
include("volume.jl")
include("stress.jl")
include("trace.jl")
include("mean_compliance.jl")
include("block_compliance.jl")

# buckling
include("apply_boundary.jl")
include("assemble_K.jl")
include("element_ksigma.jl")
include("element_k.jl")

# TODO no rrules yet
include("truss_stress.jl")

include("neural.jl")

end


text:
module PhysicalConstants

export c_l, q_e, m_e, h_P, k_B, hbar, epsilon_0, a_0
export 𝒄, 𝒒ₑ, 𝒎ₑ, 𝒉, 𝒌, ħ, ħ², ɛ₀, a₀ , Ha, π², ⅈ

"Pi squared. Type as \\pi<tab>\\^2<tab>"
const π² = π * π

"Imaginary unit, √(-1)"
const ⅈ = im;

"Speed of light in vacuum, 𝒄 · s / m. Type as c_l or \\mbfitc<tab>"
const 𝒄 = c_l = 299792458

"Electron charge, 𝒒ₑ / C. Type as q_e or \\mbfitq<tab>\\_e<tab>"
const 𝒒ₑ = q_e = 1.6021766208e-19;

"Electron mass, 𝒎ₑ / kg. Type as m_e or \\mbfitm<tab>\\_e<tab>"
const 𝒎ₑ = m_e = 9.10938356e-31;

"Plank constant, 𝒉 / (J · s). Type as h_P or \\mbfith<tab>"
const 𝒉 = h_P = 6.626070040e-34;

"Boltzmann constant, 𝒌 · K / eV. Type as k_B or \\mbfitk<tab>"
const 𝒌 = k_B = 8.6173303e-5;

"Plank constant over 2π, ħ / (J · s). Type as hbar or \\hbar<tab>"
const ħ = hbar = 1.0545718e-34;
const ħ² = ħ * ħ;

"Electric constant, ɛ₀ · m / F. Type as epsilon_0 or \\varepsilon<tab>\\_0<tab>"
const ɛ₀  = epsilon_0 = 1 / (4π * 1e-7 * 𝒄^2);

"Bohr radius in m"
const a_0 = a₀  = 0.52917721067e-10;

"Hartree energy in eV"
const Ha = 27.21138602;

end # End of module


text:
module KinematicRepresentations

using LinearAlgebra, StaticArrays, Rotations, CoordinateTransformations

include("motion.jl")
include("angularvelocity.jl")
include("twist.jl")

export RigidMotion, AngularVelocity, Twist
export groupexp, grouplog

end


text: @inplacerule (y = conv2d(x, w; _opts)) (conv2d!(y, x, w; _opts))


text:
# Autogenerated wrapper script for XNNPACK_jll for x86_64-apple-darwin
export libxnnpack

using CPUInfo_jll
using PThreadPool_jll
JLLWrappers.@generate_wrapper_header("XNNPACK")
JLLWrappers.@declare_library_product(libxnnpack, "@rpath/libXNNPACK.0.dylib")
function __init__()
    JLLWrappers.@generate_init_header(CPUInfo_jll, PThreadPool_jll)
    JLLWrappers.@init_library_product(
        libxnnpack,
        "lib/libXNNPACK.0.0.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()