"""
Voxel corner and edge indexing conventions

        Z
        |

        5------5------6              Extra edges not drawn
       /|            /|              -----------
      8 |           6 |              - face diagonals
     /  9          /  10                - 13: 1 to 3
    8------7------7   |                 - 14: 1 to 8
    |   |         |   |                 - 15: 1 to 6
    |   1------1--|---2  -- Y           - 16: 5 to 7
    12 /          11 /                  - 17: 2 to 7
    | 4           | 2                   - 18: 4 to 7
    |/            |/                 - body diagonal
    4------3------3                     - 19: 1 to 7

  /
 X
"""
# this gets vectorized so we want to ensure it is the
# same type as out vertex
@inline function voxCrnrPos(::Type{PT}) where {PT}
    (PT(0, 0, 0),
    PT(0, 1, 0),
    PT(1, 1, 0),
    PT(1, 0, 0),
    PT(0, 0, 1),
    PT(0, 1, 1),
    PT(1, 1, 1),
    PT(1, 0, 1))
end

const voxCrnrPosInt = voxCrnrPos(SVector{3,UInt8})

# the voxel IDs at either end of the tetrahedra edges, by edge ID
const voxEdgeCrnrs = ((0x01, 0x02),
                      (0x02, 0x03),
                      (0x04, 0x03),
                      (0x01, 0x04),
                      (0x05, 0x06),
                      (0x06, 0x07),
                      (0x08, 0x07),
                      (0x05, 0x08),
                      (0x01, 0x05),
                      (0x02, 0x06),
                      (0x03, 0x07),
                      (0x04, 0x08),
                      (0x01, 0x03),
                      (0x01, 0x08),
                      (0x01, 0x06),
                      (0x05, 0x07),
                      (0x02, 0x07),
                      (0x04, 0x07),
                      (0x01, 0x07))

# direction codes:
# 0 => +x, 1 => +y, 2 => +z,
# 3 => +xy, 4 => +xz, 5 => +yz, 6 => +xyz
const voxEdgeDir = (0x01, 0x00, 0x01, 0x00, 0x01, 0x00, 0x01, 0x00, 0x02, 0x02, 0x02, 0x02, 0x03, 0x04, 0x05, 0x03, 0x04, 0x05, 0x06)

# For a pair of corner IDs, the edge ID joining them
# 0 denotes a pair with no edge
const voxEdgeIx = ((0x00, 0x01, 0x0d, 0x04, 0x09, 0x0f, 0x13, 0x0e),
                   (0x01, 0x00, 0x02, 0x00, 0x00, 0x0a, 0x11, 0x00),
                   (0x0d, 0x02, 0x00, 0x03, 0x00, 0x00, 0x0b, 0x00),
                   (0x04, 0x00, 0x03, 0x00, 0x00, 0x00, 0x12, 0x0c),
                   (0x09, 0x00, 0x00, 0x00, 0x00, 0x05, 0x10, 0x08),
                   (0x0f, 0x0a, 0x00, 0x00, 0x05, 0x00, 0x06, 0x00),
                   (0x13, 0x11, 0x0b, 0x12, 0x10, 0x06, 0x00, 0x07),
                   (0x0e, 0x00, 0x00, 0x0c, 0x08, 0x00, 0x07, 0x00))

# voxel corners that comprise each of the six tetrahedra
const subTets = ((0x01, 0x03, 0x02, 0x07),
                 (0x01, 0x08, 0x04, 0x07),
                 (0x01, 0x04, 0x03, 0x07),
                 (0x01, 0x02, 0x06, 0x07),
                 (0x01, 0x05, 0x08, 0x07),
                 (0x01, 0x06, 0x05, 0x07))
# tetrahedron corners for each edge (indices 1-4)
const tetEdgeCrnrs = ((0x01, 0x02),
                      (0x02, 0x03),
                      (0x01, 0x03),
                      (0x01, 0x04),
                      (0x02, 0x04),
                      (0x03, 0x04))

# triangle cases for a given tetrahedron edge code
const tetTri = ((0x01, 0x03, 0x04, 0x00, 0x00, 0x00),
                (0x01, 0x05, 0x02, 0x00, 0x00, 0x00),
                (0x03, 0x05, 0x02, 0x03, 0x04, 0x05),
                (0x02, 0x06, 0x03, 0x00, 0x00, 0x00),
                (0x01, 0x06, 0x04, 0x01, 0x02, 0x06),
                (0x01, 0x05, 0x06, 0x01, 0x06, 0x03),
                (0x04, 0x05, 0x06, 0x00, 0x00, 0x00),
                (0x04, 0x06, 0x05, 0x00, 0x00, 0x00),
                (0x01, 0x06, 0x05, 0x01, 0x03, 0x06),
                (0x01, 0x04, 0x06, 0x01, 0x06, 0x02),
                (0x02, 0x03, 0x06, 0x00, 0x00, 0x00),
                (0x03, 0x02, 0x05, 0x03, 0x05, 0x04),
                (0x01, 0x02, 0x05, 0x00, 0x00, 0x00),
                (0x01, 0x04, 0x03, 0x00, 0x00, 0x00))


export
    Event,
    EventList,
    AllEventList,
    getevent,
    getlastevent,
    pushevent!,
    samplelocalpath,
    quadpathpoly,
    pathmean,
    tmax

"""
    AllowedVarType

Syntactic sugar to make it explicit what variable types are handled in the
local BPS (i.e, the type of the variables on each node of the factor graph).
See also, Event, EventList.
"""
const AllowedVarType = Union{Real, Vector{<:Real}}

"""
    Event{T <: AllowedVarType}

Encapsulation of a local event: a tuple made out of (position, velocity, time)
for a node in the local BPS algorithm.
See also: `EventList`.
"""
struct Event{T <: AllowedVarType}
    x::T
    v::T
    t::Real
end

"""
    EventList{T <: AllowedVarType}

Storage for all events that happened for a variable in the local BPS.
Each variable is stored explicitly (as opposed to storing a list of Event)
so that the times can be traversed efficiently. The xs, vs must all have the
same type within the list.
See also: `Event`, `AllEventList`.
"""
mutable struct EventList{T <: AllowedVarType}
    xs::Vector{T}     # coordinates
    vs::Vector{T}     # associated velocities
    ts::Vector{Real}  # associated times
end

function EventList(T::Type)
    @assert T <: AllowedVarType
    EventList(Vector{T}(undef, 0), Vector{T}(undef, 0), Vector{Real}(undef, 0))
end

"""
    AllEventList

Storage for all the `EventList` objects corresponding to a factor graph in a
local BPS setting. The element `evl` can be accessed by the variable index so
`a.evl[k]` will return the eventlist associated with variable `k`. Note that
the different `EventList` may have different types (in particular, different
dimensions). See also: `Event`, `EventList`.
"""
mutable struct AllEventList
    evl::Vector{EventList}
    types::Vector{DataType}
    function AllEventList(types::Vector{DataType})
        new([EventList(types[i]) for i ∈ 1:length(types)], types)
    end
end
AllEventList(T::DataType, nvars::Int) = AllEventList([T for i ∈ 1:nvars])

tmax(aev::AllEventList) = maximum(aev.evl[i].ts[end] for i ∈ 1:length(aev.evl))

"""
    getevent(evl, evidx)

Retrieve a specfic event (at index `evidx`) in an `EventList` object `evl` and
return it as an `Event` object.
See also: `Event`, `EventList`.
"""
function getevent(evl::EventList, evidx::Int)
    Event(evl.xs[evidx], evl.vs[evidx], evl.ts[evidx])
end

"""
    getlastevent(evl)

Retrieve the last `Event` in an `EventList` object `evl`.
See also: `getevent`, `EventList`.
"""
function getlastevent(evl::EventList)
    Event(evl.xs[end], evl.vs[end], evl.ts[end])
end

"""
    pushevent!(evl, ev)

Append an `Event` object `ev` to an `EventList` object `evl`. The `Event`
object is unpacked and pushed to the vectors encapsulated by `evl`.
See also: `EventList`.
"""
function pushevent!(evl::EventList, ev::Event)
    push!(evl.xs, ev.x)
    push!(evl.vs, ev.v)
    push!(evl.ts, ev.t)
    evl
end

"""
    getlocalsegment(evl, j)

Retrieve the `j`th event in an `EventList` object `evl` and return as a tuple
that event and the next event. If there is no next event (end of event list),
an "infinity event" is returned (in order to allow extrapolation).
See also: `samplelocalpath`.
"""
function getlocalsegment(evl::EventList, j::Int)
    if j == length(evl.ts)
        return (getevent(evl,j), Event(NaN, NaN, Inf))
    end
    (getevent(evl,j), getevent(evl,j+1))
end

"""
    samplelocalpath(evl, t)

Sample x_k(t) given the corresponding `EventList` object `evl`. The times `t`
must be a list of sorted times (or a single positive time).
See also: `getlocalsegment`.
"""
function samplelocalpath(evl::EventList, t::AllowedTimeType)
    @assert issorted(t) "list of times should be sorted"
    ti  = t[1]
    @assert ti >= 0.0 "the first time should be greater than 0.0"
    # find the starting local segment
    j = searchsortedlast(evl.ts, ti)
    # event and next event (describes the "segment")
    eva, evb = getlocalsegment(evl,j)
    # allocation of samples
    samples = Vector{typeof(eva.x)}(undef, length(t))
    # recursively work with local segments
    i=1
    while i<= length(t)
        ti = t[i]
        if ti <= evb.t
            # stay in the current segment
            samples[i] = eva.x + (ti - eva.t) * eva.v
            i += 1
        else
            # get the next segment
            j += searchsortedlast(evl.ts[j+1:end], ti)
            eva, evb = getlocalsegment(evl, j)
        end
    end
    samples
end
samplelocalpath(evl::EventList, t::Real) = samplelocalpath(evl, [t])[1]

"""
    quadpathpoly(evl, pol, T)

Integrate a polynomial on the elements along a path for a single node.
So if we are sampling in Rd and have phi(X)=Poly(xk), then we can integrate the
polynomial along the path. Note that we can't do the same thing if phi(X) mixes
2 or more nodes. See also syntactic sugar to perform the same operation on all
nodes.
"""
function quadpathpoly(evl::EventList, pol::Poly, T::Real)
    dim  = length(evl.xs[1])
    res  = zeros(dim)
    nseg = searchsortedlast(evl.ts,T)
    for segidx = 1:nseg-1
        # for each segment in the path
        eva, evb = getlocalsegment(evl, segidx)
        # segment time, initial point and velocity
        tau = evb.t - eva.t
        xa  = eva.x
        v   = (evb.x - xa) / tau
        # integrate along that segment
        for d = 1:dim
            res[d] += polyint(polyval(pol,Poly([xa[d],v[d]])))(tau)
        end
    end
    # last segment
    eva = getevent(evl, nseg)
    # characteristics
    tau = T - eva.t
    for d = 1:dim
        res[d] += polyint(polyval(pol,Poly([eva.x[d],eva.v[d]])))(tau)
    end
    return res / T
end
function quadpathpoly(aev::AllEventList, pol::Poly, T::Real)
    res = Vector{Vector{<:Real}}(undef, length(aev.evl))
    for (d, evl) in enumerate(aev.evl)
        res[d] = quadpathpoly(evl, pol, T)
    end
    return res
 end
pathmean(evl::EventList, T::Real) = quadpathpoly(evl, Poly([0.0,1.0]), T)
pathmean(aev::AllEventList, T::Real) = quadpathpoly(aev, Poly([0.0,1.0]), T)
pathmean(aev::AllEventList) = pathmean(aev, tmax(aev))


"""
Type for storing data acquired with a steady state response (SSR) experimental paradigm.

In addition to the functions available for all EEG types,
the SSR type supports:

* `modulationrate()`

The following standard names are used when saving data to the processing dictionary.

* `Name`: The identifier for the participant
* `Side`: Side of stimulation
* `Carrier_Frequency`: Carrier frequency of the stimulus
* `Amplitude`: Amplitude of the stimulus
* `epochs`: The epochs extracted from the recording
* `sweeps`: The extracted sweeps from the recording

#### Example

Put an example here

```julia
s = read_SSR("filename")
s.modulationrate = 40.0391u"Hz"
s = rereference(s, "Cz")
```

"""
mutable struct SSR <: EEG
    data::Array
    sensors::Array{Sensor}
    triggers::Dict
    system_codes::Dict
    samplingrate::typeof(1.0u"Hz")
    modulationrate::typeof(1.0u"Hz")
    reference_channel::Array{AbstractString,1}
    file_path::AbstractString
    file_name::AbstractString
    processing::Dict
    header::Dict
end

#######################################
#
# SSR info
#
#######################################

"""
Return the modulation rate of a steady state type.
If no type is provided, the modulation rate is returned as a floating point.

#### Example

Return the modulation rate of a recording

```julia
s = read_SSR(filename)
modulationrate(s)
```
"""
modulationrate(t, s::SSR) = convert(t, ustrip(s.modulationrate))
modulationrate(s::SSR) = modulationrate(AbstractFloat, s)

#######################################
#
# Show
#
#######################################

import Base.show
function Base.show(io::IO, a::SSR)
    time_length = round.(size(a.data, 1) / samplingrate(a) / 60)
    println(
        io,
        "SSR measurement of $time_length mins with $(size(a.data,2)) channels sampled at $(a.samplingrate)",
    )
    println(io, "  Modulation frequency: $(a.modulationrate )")

    if haskey(a.processing, "Amplitude")
        println(io, "  Stimulation amplitude: $(a.processing["Amplitude"]) dB")
    end
    if haskey(a.processing, "Name")
        println(io, "  Participant name: $(a.processing["Name"] )")
    end
    if haskey(a.processing, "Side")
        println(io, "  Stimulation side: $(a.processing["Side"] )")
    end
    if haskey(a.processing, "Carrier_Frequency")
        println(io, "  Carrier frequency: $(a.processing["Carrier_Frequency"] ) Hz")
    end

end

#######################################
#
# Helper functions
#
#######################################

