text:
# This file contains a library of functions for splitting/casting/merging
# partition types (PTs). Any `pt_lib.jl` should have a corresponding
# `pt_lib_info.json` that contains an annotation for each
# splitting/casting/merging that describes how data should be partitioned
# in order for that function to be applicable.

###################################
# Splitting and merging functions #
###################################

ReturnNull(
    src,
    params,
    batch_idx::Integer,
    nbatches::Integer,
    comm::MPI.Comm,
    loc_name,
    loc_params,
) = nothing

format_available_memory() =
    format_bytes(Sys.free_memory()) * " / " * format_bytes(Sys.total_memory())

function sortablestring(val, maxval)
    s = string(val)
    maxs = string(maxval)
    res = fill('0', length(maxs))
    res[length(res)-length(s)+1:length(res)] .= collect(s)
    join(res)
end

function ReadBlock(
    src,
    params,
    batch_idx::Integer,
    nbatches::Integer,
    comm::MPI.Comm,
    loc_name,
    loc_params,
)
    # TODO: Implement a Read for balanced=false where we can avoid duplicate
    # reading of the same range in different reads

    # Handle single-file nd-arrays
    # We check if it's a file because for items on disk, files are HDF5
    # datasets while directories contain Parquet, CSV, or Arrow datasets
    path = getpath(loc_params["path"])
    if (loc_name == "Remote" && (occursin(".h5", loc_params["path"]) || occursin(".hdf5", loc_params["path"]))) ||
        (loc_name == "Disk" && HDF5.ishdf5(path))
        # @show isfile(path)
        f = h5open(path, "r")
        dset = loc_name == "Disk" ? f["part"] : f[loc_params["subpath"]]

        ismapping = false
        # TODO: Use `view` instead of `getindex` in the call to
        # `split_on_executor` here if HDF5 doesn't support this kind of usage
        # TODO: Support modifying a memory-mappable file here without having
        # to read and then write back
        # if ismmappable(dset)
        #     ismapping = true
        #     dset = readmmap(dset)
        #     close(f)
        #     dset = split_on_executor(dset, params["key"], batch_idx, nbatches, comm)
        # else
        dim = params["key"]
        dimsize = size(dset, dim)
        dimrange = split_len(dimsize, batch_idx, nbatches, comm)
        dset = if length(dimrange) == 0
            # If we want to read in an emoty dataset, it's a little tricky to
            # do that with HDF5.jl. But this is how we do it:
            if dimsize == 0
                dset[[Colon() for _ in 1:ndims(dset)]...]
            else
                dset[[
                    # We first read in the first slice into memory. This is
                    # because HDF5.jl (unlike h5py) does not support just
                    # reading in an empty `1:0` slice.
                    if i == dim
                        1:1
                    else
                        Colon()
                    end for i = 1:ndims(dset)
                ]...][[
                    # Then once that row is in memory we just remove it so
                    # that we have the appropriate empty slice.
                    if i == dim
                        1:0
                    else
                        Colon()
                    end for i = 1:ndims(dset)
                ]...]
            end
        else
            # If it's not an empty slice that we want to read, it's pretty
            # straightforward - we just specify the slice.
            dset[[
                if i == dim
                    dimrange
                else
                    Colon()
                end for i = 1:ndims(dset)
            ]...]
        end
        close(f)
        return dset
    end

    # Handle single-file replicated objects
    if loc_name == "Disk" && isfile(path)
        res = deserialize(path)
        return res
    end

    # Handle multi-file tabular datasets

    # Handle None location by finding all files in directory used for spilling
    # this value to disk
    if loc_name == "Disk"
        # TODO: Only collect files and nrows info for this location associated
        # with a unique name corresponding to the value ID - only if this is
        # the first batch or loop iteration.
        name = loc_params["path"]
        name_path = getpath(name)
        # TODO: isdir might not work for S3FS
        if isdir(name_path)
            files = []
            nrows = 0
            for partfilename in readdir(name_path)
                part_nrows = parse(
                    Int64,
                    replace(split(partfilename, "_nrows=")[end], ".arrow" => ""),
                )
                push!(
                    files,
                    Dict("nrows" => part_nrows, "path" => joinpath(name, partfilename)),
                )
                nrows += part_nrows
            end
            loc_params["files"] = files
            loc_params["nrows"] = nrows
        else
            # This is the case where no data has been spilled to disk and this
            # is maybe just an intermediate variable only used for this stage.
            # We never spill tabular data to a single file - it's always a
            # directory of Arrow files.
            return nothing
        end
    end

    # Iterate through files and identify which ones correspond to the range of
    # rows for the batch currently being processed by this worker
    nrows = loc_params["nrows"]
    rowrange = split_len(nrows, batch_idx, nbatches, comm)
    dfs::Vector{DataFrames.DataFrame} = []
    rowsscanned = 0
    for file in sort(loc_params["files"], by = filedict -> filedict["path"])
        newrowsscanned = rowsscanned + file["nrows"]
        filerowrange = (rowsscanned+1):newrowsscanned
        # Check if the file corresponds to the range of rows for the batch
        # currently being processed by this worker
        if isoverlapping(filerowrange, rowrange)
            # Deterine path to read from
            file_path = file["path"]
            path = getpath(file_path)

            # Read from location depending on data format
            readrange =
                max(rowrange.start, filerowrange.start):min(
                    rowrange.stop,
                    filerowrange.stop,
                )
            header = 1
            # TODO: Scale the memory usage appropriately when splitting with
            # this and garbage collect if too much memory is used.
            if endswith(file_path, ".csv")
                f = CSV.File(
                    path,
                    header = header,
                    skipto = header + readrange.start - filerowrange.start + 1,
                    footerskip = filerowrange.stop - readrange.stop,
                )
                push!(dfs, DataFrames.DataFrame(f, copycols=false))
            elseif endswith(file_path, ".parquet")
                f = Parquet.read_parquet(
                    path,
                    rows = (readrange.start-filerowrange.start+1):(readrange.stop-filerowrange.start+1),
                )
                push!(dfs, DataFrames.DataFrame(f, copycols=false))
            elseif endswith(file_path, ".arrow")
                rbrowrange = filerowrange.start:(filerowrange.start-1)
                for tbl in Arrow.Stream(path)
                    rbrowrange = (rbrowrange.stop+1):(rbrowrange.stop+Tables.rowcount(tbl))
                    if isoverlapping(rbrowrange, rowrange)
                        readrange =
                            max(rowrange.start, rbrowrange.start):min(
                                rowrange.stop,
                                rbrowrange.stop,
                            )
                        df = let unfiltered = DataFrames.DataFrame(tbl, copycols=false)
                            unfiltered[
                                (readrange.start-rbrowrange.start+1):(readrange.stop-rbrowrange.start+1),
                                :,
                            ]
                        end
                        push!(dfs, df)
                    end
                end
            else
                error("Expected CSV or Parquet or Arrow format")
            end
        end
        rowsscanned = newrowsscanned
    end

    # Concatenate and return
    # NOTE: If this partition is empty, it is possible that the result is
    # schemaless (unlike the case with HDF5 where the resulting array is
    # guaranteed to have its ndims correct) and so if a split/merge/cast
    # function requires the schema (for example for grouping) then it must be
    # sure to take that account
    if isempty(dfs)
        # Note that if we are reading disk-spilled Arrow data, we would have
        # files for each of the workers that wrote that data. So there should
        # be files but they might be empty.
        if loc_name == "Disk"
            files_sorted_by_nrow = sort(loc_params["files"], by = filedict -> filedict["nrows"])
            if isempty(files_sorted_by_nrow)
                # This should not be empty for disk-spilled data
                DataFrame()
            else
                empty(DataFrames.DataFrame(Arrow.Table(getpath(first(files_sorted_by_nrow)["path"])), copycols=false))
            end
        else
            # When we construct the location, we store an empty data frame with The
            # correct schema.
            from_jl_value_contents(loc_params["emptysample"])
        end
    elseif length(dfs) == 1
        dfs[1]
    else
        vcat(dfs...)
    end
end

splitting_divisions = IdDict()

function ReadGroup(
    src,
    params,
    batch_idx::Integer,
    nbatches::Integer,
    comm::MPI.Comm,
    loc_name,
    loc_params,
)
    # TODO: Store filters in parameters of the PT and use them to do
    # partition pruning, avoiding reads that are unnecessary

    # Get information needed to read in the appropriate group
    divisions = params["divisions"]
    key = params["key"]
    rev = params["rev"] # Passed in ReadBlock
    nworkers = get_nworkers(comm)
    npartitions = nworkers * nbatches
    partition_divisions = get_divisions(divisions, npartitions)

    # TODO: Do some reversing here instead of only doing it later in Shuffle
    # to ensure that sorting in reverse order works correctly

    # The first and last partitions (used if this lacks a lower or upper bound)
    # must have actual division(s) associated with them. If there is no
    # partition that has divisions, then they will all be skipped and -1 will
    # be returned. So these indices are only used if there are nonempty
    # divisions.
    hasdivision = any(x->!isempty(x), partition_divisions)
    firstdivisionidx = findfirst(x->!isempty(x), partition_divisions)
    lastdivisionidx = findlast(x->!isempty(x), partition_divisions)
    firstbatchidx = nothing
    lastbatchidx = nothing

    # Get the divisions that are relevant to this batch by iterating
    # through the divisions in a stride and consolidating the list of divisions
    # for each partition. Then, ensure we use boundedlower=true only for the
    # first batch and boundedupper=true for the last batch.
    curr_partition_divisions = []
    for worker_division_idx = 1:nworkers
        for batch_division_idx = 1:nbatches
            # partition_division_idx =
            #     (worker_division_idx - 1) * nbatches + batch_division_idx
            partition_division_idx =
                get_partition_idx(batch_division_idx, nbatches, worker_division_idx)
            if batch_division_idx == batch_idx
                # Get the divisions for this partition
                p_divisions = partition_divisions[partition_division_idx]

                # We've already used `get_divisions` to get a list of min-max
                # tuples (we call these tuples "divisions") for each partition
                # that `ReadGroup` produces. But we only want to read in all
                # the partitions relevant for this batch. But it is important
                # then that `curr_partition_divisions` has an element for each
                # worker. That way, when we call `Shuffle`, it will properly
                # read data onto each worker that is in the appropriate
                # partition.
                push!(
                    curr_partition_divisions,
                    p_divisions,
                )
            end

            # Find the batches that have the first and last divisions
            if partition_division_idx == firstdivisionidx
                firstbatchidx = batch_division_idx
            end
            if partition_division_idx == lastdivisionidx
                lastbatchidx = batch_division_idx
            end
        end
    end

    # Read in each batch and shuffle it to get the data for this partition
    parts = []
    for i = 1:nbatches
        # Read in data for this batch
        part = ReadBlock(src, params, i, nbatches, comm, loc_name, loc_params)

        # Shuffle the batch and add it to the set of data for this partition
        params["divisions_by_worker"] = curr_partition_divisions
        push!(
            parts,
            Shuffle(
                part,
                Dict{String,Any}(),
                params,
                comm,
                boundedlower = !hasdivision || batch_idx != firstbatchidx,
                boundedupper = !hasdivision || batch_idx != lastbatchidx,
                store_splitting_divisions = false
            ),
        )
        delete!(params, "divisions_by_worker")
    end

    # Concatenate together the data for this partition
    res = merge_on_executor(parts...; key = key)

    # If there are no divisions for any of the partitions, then they are all
    # bounded. For a partition to be unbounded on one side, there must be a
    # division(s) for that partition.

    # Store divisions
    global splitting_divisions
    partition_idx = get_partition_idx(batch_idx, nbatches, comm)
    splitting_divisions[res] =
        (partition_divisions[partition_idx], !hasdivision || partition_idx != firstdivisionidx, !hasdivision || partition_idx != lastdivisionidx)

    res
