text:
module Tar

using ArgTools

const true_predicate = _ -> true

# 2 MiB to take advantage of THP if enabled
const DEFAULT_BUFFER_SIZE = 2 * 1024 * 1024

# TODO: add some version of this method to Base
function Base.skip(io::Union{Base.Process, Base.ProcessChain}, n::Integer)
    n < 0 && throw(ArgumentError("cannot skip backwards when reading from a process"))
    isempty(skip_buffer) && resize!(skip_buffer, DEFAULT_BUFFER_SIZE)
    while n > 0
        n -= readbytes!(io, skip_buffer, min(n, length(skip_buffer)))
    end
    return io
end
const skip_buffer = UInt8[]

function can_symlink(dir::AbstractString)
    # guaranteed to be an empty directory
    link_path = joinpath(dir, "link")
    return try
        symlink("target", link_path)
        true
    catch err
        err isa Base.IOError || rethrow()
        false
    finally
        rm(link_path, force=true)
    end
end

include("header.jl")
include("create.jl")
include("extract.jl")

## official API: create, list, extract, rewrite, tree_hash

"""
    create([ predicate, ] dir, [ tarball ]; [ skeleton ]) -> tarball

        predicate :: String --> Bool
        dir       :: AbstractString
        tarball   :: Union{AbstractString, AbstractCmd, IO}
        skeleton  :: Union{AbstractString, AbstractCmd, IO}

Create a tar archive ("tarball") of the directory `dir`. The resulting archive
is written to the path `tarball` or if no path is specified, a temporary path is
created and returned by the function call. If `tarball` is an IO object then the
tarball content is written to that handle instead (the handle is left open).

If a `predicate` function is passed, it is called on each system path that is
encountered while recursively searching `dir` and `path` is only included in the
tarball if `predicate(path)` is true. If `predicate(path)` returns false for a
directory, then the directory is excluded entirely: nothing under that directory
will be included in the archive.

If the `skeleton` keyword is passed then the file or IO handle given is used as
a "skeleton" to generate the tarball. You create a skeleton file by passing the
`skeleton` keyword to the `extract` command. If `create` is called with that
skeleton file and the extracted files haven't changed, an identical tarball is
recreated. The `skeleton` and `predicate` arguments cannot be used together.
"""
function create(
    predicate::Function,
    dir::AbstractString,
    tarball::Union{ArgWrite, Nothing} = nothing;
    skeleton::Union{ArgRead, Nothing} = nothing,
)
    check_create_dir(dir)
    if skeleton === nothing
        arg_write(tarball) do tar
            create_tarball(predicate, tar, dir)
        end
    else
        predicate === true_predicate ||
            error("create: predicate and skeleton cannot be used together")
        check_create_skeleton(skeleton)
        arg_read(skeleton) do skeleton
            arg_write(tarball) do tar
                recreate_tarball(tar, dir, skeleton)
            end
        end
    end
end

function create(
    dir::AbstractString,
    tarball::Union{ArgWrite, Nothing} = nothing;
    skeleton::Union{ArgRead, Nothing} = nothing,
)
    create(true_predicate, dir, tarball, skeleton=skeleton)
end

"""
    list(tarball; [ strict = true ]) -> Vector{Header}
    list(callback, tarball; [ strict = true ])

        callback  :: Header --> Bool
        tarball   :: Union{AbstractString, AbstractCmd, IO}
        strict    :: Bool

List the contents of a tar archive ("tarball") located at the path `tarball`. If
`tarball` is an IO handle, read the tar contents from that stream. Returns a
vector of `Header` structs. See [`Header`](@ref) for details. If a `callback` is
provided then instead of returning a vector of headers, the callback is called
on each `Header`. This can be useful if the number of items in the tarball is
large or if you want examine items prior to an error in the tarball.

By default `list` will error if it encounters any tarball contents which the
`extract` function would refuse to extract. With `strict=false` it will skip
these checks and list all the the contents of the tar file whether `extract`
would extract them or not. Beware that malicious tarballs can do all sorts of
crafty and unexpected things to try to trick you into doing something bad.

If the `tarball` argument is a skeleton file (see `extract` and `create`) then
`list` will detect that from the file header and appropriately list or iterate
the headers of the skeleton file.
"""
function list(
    callback::Function,
    tarball::ArgRead;
    raw::Bool = false,
    strict::Bool = !raw,
)
    raw && strict &&
        error("`raw=true` and `strict=true` options are incompatible")
    arg_read(tarball) do tar
        iterate_headers(callback, tar, raw=raw, strict=strict)
    end
end

function list(
    tarball::ArgRead;
    raw::Bool = false,
    strict::Bool = !raw,
)
    headers = Header[]
    list(tarball, raw=raw, strict=strict) do hdr
        push!(headers, hdr)
    end
    return headers
end

"""
    extract([ predicate, ] tarball, [ dir ];
            [ skeleton, ] [ copy_symlinks ]) -> dir

        predicate     :: Header --> Bool
        tarball       :: Union{AbstractString, AbstractCmd, IO}
        dir           :: AbstractString
        skeleton      :: Union{AbstractString, AbstractCmd, IO}
        copy_symlinks :: Bool

Extract a tar archive ("tarball") located at the path `tarball` into the
directory `dir`. If `tarball` is an IO object instead of a path, then the
archive contents will be read from that IO stream. The archive is extracted to
`dir` which must either be an existing empty directory or a non-existent path
which can be created as a new directory. If `dir` is not specified, the archive
is extracted into a temporary directory which is returned by `extract`.

If a `predicate` function is passed, it is called on each `Header` object that
is encountered while extracting `tarball` and the entry is only extracted if the
`predicate(hdr)` is true. This can be used to selectively extract only parts of
an archive, to skip entries that cause `extract` to throw an error, or to record
what is extracted during the extraction process.

If the `skeleton` keyword is passed then a "skeleton" of the extracted tarball
is written to the file or IO handle given. This skeleton file can be used to
recreate an identical tarball by passing the `skeleton` keyword to the `create`
function. The `skeleton` and `predicate` arguments cannot be used together.

If `copy_symlinks` is `true` then instead of extracting symbolic links as such,
they will be extracted as copies of what they link to if they are internal to
the tarball and if it is possible to do so. Non-internal symlinks, such as a
link to `/etc/passwd` will not be copied. Symlinks which are in any way cyclic
will also not be copied and will instead be skipped. By default, `extract` will
detect whether symlinks can be created in `dir` or not and will automatically
copy symlinks if they cannot be created.
"""
function extract(
    predicate::Function,
    tarball::ArgRead,
    dir::Union{AbstractString, Nothing} = nothing;
    skeleton::Union{ArgWrite, Nothing} = nothing,
    copy_symlinks::Union{Bool, Nothing} = nothing,
)
    predicate === true_predicate || skeleton === nothing ||
        error("extract: predicate and skeleton cannot be used together")
    skeleton = something(skeleton, devnull)
    check_extract_tarball(tarball)
    check_extract_dir(dir)
    arg_read(tarball) do tar
        arg_mkdir(dir) do dir
            if copy_symlinks === nothing
                copy_symlinks = !can_symlink(dir)
            end
            arg_write(skeleton) do skeleton
                extract_tarball(
                    predicate, tar, dir,
                    skeleton = skeleton,
                    copy_symlinks = copy_symlinks,
                )
            end
        end
    end
end

function extract(
    tarball::ArgRead,
    dir::Union{AbstractString, Nothing} = nothing;
    skeleton::Union{ArgWrite, Nothing} = nothing,
    copy_symlinks::Union{Bool, Nothing} = nothing,
)
    extract(
        true_predicate, tarball, dir,
        skeleton = skeleton,
        copy_symlinks = copy_symlinks,
    )
end

"""
    rewrite([ predicate, ], old_tarball, [ new_tarball ]) -> new_tarball

        predicate   :: Header --> Bool
        old_tarball :: Union{AbstractString, AbtractCmd, IO}
        new_tarball :: Union{AbstractString, AbtractCmd, IO}

Rewrite `old_tarball` to the standard format that `create` generates, while also
checking that it doesn't contain anything that would cause `extract` to raise an
error. This is functionally equivalent to doing

    Tar.create(Tar.extract(predicate, old_tarball), new_tarball)

However, it never extracts anything to disk and instead uses the `seek` function
to navigate the old tarball's data. If no `new_tarball` argument is passed, the
new tarball is written to a temporary file whose path is returned.

If a `predicate` function is passed, it is called on each `Header` object that
is encountered while extracting `old_tarball` and the entry is skipped unless
`predicate(hdr)` is true. This can be used to selectively rewrite only parts of
an archive, to skip entries that would cause `extract` to throw an error, or to
record what content is encountered during the rewrite process.
"""
function rewrite(
    predicate::Function,
    old_tarball::ArgRead,
    new_tarball::Union{ArgWrite, Nothing} = nothing,
)
    old_tarball = check_rewrite_old_tarball(old_tarball)
    arg_read(old_tarball) do old_tar
        arg_write(new_tarball) do new_tar
            rewrite_tarball(predicate, old_tar, new_tar)
        end
    end
end

function rewrite(
    old_tarball::ArgRead,
    new_tarball::Union{ArgWrite, Nothing} = nothing,
)
    rewrite(true_predicate, old_tarball, new_tarball)
