text:
using DelimitedFiles

# ---------------------------------------------------------------------------------------- #

"""
$(TYPEDSIGNATURES)

Return the symmetry eigenvalues and little groups as **k**-label indexed `Dict`s for
an MPB symmetry calculation with ID `calcname`.

## Keyword arguments
- `sgnum`, `dir`: If unspecified (or `nothing`), the space group number and dimension must
  be inferrable from `calcname` (`calcname` must then follow the conventions of
  [`mpb_calcname`](@ref)).
- `dir` (default, `"."`): directory location where `calcname`-*.out files can be found.
- `αβγ` (default, `$TEST_αβγ`): free-parameters used for **k**-vectors in setting up the
  MPB calculation in [`prepare_mpbcalc`](@ref).
- `isprimitive` (default, `true`): whether the calculation is in a primitive setting.
- `flip_ksign` (default, `false`): flip the sign of the **k**-vector used in the MPB
  calculation; can be necessary to match phase conventions in Crystalline.jl.
"""
function read_symdata(calcname::AbstractString;
            sgnum::Union{Int, Nothing}=nothing,
            D::Union{Int, Nothing}=nothing,
            dir::AbstractString=".",
            αβγ::AbstractVector{<:Real}=TEST_αβγ,
            isprimitive::Bool=true,
            flip_ksign::Bool=false)

    sgnum === nothing && (sgnum = parse_sgnum(calcname))
    D === nothing     && (D = parse_dim(calcname))
    D < length(αβγ)   && (αβγ = αβγ[1:D])

    # prepare default little groups of associated space group
    lgs⁰  = littlegroups(sgnum, Val(D))
    isprimitive && map!(g -> primitivize(g, #=modw=# false), values(lgs⁰)) # primitivize

    # read mpb dispersion data; use to guarantee frequency sorting at each k-point
    dispersion_data = readdlm(joinpath(dir, calcname*"-dispersion.out"), ',')
    kvs = collect(eachrow(@view dispersion_data[:,2:2+(D-1)]))
    Nk     = length(kvs)
    freqs = dispersion_data[:,6:end]
    Nbands = size(freqs, 2)
    sortidxs = [sortperm(freqs_at_fixed_k) for freqs_at_fixed_k in eachrow(freqs)]
    # read mpb symmetry data, mostly as Strings (first column is Int)
    untyped_data = readdlm(joinpath(dir, calcname*"-symeigs.out"), ',', quotes=true)
    Nrows = size(untyped_data, 1)
    # process raw symmetry & operator data
    lgopsd   = Dict{String, Vector{SymOperation{D}}}()    # indexing: [klabel][op]
    symeigsd = Dict{String, Vector{Vector{ComplexF64}}}() # indexing: [klabel][band][op]
    rowidx = 1
    for kidx in 1:Nk
        kv = flip_ksign ? -kvs[kidx] : kvs[kidx] # `flip_ksign` is a hack to work around
                                                 # phase convention issues; there be dragons
        klab = findfirst(lg->isapprox(position(lg)(αβγ), kv, atol=1e-6), lgs⁰)
        klab === nothing && error("could not find matching KVec for loaded kv = $kv")
        lgopsd[klab]   = Vector{SymOperation{D}}()
        symeigsd[klab] = [Vector{ComplexF64}() for _ in 1:Nbands]
        while rowidx ≤ Nrows && untyped_data[rowidx, 1] == kidx
            op = SymOperation{D}(strip(untyped_data[rowidx, 2], '"'))
            push!(lgopsd[klab], op)
            # symmetry eigenvalues at `op` (and `klab`) across all bands (frequency sorted)
            symeigs = parse.(ComplexF64, @view untyped_data[rowidx, 3:end][sortidxs[kidx]])
            push!.(symeigsd[klab], symeigs) # broadcast over band indices
            rowidx += 1
        end
    end
    # build little groups
    lgd = Dict(klab => LittleGroup{D}(sgnum, position(lgs⁰[klab]), klab, ops) for (klab, ops) in lgopsd)

    return symeigsd, lgd
end

# ---------------------------------------------------------------------------------------- #

function symeigs2irreps(symeigs::AbstractVector, lgirs::Vector{LGIrrep{D}}, bands;
            atol::Real=DEFAULT_ATOL, αβγ::AbstractVector{<:Real}=TEST_αβγ) where D
                                                                        # ... root accessor
    D < length(αβγ) && (αβγ = αβγ[1:D])

    # return multiplicities over provided irreps (and across `bands`)
    symeigs_bands = sum(@view symeigs[bands])
    return find_representation(symeigs_bands, lgirs, αβγ, Float64, atol=atol)
end

"""
$(TYPEDSIGNATURES) --> Dict{String, Union{Nothing, Vector{Float64}}}

Return the irrep multiplicities for provided symmetry eigenvalues `symeigsd` with associated
little group irreps `lgirsd`, with symmetry eigenvalues aggregated (i.e. summed) over the
band-indices in `bands`.

`symeigsd` and `lgirsd` must be `Dict`s referencing the same **k**-points and assuming the
same operator sorting (and setting).

Result is returned as a `Dict` over **k**-labels. If the eigenvalues at this **k**-point are
expandable in the irreps in `lgirsd`, the `Dict` value is a vector of multiplicities.
If not (e.g., due to `bands` "splitting" an irrep or overly tight `atol` tolerances),
the value is `nothing`.

## Keyword arguments
- `kwargs...`: optional keyword arguments (including `atol`) forwarded to
  [`symeigs2irreps`](@ref) (e.g., `atol` and `αβγ`).

If `bands` is unspecified, it assumes the maximal possible bandrange, as inferrable from the
size of `symeigsd`'s elements.

## Workflow
Symmetry eigenvalues `symeigsd` and little groups `lgd` can be loaded via
[`read_symdata`](@ref); `lgirsd` can subsequently be generated via [`pick__lgirreps`](@ref).
To extract the multiplicities of the first 5 bands, we then call:
```jl
julia> bands = 1:5
julia> extract_multiplicities(symeigsd, lgirsd, bands)
```
"""
function extract_multiplicities(symeigsd::Dict{String,<:Any},
            lgirsd::Dict{String,Vector{LGIrrep{D}}},
            bands=eachindex(first(values(symeigsd)));
            kwargs...) where D                                  # ... main accessor

    nd = Dict{String, Union{Nothing, Vector{Float64}}}()
    for klab in keys(symeigsd)
        nd[klab] = symeigs2irreps(symeigsd[klab], lgirsd[klab], bands; kwargs...)
    end
    return nd
end

"""
$(TYPEDSIGNATURES) --> Dict{String, Union{Nothing, Vector{Float64}}}

Return the irrep multiplicities for provided symmetry eigenvalues `symeigsd` with associated
little groups `lgd`, with symmetry eigenvalues aggregated (i.e. summed) over the
band-indices in `bands`.

`symeigsd` and `lgd` must be `Dict`s referencing the same **k**-points and assuming the
same operator sorting (and setting). `lgd` is used as an intermediary to determine the
relevant little group irreps via Crystalline.

Result is returned as a `Dict` over **k**-labels. If the eigenvalues at this **k**-point are
expandable in the associated little group irreps, the `Dict` value is a vector of
multiplicities. If not (e.g., due to `bands` "splitting" an irrep or overly tight `atol`
tolerances), the value is `nothing`.

## Keyword arguments
- `timereversal` (default, `true`): whether the underlying MPB calculation has time-reversal
  symmetry; if `true`, physically real irreps (coreps) will be used. See Crystalline's
  `realify`.
- `isprimitive` (default, `true`): whether the underlying MPB calculation was performed in a
  primitive unit cell; if `true`, `lgd` is assumed to be provided in a primitive setting as
  well. Used as a consistency check in constructing associated little group irreps.
- `kwargs...`: optional keyword arguments (including `atol`) forwarded to
  [`symeigs2irreps`](@ref) (e.g., `atol` and `αβγ`).

If `bands` is unspecified, it assumes the maximal possible bandrange, as inferrable from the
size of `symeigsd`'s elements.

## Workflow
Given a the "base" name of an MPB calculation, `calcname`, we can extract associated
multiplicities by:
```jl
julia> symeigsd, lgd = read_symdata(calcname, dir = "location/of/calcname/out/files")
julia> extract_multiplicities(symeigsd, lgd, 1:5)
```
where we consider the joint multiplicities of the first 5 bands. See also
[`extract_individual_multiplicities`](@ref).
"""
function extract_multiplicities(symeigsd::Dict{String, <:Any},
            lgd::Dict{String,LittleGroup{D}},
            bands=eachindex(first(values(symeigsd)));
            timereversal::Bool=true, isprimitive::Bool=true,
            kwargs...) where D                                  # ... main accessor

    if length(symeigsd) ≠ length(lgd) || !(keys(symeigsd) ⊆ keys(lgd))
        error("`symeigsd` and `lgd` must have identical keys")
    end

    lgirsd = pick_lgirreps(lgd; timereversal, isprimitive) # irreps for k-points in `lgd`

    return extract_multiplicities(symeigsd, lgirsd, bands; kwargs...)
end

function extract_multiplicities(calcname::String, bands;
            timereversal::Bool=true, isprimitive::Bool=true,
            atol::Real=DEFAULT_ATOL, αβγ::AbstractVector{<:Real}=TEST_αβγ,
            read_kwargs...)                                     # ... convenience accessor

    symeigsd, lgd = read_symdata(calcname; αβγ, isprimitive, read_kwargs...)
    return extract_multiplicities(symeigsd, lgd, bands; timereversal, isprimitive, atol, αβγ)
end

# ---------------------------------------------------------------------------------------- #

function pick_lgirreps(lgd::Dict{String, LittleGroup{D}};
            timereversal::Bool=true, isprimitive::Bool=true) where D
    sgnum = num(first(values(lgd)))
    lgirsd = lgirreps(sgnum, Val(D))
    # filter out k-points not in `lgd`
    filter!(((klab,_),) -> klab ∈ keys(lgd), lgirsd)

    # incorporate timereveral & primitivize if relevant
    if timereversal
        map!(realify, values(lgirsd))
    end
    if isprimitive
        cntr = centering(sgnum, D)
        if cntr ≠ 'P' || cntr ≠ 'p'
            for (klab, lgirs) in lgirsd
                # all elements in `lgirs` point to the same group, so we should only change
                # _one_ such element - pick the first one (bit icky, yeah...)
                lg = group(first(lgirs))
                lg.operations .= primitivize.(lg.operations, cntr, #=modw=# false)
            end
        end
    end

    # ensure that `lgirsd` and `lgd` have identical operator sorting
    for (klab, lg) in lgd
        align_operators!(lgirsd[klab], lg) # mutates `lgirsd` if operator sorting differs
    end

    return lgirsd
end

function align_operators!(lgirs::Vector{LGIrrep{D}}, lg::LittleGroup{D}) where D
    lg′ = group(first(lgirs))
    lg == lg′ && return lgirs
    length(lg) == length(lg′) || error("mismatched little groups")
    # find permutation to "realign" operator sorting between `lgd[klab]` and `lgirsd[klab]`
    perm = Vector{Int}(undef, length(lg))
    for op in lg
        idx = findfirst(==(op), lg′)
        idx === nothing && error("could not align operators")
        perm[i] = idx
    end
    # apply permutation to `lgirs`
    for lgir in lgirs
        permute!.(lgir.matrices,     Ref(perm))
        permute!.(lgir.translations, Ref(perm))
    end
    permute!(first(lgirs).g.operations, perm) # NB: the `group` of each `lgirs` element is
                        # actually the same struct, so we only permute it _once_
                        # (admittedly, a bit iffy to rely on...)
    @info "Permuted little group irreps to ensure aligned sorting" klab
    return lgirs
end

# ---------------------------------------------------------------------------------------- #

"""
$(TYPEDSIGNATURES) --> bandirsd, lgirsd

Return band-groupings and **k**-projected symmetry vectors at individual **k**-points
(`bandirsd`) and associated little group irreps (`lgirsd`).

To extract the associated **k**-point projected symmetry vectors of potentially separable
bands, see [`collect_separable`](@ref).

## Keyword arguments

- `latestarts` (default, `Dict("Γ" => D)`): allow user to specify that certain early bands
  be avoided for specific **k**-points. This is mainly useful to avoid considering the two
  singular Γ-point bands in 3D photonic crystals (which is the default behavior). To skip
  nothing, set to `Dict{String,Int}()`. Defaults to the spatial dimension `D` at Γ
  (meaning, start at band `D` and skip the first `D-1` bands at Γ).
- `timereversal` & `isprimitive`: forwarded to [`read_symdata`](@ref) and
  [`extract_multiplicities`](@ref).
- `atol` & `αβγ`: forwarded to [`read_symdata`](@ref) and [`find_representation`](@ref).
- `read_kwargs`: additional keyword arguments forwarded to [`read_symdata`](@ref).

## Example

```jl
julia> calcname = "dim3-sg147-symeigs_6936-res32"
julia> bandirsd, lgirsd = extract_individual_multiplicities(calcname,
                        timereversal=true, dir = "../../mpb-ctl/output/", atol=1e-3)
```
The result can be pretty-printed by e.g.:
```
julia> using Crystalline: label, formatirreplabel, symvec2string
julia> irlabs = Dict(klab => formatirreplabel.(label.(lgirs)) for (klab, lgirs) in lgirsd)
julia> Dict(klab => [bands => symvec2string(n, irlabs[klab]; braces=false)
                     for (bands, n) in bandirs]         for (klab, bandirs) in bandirsd)
```
"""
function extract_individual_multiplicities(calcname::String;
            timereversal::Bool=true, isprimitive::Bool=true,
            atol::Real=DEFAULT_ATOL, αβγ::AbstractVector{<:Real}=TEST_αβγ,
            latestarts::Dict{String,Int}=Dict("Γ" => parse_dim(calcname)),
            kwargs...)

    symeigsd, lgd = read_symdata(calcname; αβγ, isprimitive, kwargs...)
    lgirsd = pick_lgirreps(lgd; timereversal, isprimitive)

    # determine how we can consistently group up irreps and bands across distinct k-points
    bandirsd = find_individual_multiplicities(symeigsd, lgirsd; atol, αβγ, latestarts)

    return bandirsd, lgirsd
end

function find_individual_multiplicities(symeigsd::Dict{String,<:AbstractVector},
            lgirsd::Dict{String,Vector{LGIrrep{D}}};
            atol::Real=DEFAULT_ATOL,
            αβγ::AbstractVector{<:Real}=TEST_αβγ,
            latestarts::Dict{String,Int}=Dict("Γ" => D)) where D

    Nbands = length(first(values(symeigsd)))

    bandirsd = Dict(klab => Vector{Pair{UnitRange{Int}, Vector{Int}}}() for klab in keys(lgirsd))
    for (klab, lgirs) in lgirsd
        symeigs = symeigsd[klab]
        start = stop = get(latestarts, klab, 1)
        while stop ≤ Nbands
            bands = start:stop
            n = symeigs2irreps(symeigs, lgirs, bands; atol, αβγ)
            if n !== nothing
                idxs = findall(nᵢ -> nᵢ > atol && abs(round(nᵢ) - nᵢ) < atol, n)
                if length(idxs) ≥ 1
                    # `bands` makes up a valid whole-multiple irrep combination at `klab`
                    # we emphasize that this can be a _combination_ i.e. multiple whole
                    # irreps; this can arise in cases of near-degeneracies, where numerical
                    # "mode-mixing" effectively spreads two irreps over two bands bands
                    # (even if they _should_ be each in one band); in principle, this can
                    # always be circumvented by increasing the resolution, but this is not
                    # a practical solution in general
                    push!(bandirsd[klab], bands => round.(Int, n))
                    start = stop + 1 # prepare for next band grouping
                end
            end
            stop += 1
        end
    end

    return bandirsd
end

"""
$(TYPEDSIGNATURES)

Return the "projected" symmetry-vectors for potentially separable bands using `bandirsd`
(see [`extract_individual_multiplicities`](@ref)) and `lgirsd` (see[`read_symdata`](@ref)
and [`pick_lgirreps`).
"""
function collect_separable(bandirsd::Dict{String, Vector{Pair{UnitRange{Int}, Vector{Int}}}},
                lgirsd::Dict{String, Vector{LGIrrep{D}}};
                latestarts::Dict{String, Int}=Dict("Γ" => D)) where D

    Nbands = mapreduce(last∘first∘last, min, values(bandirsd)) # smallest "last" band-index
    include = Dict(klab => Int[] for klab in keys(bandirsd))
    collectibles_bands   = Vector{UnitRange{Int}}()
    collectibles_include = Vector{typeof(include)}()
    start = stop = 1
    while stop ≤ Nbands
        stable = true
        for (klab, bandirs) in bandirsd
            idxs = include[klab]
            latestart = get(latestarts, klab, nothing)
            for (i, (bands, _)) in enumerate(bandirs)
                minband, maxband = extrema(bands)
                minband < start && continue

                if latestart !== nothing && stop < latestart && minband ≤ latestart
                    break # allow `bands` to "not count" if `latestarts` indicates skips
                end
                i ∉ idxs && push!(idxs, i)

                if maxband > stop
                    stop = maxband
                    stable = false
                    break
                elseif maxband == stop
                    break
                end
            end
        end

        if stable
            push!(collectibles_bands,   start:stop)
            push!(collectibles_include, include)
            start = (stop += 1)
            include = Dict(klab => Int[] for klab in keys(bandirsd))
        end
    end

    # create projected symmetry vectors for each, using knowledge of number of irreps in
    # each little group
    Nirs = Dict(klab => length(lgirsd[klab]) for klab in keys(bandirsd))
    collectibles_symvecs = Vector{Dict{String, Vector{Int}}}()
    for (bands, include) in zip(collectibles_bands, collectibles_include)
        symvecs = Dict(klab => zeros(Int, Nir) for (klab, Nir) in Nirs)
        for (klab, idxs) in include
            for i in idxs
                symvecs[klab] += last(bandirsd[klab][i])
            end
        end
        push!(collectibles_symvecs, symvecs)
    end

    return collectibles_bands, collectibles_symvecs
end

# ---------------------------------------------------------------------------------------- #

"""
$(TYPEDSIGNATURES)

Return the space group number of an MPB calculation `calcname` (in the conventions of
[`mpb_calcname`](@ref)).
"""
function parse_sgnum(calcname::AbstractString)
    idxs = findfirst("sg", calcname)
    idxs === nothing && error("could not infer space group number from calcname")
    # finds `sgnum` start & stop position in `calcname`, then pulls out `sgnum` as a string
    sgstart = idxs[end] + 1
    sgstop  = findnext(!isdigit, calcname, sgstart) - 1
    sgnum_str = calcname[sgstart:sgstop]

    return parse(Int, sgnum_str)
end

"""
$(TYPEDSIGNATURES)

Return the dimension of an MPB calculation `calcname` (in the conventions of
[`mpb_calcname`](@ref)).
"""
function parse_dim(calcname::AbstractString)
    idxs = findfirst("dim", calcname)
    idxs === nothing && error("could not infer dimension from calcname")

    Dstart = idxs[end] + 1
    D = parse(Int, calcname[Dstart]) # return the dimension as an integer
end


text:
#
# Formatted operations preserve the most stringent array structure
#

import ApproxFun: ⊕
export ⊖,⊛,⊘

for (op,opformatted) in ((:+,:⊕),(:-,:⊖),(:*,:⊛),(:/,:⊘))
    @eval begin
        $opformatted(A) = $op(A)
        $opformatted(A,B) = $op(A,B)
    end
end

for (op,opformatted) in ((:+,:⊕),(:-,:⊖))
    @eval begin

        function $opformatted(H::HierarchicalVector,J::Vector)
            @assert length(H) == length(J)
            H1,H2 = data(H)
            n1,n2 = length(H1),length(H2)
            HierarchicalVector(($opformatted(H1,J[1:n1]),$opformatted(H2,J[1+n1:n1+n2])))
        end
        function $opformatted(J::Vector,H::HierarchicalVector)
            @assert length(H) == length(J)
            H1,H2 = data(H)
            n1,n2 = length(H1),length(H2)
            HierarchicalVector(($opformatted(J[1:n1],H1),$opformatted(J[1+n1:n1+n2],H2)))
        end

        $opformatted(L::LowRankMatrix,a::Number) = $opformatted(L,LowRankMatrix(a,size(L)...))
        $opformatted(a::Number,L::LowRankMatrix) = $opformatted(LowRankMatrix(a,size(L)...),L)

        function $opformatted(L::LowRankMatrix,M::LowRankMatrix)
            N = $op(L,M)
            T = eltype(N)
            QU,RU = qr(N.U)
            QV,RV = qr(N.V)
            U,Σ,V = svd(RU*RV.')
            r = refactorsvd!(U,Σ,V)
            LowRankMatrix(QU[:,1:r]*U[1:r,1:r],QV[:,1:r]*V[1:r,1:r])
        end
        $opformatted(L::LowRankMatrix,A::AbstractMatrix) = $opformatted(L,LowRankMatrix(A))
        $opformatted(A::AbstractMatrix,L::LowRankMatrix) = $opformatted(LowRankMatrix(A),L)
    end
end

ext: jl


text:
# Submodule to compute the coefficients of shearlet recovery and decoding in 2D

##############################################################################
## Function that compute the coefficient matrix of the Shearlet Transform of
## some array
"""
...
sheardec2D(X,shearletSystem) compute the coefficient matrix of the Shearlet
transform of the array X
...
"""
function sheardec2D(X,shearletSystem)
    #Read the GPU info of the system
    gpu = shearletSystem.gpu;
    if gpu == 1
        coeffs = AFArray(zeros(Complex{Float32},size(shearletSystem.shearlets)));
    else
        coeffs = zeros(Complex{Float64},size(shearletSystem.shearlets));
    end
    # The fourier transform of X
    Xfreq = fftshift(fft(ifftshift(X)));
    #compute shearlet coefficients at each scale
    #not that pointwise multiplication in the fourier domain equals convolution
    #in the time-domain
    for j = 1:shearletSystem.nShearlets
        coeffs[:,:,j] = fftshift(ifft(ifftshift(Xfreq.*conj(shearletSystem.shearlets[:,:,j]))));
    end
    return coeffs
end # sheardec2D

##############################################################################
## Function that recovers the image with the Shearlet transform with some shearletSystem
"""
...
shearrec2D(coeffs,shearletSystem) function that recovers the image with the Shearlet transform with some shearletSystem
...
"""
function shearrec2D(coeffs,shearletSystem)
    # Initialize reconstructed data
    gpu = shearletSystem.gpu;
    if gpu == 1
        X = AFArray(zeros(Complex{Float32},size(coeffs,1),size(coeffs,2)));
    else
        X = zeros(Complex{Float64},size(coeffs,1),size(coeffs,2));
    end
    for j = 1:shearletSystem.nShearlets
        X = X+fftshift(fft(ifftshift(coeffs[:,:,j]))).*shearletSystem.shearlets[:,:,j];
    end
    return real(fftshift(ifft(ifftshift((1 ./ shearletSystem.dualFrameWeights).*X))))
end# shearrec2D

##############################################################################
## Function that compute the array of the adjoint Shearlet Transform of
## some coefficients matrix
"""
...
sheardecadjoint2D(coeffs,shearletSystem) compute the adjoint of the decomposition operator
...
"""
function sheardecadjoint2D(coeffs,shearletSystem)
		# Initialize reconstructed data
    gpu = shearletSystem.gpu;
    if gpu == 1
        X = AFArray(zeros(Complex{Float32},size(coeffs,1),size(coeffs,2)));
    else
        X = zeros(Complex{Float64},size(coeffs,1),size(coeffs,2));
    end
    for j = 1:shearletSystem.nShearlets
        X = X+fftshift(fft(ifftshift(coeffs[:,:,j]))).*conj(shearletSystem.shearlets[:,:,j]);
    end
    return real(fftshift(ifft(ifftshift((1 ./shearletSystem.dualFrameWeights).*X))))

end # sheardecadjoint2D

##########################################################################################
## Function that compute the coefficient matrix of the adjoint inverse Shearlet Transform of
## some array
"""
...
shearrecadjoint2D(X,shearletSystem) compute the coefficient matrix of the adjoint inverse Shearlet
transform of the array X
...
"""
function shearrecadjoint2D(X,shearletSystem)
    #Read the GPU info of the system
    gpu = shearletSystem.gpu;
    if gpu == 1
        coeffs = AFArray(zeros(Complex{Float32},size(shearletSystem.shearlets)));
    else
        coeffs = zeros(Complex{Float64},size(shearletSystem.shearlets));
    end
    # The fourier transform of X
    Xfreq = fftshift(fft(ifftshift(X)));
    #compute shearlet coefficients at each scale
    #not that pointwise multiplication in the fourier domain equals convolution
    #in the time-domain
    for j = 1:shearletSystem.nShearlets
        coeffs[:,:,j] = fftshift(ifft(ifftshift(Xfreq.*shearletSystem.shearlets[:,:,j])));
    end
    return coeffs
end # shearrecadjoint2D


text:
using ContinuousPOMDPTreeSearchExperiments
using ParticleFilters
using ARDESPOT
using BasicPOMCP
using POMCPOW
using POMDPs
using QMDP
using MCTS
using DiscreteValueIteration
using POMDPToolbox
using ProgressMeter
using PmapProgressMeter

@everywhere using ContinuousPOMDPTreeSearchExperiments
@everywhere using POMDPs
@everywhere using ParticleFilters

file_contents = readstring(@__FILE__())

@show max_time = 1.0
@show max_depth = 20
pomdp = SimpleLightDark()

solvers = Dict{String, Union{Solver,Policy}}(

    "pomcpow" => begin
        rng = MersenneTwister(13)
        ro = ValueIterationSolver()
        solver = POMCPOWSolver(tree_queries=10_000_000,
                               criterion=MaxUCB(90.0),
                               final_criterion=MaxTries(),
                               max_depth=max_depth,
                               max_time=max_time,
                               enable_action_pw=false,
                               k_observation=5.0,
                               alpha_observation=1/15.0,
                               estimate_value=FOValue(ro),
                               check_repeat_obs=false,
                               tree_in_info=false,
                               # default_action=ReportWhenUsed(-1),
                               rng=rng
                              )
    end,

    "d_pomcp" => begin
        rng = MersenneTwister(13)
        ro = ValueIterationSolver()
        sol = POMCPSolver(max_depth=max_depth,
                    max_time=max_time,
                    c=100.0,
                    tree_queries=typemax(Int),
                    # default_action=ro,
                    estimate_value=FORollout(ro),
                    tree_in_info=false,
                    rng=rng
                   )
    end,

    "d_despot" => begin
        rng = MersenneTwister(13)
        ro = QMDPSolver()
        b = IndependentBounds(-100.0, 100.0, check_terminal=true)
        sol = DESPOTSolver(lambda=0.01,
                     epsilon_0=0.0,
                     K=5000,
                     D=max_depth,
                     max_trials=1_000_000,
                     T_max=max_time,
                     bounds=b,
                     default_action=ReportWhenUsed(solve(ro, pomdp)),
                     rng=rng)
    end,

    "qmdp" => QMDPSolver(),
    "side" => LDSide()
)

@show N=1000

alldata = DataFrame()

# for (k, solver) in solvers
test = keys(solvers)
for d in logspace(-2, 1, 7)
    dpomdp = DSimpleLightDark(pomdp, d)
    for (k, solver) in [(k, solvers[k]) for k in test]
        @show k, d
        if isa(solver, Solver)
            planner = solve(solver, dpomdp)
        else
            planner = solver
        end
        # prog = Progress(N, desc="Creating Simulations...")
        # sims = pmap(prog, 1:N) do i
        sims = []
        for i in 1:N
            srand(planner, i+50_000)
            filter = ObsAdaptiveParticleFilter(deepcopy(pomdp),
                                               LowVarianceResampler(10_000),
                                               0.05, MersenneTwister(i+90_000))

            sim = Sim(deepcopy(pomdp),
                planner,
                filter,
                rng=MersenneTwister(i+70_000),
                max_steps=100,
                metadata=Dict(:d=>d, :solver=>k, :i=>i)
               )

            push!(sims, sim)
        end

        data = run_parallel(sims)
        # data = run(sims)

        rs = data[:reward]
        println(@sprintf("reward: %6.3f ± %6.3f", mean(rs), std(rs)/sqrt(length(rs))))
        if isempty(alldata)
            alldata = data
        else
            alldata = vcat(alldata, data)
        end
    end
end

datestring = Dates.format(now(), "E_d_u_HH_MM")
copyname = Pkg.dir("ContinuousPOMDPTreeSearchExperiments", "icaps_2018", "data", "ld_discretization_$(datestring).jl")
write(copyname, file_contents)
filename = Pkg.dir("ContinuousPOMDPTreeSearchExperiments", "icaps_2018", "data", "ld_discretization_$(datestring).csv")
println("saving to $filename...")
CSV.write(filename, alldata)
println("done.")


text:
using Rhapsodie

if prod(readdir() .!= "test_results")
    mkdir("test_results")
end

DSIZE=128;
NTOT=64;
Nframe=2;
Nrot=1
Nangle=Int64.(NTOT/(Nrot*Nframe*4))
Center=[DSIZE/2;  DSIZE/2];
DerotAng=Vector{Float64}();
Epsilon=Vector{Tuple{Array{Float64,1},Array{Float64,1}}}();

for iter=1:NTOT
    ind=div(iter-1, NTOT/4)+1
    push!(Epsilon,([0. ,0. ],[10.7365 , -1.39344]));
    push!(DerotAng, deg2rad(-300.0));
end

psf_center=readdlm("../data/PSF_centers_Airy.txt");
load_parameters((DSIZE, 2*DSIZE, NTOT), Nframe, Nrot, Nangle, Center, (psf_center[1:2], psf_center[3:4]), Epsilon, DerotAng)

writedlm("test_results/Parameters.txt", [DSIZE; NTOT; Nframe; Nrot; Center; 300; [10.7365 , -1.39344]]);

psf=read(FitsArray,  "../data/PSF_parametered_Airy.fits");
const A=set_fft_op((psf[1:end÷2,:]'), get_par().psf_center[1]);

BadPixMap=rand(0.0:1e-16:1.0,(DSIZE, 2*DSIZE)).< 0.9;

for tau in [0.03]#, 0.07, 0.1, 0.15, 0.25, 0.5]

    data, weight, S, S_convolved=data_simulator(BadPixMap, tau, A);

    write(FitsFile, "test_results/DATA_$tau-$DSIZE.fits",
          mapslices(transpose,data,dims=[1,2]), overwrite=true)
    write(FitsFile, "test_results/WEIGHT_$tau-$DSIZE.fits",
          mapslices(transpose,weight,dims=[1,2]), overwrite=true)

    write(S, "test_results/TRUE_$tau-$DSIZE.fits")
    write(S_convolved, "test_results/TRUE_convolved_$tau-$DSIZE.fits")

end

max_stars_repo_name: LaurenceDenneulin/Rhapsodie

max_issues_repo_name: LaurenceDenneulin/Rhapsodie

max_forks_repo_name: LaurenceDenneulin/Rhapsodie


text:
module JSON3

using Parsers, Mmap, UUIDs, Dates

struct Object{S <: AbstractVector{UInt8}, TT <: AbstractVector{UInt64}} <: AbstractDict{Symbol, Any}
    buf::S
    tape::TT
end

Object() = Object(codeunits(""), UInt64[object(Int64(2)), 0])

struct Array{T, S <: AbstractVector{UInt8}, TT <: AbstractVector{UInt64}} <: AbstractVector{T}
    buf::S
    tape::TT
end

Array{T}(buf::S, tape::TT) where {T, S, TT} = Array{T, S, TT}(buf, tape)

getbuf(j::Union{Object, Array}) = getfield(j, :buf)
gettape(j::Union{Object, Array}) = getfield(j, :tape)

include("utils.jl")

@noinline invalid(error, buf, pos, T) = throw(ArgumentError("""
invalid JSON at byte position $pos while parsing type $T: $error
$(String(buf[max(1, pos-25):min(end, pos+25)]))
"""))

@enum Error UnexpectedEOF ExpectedOpeningObjectChar ExpectedOpeningQuoteChar ExpectedOpeningArrayChar ExpectedClosingArrayChar ExpectedComma ExpectedSemiColon InvalidChar

# AbstractDict interface
function Base.length(obj::Object)
    @inbounds len = getnontypemask(gettape(obj)[2])
    return len
end

@inline function Base.iterate(obj::Object, (i, tapeidx)=(1, 3))
    i > length(obj) && return nothing
    tape = gettape(obj)
    @inbounds t = tape[tapeidx]
    key = getvalue(Symbol, getbuf(obj), tape, tapeidx, t)
    tapeidx += 2
    @inbounds t = tape[tapeidx]
    x = Pair{Symbol, Any}(key, getvalue(Any, getbuf(obj), tape, tapeidx, t))
    tapeidx += gettapelen(Any, t)
    return x, (i + 1, tapeidx)
end

function Base.get(obj::Object, key::Symbol)
    for (k, v) in obj
        k == key && return v
    end
    throw(KeyError(key))
end

function Base.get(obj::Object, key)
    k2 = Base.string(key)
    for (k, v) in obj
        String(k) == k2 && return v
    end
    throw(KeyError(key))
end

function Base.get(obj::Object, ::Type{T}, key::Symbol)::T where {T}
    for (k, v) in obj
        k == key && return v
    end
    throw(KeyError(key))
end

function Base.get(obj::Object, ::Type{T}, key)::T where {T}
    k2 = Base.string(key)
    for (k, v) in obj
        String(k) == k2 && return v
    end
    throw(KeyError(key))
end

function Base.get(obj::Object, key::Symbol, default)
    for (k, v) in obj
        k == key && return v
    end
    return default
end

function Base.get(obj::Object, key, default)
    k2 = Base.string(key)
    for (k, v) in obj
        String(k) == k2 && return v
    end
    return default
end

function Base.get(obj::Object, ::Type{T}, key::Symbol, default::T)::T where {T}
    for (k, v) in obj
        k == key && return v
    end
    return default
end

function Base.get(obj::Object, ::Type{T}, key, default::T)::T where {T}
    k2 = Base.string(key)
    for (k, v) in obj
        String(k) == k2 && return v
    end
    return default
end

function Base.get(default::Base.Callable, obj::Object, key::Symbol)
    for (k, v) in obj
        k == key && return v
    end
    return default()
end

function Base.get(default::Base.Callable, obj::Object, key)
    k2 = Base.string(key)
    for (k, v) in obj
        String(k) == k2 && return v
    end
    return default()
end

Base.propertynames(obj::Object) = collect(keys(obj))

Base.getproperty(obj::Object, prop::Symbol) = get(obj, prop)
Base.getindex(obj::Object, str::String) = get(obj, Symbol(str))

function Base.copy(obj::Object)
    dict = Dict{Symbol, Any}()
    for (k, v) in obj
        dict[k] = v isa Object || v isa Array ? copy(v) : v
    end
    return dict
end

# AbstractArray interface
Base.IndexStyle(::Type{<:Array}) = Base.IndexLinear()

function Base.size(arr::Array)
    @inbounds len = getnontypemask(gettape(arr)[2])
    return (len,)
end

function Base.iterate(arr::Array{T}, (i, tapeidx)=(1, 3)) where {T}
    i > length(arr) && return nothing
    tape = gettape(arr)
    @inbounds t = tape[tapeidx]
    val = getvalue(T, getbuf(arr), tape, tapeidx, t)
    tapeidx += gettapelen(T, t)
    return val, (i + 1, tapeidx)
end

@inline Base.@propagate_inbounds function Base.getindex(arr::Array{T}, i::Int) where {T}
    @boundscheck checkbounds(arr, i)
    tape = gettape(arr)
    buf = getbuf(arr)
    if regularstride(T)
        tapeidx = 1 + 2 * i
        @inbounds t = tape[tapeidx]
        return getvalue(T, buf, tape, tapeidx, t)
    else
        tapeidx = 3
        idx = 1
        while true
            @inbounds t = tape[tapeidx]
            if i == idx
                return getvalue(T, buf, tape, tapeidx, t)
            else
                tapeidx += gettapelen(T, t)
                idx += 1
            end
        end
    end
end

Base.copy(arr::Array) = map(x->x isa Object || x isa Array ? copy(x) : x, arr)

include("read.jl")
include("strings.jl")
include("show.jl")
include("structs.jl")
include("write.jl")

end # module


text:
module Images

import Base.Order: Ordering, ForwardOrdering, ReverseOrdering
import Base.Graphics: width, height
importall Base

using Color
using Cartesian

# We'll use Grid's Counter, but we don't need the rest of it
include(joinpath(Pkg.dir(),"Grid","src","counter.jl"))

include("core.jl")
include("config.jl")
include("iterator.jl")
include("io.jl")
include("scaling.jl")
include("labeledarrays.jl")
include("algorithms.jl")

export # types
    AbstractImage,
    AbstractImageDirect,
    AbstractImageIndexed,
    Image,
    ImageCmap,
    BitShift,
    ClipMin,
    ClipMax,
    ClipMinMax,
    LabeledArray,
    Overlay,
    OverlayImage,
    ScaleAutoMinMax,
    ScaleInfo,
    ScaleMinMax,
    ScaleNone,
    ScaleSigned,
    SliceData,

    # constants
    palette_fire,
    palette_gray32,
    palette_gray64,
    palette_rainbow,

    # core functions
    assert2d,
    assert_scalar_color,
    assert_timedim_last,
    assert_xfirst,
    assert_yfirst,
    colordim,
    colorspace,
    coords,
    coords_spatial,
    data,
    dimindex,
    dimindexes,
    getindexim,
    grayim,
    height,
    isdirect,
    isxfirst,
    isyfirst,
    limits,
    maxabsfinite,
    maxfinite,
    minfinite,
    nimages,
    pixelspacing,
    properties,
    refim,
    rerange!,
    reslice!,
    sdims,
    size_spatial,
    share,
    sliceim,
    spatialorder,
    spatialpermutation,
    spatialproperties,
    storageorder,
    subim,
    timedim,
    width,
    widthheight,

    # io functions
    add_image_file_format,
    imread,
    imwrite,
    loadformat,

    # iterator functions
    first_index,
    iterate_spatial,
    parent,

    # color-related functions
    alphaval,
    blueval,
    greenval,
    redval,
    red,
    green,
    blue,
    hsi2rgb,
    imadjustintensity,
    indexedcolor,
    lut,
    ntsc2rgb,
    rgb2gray,
    rgb2hsi,
    rgb2ntsc,
    rgb2ycbcr,
    uint32color,
    uint32color!,
    ycbcr2rgb,

    # Scaling of intensity
    climdefault,
    float32sc,
    float64sc,
    sc,
    scale,
    scaleinfo,
    scaleminmax,
    scalesigned,
    uint8sc,
    uint16sc,
    uint32sc,

    # algorithms
    backdiffx,
    backdiffy,
    dilate,
    erode,
    opening,
    closing,
    forwarddiffx,
    forwarddiffy,
    gaussian2d,
    imaverage,
    imcomplement,
    imdog,
    imedge,
    imfilter,
    imfilter_gaussian,
    imfilter_gaussian!,
    imgaussiannoise,
    imlaplacian,
    imlineardiffusion,
    imlog,
    imROF,
    imstretch,
#     imthresh,
    label_components,
    label_components!,
    ncc,
    padarray,
    prewitt,
    sad,
    sadn,
    sobel,
    ssd,
    ssdn,

    # phantoms
    shepp_logan

@deprecate cairoRGB     uint32color!
@deprecate refim        getindexim
@deprecate scaledefault climdefault

end


text:
import sbp
using MAT
using LinearAlgebra
include("../common/wave_equation.jl")
include("../common/bound.jl")
include("../common/discretization.jl")
disc = Discretization
pde=WaveEquation

datafile=ARGS[1]
dR = parse(Float64, ARGS[2]) * ones(1,1)
nsamples = 4
nrefinements = 6
nx = 20
T = 1.2
dt = 0.25 / (nx - 1)
nt = T / dt
order = 4

mms = pde.manufactured_solution

function data()
        return Dict("p" => [], "v" => [] ,
                    "Fp" => [], "Fv" => [])
end

function discretize(nx, dRi)
    ops = disc.OperatorsType(nx, order)
    M = disc.discretization(pde.A(), ops)
    Rl = -ones(1,1)
    Rr = -ones(1,1)
    Pl = disc.projection_matrix(pde.A(pde.LEFT), Rl, dRi)
    Pr = disc.projection_matrix(pde.A(pde.RIGHT), Rr, dRi)
    Sl, Slg = pde.sat(pde.A(pde.LEFT), pde.LEFT, ops, Pl)
    Sr, Srg = pde.sat(pde.A(pde.RIGHT), pde.RIGHT, ops, Pr)
    gl = t -> pack(mms(ops.x[1], t))
    gr = t -> pack(mms(ops.x[end], t))
    f = (u,t) -> (M + Sl + Sr) * u - Slg * gl(t) - Srg * gr(t)
    return f, ops
end

function pack(mms_)
   p = mms_["p"]
   v = mms_["v"]
   return [p; v]
end

function initialize(ops, nx)
   return pack(mms(ops.x, 0))
end

function error(d, u, ops, t)
    mms = pde.manufactured_solution(ops.x,  t)
    F_mms = pde.functionals(t)

    soln = Dict(sbp.MMS.extract_vars(u, pde.sizes(ops.nx)))
    norms = Dict("p" => ops.Hx, "v" => ops.Hx)
    err = sbp.MMS.l2_errors(soln, mms, norms)
    F_err = sbp.MMS.functional_errors(soln, F_mms, norms)
    append!(d["p"], err["p"])
    append!(d["v"], err["v"])
    append!(d["Fp"], F_err["p"])
    append!(d["Fv"], F_err["v"])
end

function info(i, nx, nt, dt, ops, err)
        println("Grid = $i, nx = $nx, p = ", err["p"][end],
                " v = ", err["v"][end],
                ", rate p = ", sbp.MMS.log2_convergence_rates(err["p"])[end],
                ", rate v = ", sbp.MMS.log2_convergence_rates(err["v"])[end],
                ", rate Fp = ", sbp.MMS.log2_convergence_rates(err["Fp"])[end],
                ", rate Fv = ", sbp.MMS.log2_convergence_rates(err["Fv"])[end]
                )
end

errors = data()
rates = data()

x = disc.convergence_study(nx, nt, dt, nrefinements, errors, rates,
                           initialize, x -> discretize(x, dR), error,
                           info)
println("")
matwrite(datafile,
         Dict("dR" => dR,
              "x" => x,
              "p" => errors["p"],
              "v" => errors["v"],
              "Fp" => errors["Fp"],
              "Fv" => errors["Fv"],
              "q_p" => rates["p"],
              "q_v" => rates["v"],
              "q_Fp" => rates["Fp"],
              "q_Fv" => rates["Fv"]
               ))

max_stars_repo_name: ooreilly/sbp

max_issues_repo_name: ooreilly/sbp

max_forks_repo_name: ooreilly/sbp


text:
# Compressor Codec
# ================

struct Bzip2Compressor <: TranscodingStreams.Codec
    stream::BZStream
    blocksize100k::Int
    workfactor::Int
    verbosity::Int
end

function Base.show(io::IO, codec::Bzip2Compressor)
    print(io, summary(codec), "(blocksize100k=$(codec.blocksize100k), workfactor=$(codec.workfactor), verbosity=$(codec.verbosity))")
end

const DEFAULT_BLOCKSIZE100K = 9
const DEFAULT_WORKFACTOR = 30
const DEFAULT_VERBOSITY = 0

"""
    Bzip2Compressor(;blocksize100k=$(DEFAULT_BLOCKSIZE100K), workfactor=$(DEFAULT_WORKFACTOR), verbosity=$(DEFAULT_VERBOSITY))

Create a bzip2 compression codec.

Arguments
---------
- `blocksize100k`: block size to be use for compression (1..9)
- `workfactor`: amount of effort the standard algorithm will expend before resorting to the fallback (0..250)
- `verbosity`: verbosity level (0..4)
"""
function Bzip2Compressor(;blocksize100k::Integer=DEFAULT_BLOCKSIZE100K,
                           workfactor::Integer=DEFAULT_WORKFACTOR,
                           verbosity::Integer=DEFAULT_VERBOSITY)
    if !(1 ≤ blocksize100k ≤ 9)
        throw(ArgumentError("blocksize100k must be within 1..9"))
    elseif !(0 ≤ workfactor ≤ 250)
        throw(ArgumentError("workfactor must be within 0..250"))
    elseif !(0 ≤ verbosity ≤ 4)
        throw(ArgumentError("verbosity must be within 0..4"))
    end
    return Bzip2Compressor(BZStream(), blocksize100k, workfactor, verbosity)
end

const Bzip2CompressorStream{S} = TranscodingStream{Bzip2Compressor,S} where S<:IO

"""
    Bzip2CompressorStream(stream::IO; kwargs...)

Create a bzip2 compression stream (see `Bzip2Compressor` for `kwargs`).
"""
function Bzip2CompressorStream(stream::IO; kwargs...)
    x, y = splitkwargs(kwargs, (:blocksize100k, :workfactor, :verbosity))
    return TranscodingStream(Bzip2Compressor(;x...), stream; y...)
end

# Methods
# -------

function TranscodingStreams.finalize(codec::Bzip2Compressor)
    if codec.stream.state != C_NULL
        code = compress_end!(codec.stream)
        if code != BZ_OK
            bzerror(codec.stream, code)
        end
    end
    return
end

function TranscodingStreams.startproc(codec::Bzip2Compressor, ::Symbol, error::Error)
    if codec.stream.state != C_NULL
        code = compress_end!(codec.stream)
        if code != BZ_OK
            error[] = BZ2Error(code)
            return :error
        end
    end
    code = compress_init!(codec.stream, codec.blocksize100k, codec.verbosity, codec.workfactor)
    if code != BZ_OK
        error[] = BZ2Error(code)
        return :error
    end
    return :ok
end

function TranscodingStreams.process(codec::Bzip2Compressor, input::Memory, output::Memory, error::Error)
    stream = codec.stream
    stream.next_in = input.ptr
    stream.avail_in = input.size
    stream.next_out = output.ptr
    stream.avail_out = output.size
    code = compress!(stream, input.size > 0 ? BZ_RUN : BZ_FINISH)
    Δin = Int(input.size - stream.avail_in)
    Δout = Int(output.size - stream.avail_out)
    if code == BZ_RUN_OK || code == BZ_FINISH_OK
        return Δin, Δout, :ok
    elseif code == BZ_STREAM_END
        return Δin, Δout, :end
    else
        error[] = BZ2Error(code)
        return Δin, Δout, :error
    end
end


text:
using OrdinaryDiffEq
using Trixi

###############################################################################
# semidiscretization of the compressible Euler equations

equations = CompressibleEulerEquations2D(1.4)

"""
    initial_condition_blast_wave(x, t, equations::CompressibleEulerEquations2D)

A medium blast wave taken from
- Sebastian Hennemann, Gregor J. Gassner (2020)
  A provably entropy stable subcell shock capturing approach for high order split form DG
  [arXiv: 2008.12044](https://arxiv.org/abs/2008.12044)
"""
function initial_condition_blast_wave(x, t, equations::CompressibleEulerEquations2D)
  # Modified From Hennemann & Gassner JCP paper 2020 (Sec. 6.3) -> "medium blast wave"
  # Set up polar coordinates
  inicenter = SVector(0.0, 0.0)
  x_norm = x[1] - inicenter[1]
  y_norm = x[2] - inicenter[2]
  r = sqrt(x_norm^2 + y_norm^2)
  phi = atan(y_norm, x_norm)
  sin_phi, cos_phi = sincos(phi)

  # Calculate primitive variables
  rho = r > 0.5 ? 1.0 : 1.1691
  v1  = r > 0.5 ? 0.0 : 0.1882 * cos_phi
  v2  = r > 0.5 ? 0.0 : 0.1882 * sin_phi
  p   = r > 0.5 ? 1.0E-3 : 1.245

  return prim2cons(SVector(rho, v1, v2, p), equations)
end
initial_condition = initial_condition_blast_wave

surface_flux = flux_lax_friedrichs
volume_flux  = flux_ranocha
basis = LobattoLegendreBasis(3)
indicator_sc = IndicatorHennemannGassner(equations, basis,
                                         alpha_max=0.5,
                                         alpha_min=0.001,
                                         alpha_smooth=true,
                                         variable=density_pressure)
volume_integral = VolumeIntegralShockCapturingHG(indicator_sc;
                                                 volume_flux_dg=volume_flux,
                                                 volume_flux_fv=surface_flux)
solver = DGSEM(basis, surface_flux, volume_integral)

coordinates_min = (-2.0, -2.0)
coordinates_max = ( 2.0,  2.0)
mesh = TreeMesh(coordinates_min, coordinates_max,
                initial_refinement_level=6,
                n_cells_max=10_000)

semi = SemidiscretizationHyperbolic(mesh, equations, initial_condition, solver)

###############################################################################
# ODE solvers, callbacks etc.

tspan = (0.0, 12.5)
ode = semidiscretize(semi, tspan)

summary_callback = SummaryCallback()

analysis_interval = 100
analysis_callback = AnalysisCallback(semi, interval=analysis_interval)

alive_callback = AliveCallback(analysis_interval=analysis_interval)

save_solution = SaveSolutionCallback(interval=100,
                                     save_initial_solution=true,
                                     save_final_solution=true,
                                     solution_variables=cons2prim)

stepsize_callback = StepsizeCallback(cfl=0.9)

callbacks = CallbackSet(summary_callback,
                        analysis_callback, alive_callback,
                        save_solution,
                        stepsize_callback)

###############################################################################
# run the simulation

sol = solve(ode, CarpenterKennedy2N54(williamson_condition=false),
            dt=1.0, # solve needs some value here but it will be overwritten by the stepsize_callback
            save_everystep=false, callback=callbacks);
summary_callback() # print the timer summary


text:
####################
#
# Polygon
#
####################

# Union of Shapefile Polygon types
ShpPolygon = Union{Shapefile.Polygon, Shapefile.PolygonZ, Shapefile.PolygonM}

#=
For an array of ESRIShape points
1. convert to a different coordinate system (optional)
2. convert to a Compose polygon

Note:
Units in Compose can be either:
- Context - cx, cy - using the axis coordinates
- Relative - w (width), h (height) - using the size of the image
- Absolute - pt, mm, cm, inch
If you just provide numbers e.g. if pts below was an array of tuples of Float64
the units would default to context units.
=#
function ESRItoComposePolygon(points::AbstractArray{T, 1},
                              convertcoords) where {T<:ShpPoint}
    pts = Array{Tuple{Measures.Length{:cx, Float64}, Measures.Length{:cy, Float64}}}(undef, 0)
    for pt in points
        x, y = convertcoords(pt.x, pt.y)
        push!(pts, (x * cx, y * cy))
    end
    Compose.polygon(pts)
end

# Plot a single part polygon
function compose_single(polygon::T, canvas, convertcoords, line_width,
                        line_color, fill_color) where {T<:ShpPolygon}
    compose(canvas, (context(), ESRItoComposePolygon(polygon.points, convertcoords),
                     linewidth(line_width), stroke(line_color), fill(fill_color)))
end

# Plot a multi part polygon
function compose_multi(polygon::T, canvas, convertcoords, line_width,
                       line_color, fill_color) where {T<:ShpPolygon}
    start = 1
    finish = length(polygon.points)
    for idx in polygon.parts[2:end]
        canvas = compose(canvas, (context(),
                                  ESRItoComposePolygon(polygon.points[start:idx], convertcoords),
                                  linewidth(line_width), stroke(line_color), fill(fill_color)))
        start = idx + 1
    end
    canvas = compose(canvas, (context(),
                              ESRItoComposePolygon(polygon.points[start:finish], convertcoords),
                              linewidth(line_width), stroke(line_color), fill(fill_color)))
end

# Plot an array of polygons
function draw_shp(shapes::AbstractArray{T, 1}, canvas, convertcoords,
                  line_width, line_color, fill_color, radius) where {T<:ShpPolygon}
    for polygon in shapes
        if length(polygon.parts) == 1
            canvas = compose_single(polygon, canvas, convertcoords, line_width, line_color,
                                    fill_color)
        else
            canvas = compose_multi(polygon, canvas, convertcoords, line_width, line_color,
                                   fill_color)
        end
    end
    canvas
end


text:
using BandedMatrices, LinearAlgebra, Test, Random
import BandedMatrices: _BandedMatrix

# set prng to some value that avoids test failure
Random.seed!(0)
struct _foo <: Number end

tldiv!(A, b) = ldiv!(transpose(A), b)
cldiv!(A, b) = ldiv!(adjoint(A), b)

@testset "Conversion to blas type" begin
    @testset "_promote_to_blas_type" begin
        typ = Float64
        @test BandedMatrices._promote_to_blas_type(typ, ComplexF64) == ComplexF64
        @test BandedMatrices._promote_to_blas_type(typ, ComplexF32)  == ComplexF64
        @test BandedMatrices._promote_to_blas_type(typ, Float64)    == typ
        @test BandedMatrices._promote_to_blas_type(typ, Float32)    == typ
        @test BandedMatrices._promote_to_blas_type(typ, Int64)      == typ
        @test BandedMatrices._promote_to_blas_type(typ, Int32)      == typ
        @test BandedMatrices._promote_to_blas_type(typ, Int16)      == typ
        @test BandedMatrices._promote_to_blas_type(typ, Int8)       == typ

        typ = Float32
        @test BandedMatrices._promote_to_blas_type(typ, ComplexF64) == ComplexF64
        @test BandedMatrices._promote_to_blas_type(typ, ComplexF32)  == ComplexF32
        @test BandedMatrices._promote_to_blas_type(typ, Float64)    == Float64
        @test BandedMatrices._promote_to_blas_type(typ, Float32)    == typ
        @test BandedMatrices._promote_to_blas_type(typ, Int64)      == typ
        @test BandedMatrices._promote_to_blas_type(typ, Int32)      == typ
        @test BandedMatrices._promote_to_blas_type(typ, Int16)      == typ
        @test BandedMatrices._promote_to_blas_type(typ, Int8)       == typ

        @test_throws ErrorException BandedMatrices._promote_to_blas_type(_foo, Float64)
    end

    @testset "ldiv!" begin
        As   = Any[_BandedMatrix(rand(1:10, 3, 5), 5, 1, 1),
                   _BandedMatrix(rand(3, 5)*im,    5, 1, 1),
                   _BandedMatrix(rand(3, 5),       5, 1, 1)
                  ]
        bs   = Any[rand(1:10,   5),
                   rand(1:10,   5),
                   rand(1:10.0, 5)*im,
                  ]
        typs = Any[Float64,
                   ComplexF64,
                   ComplexF64]

        for (A, b, typ) in zip(As, bs, typs)
            AA,   bb   = BandedMatrices._convert_to_blas_type(A,         b)
            AAlu, bblu = BandedMatrices._convert_to_blas_type(lu(A), b)
            @test eltype(AA) == eltype(bb) == eltype(AAlu) == eltype(bblu) == typ
            @test Matrix(A)\copy(b)             ≈ A\copy(b)
            @test Matrix(A)\copy(b)             ≈ lu(A)\copy(b)
            @test Matrix(A)\copy(b)             ≈ ldiv!(A, copy(b))
            @test Matrix(A)\copy(b)             ≈ ldiv!(lu(A), copy(b))
            @test transpose(Matrix(A))\copy(b)  ≈ tldiv!(lu(A), copy(b))
            @test transpose(Matrix(A))\copy(b)  ≈ ldiv!(lu(transpose(A)), copy(b))
            @test Matrix(A)'\copy(b)            ≈ cldiv!(lu(A), copy(b))
        end
    end
end

@testset "Banded A\b" begin
    let
        # banded
        A = brand(5, 1, 1)
        b = Float64[1, 2, 3, 4, 5]

        # dense storage
        Af = Matrix(A)
        bf = copy(b)

        @test Af\bf ≈ A\b

        # check A\b has not overwritten b, but has made a copy
        @test b == [1, 2, 3, 4, 5]
    end

    # advanced interface
    let
        # banded
        A = brand(5, 1, 1)
        b = rand(5)

        # dense storage
        Af = Matrix(A)
        bf = copy(b)

        # note lu makes copies; these need revision
        # once the lapack storage is built in to a BandedMatrix
        @test Af\bf ≈ lu(A)\copy(b)
        @test Af\bf ≈ ldiv!(A, copy(b))
        @test Af\bf ≈ ldiv!(lu(A), copy(b))
    end

    # conversion of inputs if needed
    let
        # factorisation performs conversion
        Ai = _BandedMatrix(rand(1:10, 3, 5), 5, 1, 1)
        @test eltype(lu(Ai)) == Float64

        # no op
        Af = _BandedMatrix(rand(Float32, 3, 5), 5, 1, 1)
        @test eltype(lu(Af)) == Float32

        # linear systems of integer data imply promotion
        Ai = _BandedMatrix(rand(1:10, 3, 5), 5, 1, 1)
        bi = collect(1:5)
        @test eltype(Ai\bi) == Float64
        # this code                     ≈ julia base
        @test Ai\bi                     ≈ Matrix(Ai)\copy(bi)
        @test lu(Ai)\bi             ≈ Matrix(Ai)\copy(bi)
        @test ldiv!(Ai, bi)         ≈ Matrix(Ai)\copy(bi)
        @test ldiv!(lu(Ai), bi) ≈ Matrix(Ai)\copy(bi)

        # check A\b makes a copy of b
        Ai = _BandedMatrix(rand(1:10, 3, 5), 5, 1, 1)
        bi = collect(1:5)
        Ai\bi
        @test bi == [1, 2, 3, 4, 5]
    end

    # matrix must be square error
    let
        for (n, m) in zip([7, 5], [5, 7])
            A = brand(m, n, 1, 1)
            b = rand(m)
            @test_throws DimensionMismatch A\b
            @test_throws DimensionMismatch lu(A)\b
            @test_throws DimensionMismatch ldiv!(A, b)
            @test_throws DimensionMismatch ldiv!(lu(A), b)
        end
    end

    # test transposed algorithm
    let
        # banded
        A = brand(5, 1, 1)
        b = rand(5)

        # dense storage
        Af = Matrix(A)
        bf = copy(b)

        # note lu makes copies; these need revision
        # once the lapack storage is built in to a BandedMatrix
        @test Af'\bf ≈ tldiv!(lu(A),  copy(b))
        @test Af'\bf ≈ ldiv!(lu(A'), copy(b))
    end

    @testset "complex input" begin
        # banded
        A = brand(5, 1, 1) + brand(5, 1, 1)*im
        b = rand(5) + rand(5)*im

        # dense storage
        Af = Matrix(A)
        bf = copy(b)

        # note lu makes copies; these need revision
        # once the lapack storage is built in to a BandedMatrix
        @test Af\bf  ≈ A\copy(b)
        @test Af\bf  ≈ lu(A)\copy(b)
        @test Af'\bf ≈ A'\copy(b)
        @test Af'\bf ≈ lu(A')\copy(b)
        @test Af\bf  ≈ ldiv!(lu(A),  copy(b))
        @test Af'\bf ≈ ldiv!(lu(A'), copy(b))
        @test transpose(Af)\bf ≈ tldiv!(lu(A),  copy(b))
        @test adjoint(Af)\bf ≈ ldiv!(lu(adjoint(A)), copy(b))
        @test adjoint(Af)\bf ≈ cldiv!(lu(A), copy(b))
    end

    @testset "multiple rhs" begin
        # banded
        A = brand(5, 1, 1)
        b = rand(5, 10)

        # dense storage
        Af = Matrix(A)
        bf = copy(b)

        # note lu makes copies; these need revision
        # once the lapack storage is built in to a BandedMatrix
        @test Af\bf  ≈ A\copy(b)
        @test Af\bf  ≈ lu(A)\copy(b)
        @test Af'\bf ≈ A'\copy(b)
        @test Af'\bf ≈ lu(A')\copy(b)
        @test Af\bf  ≈ ldiv!(lu(A),  copy(b))
        @test Af'\bf ≈ ldiv!(lu(A'), copy(b))
        @test Af'\bf ≈ tldiv!(lu(A),  copy(b))
        @test Af'\bf ≈ ldiv!(lu(A'), copy(b))
    end

    @testset "properties of factorisation" begin
        BLU = lu(brand(5, 4, 1, 1))
        @test size(BLU) == (5, 4)
        @test size(BLU, 1) == 5
        @test size(BLU, 2) == 4
        @test size(BLU, 3) == 1
        @test_throws ErrorException size(BLU, -1)
        @test_throws ErrorException size(BLU,  0)
    end