end

function rmdir_on_nfs(actualpath)
    if isdir(actualpath)
        for actualpath_f in readdir(actualpath, join=true)
            if isfile(actualpath_f)
                rm(actualpath_f, force=true)
            end
        end
    end
    # TODO: Also try to remove the directory itself right away although there
    # might still be .nfs files in it. This isn't too much of a problem since we _do_ try
    # to remove all directories at the end of the job.
end

function Write(
    src,
    part,
    params,
    batch_idx::Integer,
    nbatches::Integer,
    comm::MPI.Comm,
    loc_name,
    loc_params,
)
    # Get rid of splitting divisions if they were used to split this data into
    # groups
    global splitting_divisions
    delete!(splitting_divisions, part)

    # Get path of directory to write to
    path = loc_params["path"]
    if startswith(path, "http://") || startswith(path, "https://")
        error("Writing to http(s):// is not supported")
    elseif startswith(path, "s3://")
        path = getpath(path)
        # NOTE: We expect that the ParallelCluster instance was set up
        # to have the S3 filesystem mounted at ~/s3fs/<bucket name>
    else
        # Prepend "efs/" for local paths
        path = getpath(path)
    end

    # Write file for this partition
    worker_idx = get_worker_idx(comm)
    idx = get_partition_idx(batch_idx, nbatches, comm)
    if isa_df(part)
        actualpath = deepcopy(path)
        if nbatches > 1
            # Add _tmp to the end of the path
            if endswith(path, ".parquet")
                path = replace(path, ".parquet" => "_tmp.parquet")
            elseif endswith(path, ".csv")
                path = replace(path, ".csv" => "_tmp.csv")
            elseif endswith(path, ".arrow")
                path = replace(path, ".arrow" => "_tmp.arrow")
            else
                path = path * "_tmp"
            end
        end

        # TODO: Delete existing files that might be in the directory but first
        # finish writing to a path*"_new" directory and then linking ... or
        # something like that. Basically we need to handle the case where we
        # have batching in the first PT and we are reading from and writing to
        # the same directory.
        # 1. Write all output to a new directory
        # 2. On the last partition, do a barrier (or a gather) and then delete
        # the old directory
        # 3. Do another barrier on the last batch and delete the old directory
        # and link to the new one

        # NOTE: This is only needed because we might be writing to the same
        # place we are reading from. And so we want to make sure we finish
        # reading before we write the last batch
        if batch_idx == nbatches
            MPI.Barrier(comm)
        end

        if worker_idx == 1
            if nbatches == 1
                # If there is no batching we can delete the original directory
                # right away. Otherwise, we must delete the original directory
                # only at the end.
                # TODO: When refactoring the locations, think about how to allow
                # stuff in the directory
                rmdir_on_nfs(actualpath)
            end

            # Create directory if it doesn't exist
            # TODO: Avoid this and other filesystem operations that would be costly
            # since S3FS is being used
            if batch_idx == 1
                rmdir_on_nfs(path)
                mkpath(path)
            end
        end
        MPI.Barrier(comm)

        nrows = size(part, 1)
        sortableidx = sortablestring(idx, get_npartitions(nbatches, comm))
        if endswith(path, ".parquet")
            # Parquet.jl doesn't support writing empty data frames.
            if nrows > 0
                partfilepath = joinpath(path, "part$sortableidx" * "_nrows=$nrows.parquet")
                Parquet.write_parquet(partfilepath, part)
            end
        elseif endswith(path, ".csv")
            partfilepath = joinpath(path, "part$sortableidx" * "_nrows=$nrows.csv")
            CSV.write(partfilepath, part)
        else
            partfilepath = joinpath(path, "part$sortableidx" * "_nrows=$nrows.arrow")
            Arrow.write(partfilepath, part)
        end
        MPI.Barrier(comm)
        if nbatches > 1 && batch_idx == nbatches
            tmpdir = readdir(path)
            if worker_idx == 1
                rmdir_on_nfs(actualpath)
                mkpath(actualpath)
            end
            MPI.Barrier(comm)
            for batch_i = 1:nbatches
                idx = get_partition_idx(batch_i, nbatches, worker_idx)
                tmpdir_idx = findfirst(fn -> startswith(fn, "part$idx"), tmpdir)
                if !isnothing(tmpdir_idx)
                    tmpsrc = joinpath(path, tmpdir[tmpdir_idx])
                    actualdst = joinpath(actualpath, tmpdir[tmpdir_idx])
                    cp(tmpsrc, actualdst, force=true)
                end
            end
            MPI.Barrier(comm)
            if worker_idx == 1
                rmdir_on_nfs(path)
            end
            MPI.Barrier(comm)
        end
        src
        # TODO: Delete all other part* files for this value if others exist
    elseif isa_array(part) && hasmethod(HDF5.datatype, (eltype(part),))
        # TODO: Use Julia serialization to write arrays as well as other
        # objects to disk. This way, we won't trip up when we come across
        # a distributed array that we want to write to disk but can't because
        # of an unsupported HDF5 data type.
        # TODO: Support missing values in the array for locations that use
        # Julia serialized objects
        part = DataFrames.disallowmissing(part)

        dim = params["key"]
        # TODO: Ensure that wherever we are using MPI for reduction or
        # broadcasts, we should always ensure that we cast back into
        # the original data type. Even if we have an isbits type like a tuple,
        # we will still want to cast to a tuple
        # TODO: Ensure this works where some partitions are empty
        # If each worker has a single batch, compute the total size
        # TODO: Compute exscan of size across the nodes, on the last node,
        # create the dataset, apply a barrier, and then synchronize so that
        # we can have each worker write on the range from its scanned value
        # with the size of data

        # TODO: Check if HDF5 dataset is created. If not, wait for the master
        # node to create it.
        # If each worker has a single batch, create a file on the first worker.
        # Otherwise, if we need multiple batches for each partition, write each
        # batch to a separate group. TODO: If we need multiple batches, don't
        # write the last batch to its own group. Instead just write it into the
        # aggregated group.
        group_prefix = loc_name == "Disk" ? "part" : loc_params["subpath"]
        partition_idx = get_partition_idx(batch_idx, nbatches, comm)
        worker_idx = get_worker_idx(comm)
        nworkers = get_nworkers(comm)
        group = nbatches == 1 ? group_prefix : group_prefix * "_part$idx" * "_dim=$dim"

        # TODO: Have an option in the location to set this to either "w" or
        # "cw". Both will create a new file if it's not already there but
        # "w" will delete all existing datasets while "cw" will keep all
        # existing datasets (and fail if you try to write over anything). We
        # should also maybe have an option for "cw" but if you try to write
        # to an existing dataset, we will simply delete that dataset first.
        # dataset_writing_permission = "cw"
        force_overwrite = true

        info = MPI.Info()

        # Write out to an HDF5 dataset differently depending on whether there
        # are multiple batches per worker or just one per worker
        if nbatches == 1
            # Determine the offset into the resulting HDF5 dataset where this
            # worker should write
            offset = MPI.Exscan(size(part, dim), +, comm)
            if worker_idx == 1
                offset = 0
            end

            # Create file if not yet created
            # TODO: Figure out why sometimes a deleted file still `isfile`
            f = h5open(
                path,
                "cw",
                fapl_mpio = (comm, info),
                dxpl_mpio = HDF5.H5FD_MPIO_COLLECTIVE,
            )
            close(f)

            MPI.Barrier(comm)

            # Open file for writing data
            f = h5open(path, "r+", comm, info)

            # Overwrite existing dataset if found
            # TODO: Return error on client side if we don't want to allow this
            if force_overwrite && haskey(f, group)
                delete_object(f[group])
            end

            # Create dataset
            whole_size = indexapply(+, size(part), offset, index = dim)
            whole_size = MPI.bcast(whole_size, nworkers - 1, comm) # Broadcast dataset size to all workers
            dset = create_dataset(f, group, eltype(part), (whole_size, whole_size))

            # Write out each partition
            setindex!(
                dset,
                part,
                [
                    # d == dim ? split_len(whole_size[dim], batch_idx, nbatches, comm) :
                    if d == dim
                        (offset+1):(offset+size(part, dim))
                    else
                        Colon()
                    end for d = 1:ndims(dset)
                ]...,
            )

            # Close file
            close(dset)
            close(f)
            MPI.Barrier(comm)
        else
            # TODO: See if we have missing `close`s or missing `fsync`s or extra `MPI.Barrier`s
            # fsync_file(p) =
            #     open(p) do f
            #         # TODO: Maybe use MPI I/O method for fsync instead
            #         ccall(:fsync, Cint, (Cint,), fd(f))
            #     end

            # Create the file if not yet created
            if batch_idx == 1
                f = h5open(
                    path,
                    "cw",
                    fapl_mpio = (comm, info),
                    dxpl_mpio = HDF5.H5FD_MPIO_COLLECTIVE,
                )
                close(f)
                MPI.Barrier(comm)
            end

            # TODO: Maybe use fsync to flush out after the file is closed
            # TODO: Use a separate file for intermediate datasets so that we
            # don't create a bunch of extra datasets everytime we try to write

            # TODO: Maybe pass in values for fapl_mpi and
            # dxpl_mpio = HDF5.H5FD_MPIO_COLLECTIVE,
            f = h5open(path, "r+", comm, info)
            # Allocate all datasets needed by gathering all sizes to the head
            # node and making calls from there
            part_lengths = MPI.Allgather(size(part, dim), comm)

            partdsets = [
                begin
                    idx = get_partition_idx(batch_idx, nbatches, worker_i)
                    group = group_prefix * "_part$idx" * "_dim=$dim"
                    # If there are multiple batches, each batch just gets written
                    # to its own group
                    dataspace_size =
                        indexapply(_ -> part_length, size(part), index = dim)
                    # TODO: Maybe pass in values for fapl_mpi and
                    # dxpl_mpio = HDF5.H5FD_MPIO_COLLECTIVE,
                    new_dset = create_dataset(
                        f,
                        group,
                        eltype(part),
                        (dataspace_size, dataspace_size),
                    )
                    new_dset
                end for (worker_i, part_length) in enumerate(part_lengths)
            ]

            # Wait for the head node to allocate all the datasets for this
            # batch index
            # TODO: Ensure that nothing doesn't cause bcast to just become a
            # no-op or something like that
            # MPI.bcast(nothing, 0, comm)
            # TODO: Try removing this barrier
            MPI.Barrier(comm)

            # Each worker then writes their partition to a separate dataset
            # in parallel
            partdsets[worker_idx][fill(Colon(), ndims(part))...] = part

            # Close (flush) all the intermediate datasets that we have created
            # TODO: Try removing this barrier
            MPI.Barrier(comm)
            for partdset in partdsets
                close(partdset)
            end
            # TODO: Try removing this barrier
            MPI.Barrier(comm)

            # Collect datasets from each batch and write into the final result dataset
            if batch_idx == nbatches
                # Get all intermediate datasets that have been written to by this worker
                partdsets = [
                    begin
                        # Determine what index partition this batch is
                        idx = get_partition_idx(batch_idx, nbatches, comm)

                        # Get the dataset
                        group = group_prefix * "_part$idx" * "_dim=$dim"
                        f[group]
                    end for batch_idx = 1:nbatches
                ]

                # Compute the size of all the batches on this worker
                # concatenated
                whole_batch_length = sum([size(partdset, dim) for partdset in partdsets])

                # Determine the offset into the resulting HDF5 dataset where this
                # worker should write
                offset = MPI.Exscan(whole_batch_length, +, comm)
                if worker_idx == 1
                    offset = 0
                end

                # Make the last worker create the dataset (since it can compute
                # the total dataset size using its offset)
                # NOTE: It's important that we use the last node since the
                # last node has the scan result
                whole_size =
                    indexapply(_ -> offset + whole_batch_length, size(part), index = dim)
                whole_size = MPI.bcast(whole_size, nworkers - 1, comm) # Broadcast dataset size to all workers
                # The permission used here is "r+" because we already
                # created the file on the head node
                # Delete the dataset if needed before we write the
                # dataset to which each batch will write its chunk
                if force_overwrite && haskey(f, group_prefix)
                    delete_object(f[group_prefix])
                end

                # If there are multiple batches, each batch just gets written
                # to its own group
                dset =
                    create_dataset(f, group_prefix, eltype(part), (whole_size, whole_size))

                # Wait until all workers have the file
                # TODO: Maybe use a broadcast so that each node is only blocked on
                # the last node which is where the file is creating
                # TODO: Try removing this barrier
                MPI.Barrier(comm)

                # Write out each batch
                batchoffset = offset
                for batch_i = 1:nbatches
                    partdset = partdsets[batch_i]

                    # Determine what index partition this batch is
                    idx = get_partition_idx(batch_i, nbatches, comm)

                    # Write
                    group = group_prefix * "_part$idx" * "_dim=$dim"
                    partdset_reading = partdset[fill(Colon(), ndims(dset))...]

                    # # println("In writing worker_idx=$worker_idx, batch_idx=$batch_idx/$nbatches: after reading batch $batch_i with available memory: $(format_available_memory())")
                    setindex!(
                        # We are writing to the whole dataset that was just
                        # created
                        dset,
                        # We are copying from the written HDF5 dataset for a
                        # particular batch
                        partdset_reading,
                        # We write to the appropriate split of the whole
                        # dataset
                        [
                            if d == dim
                                (batchoffset+1):batchoffset+size(partdset, dim)
                            else
                                Colon()
                            end
                            # split_len(whole_size[dim], batch_idx, nbatches, comm) : Colon()
                            for d = 1:ndims(dset)
                        ]...,
                    )
                    partdset_reading = nothing

                    # Update the offset of this batch
                    batchoffset += size(partdset, dim)
                    close(partdset)
                    partdset = nothing
                end
                close(dset)
                dset = nothing
                # fsync_file()

                # TODO: Delete data by keeping intermediates in separate file
                # Wait until all the data is written
                # TODO: Try removing this barrier
                MPI.Barrier(comm)
                # NOTE: Issue is that the barrier here doesn't ensure that all
                # processes have written in the previous step
                # TODO: Use a broadcast here

                # Then, delete all data for all groups on the head node
                for worker_i = 1:nworkers
                    for batch_i = 1:nbatches
                        idx = get_partition_idx(batch_i, nbatches, worker_i)
                        group = group_prefix * "_part$idx" * "_dim=$dim"
                        delete_object(f[group])
                    end
                end
                # TODO: Ensure that closing (flushing) HDF5 datasets
                # and files is sufficient. We might additionally have
                # to sync to ensure that the content is actually
                # written to disk or to S3
                # TODO: Use create_dataset passing in dtype and dimensions
                # fsync_file()

                # # TODO: Determine whether this is necessary. This barrier might
                # # be necessary to ensure that all groups are deleted before we
                # # continue.
                MPI.Barrier(comm)
            end
            close(f)
            f = nothing
            # TODO: Ensure that we are closing stuff everywhere before trying
            # to write
            MPI.Barrier(comm)
        end
    elseif get_partition_idx(batch_idx, nbatches, comm) == 1
        # This must be on disk; we don't support Julia serialized objects
        # as a remote location yet. We will need to first refactor locations
        # before we add support for that.
        if isa_gdf(part)
            part = nothing
        end
        serialize(path, part)
    end