function assr_frequency(
    rounded_freq::Number;
    stimulation_samplingrate::Number = 32000,
    stimulation_frames_per_epoch::Number = 32768,
)

    round.(rounded_freq / (stimulation_samplingrate / stimulation_frames_per_epoch)) *
    stimulation_samplingrate / stimulation_frames_per_epoch
end

function assr_frequency(rounded_freq::AbstractVector)

    [assr_frequency(f) for f in rounded_freq]
end


include("DGFVModel_kernels.jl")
struct DGFVModel{BL, G, FVR, NFND, NFD, GNF, AS, DS, HDS, D, DD, MD, GF, TF} <:
       SpaceDiscretization
    balance_law::BL
    grid::G
    fv_reconstruction::FVR
    numerical_flux_first_order::NFND
    numerical_flux_second_order::NFD
    numerical_flux_gradient::GNF
    state_auxiliary::AS
    state_gradient_flux::DS
    states_higher_order::HDS
    direction::D
    diffusion_direction::DD
    modeldata::MD
    gradient_filter::GF
    tendency_filter::TF
    check_for_crashes::Bool
end

function DGFVModel(
    balance_law,
    grid,
    fv_reconstruction,
    numerical_flux_first_order,
    numerical_flux_second_order,
    numerical_flux_gradient;
    fill_nan = false,
    check_for_crashes = false,
    state_auxiliary = create_state(
        balance_law,
        grid,
        Auxiliary(),
        fill_nan = fill_nan,
    ),
    state_gradient_flux = create_state(balance_law, grid, GradientFlux()),
    states_higher_order = (
        create_state(balance_law, grid, GradientLaplacian()),
        create_state(balance_law, grid, Hyperdiffusive()),
    ),
    direction = EveryDirection(),
    diffusion_direction = direction,
    modeldata = nothing,
    gradient_filter = nothing,
    tendency_filter = nothing,
)
    # Make sure we are FVM in the vertical
    @assert polynomialorders(grid)[end] == 0
    @assert isstacked(grid.topology)
    state_auxiliary =
        init_state(state_auxiliary, balance_law, grid, direction, Auxiliary())
    DGFVModel(
        balance_law,
        grid,
        fv_reconstruction,
        numerical_flux_first_order,
        numerical_flux_second_order,
        numerical_flux_gradient,
        state_auxiliary,
        state_gradient_flux,
        states_higher_order,
        direction,
        diffusion_direction,
        modeldata,
        gradient_filter,
        tendency_filter,
        check_for_crashes,
    )
end

"""
    (dgfvm::DGFVModel)(tendency, state_prognostic, _, t, α, β)

Uses spectral element discontinuous Galerkin in the horizontal and finite volume
in the vertical to compute the tendency.

Computes the tendency terms compatible with `IncrementODEProblem`

    tendency .= α .* dQdt(state_prognostic, p, t) .+ β .* tendency

The 4-argument form will just compute

    tendency .= dQdt(state_prognostic, p, t)
"""
function (dgfvm::DGFVModel)(tendency, state_prognostic, _, t, α, β)
    device = array_device(state_prognostic)

    FT = eltype(state_prognostic)
    num_state_prognostic = number_states(dgfvm.balance_law, Prognostic())
    num_state_gradient_flux = number_states(dgfvm.balance_law, GradientFlux())
    @assert 0 == number_states(dgfvm.balance_law, Hyperdiffusive())
    num_state_tendency = size(tendency, 2)

    if num_state_prognostic < num_state_tendency && β != 1
        # If we don't operate on the full state, then we need to scale here instead of volume_tendency!
        tendency .*= β
        β = β != 0 # if β==0 then we can avoid the memory load in volume_tendency!
    end

    communicate =
        !(
            isstacked(dgfvm.grid.topology) &&
            typeof(dgfvm.direction) <: VerticalDirection
        )

    update_auxiliary_state!(
        dgfvm,
        dgfvm.balance_law,
        state_prognostic,
        t,
        dgfvm.grid.topology.realelems,
    )

    exchange_state_gradient_flux = NoneEvent()
    exchange_state_prognostic = NoneEvent()

    comp_stream = Event(device)

    if communicate
        exchange_state_prognostic = MPIStateArrays.begin_ghost_exchange!(
            state_prognostic;
            dependencies = comp_stream,
        )
    end

    if num_state_gradient_flux > 0
        ########################
        # Gradient Computation #
        ########################

        comp_stream = launch_volume_gradients!(
            dgfvm,
            state_prognostic,
            t;
            dependencies = comp_stream,
        )

        comp_stream = launch_interface_gradients!(
            dgfvm,
            state_prognostic,
            t;
            surface = :interior,
            dependencies = comp_stream,
        )

        if communicate
            exchange_state_prognostic = MPIStateArrays.end_ghost_exchange!(
                state_prognostic;
                dependencies = exchange_state_prognostic,
                check_for_crashes = dgfvm.check_for_crashes,
            )

            # Update_aux may start asynchronous work on the compute device and
            # we synchronize those here through a device event.
            checked_wait(
                device,
                exchange_state_prognostic,
                nothing,
                dgfvm.check_for_crashes,
            )
            update_auxiliary_state!(
                dgfvm,
                dgfvm.balance_law,
                state_prognostic,
                t,
                dgfvm.grid.topology.ghostelems,
            )
            exchange_state_prognostic = Event(device)
        end

        comp_stream = launch_interface_gradients!(
            dgfvm,
            state_prognostic,
            t;
            surface = :exterior,
            dependencies = (comp_stream, exchange_state_prognostic),
        )

        if dgfvm.gradient_filter !== nothing
            comp_stream = Filters.apply_async!(
                dgfvm.state_gradient_flux,
                1:num_state_gradient_flux,
                dgfvm.grid,
                dgfvm.gradient_filter;
                dependencies = comp_stream,
            )
        end

        if communicate
            if num_state_gradient_flux > 0
                exchange_state_gradient_flux =
                    MPIStateArrays.begin_ghost_exchange!(
                        dgfvm.state_gradient_flux,
                        dependencies = comp_stream,
                    )
            end
        end

        if num_state_gradient_flux > 0
            # Update_aux_diffusive may start asynchronous work on the compute device
            # and we synchronize those here through a device event.
            checked_wait(device, comp_stream, nothing, dgfvm.check_for_crashes)
            update_auxiliary_state_gradient!(
                dgfvm,
                dgfvm.balance_law,
                state_prognostic,
                t,
                dgfvm.grid.topology.realelems,
            )
            comp_stream = Event(device)
        end
    end

    ###################
    # RHS Computation #
    ###################
    comp_stream = launch_volume_tendency!(
        dgfvm,
        tendency,
        state_prognostic,
        t,
        α,
        β;
        dependencies = comp_stream,
    )

    comp_stream = launch_interface_tendency!(
        dgfvm,
        tendency,
        state_prognostic,
        t,
        α,
        β;
        surface = :interior,
        dependencies = comp_stream,
    )

    if communicate
        if num_state_gradient_flux > 0
            exchange_state_gradient_flux = MPIStateArrays.end_ghost_exchange!(
                dgfvm.state_gradient_flux;
                dependencies = exchange_state_gradient_flux,
                check_for_crashes = dgfvm.check_for_crashes,
            )

            # Update_aux_diffusive may start asynchronous work on the
            # compute device and we synchronize those here through a device
            # event.
            checked_wait(
                device,
                exchange_state_gradient_flux,
                nothing,
                dgfvm.check_for_crashes,
            )
            update_auxiliary_state_gradient!(
                dgfvm,
                dgfvm.balance_law,
                state_prognostic,
                t,
                dgfvm.grid.topology.ghostelems,
            )
            exchange_state_gradient_flux = Event(device)
        else
            exchange_state_prognostic = MPIStateArrays.end_ghost_exchange!(
                state_prognostic;
                dependencies = exchange_state_prognostic,
                check_for_crashes = dgfvm.check_for_crashes,
            )

            # Update_aux may start asynchronous work on the compute device and
            # we synchronize those here through a device event.
            checked_wait(
                device,
                exchange_state_prognostic,
                nothing,
                dgfvm.check_for_crashes,
            )
            update_auxiliary_state!(
                dgfvm,
                dgfvm.balance_law,
                state_prognostic,
                t,
                dgfvm.grid.topology.ghostelems,
            )
            exchange_state_prognostic = Event(device)
        end
    end

    comp_stream = launch_interface_tendency!(
        dgfvm,
        tendency,
        state_prognostic,
        t,
        α,
        β;
        surface = :exterior,
        dependencies = (
            comp_stream,
            exchange_state_prognostic,
            exchange_state_gradient_flux,
            # XXX: This is disabled until FVM with hyperdiffusion for DG is implemented: exchange_Qhypervisc_grad,
        ),
    )

    if dgfvm.tendency_filter !== nothing
        comp_stream = Filters.apply_async!(
            tendency,
            1:num_state_tendency,
            dgfvm.grid,
            dgfvm.tendency_filter;
            dependencies = comp_stream,
        )
    end

    # The synchronization here through a device event prevents CuArray based and
    # other default stream kernels from launching before the work scheduled in
    # this function is finished.
    checked_wait(device, comp_stream, nothing, dgfvm.check_for_crashes)
end

function fvm_balance!(
    balance_func,
    m::BalanceLaw,
    state_auxiliary::MPIStateArray,
    grid,
)
    device = array_device(state_auxiliary)

    dim = dimensionality(grid)
    N = polynomialorders(grid)
    Nq = N .+ 1
    Nqj = dim == 2 ? 1 : Nq[2]

    topology = grid.topology
    elems = topology.elems
    nelem = length(elems)
    nvertelem = topology.stacksize
    horzelems = fld1(first(elems), nvertelem):fld1(last(elems), nvertelem)

    event = Event(device)
    event = kernel_fvm_balance!(device, (Nq[1], Nqj))(
        balance_func,
        m,
        Val(nvertelem),
        state_auxiliary.data,
        grid.vgeo,
        horzelems;
        ndrange = (length(horzelems) * Nq[1], Nqj),
        dependencies = (event,),
    )
    wait(device, event)
end


module Inference

using ..Core
using ..Utilities
using DynamicPPL: Metadata, _tail, VarInfo, TypedVarInfo,
    islinked, invlink!, getlogp, tonamedtuple, VarName, getsym, vectorize,
    settrans!, _getvns, getdist, CACHERESET,
    Model, Sampler, SampleFromPrior, SampleFromUniform,
    Selector, DefaultContext, PriorContext,
    LikelihoodContext, MiniBatchContext, set_flag!, unset_flag!, NamedDist, NoDist,
    getspace, inspace
using Distributions, Libtask, Bijectors
using DistributionsAD: VectorOfMultivariate
using LinearAlgebra
using ..Turing: PROGRESS, Turing
using StatsFuns: logsumexp
using Random: AbstractRNG
using DynamicPPL
using AbstractMCMC: AbstractModel, AbstractSampler
using DocStringExtensions: TYPEDEF, TYPEDFIELDS

import AbstractMCMC
import AdvancedHMC; const AHMC = AdvancedHMC
import AdvancedMH; const AMH = AdvancedMH
import AdvancedPS
import BangBang
import ..Core: getchunksize, getADbackend
import DynamicPPL: get_matching_type,
    VarName, _getranges, _getindex, getval, _getvns
import EllipticalSliceSampling
import Random
import MCMCChains
import StatsBase: predict

export  InferenceAlgorithm,
        Hamiltonian,
        GibbsComponent,
        StaticHamiltonian,
        AdaptiveHamiltonian,
        SampleFromUniform,
        SampleFromPrior,
        MH,
        ESS,
        Emcee,
        Gibbs,      # classic sampling
        GibbsConditional,
        HMC,
        SGLD,
        SGHMC,
        HMCDA,
        NUTS,       # Hamiltonian-like sampling
        DynamicNUTS,
        IS,
        SMC,
        CSMC,
        PG,
        Prior,
        assume,
        dot_assume,
        observe,
        dot_observe,
        resume,
        predict,
        isgibbscomponent

#######################
# Sampler abstraction #
#######################
abstract type AbstractAdapter end
abstract type InferenceAlgorithm end
abstract type ParticleInference <: InferenceAlgorithm end
abstract type Hamiltonian{AD} <: InferenceAlgorithm end
abstract type StaticHamiltonian{AD} <: Hamiltonian{AD} end
abstract type AdaptiveHamiltonian{AD} <: Hamiltonian{AD} end

getchunksize(::Type{<:Hamiltonian{AD}}) where AD = getchunksize(AD)
getADbackend(::Hamiltonian{AD}) where AD = AD()

# Algorithm for sampling from the prior
struct Prior <: InferenceAlgorithm end

"""
    mh_accept(logp_current::Real, logp_proposal::Real, log_proposal_ratio::Real)

Decide if a proposal ``x'`` with log probability ``\\log p(x') = logp_proposal`` and
log proposal ratio ``\\log k(x', x) - \\log k(x, x') = log_proposal_ratio`` in a
Metropolis-Hastings algorithm with Markov kernel ``k(x_t, x_{t+1})`` and current state
``x`` with log probability ``\\log p(x) = logp_current`` is accepted by evaluating the
Metropolis-Hastings acceptance criterion
```math
\\log U \\leq \\log p(x') - \\log p(x) + \\log k(x', x) - \\log k(x, x')
```
for a uniform random number ``U \\in [0, 1)``.
"""
function mh_accept(logp_current::Real, logp_proposal::Real, log_proposal_ratio::Real)
    # replacing log(rand()) with -randexp() yields test errors
    return log(rand()) + logp_current ≤ logp_proposal + log_proposal_ratio
end

######################
# Default Transition #
######################

struct Transition{T, F<:AbstractFloat}
    θ  :: T
    lp :: F
end

function Transition(vi::AbstractVarInfo, nt::NamedTuple=NamedTuple())
    theta = merge(tonamedtuple(vi), nt)
    lp = getlogp(vi)
    return Transition{typeof(theta), typeof(lp)}(theta, lp)