end

"""
    tree_hash([ predicate, ] tarball;
              [ algorithm = "git-sha1", ]
              [ skip_empty = false ]) -> hash::String

        predicate  :: Header --> Bool
        tarball    :: Union{AbstractString, AbstractCmd, IO}
        algorithm  :: AbstractString
        skip_empty :: Bool

Compute a tree hash value for the file tree that the tarball contains. By
default, this uses git's tree hashing algorigthm with the SHA1 secure hash
function (like current versions of git). This means that for any tarball whose
file tree git can represent—i.e. one with only files, symlinks and non-empty
directories—the hash value computed by this function will be the same as the
hash value git would compute for that file tree. Note that tarballs can
represent file trees with empty directories, which git cannot store, and this
function can generate hashes for those, which will, by default (see `skip_empty`
below for how to change this behavior), differ from the hash of a tarball which
omits those empty directories. In short, the hash function agrees with git on
all trees which git can represent, but extends (in a consistent way) the domain
of hashable trees to other trees which git cannot represent.

If a `predicate` function is passed, it is called on each `Header` object that
is encountered while processing `tarball` and an entry is only hashed if
`predicate(hdr)` is true. This can be used to selectively hash only parts of an
archive, to skip entries that cause `extract` to throw an error, or to record
what is extracted during the hashing process.

Currently supported values for `algorithm` are `git-sha1` (the default) and
`git-sha256`, which uses the same basic algorithm as `git-sha1` but replaces the
SHA1 hash function with SHA2-256, the hash function that git will transition to
using in the future (due to known attacks on SHA1). Support for other file tree
hashing algorithms may be added in the future.

The `skip_empty` option controls whether directories in the tarball which
recursiviely contain no files or symlinks are included in the hash or ignored.
In general, if you are hashing the content of a tarball or a file tree, you care
about all directories, not just non-empty ones, so including these in the
computed hash is the default. So why does this function even provide the option
to skip empty directories? Because git refuses to store empty directories and
will ignore them if you try to add them to a repo. So if you compute a reference
tree hash by by adding files to a git repo and then asking git for the tree
hash, the hash value that you get will match the hash value computed by
`tree_hash` with `skip_empty=true`. In other words, this option allows
`tree_hash` to emulate how git would hash a tree with empty directories. If you
are hashing trees that may contain empty directories (i.e. do not come from a
git repo), however, it is recommended that you hash them using a tool (such as
this one) that does not ignore empty directories.
"""
function tree_hash(
    predicate::Function,
    tarball::ArgRead;
    algorithm::AbstractString = "git-sha1",
    skip_empty::Bool = false,
)
    HashType =
        algorithm == "git-sha1"   ? SHA.SHA1_CTX :
        algorithm == "git-sha256" ? SHA.SHA256_CTX :
            error("invalid tree hashing algorithm: $algorithm")

    check_tree_hash_tarball(tarball)
    arg_read(tarball) do tar
        git_tree_hash(predicate, tar, HashType, skip_empty)
    end
end

function tree_hash(
    tarball::ArgRead;
    algorithm::AbstractString = "git-sha1",
    skip_empty::Bool = false,
)
    tree_hash(
        true_predicate,
        tarball,
        algorithm = algorithm,
        skip_empty = skip_empty,
    )
end

## error checking utility functions

check_create_dir(dir::AbstractString) =
    isdir(dir) || error("""
    `dir` not a directory: $dir
    USAGE: create([predicate,] dir, [tarball])
    """)

check_create_skeleton(skeleton::AbstractString) =
    isfile(skeleton) || error("""
    `skeleton` not a file: $skeleton
    USAGE: create([predicate,] dir, [tarball]; [skeleton = <file>])
    """)

check_create_skeleton(skeleton::ArgRead) = nothing

check_list_tarball(tarball::AbstractString) =
    isfile(tarball) || error("""
    `tarball` not a file: $tarball
    USAGE: list([predicate,] tarball)
    """)

check_extract_tarball(tarball::AbstractString) =
    isfile(tarball) || error("""
    `tarball` not a file: $tarball
    USAGE: extract([predicate,] tarball, [dir])
    """)

check_extract_tarball(tarball::ArgWrite) = nothing

function check_extract_dir(dir::AbstractString)
    st = stat(dir)
    ispath(st) && !isdir(st) &&
        error("""
        `dir` not a directory: $dir
        USAGE: extract([predicate,] tarball, [dir])
        """)
    isdir(st) && !isempty(readdir(dir)) &&
        error("""
        `dir` not empty: $dir
        USAGE: extract([predicate,] tarball, [dir])
        """)
end

check_extract_dir(dir::Nothing) = nothing

check_rewrite_old_tarball(old_tarball::AbstractString) =
    isfile(old_tarball) ? old_tarball : error("""
    `old_tarball` not a file: $old_tarball
    USAGE: rewrite([predicate,] old_tarball, [new_tarball])
    """)

check_rewrite_old_tarball(old_tarball::ArgRead) =
    applicable(seek, old_tarball, 0) ?
        old_tarball : IOBuffer(read(old_tarball))

check_tree_hash_tarball(tarball::AbstractString) =
    isfile(tarball) || error("""
    `tarball` not a file: $tarball
    USAGE: tree_hash([predicate,] tarball)
    """)

check_tree_hash_tarball(tarball::ArgRead) = nothing

end # module


text:
export load_checkpoints, checkpoint!, recent_model, best_model

struct Checkpoint
    model
    epoch
    loss_value
    loss_error
end

const MaybeCheckpoint = Union{Missing, Checkpoint}

mutable struct Checkpoints
    recent::MaybeCheckpoint
    top::Vector{Checkpoint}
    top_num::Integer
end

Checkpoints() = Checkpoints(missing, Vector{Checkpoint}(), 3)

Base.isless(c1::Checkpoint, c2::Checkpoint) = isless(c1.loss_value, c2.loss_value)

function _repeated_load(bson)
    i = 0
    e = nothing
    while i < 10
        try
            return BSON.load(bson)
        catch e
            @warn("Loading BSON failed. Trying again. Exception: $e")
            i += 1
        end
    end
    throw(e)
end

function load_checkpoints(bson)
    if !isfile(bson)
        return Checkpoints()
    else
        content = _repeated_load(bson)
        return haskey(content, :checkpoints) ? content[:checkpoints] : Checkpoints()
    end
end

function checkpoint!(bson, model, epoch, loss_value, loss_error)
    checkpoints = load_checkpoints(bson)

    # Construct current checkpoint.
    current_checkpoint = Checkpoint(cpu(model), epoch, loss_value, loss_error)

    # Always update most recent model.
    checkpoints.recent = current_checkpoint

    # Update top models.
    extended_top = sort(vcat(checkpoints.top, current_checkpoint))
    checkpoints.top = extended_top[1:min(length(extended_top), checkpoints.top_num)]

    # Write changes.
    BSON.bson(bson, checkpoints=checkpoints)
end

recent_model(checkpoints::Checkpoints) = checkpoints.recent.model
recent_model(bson::String) = recent_model(load_checkpoints(bson))

best_model(checkpoints::Checkpoints, position::Integer=1) = checkpoints.top[position].model
best_model(bson::String, position::Integer=1) = best_model(load_checkpoints(bson), position)

ext: jl


text:
abstract type AbstractIndices{T,N,I<:Tuple,L} <: AbstractArray{T,N} end

_axeltype(inds::Tuple) = (eltype(first(inds)), _axeltype(Base.tail(inds))...)
_axeltype(inds::Tuple{R1,R2}) where {R1,R2} = (eltype(first(inds)), eltype(last(inds)))
_axeltype(inds::Tuple{R}) where R<:AbstractRange = (eltype(first(inds)),)

_axlength(inds::Tuple) = length(first(inds)) * _axlength(Base.tail(inds))
_axlength(inds::Tuple{R}) where R = length(first(inds))

"""
    StaticCartesianIndices{T,N,I,L}

Indexing returns a tuple of length N and type T

# Examples
```jldoctest

julia> x = SOne:SInt(10)
SInteger(1):SInteger(10)

julia> y = SOne:SInt(5)
SInteger(1):SInteger(5)

julia> a = CartesianSIndices((x,y));

```
"""
abstract type StaticCartesianIndices{T<:Tuple,N,I,L} <: AbstractIndices{T,N,I,L} end

struct CartesianSIndices{T,N,I,L} <: StaticCartesianIndices{T,N,I,L} end

CartesianSIndices(inds::SRange...) = CartesianSIndices(inds)
CartesianSIndices(inds::Tuple{Vararg{<:SRange,N}}) where N =
    CartesianSIndices{Tuple{_axeltype(inds)...},N,typeof(inds),_axlength(inds)}()

mutable struct CartesianMIndices{T,N,I} <: StaticCartesianIndices{T,N,I,Dynamic}
    indices::I
end

CartesianMIndices(inds::Vararg{<:AbstractRange,N}) where N = CartesianMIndices(inds)

function CartesianMIndices(inds::Tuple{Vararg{<:AbstractRange,N}}) where N
    CartesianMIndices{Tuple{_axeltype(inds)...},N,typeof(inds)}(inds)
end

"""
    StaticLinearIndices

Indexing returns a single value.
"""
abstract type StaticLinearIndices{T,N,I,L} <: AbstractIndices{T,N,I,L} end