end


text:
using MLJ
using PrettyPrinting

KNNRegressor = @load KNNRegressor

X = (age    = [23, 45, 34, 25, 67],
     gender = categorical(['m', 'm', 'f', 'm', 'f']))

height = [178, 194, 165, 173, 168];

scitype(X.age)

pipe = Pipeline(
    coercer = X -> coerce(X, :age=>Continuous),
    one_hot_encoder = OneHotEncoder(),
    transformed_target_model = TransformedTargetModel(
        model = KNNRegressor(K=3);
        target=UnivariateStandardizer()
    )
)

pipe.transformed_target_model.model.K = 2
pipe.one_hot_encoder.drop_last = true;

evaluate(
    pipe,
    X,
    height,
    resampling=Holdout(),
    measure=rms
) |> pprint

# This file was generated using Literate.jl, https://github.com/fredrikekre/Literate.jl

max_issues_repo_name: alan-turing-institute/MLJTutorials

max_forks_repo_name: alan-turing-institute/MLJTutorials


text:
struct Disconnect <: Packet
    header::UInt8
end
Disconnect() = Disconnect(DISCONNECT)
Base.show(io::IO, x::Disconnect) = print(io, "disconnect")

max_issues_repo_name: kivaari/jlMQTT

max_forks_repo_name: kivaari/jlMQTT