end

metadata(t::Transition) = (lp = t.lp,)

DynamicPPL.getlogp(t::Transition) = t.lp

# Metadata of VarInfo object
metadata(vi::AbstractVarInfo) = (lp = getlogp(vi),)

#########################################
# Default definitions for the interface #
#########################################

function AbstractMCMC.sample(
    model::AbstractModel,
    alg::InferenceAlgorithm,
    N::Integer;
    kwargs...
)
    return AbstractMCMC.sample(Random.GLOBAL_RNG, model, alg, N; kwargs...)
end

function AbstractMCMC.sample(
    rng::AbstractRNG,
    model::AbstractModel,
    alg::InferenceAlgorithm,
    N::Integer;
    kwargs...
)
    return AbstractMCMC.sample(rng, model, Sampler(alg, model), N; kwargs...)
end

function AbstractMCMC.sample(
    rng::AbstractRNG,
    model::AbstractModel,
    sampler::Sampler{<:InferenceAlgorithm},
    N::Integer;
    chain_type=MCMCChains.Chains,
    resume_from=nothing,
    progress=PROGRESS[],
    kwargs...
)
    if resume_from === nothing
        return AbstractMCMC.mcmcsample(rng, model, sampler, N;
                                       chain_type=chain_type, progress=progress, kwargs...)
    else
        return resume(resume_from, N; chain_type=chain_type, progress=progress, kwargs...)
    end
end

function AbstractMCMC.sample(
    rng::AbstractRNG,
    model::AbstractModel,
    alg::Prior,
    N::Integer;
    chain_type=MCMCChains.Chains,
    resume_from=nothing,
    progress=PROGRESS[],
    kwargs...
)
    if resume_from === nothing
        return AbstractMCMC.mcmcsample(rng, model, SampleFromPrior(), N;
                                       chain_type=chain_type, progress=progress, kwargs...)
    else
        return resume(resume_from, N; chain_type=chain_type, progress=progress, kwargs...)
    end
end

function AbstractMCMC.sample(
    model::AbstractModel,
    alg::InferenceAlgorithm,
    parallel::AbstractMCMC.AbstractMCMCParallel,
    N::Integer,
    n_chains::Integer;
    kwargs...
)
    return AbstractMCMC.sample(Random.GLOBAL_RNG, model, alg, parallel, N, n_chains;
                               kwargs...)
end

function AbstractMCMC.sample(
    rng::AbstractRNG,
    model::AbstractModel,
    alg::InferenceAlgorithm,
    parallel::AbstractMCMC.AbstractMCMCParallel,
    N::Integer,
    n_chains::Integer;
    kwargs...
)
    return AbstractMCMC.sample(rng, model, Sampler(alg, model), parallel, N, n_chains;
                               kwargs...)
end

function AbstractMCMC.sample(
    rng::AbstractRNG,
    model::AbstractModel,
    sampler::Sampler{<:InferenceAlgorithm},
    parallel::AbstractMCMC.AbstractMCMCParallel,
    N::Integer,
    n_chains::Integer;
    chain_type=MCMCChains.Chains,
    progress=PROGRESS[],
    kwargs...
)
    return AbstractMCMC.mcmcsample(rng, model, sampler, parallel, N, n_chains;
                                   chain_type=chain_type, progress=progress, kwargs...)
end

function AbstractMCMC.sample(
    rng::AbstractRNG,
    model::AbstractModel,
    alg::Prior,
    parallel::AbstractMCMC.AbstractMCMCParallel,
    N::Integer,
    n_chains::Integer;
    chain_type=MCMCChains.Chains,
    progress=PROGRESS[],
    kwargs...
)
    return AbstractMCMC.sample(rng, model, SampleFromPrior(), parallel, N, n_chains;
                               chain_type=chain_type, progress=progress, kwargs...)
end

##########################
# Chain making utilities #
##########################

"""
    getparams(t)

Return a named tuple of parameters.
"""
getparams(t) = t.θ
getparams(t::VarInfo) = tonamedtuple(TypedVarInfo(t))

function _params_to_array(ts::Vector)
    names = Vector{Symbol}()
    # Extract the parameter names and values from each transition.
    dicts = map(ts) do t
        nms, vs = flatten_namedtuple(getparams(t))
        for nm in nms
            if !(nm in names)
                push!(names, nm)
            end
        end
        # Convert the names and values to a single dictionary.
        return Dict(nms[j] => vs[j] for j in 1:length(vs))
    end
    # names = collect(names_set)
    vals = [get(dicts[i], key, missing) for i in eachindex(dicts),
        (j, key) in enumerate(names)]

    return names, vals
end

function flatten_namedtuple(nt::NamedTuple)
    names_vals = mapreduce(vcat, keys(nt)) do k
        v = nt[k]
        if length(v) == 1
            return [(Symbol(k), v)]
        else
            return mapreduce(vcat, zip(v[1], v[2])) do (vnval, vn)
                return collect(FlattenIterator(vn, vnval))
            end
        end
    end
    return [vn[1] for vn in names_vals], [vn[2] for vn in names_vals]
end

function get_transition_extras(ts::AbstractVector{<:VarInfo})
    valmat = reshape([getlogp(t) for t in ts], :, 1)
    return [:lp], valmat
end

function get_transition_extras(ts::AbstractVector)
    # Extract all metadata.
    extra_data = map(metadata, ts)
    return names_values(extra_data)
end

function names_values(extra_data::AbstractVector{<:NamedTuple{names}}) where names
    values = [getfield(data, name) for data in extra_data, name in names]
    return collect(names), values
end

function names_values(extra_data::AbstractVector{<:NamedTuple})
    # Obtain all parameter names.
    names_set = Set(Symbol[])
    for data in extra_data
        for name in names(data)
            push!(extra_names_set, name)
        end
    end
    extra_names = collect(extra_names_set)

    # Extract all values as matrix.
    values = [
        hasfield(data, name) ? missing : getfield(data, name)
        for data in extra_data, name in extra_names
    ]

    return extra_names, values
end

getlogevidence(transitions, sampler, state) = missing

# Default MCMCChains.Chains constructor.
# This is type piracy (at least for SampleFromPrior).
function AbstractMCMC.bundle_samples(
    ts::Vector,
    model::AbstractModel,
    spl::Union{Sampler{<:InferenceAlgorithm},SampleFromPrior},
    state,
    chain_type::Type{MCMCChains.Chains};
    save_state = false,
    kwargs...
)
    # Convert transitions to array format.
    # Also retrieve the variable names.
    nms, vals = _params_to_array(ts)

    # Get the values of the extra parameters in each transition.
    extra_params, extra_values = get_transition_extras(ts)

    # Extract names & construct param array.
    nms = [nms; extra_params]
    parray = hcat(vals, extra_values)

    # Get the average or final log evidence, if it exists.
    le = getlogevidence(ts, spl, state)

    # Set up the info tuple.
    if save_state
        info = (model = model, sampler = spl, samplerstate = state)
    else
        info = NamedTuple()
    end

    # Conretize the array before giving it to MCMCChains.
    parray = MCMCChains.concretize(parray)

    # Chain construction.
    return MCMCChains.Chains(
        parray,
        nms,
        (internals = extra_params,);
        evidence=le,
        info=info,
    ) |> sort
end

# This is type piracy (for SampleFromPrior).
function AbstractMCMC.bundle_samples(
    ts::Vector,
    model::AbstractModel,
    spl::Union{Sampler{<:InferenceAlgorithm},SampleFromPrior},
    state,
    chain_type::Type{Vector{NamedTuple}};
    kwargs...
)
    return map(ts) do t
        params = map(first, getparams(t))
        return merge(params, metadata(t))
    end
end

function save(c::MCMCChains.Chains, spl::Sampler, model, vi, samples)
    nt = NamedTuple{(:sampler, :model, :vi, :samples)}((spl, model, deepcopy(vi), samples))
    return setinfo(c, merge(nt, c.info))
end

function resume(chain::MCMCChains.Chains, args...; kwargs...)
    return resume(Random.GLOBAL_RNG, chain, args...; kwargs...)
end

function resume(rng::Random.AbstractRNG, chain::MCMCChains.Chains, args...;
                progress=PROGRESS[], kwargs...)
    isempty(chain.info) && error("[Turing] cannot resume from a chain without state info")

    # Sample a new chain.
    return AbstractMCMC.mcmcsample(
        rng,
        chain.info[:model],
        chain.info[:sampler],
        args...;
        resume_from = chain,
        chain_type = MCMCChains.Chains,
        progress = progress,
        kwargs...
    )
end

DynamicPPL.loadstate(chain::MCMCChains.Chains) = chain.info[:samplerstate]

#######################################
# Concrete algorithm implementations. #
#######################################

include("ess.jl")
include("hmc.jl")
include("mh.jl")
include("is.jl")
include("AdvancedSMC.jl")
include("gibbs.jl")
include("gibbs_conditional.jl")
include("../contrib/inference/sghmc.jl")
include("emcee.jl")

################
# Typing tools #
################

for alg in (:SMC, :PG, :MH, :IS, :ESS, :Gibbs, :Emcee)
    @eval DynamicPPL.getspace(::$alg{space}) where {space} = space
end
for alg in (:HMC, :HMCDA, :NUTS, :SGLD, :SGHMC)
    @eval DynamicPPL.getspace(::$alg{<:Any, space}) where {space} = space
end

function get_matching_type(
    spl::Sampler{<:Union{PG, SMC}},
    vi,
    ::Type{TV},
) where {T, N, TV <: Array{T, N}}
    return TArray{T, N}
end

##############
# Utilities  #
##############

DynamicPPL.getspace(spl::Sampler) = getspace(spl.alg)
DynamicPPL.inspace(vn::VarName, spl::Sampler) = inspace(vn, getspace(spl.alg))

"""

    predict([rng::AbstractRNG,] model::Model, chain::MCMCChains.Chains; include_all=false)

Execute `model` conditioned on each sample in `chain`, and return the resulting `Chains`.

If `include_all` is `false`, the returned `Chains` will contain only those variables
sampled/not present in `chain`.

# Details
Internally calls `Turing.Inference.transitions_from_chain` to obtained the samples
and then converts these into a `Chains` object using `AbstractMCMC.bundle_samples`.

# Example
```jldoctest
julia> using Turing; Turing.turnprogress(false);
[ Info: [Turing]: progress logging is disabled globally

julia> @model function linear_reg(x, y, σ = 0.1)
           β ~ Normal(0, 1)

           for i ∈ eachindex(y)
               y[i] ~ Normal(β * x[i], σ)
           end
       end;

julia> σ = 0.1; f(x) = 2 * x + 0.1 * randn();

julia> Δ = 0.1; xs_train = 0:Δ:10; ys_train = f.(xs_train);

julia> xs_test = [10 + Δ, 10 + 2 * Δ]; ys_test = f.(xs_test);

julia> m_train = linear_reg(xs_train, ys_train, σ);

julia> chain_lin_reg = sample(m_train, NUTS(100, 0.65), 200);
┌ Info: Found initial step size
└   ϵ = 0.003125

julia> m_test = linear_reg(xs_test, Vector{Union{Missing, Float64}}(undef, length(ys_test)), σ);

julia> predictions = predict(m_test, chain_lin_reg)
Object of type Chains, with data of type 100×2×1 Array{Float64,3}

Iterations        = 1:100
Thinning interval = 1
Chains            = 1
Samples per chain = 100
parameters        = y[1], y[2]

2-element Array{ChainDataFrame,1}

Summary Statistics
  parameters     mean     std  naive_se     mcse       ess   r_hat
  ──────────  ───────  ──────  ────────  ───────  ────────  ──────
        y[1]  20.1974  0.1007    0.0101  missing  101.0711  0.9922
        y[2]  20.3867  0.1062    0.0106  missing  101.4889  0.9903

Quantiles
  parameters     2.5%    25.0%    50.0%    75.0%    97.5%
  ──────────  ───────  ───────  ───────  ───────  ───────
        y[1]  20.0342  20.1188  20.2135  20.2588  20.4188
        y[2]  20.1870  20.3178  20.3839  20.4466  20.5895

julia> ys_pred = vec(mean(Array(group(predictions, :y)); dims = 1));

julia> sum(abs2, ys_test - ys_pred) ≤ 0.1
true
```
"""
function predict(model::Model, chain::MCMCChains.Chains; kwargs...)
    return predict(Random.GLOBAL_RNG, model, chain; kwargs...)
end
function predict(rng::AbstractRNG, model::Model, chain::MCMCChains.Chains; include_all = false)
    spl = DynamicPPL.SampleFromPrior()

    # Sample transitions using `spl` conditioned on values in `chain`
    transitions = [
        transitions_from_chain(rng, model, chain[:, :, chn_idx]; sampler = spl)
        for chn_idx = 1:size(chain, 3)
    ]

    # Let the Turing internals handle everything else for you
    chain_result = reduce(
        MCMCChains.chainscat, [
            AbstractMCMC.bundle_samples(
                transitions[chn_idx],
                model,
                spl,
                nothing,
                MCMCChains.Chains
            ) for chn_idx = 1:size(chain, 3)
        ]
    )

    parameter_names = if include_all
        names(chain_result, :parameters)
    else
        filter(k -> ∉(k, names(chain, :parameters)), names(chain_result, :parameters))
    end

    return chain_result[parameter_names]
end