end

mutable struct PartiallyMerged
    pieces::Vector{Any}
end

SplitBlock(
    src::Nothing,
    params::Dict{String,Any},
    batch_idx::Integer,
    nbatches::Integer,
    comm::MPI.Comm,
    loc_name::String,
    loc_params::Dict{String,Any},
) = nothing

SplitBlock(
    src::PartiallyMerged,
    params::Dict{String,Any},
    batch_idx::Integer,
    nbatches::Integer,
    comm::MPI.Comm,
    loc_name::String,
    loc_params::Dict{String,Any},
) = nothing

function SplitBlock(
    src::AbstractArray,
    params::Dict{String,Any},
    batch_idx::Integer,
    nbatches::Integer,
    comm::MPI.Comm,
    loc_name::String,
    loc_params::Dict{String,Any},
)
    split_on_executor(
        src,
        params["key"],
        batch_idx,
        nbatches,
        comm,
    )
end

function SplitBlock(
    src::T,
    params::Dict{String,Any},
    batch_idx::Integer,
    nbatches::Integer,
    comm::MPI.Comm,
    loc_name::String,
    loc_params::Dict{String,Any},
) where {T}
    split_on_executor(
        src,
        1,
        batch_idx,
        nbatches,
        comm,
    )
end

function SplitGroup(
    src,
    params,
    batch_idx::Integer,
    nbatches::Integer,
    comm::MPI.Comm,
    loc_name,
    loc_params;
    store_splitting_divisions = false
)
    # NOTE: The way we have `partial_merges` requires us to be splitting from
    # `nothing` and then merging back. If we are splitting from some value and
    # then expecting to merge back in some way then that won't work. If we are
    # splitting from a value we assume that we don't have to merge back either
    # because we split with a view (so the source was directly mutated) or we
    # didn't mutate this value at all. If we are doing in-place mutations where
    # we split from some value and then merge back up, then we might have to
    # add support for that. Right now, because of the way `SplitBlock`,
    # `SplitGroup`, and `Merge` are implemented, we unnecessarily concatenate
    # in the case where we are doing things like `setindex!` with a somewhat
    # faked mutation.
    if isnothing(src) || src isa PartiallyMerged
        # src is [] if we are partially merged (because as we iterate over
        # batches we take turns between splitting and merging)
        return nothing
    end

    partition_idx = get_partition_idx(batch_idx, nbatches, comm)
    npartitions = get_npartitions(nbatches, comm)

    # Ensure that this partition has a schema that is suitable for usage
    # here. We have to do this for `Shuffle` and `SplitGroup` (which is
    # used by `DistributeAndShuffle`)
    if isempty(src) || npartitions == 1
        # TODO: Ensure we can return here like this and don't need the above
        # (which is copied from `Shuffle`)
        return src
    end

    # Get divisions stored with src
    global splitting_divisions
    src_divisions, boundedlower, boundedupper = get!(splitting_divisions, src) do
        # This case lets us use `SplitGroup` in `DistributeAndShuffle`
        (params["divisions"], false, false)
    end
    divisions_by_partition = get_divisions(src_divisions, npartitions)

    # Get the divisions to apply
    key = params["key"]
    rev = params["rev"]
    if rev
        reverse!(divisions_by_partition)
    end

    # Create helper function for getting index of partition that owns a given
    # value
    partition_idx_getter(val) = get_partition_idx_from_divisions(
        val,
        divisions_by_partition,
        boundedlower = boundedlower,
        boundedupper = boundedupper,
    )

    # Apply divisions to get only the elements relevant to this worker
    res = if isa_df(src)
        # TODO: Do the groupby and filter on batch_idx == 1 and then share
        # among other batches
        filter(row -> partition_idx_getter(row[key]) == partition_idx, src)
    elseif isa_array(src)
        cat(
            filter(
                e -> partition_idx_getter(e) == partition_idx,
                eachslice(src, dims = key),
            )...;
            dims = key,
        )
    else
        throw(ArgumentError("Expected array or dataframe to distribute and shuffle"))
    end

    if store_splitting_divisions
        # The first and last partitions (used if this lacks a lower or upper bound)
        # must have actual division(s) associated with them. If there is no
        # partition that has divisions, then they will all be skipped and -1 will
        # be returned. So these indices are only used if there are nonempty
        # divisions.
        hasdivision = any(x->!isempty(x), divisions_by_partition)
        firstdivisionidx = findfirst(x->!isempty(x), divisions_by_partition)
        lastdivisionidx = findlast(x->!isempty(x), divisions_by_partition)

        # Store divisions
        global splitting_divisions
        splitting_divisions[res] = (
            divisions_by_partition[partition_idx],
            !hasdivision || boundedlower || partition_idx != firstdivisionidx,
            !hasdivision || boundedupper || partition_idx != lastdivisionidx,
        )
    end

    res
end

function Merge(
    src::Union{Nothing,PartiallyMerged},
    part,
    params::Dict{String,Any},
    batch_idx::Integer,
    nbatches::Integer,
    comm::MPI.Comm,
    loc_name::String,
    loc_params::Dict{String,Any},
)
    # TODO: Ensure we can merge grouped dataframes if computing them

    global splitting_divisions

    # TODO: To allow for mutation of a value, we may want to remove this
    # condition
    # We only need to concatenate partitions if the source is nothing.
    # Because if the source is something, then part must be a view into it
    # and no data movement is needed.

    key = params["key"]

    # Concatenate across batches
    if batch_idx == 1
        src = PartiallyMerged(Vector{Any}(undef, nbatches))
    end
    src.pieces[batch_idx] = part
    if batch_idx == nbatches
        delete!(splitting_divisions, part)

        # Concatenate across batches
        src = merge_on_executor(src.pieces...; key = key)

        # Concatenate across workers
        nworkers = get_nworkers(comm)
        if nworkers > 1
            src = Consolidate(src, params, Dict{String,Any}(), comm)
        end
    end

    src
end

Merge(
    src,
    part,
    params::Dict{String,Any},
    batch_idx::Integer,
    nbatches::Integer,
    comm::MPI.Comm,
    loc_name::String,
    loc_params::Dict{String,Any},
) = src

function CopyFrom(
    src,
    params,
    batch_idx::Integer,
    nbatches::Integer,
    comm::MPI.Comm,
    loc_name,
    loc_params,
)
    if loc_name == "Value"
        loc_params["value"]
    elseif loc_name == "Disk"
        params["key"] = 1
        res = ReadBlock(src, params, 1, 1, MPI.COMM_SELF, loc_name, loc_params)
        res
    elseif loc_name == "Remote"
        params["key"] = 1
        res = ReadBlock(src, params, 1, 1, MPI.COMM_SELF, loc_name, loc_params)
        res
    elseif loc_name == "Client"
        received = get_worker_idx(comm) == 1 ? receive_from_client(loc_params["value_id"]) : nothing
        # TODO: Make Replicated not necessarily require it to be replicated _everywhere_
        received = MPI.bcast(received, 0, comm)
        received
    elseif loc_name == "Memory"
        src
    end
end