"""
    LinearSIndices
"""
struct LinearSIndices{T,N,I,L} <: StaticLinearIndices{T,N,I,L} end

LinearSIndices(inds::Vararg{<:SRange,N}) where N = LinearSIndices(inds)

function LinearSIndices(inds::Tuple{Vararg{<:SRange,N}}) where N
    LinearSIndices{typejoin(_axeltype(inds)...),N,typeof(inds),_axlength(inds)}()
end

"""
    LinearMIndices
"""
mutable struct LinearMIndices{T,N,I} <: StaticLinearIndices{T,N,I,Dynamic}
    indices::I
end

LinearMIndices(inds::Vararg{<:AbstractRange,N}) where N = LinearMIndices(inds)
LinearMIndices(inds::Tuple{Vararg{<:AbstractRange,N}}) where N =
    LinearMIndices{typejoin(_axeltype(inds)...),N,typeof(inds)}(inds)

const SIndices{T,N,I,L} = Union{LinearSIndices{T,N,I,L},CartesianSIndices{T,N,I,L}}

const MIndices{T,N,I} = Union{LinearMIndices{T,N,I},CartesianMIndices{T,N,I}}

@pure Base.axes(::SIndices{T,N,I,L}) where {T,N,I,L} = ntuple(i->fieldtype(I, i)(), Val(N))::I
@pure Base.axes(::Type{<:SIndices{T,N,I,L}}) where {T,N,I,L} = ntuple(i->fieldtype(I, i)(), Val(N))::I

Base.axes(inds::MIndices) = inds.indices

@inline Base.axes(inds::AbstractIndices, i::Int) = axes(inds)[i]
@inline Base.axes(inds::Type{<:SIndices}, i::Int) = axes(inds)[i]
@inline Base.axes(inds::AbstractIndices, i::SInt{I}) where I = axes(inds)[I]

@inline Base.size(inds::AbstractIndices, i::Int) = length(axes(inds, i))::Int

@inline Base.size(inds::Type{<:SIndices}, i::Int) = length(axes(inds, i))

@inline Base.size(inds::AbstractIndices{T,N}) where {T,N} = (length.(axes(inds)))

@pure Base.length(inds::SIndices{T,N,I,L}) where {T,N,I,L} = L

@inline Base.length(inds::MIndices) = prod(size(inds))

function getindex(inds::StaticLinearIndices{T,N,I,L}, i::Int) where {T,N,I,L}
    Base.@_inline_meta
    @boundscheck if firstindex(inds) > i || i > lastindex(inds)
        throw(BoundsError(inds, i))
    end
    @inbounds T(getindex(firstindex(inds):lastindex(inds), i))  # just incase someone defines crazy first/last indexing
end

function getindex(inds::StaticLinearIndices{T,N,I,L}, i::AbstractRange) where {T,N,I,L}
    Base.@_inline_meta
    @boundscheck if firstindex(inds) > first(i) || last(i) > lastindex(inds)
        throw(BoundsError(inds, i))
    end
    @inbounds getindex(firstindex(inds):lastindex(inds), i)
end

#function checkbounds(::Type{Bool}, inds::AbstractIndices, )
#end

# TODO
#@generated function to_subindices(::IndexCartesian, ::Size{S}, subinds::SAxesUnion{N,Ax}, i::Int) where {S,N,Ax}
function getindex(inds::StaticCartesianIndices, i::Integer)
    Base.@_inline_meta
    @boundscheck checkbounds(inds, i)
    @inbounds unsafe_getindex(inds, i)
end

#=
function unsafe_getindex(ci::CartesianSIndices{T,N,I,L}, i::BaseInteger)::T where {T,N,I,L}
    out = []
    ind = i - 1
    for ax in axes(ci)
        indnext = div(ind, length(ax))
        push!(out, first(ax) + (ind - length(ax) * indnext + 1 - firstindex(ax)) * step(ax))
        ind = indnext
    end
    return Tuple(out)
end

@generated function unsafe_getindex(ci::CartesianSIndices{T,N,I,L}, i::BaseInteger) where {T,N,I,L}
    out = Expr[]
    ind = :(i - 1)
    for D in 1:N
        ax = :(axes($ci, $D))
        indnext = :(div($ind, _length($ax)))
        if D == 1
            push!(out, :(first($ax) + ($ind - _length($ax) * $indnext + 1 - firstindex($ax)) * step($ax)))
        else
            push!(out, :(first($ax) + ($ind - _length($ax) *  $indnext + 1 - firstindex($ax)) * step($ax)))
        end
        ind = indnext
    end
    return quote
        ($(out...),)
    end
end

_length(::UnitSRange{T,B,E,L}) where {T,B,E,L} = values(L)

=#
@generated function unsafe_getindex(inds::CartesianSIndices{T,N,I,L}, i::SInteger) where {T,N,I,L}
    out = []
    ind = i - SOne
    for D in 1:N
        ax = axes(inds, D)
        indnext = div(ind, length(ax))
        if D == 1
            push!(out, first(ax) + (ind - length(ax) * indnext + SOne - firstindex(ax)) * step(ax))
        else
            push!(out, first(ax) + (ind - length(ax) *  indnext + SOne - firstindex(ax)) * step(ax))
        end
        ind = indnext
    end
    return quote
        ($(out...),)
    end
end

#@generated function unsafe_getindex(inds::CartesianSIndices{T,N,I}, i::Integer) where {T,N,I}
#    :(_ind2subs($(axes(inds)), i - 1))
#end

@inline function unsafe_getindex(inds::StaticCartesianIndices{T,N,I}, i::Integer) where {T,N,I}
    Base.@_inline_meta
    _ind2sub(axes(inds), i)::T
end

_ind2sub(::Tuple{}, ind::Integer) = (Base.@_inline_meta; ind == 1 ? () : throw(BoundsError()))
_ind2sub(dims::Base.DimsInteger, ind::Integer) = (Base.@_inline_meta; _ind2sub_recurse(dims, ind-1))
_ind2sub(inds::Tuple{Vararg{<:AbstractRange,N}}, ind::Integer) where N =
    (Base.@_inline_meta; _ind2sub_recurse(inds, ind-1))
_ind2sub(inds::Tuple{AbstractArray}, ind::Integer) =
    throw(ArgumentError("Linear indexing is not defined for one-dimensional arrays"))
#_ind2sub(inds::Tuple{OneTo}, ind::Integer) = (ind,)

_ind2sub_recurse(::Tuple{}, ind) = (ind+1,)
function _ind2sub_recurse(indslast::NTuple{1}, ind)
    Base.@_inline_meta
    (_lookup(ind, indslast[1]),)
end
function _ind2sub_recurse(inds, ind)
    Base.@_inline_meta
    r1 = inds[1]
    indnext, f, l = _div(ind, r1)
    (ind-l*indnext+f, _ind2sub_recurse(Base.tail(inds), indnext)...)
end

_lookup(ind, d::Integer) = ind+1
_lookup(ind, r::AbstractUnitRange) = ind+first(r)
_lookup(ind, r::StaticUnitRange) = ind+first(r)
_div(ind, d::Integer) = div(ind, d), 1, d
_div(ind, r::AbstractUnitRange) = (d = length(r); (div(ind, d), first(r), d))
_div(ind, r::StaticUnitRange) = (d = length(r); (div(ind, d), first(r), d))

# TODO
function getindex(inds::StaticLinearIndices{T,N}, i::Vararg{T,N}) where {T,N}
    Base.@_inline_meta
    @boundscheck checkbounds(inds, i...)
    @inbounds unsafe_getindex(inds, i)
end

function unsafe_getindex(li::LinearMIndices{T,N,I}, i::NTuple{N}) where {T,N,I}
    sz = 1
    out = zero(T)
    for D in 1:N
        ax = axes(li, D)
        if D == 1
            out = _inds2idx(ax, i[D])
        else
            out = out + _inds2idx(ax, sz, i[D])
        end
        sz *= size(li, D)
    end
    return out
end

_inds2idx(ax::UnitSRange{T,B}, i::Integer) where {T,B} = i
_inds2idx(ax::UnitSRange{T,B}, sz::Int, i::Integer) where {T,B} = sz * (values(B)::T + i - 2)

_inds2idx(ax::AbstractRange, i::Integer) = first(ax) + i - first(ax) * step(ax)
_inds2idx(ax::AbstractRange, sz::Int, i::Integer) = sz * (first(ax) + (i -firstindex(ax)) * step(ax) - 1)

@generated function unsafe_getindex(li::LinearSIndices{T,N,I,L}, i::NTuple{N}) where {T,N,I,L}
    sz = 1
    out = zero(T)
    for D in 1:N
        ax = :(axes(li, $D))
        if D == 1
            out = :(first($ax) + i[$D] - first($ax) * step($ax))
        else
            out = :($out + $sz * (first($ax) + (i[$D] - firstindex($ax)) * step($ax) - SOne))
        end
        sz *= size(li, D)
    end
    return quote
        $out
    end
end

# T<:Tuple for cartesian inds so easiest to index in that context
function getindex(inds::StaticCartesianIndices{T,N,I,L}, i...) where {T,N,I,L}
    Base.@_inline_meta
    @boundscheck checkbounds(inds, i...)
    @inbounds unsafe_getindex(inds, i)