"""

    transitions_from_chain(
        [rng::AbstractRNG,]
        model::Model,
        chain::MCMCChains.Chains;
        sampler = DynamicPPL.SampleFromPrior()
    )

Execute `model` conditioned on each sample in `chain`, and return resulting transitions.

The returned transitions are represented in a `Vector{<:Turing.Inference.Transition}`.

# Details

In a bit more detail, the process is as follows:
1. For every `sample` in `chain`
   1. For every `variable` in `sample`
      1. Set `variable` in `model` to its value in `sample`
   2. Execute `model` with variables fixed as above, sampling variables NOT present
      in `chain` using `SampleFromPrior`
   3. Return sampled variables and log-joint

# Example
```julia-repl
julia> using Turing

julia> @model function demo()
           m ~ Normal(0, 1)
           x ~ Normal(m, 1)
       end;

julia> m = demo();

julia> chain = Chains(randn(2, 1, 1), ["m"]); # 2 samples of `m`

julia> transitions = Turing.Inference.transitions_from_chain(m, chain);

julia> [Turing.Inference.getlogp(t) for t in transitions] # extract the logjoints
2-element Array{Float64,1}:
 -3.6294991938628374
 -2.5697948166987845

julia> [first(t.θ.x) for t in transitions] # extract samples for `x`
2-element Array{Array{Float64,1},1}:
 [-2.0844148956440796]
 [-1.704630494695469]
```
"""
function transitions_from_chain(
    model::Turing.Model,
    chain::MCMCChains.Chains;
    kwargs...
)
    return transitions_from_chain(Random.GLOBAL_RNG, model, chain; kwargs...)
end
function transitions_from_chain(
    rng::AbstractRNG,
    model::Turing.Model,
    chain::MCMCChains.Chains;
    sampler = DynamicPPL.SampleFromPrior()
)
    vi = Turing.VarInfo(model)

    transitions = map(1:length(chain)) do i
        c = chain[i]
        md = vi.metadata
        for v in keys(md)
            for vn in md[v].vns
                vn_sym = Symbol(vn)

                # Cannot use `vn_sym` to index in the chain
                # so we have to extract the corresponding "linear"
                # indices and use those.
                # `ks` is empty if `vn_sym` not in `c`.
                ks = MCMCChains.namesingroup(c, vn_sym)

                if !isempty(ks)
                    # 1st dimension is of size 1 since `c`
                    # only contains a single sample, and the
                    # last dimension is of size 1 since
                    # we're assuming we're working with a single chain.
                    val = copy(vec(c[ks].value))
                    DynamicPPL.setval!(vi, val, vn)
                    DynamicPPL.settrans!(vi, false, vn)
                else
                    DynamicPPL.set_flag!(vi, vn, "del")
                end
            end
        end
        # Execute `model` on the parameters set in `vi` and sample those with `"del"` flag using `sampler`
        model(rng, vi, sampler)

        # Convert `VarInfo` into `NamedTuple` and save
        theta = DynamicPPL.tonamedtuple(vi)
        lp = Turing.getlogp(vi)
        Transition(theta, lp)
    end

    return transitions
end

end # module


#include("../ReadWrite/Header.jl")

"""
**SeisBin**

*Bin seismic data onto regular grid using trace headers*

**IN**

* in: filename of input, irregularly sampled data
* out: filename of output, regularly sampled data
* style="sxsygxgy"
* ang=90
* gamma=1
* osx=0
* osy=0
* ogx=0
* ogy=0
* omx=0
* omy=0
* ohx=0
* ohy=0
* oh=0
* oaz=0
* dsx=1
* dsy=1
* dgx=1
* dgy=1
* dmx=1
* dmy=1
* dhx=1
* dhy=1
* dh=1
* daz=1
* min_isx=0
* max_isx=0
* min_isy=0
* max_isy=0
* min_igx=0
* max_igx=0
* min_igy=0
* max_igy=0
* min_imx=0
* max_imx=0
* min_imy=0
* max_imy=0
* min_ihx=0
* max_ihx=0
* min_ihy=0
* max_ihy=0
* min_ih=0
* max_ih=0
* min_iaz=0
* max_iaz=0

**OUT**

*Credits: Aaron Stanton, 2015*

"""