function CopyTo(
    src,
    part,
    params,
    batch_idx::Integer,
    nbatches::Integer,
    comm::MPI.Comm,
    loc_name,
    loc_params,
)
    if loc_name == "Memory"
        src = part
    else
        # If we are copying to an external location we only want to do it on
        # the first worker since assuming that `on` is either `everywhere` or
        # `head`, so any batch on the first worker is guaranteed to have the
        # value that needs to be copied (either spilled to disk if None or
        # sent to remote storage).
        if loc_name == "Disk"
            # TODO: Don't rely on ReadBlock, Write in CopyFrom, CopyTo and
            # instead do something more elegant
            if get_partition_idx(batch_idx, nbatches, comm) == 1
                params["key"] = 1
                Write(src, part, params, 1, 1, MPI.COMM_SELF, loc_name, loc_params)
            end
            # TODO: Remove this barrier if not needed to ensure correctness
            MPI.Barrier(comm)
        elseif loc_name == "Remote"
            if get_partition_idx(batch_idx, nbatches, comm) == 1
                params["key"] = 1
                Write(src, part, params, 1, 1, MPI.COMM_SELF, loc_name, loc_params)
            end
            # TODO: Remove this barrier if not needed to ensure correctness
            MPI.Barrier(comm)
        elseif loc_name == "Client"
            # TODO: Ensure this only sends once
            if get_partition_idx(batch_idx, nbatches, comm) == 1
                send_to_client(loc_params["value_id"], part)
            end
            # TODO: Remove this barrier if not needed to ensure correctness
            MPI.Barrier(comm)
        else
            error("Unexpected location")
        end
    end
    part
end

function get_op!(params::Dict{String,Any})
    op = params["reducer"]
    if params["with_key"]
        key = params["key"]
        if !haskey(params, "reducer_with_key")
            op = op(key)
            reducer_with_key = Dict(key => op)
            params["reducer_with_key"] = reducer_with_key
        else
            reducer_with_key = params["reducer_with_key"]
            if !haskey(reducer_with_key, key)
                op = op(key)
                reducer_with_key[key] = op
            else
                op = reducer_with_key[key]
            end
        end
    end
    op
end

reduce_in_memory(src::Nothing, part::T, op::Function) where {T} = part
reduce_in_memory(src, part::T, op::Function) where {T} = op(src, part)

function ReduceAndCopyTo(
    src,
    part::T,
    params,
    batch_idx::Integer,
    nbatches::Integer,
    comm::MPI.Comm,
    loc_name,
    loc_params,
) where {T}
    # Merge reductions from batches
    op = get_op!(params)
    # TODO: Ensure that we handle reductions that can produce nothing
    src = reduce_in_memory(src, part, op)

    # Merge reductions across workers
    if batch_idx == nbatches
        src = Reduce(src, params, Dict{String,Any}(), comm)

        if loc_name != "Memory"
            # We use 1 here so that it is as if we are copying from the head
            # node
            CopyTo(src, src, params, 1, nbatches, comm, loc_name, loc_params)
        end
    end

    # TODO: Ensure we don't have issues where with batched execution we are
    # merging to the thing we are splitting from
    # NOTE: We are probably okay for now because we never split and then
    # re-merge new results to the same variable. We always merge to a new
    # variable. But in the future to be more robust, we may want to store
    # partial merges in a global `IdDict` and then only mutate `src` once we
    # are finished with the last batch and we know we won't be splitting
    # from the value again.
    src
end

ReduceWithKeyAndCopyTo = ReduceAndCopyTo

function Divide(
    src::Tuple,
    params::Dict{String,Any},
    batch_idx::Integer,
    nbatches::Integer,
    comm::MPI.Comm,
    loc_name::String,
    loc_params::Dict{String,Any},
)
    dim = params["key"]
    part = CopyFrom(src, params, batch_idx, nbatches, comm, loc_name, loc_params)
    newpartdim = length(split_len(part[dim], batch_idx, nbatches, comm))
    indexapply(_ -> newpartdim, part, index = dim)
end

function Divide(
    src,
    params::Dict{String,Any},
    batch_idx::Integer,
    nbatches::Integer,
    comm::MPI.Comm,
    loc_name::String,
    loc_params::Dict{String,Any},
)
    dim = params["key"]
    part = CopyFrom(src, params, batch_idx, nbatches, comm, loc_name, loc_params)
    length(split_len(part[dim], batch_idx, nbatches, comm))
end

#####################
# Casting functions #
#####################

function Reduce(
    part::T,
    src_params::Dict{String,Any},
    dst_params::Dict{String,Any},
    comm::MPI.Comm
) where {T}
    # Get operator for reduction
    op = get_op!(src_params)

    # TODO: Handle case where different processes have differently sized
    # sendbuf and where sendbuf is not isbitstype

    # Perform reduction
    part = MPI.Allreduce(
        part,
        # sendbuf,
        # (a, b) -> begin
        #     # tobuf(op(frombuf(kind, a), frombuf(kind, b)))[2]
        #     op(a, b)
        # end,
        op,
        comm,
    )
    part
end

ReduceWithKey = Reduce

# Grouped data frames can be block-partitioned but we will have to
# redo the groupby if we try to do any sort of merging/splitting on it.
Rebalance(
    part::Union{Nothing,GroupedDataFrame},
    src_params::Dict{String,Any},
    dst_params::Dict{String,Any},
    comm::MPI.Comm
) = nothing

function Rebalance(
    part::AbstractArray,
    src_params::Dict{String,Any},
    dst_params::Dict{String,Any},
    comm::MPI.Comm
)
    # Get the range owned by this worker
    dim = dst_params["key"]
    worker_idx, nworkers = get_worker_idx(comm), get_nworkers(comm)
    len = size(part, dim)
    scannedstartidx = MPI.Exscan(len, +, comm)
    startidx = worker_idx == 1 ? 1 : scannedstartidx + 1
    endidx = startidx + len - 1

    # Get functions for serializing/deserializing
    ser = serialize
    # TODO: Use JLD for ser/de for arrays
    # TODO: Ensure that we are properly handling intermediate arrays or
    # dataframes that are empty (especially because they may not have their
    # ndims or dtype or schema). We probably are because dataframes that are
    # empty should concatenate properly. We just need to be sure to not expect
    # every partition to know what its schema is. We can however expect each
    # partition of an array to know its ndims.
    de = x -> deserialize(IOBuffer(x))

    # NOTE: Below this is all common between Rebalance for DataFrame and AbstractArray

    # Construct buffer to send parts to all workers who own in this range
    nworkers = get_nworkers(comm)
    npartitions = nworkers
    whole_len = MPI.bcast(endidx, nworkers - 1, comm)
    io = IOBuffer()
    nbyteswritten = 0
    counts::Vector{Int64} = []
    for partition_idx = 1:npartitions
        # `split_len` gives us the range that this partition needs
        partitionrange = split_len(whole_len, partition_idx, npartitions)

        # Check if the range overlaps with the range owned by this worker
        rangesoverlap =
            max(startidx, partitionrange.start) <= min(endidx, partitionrange.stop)

        # If they do overlap, then serialize the overlapping slice
        ser(
            io,
            view(
                part,
                fill(:, dim - 1)...,
                if rangesoverlap
                    max(1, partitionrange.start - startidx + 1):min(
                        size(part, dim),
                        partitionrange.stop - startidx + 1,
                    )
                else
                    # Return zero length for this dimension
                    1:0
                end,
                fill(:, ndims(part) - dim)...,
            ),
        )

        # Add the count of the size of this chunk in bytes
        push!(counts, io.size - nbyteswritten)
        nbyteswritten = io.size

    end
    sendbuf = MPI.VBuffer(view(io.data, 1:nbyteswritten), counts)

    # Create buffer for receiving pieces
    # TODO: Refactor the intermediate part starting from there if we add
    # more cases for this function
    sizes = MPI.Alltoall(MPI.UBuffer(counts, 1), comm)
    recvbuf = MPI.VBuffer(similar(io.data, sum(sizes)), sizes)

    # Perform the shuffle
    MPI.Alltoallv!(sendbuf, recvbuf, comm)

    # Return the concatenated array
    things_to_concatenate = [
        de(view(recvbuf.data, displ+1:displ+count)) for
        (displ, count) in zip(recvbuf.displs, recvbuf.counts)
    ]
    res = merge_on_executor(
        things_to_concatenate...;
        key = dim,
    )
    res
end

function Rebalance(
    part::DataFrame,
    src_params::Dict{String,Any},
    dst_params::Dict{String,Any},
    comm::MPI.Comm
)
    # Get the range owned by this worker
    dim = 1
    worker_idx, nworkers = get_worker_idx(comm), get_nworkers(comm)
    len = size(part, dim)
    scannedstartidx = MPI.Exscan(len, +, comm)
    startidx = worker_idx == 1 ? 1 : scannedstartidx + 1
    endidx = startidx + len - 1

    # Get functions for serializing/deserializing
    ser = Arrow.write
    # TODO: Use JLD for ser/de for arrays
    # TODO: Ensure that we are properly handling intermediate arrays or
    # dataframes that are empty (especially because they may not have their
    # ndims or dtype or schema). We probably are because dataframes that are
    # empty should concatenate properly. We just need to be sure to not expect
    # every partition to know what its schema is. We can however expect each
    # partition of an array to know its ndims.
    de = x -> DataFrames.DataFrame(Arrow.Table(IOBuffer(x)))

    # Construct buffer to send parts to all workers who own in this range
    nworkers = get_nworkers(comm)
    npartitions = nworkers
    whole_len = MPI.bcast(endidx, nworkers - 1, comm)
    io = IOBuffer()
    nbyteswritten = 0
    counts::Vector{Int64} = []
    for partition_idx = 1:npartitions
        # `split_len` gives us the range that this partition needs
        partitionrange = split_len(whole_len, partition_idx, npartitions)

        # Check if the range overlaps with the range owned by this worker
        rangesoverlap =
            max(startidx, partitionrange.start) <= min(endidx, partitionrange.stop)

        # If they do overlap, then serialize the overlapping slice
        ser(
            io,
            view(
                part,
                fill(:, dim - 1)...,
                if rangesoverlap
                    max(1, partitionrange.start - startidx + 1):min(
                        size(part, dim),
                        partitionrange.stop - startidx + 1,
                    )
                else
                    # Return zero length for this dimension
                    1:0
                end,
                fill(:, ndims(part) - dim)...,
            ),
        )

        # Add the count of the size of this chunk in bytes
        push!(counts, io.size - nbyteswritten)
        nbyteswritten = io.size

    end
    sendbuf = MPI.VBuffer(view(io.data, 1:nbyteswritten), counts)

    # Create buffer for receiving pieces
    # TODO: Refactor the intermediate part starting from there if we add
    # more cases for this function
    sizes = MPI.Alltoall(MPI.UBuffer(counts, 1), comm)
    recvbuf = MPI.VBuffer(similar(io.data, sum(sizes)), sizes)

    # Perform the shuffle
    MPI.Alltoallv!(sendbuf, recvbuf, comm)

    # Return the concatenated array
    things_to_concatenate = [
        de(view(recvbuf.data, displ+1:displ+count)) for
        (displ, count) in zip(recvbuf.displs, recvbuf.counts)
    ]
    res = merge_on_executor(
        things_to_concatenate...;
        key = dim,
    )
    res
end

function Distribute(part::T, src_params::Dict{String,Any}, dst_params::Dict{String,Any}, comm::MPI.Comm) where {T}
    # TODO: Determine whether copy is needed
    SplitBlock(part, dst_params, 1, 1, comm, "Memory", Dict{String,Any}())
end

# If this is a grouped data frame or nothing (the result of merging
# a grouped data frame is nothing), we consolidate by simply returning
# nothing.
Consolidate(part::Union{Nothing, GroupedDataFrame}, src_params::Dict{String,Any}, dst_params::Dict{String,Any}, comm::MPI.Comm) = nothing