text:
import Base:+,-,*,/,==
using Markdown

"""
    SqrtRational <: Real
A type to store exact value of square root of a `Rational`: ±√p. To store infinity values, use `Rational{BigInt}` to represent the `p`.
Basic functions such as `zero, one, inv, *, /` are supported for this type.
"""
struct SqrtRational <: Real
    sign::Bool
    value::Rational{BigInt}
    SqrtRational(s::Bool, x::Rational{BigInt}) = begin
        isinf(x) || x < 0 && throw(ArgumentError("invalid SqrtRational  √($x)"))
        new(s, x)
    end
end

SqrtRational(s::Bool, x::Integer) = SqrtRational(s, Rational{BigInt}(big(x)))
SqrtRational(s::Bool, x::Rational) = SqrtRational(s, big(x))
SqrtRational(x::Integer) = SqrtRational(false, Rational{BigInt}(big(x)))
SqrtRational(x::Rational) = SqrtRational(false, Rational{BigInt}(big(x)))

# some basic functions
Base.zero(::Type{SqrtRational}) = SqrtRational(zero(Rational{BigInt}))
Base.zero(x::SqrtRational) = SqrtRational(zero(Rational{BigInt}))
Base.one(::Type{SqrtRational}) = SqrtRational(one(Rational{BigInt}))
Base.one(x::SqrtRational) = SqrtRational(one(Rational{BigInt}))
Base.sign(x::SqrtRational) = iszero(x.value) ? 0 : x.sign ? -1 : 1
Base.signbit(x::SqrtRational) = x.sign
Base.inv(x::SqrtRational) = SqrtRational(x.sign, inv(x.value))