end

function unsafe_getindex(inds::CartesianMIndices{T,N,I}, i::Tuple) where {T,N,I}
    Base.@_inline_meta
    map(getindex, axes(inds), i)::T
end

@inline function unsafe_getindex(inds::CartesianSIndices{T,N,I}, i::Tuple)::T where {T,N,I}
    map(getindex, axes(inds), i)::T
end


text:
using GLMakie
export plot3d, plot3d_space_bottlenecks!

function plot3d(l::RQCLayout, vs = vertices(l))
    vs = sort!(intersect(vs, vertices(l)))
    D = depth(l)
    scene = meshscatter([l.qubit_locs[v][1] for v in vs],
        [l.qubit_locs[v][2] for v in vs],
        zeros(length(vs)) .- 0.01 * rand(length(vs)),
        color = :blue,
        markersize = 0.15
    )
    for v in vs
        loc = l.qubit_locs[v]
        text!("$v", position = (loc[1], loc[2], -0.5))
        d1 = 0
        d2 = 1
        while d1 < D
            if d2 <= D && gates(l, v)[d2] !== Noise
                d2 += 1
            else
                lines!([loc[1], loc[1]], [loc[2], loc[2]], [0.5d1, 0.5d2])
                d1 = d2 + 1
                while d1 < D && gates(l, v)[d1+1] === Noise
                    d1 += 1
                end
                d2 = d1 + 1
                if d1 === D && gates(l, v)[d1] === Noise
                    lines!([loc[1], loc[1]], [loc[2], loc[2]], [0.5d1, 0.5d2])
                end
            end
        end
    end

    for v1 in vs, d = 1:D
        if gates(l, v1)[d] === FSim
            v2 = partner(l, v1, depth_to_cycle(d))
            v2 in vs || continue
            v1 > v2 && continue
            loc1 = l.qubit_locs[v1]
            loc2 = l.qubit_locs[v2]
            lines!([loc1[1], loc2[1]], [loc1[2], loc2[2]], [d*0.5, d*0.5]; color = :red)
        end
    end

    v_xs = Float64[]
    v_ys = Float64[]
    v_zs = Float64[]
    v_sizes = Float64[]
    v_colors = Symbol[]
    for d = 1:D
        for v in vs
            push!(v_xs, l.qubit_locs[v][1])
            push!(v_ys, l.qubit_locs[v][2])
            push!(v_zs, d * 0.5)
            if gates(l, v)[d] === FSim
                push!(v_sizes, 0.1)
                push!(v_colors, :red)
            elseif gates(l, v)[d] === Noise
                push!(v_sizes, 0.1)
                if (d == 1 || d == D) || gates(l, v)[d-1] !== Noise || gates(l, v)[d+1] !== Noise
                    push!(v_colors, :white)
                else
                    push!(v_colors, :transparent)
                end
            elseif gates(l, v)[d] === Single
                push!(v_sizes, 0.05)
                push!(v_colors, :black)
            else
                push!(v_sizes, 0.05)
                push!(v_colors, :transparent)
            end
        end
    end
    meshscatter!(v_xs, v_ys, v_zs; color = v_colors, markersize = v_sizes)
    meshscatter!([l.qubit_locs[v][1] for v in vs],
        [l.qubit_locs[v][2] for v in vs],
        ones(length(vs)) * (D+1) * 0.5,
        color = :blue,
        markersize = 0.15
    )
    return scene
end

function plot3d_space_bottlenecks!(layout, ec, ids_dic, ids_loc)
    max_ids = max_space_ids(ec, ids_dic)
    scene = current_figure()
    for i in max_ids
        loc = layout.qubit_locs[ids_loc[i][1]]
        scene = lines!([loc[1], loc[1]], [loc[2], loc[2]],
            [(ids_loc[i][2])/2, (ids_loc[i][2]+1)/2],
            color = :green,
            linewidth = 10
        )
    end
    return scene
end


text:
using CFITSIO
using Documenter

makedocs(;
    modules=[CFITSIO],
    authors="JuliaAstro",
    repo="https://github.com/JuliaAstro/CFITSIO.jl/blob/{commit}{path}#L{line}",
    sitename="CFITSIO.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://juliaastro.github.io/CFITSIO.jl",
        assets=String[],
    ),
    pages=[
        "Home" => "index.md",
    ],
)

deploydocs(;
    repo="github.com/JuliaAstro/CFITSIO.jl",
)


text:
module CuArrays

using CUDAapi, CUDAdrv, CUDAnative

using GPUArrays

export CuArray, CuVector, CuMatrix, CuVecOrMat, cu
export CUBLAS, CUSPARSE, CUSOLVER, CUFFT, CURAND, CUDNN, CUTENSOR

import LinearAlgebra

using Adapt

using Libdl

using Requires

## source code includes

include("bindeps.jl")

# core array functionality
include("memory.jl")
include("array.jl")
include("gpuarrays.jl")
include("subarray.jl")
include("utils.jl")

# integrations and specialized functionality
include("indexing.jl")
include("broadcast.jl")
include("mapreduce.jl")
include("accumulate.jl")
include("linalg.jl")
include("nnlib.jl")

# vendor libraries
include("blas/CUBLAS.jl")
include("sparse/CUSPARSE.jl")
include("solver/CUSOLVER.jl")
include("fft/CUFFT.jl")
include("rand/CURAND.jl")
include("dnn/CUDNN.jl")
include("tensor/CUTENSOR.jl")

include("deprecated.jl")

## initialization

const __initialized__ = Ref(false)
functional() = __initialized__[]

export has_cudnn, has_cutensor
has_cudnn() = Libdl.dlopen_e(CUDNN.libcudnn[]) !== C_NULL
has_cutensor() = Libdl.dlopen_e(CUTENSOR.libcutensor[]) !== C_NULL

function __init__()
    precompiling = ccall(:jl_generating_output, Cint, ()) != 0
    silent = parse(Bool, get(ENV, "JULIA_CUDA_SILENT", "false")) || precompiling
    verbose = parse(Bool, get(ENV, "JULIA_CUDA_VERBOSE", "false"))

    # if any dependent GPU package failed, expect it to have logged an error and bail out
    if !CUDAdrv.functional() || !CUDAnative.functional()
        verbose && @warn "CuArrays.jl did not initialize because CUDAdrv.jl or CUDAnative.jl failed to"
        return
    end

    try
        __init_bindeps__(silent=silent, verbose=verbose)

        # package integrations
        @require ForwardDiff="f6369f11-7733-5829-9624-2563aa707210" include("forwarddiff.jl")

        __init_memory__()

        __initialized__[] = true
    catch ex
        # don't actually fail to keep the package loadable
        if !silent
            if verbose
                @error "CuArrays.jl failed to initialize" exception=(ex, catch_backtrace())
            else
                @info "CuArrays.jl failed to initialize and will be unavailable (set JULIA_CUDA_SILENT or JULIA_CUDA_VERBOSE to silence or expand this message)"
            end
        end
    end
end

end # module


text:
"""
    `get_nano_pore_mod(NANO_MOD_FILE)`

    Function that stores the signal to noise ratio for each k-mer in a tabulated file with columns

        | Unmethylated k-mer | Signal-to-noise Ratio |

    # Examples
    ```julia-repl
    julia> CpelNano.get_nano_pore_mod(nano_mod_file)
    ```
"""
function get_nano_pore_mod(nano_mod_file::String)::Dict{String,NTuple{4,Float64}}

    # Read in model file
    nano_mod = readdlm(nano_mod_file)

    # Create dictionary
    nano_dic = Dict{String,NTuple{4,Float64}}()
    for i = 2:size(nano_mod)[1]
        kmer = nano_mod[i,1]
        occursin('C', kmer) || occursin('M', kmer) || continue
        if occursin('M', kmer)
            kmer_mod = replace(kmer, "M" => "C")
            # print_log("kmer: $(kmer) => kmer_mod: $(kmer_mod)")
            if kmer_mod in keys(nano_dic)
                μu = nano_dic[kmer_mod][1]
                σu = nano_dic[kmer_mod][2]
                nano_dic[kmer_mod] = (μu, σu, nano_mod[i,2], nano_mod[i,3])
            else
                nano_dic[kmer_mod] = (NaN, NaN, nano_mod[i,2], nano_mod[i,3])
            end
        else
            if kmer in keys(nano_dic)
                μm = nano_dic[kmer][3]
                σm = nano_dic[kmer][4]
                nano_dic[kmer] = (nano_mod[i,2], nano_mod[i,3], μm, σm)
            else
                nano_dic[kmer] = (nano_mod[i,2], nano_mod[i,3], NaN, NaN)
            end
        end
    end

    # Return dictionary
    return nano_dic