function Consolidate(part::AbstractArray, src_params::Dict{String,Any}, dst_params::Dict{String,Any}, comm::MPI.Comm)
    kind, sendbuf = tobuf(part)
    recvvbuf = buftovbuf(sendbuf, comm)
    # TODO: Maybe sometimes use gatherv if all sendbuf's are known to be equally sized

    MPI.Allgatherv!(sendbuf, recvvbuf, comm)
    part = merge_on_executor(
        kind,
        recvvbuf,
        get_nworkers(comm);
        key = src_params["key"],
    )
    part
end

function Consolidate(part::DataFrame, src_params::Dict{String,Any}, dst_params::Dict{String,Any}, comm::MPI.Comm)
    kind, sendbuf = tobuf(part)
    recvvbuf = buftovbuf(sendbuf, comm)
    # TODO: Maybe sometimes use gatherv if all sendbuf's are known to be equally sized

    MPI.Allgatherv!(sendbuf, recvvbuf, comm)
    part = merge_on_executor(
        kind,
        recvvbuf,
        get_nworkers(comm);
        key = 1,
    )
    part
end

DistributeAndShuffle(part::T, src_params::Dict{String,Any}, dst_params::Dict{String,Any}, comm::MPI.Comm) where {T} =
    SplitGroup(part, dst_params, 1, 1, comm, "Memory", Dict{String,Any}(), store_splitting_divisions = true)

function Shuffle(
    part::DataFrame,
    src_params::Dict{String,Any},
    dst_params::Dict{String,Any},
    comm::MPI.Comm;
    boundedlower = false,
    boundedupper = false,
    store_splitting_divisions = true
)
    # We don't have to worry about grouped data frames since they are always
    # block-partitioned.

    # Get the divisions to apply
    key = dst_params["key"]
    rev = dst_params["rev"]
    worker_idx, nworkers = get_worker_idx(comm), get_nworkers(comm)
    divisions_by_worker = if haskey(dst_params, "divisions_by_worker")
        dst_params["divisions_by_worker"] # list of min-max tuples
    else
        get_divisions(dst_params["divisions"], nworkers)
    end # list of min-max tuple lists
    if rev
        reverse!(divisions_by_worker)
    end

    # Perform shuffle
    partition_idx_getter(val) = get_partition_idx_from_divisions(
        val,
        divisions_by_worker,
        boundedlower = boundedlower,
        boundedupper = boundedupper,
    )
    res = begin
        gdf = if !isempty(part)
            # Compute the partition to send each row of the dataframe to
            DataFrames.transform!(part, key => ByRow(partition_idx_getter) => :banyan_shuffling_key)

            # Group the dataframe's rows by what partition to send to
            gdf = DataFrames.groupby(part, :banyan_shuffling_key, sort = true)
            gdf
        else
            nothing
        end

        # Create buffer for sending dataframe's rows to all the partitions
        io = IOBuffer()
        nbyteswritten = 0
        df_counts::Vector{Int64} = []
        for partition_idx = 1:nworkers
            Arrow.write(
                io,
                if !isnothing(gdf) && (banyan_shuffling_key = partition_idx,) in keys(gdf)
                    gdf[(banyan_shuffling_key = partition_idx,)]
                else
                    empty(part)
                end,
            )
            push!(df_counts, io.size - nbyteswritten)
            nbyteswritten = io.size
        end
        sendbuf = MPI.VBuffer(view(io.data, 1:nbyteswritten), df_counts)

        # Create buffer for receiving pieces
        sizes = MPI.Alltoall(MPI.UBuffer(df_counts, 1), comm)
        recvbuf = MPI.VBuffer(similar(io.data, sum(sizes)), sizes)

        # Perform the shuffle
        MPI.Alltoallv!(sendbuf, recvbuf, comm)

        # Return the concatenated dataframe
        things_to_concatenate = [
            DataFrames.DataFrame(
                Arrow.Table(IOBuffer(view(recvbuf.data, displ+1:displ+count))),
                copycols = false,
            ) for (displ, count) in zip(recvbuf.displs, recvbuf.counts)
        ]
        res = length(things_to_concatenate) == 1 ? things_to_concatenate[1] : vcat(things_to_concatenate...)
        if :banyan_shuffling_key in propertynames(res)
            DataFrames.select!(res, Not(:banyan_shuffling_key))
        end

        res
    end

    if store_splitting_divisions
        # The first and last partitions (used if this lacks a lower or upper bound)
        # must have actual division(s) associated with them. If there is no
        # partition that has divisions, then they will all be skipped and -1 will
        # be returned. So these indices are only used if there are nonempty
        # divisions.
        hasdivision = any(x->!isempty(x), divisions_by_worker)
        firstdivisionidx = findfirst(x->!isempty(x), divisions_by_worker)
        lastdivisionidx = findlast(x->!isempty(x), divisions_by_worker)

        # Store divisions
        global splitting_divisions
        splitting_divisions[res] =
            (divisions_by_worker[worker_idx], !hasdivision || worker_idx != firstdivisionidx, !hasdivision || worker_idx != lastdivisionidx)
    end

    res
end

function Shuffle(
    part::AbstractArray,
    src_params::Dict{String,Any},
    dst_params::Dict{String,Any},
    comm::MPI.Comm;
    boundedlower = false,
    boundedupper = false,
    store_splitting_divisions = true
)
    # We don't have to worry about grouped data frames since they are always
    # block-partitioned.

    # Get the divisions to apply
    key = dst_params["key"]
    rev = dst_params["rev"]
    worker_idx, nworkers = get_worker_idx(comm), get_nworkers(comm)
    divisions_by_worker = if haskey(dst_params, "divisions_by_worker")
        dst_params["divisions_by_worker"] # list of min-max tuples
    else
        get_divisions(dst_params["divisions"], nworkers)
    end # list of min-max tuple lists
    if rev
        reverse!(divisions_by_worker)
    end

    # Perform shuffle
    partition_idx_getter(val) = get_partition_idx_from_divisions(
        val,
        divisions_by_worker,
        boundedlower = boundedlower,
        boundedupper = boundedupper,
    )
    res = begin
        # Group the data along the splitting axis (specified by the "key"
        # parameter)
        partition_idx_to_e = [[] for partition_idx = 1:nworkers]
        for e in eachslice(part, dims = key)
            partition_idx = get_partition_idx_from_divisions(e, divisions_by_worker)
            push!(partition_idx_to_e[partition_idx], e)
        end

        # Construct buffer for sending data
        io = IOBuffer()
        nbyteswritten = 0
        a_counts::Vector{Int64} = []
        for partition_idx = 1:nworkers
            # TODO: If `isbitstype(eltype(e))`, we may want to pass it in
            # directly as an MPI buffer (if there is such a thing) instead of
            # serializing
            # Append to serialized buffer
            e = partition_idx_to_e[partition_idx]
            # NOTE: We ensure that we serialize something (even if its an
            # empty array) for each partition to ensure that we can
            # deserialize each item
            serialize(
                io,
                !isempty(e) ? cat(e...; dims = key) :
                view(part, [
                    if d == key
                        1:0
                    else
                        Colon()
                    end for d = 1:ndims(part)
                ]...),
            )

            # Add the count of the size of this chunk in bytes
            push!(a_counts, io.size - nbyteswritten)
            nbyteswritten = io.size
        end
        sendbuf = MPI.VBuffer(view(io.data, 1:nbyteswritten), a_counts)

        # Create buffer for receiving pieces
        # TODO: Refactor the intermediate part starting from there if we add
        # more cases for this function
        sizes = MPI.Alltoall(MPI.UBuffer(a_counts, 1), comm)
        recvbuf = MPI.VBuffer(similar(io.data, sum(sizes)), sizes)

        # Perform the shuffle
        MPI.Alltoallv!(sendbuf, recvbuf, comm)

        # Return the concatenated array
        things_to_concatenate = [
            deserialize(IOBuffer(view(recvbuf.data, displ+1:displ+count))) for
            (displ, count) in zip(recvbuf.displs, recvbuf.counts)
        ]
        if length(things_to_concatenate) == 1
            things_to_concatenate[1]
        else
            cat(
                things_to_concatenate...;
                dims = key,
            )
        end
    end

    if store_splitting_divisions
        # The first and last partitions (used if this lacks a lower or upper bound)
        # must have actual division(s) associated with them. If there is no
        # partition that has divisions, then they will all be skipped and -1 will
        # be returned. So these indices are only used if there are nonempty
        # divisions.
        hasdivision = any(x->!isempty(x), divisions_by_worker)
        firstdivisionidx = findfirst(x->!isempty(x), divisions_by_worker)
        lastdivisionidx = findlast(x->!isempty(x), divisions_by_worker)

        # Store divisions
        global splitting_divisions
        splitting_divisions[res] =
            (divisions_by_worker[worker_idx], !hasdivision || worker_idx != firstdivisionidx, !hasdivision || worker_idx != lastdivisionidx)
    end

    res
end

max_stars_repo_name: banyan-team/banyan-julia

max_issues_repo_name: banyan-team/banyan-julia

max_forks_repo_name: banyan-team/banyan-julia


text:
using WebVisualize
using Blink

win = Window()

body!(win, WebVisualize.circles())
# tools(win)

ext: jl


text:
function perfsummary(train, test, α::Real, name = nothing)
	τtrn = quantile(subsampleanomalous(train,α)[1], 1 - α) # quantile on positive  data
	τtst = quantile(subsampleanomalous(test,α)[1], 1 - α) # quantile on positive testing

	df = DataFrame(
		τtrn = τtrn,
		τtst = τtst,
		fprate_τtrn = mean(filterclass(test,1) .>= τtrn),
		tprate_τtrn = mean(filterclass(test,2) .>= τtrn),
		fprate_τtst = mean(filterclass(test,1) .>= τtst),
		tprate_τtst = mean(filterclass(test,2) .>= τtst),
		auc_trn = auc(roccurve(train[1], train[2] .- 1)...),
		auc_tst = auc(roccurve(test[1], test[2] .- 1)...)
		)
	if name != nothing
		df = hcat(DataFrame(name = name), df)
	end
	df
end

function perfsummary(train, test, name = nothing)
	df = DataFrame(
		fprate_trn = mean(filterclass(test,1) .>= 0),
		tprate_trn = mean(filterclass(test,2) .>= 0),
		fprate_tst = mean(filterclass(test,1) .>= 0),
		tprate_tst = mean(filterclass(test,2) .>= 0),
		auc_trn = auc(roccurve(train[1], train[2] .- 1)...),
		auc_tst = auc(roccurve(test[1], test[2] .- 1)...)
		)
	if name != nothing
		df = hcat(DataFrame(name = name), df)
	end
	df
end


text:
module DebbyPacker
using RudeOil

export Package, prepare, build, test, Publisher, publish
export GitSource, TarSource, SourceList, ModifySource, FileSource
import Base: download

abstract AbstractPackage
abstract AbstractSource
abstract AbstractDebianizer
abstract AbstractBuilder
abstract AbstractTest
abstract AbstractPublisher