# override some operators
+(x::SqrtRational) = x
-(x::SqrtRational) = SqrtRational(!(x.sign), x.value)
*(x::SqrtRational, y::SqrtRational) = SqrtRational(x.sign ⊻ y.sign, x.value * y.value)
*(x::SqrtRational, y::Union{Rational, Integer}) = x * SqrtRational(signbit(y), abs(big(y))^2)
*(x::Union{Rational, Integer}, y::SqrtRational) = y * x
/(x::SqrtRational, y::Union{SqrtRational, Rational, Integer}) = x * inv(y)
/(x::Union{Rational, Integer}, y::SqrtRational) = SqrtRational(signbit(x), abs(big(x))^2) / y
==(x::SqrtRational, y::SqrtRational) = (sign(x) == sign(y)) & (x.value == y.value)

Base.float(x::SqrtRational) = sign(x) * sqrt(float(x.value))

"""
simplify a integer `n = x * t^2` to `(x, t)`
"""
function simplify(x::T) where T<:Integer
    s = sign(x)
    x = s * x
    out = one(T)
    i = 2
    while i <= floor(T, sqrt(x))
        i2 = i * i
        tp = div(x, i2)
        if tp * i2 == x
            out *= i
            x = tp
            i = 2
        else
            i += 1
        end
    end
    return out, s * x