function SeisBin(in,out;style="sxsygxgy",ang=90,gamma=1,osx=0,osy=0,ogx=0,ogy=0,omx=0,omy=0,ohx=0,ohy=0,oh=0,oaz=0,dsx=1,dsy=1,dgx=1,dgy=1,dmx=1,dmy=1,dhx=1,dhy=1,dh=1,daz=1,min_isx=0,max_isx=0,min_isy=0,max_isy=0,min_igx=0,max_igx=0,min_igy=0,max_igy=0,min_imx=0,max_imx=0,min_imy=0,max_imy=0,min_ihx=0,max_ihx=0,min_ihy=0,max_ihy=0,min_ih=0,max_ih=0,min_iaz=0,max_iaz=0,ntrace=10000)

	rad2deg = 180/pi;
	deg2rad = pi/180;
	gammainv = 1/gamma;
	if (ang > 90)
		ang2=-deg2rad*(ang-90)
	else
		ang2=deg2rad*(90-ang)
	end

	nsx = max_isx - min_isx + 1
	nsy = max_isy - min_isy + 1
	ngx = max_igx - min_igx + 1
	ngy = max_igy - min_igy + 1
	nmx = max_imx - min_imx + 1
	nmy = max_imy - min_imy + 1
	nhx = max_ihx - min_ihx + 1
	nhy = max_ihy - min_ihy + 1
	nh  = max_ih  - min_ih  + 1
	naz = max_iaz - min_iaz + 1

	if (style=="sxsygxgy")
		nx1=nsx;nx2=nsy;nx3=ngx;nx4=ngy;
		ox1=osx;ox2=osy;ox3=ogx;ox4=ogy;
		dx1=dsx;dx2=dsy;dx3=dgx;dx4=dgy;
		label2="sx";label3="sy";label4="gx";label5="gy";
		unit2="m";unit3="m";unit4="m";unit5="m";
	elseif (style=="mxmyhxhy")
		nx1=nmx;nx2=nmy;nx3=nhx;nx4=nhy;
		ox1=omx;ox2=omy;ox3=ohx;ox4=ohy;
		dx1=dmx;dx2=dmy;dx3=dhx;dx4=dhy;
		label2="mx";label3="my";label4="hx";label5="hy";
		unit2="m";unit3="m";unit4="m";unit5="m";
	elseif (style=="mxmyhaz")
		nx1=nmx;nx2=nmy;nx3=nh;nx4=naz;
		ox1=omx;ox2=omy;ox3=oh;ox4=oaz;
		dx1=dmx;dx2=dmy;dx3=dh;dx4=daz;
		label2="mx";label3="my";label4="h";label5="az";
		unit2="m";unit3="m";unit4="m";unit5="Degrees";
	elseif (style=="sxsyhxhy")
		nx1=nsx;nx2=nsy;nx3=nhx;nx4=nhy;
		ox1=osx;ox2=osy;ox3=ohx;ox4=ohy;
		dx1=dsx;dx2=dsy;dx3=dhx;dx4=dhy;
		label2="sx";label3="sy";label4="hx";label5="hy";
		unit2="m";unit3="m";unit4="m";unit5="m";
	elseif (style=="gxgyhxhy")
		nx1=ngx;nx2=ngy;nx3=nhx;nx4=nhy;
		ox1=ogx;ox2=ogy;ox3=ohx;ox4=ohy;
		dx1=dgx;dx2=dgy;dx3=dhx;dx4=dhy;
		label2="gx";label3="gy";label4="hx";label5="hy";
		unit2="m";unit3="m";unit4="m";unit5="m";
	elseif (style=="sxsyhaz")
		nx1=nsx;nx2=nsy;nx3=nh;nx4=naz;
		ox1=osx;ox2=osy;ox3=oh;ox4=oaz;
		dx1=dsx;dx2=dsy;dx3=dh;dx4=daz;
		label2="sx";label3="sy";label4="hx";label5="az";
		unit2="m";unit3="m";unit4="m";unit5="Degrees";
	elseif (style=="gxgyhaz")
		nx1=ngx;nx2=ngy;nx3=nh;nx4=naz;
		ox1=ogx;ox2=ogy;ox3=oh;ox4=oaz;
		dx1=dgx;dx2=dgy;dx3=dh;dx4=daz;
		label2="gx";label3="gy";label4="h";label5="az";
		unit2="m";unit3="m";unit4="m";unit5="Degrees";
	else
		error("style not recognized.")
	end
	nx_out = nx1*nx2*nx3*nx4

	stream = open(ParseHeaderName(in))
	@compat nx_in = int(filesize(stream)/(4*length(fieldnames(Header))))
	seek(stream, header_count["n1"])
	nt = read(stream,Int32)
	seek(stream, header_count["d1"])
	dt = read(stream,Float32)
	close(stream)
	d = zeros(Float32,nt,1)
	h = Array(Header,1)
	h[1] = InitSeisHeader()

	extent = Extent(convert(Int32,nt),convert(Int32,nx1),convert(Int32,nx2),convert(Int32,nx3),convert(Int32,nx4),
		   convert(Float32,0),convert(Float32,ox1),convert(Float32,ox2),convert(Float32,ox3),convert(Float32,ox4),
		   convert(Float32,dt),convert(Float32,dx1),convert(Float32,dx2),convert(Float32,dx3),convert(Float32,dx4),
		   "Time",label2,label3,label4,label5,
		   "s",unit2,unit3,unit4,unit5,
		   "")

	j = 1
	for ix1 = 1 : nx1
		for ix2 = 1 : nx2
			for ix3 = 1 : nx3
				for ix4 = 1 : nx4
					h[1].tracenum = convert(typeof(h[1].tracenum),j)
					h[1].o1 = convert(typeof(h[1].o1),0)
					h[1].n1 = convert(typeof(h[1].n1),nt)
					h[1].d1 = convert(typeof(h[1].d1),dt)
					if (style=="sxsygxgy")
						h[1].isx = convert(typeof(h[1].isx),ix1 - 1 + min_isx)
						h[1].isy = convert(typeof(h[1].isy),ix2 - 1 + min_isy)
						h[1].igx = convert(typeof(h[1].igx),ix3 - 1 + min_igx)
						h[1].igy = convert(typeof(h[1].igy),ix4 - 1 + min_igy)
						sx_rot = convert(Float32,(ix1 - 1 + min_isx)*dsx + osx);
						sy_rot = convert(Float32,(ix2 - 1 + min_isy)*dsy + osy);
						gx_rot = convert(Float32,(ix3 - 1 + min_igx)*dgx + ogx);
						gy_rot = convert(Float32,(ix4 - 1 + min_igy)*dgy + ogy);
						h[1].sx =  (sx_rot-osx)*cos(ang2) + (sy_rot-osy)*sin(ang2) + osx;
						h[1].sy = -(sx_rot-osx)*sin(ang2) + (sy_rot-osy)*cos(ang2) + osy;
						h[1].gx =  (gx_rot-ogx)*cos(ang2) + (gy_rot-ogy)*sin(ang2) + ogx;
						h[1].gy = -(gx_rot-ogx)*sin(ang2) + (gy_rot-ogy)*cos(ang2) + ogy;
						h[1].hx = h[1].gx - h[1].sx
						h[1].hy = h[1].gy - h[1].sy
						h[1].h = sqrt((h[1].hx^2) + (h[1].hy^2))
						h[1].az = rad2deg*atan2((h[1].gy-h[1].sy),(h[1].gx-h[1].sx))
						if (h[1].az < 0)
							h[1].az += 360.0
						end
						h[1].mx = h[1].sx + h[1].hx/(1 + gammainv);
						h[1].my = h[1].sy + h[1].hy/(1 + gammainv);
						mx_rot = (h[1].mx-omx)*cos(ang2) - (h[1].my-omy)*sin(ang2) + omx;
						my_rot = (h[1].mx-omx)*sin(ang2) + (h[1].my-omy)*cos(ang2) + omy;
						hx_rot = (h[1].hx-ohx)*cos(ang2) - (h[1].hy-ohy)*sin(ang2) + ohx;
						hy_rot = (h[1].hx-ohx)*sin(ang2) + (h[1].hy-ohy)*cos(ang2) + ohy;
						h[1].imx = convert(Int32,round((mx_rot-omx)/dmx))
						h[1].imy = convert(Int32,round((my_rot-omy)/dmy))
						h[1].ihx = convert(Int32,round((hx_rot-ohx)/dhx))
						h[1].ihy = convert(Int32,round((hy_rot-ohy)/dhy))
						h[1].ih = convert(Int32,round((h[1].h-oh)/dh))
						h[1].iaz = convert(Int32,round((h[1].az-oaz)/daz))
					elseif (style=="mxmyhxhy")
						h[1].imx = convert(typeof(h[1].imx),ix1 - 1 + min_imx)
						h[1].imy = convert(typeof(h[1].imy),ix2 - 1 + min_imy)
						h[1].ihx = convert(typeof(h[1].ihx),ix3 - 1 + min_ihx)
						h[1].ihy = convert(typeof(h[1].ihy),ix4 - 1 + min_ihy)
						mx_rot = convert(Float32,(ix1 - 1 + min_imx)*dmx + omx);
						my_rot = convert(Float32,(ix2 - 1 + min_imy)*dmy + omy);
						hx_rot = convert(Float32,(ix3 - 1 + min_ihx)*dhx + ohx);
						hy_rot = convert(Float32,(ix4 - 1 + min_ihy)*dhy + ohy);
						h[1].mx =  (mx_rot-omx)*cos(ang2) + (my_rot-omy)*sin(ang2) + omx;
						h[1].my = -(mx_rot-omx)*sin(ang2) + (my_rot-omy)*cos(ang2) + omy;
						h[1].hx =  (hx_rot-ohx)*cos(ang2) + (hy_rot-ohy)*sin(ang2) + ohx;
						h[1].hy = -(hx_rot-ohx)*sin(ang2) + (hy_rot-ohy)*cos(ang2) + ohy;
						h[1].sx = h[1].mx - h[1].hx/(1 + gammainv);
						h[1].sy = h[1].my - h[1].hy/(1 + gammainv);
						h[1].gx = h[1].mx + h[1].hx*(1-(1/(1 + gammainv)));
						h[1].gy = h[1].my + h[1].hy*(1-(1/(1 + gammainv)));
						sx_rot = (h[1].sx-osx)*cos(ang2) - (h[1].sy-osy)*sin(ang2) + osx;
						sy_rot = (h[1].sx-osx)*sin(ang2) + (h[1].sy-osy)*cos(ang2) + osy;
						gx_rot = (h[1].gx-ogx)*cos(ang2) - (h[1].gy-ogy)*sin(ang2) + ogx;
						gy_rot = (h[1].gx-ogx)*sin(ang2) + (h[1].gy-ogy)*cos(ang2) + ogy;
						h[1].isx = convert(Int32,round((sx_rot-osx)/dsx))
						h[1].isy = convert(Int32,round((sy_rot-osy)/dsy))
						h[1].igx = convert(Int32,round((gx_rot-ogx)/dgx))
						h[1].igy = convert(Int32,round((gy_rot-ogy)/dgy))
						h[1].h = sqrt((h[1].hx^2) + (h[1].hy^2))
						h[1].az = rad2deg*atan2((h[1].gy-h[1].sy),(h[1].gx-h[1].sx))
						if (h[1].az < 0)
							h[1].az += 360.0
						end
						h[1].ih = convert(Int32,round((h[1].h-oh)/dh))
						h[1].iaz = convert(Int32,round((h[1].az-oaz)/daz))
					elseif (style=="mxmyhaz")
						h[1].imx = convert(typeof(h[1].imx),ix1 - 1 + min_imx)
						h[1].imy = convert(typeof(h[1].imy),ix2 - 1 + min_imy)
						h[1].ih = convert(typeof(h[1].ih),  ix3 - 1 + min_ih)
						h[1].iaz = convert(typeof(h[1].iaz),ix4 - 1 + min_iaz)
						mx_rot = convert(Float32,(ix1 - 1 + min_imx)*dmx + omx);
						my_rot = convert(Float32,(ix2 - 1 + min_imy)*dmy + omy);
						h[1].mx =  (mx_rot-omx)*cos(ang2) + (my_rot-omy)*sin(ang2) + omx;
						h[1].my = -(mx_rot-omx)*sin(ang2) + (my_rot-omy)*cos(ang2) + omy;
						h[1].h = convert(Float32,(ix3 - 1 + min_ih)*dh + oh);
						h[1].az = convert(Float32,(ix4 - 1 + min_iaz)*daz + oaz);
						if (h[1].az <= 90)
							h[1].hx = h[1].h*cos(deg2rad*h[1].az);
							h[1].hy = h[1].h*sin(deg2rad*h[1].az);
						elseif (h[1].az > 90 && h[1].az <= 180)
							h[1].hx =-h[1].h*cos(pi-(deg2rad*h[1].az));
							h[1].hy = h[1].h*sin(pi-(deg2rad*h[1].az));
						elseif (h[1].az > 180 && h[1].az <= 270)
							h[1].hx =-h[1].h*cos((deg2rad*h[1].az)-pi);
							h[1].hy =-h[1].h*sin((deg2rad*h[1].az)-pi);
						else
							h[1].hx = h[1].h*cos(2*pi-(deg2rad*h[1].az));
							h[1].hy =-h[1].h*sin(2*pi-(deg2rad*h[1].az));
						end
						h[1].sx = h[1].mx - h[1].hx/(1 + gammainv);
						h[1].sy = h[1].my - h[1].hy/(1 + gammainv);
						h[1].gx = h[1].mx + h[1].hx*(1-(1/(1 + gammainv)));
						h[1].gy = h[1].my + h[1].hy*(1-(1/(1 + gammainv)));
						sx_rot = (h[1].sx-osx)*cos(ang2) - (h[1].sy-osy)*sin(ang2) + osx;
						sy_rot = (h[1].sx-osx)*sin(ang2) + (h[1].sy-osy)*cos(ang2) + osy;
						gx_rot = (h[1].gx-ogx)*cos(ang2) - (h[1].gy-ogy)*sin(ang2) + ogx;
						gy_rot = (h[1].gx-ogx)*sin(ang2) + (h[1].gy-ogy)*cos(ang2) + ogy;
						h[1].isx = convert(Int32,round((sx_rot-osx)/dsx))
						h[1].isy = convert(Int32,round((sy_rot-osy)/dsy))
						h[1].igx = convert(Int32,round((gx_rot-ogx)/dgx))
						h[1].igy = convert(Int32,round((gy_rot-ogy)/dgy))
						hx_rot = (h[1].hx-ohx)*cos(ang2) - (h[1].hy-ohy)*sin(ang2) + ohx;
						hy_rot = (h[1].hx-ohx)*sin(ang2) + (h[1].hy-ohy)*cos(ang2) + ohy;
						h[1].ihx = convert(Int32,round((hx_rot-ohx)/dhx))
						h[1].ihy = convert(Int32,round((hy_rot-ohy)/dhy))
					elseif (style=="sxsyhxhy")
						h[1].isx = convert(typeof(h[1].isx),ix1 - 1 + min_isx)
						h[1].isy = convert(typeof(h[1].isy),ix2 - 1 + min_isy)
						h[1].ihx = convert(typeof(h[1].ihx),ix3 - 1 + min_ihx)
						h[1].ihy = convert(typeof(h[1].ihy),ix4 - 1 + min_ihy)
						sx_rot = convert(Float32,(ix1 - 1 + min_isx)*dsx + osx);
						sy_rot = convert(Float32,(ix2 - 1 + min_isy)*dsy + osy);
						hx_rot = convert(Float32,(ix3 - 1 + min_ihx)*dhx + ohx);
						hy_rot = convert(Float32,(ix4 - 1 + min_ihy)*dhy + ohy);
						h[1].sx =  (sx_rot-osx)*cos(ang2) + (sy_rot-osy)*sin(ang2) + osx;
						h[1].sy = -(sx_rot-osx)*sin(ang2) + (sy_rot-osy)*cos(ang2) + osy;
						h[1].hx =  (hx_rot-ohx)*cos(ang2) + (hy_rot-ohy)*sin(ang2) + ohx;
						h[1].hy = -(hx_rot-ohx)*sin(ang2) + (hy_rot-ohy)*cos(ang2) + ohy;
						h[1].gx = h[1].sx + h[1].hx;
						h[1].gy = h[1].sy + h[1].hy;
						h[1].h = sqrt((h[1].hx^2) + (h[1].hy^2))
						h[1].az = rad2deg*atan2((h[1].gy-h[1].sy),(h[1].gx-h[1].sx))
						if (h[1].az < 0)
							h[1].az += 360.0
						end
						h[1].mx = h[1].sx + h[1].hx/(1 + gammainv);
						h[1].my = h[1].sy + h[1].hy/(1 + gammainv);
						mx_rot = (h[1].mx-omx)*cos(ang2) - (h[1].my-omy)*sin(ang2) + omx;
						my_rot = (h[1].mx-omx)*sin(ang2) + (h[1].my-omy)*cos(ang2) + omy;
						hx_rot = (h[1].hx-ohx)*cos(ang2) - (h[1].hy-ohy)*sin(ang2) + ohx;
						hy_rot = (h[1].hx-ohx)*sin(ang2) + (h[1].hy-ohy)*cos(ang2) + ohy;
						h[1].imx = convert(Int32,round((mx_rot-omx)/dmx))
						h[1].imy = convert(Int32,round((my_rot-omy)/dmy))
						h[1].ihx = convert(Int32,round((hx_rot-ohx)/dhx))
						h[1].ihy = convert(Int32,round((hy_rot-ohy)/dhy))
						h[1].ih = convert(Int32,round((h[1].h-oh)/dh))
						h[1].iaz = convert(Int32,round((h[1].az-oaz)/daz))
					elseif (style=="gxgyhxhy")
						h[1].igx = convert(typeof(h[1].igx),ix1 - 1 + min_igx)
						h[1].igy = convert(typeof(h[1].igy),ix2 - 1 + min_igy)
						h[1].ihx = convert(typeof(h[1].ihx),ix3 - 1 + min_ihx)
						h[1].ihy = convert(typeof(h[1].ihy),ix4 - 1 + min_ihy)
						gx_rot = convert(Float32,(ix1 - 1 + min_igx)*dgx + ogx);
						gy_rot = convert(Float32,(ix2 - 1 + min_igy)*dgy + ogy);
						hx_rot = convert(Float32,(ix3 - 1 + min_ihx)*dhx + ohx);
						hy_rot = convert(Float32,(ix4 - 1 + min_ihy)*dhy + ohy);
						h[1].gx =  (gx_rot-ogx)*cos(ang2) + (gy_rot-ogy)*sin(ang2) + ogx;
						h[1].gy = -(gx_rot-ogx)*sin(ang2) + (gy_rot-ogy)*cos(ang2) + ogy;
						h[1].hx =  (hx_rot-ohx)*cos(ang2) + (hy_rot-ohy)*sin(ang2) + ohx;
						h[1].hy = -(hx_rot-ohx)*sin(ang2) + (hy_rot-ohy)*cos(ang2) + ohy;
						h[1].sx = h[1].gx - h[1].hx;
						h[1].sy = h[1].gy - h[1].hy;
						h[1].h = sqrt((h[1].hx^2) + (h[1].hy^2))
						h[1].az = rad2deg*atan2((h[1].gy-h[1].sy),(h[1].gx-h[1].sx))
						if (h[1].az < 0)
							h[1].az += 360.0
						end
						h[1].mx = h[1].sx + h[1].hx/(1 + gammainv);
						h[1].my = h[1].sy + h[1].hy/(1 + gammainv);
						mx_rot = (h[1].mx-omx)*cos(ang2) - (h[1].my-omy)*sin(ang2) + omx;
						my_rot = (h[1].mx-omx)*sin(ang2) + (h[1].my-omy)*cos(ang2) + omy;
						hx_rot = (h[1].hx-ohx)*cos(ang2) - (h[1].hy-ohy)*sin(ang2) + ohx;
						hy_rot = (h[1].hx-ohx)*sin(ang2) + (h[1].hy-ohy)*cos(ang2) + ohy;
						h[1].imx = convert(Int32,round((mx_rot-omx)/dmx))
						h[1].imy = convert(Int32,round((my_rot-omy)/dmy))
						h[1].ihx = convert(Int32,round((hx_rot-ohx)/dhx))
						h[1].ihy = convert(Int32,round((hy_rot-ohy)/dhy))
						h[1].ih = convert(Int32,round((h[1].h-oh)/dh))
						h[1].iaz = convert(Int32,round((h[1].az-oaz)/daz))
					elseif (style=="sxsyhaz")
						error("sxsyhaz not developed yet.")
					elseif (style=="gxgyhaz")
						error("gxgyhaz not developed yet.")
					end
					h[1].selev = convert(typeof(h[1].selev),0)
					h[1].gelev = convert(typeof(h[1].gelev),0)
					h[1].trid = convert(typeof(h[1].trid),0)
					SeisWrite(out,d,h,extent,itrace=j)
					j += 1
				end
			end
		end
	end

	out_d = ParseDataName(out)
	out_h = ParseHeaderName(out)
	stream_d = open(out_d,"a")
	stream_h = open(out_h,"a")

	if (style=="sxsygxgy")
		j = 1
		while j <= nx_in
			d,h,e = SeisRead(in,group="some",itrace=j,ntrace=ntrace)
			num_traces_in = size(d[:,:],2)
			for k = 1 : num_traces_in
				itrace = (h[k].isx - min_isx)*nx2*nx3*nx4 + (h[k].isy - min_isy)*nx3*nx4 + (h[k].igx - min_igx)*nx4 + h[k].igy - min_igy + 1
				if (h[k].isx >= min_isx && h[k].isx <= max_isx && h[k].isy >= min_isy && h[k].isy <= max_isy && h[k].igx >= min_igx && h[k].igx <= max_igx && h[k].igy >= min_igy && h[k].igy <= max_igy)
					h[k].tracenum = convert(Int32,itrace)
					if (itrace > 0 && itrace <= nx_out)
						position_d = 4*nt*(itrace - 1)
						seek(stream_d,position_d)
						write(stream_d,d[:,k])
						PutHeader(stream_h,h[k],itrace)
					end
				end
			end
			j += num_traces_in
		end
	elseif (style=="mxmyhxhy")
		j = 1
		while j <= nx_in
			d,h,e = SeisRead(in,group="some",itrace=j,ntrace=ntrace)
			num_traces_in = size(d[:,:],2)
			for k = 1 : num_traces_in
				itrace = (h[k].imx - min_imx)*nx2*nx3*nx4 + (h[k].imy - min_imy)*nx3*nx4 + (h[k].ihx - min_ihx)*nx4 + h[k].ihy - min_ihy + 1
				if (h[k].imx >= min_imx && h[k].imx <= max_imx && h[k].imy >= min_imy && h[k].imy <= max_imy && h[k].ihx >= min_ihx && h[k].ihx <= max_ihx && h[k].ihy >= min_ihy && h[k].ihy <= max_ihy)
					h[k].tracenum = convert(Int32,itrace)
					if (itrace > 0 && itrace <= nx_out)
						position_d = 4*nt*(itrace - 1)
						seek(stream_d,position_d)
						write(stream_d,d[:,k])
						PutHeader(stream_h,h[k],itrace)
					end
				end
			end
			j += num_traces_in
		end
	elseif (style=="mxmyhaz")
		j = 1
		while j <= nx_in
			d,h,e = SeisRead(in,group="some",itrace=j,ntrace=ntrace)
			num_traces_in = size(d[:,:],2)
			for k = 1 : num_traces_in
				itrace = (h[k].imx - min_imx)*nx2*nx3*nx4 + (h[k].imy - min_imy)*nx3*nx4 + (h[k].ih - min_ih)*nx4 + h[k].iaz - min_iaz + 1
				if (h[k].imx >= min_imx && h[k].imx <= max_imx && h[k].imy >= min_imy && h[k].imy <= max_imy && h[k].ih >= min_ih && h[k].ih <= max_ih && h[k].iaz >= min_iaz && h[k].iaz <= max_iaz)
					h[k].tracenum = convert(Int32,itrace)
					if (itrace > 0 && itrace <= nx_out)
						position_d = 4*nt*(itrace - 1)
						seek(stream_d,position_d)
						write(stream_d,d[:,k])
						PutHeader(stream_h,h[k],itrace)
					end
				end
			end
			j += num_traces_in
		end
	elseif (style=="sxsyhxhy")
		j = 1
		while j <= nx_in
			d,h,e = SeisRead(in,group="some",itrace=j,ntrace=ntrace)
			num_traces_in = size(d[:,:],2)
			for k = 1 : num_traces_in
				itrace = (h[k].isx - min_isx)*nx2*nx3*nx4 + (h[k].isy - min_isy)*nx3*nx4 + (h[k].ihx - min_ihx)*nx4 + h[k].ihy - min_ihy + 1
				if (h[k].isx >= min_isx && h[k].isx <= max_isx && h[k].isy >= min_isy && h[k].isy <= max_isy && h[k].ihx >= min_ihx && h[k].ihx <= max_ihx && h[k].ihy >= min_ihy && h[k].ihy <= max_ihy)
					h[k].tracenum = convert(Int32,itrace)
					if (itrace > 0 && itrace <= nx_out)
						position_d = 4*nt*(itrace - 1)
						seek(stream_d,position_d)
						write(stream_d,d[:,k])
						PutHeader(stream_h,h[k],itrace)
					end
				end
			end
			j += num_traces_in
		end
	elseif (style=="gxgyhxhy")
		j = 1
		while j <= nx_in
			d,h,e = SeisRead(in,group="some",itrace=j,ntrace=ntrace)
			num_traces_in = size(d[:,:],2)
			for k = 1 : num_traces_in
				itrace = (h[k].igx - min_igx)*nx2*nx3*nx4 + (h[k].igy - min_igy)*nx3*nx4 + (h[k].ihx - min_ihx)*nx4 + h[k].ihy - min_ihy + 1
				if (h[k].igx >= min_igx && h[k].igx <= max_igx && h[k].igy >= min_igy && h[k].igy <= max_igy && h[k].ihx >= min_ihx && h[k].ihx <= max_ihx && h[k].ihy >= min_ihy && h[k].ihy <= max_ihy)
					h[k].tracenum = convert(Int32,itrace)
					if (itrace > 0 && itrace <= nx_out)
						position_d = 4*nt*(itrace - 1)
						seek(stream_d,position_d)
						write(stream_d,d[:,k])
						PutHeader(stream_h,h[k],itrace)
					end
				end
			end
			j += num_traces_in
		end
	elseif (style=="sxsyhaz")
		j = 1
		while j <= nx_in
			d,h,e = SeisRead(in,group="some",itrace=j,ntrace=ntrace)
			num_traces_in = size(d[:,:],2)
			for k = 1 : num_traces_in
				itrace = (h[k].isx - min_isx)*nx2*nx3*nx4 + (h[k].isy - min_isy)*nx3*nx4 + (h[k].ih - min_ih)*nx4 + h[k].iaz - min_iaz + 1
				if (h[k].isx >= min_isx && h[k].isx <= max_isx && h[k].isy >= min_isy && h[k].isy <= max_isy && h[k].ih >= min_ih && h[k].ih <= max_ih && h[k].iaz >= min_iaz && h[k].iaz <= max_iaz)
					h[k].tracenum = convert(Int32,itrace)
					if (itrace > 0 && itrace <= nx_out)
						position_d = 4*nt*(itrace - 1)
						seek(stream_d,position_d)
						write(stream_d,d[:,k])
						PutHeader(stream_h,h[k],itrace)
					end
				end
			end
			j += num_traces_in
		end
	elseif (style=="gxgyhaz")
		j = 1
		while j <= nx_in
			d,h,e = SeisRead(in,group="some",itrace=j,ntrace=ntrace)
			num_traces_in = size(d[:,:],2)
			for k = 1 : num_traces_in
				itrace = (h[k].igx - min_igx)*nx2*nx3*nx4 + (h[k].igy - min_igy)*nx3*nx4 + (h[k].ih - min_ih)*nx4 + h[k].iaz - min_iaz + 1
				if (h[k].igx >= min_igx && h[k].igx <= max_igx && h[k].igy >= min_igy && h[k].igy <= max_igy && h[k].ih >= min_ih && h[k].ih <= max_ih && h[k].iaz >= min_iaz && h[k].iaz <= max_iaz)
					h[k].tracenum = convert(Int32,itrace)
					if (itrace > 0 && itrace <= nx_out)
						position_d = 4*nt*(itrace - 1)
						seek(stream_d,position_d)
						write(stream_d,d[:,k])
						PutHeader(stream_h,h[k],itrace)
					end
				end
			end
			j += num_traces_in
		end
	end
	close(stream_d)
	close(stream_h)