type Package <: AbstractPackage
    name::String
    license::String
    version::VersionNumber

    homepage::String
    maintainer::String
    email::String
    build_depends::Vector
    section::String
    priority::String
    distribution::String
    binaries::Vector{Dict{Symbol, Any}}
    changes::Vector
    urgency::String

    source::AbstractSource
    debianize::AbstractDebianizer
    build::AbstractBuilder

    image::RudeOil.AbstractImage

    function Package(
            name::String, source::AbstractSource, debianize::AbstractDebianizer,
            build::AbstractBuilder;
            license::String="MIT", version::VersionNumber=v"0.0.0", homepage::String="",
            maintainer::String="", email::String="", build_depends::Vector=String[],
            section::String="devel", priority::String="optional",
            distribution::String="trusty", changes::Vector=String[], urgency::String="low",
            kwargs...)
      packages = ["dh-make", "build-essential", "devscripts", "cdbs"]
      base = RudeOil.image("packaging", "ubuntu:14.10", packages=packages)
      base.base.name = {
        "precise" => "12.04", "trusty" => "ubuntu:14.04", "utopic" => "ubuntu:14.10",
        "vivid" => "ubuntu:15.04"
      }[distribution]
      new(
        name, license, version, homepage, maintainer, email, build_depends,
        section, priority, distribution, get_binaries(name; kwargs...), changes, urgency,
        source, debianize, build, base
      )
    end
end

function prepare_source(pack::AbstractPackage, workdir::String="workspace";
    from_scratch=false, kwargs...)
  if from_scratch && isdir(build)
    rm(build; recursive=true)
  end
  download(pack, pack.source, workdir; kwargs...)
  tar(pack, pack.source, workdir; kwargs...)
end

function tar(package::AbstractPackage, ::AbstractSource, workdir::String; kwargs...)
  const name = package_name(package)
  const build = builddir(package, workdir)
  const source = sourcedir(package, workdir)
  const tar = tarfile(package, workdir)
  cd(dirname(source)) do
    run(`tar -czvf $tar --exclude .git* $name`)
  end
  cd(build) do
    run(`tar -xvf $tar`)
  end
end

function control(pack::AbstractPackage; kwargs...)
  control(pack, pack.debianize; kwargs...)
end
function install(pack::AbstractPackage; kwargs...)
  install(pack, pack.debianize; kwargs...)
end
function rules(pack::AbstractPackage; kwargs...)
  rules(pack, pack.debianize, pack.build; kwargs...)
end
function prepare(pack::AbstractPackage, workdir="workspace"; kwargs...)
  prepare_source(pack, workdir; kwargs...)
  debianize(pack, pack.debianize, workdir, kwargs...)
end

include("convenience.jl")
include("GitSource.jl")
include("TarSource.jl")
include("ModifySource.jl")
include("SourceList.jl")
include("debianize.jl")
include("builders.jl")
include("docker.jl")
include("publish.jl")

# Default debianizer is just a place holder
function Package(name::String, source::AbstractSource, builder=:CMake;
    build_flags=nothing, extra_build=[], kwargs...)
  if builder == :CMake || builder == :cmake
    buildme = CMakeBuilder(build_flags, extra_build)
  elseif builder == :Python || builder == :python
    buildme = PythonBuilder(extra_build)
  else
    buildme = MakeBuilder(extra_build)
  end
  Package(name, source, Debianizer(), buildme; kwargs...)
end

function build(machine::RudeOil.AbstractMachine, package::AbstractPackage, workdir="workspace";
    kwargs...)
  activate(machine) do vm
    container, cmds = make(package, workdir; kwargs...)
    for cmd in cmds
      vm |> container |> cmd |> run
    end
  end
end

end # module

max_stars_repo_name: UCL/DebbyPacker

max_issues_repo_name: UCL/DebbyPacker

max_forks_repo_name: UCL/DebbyPacker


text:
#script for solving a 1D BVP of the form
# -a0(x)*u''(x)+a1(x)*u(x)=f(x) for x ∈ (0,1)
# Here a0(x)=1, a1(x)=0, f(x)=4*π^2*sin(2*π*x)
# with Neumann-Robin boundary conditions u'(0) = ik, u'(1)-i*k*u(1)=0
include("../src/Solver1D.jl")

k=40
a0(x) = -1
a1(x) = k^2
f(x) = 0
exact_sol(x) = exp(im*k*x)

ptLeft = 0.
ptRight = 1.
uPtLeft = 0.
uPtRight = 1.
nrb = nrbline(ptLeft, ptRight)

bound_left = Boundary1D("Neumann", ptLeft, uPtLeft, -im*k)
α = -im*k
bound_right = Boundary1D("Robin", ptRight, uPtRight, α, 0)

numElem = 40
degP = 3
new_knots = collect(LinRange(ptLeft, ptRight, numElem+1)[2:end-1])
nrb = nrbdegelev(nrb, [degP-1])
nrb = nrbkntins(nrb, [new_knots])

prob_sin = Problem1D(f, a0, a1, [bound_left, bound_right], nrb)
IGAmesh = genMesh(prob_sin.domain)
plotBasisParam(IGAmesh)

gauss_rule = genGaussLegendre(degP+1)
stiff = assemble_stiff(IGAmesh, a0, gauss_rule)
mass = assemble_mass(IGAmesh, a1, gauss_rule)
rhs = assemble_rhs(IGAmesh, f, gauss_rule)
lhs, rhs = applyBCnurbs(prob_sin.boundary_cond, stiff, mass, rhs, nrb)
sol0 = lhs\rhs
plotSol(IGAmesh, real(sol0), "Helmholtz1D")
plotSolError(IGAmesh, real(sol0), exact_sol, "Helmholtz1D")


text:
struct TurnOutput
    next_state_1::DynamicState
    next_state_2::DynamicState
    odds::Float64
end