end

"""
simplify a SqrtRational `√(p/q)` to x/t√(p1/q1)
"""
function simplify(x::SqrtRational)
    s = signbit(x)
    nu_out, nu_in = simplify(numerator(x.value))
    de_out, de_in = simplify(denominator(x.value))
    return s, nu_out, nu_in, de_out, de_in
end

Base.show(io::IO, ::MIME"text/plain", x::SqrtRational) = begin
    s, nu_out, nu_in, de_out, de_in = simplify(x)
    if nu_in == 0
        print(io, 0)
        return
    end
    s && print(io, '-')
    if nu_in == 1
        print(io, nu_out)
    elseif nu_out == 1
        print(io, "√", nu_in)
    else
        print(io, nu_out, "√", nu_in)
    end
    if de_out == de_in == 1
        return
    elseif de_out == 1
        print(io, "/√", de_in)
    elseif de_in == 1
        print(io, "/", de_out)
    else
        print(io, "/", de_out, "√", de_in)
    end
    return
end

Base.show(io::IO, x::SqrtRational) = show(io::IO, "text/plain", x)

Base.show(io::IO, ::MIME"text/markdown", x::SqrtRational) = begin
    s, nu_out, nu_in, de_out, de_in = simplify(x)
    if nu_in == 0
        show(io, "text/markdown", Markdown.parse("``0``"))
        return
    end
    to_show::String = ""
    if nu_in == 1
        to_show = "$nu_out"
    elseif nu_out == 1
        to_show = "\\sqrt{$nu_in}"
    else
        to_show = "$nu_out\\sqrt{$nu_in}"
    end
    if de_out == de_in == 1
        nothing
    elseif de_out == 1
        to_show = "\\dfrac{$to_show}{\\sqrt{$de_in}}"
    elseif de_in == 1
        to_show = "\\dfrac{$to_show}{$de_out}"
    else
        to_show = "\\dfrac{$to_show}{$de_out\\sqrt{$de_in}}"
    end
    to_show = (s ? "-" : "") * to_show
    show(io, "text/markdown", Markdown.parse("``$to_show``"))