end


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

module Sort

import ..@__MODULE__, ..parentmodule
const Base = parentmodule(@__MODULE__)
using .Base.Order
using .Base: copymutable, LinearIndices, length, (:),
    eachindex, axes, first, last, similar, zip, OrdinalRange,
    AbstractVector, @inbounds, AbstractRange, @eval, @inline, Vector, @noinline,
    AbstractMatrix, AbstractUnitRange, isless, identity, eltype, >, <, <=, >=, |, +, -, *, !,
    extrema, sub_with_overflow, add_with_overflow, oneunit, div, getindex, setindex!,
    length, resize!, fill, Missing, require_one_based_indexing, keytype

using .Base: >>>, !==

import .Base:
    sort,
    sort!,
    issorted,
    sortperm,
    to_indices

export # also exported by Base
    # order-only:
    issorted,
    searchsorted,
    searchsortedfirst,
    searchsortedlast,
    # order & algorithm:
    sort,
    sort!,
    sortperm,
    sortperm!,
    partialsort,
    partialsort!,
    partialsortperm,
    partialsortperm!,
    # algorithms:
    InsertionSort,
    QuickSort,
    MergeSort,
    PartialQuickSort

export # not exported by Base
    Algorithm,
    DEFAULT_UNSTABLE,
    DEFAULT_STABLE,
    SMALL_ALGORITHM,
    SMALL_THRESHOLD

## functions requiring only ordering ##

function issorted(itr, order::Ordering)
    y = iterate(itr)
    y === nothing && return true
    prev, state = y
    y = iterate(itr, state)
    while y !== nothing
        this, state = y
        lt(order, this, prev) && return false
        prev = this
        y = iterate(itr, state)
    end
    return true
end

"""
    issorted(v, lt=isless, by=identity, rev:Bool=false, order::Ordering=Forward)

Test whether a vector is in sorted order. The `lt`, `by` and `rev` keywords modify what
order is considered to be sorted just as they do for [`sort`](@ref).

# Examples
```jldoctest
julia> issorted([1, 2, 3])
true

julia> issorted([(1, "b"), (2, "a")], by = x -> x[1])
true

julia> issorted([(1, "b"), (2, "a")], by = x -> x[2])
false

julia> issorted([(1, "b"), (2, "a")], by = x -> x[2], rev=true)
true
```
"""
issorted(itr;
    lt=isless, by=identity, rev::Union{Bool,Nothing}=nothing, order::Ordering=Forward) =
    issorted(itr, ord(lt,by,rev,order))

function partialsort!(v::AbstractVector, k::Union{Integer,OrdinalRange}, o::Ordering)
    inds = axes(v, 1)
    sort!(v, first(inds), last(inds), PartialQuickSort(k), o)
    maybeview(v, k)
end

maybeview(v, k) = view(v, k)
maybeview(v, k::Integer) = v[k]

"""
    partialsort!(v, k; by=<transform>, lt=<comparison>, rev=false)

Partially sort the vector `v` in place, according to the order specified by `by`, `lt` and
`rev` so that the value at index `k` (or range of adjacent values if `k` is a range) occurs
at the position where it would appear if the array were fully sorted via a non-stable
algorithm. If `k` is a single index, that value is returned; if `k` is a range, an array of
values at those indices is returned. Note that `partialsort!` does not fully sort the input
array.

# Examples
```jldoctest
julia> a = [1, 2, 4, 3, 4]
5-element Array{Int64,1}:
 1
 2
 4
 3
 4

julia> partialsort!(a, 4)
4

julia> a
5-element Array{Int64,1}:
 1
 2
 3
 4
 4

julia> a = [1, 2, 4, 3, 4]
5-element Array{Int64,1}:
 1
 2
 4
 3
 4

julia> partialsort!(a, 4, rev=true)
2

julia> a
5-element Array{Int64,1}:
 4
 4
 3
 2
 1
```
"""
partialsort!(v::AbstractVector, k::Union{Integer,OrdinalRange};
             lt=isless, by=identity, rev::Union{Bool,Nothing}=nothing, order::Ordering=Forward) =
    partialsort!(v, k, ord(lt,by,rev,order))

"""
    partialsort(v, k, by=<transform>, lt=<comparison>, rev=false)

Variant of [`partialsort!`](@ref) which copies `v` before partially sorting it, thereby returning the
same thing as `partialsort!` but leaving `v` unmodified.
"""
partialsort(v::AbstractVector, k::Union{Integer,OrdinalRange}; kws...) =
    partialsort!(copymutable(v), k; kws...)

# This implementation of `midpoint` is performance-optimized but safe
# only if `lo <= hi`.
midpoint(lo::T, hi::T) where T<:Integer = lo + ((hi - lo) >>> 0x01)
midpoint(lo::Integer, hi::Integer) = midpoint(promote(lo, hi)...)

# reference on sorted binary search:
#   http://www.tbray.org/ongoing/When/200x/2003/03/22/Binary

# index of the first value of vector a that is greater than or equal to x;
# returns length(v)+1 if x is greater than all values in v.
function searchsortedfirst(v::AbstractVector, x, lo::T, hi::T, o::Ordering)::keytype(v) where T<:Integer
    u = T(1)
    lo = lo - u
    hi = hi + u
    @inbounds while lo < hi - u
        m = midpoint(lo, hi)
        if lt(o, v[m], x)
            lo = m
        else
            hi = m
        end
    end
    return hi
end

# index of the last value of vector a that is less than or equal to x;
# returns 0 if x is less than all values of v.
function searchsortedlast(v::AbstractVector, x, lo::T, hi::T, o::Ordering)::keytype(v) where T<:Integer
    u = T(1)
    lo = lo - u
    hi = hi + u
    @inbounds while lo < hi - u
        m = midpoint(lo, hi)
        if lt(o, x, v[m])
            hi = m
        else
            lo = m
        end
    end
    return lo
end

# returns the range of indices of v equal to x
# if v does not contain x, returns a 0-length range
# indicating the insertion point of x
function searchsorted(v::AbstractVector, x, ilo::T, ihi::T, o::Ordering)::UnitRange{keytype(v)} where T<:Integer
    u = T(1)
    lo = ilo - u
    hi = ihi + u
    @inbounds while lo < hi - u
        m = midpoint(lo, hi)
        if lt(o, v[m], x)
            lo = m
        elseif lt(o, x, v[m])
            hi = m
        else
            a = searchsortedfirst(v, x, max(lo,ilo), m, o)
            b = searchsortedlast(v, x, m, min(hi,ihi), o)
            return a : b
        end
    end
    return (lo + 1) : (hi - 1)
end

function searchsortedlast(a::AbstractRange{<:Real}, x::Real, o::DirectOrdering)::keytype(a)
    require_one_based_indexing(a)
    if step(a) == 0
        lt(o, x, first(a)) ? 0 : length(a)
    else
        n = round(Integer, clamp((x - first(a)) / step(a) + 1, 1, length(a)))
        lt(o, x, a[n]) ? n - 1 : n
    end
end

function searchsortedfirst(a::AbstractRange{<:Real}, x::Real, o::DirectOrdering)::keytype(a)
    require_one_based_indexing(a)
    if step(a) == 0
        lt(o, first(a), x) ? length(a) + 1 : 1
    else
        n = round(Integer, clamp((x - first(a)) / step(a) + 1, 1, length(a)))
        lt(o, a[n] ,x) ? n + 1 : n
    end
end

function searchsortedlast(a::AbstractRange{<:Integer}, x::Real, o::DirectOrdering)::keytype(a)
    require_one_based_indexing(a)
    h = step(a)
    if h == 0
        lt(o, x, first(a)) ? 0 : length(a)
    elseif h > 0 && x < first(a)
        firstindex(a) - 1
    elseif h > 0 && x >= last(a)
        lastindex(a)
    elseif h < 0 && x > first(a)
        firstindex(a) - 1
    elseif h < 0 && x <= last(a)
        lastindex(a)
    else
        if o isa ForwardOrdering
            fld(floor(Integer, x) - first(a), h) + 1
        else
            fld(ceil(Integer, x) - first(a), h) + 1
        end
    end
end

function searchsortedfirst(a::AbstractRange{<:Integer}, x::Real, o::DirectOrdering)::keytype(a)
    require_one_based_indexing(a)
    h = step(a)
    if h == 0
        lt(o, first(a), x) ? length(a)+1 : 1
    elseif h > 0 && x <= first(a)
        firstindex(a)
    elseif h > 0 && x > last(a)
        lastindex(a) + 1
    elseif h < 0 && x >= first(a)
        firstindex(a)
    elseif h < 0 && x < last(a)
        lastindex(a) + 1
    else
        if o isa ForwardOrdering
            -fld(floor(Integer, -x) + Signed(first(a)), h) + 1
        else
            -fld(ceil(Integer, -x) + Signed(first(a)), h) + 1
        end
    end
end

function searchsortedfirst(a::AbstractRange{<:Integer}, x::Unsigned, o::DirectOrdering)::keytype(a)
    require_one_based_indexing(a)
    if lt(o, first(a), x)
        if step(a) == 0
            length(a) + 1
        else
            min(cld(x - first(a), step(a)), length(a)) + 1
        end
    else
        1
    end
end

function searchsortedlast(a::AbstractRange{<:Integer}, x::Unsigned, o::DirectOrdering)::keytype(a)
    require_one_based_indexing(a)
    if lt(o, x, first(a))
        0
    elseif step(a) == 0
        length(a)
    else
        min(fld(x - first(a), step(a)) + 1, length(a))
    end
end

searchsorted(a::AbstractRange{<:Real}, x::Real, o::DirectOrdering) =
    searchsortedfirst(a, x, o) : searchsortedlast(a, x, o)