function play_turn(state::DynamicState, static_state::StaticState,
    decision::Tuple{UInt8, UInt8})

    # unpack state
    active_1, active_2 = get_active(state)
    fm_pending_1, fm_pending_2 = get_fast_moves_pending(state)
    cmp = get_cmp(state)
    chance = 0x00
    fm_dmg_1, fm_dmg_2 = get_fm_damage(state)

    # unpack teams
    switch_cooldown_1 = state[0x01].switch_cooldown
    switch_cooldown_2 = state[0x02].switch_cooldown
    a1, d1 = get_buffs(state[0x01])
    a2, d2 = get_buffs(state[0x02])
    shields_1 = get_shields(state[0x01])
    shields_2 = get_shields(state[0x02])

    # unpack pokemon
    hp_1_1 = get_hp(state[0x01][0x01])
    hp_1_2 = get_hp(state[0x01][0x02])
    hp_1_3 = get_hp(state[0x01][0x03])
    hp_2_1 = get_hp(state[0x02][0x01])
    hp_2_2 = get_hp(state[0x02][0x02])
    hp_2_3 = get_hp(state[0x02][0x03])
    energy_1_1 = get_energy(state[0x01][0x01])
    energy_1_2 = get_energy(state[0x01][0x02])
    energy_1_3 = get_energy(state[0x01][0x03])
    energy_2_1 = get_energy(state[0x02][0x01])
    energy_2_2 = get_energy(state[0x02][0x02])
    energy_2_3 = get_energy(state[0x02][0x03])

    # if in charged move state
    if !iszero(cmp)
        agent, o_agent = isodd(cmp) ? (0x01, 0x02) : (0x02, 0x01)

        attacker_energy = agent == 0x01 ? (active_1 == 0x01 ? energy_1_1 :
            active_1 == 0x02 ? energy_1_2 : energy_1_3) :
            (active_2 == 0x01 ? energy_2_1 : active_2 == 0x02 ?
            energy_2_2 : energy_2_3)
        defender_hp = agent == 0x01 ?
            (active_2 == 0x01 ? hp_2_1 : active_2 == 0x02 ? hp_2_2 : hp_2_3) :
            (active_1 == 0x01 ? hp_1_1 : active_1 == 0x02 ? hp_1_2 : hp_1_3)

        cm_output = evaluate_charged_move(static_state, cmp,
            decision[agent] - 0x06, decision[o_agent] == 0x01, active_1,
            active_2, a1, d1, a2, d2, attacker_energy, defender_hp,
            agent == 0x02 ? shields_1 : shields_2, fm_dmg_1, fm_dmg_2)

        chance = cm_output.chance
        a1 = cm_output.a1
        d1 = cm_output.d1
        a2 = cm_output.a2
        d2 = cm_output.d2
        fm_dmg_1 = cm_output.fm_dmg_1
        fm_dmg_2 = cm_output.fm_dmg_2
        cmp = cm_output.cmp

        if agent == 0x01
            shields_2 = cm_output.shields
            fm_pending_1 = 0x00
            fm_pending_2 = min(0x01, fm_pending_2)
            active_1 == 0x01 ? energy_1_1 = cm_output.attacker_energy :
                active_1 == 0x02 ? energy_1_2 = cm_output.attacker_energy :
                energy_1_3 = cm_output.attacker_energy
            active_2 == 0x01 ? hp_2_1 = cm_output.defender_hp :
                active_2 == 0x02 ? hp_2_2 = cm_output.defender_hp :
                hp_2_3 = cm_output.defender_hp
        else
            shields_1 = cm_output.shields
            fm_pending_1 = min(0x01, fm_pending_1)
            fm_pending_2 = 0x00
            active_1 == 0x01 ? hp_1_1 = cm_output.defender_hp :
                active_1 == 0x02 ? hp_1_2 = cm_output.defender_hp :
                hp_1_3 = cm_output.defender_hp
            active_2 == 0x01 ? energy_2_1 = cm_output.attacker_energy :
                active_2 == 0x02 ? energy_2_2 = cm_output.attacker_energy :
                energy_2_3 = cm_output.attacker_energy
        end
    else
        # evaluate fast moves
        defender_hp_1 = active_1 == 0x01 ? hp_1_1 :
            active_1 == 0x02 ? hp_1_2 : hp_1_3
        defender_hp_2 = active_2 == 0x01 ? hp_2_1 :
            active_2 == 0x02 ? hp_2_2 : hp_2_3
        if fm_pending_1 == 0x01 && !iszero(defender_hp_1)
            attacker_energy = active_1 == 0x01 ? energy_1_1 :
                active_1 == 0x02 ? energy_1_2 : energy_1_3
            fm_output = evaluate_fast_move(static_state, 0x01, active_1,
                attacker_energy, defender_hp_2, fm_dmg_1)
            active_1 == 0x01 ? energy_1_1 = fm_output.attacker_energy :
                active_1 == 0x02 ? energy_1_2 = fm_output.attacker_energy :
                energy_1_3 = fm_output.attacker_energy
            active_2 == 0x01 ? hp_2_1 = fm_output.defender_hp :
                active_2 == 0x02 ? hp_2_2 = fm_output.defender_hp :
                hp_2_3 = fm_output.defender_hp
        end
        if fm_pending_2 == 0x01 && !iszero(defender_hp_2)
            attacker_energy = active_2 == 0x01 ? energy_2_1 :
                active_2 == 0x02 ? energy_2_2 : energy_2_3
            fm_output = evaluate_fast_move(static_state, 0x02,
                active_2, attacker_energy, defender_hp_1, fm_dmg_2)
            active_2 == 0x01 ? energy_2_1 = fm_output.attacker_energy :
                active_2 == 0x02 ? energy_2_2 = fm_output.attacker_energy :
                energy_2_3 = fm_output.attacker_energy
            active_1 == 0x01 ? hp_1_1 = fm_output.defender_hp :
                active_1 == 0x02 ? hp_1_2 = fm_output.defender_hp :
                hp_1_3 = fm_output.defender_hp
        end

        # step timers
        fm_cooldown_1 = decision[1] == 0x03 ?
            get_cooldown(static_state[0x01][active_1].fast_move) : 0x00
        fm_cooldown_2 = decision[2] == 0x03 ?
            get_cooldown(static_state[0x02][active_2].fast_move) : 0x00
        step_timers_output = step_timers(fm_cooldown_1, fm_cooldown_2,
            fm_pending_1, fm_pending_2, switch_cooldown_1, switch_cooldown_2)
        fm_pending_1 = step_timers_output.fm_pending_1
        fm_pending_2 = step_timers_output.fm_pending_2
        switch_cooldown_1 = step_timers_output.switch_cooldown_1
        switch_cooldown_2 = step_timers_output.switch_cooldown_2

        # evaluate switches
        if decision[1] == 0x05 || decision[1] == 0x06
            active_hp_1 = active_1 == 0x01 ? hp_1_1 :
                active_1 == 0x02 ? hp_1_2 : hp_1_3
            switch_output = evaluate_switch(static_state, 0x01,
                decision[1] - 0x04, iszero(active_hp_1) ? 0x18 : 0x00,
                active_1, active_2, switch_cooldown_1, switch_cooldown_2)

            active_1 = switch_output.active
            switch_cooldown_1 = switch_output.switch_cooldown_1
            switch_cooldown_2 = switch_output.switch_cooldown_2
            fm_dmg_1 = switch_output.fm_dmg_1
            fm_dmg_2 = switch_output.fm_dmg_2

            fm_pending_1 = 0x00
            a1, d1 = 0x04, 0x04
        end
        if decision[2] == 0x05 || decision[2] == 0x06
            active_hp_2 = active_2 == 0x01 ? hp_2_1 :
                active_2 == 0x02 ? hp_2_2 : hp_2_3
            switch_output = evaluate_switch(static_state, 0x02,
                decision[2] - 0x04, iszero(active_hp_2) ? 0x18 : 0x00,
                active_1, active_2, switch_cooldown_1, switch_cooldown_2)

            active_2 = switch_output.active
            switch_cooldown_1 = switch_output.switch_cooldown_1
            switch_cooldown_2 = switch_output.switch_cooldown_2
            fm_dmg_1 = switch_output.fm_dmg_1
            fm_dmg_2 = switch_output.fm_dmg_2

            fm_pending_2 = 0x00
            a2, d2 = 0x04, 0x04
        end

        active_hp_1 = active_1 == 0x01 ? hp_1_1 :
            active_1 == 0x02 ? hp_1_2 : hp_1_3
        active_hp_2 = active_2 == 0x01 ? hp_2_1 :
            active_2 == 0x02 ? hp_2_2 : hp_2_3
        if !iszero(active_hp_1) && !iszero(active_hp_2)
            if decision[1] == 0x04
                if decision[2] == 0x04
                    atk_cmp = Base.cmp(
                        static_state[0x01][active_1].stats.attack,
                        static_state[0x02][active_2].stats.attack
                    )
                    atk_cmp == 1 ? cmp = 0x03 : atk_cmp == -1 ?
                        cmp = 0x04 : chance = 0x05
                else
                    cmp = 0x01
                end
            elseif decision[2] == 0x04
                cmp = 0x02
            end
        end
    end

    # check chance and repack
    if iszero(chance)
        # no chance to evaluate
        next_state_1 = DynamicState(
            DynamicTeam(
                DynamicPokemon(hp_1_1, energy_1_1),
                DynamicPokemon(hp_1_2, energy_1_2),
                DynamicPokemon(hp_1_3, energy_1_3),
                switch_cooldown_1, a1, d1, shields_1
            ), DynamicTeam(
                DynamicPokemon(hp_2_1, energy_2_1),
                DynamicPokemon(hp_2_2, energy_2_2),
                DynamicPokemon(hp_2_3, energy_2_3),
                switch_cooldown_2, a2, d2, shields_2
            ), active_1, active_2, fm_pending_1, fm_pending_2, cmp, 0x00,
                fm_dmg_1, fm_dmg_2
        )
        next_state_2 = next_state_1
        odds = 1.0
    elseif chance == 0x05
        # decide randomly who gets cmp
        next_state_1 = DynamicState(
            DynamicTeam(
                DynamicPokemon(hp_1_1, energy_1_1),
                DynamicPokemon(hp_1_2, energy_1_2),
                DynamicPokemon(hp_1_3, energy_1_3),
                switch_cooldown_1, a1, d1, shields_1
            ), DynamicTeam(
                DynamicPokemon(hp_2_1, energy_2_1),
                DynamicPokemon(hp_2_2, energy_2_2),
                DynamicPokemon(hp_2_3, energy_2_3),
                switch_cooldown_2, a2, d2, shields_2
            ), active_1, active_2, fm_pending_1, fm_pending_2, 0x03, 0x00,
                fm_dmg_1, fm_dmg_2
        )
        next_state_2 = DynamicState(
            DynamicTeam(
                DynamicPokemon(hp_1_1, energy_1_1),
                DynamicPokemon(hp_1_2, energy_1_2),
                DynamicPokemon(hp_1_3, energy_1_3),
                switch_cooldown_1, a1, d1, shields_1
            ), DynamicTeam(
                DynamicPokemon(hp_2_1, energy_2_1),
                DynamicPokemon(hp_2_2, energy_2_2),
                DynamicPokemon(hp_2_3, energy_2_3),
                switch_cooldown_2, a2, d2, shields_2
            ), active_1, active_2, fm_pending_1, fm_pending_2, 0x04, 0x00,
                fm_dmg_1, fm_dmg_2
        )
        odds = 0.5
    else
        # apply or don't apply a buff from a charged move
        next_state_2 = DynamicState(
            DynamicTeam(
                DynamicPokemon(hp_1_1, energy_1_1),
                DynamicPokemon(hp_1_2, energy_1_2),
                DynamicPokemon(hp_1_3, energy_1_3),
                switch_cooldown_1, a1, d1, shields_1
            ), DynamicTeam(
                DynamicPokemon(hp_2_1, energy_2_1),
                DynamicPokemon(hp_2_2, energy_2_2),
                DynamicPokemon(hp_2_3, energy_2_3),
                switch_cooldown_2, a2, d2, shields_2
            ), active_1, active_2, fm_pending_1, fm_pending_2, cmp, 0x00,
                fm_dmg_1, fm_dmg_2
        )

        move = isodd(chance) ?
            static_state[chance < 0x03 ? 0x01 : 0x02][active_1].charged_move_1 :
            static_state[chance < 0x03 ? 0x01 : 0x02][active_2].charged_move_2
        buff_output = apply_buff(a1, d1, a2, d2, move)
            a1 = buff_output.a1
            d1 = buff_output.d1
            a2 = buff_output.a2
            d2 = buff_output.d2

        next_state_1 =  DynamicState(
            DynamicTeam(
                DynamicPokemon(hp_1_1, energy_1_1),
                DynamicPokemon(hp_1_2, energy_1_2),
                DynamicPokemon(hp_1_3, energy_1_3),
                switch_cooldown_1, a1, d1, shields_1
            ), DynamicTeam(
                DynamicPokemon(hp_2_1, energy_2_1),
                DynamicPokemon(hp_2_2, energy_2_2),
                DynamicPokemon(hp_2_3, energy_2_3),
                switch_cooldown_2, a2, d2, shields_2
            ), active_1, active_2, fm_pending_1, fm_pending_2, cmp, 0x00,
                fm_dmg_1, fm_dmg_2
        )
        odds = get_buff_chance(move)
    end

    return TurnOutput(next_state_1, next_state_2, odds)
end

function select_next_state(turn_output::TurnOutput)
    return turn_output.odds == 1.0 ? turn_output.next_state_1 :
        turn_output.odds == 0.5 ?
            rand((turn_output.next_state_1, turn_output.next_state_2)) :
        turn_output.odds == 0.1 ? (rand(rb_rng, 1:10) == 1 ?
            turn_output.next_state_1 : turn_output.next_state_2) :
        turn_output.odds == 0.125 ? rand(rb_rng,
            (turn_output.next_state_1, turn_output.next_state_2,
            turn_output.next_state_2, turn_output.next_state_2,
            turn_output.next_state_2, turn_output.next_state_2,
            turn_output.next_state_2, turn_output.next_state_2)) :
        turn_output.odds == 0.2 ? (rand(rb_rng, 1:5) == 1 ?
            turn_output.next_state_1 : turn_output.next_state_2) :
        (rand(rb_rng, 1:10) < 4 ?
            turn_output.next_state_1 : turn_output.next_state_2)
end

"""
    play_battle(state, static_state)

Play through one battle, starting from the inputted state with random,
equally weighted decisions.
"""
function play_battle(state::DynamicState, static_state::StaticState;
  allow_nothing::Bool = false, allow_overfarming::Bool = false)
    while true
        d1, d2 = get_possible_decisions(state, static_state,
            allow_nothing = allow_nothing,
            allow_overfarming = allow_overfarming)
        (iszero(d1) || iszero(d2)) &&
            return battle_score(state, static_state)
        state = select_next_state(play_turn(
            state, static_state, select_random_decision(d1, d2)))
    end
end

function play_battle(teams::NTuple{6, Union{String, Int}};
    league::String = "great", cup::String = "all",
    allow_nothing::Bool = false, allow_overfarming::Bool = false)
    static_state = StaticState(teams, league = league, cup = cup)
    dynamic_state = DynamicState(static_state)
    return play_battle(dynamic_state, static_state,
        allow_nothing = allow_nothing, allow_overfarming = allow_overfarming)
end

"""
    battle_scores(state, static_state, N)

Play through N battles, starting from the inputted state with random,
equally weighted decisions.
"""
function battle_scores(starting_state::DynamicState, static_state::StaticState,
    N::Int64; allow_nothing::Bool = false, allow_overfarming::Bool = false)
    return map(x -> play_battle(starting_state, static_state,
        allow_nothing = allow_nothing, allow_overfarming = allow_overfarming),
        1:N)
end

function battle_scores(teams::NTuple{6, Union{String, Int}}, N;
    league::String = "great", cup::String = "all",
    allow_nothing::Bool = false, allow_overfarming::Bool = false)
    static_state = StaticState(teams, league = league, cup = cup)
    dynamic_state = DynamicState(static_state)

    return map(x -> play_battle(dynamic_state, static_state,
        allow_nothing = allow_nothing, allow_overfarming = allow_overfarming),
        1:N)
end


text:
# This file is a part of Julia. License is MIT: https://julialang.org/license

function simd_loop_example_from_manual(x, y, z)
    s = zero(eltype(z))
    n = min(length(x),length(y),length(z))
    @simd for i in 1:n
        @inbounds begin
            z[i] = x[i]-y[i]
            s += z[i]*z[i]
        end
    end
    s
end

function simd_loop_with_multiple_reductions(x, y, z)
    # Use non-zero initial value to make sure reduction values include it.
    (s,t) = (one(eltype(x)),one(eltype(y)))
    @simd for i in 1:length(z)
        @inbounds begin
            s += x[i]
            t += 2*y[i]
            s += z[i]   # Two reductions go into s
        end
    end
    (s,t)
end

for T in [Int32,Int64,Float32,Float64]
    # Try various lengths to make sure "remainder loop" works
    for n in [0,1,2,3,4,255,256,257]
        local n, a, b, c, s, t
        # Dataset chosen so that results will be exact with only 24 bits of mantissa
        a = convert(Array{T},[2*j+1 for j in 1:n])
        b = convert(Array{T},[3*j+2 for j in 1:n])
        c = convert(Array{T},[5*j+3 for j in 1:n])
        s = simd_loop_example_from_manual(a,b,c)

        @test a==[2*j+1 for j in 1:n]
        @test b==[3*j+2 for j in 1:n]
        @test c==[-j-1 for j in 1:n]
        @test s==sum(c.*c)
        (s,t) = simd_loop_with_multiple_reductions(a,b,c)
        @test s==sum(a)+sum(c)+1
        @test t==2*sum(b)+1
    end