end


text:
module FESpaceWithConstantFixedTests
using Gridap
using Gridap.Geometry
using Gridap.FESpaces
using Test

domain = (0,1,0,1)
partition = (4,4)
model = CartesianDiscreteModel(domain,partition)

trian = get_triangulation(model)

order = 2

V = TestFESpace(
  model=model,
  reffe=:Lagrangian,
  valuetype=Float64,
  order=order,
  conformity=:L2)

V0 = FESpaceWithConstantFixed(V,true,rand(1:num_free_dofs(V)))
test_single_field_fe_space(V0)

@test Gridap.FESpaces.get_constant_approach(V0) == Gridap.FESpaces.FixConstant()

uh0 = interpolate(V0) do x
    sin(4*pi*(x[1]+x[2]^2)) + 3
end
#using Gridap.Visualization
#
#writevtk(trian,"trian",nsubcells=20,cellfields=["uh0"=>uh0])

end # module


text:
#
abstract type Equation end

#----------------------------------------------------------------------
export Diffusion
#----------------------------------------------------------------------
struct Diffusion{T,U} <: Equation # {T,U,D,K} # type, dimension, (bdfK order)

    fld::Field{T}

    ν  ::Array{T} # viscosity
    f  ::Array{T} # forcing
    rhs::Array{T} # RHS

    tstep::TimeStepper{T,U}

    msh::Mesh{T} # mesh
end
#--------------------------------------#
function Diffusion(bc::Array{Char,1},msh::Mesh
                  ;Ti=0.,Tf=0.,dt=0.,k=3)

    fld = Field(bc,msh)
    ν   = zero(fld.u)
    f   = zero(fld.u)
    rhs = zero(fld.u)

    tstep = TimeStepper(Ti,Tf,dt,k)

    return Diffusion(fld
                    ,ν,f,rhs
                    ,tstep
                    ,msh)
end
#----------------------------------------------------------------------
function opLHS(u::Array,dfn::Diffusion)
    @unpack fld, msh, ν = dfn
    @unpack bdfB = dfn.tstep

    lhs = hlmz(u,ν,bdfB[1],msh)

    lhs .= gatherScatter(lhs,msh)
    lhs .= mask(lhs,fld.M)
    return lhs
end

function opPrecond(u::Array,dfn::Diffusion)
    return u
end

function makeRHS!(dfn::Diffusion)
    @unpack fld, rhs, ν, f, msh = dfn
    @unpack bdfA, bdfB = dfn.tstep

    rhs  .=      mass(f     ,msh) # forcing
    rhs .-= ν .* lapl(fld.ub,msh) # boundary data

    for i=1:length(fld.uh)        # histories
        rhs .-= bdfB[1+i] .* mass(fld.uh[i],msh)
    end

    rhs .= mask(rhs,fld.M)
    rhs .= gatherScatter(rhs,msh)
    return
end

function solve!(dfn::Diffusion)
    @unpack rhs, msh, fld = dfn
    @unpack u,ub = fld

    opL(u) = opLHS(u,dfn)
    opP(u) = opPrecond(u,dfn)

    pcg!(u,rhs,opL;opM=opP,mult=msh.mult,ifv=false)
    u .+= ub
    return
end
#----------------------------------------------------------------------
export evolve!
#----------------------------------------------------------------------
function evolve!(dfn::Diffusion
                ,setBC! =fixU!
                ,setForcing! =fixU!
                ,setVisc! =fixU!)

    @unpack fld, f, ν, msh = dfn
    @unpack time, bdfA, bdfB, istep, dt = dfn.tstep

    updateHist!(fld)

    Zygote.ignore() do
        updateHist!(time)
        istep  .+= 1
        time[1] += dt[1]
        bdfExtK!(bdfA,bdfB,time)
    end

    setBC!(fld.ub,msh.x,msh.y,time[1])
    setForcing!(f,msh.x,msh.y,time[1])
    setVisc!(ν   ,msh.x,msh.y,time[1])

    makeRHS!(dfn)
    solve!(dfn)

    return
end
#----------------------------------------------------------------------
export simulate!
#----------------------------------------------------------------------
function simulate!(dfn::Diffusion,callback!::Function
                  ,setIC! =fixU!
                  ,setBC! =fixU!
                  ,setForcing! =fixU!
                  ,setVisc! =fixU!)

    @unpack fld, msh = dfn
    @unpack time, istep, dt, Tf = dfn.tstep

    setIC!(fld.u,msh.x,msh.y,time[1])

    Zygote.ignore() do
        callback!(dfn)
    end
    while time[1] <= Tf[1]

        evolve!(dfn,setBC!,setForcing!,setVisc!)

        Zygote.ignore() do
            callback!(dfn)
        end

        if(time[1] < 1e-12) break end

    end

    return
end
#----------------------------------------------------------------------
#


text:
# ## Data

import CSV

using DataFrames
using TuringModels: project_root

path = joinpath(project_root, "data", "Kline.csv")
df = CSV.read(path, DataFrame; delim=';')
df.log_pop = log.(df.population)
df.society = 1:nrow(df)
df

# ## Model

using Turing

@model function m12_6(total_tools, log_pop, society)
    N = length(total_tools)

    α ~ Normal(0, 10)
    βp ~ Normal(0, 1)

    σ_society ~ truncated(Cauchy(0, 1), 0, Inf)

    N_society = length(unique(society)) ## 10

    α_society ~ filldist(Normal(0, σ_society), N_society)

    for i in 1:N
        λ = exp(α + α_society[society[i]] + βp*log_pop[i])
        total_tools[i] ~ Poisson(λ)
    end
end;

# ## Output

chains = sample(
    m12_6(df.total_tools, df.log_pop, df.society),
    NUTS(0.95),
    1000
)

# \defaultoutput{}

# ## Original output

m12_6rethinking = "
              Mean StdDev lower 0.89 upper 0.89 n_eff Rhat
a              1.11   0.75      -0.05       2.24  1256    1
bp             0.26   0.08       0.13       0.38  1276    1
a_society[1]  -0.20   0.24      -0.57       0.16  2389    1
a_society[2]   0.04   0.21      -0.29       0.38  2220    1
a_society[3]  -0.05   0.19      -0.36       0.25  3018    1
a_society[4]   0.32   0.18       0.01       0.60  2153    1
a_society[5]   0.04   0.18      -0.22       0.33  3196    1
a_society[6]  -0.32   0.21      -0.62       0.02  2574    1
a_society[7]   0.14   0.17      -0.13       0.40  2751    1
a_society[8]  -0.18   0.19      -0.46       0.12  2952    1
a_society[9]   0.27   0.17      -0.02       0.52  2540    1
a_society[10] -0.10   0.30      -0.52       0.37  1433    1
sigma_society  0.31   0.13       0.11       0.47  1345    1
";


text:
name_and_root(b::Atom.ActualLocalBinding) = (b.name, b.root)
name(b::Atom.ActualLocalBinding) = b.name

let str = """
    function local_bindings(expr, bindings = [], pos = 1)
        bind = CSTParser.bindingof(expr)
        scope = CSTParser.scopeof(expr)
        if bind !== nothing && scope === nothing
            push!(bindings, LocalBinding(bind.name, pos:pos+expr.span))
        end
        if scope !== nothing
            if expr.typ == CSTParser.Kw
                return bindings
            end
            range = pos:pos+expr.span
            localbindings = []
            if expr.args !== nothing
                for (i, arg) in expr.args
                    ### L15 -- FOR THE FIRST TEST ###
                    local_bindings(arg, text, localbindings, pos, line)
                    line += count(c -> c === '\n', text[nextind(text, pos - 1):prevind(text, pos + arg.fullspan)])
                    pos += arg.fullspan
                end
            end
            if expr.typ == CSTParser.TupleH # deconstruct tuple expression
                ### L21 - FOR THE SECOND TEST ###
                for leftbind in localbindings
                    push!(bindings, leftbind)
                end
            else
                push!(bindings, LocalScope(bind === nothing ? "" : bind.name, range, line, localbindings, expr))
            end
            return bindings
        elseif expr.args !== nothing
            for arg in expr.args
                local_bindings(arg, bindings, pos)
                pos += arg.fullspan
            end
        end
        return bindings
    end
    """
    let ls = Set(name_and_root.(Atom.locals(str, 15, 1)))
        for l in [
                ("i", "")
                ("arg", "")
                ("localbindings", "local_bindings")
                ("range", "local_bindings")
                ("scope", "local_bindings")
                ("bind", "local_bindings")
                ("pos", "local_bindings")
                ("bindings", "local_bindings")
                ("expr", "local_bindings")
                ("local_bindings", "")
            ]
            @test l in ls
        end
    end

    let ls = Set(name_and_root.(Atom.locals(str, 21, 100)))
        for l in [
                ("localbindings", "local_bindings")
                ("pos", "local_bindings")
                ("scope", "local_bindings")
                ("bind", "local_bindings")
                ("bindings", "local_bindings")
                ("expr", "local_bindings")
                ("local_bindings", "")
                ("range", "local_bindings")
                ("local_bindings", "")
            ]
            @test l in ls
        end
    end
end

let str = """
    const bar = 2
    function f(x=2) # `x` should show up in completions
        ff = function (x, xxx)
            z = 3
        end
        return y+x
    end
    kwfun(kw = 3) # `kw` should not show up in completions
    function foo(x)
        asd = 3
    end
    """
    outers = (("bar", ""), ("f", ""), ("foo", ""))
    let ls = Set(name_and_root.(Atom.locals(str, 1, 1)))
        for o in outers; @test o in ls; end
    end
    let ls = Set(name_and_root.(Atom.locals(str, 2, 1)))
        for o in outers; @test o in ls; end
        @test ("x", "f") in ls
        @test ("ff", "f") in ls
    end
    let ls = Set(name_and_root.(Atom.locals(str, 4, 100)))
        for o in outers; @test o in ls; end
        @test ("ff", "f") in ls
        @test ("x", "") in ls
        @test ("xxx", "") in ls
        @test ("z", "") in ls
    end
    let ls = Set(name_and_root.(Atom.locals(str, 10, 100)))
        for o in outers; @test o in ls; end
        @test ("x", "foo") in ls
        @test ("asd", "foo") in ls
    end
end

# destructure multiple return expression
let str = """
    function foo()
        tpl1, tpl2 = (1, 2)
        shown1, shown2 = (ntpl1 = 1, ntpl2 = 2)
    end
    """

    ls = Atom.locals(str, 1, 1)
    # basic
    for l in filter(l -> l.line == 2, ls)
        @test l.name in ("tpl1", "tpl2")
        @test l.root == "foo"
        @test l.verbatim == "tpl1, tpl2 = (1, 2)"
    end
    # named tuple elements shouldn't show up in completions
    for l in filter(l -> l.line == 3, ls)
        @test l.name in ("shown1", "shown2")
        @test l.root == "foo"
        @test l.verbatim == "shown1, shown2 = (ntpl1 = 1, ntpl2 = 2)"
    end
end

# detect all the method parameters of a method with a `where` clause
let str = """
    function push!′(ary::Vector{T}, item::S) where {T, S<:T}
        tmpvec, tmpitem = ary, item
        push!(tmpvec, tmpitem)
    end
    """

    binds = name.(Atom.locals(str, 1, 0))
    @test "ary" in binds
    @test "item" in binds
    @test "T" in binds
    @test "S" in binds
    @test "tmpvec" in binds
    @test "tmpitem" in binds
    @test "push!′" in binds
end

# finds function arguments when the function is type-declared
# ref: https://github.com/JunoLab/Atom.jl/issues/223
let str = """
    function func(ary::Vector{T}, i::Int)::T where {T<:Number}
        # here
        s = sum(ary[1:i])
        return s
    end
    """

    binds = name.(Atom.locals(str, 2, 4))
    @test "ary" in binds
    @test "i" in binds
    @test "T" in binds
    @test "func" in binds
end

let str = """
    function foo(x)
        @macrocall begin
            xyz = 2
        end
        xxx = 3
        return 12
    end
    """
    let ls = Set(name_and_root.(Atom.locals(str, 2, 1)))
        for l in [
                ("foo", "")
                ("xxx", "foo")
                ("x", "foo")
            ]
            @test l in ls
        end
    end
    let ls = Set(name_and_root.(Atom.locals(str, 3, 1)))
        for l in [
                ("foo", "")
                ("xxx", "foo")
                ("xyz", "")
                ("x", "foo")
            ]
            @test l in ls
        end
    end
    let ls = Set(name_and_root.(Atom.locals(str, 5, 1)))
        for l in [
                ("foo", "")
                ("xxx", "foo")
                ("x", "foo")
            ]
            @test l in ls
        end
    end
end


text:
# syntax: proto2
using Compat
using ProtoBuf
import ProtoBuf.meta

struct __enum_Status <: ProtoEnum
    Idle::Int32
    SolivingODEInProgress::Int32
    __enum_Status() = new(0,1)
end #struct __enum_Status
const Status = __enum_Status()

struct __enum_ProblemType <: ProtoEnum
    ODEProblem::Int32
    __enum_ProblemType() = new(0)
end #struct __enum_ProblemType
const ProblemType = __enum_ProblemType()