for s in [:searchsortedfirst, :searchsortedlast, :searchsorted]
    @eval begin
        $s(v::AbstractVector, x, o::Ordering) = (inds = axes(v, 1); $s(v,x,first(inds),last(inds),o))
        $s(v::AbstractVector, x;
           lt=isless, by=identity, rev::Union{Bool,Nothing}=nothing, order::Ordering=Forward) =
            $s(v,x,ord(lt,by,rev,order))
    end
end

"""
    searchsorted(a, x; by=<transform>, lt=<comparison>, rev=false)

Return the range of indices of `a` which compare as equal to `x` (using binary search)
according to the order specified by the `by`, `lt` and `rev` keywords, assuming that `a`
is already sorted in that order. Return an empty range located at the insertion point
if `a` does not contain values equal to `x`.

# Examples
```jldoctest
julia> searchsorted([1, 2, 4, 5, 5, 7], 4) # single match
3:3

julia> searchsorted([1, 2, 4, 5, 5, 7], 5) # multiple matches
4:5

julia> searchsorted([1, 2, 4, 5, 5, 7], 3) # no match, insert in the middle
3:2

julia> searchsorted([1, 2, 4, 5, 5, 7], 9) # no match, insert at end
7:6

julia> searchsorted([1, 2, 4, 5, 5, 7], 0) # no match, insert at start
1:0
```
""" searchsorted

"""
    searchsortedfirst(a, x; by=<transform>, lt=<comparison>, rev=false)

Return the index of the first value in `a` greater than or equal to `x`, according to the
specified order. Return `length(a) + 1` if `x` is greater than all values in `a`.
`a` is assumed to be sorted.

# Examples
```jldoctest
julia> searchsortedfirst([1, 2, 4, 5, 5, 7], 4) # single match
3

julia> searchsortedfirst([1, 2, 4, 5, 5, 7], 5) # multiple matches
4

julia> searchsortedfirst([1, 2, 4, 5, 5, 7], 3) # no match, insert in the middle
3

julia> searchsortedfirst([1, 2, 4, 5, 5, 7], 9) # no match, insert at end
7

julia> searchsortedfirst([1, 2, 4, 5, 5, 7], 0) # no match, insert at start
1
```
""" searchsortedfirst

"""
    searchsortedlast(a, x; by=<transform>, lt=<comparison>, rev=false)

Return the index of the last value in `a` less than or equal to `x`, according to the
specified order. Return `0` if `x` is less than all values in `a`. `a` is assumed to
be sorted.

# Examples
```jldoctest
julia> searchsortedlast([1, 2, 4, 5, 5, 7], 4) # single match
3

julia> searchsortedlast([1, 2, 4, 5, 5, 7], 5) # multiple matches
5

julia> searchsortedlast([1, 2, 4, 5, 5, 7], 3) # no match, insert in the middle
2

julia> searchsortedlast([1, 2, 4, 5, 5, 7], 9) # no match, insert at end
6

julia> searchsortedlast([1, 2, 4, 5, 5, 7], 0) # no match, insert at start
0
```
""" searchsortedlast

## sorting algorithms ##

abstract type Algorithm end

struct InsertionSortAlg <: Algorithm end
struct QuickSortAlg     <: Algorithm end
struct MergeSortAlg     <: Algorithm end

"""
    PartialQuickSort{T <: Union{Integer,OrdinalRange}}

Indicate that a sorting function should use the partial quick sort
algorithm. Partial quick sort returns the smallest `k` elements sorted from smallest
to largest, finding them and sorting them using [`QuickSort`](@ref).

Characteristics:
  * *not stable*: does not preserve the ordering of elements which
    compare equal (e.g. "a" and "A" in a sort of letters which
    ignores case).
  * *in-place* in memory.
  * *divide-and-conquer*: sort strategy similar to [`MergeSort`](@ref).
"""
struct PartialQuickSort{T <: Union{Integer,OrdinalRange}} <: Algorithm
    k::T
end

"""
    InsertionSort

Indicate that a sorting function should use the insertion sort
algorithm. Insertion sort traverses the collection one element
at a time, inserting each element into its correct, sorted position in
the output list.

Characteristics:
  * *stable*: preserves the ordering of elements which
    compare equal (e.g. "a" and "A" in a sort of letters
    which ignores case).
  * *in-place* in memory.
  * *quadratic performance* in the number of elements to be sorted:
    it is well-suited to small collections but should not be used for large ones.
"""
const InsertionSort = InsertionSortAlg()
"""
    QuickSort

Indicate that a sorting function should use the quick sort
algorithm, which is *not* stable.

Characteristics:
  * *not stable*: does not preserve the ordering of elements which
    compare equal (e.g. "a" and "A" in a sort of letters which
    ignores case).
  * *in-place* in memory.
  * *divide-and-conquer*: sort strategy similar to [`MergeSort`](@ref).
  * *good performance* for large collections.
"""
const QuickSort     = QuickSortAlg()
"""
    MergeSort

Indicate that a sorting function should use the merge sort
algorithm. Merge sort divides the collection into
subcollections and repeatedly merges them, sorting each
subcollection at each step, until the entire
collection has been recombined in sorted form.

Characteristics:
  * *stable*: preserves the ordering of elements which compare
    equal (e.g. "a" and "A" in a sort of letters which ignores
    case).
  * *not in-place* in memory.
  * *divide-and-conquer* sort strategy.
"""
const MergeSort     = MergeSortAlg()

const DEFAULT_UNSTABLE = QuickSort
const DEFAULT_STABLE   = MergeSort
const SMALL_ALGORITHM  = InsertionSort
const SMALL_THRESHOLD  = 20

function sort!(v::AbstractVector, lo::Integer, hi::Integer, ::InsertionSortAlg, o::Ordering)
    @inbounds for i = lo+1:hi
        j = i
        x = v[i]
        while j > lo
            if lt(o, x, v[j-1])
                v[j] = v[j-1]
                j -= 1
                continue
            end
            break
        end
        v[j] = x
    end
    return v
end

# selectpivot!
#
# Given 3 locations in an array (lo, mi, and hi), sort v[lo], v[mi], v[hi]) and
# choose the middle value as a pivot
#
# Upon return, the pivot is in v[lo], and v[hi] is guaranteed to be
# greater than the pivot

@inline function selectpivot!(v::AbstractVector, lo::Integer, hi::Integer, o::Ordering)
    @inbounds begin
        mi = midpoint(lo, hi)

        # sort v[mi] <= v[lo] <= v[hi] such that the pivot is immediately in place
        if lt(o, v[lo], v[mi])
            v[mi], v[lo] = v[lo], v[mi]
        end

        if lt(o, v[hi], v[lo])
            if lt(o, v[hi], v[mi])
                v[hi], v[lo], v[mi] = v[lo], v[mi], v[hi]
            else
                v[hi], v[lo] = v[lo], v[hi]
            end
        end

        # return the pivot
        return v[lo]
    end
end

# partition!
#
# select a pivot, and partition v according to the pivot

function partition!(v::AbstractVector, lo::Integer, hi::Integer, o::Ordering)
    pivot = selectpivot!(v, lo, hi, o)
    # pivot == v[lo], v[hi] > pivot
    i, j = lo, hi
    @inbounds while true
        i += 1; j -= 1
        while lt(o, v[i], pivot); i += 1; end;
        while lt(o, pivot, v[j]); j -= 1; end;
        i >= j && break
        v[i], v[j] = v[j], v[i]
    end
    v[j], v[lo] = pivot, v[j]

    # v[j] == pivot
    # v[k] >= pivot for k > j
    # v[i] <= pivot for i < j
    return j
end

function sort!(v::AbstractVector, lo::Integer, hi::Integer, a::QuickSortAlg, o::Ordering)
    @inbounds while lo < hi
        hi-lo <= SMALL_THRESHOLD && return sort!(v, lo, hi, SMALL_ALGORITHM, o)
        j = partition!(v, lo, hi, o)
        if j-lo < hi-j
            # recurse on the smaller chunk
            # this is necessary to preserve O(log(n))
            # stack space in the worst case (rather than O(n))
            lo < (j-1) && sort!(v, lo, j-1, a, o)
            lo = j+1
        else
            j+1 < hi && sort!(v, j+1, hi, a, o)
            hi = j-1
        end
    end
    return v
end

function sort!(v::AbstractVector, lo::Integer, hi::Integer, a::MergeSortAlg, o::Ordering, t=similar(v,0))
    @inbounds if lo < hi
        hi-lo <= SMALL_THRESHOLD && return sort!(v, lo, hi, SMALL_ALGORITHM, o)

        m = midpoint(lo, hi)
        (length(t) < m-lo+1) && resize!(t, m-lo+1)

        sort!(v, lo,  m,  a, o, t)
        sort!(v, m+1, hi, a, o, t)

        i, j = 1, lo
        while j <= m
            t[i] = v[j]
            i += 1
            j += 1
        end

        i, k = 1, lo
        while k < j <= hi
            if lt(o, v[j], t[i])
                v[k] = v[j]
                j += 1
            else
                v[k] = t[i]
                i += 1
            end
            k += 1
        end
        while k < j
            v[k] = t[i]
            k += 1
            i += 1
        end
    end

    return v
end

function sort!(v::AbstractVector, lo::Integer, hi::Integer, a::PartialQuickSort{<:Integer},
               o::Ordering)
    @inbounds while lo < hi
        hi-lo <= SMALL_THRESHOLD && return sort!(v, lo, hi, SMALL_ALGORITHM, o)
        j = partition!(v, lo, hi, o)
        if j >= a.k
            # we don't need to sort anything bigger than j
            hi = j-1
        elseif j-lo < hi-j
            # recurse on the smaller chunk
            # this is necessary to preserve O(log(n))
            # stack space in the worst case (rather than O(n))
            lo < (j-1) && sort!(v, lo, j-1, a, o)
            lo = j+1
        else
            (j+1) < hi && sort!(v, j+1, hi, a, o)
            hi = j-1
        end
    end
    return v
end

function sort!(v::AbstractVector, lo::Integer, hi::Integer, a::PartialQuickSort{T},
               o::Ordering) where T<:OrdinalRange
    @inbounds while lo < hi
        hi-lo <= SMALL_THRESHOLD && return sort!(v, lo, hi, SMALL_ALGORITHM, o)
        j = partition!(v, lo, hi, o)

        if j <= first(a.k)
            lo = j+1
        elseif j >= last(a.k)
            hi = j-1
        else
            if j-lo < hi-j
                lo < (j-1) && sort!(v, lo, j-1, a, o)
                lo = j+1
            else
                hi > (j+1) && sort!(v, j+1, hi, a, o)
                hi = j-1
            end
        end
    end
    return v
end

## generic sorting methods ##

defalg(v::AbstractArray) = DEFAULT_STABLE
defalg(v::AbstractArray{<:Union{Number, Missing}}) = DEFAULT_UNSTABLE

function sort!(v::AbstractVector, alg::Algorithm, order::Ordering)
    inds = axes(v,1)
    sort!(v,first(inds),last(inds),alg,order)
end

"""
    sort!(v; alg::Algorithm=defalg(v), lt=isless, by=identity, rev::Bool=false, order::Ordering=Forward)

Sort the vector `v` in place. [`QuickSort`](@ref) is used by default for numeric arrays while
[`MergeSort`](@ref) is used for other arrays. You can specify an algorithm to use via the `alg`
keyword (see [Sorting Algorithms](@ref) for available algorithms). The `by` keyword lets you provide
a function that will be applied to each element before comparison; the `lt` keyword allows
providing a custom "less than" function; use `rev=true` to reverse the sorting order. These
options are independent and can be used together in all possible combinations: if both `by`
and `lt` are specified, the `lt` function is applied to the result of the `by` function;
`rev=true` reverses whatever ordering specified via the `by` and `lt` keywords.

# Examples
```jldoctest
julia> v = [3, 1, 2]; sort!(v); v
3-element Array{Int64,1}:
 1
 2
 3

julia> v = [3, 1, 2]; sort!(v, rev = true); v
3-element Array{Int64,1}:
 3
 2
 1

julia> v = [(1, "c"), (3, "a"), (2, "b")]; sort!(v, by = x -> x[1]); v
3-element Array{Tuple{Int64,String},1}:
 (1, "c")
 (2, "b")
 (3, "a")

julia> v = [(1, "c"), (3, "a"), (2, "b")]; sort!(v, by = x -> x[2]); v
3-element Array{Tuple{Int64,String},1}:
 (3, "a")
 (2, "b")
 (1, "c")
```
"""
function sort!(v::AbstractVector;
               alg::Algorithm=defalg(v),
               lt=isless,
               by=identity,
               rev::Union{Bool,Nothing}=nothing,
               order::Ordering=Forward)
    ordr = ord(lt,by,rev,order)
    if (ordr === Forward || ordr === Reverse) && eltype(v)<:Integer
        n = length(v)
        if n > 1
            min, max = extrema(v)
            (diff, o1) = sub_with_overflow(max, min)
            (rangelen, o2) = add_with_overflow(diff, oneunit(diff))
            if !o1 && !o2 && rangelen < div(n,2)
                return sort_int_range!(v, rangelen, min, ordr === Reverse ? reverse : identity)
            end
        end
    end
    sort!(v, alg, ordr)
end

# sort! for vectors of few unique integers
function sort_int_range!(x::AbstractVector{<:Integer}, rangelen, minval, maybereverse)
    offs = 1 - minval

    where = fill(0, rangelen)
    @inbounds for i = eachindex(x)
        where[x[i] + offs] += 1
    end

    idx = firstindex(x)
    @inbounds for i = maybereverse(1:rangelen)
        lastidx = idx + where[i] - 1
        val = i-offs
        for j = idx:lastidx
            x[j] = val
        end
        idx = lastidx + 1
    end

    return x
end