end
"""
    `pmap_split_nanopolish_file(NANO_FILE,N_FILES,HEADER,PREF,READ_FILE_IND,READ_NAME_VEC,IND_FILE)`

    Function writes one of the files in `split_nanopolish_file`.

    # Examples
    ```julia-repl
    julia> CpelNano.pmap_split_nanopolish_file(nano_file,n_files,header,pref,read_file_ind,read_name_vec,ind_file)
    ```
"""
function pmap_split_nanopolish_file(nano_file::String, n_files::Int64, header::String, pref::SubString{String}, read_file_ind::Vector{Int64}, read_nm_vec::Vector{String}, ind_file::Int64)::Nothing

    # Name of file
    nth_file = "$(pref).rep_$(ind_file).tsv"
    print_log("Writing file $(ind_file) out of $(n_files)...")
    # Write individual file
    open(nth_file, "w") do io
        # Add header
        write(io, "$(header)\n")
        # Add reads
        open(nano_file) do f
            for line in enumerate(eachline(f))
                # Skip header
                line[1] == 1 && continue
                # Get ID
                line_id = split(line[2], "\t")[5]
                # Get index of read in permuted vector
                read_ind_vec = findfirst(x -> x == line_id, read_nm_vec)
                # If read belongs to this file print
                if read_file_ind[read_ind_vec] == ind_file
                    write(io, "$(line[2])\n")
                end
            end
        end
    end

    # Return nothing
    return nothing

end
"""
    `split_nanopolish_file(NANO_FILE,N_FILES)`

    Function that splits nanopolish file into `N_FILES` files.

    # Examples
    ```julia-repl
    julia> CpelNano.split_nanopolish_file(nano_file,n_files)
    ```
"""
function split_nanopolish_file(nano_file::String, n_files::Int64)::Nothing

    # Prefix output file
    pref = split(nano_file, ".tsv")[1]

    print_log("Getting read IDs...")

    # Obtain all read names
    header = ""
    read_nm_vec = fill("", 200000)
    open(nano_file) do f
        next_entry = 1
        for line in enumerate(eachline(f))
            # Record and skip header
            if line[1] == 1
                header =  line[2]
                continue
            end
            # Get ID
            line_id = split(line[2], "\t")[5]
            # Add if necessary
            min_ind = max(1, next_entry - 100)
            ind_eq = findfirst(isequal(line_id), read_nm_vec[min_ind:next_entry])
            if isnothing(ind_eq)
                read_nm_vec[next_entry] = line_id
                mod(next_entry, 1000) == 0 && println(next_entry)
                next_entry += 1
            end
        end
    end

    # Remove empty
    deleteat!(read_nm_vec, read_nm_vec .== "")

    # Remove duplicates
    read_nm_vec = unique(read_nm_vec)

    # Number of reads per file
    base_num = div(length(read_nm_vec), n_files)
    extra_num = mod(length(read_nm_vec), n_files)

    print_log("Permuting reads...")

    # Permute order of reads so to not create bias
    permute!(read_nm_vec, sample(1:length(read_nm_vec), length(read_nm_vec);replace=false))

    print_log("Randomly assigning read IDs to files...")

    # Assign read to each file
    read_file_ind = Vector{Int64}()
    file_bnds = [i * base_num + extra_num for i = 1:n_files]
    @inbounds for i = 1:length(read_nm_vec)
        push!(read_file_ind, findfirst(x -> i <= x, file_bnds))
    end

    # Write output files
    pmap(x -> pmap_split_nanopolish_file(nano_file, n_files, header, pref, read_file_ind, read_nm_vec, x), 1:n_files)

    # Return nothing
    return nothing

end
"""
    `get_calls_nanopolish(CALL_DIC,REG_DATA)`

    Function that using the takes overlapping the analysis region and extracts information for each CpG
    group.

    # Examples
    ```julia-repl
    julia> CpelNano.get_calls_nanopolish(call_dic,rs)
    ```
"""
function get_calls_nanopolish!(call_dic::Dict{String,Vector{Vector{SubString}}}, rs::RegStruct, mod_nse::Bool)::Nothing

    # For each read in call_dic
    calls = Vector{Vector{MethCallCpgGrp}}()
    @inbounds for kk in keys(call_dic)

        # Initizalize vector of group calls for the read
        call = [MethCallCpgGrp() for l = 1:length(rs.cpg_grps)]

        @inbounds for cpg_grp_dat in call_dic[kk]

            # Get chr_start (0-based in nanopolish)
            # Note: nanopolish does not consider the ±5 flanking bp
            cpg_grp_st = parse(Int64, cpg_grp_dat[3]) + 1 - 5
            cpg_grp_end = parse(Int64, cpg_grp_dat[4]) + 1 + 5

            # Find correct index of CpG group in region
            x_index_st = findfirst(x -> minimum(x.grp_int) == cpg_grp_st, rs.cpg_grps)
            x_index_end = findfirst(x -> maximum(x.grp_int) == cpg_grp_end, rs.cpg_grps)
            (isnothing(x_index_st) || isnothing(x_index_end)) && continue
                # print_log("x_index_st: $(x_index_st)")

            # Check if there is a discrepancy in CpG group coordinates
            if x_index_st != x_index_end
                print_log("[WARNING] CpG group coordinates do not match nanopolish output ... ")
                continue
            end

            # Record log p(y|x̄)
            log_pyx_m = parse(Float64, cpg_grp_dat[7])
            log_pyx_u = parse(Float64, cpg_grp_dat[8])
            if mod_nse
                # Modeling noise uses log p's given by nanopolish
                call[x_index_st].log_pyx_m = log_pyx_m
                call[x_index_st].log_pyx_u = log_pyx_u
            else
            # Not modeling noise does not use log p's given by nanopolish
                call[x_index_st].log_pyx_m = log_pyx_m > log_pyx_u ? log_pyx_right_x : log_pyx_wrong_x
        call[x_index_st].log_pyx_u = log_pyx_m > log_pyx_u ? log_pyx_wrong_x : log_pyx_right_x
            end

            # CpG group is observed in read
            call[x_index_st].obs = true

        end

        # Push call
        push!(calls, call)

    end

    # Store call info
    rs.calls = calls
rs.m = length(calls)

# Return nothing
    return nothing

end
"""
    `get_calls!(REG,CALL_FILE,CONFIG)`

    Function that stores observations in REG from calls in CALL_FILE.

    # Examples
    ```julia-repl
    julia> CpelNano.get_calls!(reg,call_file,config)
    ```
"""
function get_calls!(rs::RegStruct, call_file::String, config::CpelNanoConfig)::Nothing

    # Get calls
    if config.caller == "nanopolish"
        # nanopolish (sorted)
        data_dic = get_dic_nanopolish(rs, call_file)
        get_calls_nanopolish!(data_dic, rs, config.mod_nse)
else
        # exit if not valid caller
        print_log("Invalid option for config.caller ...")
        sleep(5)
        exit(0)
    end

    # Return nothing
    return nothing

end
"""
    `rscle_grp_mod!(RS)`

    Function that rescales group model (coarse) to per cpg model (fine).

    # Examples
    ```julia-repl
    julia> CpelNano.rscle_grp_mod!(rs)
    ```
"""
function rscle_grp_mod!(rs::RegStruct)::Nothing

    # Redefine groups as singletons
    rs.L = rs.N
    rs.dl = rs.dn
    rs.ρl = rs.ρn
    rs.Nl = fill(1, rs.N)

    # Return nothing
    return nothing

end
"""
    `pmap_analyze_reg(REG_INT,CHR,NANO,FA_REC,CONFIG)`

    Function that estimates ϕ for a given analysis region starting at REG_ST and stores MML, NME, as well as the
    parameter vector ϕ.

    # Examples
    ```julia-repl
    julia> CpelNano.pmap_analyze_reg(reg_int,chr,nano,fa_rec,config)
    ```
"""
function pmap_analyze_reg(reg_int::UnitRange{Int64}, chr::String, nano::String, fa_rec::FASTA.Record, config::CpelNanoConfig)::RegStruct

    ## Get info about analysis region

    # Initizalize empty region structure
    rs = RegStruct()

    # Store chromosomal info (1-based)
    rs.chr = chr
    rs.chrst = minimum(reg_int)
    rs.chrend = maximum(reg_int)

    # Print out region being processed
    print_log("Processing $(rs.chr):$(rs.chrst)-$(rs.chrend)")

    # Get CpG site information (1-based)
    get_grp_info!(rs, fa_rec, config.min_grp_dist)
    length(rs.cpg_pos) > 9 || return rs
    config.verbose && print_log("CpG sites: $(rs.cpg_pos)")
    config.verbose && print_log("CpG groups: $(rs.cpg_grps)")

    # Retrieve region calls
    get_calls!(rs, nano, config)
    config.verbose && print_log("Number of reads: $(rs.m)")
    config.verbose && print_log("Average depth: $(get_ave_depth(rs))")
    config.verbose && print_log("Percentage groups observed: $(perc_gprs_obs(rs) * 100)%")
        # config.verbose && print_log("Calls: $(rs.calls)")
        # readline()

    # Skip region if not enough data
    get_ave_depth(rs) >= config.min_cov || return rs
    perc_gprs_obs(rs) >= 0.66 || return rs

    # Skip region if only one group
    rs.L > 1 || return rs

    ## Estimate ϕs

    # Estimate parameter vector ϕ
    get_ϕhat!(rs, config)

    # Leave if we couldn't estimate ϕ
    length(rs.ϕhat) > 0 || return rs
    config.verbose && print_log("ϕhat: $(rs.ϕhat)")
        # readline()

    ## Re-scale to per cpg site resolution
    rscle_grp_mod!(rs)
    config.verbose && print_log("rs.Nl: $(rs.Nl)")

    ## Statistical summaries
    get_stat_sums!(rs, config)

    # Set estimation region as processed
    rs.proc = true

    # Return
    return rs