struct __enum_ReturnCode <: ProtoEnum
    Default::Int32
    Success::Int32
    MaxIters::Int32
    DtLessThanMin::Int32
    Unstable::Int32
    InitialFailure::Int32
    ConvergenceFailure::Int32
    Failure::Int32
    __enum_ReturnCode() = new(0,1,2,3,4,5,6,7)
end #struct __enum_ReturnCode
const ReturnCode = __enum_ReturnCode()

struct __enum_MessageType <: ProtoEnum
    SolveODE::Int32
    ODEResult::Int32
    StatusUpdate::Int32
    StopServer::Int32
    Error::Int32
    __enum_MessageType() = new(1,2,3,4,5)
end #struct __enum_MessageType
const MessageType = __enum_MessageType()

mutable struct Nuclide <: ProtoType
    A::UInt32
    Z::UInt32
    q::UInt32
    i::UInt32
    Nuclide(; kwargs...) = (o=new(); fillunset(o); isempty(kwargs) || ProtoBuf._protobuild(o, kwargs); o)
end #mutable struct Nuclide
const __req_Nuclide = Symbol[:A,:Z,:q,:i]
meta(t::Type{Nuclide}) = meta(t, __req_Nuclide, ProtoBuf.DEF_FNUM, ProtoBuf.DEF_VAL, true, ProtoBuf.DEF_PACK, ProtoBuf.DEF_WTYPES, ProtoBuf.DEF_ONEOFS, ProtoBuf.DEF_ONEOF_NAMES, ProtoBuf.DEF_FIELD_TYPES)

mutable struct MatrixValue <: ProtoType
    value::Float64
    row::UInt32
    column::UInt32
    MatrixValue(; kwargs...) = (o=new(); fillunset(o); isempty(kwargs) || ProtoBuf._protobuild(o, kwargs); o)
end #mutable struct MatrixValue
const __req_MatrixValue = Symbol[:value,:row,:column]
meta(t::Type{MatrixValue}) = meta(t, __req_MatrixValue, ProtoBuf.DEF_FNUM, ProtoBuf.DEF_VAL, true, ProtoBuf.DEF_PACK, ProtoBuf.DEF_WTYPES, ProtoBuf.DEF_ONEOFS, ProtoBuf.DEF_ONEOF_NAMES, ProtoBuf.DEF_FIELD_TYPES)

mutable struct Update <: ProtoType
    status::Int32
    message::AbstractString
    Update(; kwargs...) = (o=new(); fillunset(o); isempty(kwargs) || ProtoBuf._protobuild(o, kwargs); o)
end #mutable struct Update
const __req_Update = Symbol[:status,:message]
meta(t::Type{Update}) = meta(t, __req_Update, ProtoBuf.DEF_FNUM, ProtoBuf.DEF_VAL, true, ProtoBuf.DEF_PACK, ProtoBuf.DEF_WTYPES, ProtoBuf.DEF_ONEOFS, ProtoBuf.DEF_ONEOF_NAMES, ProtoBuf.DEF_FIELD_TYPES)

mutable struct Progress <: ProtoType
    time::Float64
    Progress(; kwargs...) = (o=new(); fillunset(o); isempty(kwargs) || ProtoBuf._protobuild(o, kwargs); o)
end #mutable struct Progress
const __req_Progress = Symbol[:time]
meta(t::Type{Progress}) = meta(t, __req_Progress, ProtoBuf.DEF_FNUM, ProtoBuf.DEF_VAL, true, ProtoBuf.DEF_PACK, ProtoBuf.DEF_WTYPES, ProtoBuf.DEF_ONEOFS, ProtoBuf.DEF_ONEOF_NAMES, ProtoBuf.DEF_FIELD_TYPES)

mutable struct TimeSpan <: ProtoType
    start::Float64
    stop::Float64
    TimeSpan(; kwargs...) = (o=new(); fillunset(o); isempty(kwargs) || ProtoBuf._protobuild(o, kwargs); o)
end #mutable struct TimeSpan
const __req_TimeSpan = Symbol[:start,:stop]
meta(t::Type{TimeSpan}) = meta(t, __req_TimeSpan, ProtoBuf.DEF_FNUM, ProtoBuf.DEF_VAL, true, ProtoBuf.DEF_PACK, ProtoBuf.DEF_WTYPES, ProtoBuf.DEF_ONEOFS, ProtoBuf.DEF_ONEOF_NAMES, ProtoBuf.DEF_FIELD_TYPES)

mutable struct InitialValue <: ProtoType
    index::UInt32
    number_of_particles::Float64
    temperature_in_ev::Float64
    InitialValue(; kwargs...) = (o=new(); fillunset(o); isempty(kwargs) || ProtoBuf._protobuild(o, kwargs); o)
end #mutable struct InitialValue
const __req_InitialValue = Symbol[:index,:number_of_particles,:temperature_in_ev]
meta(t::Type{InitialValue}) = meta(t, __req_InitialValue, ProtoBuf.DEF_FNUM, ProtoBuf.DEF_VAL, true, ProtoBuf.DEF_PACK, ProtoBuf.DEF_WTYPES, ProtoBuf.DEF_ONEOFS, ProtoBuf.DEF_ONEOF_NAMES, ProtoBuf.DEF_FIELD_TYPES)

mutable struct DiffEqParameters <: ProtoType
    qVe::Base.Vector{Float64}
    qVt::Base.Vector{Float64}
    mass_number::Base.Vector{Float64}
    spitzer_divided_by_overlap::Base.Vector{Float64}
    q::Base.Vector{Float64}
    inverted_collision_constant::Base.Vector{MatrixValue}
    dCharge_ex_divided_by_N_times_tau::Base.Vector{MatrixValue}
    rate_of_change_divided_by_N::Base.Vector{MatrixValue}
    no_dimensions::UInt32
    V_0::Float64
    r_e::Float64
    r_dt::Float64
    l_dt::Float64
    initial_values::Base.Vector{InitialValue}
    initial_temperature::Float64
    minimum_N::Float64
    source_terms_n::Base.Vector{Float64}
    source_terms_kt::Base.Vector{Float64}
    DiffEqParameters(; kwargs...) = (o=new(); fillunset(o); isempty(kwargs) || ProtoBuf._protobuild(o, kwargs); o)
end #mutable struct DiffEqParameters
const __req_DiffEqParameters = Symbol[:no_dimensions,:V_0,:r_e,:r_dt,:l_dt,:initial_temperature,:minimum_N]
const __fnum_DiffEqParameters = Int[1,2,3,4,15,6,7,8,9,16,17,18,19,10,11,12,13,14]
meta(t::Type{DiffEqParameters}) = meta(t, __req_DiffEqParameters, __fnum_DiffEqParameters, ProtoBuf.DEF_VAL, true, ProtoBuf.DEF_PACK, ProtoBuf.DEF_WTYPES, ProtoBuf.DEF_ONEOFS, ProtoBuf.DEF_ONEOF_NAMES, ProtoBuf.DEF_FIELD_TYPES)

mutable struct ProblemParameters <: ProtoType
    problem_type::Int32
    time_span::TimeSpan
    ProblemParameters(; kwargs...) = (o=new(); fillunset(o); isempty(kwargs) || ProtoBuf._protobuild(o, kwargs); o)
end #mutable struct ProblemParameters
const __req_ProblemParameters = Symbol[:problem_type,:time_span]
meta(t::Type{ProblemParameters}) = meta(t, __req_ProblemParameters, ProtoBuf.DEF_FNUM, ProtoBuf.DEF_VAL, true, ProtoBuf.DEF_PACK, ProtoBuf.DEF_WTYPES, ProtoBuf.DEF_ONEOFS, ProtoBuf.DEF_ONEOF_NAMES, ProtoBuf.DEF_FIELD_TYPES)

mutable struct SolverParameters <: ProtoType
    saveat::Base.Vector{Float64}
    SolverParameters(; kwargs...) = (o=new(); fillunset(o); isempty(kwargs) || ProtoBuf._protobuild(o, kwargs); o)
end #mutable struct SolverParameters

mutable struct SolveODEProblem <: ProtoType
    nuclides::Base.Vector{Nuclide}
    problem_parameters::ProblemParameters
    diff_eq_parameters::DiffEqParameters
    solver_parameters::SolverParameters
    SolveODEProblem(; kwargs...) = (o=new(); fillunset(o); isempty(kwargs) || ProtoBuf._protobuild(o, kwargs); o)
end #mutable struct SolveODEProblem
const __req_SolveODEProblem = Symbol[:problem_parameters,:diff_eq_parameters,:solver_parameters]
meta(t::Type{SolveODEProblem}) = meta(t, __req_SolveODEProblem, ProtoBuf.DEF_FNUM, ProtoBuf.DEF_VAL, true, ProtoBuf.DEF_PACK, ProtoBuf.DEF_WTYPES, ProtoBuf.DEF_ONEOFS, ProtoBuf.DEF_ONEOF_NAMES, ProtoBuf.DEF_FIELD_TYPES)

mutable struct ValuesPerNuclide <: ProtoType
    nuclide::Nuclide
    values::Base.Vector{Float64}
    ValuesPerNuclide(; kwargs...) = (o=new(); fillunset(o); isempty(kwargs) || ProtoBuf._protobuild(o, kwargs); o)
end #mutable struct ValuesPerNuclide
const __req_ValuesPerNuclide = Symbol[:nuclide]
meta(t::Type{ValuesPerNuclide}) = meta(t, __req_ValuesPerNuclide, ProtoBuf.DEF_FNUM, ProtoBuf.DEF_VAL, true, ProtoBuf.DEF_PACK, ProtoBuf.DEF_WTYPES, ProtoBuf.DEF_ONEOFS, ProtoBuf.DEF_ONEOF_NAMES, ProtoBuf.DEF_FIELD_TYPES)

mutable struct Result <: ProtoType
    problem::SolveODEProblem
    return_code::Int32
    start_time::Float64
    stop_time::Float64
    times::Base.Vector{Float64}
    n::Base.Vector{ValuesPerNuclide}
    kT::Base.Vector{ValuesPerNuclide}
    Result(; kwargs...) = (o=new(); fillunset(o); isempty(kwargs) || ProtoBuf._protobuild(o, kwargs); o)
end #mutable struct Result
const __req_Result = Symbol[:problem,:return_code,:start_time,:stop_time]
meta(t::Type{Result}) = meta(t, __req_Result, ProtoBuf.DEF_FNUM, ProtoBuf.DEF_VAL, true, ProtoBuf.DEF_PACK, ProtoBuf.DEF_WTYPES, ProtoBuf.DEF_ONEOFS, ProtoBuf.DEF_ONEOF_NAMES, ProtoBuf.DEF_FIELD_TYPES)

mutable struct ErrorEncountered <: ProtoType
    msg::AbstractString
    ErrorEncountered(; kwargs...) = (o=new(); fillunset(o); isempty(kwargs) || ProtoBuf._protobuild(o, kwargs); o)
end #mutable struct ErrorEncountered
const __req_ErrorEncountered = Symbol[:msg]
meta(t::Type{ErrorEncountered}) = meta(t, __req_ErrorEncountered, ProtoBuf.DEF_FNUM, ProtoBuf.DEF_VAL, true, ProtoBuf.DEF_PACK, ProtoBuf.DEF_WTYPES, ProtoBuf.DEF_ONEOFS, ProtoBuf.DEF_ONEOF_NAMES, ProtoBuf.DEF_FIELD_TYPES)

mutable struct Message <: ProtoType
    msg_type::Int32
    ode_problem::SolveODEProblem
    ode_result::Result
    status::Int32
    err::ErrorEncountered
    progress::Progress
    Message(; kwargs...) = (o=new(); fillunset(o); isempty(kwargs) || ProtoBuf._protobuild(o, kwargs); o)
end #mutable struct Message
const __req_Message = Symbol[:msg_type]
meta(t::Type{Message}) = meta(t, __req_Message, ProtoBuf.DEF_FNUM, ProtoBuf.DEF_VAL, true, ProtoBuf.DEF_PACK, ProtoBuf.DEF_WTYPES, ProtoBuf.DEF_ONEOFS, ProtoBuf.DEF_ONEOF_NAMES, ProtoBuf.DEF_FIELD_TYPES)

export Status, ProblemType, ReturnCode, MessageType, Nuclide, MatrixValue, Update, Progress, TimeSpan, InitialValue, DiffEqParameters, ProblemParameters, SolverParameters, SolveODEProblem, ValuesPerNuclide, Result, ErrorEncountered, Message


text:
@testset "Test freecf $mat, $obj" for mat in ["her", "rec"], obj in ["kurt", "ent"]

    # test for the hermitian matrices
    if mat == "her"
        for idx = 1: 5
            # set up
            N = 300
            G1, G2 = randn(N, N), randn(N, 2N);
            X1, X2 = (G1 + G1') / sqrt(2*N), (G2 * G2') / (2*N)
            X1 = free_whiten([X1]; mat = "her")[1][1]
            X2 = free_whiten([X2]; mat = "her")[1][1]
            X = [X1, X2];
            A = randn(2,2)
            Z = A*X

            # apply freecf
            Aest = freecf(Z; mat = mat, obj = obj, opt = "orth")[1]
            @test err_pd(pinv(Aest)*A) < 5e-2
        end
    end

    # test for the rectangular matrices
    if mat == "rec"
        for idx = 1: 5
            # set up
            N, M = 300, 500
            X1 = randn(N, M) / sqrt(M)
            X1 = free_whiten([X1], mat = "rec")[1][1]
            U, V = Matrix(qr(randn(N,N)).Q), Matrix(qr(randn(M,M)).Q)
            D = [Diagonal((collect(range(1,stop = 0,length = N)) .- 1).^4) zeros(N, M - N)]
            X2 = U*D*V'
            X2 = free_whiten([X2], mat = "rec")[1][1]
            X = [X1, X2];
            A = randn(2,2)
            Z = A*X

            # apply freecf
            Aest = freecf(Z; mat = mat, obj = obj, opt = "orth")[1]
            @test err_pd(pinv(Aest)*A) < 5e-2
        end
    end