end

# Test that scope rules match regular for
let j=4
    # Use existing local variable.
    @simd for j=1:0 end
    @test j==4
    @simd for j=1:3 end
    @test j==3

    # Use global variable
    global simd_glob = 4
    @simd for simd_glob=1:0 end
    @test simd_glob==4
    @simd for simd_glob=1:3 end
    @test simd_glob==3

    # Index that is local to loop
    @simd for simd_loop_local=1:0 end
    simd_loop_local_present = true
    try
        simd_loop_local += 1
    catch
        simd_loop_local_present = false
    end
    @test !simd_loop_local_present
end

import Base.SimdLoop.SimdError

# Test that @simd rejects inner loop body with invalid control flow statements
# issue #8613
macro test_throws(ty, ex)
    return quote
        Test.@test_throws $(esc(ty)) try
            $(esc(ex))
        catch err
            @test err isa LoadError
            @test err.file === $(string(__source__.file))
            @test err.line === $(__source__.line + 1)
            rethrow(err.error)
        end
    end
end

@test_throws SimdError("break is not allowed inside a @simd loop body") @macroexpand begin
    @simd for x = 1:10
        x == 1 && break
    end
end

@test_throws SimdError("continue is not allowed inside a @simd loop body") @macroexpand begin
    @simd for x = 1:10
        x < 5 && continue
    end
end

@test_throws SimdError("@goto is not allowed inside a @simd loop body") @macroexpand begin
    @simd for x = 1:10
        x == 1 || @goto exit_loop
    end
    @label exit_loop
end

# @simd with cartesian iteration
function simd_cartesian_range!(indices, crng)
    @simd for I in crng
        push!(indices, I)
    end
    indices
end

crng = CartesianIndices(2:4, 0:1, 1:1, 3:5)
indices = simd_cartesian_range!(Vector{eltype(crng)}(), crng)
@test indices == vec(collect(crng))

crng = CartesianIndices(-1:1, 1:3)
indices = simd_cartesian_range!(Vector{eltype(crng)}(), crng)
@test indices == vec(collect(crng))

crng = CartesianIndices(-1:-1, 1:3)
indices = simd_cartesian_range!(Vector{eltype(crng)}(), crng)
@test indices == vec(collect(crng))

crng = CartesianIndices(2:4)
indices = simd_cartesian_range!(Vector{eltype(crng)}(), crng)
@test indices == collect(crng)

crng = CartesianIndices()
indices = simd_cartesian_range!(Vector{eltype(crng)}(), crng)
@test indices == vec(collect(crng))

# @simd with array as "range"
# issue #13869
function simd_sum_over_array(a)
    s = zero(eltype(a))
    @inbounds @simd for x in a
        s += x
    end
    s
end
@test 2001000 == simd_sum_over_array(collect(1:2000))
@test 2001000 == simd_sum_over_array(Float32[i+j*500 for i=1:500, j=0:3])

max_stars_repo_name: YingboMa/julia

max_issues_repo_name: YingboMa/julia

max_forks_repo_name: YingboMa/julia


text:
# Includes basic risk measures and statistics without optimization

"""
    Expectation(states,prob)

Computes the expectation for the random variable ``Y`` with values given by `states` under
the probability measure ``Q`` given by `prob`:

```math
\\mathbb{E} Y = states \\cdot prob
```
"""
function Expectation(states::Vector{Float64},prob::Vector{Float64})
if sum(prob) == 0.0
    prob = ones(length(states))./ length(states)
end
return(dot(states,prob))
end

function Expectation(d::Distribution)
    return(mean(d))
end

"""
    entropic(states, prob,theta::Float32)

implements the entropic risk measure defined as
```math
\\rho_\\theta(Y) = \\theta \\cdot \\log \\mathbb{E} e^{\\frac{Y}{\\theta}},
```
where ``\\theta`` is greater 0 and ``Y`` is the random variable defined by `states` and `prob`.
"""
function entropic(states::Vector{Float64}, prob::Vector{Float64},theta::Float32)
    if theta <= 0
        return(nothing)
    end
	if sum(prob) == 0
		prob = ones(length(states))./ length(states)
        # improve this with isprob and onto simplex?
	end

    w = states./theta
    offset = maximum(w)
    we = exp.(w .- offset)
    s = dot(prob, we)
    w = log(s) + offset
    tmp = theta*w

	return(tmp)
end

function entropic(d::Distribution,t::Float32)
    if t <= 0
        return(nothing)
    end
	tmp = mgf(d,1.0 ./t);
    return ( t* log(tmp))
end

#########################################################
#Mean Semi-deviation
#########################################################
"""
    mSD(states,prob,beta::Float32,p::Float32)

implements the mean semi-deviation of order `p` which is a coherent risk measure defined by
```math
mSD_\\beta^p (Y) = \\mathbb{E} Y  + \\beta \\lvert \\left( Y - \\mathbb{E}Y \\right)_+ \\rvert_p,
```
for the random variable ``Y`` defined by `states` and `prob`.
"""
function mSD(states::Vector{Float64},prob::Vector{Float64},beta::Float32,p::Float32)
	if sum(prob) == 0
		prob = ones(length(states))./ length(states)
	end

	if beta == 0.0
		mSD = dot(prob, states)
	else
		tmp = max.(states .- dot(prob,states),0)
		mSD = dot(prob,states) + beta*pnorm(p,prob,tmp)
	end
	return(mSD)
end

"""
    meanVariance(states, prob,c::Float32)

implements the mean Variance risk measure defined by
```math
\\rho_c(Y) := \\mathbb{E}Y + c \\cdot \\mathbb{E} \\left( Y- \\mathbb{E}Y)^2 \\right),
```
where ``c >0`` and ``Y`` is the random variable defined by `states` and `prob`.
"""
function meanVariance(states::Vector{Float64}, prob::Vector{Float64},c::Float32)
    if c<0
        return(nothing)
    end
	if sum(prob) == 0
		prob = ones(length(states))./ length(states)
        # improve this with isprob and onto simplex?
	end

    tmp1 = dot(states,prob)
    tmp2 = dot(prob, (states .- tmp1).^2);
	return(tmp1 + c*tmp2)
end

"""
    meanDeviation(states, prob,c::Float32,p::Float32)

implements the mean deviation risk measure of order ``p\\geq 1`` defined by
```math
\\rho_c^p(Y) := \\mathbb{E}Y + c \\cdot \\lVert \\left( Y- \\mathbb{E}Y \\right)^2 \\rVert_p,
```
where ``c >0`` and ``Y`` is the random variable defined by `states` and `prob`.
"""

function meanDeviation(states::Vector{Float64}, prob::Vector{Float64},c::Float32,p::Float32)
    if c<0
        return(nothing)
    end
    if p<1.0
        return(nothing)
    end
	if sum(prob) == 0
		prob = ones(length(states))./ length(states)
        # improve this with isprob and onto simplex?
	end

    tmp1 = dot(states,prob)
    tmp2 = pnorm(p,prob, states .- tmp1);
	return(tmp1 + c*tmp2)
end

"""
    meanSemiVariance(states, prob,c::Float32,t::Float32)

implements the mean upper-semi variance risk measure from a target ``t`` defined by
```math
\\rho_{c,t}(Y) = \\mathbb{E}Y + c \\cdot \\mathbb{E} \\left( Y - t \\right)^2_+ ,
```
where ``c >0`` and ``Y`` is the random variable defined by `states` and `prob`.
"""
function meanSemiVariance(states::Vector{Float64}, prob::Vector{Float64},c::Float32,t::Float32)
    if c<0
        return(nothing)
    end
	if sum(prob) == 0
		prob = ones(length(states))./ length(states)
        # improve this with isprob and onto simplex?
	end

    tmp1 = dot(states,prob)
    tmp2 = pnorm(2,prob, max.(states .- t,0))^2;
	return(tmp1 + c*tmp2)
end

"""
    meanSemiDevi(states, prob,c::Float32,target::Float32,p::Float32)

implements the mean upper-semi variance risk measure of order ``p \\geq 1`` from a target ``t`` defined by
```math
\\rho_{c,t}^p(Y) = \\mathbb{E}Y + c \\cdot \\lVert \\left( Y - t \\right)_+ \\rVert_p ,
```
where ``c >0`` and ``Y`` is the random variable defined by `states` and `prob`.
"""
function meanSemiDevi(states::Vector{Float64}, prob::Vector{Float64},c::Float32,target::Float32,p::Float32)
    if c<0
        return(nothing)
    end
    if p<1.0
        return(nothing)
    end
	if sum(prob) == 0
		prob = ones(length(states))./ length(states)
        # improve this with isprob and onto simplex?
	end

    tmp1 = dot(states,prob)
    tmp2 = pnorm(p,prob, max.(states .- target,0));
	return(tmp1 + c*tmp2)
end


text:
module ApproxManifoldProducts

using Reexport
@reexport using KernelDensityEstimate
@reexport using TransformUtils

using DocStringExtensions

using NLsolve
using Optim
using CoordinateTransformations
using Requires

import Base: *
# import KernelDensityEstimate: kde!

const KDE = KernelDensityEstimate
const TUs = TransformUtils
const CTs = CoordinateTransformations

# TODO temporary for initial version of on-manifold products
KDE.setForceEvalDirect!(true)

export
  get2DLambda,
  get2DMu,
  get2DMuMin,
  resid2DLinear,
  ensurePeriodicDomains!,
  solveresid2DLinear!,
  solveresid2DLinear,
  ManifoldBelief,
  MB,
  *,

  # APi and util functions
  buildHybridManifoldCallbacks,
  getKDEManifoldBandwidths,
  manifoldProduct,
  manikde!,

  # Supported manifolds
  Manifold,
  EuclideanManifold,
  Euclid2,
  Circular

include("Interface.jl")
include("CommonUtils.jl")
include("Euclidean.jl")
include("CircularUtils.jl")
include("Circular.jl")
include("KernelHilbertEmbeddings.jl")

include("Legacy.jl")
include("API.jl")

function __init__()
  @require Gadfly="c91e804a-d5a3-530f-b6f0-dfbca275c004" begin
    @require Colors="5ae59095-9a9b-59fe-a467-6f913c188581" include("plotting/CircularPlotting.jl")
  end
end

end


text:
# module Bukdu

import Mustache

function template(path::String, options::Dict)::String
    contents = chomp(readstring(path))
    Mustache.render(contents, options)
end

function render{AV<:ApplicationView}(::Type{AV}, args...; kw...)::Conn
    filtering(render, AV, args...) do
        options = Dict(kw)
        path = options[:path]
        Conn(200, Dict("Content-Type"=>"text/html"), template(path, options))
    end
end


text:
function _write_task(gwt::GWTask, ex::Expr;
        desc = "", lang = "",
        src = _expr_src(ex),
        readme = false,
        vtime = _GW_DFLT_TASK_VTIME
    )

    # dir
    _mkdir(task_dir(gwt))

    # task.toml
    extime = time() + vtime
    _write_task_toml!(gwt; extime, desc)

    # taskdat.jls
    _write_task_dat!(gwt; ex, src, desc, lang)

    # readme.md
    readme && _do_readme!(gwt)

    # runme.jl
    _write_runme(gwt)

    return gwt
end

max_stars_repo_name: josePereiro/GitWorkers

max_issues_repo_name: josePereiro/GitWorkers

max_forks_repo_name: josePereiro/GitWorkers


text:
using Pkg

# Reset Julia to use Conda environment.
ENV["PYTHON"] = ""

# Re-build PyCall.jl with updated Python version.
Pkg.build("PyCall")