end
"""
    `analyze_nano(NANO,FASTA,CONFIG)`

    Function that estimates ϕ for all analysis regions and stores MML, NME, as well as the parameter vector ϕ.

    # Examples
    ```julia-repl
    julia> CpelNano.analyze_nano(nano,fasta,config)
    ```
"""
function analyze_nano(nano::String, fasta::String, config::CpelNanoConfig)::Nothing

    # Find chromosomes
    chr_names, chr_sizes = get_genome_info(fasta)

    # Create output directory if not existant
    isdir(config.out_dir) || mkdir(config.out_dir)

    # Get output file names
    config.out_files = OutputFiles(config.out_dir, config.out_prefix)

    # Check if output file exists
    check_output_exists(config) && return nothing

    # If BED file provided, identify intervals of interest
    print_log("Checking for target regions ...")
    targ_regs = isempty(config.bed_reg) ? nothing : read_bed_reg(config)

    # Loop over chromosomes
    for i = 1:length(chr_names)

        # Get chromosome name and size
        chr = chr_names[i]
        chr_size = chr_sizes[i]

        # Print info
        print_log("Processing chr $(chr) of length $(chr_size) ...")

        # Get analysis region start positions
        if isnothing(targ_regs)
            # If no BED file provided partition entire chromosome
            print_log("Partitioning chr $(chr) ...")
            chr_part = get_chr_part(chr, config, fasta)
        else
            # If BED file provided, partition regions of interest. NOTE: ensure no split CpG groups
            print_log("Partitioning target regions in chr $(chr) ...")
            haskey(targ_regs, chr) || continue
            chr_part = get_chr_part(chr, config, fasta, targ_regs[chr])
        end

        # Get FASTA record for chr
        fa_rec = get_chr_fa_rec(chr, fasta)

        # Process each analysis region in chr
        out_pmap = pmap(x -> pmap_analyze_reg(x, chr, nano, fa_rec, config), chr_part)

        # Write chromosome
        write_output(out_pmap, config)

    end

    # Return nothing
    return nothing

end
"""
    `pmap_nano_smp_est(REG_INT,CHR,NANO,FA_REC,CONFIG)`

    Function that estimates computes sample statistics E[X] and E[XX] for a given analysis region.

    # Examples
    ```julia-repl
    julia> CpelNano.pmap_nano_smp_est(reg_int,chr,nano,fa_rec,config)
    ```
"""
function pmap_nano_smp_est(reg_int::UnitRange{Int64}, chr::String, nano::String, fa_rec::FASTA.Record, config::CpelNanoConfig)::RegStruct

    ## Get info about analysis region

    # Initizalize empty region structure
    rs = RegStruct()

    # Store chromosomal info (1-based)
    rs.chr = chr
    rs.chrst = minimum(reg_int)
    rs.chrend = maximum(reg_int)

    # Print out region being processed
    print_log("Processing $(rs.chr):$(rs.chrst)-$(rs.chrend)")

    # Get CpG site information (1-based)
    get_grp_info!(rs, fa_rec, config.min_grp_dist)
    length(rs.cpg_pos) > 9 || return rs
    config.verbose && print_log("CpG sites: $(rs.cpg_pos)")
    config.verbose && print_log("CpG groups: $(rs.cpg_grps)")

    # Retrieve region calls
    get_calls!(rs, nano, config)
    config.verbose && print_log("Number of reads: $(rs.m)")
    config.verbose && print_log("Average depth: $(get_ave_depth(rs))")
    config.verbose && print_log("Percentage groups observed: $(perc_gprs_obs(rs) * 100)%")
        # config.verbose && print_log("Calls: $(rs.calls)")
        # readline()

    # Skip region if not enough data
    get_ave_depth(rs) >= config.min_cov || return rs
    perc_gprs_obs(rs) >= 0.66 || return rs

    # Skip region if only one group
    rs.L > 1 || return rs

    ## Compute sample averages
    comp_smp_exps!(rs)
    config.verbose && print_log("E[X]: $(rs.exps.ex)")
    config.verbose && print_log("E[XX]: $(rs.exps.exx)")
        # config.verbose && readline()

    # Set estimation region as processed
    rs.proc = true

    # Return
    return rs

end
"""
    `nano_smp_est(NANO,FASTA,CONFIG)`

    Function that computes sample statistics E[X] and E[XX].

    # Examples
    ```julia-repl
    julia> CpelNano.nano_smp_est(nano,fasta,config)
    ```
"""
function nano_smp_est(nano::String, fasta::String, config::CpelNanoConfig)::Nothing

    # Find chromosomes
    chr_names, chr_sizes = get_genome_info(fasta)

    # Create output directory if not existant
    isdir(config.out_dir) || mkdir(config.out_dir)

    # Get output file names
    config.out_files = OutputFiles(config.out_dir, config.out_prefix)

    # Check if output file exists
    check_output_exists(config) && return nothing

    # If BED file provided, identify intervals of interest
    print_log("Checking for target regions ...")
    targ_regs = isempty(config.bed_reg) ? nothing : read_bed_reg(config)

    # Loop over chromosomes
    for i = 1:length(chr_names)

        # Get chromosome name and size
        chr = chr_names[i]
        chr_size = chr_sizes[i]

        # Print info
        print_log("Processing chr $(chr) of length $(chr_size) ...")

        # Get analysis region start positions
        print_log("Partitioning chr $(chr) ...")
        chr_part = get_chr_part(chr, config, fasta)

        # Get FASTA record for chr
        fa_rec = get_chr_fa_rec(chr, fasta)

        # Process each analysis region in chr
        out_pmap = pmap(x -> pmap_nano_smp_est(x, chr, nano, fa_rec, config), chr_part)

        # Write output
        write_output_ex(config.out_files.ex_file, out_pmap)
        write_output_exx(config.out_files.exx_file, out_pmap)
        write_output_ϕ_marg(config.out_files.theta_file, out_pmap)

    end

    # Return nothing
    return nothing

end


text:
# #############################################################################
# #############################################################################
# Plotting commands

#This file contains the following functions for common plots, used to render low-res images of ray paths for easy qualitative identification of modes.

#   plotbnd(bnd::Boundary;axes=plt.gca(),normvec::Bool=false)
#   Plots the boundary <bnd> on the specified matplotlib axes <axes>. If <normvec> is true, also plots some arrows showing the normal vector at some points along the boundary, for a visual verification of the derivative functions.

#   plotidx(idx::RefractiveIndex,[bnd::Boundary];axes=plt.gca())
#   Plots a filled labelled contour plot of the refractive index distribution <idx> on the specified matplotlib axes <axes>. If the boundary <bnd> is specified, the contour plot is not plotted in the entire axes but limited to the region within the boundary, so this can be plotted alongside the boundary plot for a nice image of the cavity.

#   rayline = plotrays(raypath::Array{Float64,2};axes=plt.gca())
#   Plots on the specified matplotlib axes <axes> the rays along the path specified by <raypath>, which has rows corresponding to the polar (<r>,<theta>) coordinates of the ray position. Returns the matplotlib line object for the plotted rays.

#The following functions are used for storing or retrieving the cavity and ray plots as bitmap images, which are sufficient for obtaining a qualitative idea of any mode present, but much less disk space intensive than storing the entire raypath data directly.

#   cavityimgpath = getcavityimgfile(resultsdir::AbstractString)
#   Returns the path <cavityimgpath> to the cavity image file for a specified simulation run in the results directory <resultsdir>.

#   rayimgpath = getrayimgfile(resultsdir::AbstractString,resultid::Int64)
#   Returns the path <rayimgpath> to the ray image file for the <resultid>th computation of a specified simulation run in the results directory <resultsdir>.

#   range = writecavityimg(resultsdir::AbstractString,bnd::Boundary,idx::RefractiveIndex)
#   Writes an image of the cavity with boundary <bnd> and index distribution <idx> in the results directory <resultsdir>. Returns the range <range> of the matplotlib axes plot for calibrating the ray images.

#   rayline = writerayimg(resultsdir::AbstractString,resultid::Int64, raypath::Array{Float64,2},range::Array{Float64,1}=[-1.0,1.0,-1.0,1.0])
#   Writes an image of the rays following the path <raypath> from the <resultid>th computation of a specified simulation run in the results directory <resultsdir>, adjusting the matplotlib axes range to <range>. Returns the matplotlib line object <rayline> from the plot.

#   writerayimg(resultsdir::AbstractString,resultid::Int64,raypath::Array{Float64,2},rayline::PyObject)
#   Writes an image of the rays following the path <raypath> from the <resultid>th computation of a specified simulation run in the results directory <resultsdir>. Obtains image by altering the line data in the matplotlib line object <rayline> of an existing ray image, for a speedier plot.

# #############################################################################
# #############################################################################
#Initiate

#Dependencies
#require("util.jl")
#require("boundary.jl")
#require("refractive_index.jl")
#require("io.jl")
import PyPlot.plt
plt[:ioff]()
using PyCall

# #############################################################################
# #############################################################################
#Plotting tools