"""
    sort(v; alg::Algorithm=defalg(v), lt=isless, by=identity, rev::Bool=false, order::Ordering=Forward)

Variant of [`sort!`](@ref) that returns a sorted copy of `v` leaving `v` itself unmodified.

# Examples
```jldoctest
julia> v = [3, 1, 2];

julia> sort(v)
3-element Array{Int64,1}:
 1
 2
 3

julia> v
3-element Array{Int64,1}:
 3
 1
 2
```
"""
sort(v::AbstractVector; kws...) = sort!(copymutable(v); kws...)

## partialsortperm: the permutation to sort the first k elements of an array ##

"""
    partialsortperm(v, k; by=<transform>, lt=<comparison>, rev=false)

Return a partial permutation `I` of the vector `v`, so that `v[I]` returns values of a fully
sorted version of `v` at index `k`. If `k` is a range, a vector of indices is returned; if
`k` is an integer, a single index is returned. The order is specified using the same
keywords as `sort!`. The permutation is stable, meaning that indices of equal elements
appear in ascending order.

Note that this function is equivalent to, but more efficient than, calling `sortperm(...)[k]`.

# Examples
```jldoctest
julia> v = [3, 1, 2, 1];

julia> v[partialsortperm(v, 1)]
1

julia> p = partialsortperm(v, 1:3)
3-element view(::Array{Int64,1}, 1:3) with eltype Int64:
 2
 4
 3

julia> v[p]
3-element Array{Int64,1}:
 1
 1
 2
```
"""
partialsortperm(v::AbstractVector, k::Union{Integer,OrdinalRange}; kwargs...) =
    partialsortperm!(similar(Vector{eltype(k)}, axes(v,1)), v, k; kwargs..., initialized=false)

"""
    partialsortperm!(ix, v, k; by=<transform>, lt=<comparison>, rev=false, initialized=false)

Like [`partialsortperm`](@ref), but accepts a preallocated index vector `ix` the same size as
`v`, which is used to store (a permutation of) the indices of `v`.

If the index vector `ix` is initialized with the indices of `v` (or a permutation thereof), `initialized` should be set to
`true`.

If `initialized` is `false` (the default), then `ix` is initialized to contain the indices of `v`.

If `initialized` is `true`, but `ix` does not contain (a permutation of) the indices of `v`, the behavior of
`partialsortperm!` is undefined.

(Typically, the indices of `v` will be `1:length(v)`, although if `v` has an alternative array type
with non-one-based indices, such as an `OffsetArray`, `ix` must also be an `OffsetArray` with the same
indices, and must contain as values (a permutation of) these same indices.)

Upon return, `ix` is guaranteed to have the indices `k` in their sorted positions, such that

```julia
partialsortperm!(ix, v, k);
v[ix[k]] == partialsort(v, k)
```

The return value is the `k`th element of `ix` if `k` is an integer, or view into `ix` if `k` is
a range.

# Examples
```jldoctest
julia> v = [3, 1, 2, 1];

julia> ix = Vector{Int}(undef, 4);

julia> partialsortperm!(ix, v, 1)
2

julia> ix = [1:4;];

julia> partialsortperm!(ix, v, 2:3, initialized=true)
2-element view(::Array{Int64,1}, 2:3) with eltype Int64:
 4
 3
```
 """
function partialsortperm!(ix::AbstractVector{<:Integer}, v::AbstractVector,
                          k::Union{Integer, OrdinalRange};
                          lt::Function=isless,
                          by::Function=identity,
                          rev::Union{Bool,Nothing}=nothing,
                          order::Ordering=Forward,
                          initialized::Bool=false)
    if axes(ix,1) != axes(v,1)
        throw(ArgumentError("The index vector is used as a workspace and must have the " *
                            "same length/indices as the source vector, $(axes(ix,1)) != $(axes(v,1))"))
    end
    if !initialized
        @inbounds for i = axes(ix,1)
            ix[i] = i
        end
    end

    # do partial quicksort
    sort!(ix, PartialQuickSort(k), Perm(ord(lt, by, rev, order), v))

    maybeview(ix, k)
end

## sortperm: the permutation to sort an array ##

"""
    sortperm(v; alg::Algorithm=DEFAULT_UNSTABLE, lt=isless, by=identity, rev::Bool=false, order::Ordering=Forward)

Return a permutation vector `I` that puts `v[I]` in sorted order. The order is specified
using the same keywords as [`sort!`](@ref). The permutation is guaranteed to be stable even
if the sorting algorithm is unstable, meaning that indices of equal elements appear in
ascending order.

See also [`sortperm!`](@ref).

# Examples
```jldoctest
julia> v = [3, 1, 2];

julia> p = sortperm(v)
3-element Array{Int64,1}:
 2
 3
 1

julia> v[p]
3-element Array{Int64,1}:
 1
 2
 3
```
"""
function sortperm(v::AbstractVector;
                  alg::Algorithm=DEFAULT_UNSTABLE,
                  lt=isless,
                  by=identity,
                  rev::Union{Bool,Nothing}=nothing,
                  order::Ordering=Forward)
    ordr = ord(lt,by,rev,order)
    if ordr === Forward && isa(v,Vector) && eltype(v)<:Integer
        n = length(v)
        if n > 1
            min, max = extrema(v)
            (diff, o1) = sub_with_overflow(max, min)
            (rangelen, o2) = add_with_overflow(diff, oneunit(diff))
            if !o1 && !o2 && rangelen < div(n,2)
                return sortperm_int_range(v, rangelen, min)
            end
        end
    end
    ax = axes(v, 1)
    p = similar(Vector{eltype(ax)}, ax)
    for (i,ind) in zip(eachindex(p), ax)
        p[i] = ind
    end
    sort!(p, alg, Perm(ordr,v))
end

"""
    sortperm!(ix, v; alg::Algorithm=DEFAULT_UNSTABLE, lt=isless, by=identity, rev::Bool=false, order::Ordering=Forward, initialized::Bool=false)

Like [`sortperm`](@ref), but accepts a preallocated index vector `ix`.  If `initialized` is `false`
(the default), `ix` is initialized to contain the values `1:length(v)`.

# Examples
```jldoctest
julia> v = [3, 1, 2]; p = zeros(Int, 3);

julia> sortperm!(p, v); p
3-element Array{Int64,1}:
 2
 3
 1

julia> v[p]
3-element Array{Int64,1}:
 1
 2
 3
```
"""
function sortperm!(x::AbstractVector{<:Integer}, v::AbstractVector;
                   alg::Algorithm=DEFAULT_UNSTABLE,
                   lt=isless,
                   by=identity,
                   rev::Union{Bool,Nothing}=nothing,
                   order::Ordering=Forward,
                   initialized::Bool=false)
    if axes(x,1) != axes(v,1)
        throw(ArgumentError("index vector must have the same length/indices as the source vector, $(axes(x,1)) != $(axes(v,1))"))
    end
    if !initialized
        @inbounds for i = axes(v,1)
            x[i] = i
        end
    end
    sort!(x, alg, Perm(ord(lt,by,rev,order),v))
end

# sortperm for vectors of few unique integers
function sortperm_int_range(x::Vector{<:Integer}, rangelen, minval)
    offs = 1 - minval
    n = length(x)

    where = fill(0, rangelen+1)
    where[1] = 1
    @inbounds for i = 1:n
        where[x[i] + offs + 1] += 1
    end

    #cumsum!(where, where)
    @inbounds for i = 2:length(where)
        where[i] += where[i-1]
    end

    P = Vector{Int}(undef, n)
    @inbounds for i = 1:n
        label = x[i] + offs
        P[where[label]] = i
        where[label] += 1
    end

    return P
end

## sorting multi-dimensional arrays ##

"""
    sort(A; dims::Integer, alg::Algorithm=DEFAULT_UNSTABLE, lt=isless, by=identity, rev::Bool=false, order::Ordering=Forward)

Sort a multidimensional array `A` along the given dimension.
See [`sort!`](@ref) for a description of possible
keyword arguments.

To sort slices of an array, refer to [`sortslices`](@ref).

# Examples
```jldoctest
julia> A = [4 3; 1 2]
2×2 Array{Int64,2}:
 4  3
 1  2

julia> sort(A, dims = 1)
2×2 Array{Int64,2}:
 1  2
 4  3

julia> sort(A, dims = 2)
2×2 Array{Int64,2}:
 3  4
 1  2
```
"""
function sort(A::AbstractArray;
              dims::Integer,
              alg::Algorithm=DEFAULT_UNSTABLE,
              lt=isless,
              by=identity,
              rev::Union{Bool,Nothing}=nothing,
              order::Ordering=Forward)
    dim = dims
    order = ord(lt,by,rev,order)
    n = length(axes(A, dim))
    if dim != 1
        pdims = (dim, setdiff(1:ndims(A), dim)...)  # put the selected dimension first
        Ap = permutedims(A, pdims)
        Av = vec(Ap)
        sort_chunks!(Av, n, alg, order)
        permutedims(Ap, invperm(pdims))
    else
        Av = A[:]
        sort_chunks!(Av, n, alg, order)
        reshape(Av, axes(A))
    end
end

@noinline function sort_chunks!(Av, n, alg, order)
    inds = LinearIndices(Av)
    for s = first(inds):n:last(inds)
        sort!(Av, s, s+n-1, alg, order)
    end
    Av
end

"""
    sort!(A; dims::Integer, alg::Algorithm=defalg(A), lt=isless, by=identity, rev::Bool=false, order::Ordering=Forward)

Sort the multidimensional array `A` along dimension `dims`.
See [`sort!`](@ref) for a description of possible keyword arguments.

To sort slices of an array, refer to [`sortslices`](@ref).

!!! compat "Julia 1.1"
    This function requires at least Julia 1.1.

# Examples
```jldoctest
julia> A = [4 3; 1 2]
2×2 Array{Int64,2}:
 4  3
 1  2

julia> sort!(A, dims = 1); A
2×2 Array{Int64,2}:
 1  2
 4  3

julia> sort!(A, dims = 2); A
2×2 Array{Int64,2}:
 1  2
 3  4
```
"""
function sort!(A::AbstractArray;
               dims::Integer,
               alg::Algorithm=defalg(A),
               lt=isless,
               by=identity,
               rev::Union{Bool,Nothing}=nothing,
               order::Ordering=Forward)
    ordr = ord(lt, by, rev, order)
    nd = ndims(A)
    k = dims

    1 <= k <= nd || throw(ArgumentError("dimension out of range"))

    remdims = ntuple(i -> i == k ? 1 : size(A, i), nd)
    for idx in CartesianIndices(remdims)
        Av = view(A, ntuple(i -> i == k ? Colon() : idx[i], nd)...)
        sort!(Av, alg, ordr)
    end
    A
end

## fast clever sorting for floats ##

module Float
using ..Sort
using ...Order
using ..Base: @inbounds, AbstractVector, Vector, last, axes

import Core.Intrinsics: slt_int
import ..Sort: sort!
import ...Order: lt, DirectOrdering

const Floats = Union{Float32,Float64}

struct Left <: Ordering end
struct Right <: Ordering end

left(::DirectOrdering) = Left()
right(::DirectOrdering) = Right()

left(o::Perm) = Perm(left(o.order), o.data)
right(o::Perm) = Perm(right(o.order), o.data)

lt(::Left, x::T, y::T) where {T<:Floats} = slt_int(y, x)
lt(::Right, x::T, y::T) where {T<:Floats} = slt_int(x, y)

isnan(o::DirectOrdering, x::Floats) = (x!=x)
isnan(o::Perm, i::Integer) = isnan(o.order,o.data[i])

function nans2left!(v::AbstractVector, o::Ordering, lo::Integer=first(axes(v,1)), hi::Integer=last(axes(v,1)))
    i = lo
    @inbounds while i <= hi && isnan(o,v[i])
        i += 1
    end
    j = i + 1
    @inbounds while j <= hi
        if isnan(o,v[j])
            v[i], v[j] = v[j], v[i]
            i += 1
        end
        j += 1
    end
    return i, hi
end
function nans2right!(v::AbstractVector, o::Ordering, lo::Integer=first(axes(v,1)), hi::Integer=last(axes(v,1)))
    i = hi
    @inbounds while lo <= i && isnan(o,v[i])
        i -= 1
    end
    j = i - 1
    @inbounds while lo <= j
        if isnan(o,v[j])
            v[i], v[j] = v[j], v[i]
            i -= 1
        end
        j -= 1
    end
    return lo, i
end

nans2end!(v::AbstractVector, o::ForwardOrdering) = nans2right!(v,o)
nans2end!(v::AbstractVector, o::ReverseOrdering) = nans2left!(v,o)
nans2end!(v::AbstractVector{<:Integer}, o::Perm{<:ForwardOrdering}) = nans2right!(v,o)
nans2end!(v::AbstractVector{<:Integer}, o::Perm{<:ReverseOrdering}) = nans2left!(v,o)

issignleft(o::ForwardOrdering, x::Floats) = lt(o, x, zero(x))
issignleft(o::ReverseOrdering, x::Floats) = lt(o, x, -zero(x))
issignleft(o::Perm, i::Integer) = issignleft(o.order, o.data[i])

function fpsort!(v::AbstractVector, a::Algorithm, o::Ordering)
    i, j = lo, hi = nans2end!(v,o)
    @inbounds while true
        while i <= j &&  issignleft(o,v[i]); i += 1; end
        while i <= j && !issignleft(o,v[j]); j -= 1; end
        i <= j || break
        v[i], v[j] = v[j], v[i]
        i += 1; j -= 1
    end
    sort!(v, lo, j,  a, left(o))
    sort!(v, i,  hi, a, right(o))
    return v
end

fpsort!(v::AbstractVector, a::Sort.PartialQuickSort, o::Ordering) =
    sort!(v, first(axes(v,1)), last(axes(v,1)), a, o)

sort!(v::AbstractVector{<:Floats}, a::Algorithm, o::DirectOrdering) = fpsort!(v,a,o)
sort!(v::Vector{Int}, a::Algorithm, o::Perm{<:DirectOrdering,<:Vector{<:Floats}}) = fpsort!(v,a,o)

end # module Sort.Float

end # module Sort