end

@testset "Test freecf $mat, $obj" for mat in ["her", "rec"], obj in ["kurt"]

    # sphe opt is not designed for entropy based fcf

    # test for the hermitian matrices
    if mat == "her"
        for idx = 1: 5
            # set up
            N = 300
            G1, G2 = randn(N, N), randn(N, 2N);
            X1, X2 = (G1 + G1') / sqrt(2*N), (G2 * G2') / (2*N)
            X1 = free_whiten([X1]; mat = "her")[1][1]
            X2 = free_whiten([X2]; mat = "her")[1][1]
            X = [X1, X2];
            A = randn(2,2)
            Z = A*X

            # apply freecf
            Aest = freecf(Z; mat = mat, obj = obj, opt = "sphe")[1]
            @test err_pd(pinv(Aest)*A) < 5e-2
        end
    end

    # test for the rectangular matrices
    if mat == "rec"
        for idx = 1: 5
            # set up
            N, M = 300, 500
            X1 = randn(N, M) / sqrt(M)
            X1 = free_whiten([X1], mat = "rec")[1][1]
            U, V = Matrix(qr(randn(N,N)).Q), Matrix(qr(randn(M,M)).Q)
            D = [Diagonal((collect(range(1,stop = 0,length = N)) .- 1).^4) zeros(N, M - N)]
            X2 = U*D*V'
            X2 = free_whiten([X2], mat = "rec")[1][1]
            X = [X1, X2];
            A = randn(2,2)
            Z = A*X

            # apply freecf
            Aest = freecf(Z; mat = mat, obj = obj, opt = "sphe")[1]
            @test err_pd(pinv(Aest)*A) < 5e-2
        end
    end
end


text:
# Test task copying

using  Turing
import Turing.Traces: produce, consume

# Test case 1: stack allocated objects are deep copied.
function f_ct()
  t = 0;
  while true
    produce(t)
    t = 1 + t
  end
end

t = Task(f_ct)

@test consume(t) == 0
@test consume(t) == 1
a = copy(t);
@test consume(a) == 2
@test consume(a) == 3

# Test case 2: heap allocated objects are shallowly copied.

function f_ct2()
  t = [0 1 2];
  while true
    produce(t[1])
    t[1] = 1 + t[1]
  end
end

t = Task(f_ct2)

@test consume(t) == 0
@test consume(t) == 1
a = copy(t);
@test consume(a) == 2
@test consume(a) == 3
@test consume(t) == 4
@test consume(t) == 5
@test consume(a) == 6
@test consume(a) == 7


text:
using GCSPET
using Test

fname = instancepath("GCSPET_Guo", "20_2_0.4_1.dat")

inst = GCSPET.read(fname, Instance)

GCSPET.draw(inst, 1000, 800, "test.png")
GCSPET.draw(Instance, fname, 1000, 800, "test2.png")
"""
    randsolution(inst)

Construct a random (infeasible!) solution for testing.
"""
function randsolution(inst)
    J = jobs(inst)
    Q = cranes(inst)
    trajs = [GCSPET.Trajectory(id(q), [starting_position(q)], [0], [-1]) for q in Q]
    for j in J
        i = rand(eachindex(trajs))
        traj = trajs[i]
        q = Q[i]
        t_start = last(trajs[i].T) + abs(last(trajs[i].X)-loc(j))*speed(q)

        push!(traj.X, loc(j))
        push!(traj.T, t_start)
        if t_arrival(j) > t_start
            push!(traj.J, -1)
            push!(traj.T, t_arrival(j))
            push!(traj.X, loc(j))
            push!(traj.J, id(j))
        else
            push!(traj.J, id(j))
        end
        push!(traj.X, loc(j))
        push!(traj.T, last(traj.T) + t_processing(j))
        push!(traj.J, id(j))
    end
    return GCSPET.Solution(name(inst), jobs(inst), cranes(inst), trajs)
end

s = randsolution(inst)

@testset "GCSPET.jl" begin
    @test typeof(precedence_graph(jobs(inst))) <: SimpleDiGraph
    @test typeof(precedence_matrix(jobs(inst))) <: Matrix
    @test typeof(assignment_matrix(jobs(inst), cranes(inst))) <: Matrix
    @test GCSPET.makespan(s) > 0
end


text:
using SparseArrays, StatPlots
include("util.jl")

mutable struct Point
    position::Tuple{Int64,Int64}
    velocity::Tuple{Int64,Int64}
end

function make_sky()::Array{Point}
    points::Array{Point} =  []
    for line in getlines("inputs/10_test.txt")
        m = match(r"position=<\s*(-?\d+),\s*(-?\d+)>\s*velocity=<\s*(-?\d+),\s*(-?\d+)>", line)
        p = Point((parse(Int64,m[1]),parse(Int64,m[2])),(parse(Int64,m[3]),parse(Int64,m[4])))
        push!(points, p)
    end
    return points
end

function make_all_positive(a::Array{Tuple{Int64,Int64}})::Array{Tuple{Int64,Int64}}
    xmin = Inf
    ymin = Inf
    for t in a
        if t[1] < xmin
            xmin = t[1]
        end
        if t[2] < ymin
            ymin = t[2]
        end
    end
    newpoints::Array{Tuple{Int64,Int64}} = []
    for t in a
        push!(newpoints, ((t[1]-xmin)+1, (t[2]-ymin)+1))
    end
    return newpoints
end

function fast_forward_point(p::Point,seconds::Int64)::Tuple{Int64,Int64}
    x = p.position[1] + seconds*p.velocity[1]
    y = p.position[2] + seconds*p.velocity[2]
    return (x,y)
end

function make_grid(a::Array{Tuple{Int64,Int64}})::Array{Int64,2}
    xmax = 0
    ymax = 0
    for t in a
        if t[1] > xmax
            xmax = t[1]
        end
        if t[2] > ymax
            ymax = t[2]
        end
    end
    grid = zeros(ymax,xmax)
    for t in a
        grid[1+ymax-t[2],t[1]] = 1
    end
    return grid
end

function get_sky_at(init::Array{Point}, seconds::Int64)::Array{Tuple{Int64,Int64}}
    newpoints = fast_forward_point.(init, seconds)
    newview = make_all_positive(newpoints)
    return newview
end

function get_size(a::Array{Tuple{Int64,Int64}})::Int64
    xmax = 0
    ymax = 0
    for p in a
        if p[1] > xmax
            xmax = p[1]
        end
        if p[2] > ymax
            ymax = p[2]
        end
    end
    return xmax*ymax
end

function get_message(init::Array{Point})::Array{Int64,2}
    skies = []
    for n in 0:100000
        if n % 1000 == 0
            println(n)
        end
        s = get_sky_at(init, n)
        area = get_size(s)
        push!(skies, (area,s))
    end
    smallest_sky = sort(skies, by=x->x[1])[1]
    return make_grid(smallest_sky[2])
end

sky = make_sky()
message = get_message(sky)
heatmap(message)


text:
module TimeCorrelationFunctions

using NQCDynamics: AbstractSimulation, Estimators
using NQCDynamics.NonadiabaticDistributions: Diabatic, Adiabatic
using NQCDynamics.NQCModels: nstates

export PopulationCorrelationFunction

"""
    TimeCorrelationFunction

Abstract type for defining time correlation functions
"""
abstract type TimeCorrelationFunction end

function (correlation::TimeCorrelationFunction)(sol, _)
    sim = correlation.sim
    template = correlation_template(correlation)
    out = [similar(template) for _ in 1:length(sol.u)]

    normalisation = evaluate_normalisation(sim, correlation)
    initial_value = evaluate_initial_value(sim, correlation, first(sol.u))

    for i in 1:length(sol.u)
        final_value = evaluate_final_value(sim, correlation, sol.u[i])
        correlate!(out[i], normalisation, initial_value, final_value, correlation)
    end

    return (out, false)
end

evaluate_normalisation(sim::AbstractSimulation, correlation::TimeCorrelationFunction) = 1
evaluate_initial_value(sim::AbstractSimulation, correlation::TimeCorrelationFunction, u) =
    evaluate_final_value(sim, correlation, u)

function correlate!(out, normalisation, initial_value, final_value, ::TimeCorrelationFunction)
    @. out = normalisation * initial_value * final_value
end

function correlation_template(::TimeCorrelationFunction)
    error("Implement this for your correlation function type.")
end

"""
    PopulationCorrelationFunction{T,S<:AbstractSimulation} <: TimeCorrelationFunction

Output type for computing the population correlation function.
The `statetype` determines the population type (diabatic or adiabatic).
`sim` must also be provided to access the parameters to compute the population.
"""
struct PopulationCorrelationFunction{T,S<:AbstractSimulation} <: TimeCorrelationFunction
    sim::S
    statetype::T
end

function correlation_template(correlation::PopulationCorrelationFunction)
    zeros(nstates(correlation.sim), nstates(correlation.sim))
end

function evaluate_final_value(sim::AbstractSimulation, ::PopulationCorrelationFunction{Diabatic}, u)
    Estimators.diabatic_population(sim, u)
end

function evaluate_initial_value(sim::AbstractSimulation, ::PopulationCorrelationFunction{Diabatic}, u)
    Estimators.initial_diabatic_population(sim, u)
end

function evaluate_final_value(sim::AbstractSimulation, ::PopulationCorrelationFunction{Adiabatic}, u)
    Estimators.adiabatic_population(sim, u)
end

function evaluate_initial_value(sim::AbstractSimulation, ::PopulationCorrelationFunction{Adiabatic}, u)
    Estimators.initial_adiabatic_population(sim, u)
end

function correlate!(out, normalisation::Real, initial_value::Vector, final_value::Vector, ::PopulationCorrelationFunction)
    out .= normalisation .* initial_value * final_value'
end

end # module


text:
struct SRDirectCache{Tc,T,G,E,TT}
    S_c::Tc
    S::T
    F::G
    ∇out::E

    time_evolving::TT
end

function _sr_direct_cache(algo::SR, prob, net, par_cache)
    T = eltype(trainable_first(net))
    g = grad_cache(T, net)
    gv = vec_data(g)

    S = [similar(gv, T, length(gv), length(gv)) for gv in vec_data(g)]
    if T <: Real
        Sc = [similar(gv, Complex{T}, length(gv), length(gv)) for gv in vec_data(g)]
    else # is real
        Sc = S
    end

    return SRDirectCache(Sc, S, g, grad_cache(T, net),
                         algo.time_evo)
end

function setup_algorithm!(g::SRDirectCache, ∇C, Ô, par_cache)
    α = g.time_evolving isa Val{true} ? true * im : true

    O = Ô
    for (Sc, S, F) in zip(g.S_c, g.S, vec_data(g.F))
        T = eltype(S)
        N = size(Ô, 2)

        mul!(Sc, O, O')

        if T <: Real
            S .= real.(Sc ./ N)
            F .= real.(α .* ∇C)
        else
            # We take the conjugate because O is actually Oconj matrix respect
            # to the standard SR implementation..
            S .= conj.(Sc) ./ N
            F .= α .* ∇C
        end

        workers_sum!(S, par_cache)
    end
end

function precondition!(data::SRDirectCache, params::SR, iter_n)
    ϵ = params.sr_diag_shift

    success = true

    for (Δw, S, F) in zip(vec_data(data.∇out), data.S, vec_data(data.F))
        # new matrix
        if params.precondition_type == sr_none
            Sprecond = S
        elseif params.precondition_type == sr_shift
            #Sprecond = S + convert(eltype(S), ϵ)*I
            @inbounds @simd for i = 1:size(S, 1)
                S[i, i] += convert(eltype(S), ϵ)
            end
            Sprecond = S
        elseif params.precondition_type == sr_multiplicative
            λ0 = params.λ0
            b = params.b
            λmin = params.λmin
            λ = convert(eltype(S), max(λ0 * b^iter_n, λmin))
            Sprecond = S + λ * Diagonal(diag(S))
        end

        @assert is_iterative(params) == false
        if params.algorithm == sr_diag
            Sinv = pinv(Sprecond)
            mul!(Δw, Sinv, F)
        elseif params.algorithm == sr_cholesky
            C = cholesky!(Hermitian(S), check = true)
            copyto!(Δw, F)
            ldiv!(C, Δw)
        elseif params.algorithm == sr_div
            copyto!(Δw, F)
            ldiv!(Hermitian(S), Δw)
        else
            throw("Alg not known")
        end
    end
    return data.∇out
end
###########


text:
module ConsistencyResampling

using Bootstrap
using Bootstrap: BootstrapSampling
using StatsBase

using Random

import Bootstrap: bootstrap

export bootstrap, ConsistentSampling

include("bootstrap.jl")
include("utils.jl")

end # module


text:
module FBDCompare

using FBD

export ER, Kronecker, DCHSBM, HyperPA

include("er.jl")

module Kronecker
    include("hyperkron.jl")
end

module DCHSBM
    include("DCHSBM/utils.jl")
    include("DCHSBM/omega.jl")
    include("DCHSBM/HSBM.jl")
end

include("hyper_pa.jl")

include("evaluation/test_functions.jl")
include("evaluation/statistics.jl")

include("figures.jl")

end