#Common preparation of empty plot for cavity and ray plotting functions
function prepareplot(;figwidth::Float64=5.0,figheight::Float64=5.0)
    ax = plt[:figure](figsize=(figwidth,figheight))[:gca](aspect="equal")
    plt[:gcf]()[:tight_layout]()
    return ax
end

#Plot the boundary
function plotbnd(bnd::Boundary;axes=plt[:gca](),normvec::Bool=false)
    #Plots a boundary, with the option of showing its normal vectors as red emanating lines.
    theta::Array{Float64,1} = linspace(0,2*pi,500)
    r::Array{Float64,1} = rfunc(bnd,theta)
    plt[:plot](r.*cos(theta),r.*sin(theta),"g",axes=axes)
    plt[:axis]([-r[1],r[1],-r[1],r[1]],axes=axes)

    if normvec
        #Plot the norm vector along boundary
        theta_n::Array{Float64,1} = linspace(0,2*pi,51)
        r_n::Array{Float64,1},alpha::Array{Float64,1} = rsys(bnd,theta_n)
        for i = 1:length(theta_n)
            plt[:arrow](r_n[i]*cos(theta_n[i]),r_n[i]*sin(theta_n[i]),0.1*cos(alpha[i]),0.1*sin(alpha[i]);color="r",axes=axes)
        end
    end

    return nothing
end

#Plot the refractive index distribution
function plotidx(idx::RefractiveIndex,bnd::Boundary;axes=plt[:gca]())
    #Plots a filled labelled contour plot of the cavity interior's refractive index distribution, where darker shades indicate higher refractive index.
    thetares::Int64 = 100; rres::Int64 = 100

    theta::Array{Float64,2} = transpose(linspace(0,2*pi,thetares))
    R::Array{Float64,2} = rfunc(bnd,theta)
    rs::Array{Float64,2} = linspace(0.0,1.0,rres)*R
    thetas::Array{Float64,2} = ones(Float64,rres)*theta

    x::Array{Float64,2} = rs.*cos(thetas)
    y::Array{Float64,2} = rs.*sin(thetas)
    n::Array{Float64,2} = nfunc(idx,rs,thetas)

    plt[:contourf](x,y,n,100,cmap="bone_r",axes=axes)
    contobj = plt[:contour](x,y,n,100,cmap="bone_r",alpha=0.2,axes=axes)
    levels = contobj[:levels][1:10:end]
    plt[:clabel](contobj,levels,colors="k")
    return nothing
end

function plotidx(idx::RefractiveIndex;axes=plt[:gca]())
    #Plots a filled labelled contour plot of the refractive index distribution
    res::Int64 = 200
    x::Array{Float64,2} = ones(Float64,res)*transpose(linspace(-1.0,1.0,res))
    y::Array{Float64,2} = linspace(-1.0,1.0,res)*ones(Float64,(1,res))
    rs::Array{Float64,2} = hypot(x,y)
    thetas::Array{Float64,2} = atan2(y,x)
    n::Array{Float64,2} = nfunc(idx,rs,thetas)

    plt[:contourf](x,y,n,100,cmap="bone_r",axes=axes)
    contobj = plt[:contour](x,y,n,100,cmap="bone_r",alpha=0.2,axes=axes)
    levels = contobj[:levels][1:10:end]
    plt[:clabel](contobj,levels,colors="k")
    plt[:title]("Plot of $idx")
    return nothing
end

#Draw rays
function plotrays(raypath::Array{Float64,2};axes=plt[:gca]())
    return plt[:plot](raypath[:,1].*cos(raypath[:,2]),raypath[:,1].*sin(raypath[:,2]),"y-",lw=0.4,aa=false,axes=axes)[1]
end

# #############################################################################
# #############################################################################
#Ray image I/O

#Get cavity image file location
function getcavityimgfile(resultsdir::AbstractString)
    return joinpath(resultsdir,"cavity.png")
end

#Get ray image file location
function getrayimgfile(resultsdir::AbstractString,resultid::Int64)
    return joinpath(resultsdir,getresultfname(resultid,"ray")*".png")
end

#Attempts to store image of cavity if not already present. Returns the axis range of the cavity.
function writecavityimg(resultsdir::AbstractString,bnd::Boundary,idx::RefractiveIndex)
    #Create new figure
    cavityplot = prepareplot()

    #Plot boundary to obtain axis range
    plotbnd(bnd,axes=cavityplot)
    range = [plt[:axis]()...]

    #Continue plotting if cavity image file does not already exist
    cavityimgfile = getcavityimgfile(resultsdir)
    if !isfile(cavityimgfile)
        #Plot index distribution
        plotidx(idx,bnd,axes=cavityplot)
        plt[:savefig](cavityimgfile)
    end

    plt[:close]()
    return range
end

#Stores newly created image of ray and returns the matplotlib line object of the ray plot for more efficient future editing.
function writerayimg(resultsdir::AbstractString,resultid::Int64,raypath::Array{Float64,2},range::Array{Float64,1}=[-1.0,1.0,-1.0,1.0])
    #Create new figure
    rayplot = prepareplot()

    #Plot
    rayline = plotrays(raypath,axes=rayplot)
    plt[:axis](range)
    plt[:axis]("off")

    #Save image
    plt[:savefig](getrayimgfile(resultsdir,resultid),transparent=true)
    return rayline
end

#Stores image of ray after editing the matplotlib line object of an existing ray plot.
function writerayimg(resultsdir::AbstractString,resultid::Int64,raypath::Array{Float64,2},rayline::PyObject)
    #Edit existing ray line
    rayline[:set_xdata](raypath[:,1].*cos(raypath[:,2]))
    rayline[:set_ydata](raypath[:,1].*sin(raypath[:,2]))

    #Save image
    plt[:savefig](getrayimgfile(resultsdir,resultid),transparent=true)
    return nothing
end


text:
function generate(rng::AbstractRNG,
                  ::Type{Event},
                  rates::EventRates{T},
                  time::Float64) where T <: DiseaseStateSequence
  totals = Weights([sum(rates[state]) for state in convert(DiseaseStates, T)[2:end]])
  if sum(totals) == Inf
    new_state = sample(rng, convert(DiseaseStates, T)[2:end], totals)
    id = sample(rng, 1:individuals(rates), Weights(rates[new_state]))
    return Event{T}(time, id, new_state)
  elseif sum(totals) == 0.0
    return NoEvent()
  else
    # Generate new state
    new_state = sample(rng, convert(DiseaseStates, T)[2:end], totals)
    # Generate event individual
    id = sample(rng, 1:individuals(rates), Weights(rates[new_state]))
    return Event{T}(time + rand(rng, Exponential(1.0 / sum(totals))), id, new_state)
  end
end


text:
using DSGE, ModelConstructors, CSV, DataFrames, Dates, Test

# Data vintages and settings
fcast_date      = DSGE.quartertodate("2021-Q1")
date_fcast_end  = iterate_quarters(fcast_date, 40)
cond_type       = :full

# Flexible AIT Rule
tvcred_n_qtrs          = 24 # length of tv cred period
start_tvcred_level     = 0.
end_tvcred_level       = 1.
φ_π                    = 15.
φ_y                    = 3.
Thalf                  = 10.
ρ_smooth               = .75
pgap_init              = 0.125 # NOTE: positive pgap and ygap values imply
ygap_init              = 12.   # NEGATIVE initial pgap and ygaps (i.e. ygap_init = 12 implies
pgap_ygap_init_date    = Date(2020, 6, 30)  # an initial ygap of -12 in 2020:Q2)
start_zlb_date         = Date(2020, 12, 31) # start and end dates should be inclusive, i.e. first period w/zero rate
end_zlb_date           = Date(2023, 9, 30)  # is start_zlb_date, and the last period is end_zlb_date

# Initialize model objects
m = Model1002("ss62")
m <= Setting(:date_forecast_start,  fcast_date)
m <= Setting(:date_conditional_end, fcast_date)
m <= Setting(:forecast_horizons, subtract_quarters(date_fcast_end, fcast_date))

# Load data
df = CSV.read(joinpath(dirname(@__FILE__), "..", "test", "reference", "ss62_modal_forecast_data.csv"), DataFrame)

# Set up Flexible AIT with a temporary ZLB. Note that df will be altered
start_reg_baseline = 6

θ = (φ_π = φ_π, φ_y = φ_y, Thalf = Thalf, pgap = pgap_init,
     ygap = ygap_init, ρ_smooth = ρ_smooth,
     historical_policy = default_policy()) # NamedTuple with parameters of the Flexible AIT policy

m <= Setting(:flexible_ait_φ_π, θ[:φ_π])
m <= Setting(:flexible_ait_φ_y, θ[:φ_y])
m <= Setting(:ait_Thalf, θ[:Thalf])
m <= Setting(:gdp_Thalf, θ[:Thalf])
m <= Setting(:pgap_value, θ[:pgap])
m <= Setting(:pgap_type, :flexible_ait)
m <= Setting(:ygap_value, θ[:ygap])
m <= Setting(:ygap_type, :flexible_ait)
m <= Setting(:flexible_ait_ρ_smooth, θ[:ρ_smooth])
m <= Setting(:alternative_policies, AltPolicy[θ[:historical_policy]])

# Define a function to help update the mappings
# from model regimes to parameter regimes
function set_regime_vals_fnct(m, n)
    m2p = get_setting(m, :model2para_regime)

    for (k, v) in m2p
        for i in 7:n
            v[i] = 1 # For regimes 7 onward, use parameter regime 1
        end
    end

    m
end

# Set up credibility for ZLB and alternative policy
m <= Setting(:gensys2, true)
m <= Setting(:replace_eqcond, true)
m <= Setting(:uncertain_temporary_altpolicy, true)
m <= Setting(:uncertain_altpolicy, true)

# Find the first regime in which gensys2 should apply
gensys2_first_regime = findfirst([get_setting(m, :regime_dates)[i] .== start_zlb_date for i in 1:get_setting(m, :n_regimes)])
n_zlb_reg = DSGE.subtract_quarters(end_zlb_date, start_zlb_date) + 1
if end_zlb_date < get_setting(m, :regime_dates)[get_setting(m, :n_regimes)]
    set_regime_vals_fnct(m, get_setting(m, :n_regimes))
else
    # technically the second input is `gensys2_first_regime + (n_zlb_reg - 1) + 1` b/c
    # gensys2_first_regime is the first regime for the temporary ZLB, so we need to subtract 1
    # from n_zlb_reg, but we need to also add 1 for the lift-off regime
    set_regime_vals_fnct(m, gensys2_first_regime + n_zlb_reg)
end

# Construct regime_eqcond_info
reg_dates = deepcopy(get_setting(m, :regime_dates))
regime_eqcond_info = Dict{Int, EqcondEntry}()
weights = [end_tvcred_level, 1. - end_tvcred_level] # dummy values for now, time-varying credibility added later
for (regind, date) in zip(gensys2_first_regime:(n_zlb_reg - 1 + gensys2_first_regime), # See comments starting at line 75
                          DSGE.quarter_range(reg_dates[gensys2_first_regime], # looping over the date range of the ZLB
                                                iterate_quarters(reg_dates[gensys2_first_regime], n_zlb_reg - 1)))
    reg_dates[regind] = date # updating the regime date dictionary
    regime_eqcond_info[regind] = EqcondEntry(zlb_rule(), weights) # specify which regimes have a ZLB rule applying
end
reg_dates[n_zlb_reg + gensys2_first_regime] = iterate_quarters(reg_dates[gensys2_first_regime], n_zlb_reg) # add liftoff regime date
regime_eqcond_info[n_zlb_reg + gensys2_first_regime] = EqcondEntry(flexible_ait(), weights) # liftoff policy is flexible AIT
m <= Setting(:regime_dates,               reg_dates)
m <= Setting(:regime_eqcond_info,         regime_eqcond_info)
m <= Setting(:temporary_altpolicy_length, n_zlb_reg) # length of the temporary ZLB, otherwise a guess is made. Since we have time-varying credibility, this guess would be wrong
m <= Setting(:zlb_rule_value,             .1) # (optional) set the interest rate during the ZLB, defaults to forecast_zlb_value(m) if not available
m <= Setting(:temporary_altpolicy_names,  [:zlb_rule]) # Ensures compute_system knows the regimes w/ a zlb_rule in EqcondEntry are gensys2 regimes

# Set up regime indices
setup_regime_switching_inds!(m; cond_type = cond_type)

# Set up time-varying credibility (which extends beyond the ZLB, so we need to add more regimes)
tvcred_dates = (start_zlb_date, DSGE.iterate_quarters(start_zlb_date, tvcred_n_qtrs))
n_tvcred_reg = DSGE.subtract_quarters(tvcred_dates[2], tvcred_dates[1]) + 1 # number of regimes for time-varying credibility
reg_start = DSGE.subtract_quarters(tvcred_dates[1], start_zlb_date) + gensys2_first_regime # first regime for time-varying credibility

# Add extra entries for regime_eqcond_info with dummy credibility weights
for (regind, date) in zip(reg_start:(reg_start + n_tvcred_reg),
                          DSGE.quarter_range(reg_dates[reg_start], DSGE.iterate_quarters(reg_dates[reg_start], n_tvcred_reg - 1)))
    reg_dates[regind] = date
    if date > end_zlb_date # Add extra regimes with flexible AIT as the policy after the ZLB's end b/c credibility is still changing
        regime_eqcond_info[regind] = DSGE.EqcondEntry(flexible_ait(), [end_tvcred_level, 1. - end_tvcred_level])
    end
end
m <= Setting(:regime_dates,             reg_dates)
m <= Setting(:regime_eqcond_info, regime_eqcond_info)
setup_regime_switching_inds!(m; cond_type = cond_type)

# Update mapping of model regimes to parameter regimes
set_regime_vals_fnct(m, get_setting(m, :n_regimes))

# Set up TVIS information set
m <= Setting(:tvis_information_set, vcat([i:i for i in 1:(gensys2_first_regime - 1)],
                                         [i:17 for i in gensys2_first_regime:17],
                                         [i:i for i in 18:get_setting(m, :n_regimes)]))

# Now actually populate the weights in regime_eqcond_info so that
# credibility/awareness increases linearly over time
credvec = collect(range(start_tvcred_level, stop = end_tvcred_level, length = n_tvcred_reg))
for (i, k) in enumerate(sort!(collect(keys(regime_eqcond_info))))
    if tvcred_dates[1] <= reg_dates[k] <= tvcred_dates[2]
        get_setting(m, :regime_eqcond_info)[k].weights = [credvec[i], 1. - credvec[i]]
    end
end

# We implement a shorter ZLB with a return to the historical
# monetary policy reaction function as the other policy
# agents believe could occur. This is where the `MultiPeriodAltPolicy` comes into play.
altpol_eqcond_info = deepcopy(get_setting(m, :regime_eqcond_info))
for (reg, eq_entry) in altpol_eqcond_info
    if reg >= start_reg_baseline + 5
        altpol_eqcond_info[reg] = DSGE.EqcondEntry(DSGE.default_policy(), [1.])
    end
    altpol_eqcond_info[reg].weights = [1.]
end

# The key for the MultiPeriodAltPolicy can be anything; it is never used in the code
# to infer some property. Of course, it should still be a descriptive name.
altpol = MultiPeriodAltPolicy(:other_altpol, get_setting(m, :n_regimes), altpol_eqcond_info, gensys2 = true,
                              temporary_altpolicy_names = [:zlb_rule],
                              temporary_altpolicy_length = 6,
                              infoset = copy(get_setting(m, :tvis_information_set)))
delete!(m.settings, :alternative_policies) # if this setting already exists, the type may be wrong
m <= Setting(:alternative_policies, DSGE.AbstractAltPolicy[altpol])

# Run forecasts
output_vars = [:histobs, :histpseudo, :forecastpseudo, :forecastobs]
modal_params = map(x -> x.value, m.parameters)
fcast = DSGE.forecast_one_draw(m, :mode, cond_type, output_vars, modal_params, df,
                               regime_switching = true, n_regimes = get_setting(m, :n_regimes))

nothing


text:
using Distributions
using Plots

population = sort(convert(Array{Int64}, rand(Bernoulli(.7), 100)))
population_over_time = []
n_iter = 100
gr()

function single_reproduction(population::AbstractArray)
    n = length(population)
    pₐ = sum(population) / length(population)
    fₐ = Bernoulli(pₐ)
    sort(convert(Array{Int64}, rand(fₐ, n)))
end

function multiple_reproductions(seed_array::AbstractArray, population_over_time::AbstractArray, n_iterations::Int64)
    seed_array = single_reproduction(seed_array)
    push!(population_over_time, seed_array)
    if n_iterations == 1
        hcat(population_over_time...)
    else
        multiple_reproductions(seed_array, population_over_time, n_iterations-1)
    end
end

function plot_many_samples(n_samples::Int64, n_iterations::Int64, Nₑ::Int64, pₐ::Float64)
    p = false
    for i in 1:n_samples
        population = sort(convert(Array{Int64}, rand(Bernoulli(pₐ), Nₑ)))
        population_at_time = multiple_reproductions(population, [], n_iterations)
        if p
            display(plot!(dropdims(sum(population_at_time, dims=1), dims=1), legend=false))
        else
            display(plot(dropdims(sum(population_at_time, dims=1), dims=1), xlims=(0, n_iterations), ylims=(0, Nₑ)))
            p = true
        end
    end
end

plot_many_samples(10, 10000, 10000, .5)
# j = multiple_reproductions(population, population_over_time, n_iter)
# plot(dropdims(sum(j, dims=1), dims=1))
# j = multiple_reproductions(population, [], n_iter)
# plot!(dropdims(sum(j, dims=1), dims=1))
# j = multiple_reproductions(population, [], n_iter)
# plot!(dropdims(sum(j, dims=1), dims=1))

max_stars_repo_name: biothomme/Geosmin

max_issues_repo_name: biothomme/Geosmin

max_forks_repo_name: biothomme/Geosmin


text:
export rule

@rule GammaShapeRate(:out, Marginalisation) (m_α::PointMass, m_β::PointMass) = GammaShapeRate(mean(m_α), mean(m_β))

@rule GammaShapeRate(:out, Marginalisation) (q_α::PointMass, q_β::PointMass) = GammaShapeRate(mean(q_α), mean(q_β))


text:
@benchgroup "insertions" begin
  n = 10000
  @bench "ER Generation" g = Graph($n, 16 * $n)
end