text:
################################################################################
### FileIO
################################################################################

to_tag(::Type{<: DQMC}) = Val(:DQMC)
to_tag(::Type{<: DQMCParameters}) = Val(:DQMCParameters)
to_tag(::Type{<: DQMCAnalysis}) = Val(:DQMCAnalysis)
to_tag(::Type{<: MagnitudeStats}) = Val(:MagnitudeStats)

#     save_mc(filename, mc, entryname)
#
# Saves (minimal) information necessary to reconstruct a given `mc::DQMC` to a
# JLD-file `filename` under group `entryname`.
#
# When saving a simulation the default `entryname` is `MC`
function save_mc(file::JLDFile, mc::DQMC, entryname::String="DQMC")
    write(file, entryname * "/VERSION", 2)
    write(file, entryname * "/tag", "DQMC")
    write(file, entryname * "/CB", mc isa DQMC_CBTrue)
    save_parameters(file, mc.parameters, entryname * "/Parameters")
    save_analysis(file, mc.analysis, entryname * "/Analysis")
    write(file, entryname * "/conf", mc.conf)
    _save(file, mc.recorder, entryname * "/configs")
    write(file, entryname * "/last_sweep", mc.last_sweep)
    save_measurements(file, mc, entryname * "/Measurements")
    save_model(file, mc.model, entryname * "/Model")
    save_scheduler(file, mc.scheduler, entryname * "/Scheduler")
    nothing
end

CB_type(T::UnionAll) = T.body.parameters[2]
CB_type(T::DataType) = T.parameters[2]

#     load_mc(data, ::Type{<: DQMC})
#
# Loads a DQMC from a given `data` dictionary produced by `JLD.load(filename)`.
function _load(data, ::Val{:DQMC})
    if data["VERSION"] > 2
        throw(ErrorException("Failed to load DQMC version $(data["VERSION"])"))
    end

    CB = if haskey(data, "CB")
        data["CB"] ? CheckerboardTrue : CheckerboardFalse
    else CB_type(data["type"]) end
    @assert CB <: Checkerboard
    parameters = _load(data["Parameters"], Val(:DQMCParameters))
    analysis = _load(data["Analysis"], Val(:DQMCAnalysis))
    conf = data["conf"]
    recorder = _load(data["configs"], to_tag(data["configs"]))
    last_sweep = data["last_sweep"]
    model = _load(data["Model"], to_tag(data["Model"]))
    scheduler = if haskey(data, "Scheduler")
        _load(data["Scheduler"], to_tag(data["Scheduler"]))
    else
        if haskey(data["Parameters"], "global_moves") && Bool(data["Parameters"]["global_moves"])
            rate = get(data["Parameters"], "global_rate", 10)
            @warn "Replacing `global_moves = true` with GlobalFlip"
            SimpleScheduler(LocalSweep(rate), GlobalFlip())
        else
            SimpleScheduler(LocalSweep())
        end
    end

    combined_measurements = _load(data["Measurements"], Val(:Measurements))
    thermalization_measurements = combined_measurements[:TH]
    measurements = combined_measurements[:ME]

    HET = hoppingeltype(DQMC, model)
    GET = greenseltype(DQMC, model)
    HMT = hopping_matrix_type(DQMC, model)
    GMT = greens_matrix_type(DQMC, model)
    IMT = interaction_matrix_type(DQMC, model)
    stack = DQMCStack{GET, HET, GMT, HMT, IMT}()
    ut_stack = UnequalTimeStack{GET, GMT}()

    DQMC(
        CB,
        model, conf, deepcopy(conf), last_sweep,
        stack, ut_stack, scheduler,
        parameters, analysis,
        recorder, thermalization_measurements, measurements
    )
end

#   save_parameters(file::JLDFile, p::DQMCParameters, entryname="Parameters")
#
# Saves (minimal) information necessary to reconstruct a given
# `p::DQMCParameters` to a JLD-file `filename` under group `entryname`.
#
# When saving a simulation the default `entryname` is `MC/Parameters`
function save_parameters(file::JLDFile, p::DQMCParameters, entryname::String="Parameters")
    write(file, entryname * "/VERSION", 1)
    write(file, entryname * "/tag", "DQMCParameters")

    write(file, entryname * "/thermalization", p.thermalization)
    write(file, entryname * "/sweeps", p.sweeps)
    write(file, entryname * "/silent", Int(p.silent))
    write(file, entryname * "/check_sign_problem", Int(p.check_sign_problem))
    write(file, entryname * "/check_propagation_error", Int(p.check_propagation_error))
    write(file, entryname * "/safe_mult", p.safe_mult)
    write(file, entryname * "/delta_tau", p.delta_tau)
    write(file, entryname * "/beta", p.beta)
    write(file, entryname * "/slices", p.slices)
    write(file, entryname * "/measure_rate", p.measure_rate)
    write(file, entryname * "/print_rate", p.print_rate)

    nothing
end

#     load_parameters(data, ::Type{<: DQMCParameters})
#
# Loads a DQMCParameters object from a given `data` dictionary produced by
# `JLD.load(filename)`.
function _load(data, ::Val{:DQMCParameters})
    if !(data["VERSION"] == 1)
        throw(ErrorException("Failed to load DQMCParameters version $(data["VERSION"])"))
    end

    DQMCParameters(
        data["thermalization"],
        data["sweeps"],

        Bool(data["silent"]),
        Bool(data["check_sign_problem"]),
        Bool(data["check_propagation_error"]),

        data["safe_mult"],
        data["delta_tau"],
        data["beta"],
        data["slices"],

        data["measure_rate"],
        haskey(data, "print_rate") ? data["print_rate"] : 10,
    )
end

function save_analysis(file::JLDFile, a::DQMCAnalysis, entryname::String="Analysis")
    write(file, entryname * "/VERSION", 1)
    write(file, entryname * "/type", typeof(a))

    write(file, entryname * "/th_runtime", a.th_runtime)
    write(file, entryname * "/me_runtime", a.me_runtime)
    save_stats(file, a.imaginary_probability, entryname * "/imag_prob")
    save_stats(file, a.negative_probability, entryname * "/neg_prob")
    save_stats(file, a.propagation_error, entryname * "/propagation")
end
function save_stats(file::JLDFile, ms::MagnitudeStats, entryname::String="MStats")
    write(file, entryname * "/max", ms.max)
    write(file, entryname * "/min", ms.min)
    write(file, entryname * "/sum", ms.sum)
    write(file, entryname * "/count", ms.count)
end

function _load(data, ::Val{:DQMCAnalysis})
    if !(data["VERSION"] == 1)
        throw(ErrorException("Failed to load DQMCAnalysis version $(data["VERSION"])"))
    end

    DQMCAnalysis(
        th_runtime = get(data, "th_runtime", 0.0),
        me_runtime = get(data, "me_runtime", 0.0),
        imaginary_probability = load_stats(data["imag_prob"]),
        negative_probability  = load_stats(data["neg_prob"]),
        propagation_error     = load_stats(data["propagation"])
    )
end
function load_stats(data)
    MagnitudeStats(data["max"], data["min"], data["sum"], data["count"])
end


text:
 using Pkg
 Pkg.instantiate()
 Pkg.activate(".")
 using Psychometrics
 using Distributions
 using JLD2
 using StatsBase
 using DelimitedFiles
 using FileIO

 function generate_true_values()
I_operational = 1000
I_field = 25
I_total = I_field + I_operational
N = 4_000

test_operational  = 25
test_field = 5
test_length = test_field + test_operational

iter_mcmc_latent = 2_000
iter_mcmc_item = 4_000
#after how many responses update item parameter estimates
batch_size = 5
N_T = 500

    # TRUE VALUES

    ## ITEMS

    ### operational items

    log_a_dist = Normal(0.4, sqrt(0.1)); # mu= exp(0.4 + (0.5*0.1^2)), sigma^2= exp(0.4 + 0.1^2)*(exp(0.1^2) - 1)
    a_bounds = [1e-5,Inf];
    b_dist = Normal(0, 1);
    b_bounds = [-6, 6];
    items_operational = [Item(i, string("item_", i), ["math"], Parameters2PL(Product([log_a_dist, b_dist]), a_bounds, b_bounds)) for i = 1:I_operational];
    map(i ->
        begin
            i.parameters.a = exp(i.parameters.a)
            i.parameters.calibrated = true
        end,
    items_operational)

    ### field items
    quantiles_a = quantile(map(i -> i.parameters.a, items_operational), [0.2, 0.35, 0.5, 0.65, 0.8])
    quantiles_b = quantile(map(i -> i.parameters.b, items_operational), [0.2, 0.35, 0.5, 0.65, 0.8] )
    a_field = vcat([fill(q,5) for q in quantiles_a]...)
    b_field = vcat([quantiles_b for q in quantiles_a]...)

    items_field = [Item(i, string("item_", i), ["math"], Parameters2PL(Product([log_a_dist, b_dist]), a_bounds, b_bounds)) for i = (I_operational + 1) : I_total];
    map((i, a, b) ->
        begin
            i.parameters.a = a
            i.parameters.b = b
        end,
        items_field,
        a_field,
        b_field
    )
    items = vcat(items_operational, items_field)

    @save "examples/online-calibration/true values/true_items.jld2" items
    @save "examples/online-calibration/settings.jld2" I_total  I_field N test_length test_field test_operational iter_mcmc_latent iter_mcmc_item batch_size N_T

    # ## EXAMINEES

    latent_dist = Normal(0, 1)
    latent_bounds = [-Inf, Inf]
    examinees = Examinee[]
    global n=1
    for n in 1:N
        push!(examinees,  Examinee(n, string("examinee_", n), Latent1D(latent_dist, latent_bounds)))
    end

    @save "examples/online-calibration/true values/true_examinees.jld2" examinees

end

generate_true_values()

ext: jl


text:
################################################################################
#
#  Unit rank
#
################################################################################

@doc Markdown.doc"""
    unit_group_rank(O::NfOrd) -> Int

Returns the unit rank of $\mathcal O$, that is, the rank of the unit group
$\mathcal O^\times$.
"""
function unit_group_rank(O::NfOrd)
  return unit_group_rank(nf(O))
end

################################################################################
#
#  Test if units are independent
#
################################################################################

@doc Markdown.doc"""
    isindependent{T}(x::Vector{T})

Given an array of non-zero units in a number field, returns whether they
are multiplicatively independent.
"""
function isindependent(x::Vector{T}, p::Int = 32) where T
  return _isindependent(x, p)
end

function _isindependent(x::Vector{T}, p::Int = 32) where T
  K = _base_ring(x[1])

  deg = degree(K)
  r1, r2 = signature(K)
  rr = r1 + r2
  r = rr - 1 # unit rank

  conlog = Vector{Vector{arb}}(undef, length(x))

  # This can be made more memory friendly
  while true
    @assert p != 0

    q = 2
    for i in 1:length(x)
      @vprint :UnitGroup 3 "Computing conjugates with precision $p of ($i) $(length(x[i].fac))...\n"
      conlog[i] = conjugates_arb_log(x[i], p)
      for j in 1:rr
        q = max(q, bits(conlog[i][j]))
      end
    end

    A = zero_matrix(ArbField(q, cached = false), length(x), rr)

    for i in 1:length(x)
      for j in 1:rr
        A[i, j] = conlog[i][j]
        @assert radiuslttwopower(A[i, j], -p)
      end
    end

    Ar = base_ring(A)

    B = A*transpose(A)
    @vprint :UnitGroup 1 "Computing det of $(nrows(B))x$(ncols(B)) matrix with precision $(p) ... \n"
    d = det(B)

    y = (Ar(1)//Ar(r))^r * (Ar(21)//Ar(128) * log(Ar(deg))//(Ar(deg)^2))^(2*r)
    if isfinite(d) && ispositive(y - d)
      return false, p
    elseif isfinite(d) && ispositive(d)
      return true, p
    end
    p = 2*p
  end
end


text:
```

ДАНО: Где-то на неограниченном со всех сторон поле и без внутренних перегородок имеется единственный маркер.
Робот - в произвольной клетке поля.

РЕЗУЛЬТАТ: Робот - в клетке с тем маркером.

```
using HorizonSideRobots
originRobot = Robot(animate = true, "9example.sit")

# будем ходить кругом, пока не найжем этот маркер.
# 1 Вниз, 2 Вправо, 2 Вверх, 3 Влево, 3 Вниз, 4 Вправо, 4 Вверх и тд.
function master9!(robot::Robot)
    flag = false
    steps = 0
    if (ismarker(robot))
        return true
    else
        while !ismarker(robot)
            for side in [Ost, Nord, West, Sud]
                steps += 1*(side == Ost || side == West)
                for step in 1:steps
                    move!(robot,side)
                    if (ismarker(robot))
                        return true
                    end
                end
            end
        end
    end
end

master9!(originRobot)

max_stars_repo_name: Igetsuga/AlgorithmicsHomework

max_issues_repo_name: Igetsuga/AlgorithmicsHomework

max_forks_repo_name: Igetsuga/AlgorithmicsHomework


text:
using LinearAlgebra: eigen
using NLsolve: nlsolve
using .PhysicalConstants: e, ϵ₀

export IonConfiguration, ions, linear_equilibrium_positions, Anm, LinearChain,
       characteristic_length_scale, get_vibrational_modes

"""
    IonConfiguration
Physical configuration of ions. Stores a collection of ions and  information about the
interactions of their center of mass motion.
"""
abstract type IonConfiguration end

# required functions
ions(I::IonConfiguration)::Vector{Ion} = I.ions

#############################################################################################
# LinearChain - a linear Coulomb crystal
#############################################################################################

"""
    linear_equilibrium_positions(N::Int)
Returns the scaled equilibrium positions of `N` ions in a harmonic potential, assuming that
all ions have the same mass.
[ref](https://link.springer.com/content/pdf/10.1007%2Fs003400050373.pdf)
"""
function linear_equilibrium_positions(N::Int)
    function f!(F, x, N)
        for i in 1:N
            F[i] = (
                x[i]
              - sum([1 / (x[i] - x[j])^2 for j in 1:i-1])
              + sum([1 / (x[i] - x[j])^2 for j in i+1:N])
            )
        end
    end

    function j!(J, x, N)
        for i in 1:N, j in 1:N
            if i ≡ j
                J[i, j] = (
                    1
                  + 2 * sum([1 / (x[i] - x[j])^3 for j in 1:i-1])
                  - 2 * sum([1 / (x[i] - x[j])^3 for j in i+1:N])
            )
            else
                J[i, j] = (
                  - 2 * sum([1 / (x[i] - x[j])^3 for j in 1:i-1])
                  + 2 * sum([1 / (x[i] - x[j])^3 for j in i+1:N])
            )
            end
        end
    end
    # see eq.8 in the ref to see where (2.018/N^0.559) comes from
    if isodd(N)
        initial_x = [(2.018 / N^0.559)*i for i in -N÷2:N÷2]
    else
        initial_x = [(2.018 / N^0.559)*i for i in filter(x->x≠0, collect(-N÷2:N÷2))]
    end
    sol = nlsolve((F,x) -> f!(F,x,N), (J,x) -> j!(J,x,N), initial_x, method=:newton)
    sol.zero
end

"""
    characteristic_length_scale(M::Real, ν::Real)
Returns the characteristic length scale for a linear chain of identical ions of mass `M`
and with axial trap frequency 2π × ν.
"""
characteristic_length_scale(M::Real, ν::Real) = (e^2 / (4π * ϵ₀ * M * (2π * ν)^2))^(1/3)

"""
    Anm(N::Real, com::NamedTuple{(:x,:y,:z)}, axis::NamedTuple{(:x,:y,:z)})
Computes the normal modes and corresponding trap frequencies along a particular `axis` for a
collection of `N` ions in a linear Coloumb crystal and returns an array of tuples with first
element the frequency of the normal mode and 2nd element the corresponding eigenvector.

`com` should be a `NamedTuple` of COM frequences for the different axes:
`(x<:Real, y<:Real, z<:Real)`, where the ``z``-axis is taken to be parallel to the axis of
the crystal.
"""
function Anm(N::Int, com::NamedTuple{(:x,:y,:z)}, axis::NamedTuple{(:x,:y,:z)})
    @assert axis in [x̂, ŷ, ẑ] "axis can be x̂, ŷ, ẑ"
    axis == ẑ ? a = 2 : a = -1
    l = linear_equilibrium_positions(N)
    axis_dict = Dict([(x̂, :x), (ŷ, :y), (ẑ, :z)])
    sa = axis_dict[axis]
    β = com[sa] / com.z
    A = Array{Real}(undef, N, N)
    for n in 1:N, j in 1:N
        if n ≡ j
            A[n, j] = β^2 + a * sum([ 1 / abs(l[j] - l[p])^3 for p in 1:N if p != j])
        else
            A[n, j] = -a / abs(l[j] - l[n])^3
        end
    end
    a = eigen(A)
    for i in a.values
        @assert i > 0 """
            ($(axis_dict[axis])=$(com[sa]), z=$(com.z)) outside stability region for $N ions
        """
    end
    a1 = [sqrt(i) * com.z for i in a.values]
    a2 = a.vectors
    for i in 1:size(a2, 2)
        _sparsify!(view(a2, :,i), 1e-5)
    end
    if axis == ẑ
        return [(a1[i], a2[:,i]) for i in 1:length(a1)]
    else
        a = reverse([(a1[i], a2[:,i]) for i in 1:length(a1)])
        return a
    end
end

_sparsify!(x, eps) = @. x[abs(x) < eps] = 0

"""
    LinearChain(;
            ions::Vector{Ion}, com_frequencies::NamedTuple{(:x,:y,:z)},
            vibrational_modes::NamedTuple{(:x,:y,:z),Tuple{Vararg{Vector{VibrationalMode},3}}}
        )

Contains all of the information necessary to describe a collection of ions trapped in a 3D
harmonic potential and forming a linear coulomb crystal.

**user-defined fields**
* `ions::Vector{Ion}`: a list of ions that compose the linear Coulomb crystal
* `com_frequencies::NamedTuple{(:x,:y,:z),Tuple{Vararg{Vector{VibrationalMode},3}}}`:
        Describes the COM frequencies `(x=ν_x, y=ν_y, z=ν_z)`. The ``z``-axis is taken to be
        parallel to the crystal's symmetry axis.
* `vibrational_modes::NamedTuple{(:x,:y,:z)}`:  e.g. `vibrational_modes=(x=[1], y=[2], z=[1,2])`.
    Specifies the axis and a list of integers which correspond to the ``i^{th}`` farthest
    mode away from the COM for that axis. For example, `vibrational_modes=(z=[2])` would
    specify the axial stretch mode. These are the modes that will be modeled in the chain.
    Note: `vibrational_modes=(x=[],y=[],z=[1])`, `vibrational_modes=(y=[],z=[1])`
    and `vibrational_modes=(;z=[1])` are all acceptable and equivalent.
**derived fields**
* `full_normal_mode_description::NamedTuple{(:x,:y,:z)}`: For each axis, this contains an
    array of tuples where the first element is a vibrational frequency [Hz] and the second
    element is a vector describing the corresponding normalized normal mode structure.
"""
struct LinearChain <: IonConfiguration  # Note: this is not a mutable struct
    ions::Vector{<:Ion}
    com_frequencies::NamedTuple{(:x,:y,:z)}
    vibrational_modes::NamedTuple{(:x,:y,:z),Tuple{Vararg{Vector{VibrationalMode},3}}}
    full_normal_mode_description::NamedTuple{(:x,:y,:z)}
    function LinearChain(;
            ions, com_frequencies, vibrational_modes::NamedTuple,
        )
        vibrational_modes = _construct_vibrational_modes(vibrational_modes)
        warn = nothing
        for i in 1:length(ions), j in i+1:length(ions)
            if ions[j] ≡ ions[i]
                ions[j] = copy(ions[i])
                if isnothing(warn)
                    warn = "Some ions point to the same thing. Making copies."
                    @warn warn
                end
            end
        end
        N = length(ions)
        A = (x=Anm(N, com_frequencies, x̂),
             y=Anm(N, com_frequencies, ŷ),
             z=Anm(N, com_frequencies, ẑ))
        vm = (x=Vector{VibrationalMode}(undef, 0),
              y=Vector{VibrationalMode}(undef, 0),
              z=Vector{VibrationalMode}(undef, 0))
        r = [x̂, ŷ, ẑ]
        for (i, modes) in enumerate(vibrational_modes), mode in modes
            push!(vm[i], VibrationalMode(A[i][mode]..., axis=r[i]))
        end
        l = linear_equilibrium_positions(length(ions))
        l0 = characteristic_length_scale(m_ca40, com_frequencies.z)
        for (i, ion) in enumerate(ions)
            Core.setproperty!(ion, :number, i)
            Core.setproperty!(ion, :position, l[i] * l0)
        end
        new(ions, com_frequencies, vm, A)
    end
end

"""
    get_vibrational_modes(lc::LinearChain)
Returns an array of all of the selected `VibrationalModes` in the `LinearChain`.
The order is `[lc.x..., lc.y..., lc.z...]`.
"""
function get_vibrational_modes(lc::LinearChain)
    collect(Iterators.flatten(lc.vibrational_modes))
end

function Base.print(lc::LinearChain)
    println("$(length(lc.ions)) ions")
    println("com frequencies: $(lc.com_frequencies)")
    println("selected vibrational_modes: $(lc.vibrational_modes)")
end

function Base.show(io::IO, lc::LinearChain)  # suppress long output
    print(io, "LinearChain($(length(lc.ions)) ions)")
end

# Takes e.g. (y=[1]) to (x=[], y=[1], z=[])
function _construct_vibrational_modes(x)
    k = collect(keys(x))
    xyz = [:x, :y, :z]
    @assert isnothing(findfirst(x -> x ∉ xyz, k)) "keys of `vibrational_modes` must be `:x`, `:y` or `:z`"
    indxs = findall(x -> x ∉ k, xyz)
    values = []
    for i in 1:3
        if i in indxs
            push!(values, Int[])
        else
            push!(values, x[xyz[i]])
        end
    end
    (; zip(xyz, values)...)
end


text:
@adjoint (::Type{T})(::UndefInitializer, args...) where T<:Array = T(undef, args...), Δ -> nothing

@nograd size, length, eachindex, Colon(), findfirst, randn, ones, zeros, one, zero,
  print, println

@adjoint Base.vect(xs...) = Base.vect(xs...), Δ -> (Δ...,)

@adjoint copy(x::AbstractArray) = copy(x), ȳ -> (ȳ,)

@adjoint (::Type{T})(x::T) where T<:Array = T(x), ȳ -> (ȳ,)

_zero(xs::AbstractArray{<:Integer}) = fill!(similar(xs, float(eltype(xs))), false)
_zero(xs::AbstractArray{<:Number}) = zero(xs)
_zero(xs::AbstractArray) = Any[nothing for x in xs]

@adjoint function getindex(xs::Array, i...)
  xs[i...], function (Δ)
    Δ′ = _zero(xs)
    Δ′[i...] = Δ
    (Δ′, map(_ -> nothing, i)...)
  end
end

@adjoint! setindex!(xs::AbstractArray, x...) = setindex!(xs, x...),
  _ -> error("Mutating arrays is not supported")

# General

@adjoint collect(x::Array) = collect(x), Δ -> (Δ,)

@adjoint permutedims(xs, dims) = permutedims(xs, dims),
  Δ -> (permutedims(Δ, invperm(dims)), nothing)

@adjoint reshape(xs, dims...) = reshape(xs, dims...),
  Δ -> (reshape(Δ, size(xs)),map(_->nothing,dims)...)

@adjoint function hvcat(rows::Tuple{Vararg{Int}}, xs::T...) where T<:Number
  hvcat(rows, xs...), ȳ -> (nothing, ȳ...)
end

pull_block_vert(sz, Δ, A::AbstractVector) = Δ[sz-length(A)+1:sz]
pull_block_vert(sz, Δ, A::AbstractMatrix) = Δ[sz-size(A, 1)+1:sz, :]
@adjoint function vcat(A::Union{AbstractVector, AbstractMatrix}...)
  sz = cumsum([size.(A, 1)...])
  return vcat(A...), Δ->(map(n->pull_block_vert(sz[n], Δ, A[n]), eachindex(A))...,)
end

pull_block_horz(sz, Δ, A::AbstractVector) = Δ[:, sz]
pull_block_horz(sz, Δ, A::AbstractMatrix) = Δ[:, sz-size(A, 2)+1:sz]
@adjoint function hcat(A::Union{AbstractVector, AbstractMatrix}...)
  sz = cumsum([size.(A, 2)...])
  return hcat(A...), Δ->(map(n->pull_block_horz(sz[n], Δ, A[n]), eachindex(A))...,)
end

@adjoint function repeat(xs; inner=ntuple(_->1, ndims(xs)), outer=ntuple(_->1, ndims(xs)))
  repeat(xs, inner = inner, outer = outer), function (Δ)
    Δ′ = zero(xs)
    S = size(xs)

    # Loop through each element of Δ, calculate source dimensions, accumulate into Δ′
    for (dest_idx, val) in pairs(IndexCartesian(), Δ)
        # First, round dest_idx[dim] to nearest gridpoint defined by inner[dim], then
        # wrap around based on original size S.
        src_idx = [mod1(div(dest_idx[dim] - 1, inner[dim]) + 1, S[dim]) for dim in 1:length(S)]
        Δ′[src_idx...] += val
    end
    return (Δ′,)
  end
end

@adjoint getindex(i::Int, j::Int) = i[j], _ -> nothing

function unzip(tuples)
  map(1:length(first(tuples))) do i
      map(tuple -> tuple[i], tuples)
  end
end
@adjoint function map(f, args::AbstractArray...)
  ys_and_backs = map((args...) -> _forward(__context__, f, args...), args...)
  ys, backs = unzip(ys_and_backs)
  ys, function (Δ)
    Δf_and_args_zipped = map((f, δ) -> f(δ), backs, Δ)
    Δf_and_args = unzip(Δf_and_args_zipped)
    Δf = reduce(accum, Δf_and_args[1])
    (Δf, Δf_and_args[2:end]...)
  end
end

function _forward(cx::Context, ::typeof(collect), g::Base.Generator)
  y, back = _forward(cx, map, g.f, g.iter)
  y, function (ȳ)
    _, f̄, x̄ = back(ȳ)
    (nothing, (f = f̄, iter = x̄),)
  end
end

@adjoint iterate(r::UnitRange, i...) = iterate(r, i...), _ -> nothing

# Reductions

@adjoint function sum(xs::AbstractArray; dims = :)
  if dims === (:)
    sum(xs), Δ -> (FillArray(Δ, size(xs)),)
  else
    sum(xs, dims = dims), Δ -> (similar(xs) .= Δ,)
  end
end

function _forward(cx::Context, ::typeof(sum), f, xs::AbstractArray)
  y, back = forward(cx, (xs -> sum(f.(xs))), xs)
  y, ȳ -> (nothing, nothing, back(ȳ)...)
end

@adjoint function sum(::typeof(abs2), X::AbstractArray; dims = :)
  return sum(abs2, X; dims=dims), Δ::Union{Number, AbstractArray}->(nothing, ((2Δ) .* X))
end

@adjoint function prod(xs; dims = :)
  if dims === (:)
    prod(xs), Δ -> (prod(xs) ./ xs .* Δ,)
  else
    prod(xs, dims = dims),
      Δ -> (reshape(.*(circshift.([reshape(xs, length(xs))], 1:length(xs)-1)...), size(xs)) .* Δ,)
  end
end

function _forward(cx::Context, ::typeof(prod), f, xs::AbstractArray)
  y, back = forward(cx, (xs -> prod(f.(xs))), xs)
  y, ȳ -> (nothing, nothing, back(ȳ)...)
end

@adjoint function maximum(xs; dims = :)
  max, i = findmax(xs, dims = dims)
  max, function (Δ)
    Δ isa Real && Δ <= sqrt(eps(float(Δ))) && return nothing
    Δ′ = zero(xs)
    Δ′[i] = Δ
    return (Δ′,)
  end
end

@adjoint function minimum(xs; dims = :)
  min, i = findmin(xs, dims = dims)
  min, function (Δ)
    Δ′ = zero(xs)
    Δ′[i] = Δ
    return (Δ′,)
  end
end

@adjoint function mean(xs::AbstractArray; dims = :)
  return mean(xs, dims=dims), Δ -> (_backmean(xs,Δ,dims),)
end
_backmean(xs, Δ, ::Colon) = zero(xs) .+ Δ ./ length(xs)
_backmean(xs, Δ, dims) = zero(xs) .+ Δ ./ mapreduce(i -> size(xs,i),*,dims)

# LinAlg
# ======

@adjoint function(a::AbstractVecOrMat * b::AbstractVecOrMat)
  return a * b, function(Δ)
    return (reshape(Δ * b', size(a)), reshape(a' * Δ, size(b)))
  end
end

@adjoint transpose(x) = transpose(x), Δ -> (transpose(Δ),)
@adjoint Base.adjoint(x) = x', Δ -> (Δ',)
@adjoint parent(x::LinearAlgebra.Adjoint) = parent(x), ȳ -> (LinearAlgebra.Adjoint(ȳ),)

@adjoint dot(x::AbstractArray, y::AbstractArray) = dot(x, y), Δ->(Δ .* y, Δ .* x)

function _kron(mat1::AbstractMatrix,mat2::AbstractMatrix)
    m1, n1 = size(mat1)
    mat1_rsh = reshape(mat1,(1,m1,1,n1))

    m2, n2 = size(mat2)
    mat2_rsh = reshape(mat2,(m2,1,n2,1))

    return reshape(mat1_rsh.*mat2_rsh, (m1*m2,n1*n2))
end

@adjoint kron(a::AbstractMatrix, b::AbstractMatrix) = forward(_kron, a, b)

@adjoint function Diagonal(d::AbstractVector)
  back(Δ::NamedTuple) = (Δ.diag,)
  back(Δ::AbstractMatrix) = (diag(Δ),)
  return Diagonal(d), back
end

@adjoint diag(A::AbstractMatrix) = diag(A), Δ->(Diagonal(Δ),)

@adjoint det(xs) = det(xs), Δ -> (Δ * det(xs) * transpose(inv(xs)),)

@adjoint logdet(xs) = logdet(xs), Δ -> (Δ * transpose(inv(xs)),)

@adjoint logabsdet(xs) = logabsdet(xs), Δ -> (Δ[1] * transpose(inv(xs)),)

@adjoint function inv(A)
    return inv(A), function (Δ)
        Ainv = inv(A)
        ∇A = - Ainv' * Δ * Ainv'
        return (∇A, )
    end
end

@adjoint function \(A::AbstractMatrix, B::AbstractVecOrMat)
  Y = A \ B
  return Y, function(Ȳ)
      B̄ = A' \ Ȳ
      return (-B̄ * Y', B̄)
  end
end

function _forward(cx::Context, ::typeof(norm), x::AbstractArray, p::Real = 2)
  fallback = (x, p) -> sum(abs.(x).^p .+ eps(0f0))^(1/p) # avoid d(sqrt(x))/dx == Inf at 0
  _forward(cx, fallback, x, p)
end

# LinAlg Matrix Types
# ===================

# This is basically a hack while we don't have a working `ldiv!`.
@adjoint function \(A::Cholesky, B::AbstractVecOrMat)
  Y, back = Zygote.forward((U, B)->U \ (U' \ B), A.U, B)
  return Y, function(Ȳ)
    Ā_factors, B̄ = back(Ȳ)
    return ((uplo=nothing, status=nothing, factors=Ā_factors), B̄)
  end
end

@adjoint function /(A::AbstractMatrix, B::AbstractMatrix)
  Y = A / B
  return Y, function(Ȳ)
      Ā = Ȳ / B'
      return (Ā, -Y' * Ā)
  end
end

_symmetric_back(Δ) = UpperTriangular(Δ) + LowerTriangular(Δ)' - Diagonal(Δ)
_symmetric_back(Δ::Union{Diagonal, UpperTriangular}) = Δ
@adjoint function Symmetric(A::AbstractMatrix)
  back(Δ::AbstractMatrix) = (_symmetric_back(Δ),)
  back(Δ::NamedTuple) = (_symmetric_back(Δ.data),)
  return Symmetric(A), back
end

@adjoint function cholesky(Σ::Real)
  C = cholesky(Σ)
  return C, Δ::NamedTuple->(Δ.factors[1, 1] / (2 * C.U[1, 1]),)
end

@adjoint function cholesky(Σ::Diagonal)
  C = cholesky(Σ)
  return C, Δ::NamedTuple->(Diagonal(diag(Δ.factors) .* inv.(2 .* C.factors.diag)),)
end

# Implementation due to Seeger, Matthias, et al. "Auto-differentiating linear algebra."
@adjoint function cholesky(Σ::Union{StridedMatrix, Symmetric{<:Real, <:StridedMatrix}})
  C = cholesky(Σ)
  return C, function(Δ::NamedTuple)
    U, Ū = C.U, Δ.factors
    Σ̄ = Ū * U'
    Σ̄ = copytri!(Σ̄, 'U')
    Σ̄ = ldiv!(U, Σ̄)
    BLAS.trsm!('R', 'U', 'T', 'N', one(eltype(Σ)), U.data, Σ̄)
    @inbounds for n in diagind(Σ̄)
      Σ̄[n] /= 2
    end
    return (UpperTriangular(Σ̄),)
  end
end

@adjoint function lyap(A::AbstractMatrix, C::AbstractMatrix)
  X = lyap(A, C)
  return X, function (X̄)
    C̄ = lyap(collect(A'), X̄)
    Ā = C̄*X' + C̄'*X
    return (Ā, C̄)
  end
end

# Adjoint based on the Theano implementation, which uses the differential as described
# in Brančík, "Matlab programs for matrix exponential function derivative evaluation"
@adjoint exp(A::AbstractMatrix) = exp(A), function(F̄)
  n = size(A, 1)
  E = eigen(A)
  w = E.values
  ew = exp.(w)
  X = [i==j ? ew[i] : (ew[i]-ew[j])/(w[i]-w[j]) for i in 1:n,j=1:n]
  VT = transpose(E.vectors)
  VTF = factorize(collect(VT))
  Ā = real.(VTF\(VT*F̄/VTF.*X)*VT)
  (Ā, )
end

Zygote.@adjoint function LinearAlgebra.tr(x::AbstractMatrix)
  # x is a squre matrix checked by tr,
  # so we could just use Eye(size(x, 1))
  # to create a Diagonal
  tr(x), function (Δ::Number)
    (Diagonal(FillArray(Δ, (size(x, 1), ))), )
  end
end

# Various sensitivities for `literal_getproperty`, depending on the 2nd argument.
@adjoint function literal_getproperty(C::Cholesky, ::Val{:uplo})
  return literal_getproperty(C, Val(:uplo)), function(Δ)
    return ((uplo=nothing, info=nothing, factors=nothing),)
  end
end
@adjoint function literal_getproperty(C::Cholesky, ::Val{:info})
  return literal_getproperty(C, Val(:info)), function(Δ)
    return ((uplo=nothing, info=nothing, factors=nothing),)
  end
end
@adjoint function literal_getproperty(C::Cholesky, ::Val{:U})
  return literal_getproperty(C, Val(:U)), function(Δ)
    Δ_factors = C.uplo == 'U' ? UpperTriangular(Δ) : LowerTriangular(copy(Δ'))
    return ((uplo=nothing, info=nothing, factors=Δ_factors),)
  end
end
@adjoint function literal_getproperty(C::Cholesky, ::Val{:L})
  return literal_getproperty(C, Val(:L)), function(Δ)
    Δ_factors = C.uplo == 'L' ? LowerTriangular(Δ) : UpperTriangular(copy(Δ'))
    return ((uplo=nothing, info=nothing, factors=Δ_factors),)
  end
end

@adjoint function logdet(C::Cholesky)
  return logdet(C), function(Δ)
    return ((uplo=nothing, info=nothing, factors=Diagonal(2 .* Δ ./ diag(C.factors))),)
  end
end

@adjoint function +(A::AbstractMatrix, S::UniformScaling)
  return A + S, Δ->(Δ, (λ=sum(view(Δ, diagind(Δ))),))
end


text:
# code reflection entry-points

using .CUPTI
using .CUPTI: CUpti_CallbackDomain, CUpti_CallbackId, CUpti_SubscriberHandle,
              CUpti_ResourceData, CUpti_ModuleResourceData

#
# code_* replacements
#

function code_sass_callback(userdata::Ptr{Cvoid}, domain::CUpti_CallbackDomain,
                            cbid::CUpti_CallbackId, cbdada::Ptr{Cvoid})
    dest = Base.unsafe_pointer_to_objref(userdata)::Ref{Any}

    if domain == CUPTI.CUPTI_CB_DOMAIN_RESOURCE
        cbdada = unsafe_load(reinterpret(Ptr{CUpti_ResourceData}, cbdada))
        if cbid == CUPTI.CUPTI_CBID_RESOURCE_MODULE_LOADED
            resourceDescriptor =
                unsafe_load(reinterpret(Ptr{CUpti_ModuleResourceData}, cbdada.resourceDescriptor))
            cubin = unsafe_wrap(Vector{Cchar}, pointer(resourceDescriptor.pCubin),
                                resourceDescriptor.cubinSize)
            dest[] = copy(cubin)
        end
    end

    return
end

"""
    code_sass([io], f, types, cap::VersionNumber)

Prints the SASS code generated for the method matching the given generic function and type
signature to `io` which defaults to `stdout`.

The following keyword arguments are supported:

- `cap` which device to generate code for
- `kernel`: treat the function as an entry-point kernel
- `verbose`: enable verbose mode, which displays code generation statistics

See also: [`@device_code_sass`](@ref)
"""
function code_sass(io::IO, @nospecialize(func), @nospecialize(types), kernel::Bool=true;
                   verbose::Bool=false, kwargs...)
    tt = Base.to_tuple_type(types)
    target = CUDACompilerTarget(device(); kwargs...)
    params = CUDACompilerParams()
    job = CompilerJob(target, FunctionSpec(func, tt, kernel), params)
    code_sass(io, job; verbose=verbose)
end

# multiple subscribers aren't supported, so make sure we only call CUPTI once
const cupti_lock = ReentrantLock()

function code_sass(io::IO, job::CUDACompilerJob; verbose::Bool=false)
    if !job.source.kernel
        error("Can only generate SASS code for kernel functions")
    end

    compiled = cufunction_compile(job)

    cubin = Ref{Any}()
    callback = @cfunction(code_sass_callback, Cvoid,
                          (Ptr{Cvoid}, CUpti_CallbackDomain, CUpti_CallbackId, Ptr{Cvoid}))

    # JIT compile and capture the generated object file
    lock(cupti_lock) do
        subscriber_ref = Ref{CUpti_SubscriberHandle}()
        res = CUPTI.unsafe_cuptiSubscribe(subscriber_ref, callback, Base.pointer_from_objref(cubin))
        if res === CUPTI.CUPTI_ERROR_INSUFFICIENT_PRIVILEGES
            error("""Insufficient priviliges: You don't have permissions to profile GPU code, which is required for `code_sass`.
                    Get administrative priviles or allow all users to profile: https://developer.nvidia.com/ERR_NVGPUCTRPERM#SolnAdminTag""")
        elseif res != CUPTI.CUPTI_SUCCESS
            throw(CUPTIError(res))
        end
        subscriber = subscriber_ref[]
        try
            CUPTI.cuptiEnableDomain(1, subscriber, CUPTI.CUPTI_CB_DOMAIN_RESOURCE)
            cufunction_link(job, compiled)
        finally
            CUPTI.cuptiUnsubscribe(subscriber)
        end
    end

    # disassemble to SASS
    isassigned(cubin) || error("No kernels compiled")
    mktemp() do cubin_path,cubin_io
        write(cubin_io, cubin[])
        flush(cubin_io)

        cmd = `$(nvdisasm()) --print-code --print-line-info $cubin_path`
        for line in readlines(cmd)
            # nvdisasm output is pretty verbose;
            # perform some clean-up and make it look like @code_native
            line = replace(line, r"/\*[0-9a-f]{4}\*/" => "        ") # strip inst addr
            line = replace(line, r"^[ ]{30}" => "   ")               # reduce leading spaces
            line = replace(line, r"[\s+]//##" => ";")                # change line info tag
            line = replace(line, r"^\." => "\n.")                    # break before new BBs
            line = replace(line, r"; File \"(.+?)\", line (\d+)" => s"; Location \1:\2") # rename line info
            println(io, line)
        end
    end
end

code_sass(@nospecialize(func), @nospecialize(types); kwargs...) =
    code_sass(stdout, func, types; kwargs...)

# forward the rest to GPUCompiler with an appropriate CompilerJob

for method in (:code_typed, :code_warntype, :code_llvm, :code_native)
    # only code_typed doesn't take a io argument
    args = method == :code_typed ? (:job,) : (:io, :job)

    @eval begin
        function $method(io::IO, @nospecialize(func), @nospecialize(types);
                         kernel::Bool=false, minthreads=nothing, maxthreads=nothing,
                         blocks_per_sm=nothing, maxregs=nothing, kwargs...)
            source = FunctionSpec(func, Base.to_tuple_type(types), kernel)
            target = CUDACompilerTarget(device(); minthreads, maxthreads, blocks_per_sm, maxregs)
            params = CUDACompilerParams()
            job = CompilerJob(target, source, params)
            GPUCompiler.$method($(args...); kwargs...)
        end
        $method(@nospecialize(func), @nospecialize(types); kwargs...) =
            $method(stdout, func, types; kwargs...)
    end
end

const code_ptx = code_native

#
# @device_code_* functions
#

export @device_code_lowered, @device_code_typed, @device_code_warntype,
       @device_code_llvm, @device_code_ptx, @device_code_sass,
       @device_code

"""
    @device_code_sass [io::IO=stdout, ...] ex

Evaluates the expression `ex` and prints the result of [`CUDA.code_sass`](@ref) to
`io` for every compiled CUDA kernel. For other supported keywords, see
[`CUDA.code_sass`](@ref).
"""
macro device_code_sass(ex...)
    function hook(job::CUDACompilerJob; io::IO=stdout, kwargs...)
        println(io, "// $job")
        println(io)
        code_sass(io, job; kwargs...)
    end
    GPUCompiler.emit_hooked_compilation(hook, ex...)
end

# forward the rest to GPUCompiler
@eval $(Symbol("@device_code_lowered")) = $(getfield(GPUCompiler, Symbol("@device_code_lowered")))
@eval $(Symbol("@device_code_typed")) = $(getfield(GPUCompiler, Symbol("@device_code_typed")))
@eval $(Symbol("@device_code_warntype")) = $(getfield(GPUCompiler, Symbol("@device_code_warntype")))
@eval $(Symbol("@device_code_llvm")) = $(getfield(GPUCompiler, Symbol("@device_code_llvm")))
@eval $(Symbol("@device_code_ptx")) = $(getfield(GPUCompiler, Symbol("@device_code_native")))
@eval $(Symbol("@device_code")) = $(getfield(GPUCompiler, Symbol("@device_code")))


text:
#######################################################################
#                        BASIC RULES                                  #
#######################################################################

@nodiff length(u)    u

@nodiff size(u)      u
@nodiff size(u,v)    u
@nodiff size(u,v)    v

@nodiff fill(u,v)    u
@nodiff fill(u,v)    v

@nodiff similar(u,v) u
@nodiff similar(u,v) v

@nodiff zeros(u)             u
@nodiff zero(u)              u
@nodiff zeros(i, j)          i
@nodiff zeros(i, j)          j
@nodiff zeros(i, j, k)       i
@nodiff zeros(i, j, k)       j
@nodiff zeros(i, j, k)       k
@nodiff zeros(i, j, k, _l)   i
@nodiff zeros(i, j, k, _l)   j
@nodiff zeros(i, j, k, _l)   k
@nodiff zeros(i, j, k, _l)   _l

@nodiff ones(u)             u
@nodiff one(u)              u
@nodiff ones(i, j)          i
@nodiff ones(i, j)          j
@nodiff ones(i, j, k)       i
@nodiff ones(i, j, k)       j
@nodiff ones(i, j, k)       k
@nodiff ones(i, j, k, _l)   i
@nodiff ones(i, j, k, _l)   j
@nodiff ones(i, j, k, _l)   k
@nodiff ones(i, j, k, _l)   _l

@nodiff sign(u)      u

# @diffrule reverse(u)   u     0.0

# tuple
@diffrule tuple(u)        u     dy[1]
@diffrule tuple(u,v)      u     dy[1]
@diffrule tuple(u,v)      v     dy[2]
@diffrule tuple(u,v,w)    u     dy[1]
@diffrule tuple(u,v,w)    v     dy[2]
@diffrule tuple(u,v,w)    w     dy[3]
@diffrule tuple(u,v,w,t)  u     dy[1]
@diffrule tuple(u,v,w,t)  v     dy[2]
@diffrule tuple(u,v,w,t)  w     dy[3]
@diffrule tuple(u,v,w,t)  t     dy[4]

# __tuple__ (cassette tracer implemetation uses it instead of normal tuple)
@diffrule __tuple__(u)        u     dy[1]
@diffrule __tuple__(u,v)      u     dy[1]
@diffrule __tuple__(u,v)      v     dy[2]
@diffrule __tuple__(u,v,w)    u     dy[1]
@diffrule __tuple__(u,v,w)    v     dy[2]
@diffrule __tuple__(u,v,w)    w     dy[3]
@diffrule __tuple__(u,v,w,t)  u     dy[1]
@diffrule __tuple__(u,v,w,t)  v     dy[2]
@diffrule __tuple__(u,v,w,t)  w     dy[3]
@diffrule __tuple__(u,v,w,t)  t     dy[4]

# vcat
@diffrule vcat(u,v)      u     unvcat(dy, 1, u, v)
@diffrule vcat(u,v)      v     unvcat(dy, 2, u, v)
@diffrule vcat(u,v,w)    u     unvcat(dy, 1, u, v, w)
@diffrule vcat(u,v,w)    v     unvcat(dy, 2, u, v, w)
@diffrule vcat(u,v,w)    w     unvcat(dy, 3, u, v, w)
@diffrule vcat(u,v,w,t)  u     unvcat(dy, 1, u, v, w, t)
@diffrule vcat(u,v,w,t)  v     unvcat(dy, 2, u, v, w, t)
@diffrule vcat(u,v,w,t)  w     unvcat(dy, 3, u, v, w, t)
@diffrule vcat(u,v,w,t)  t     unvcat(dy, 4, u, v, w, t)

#  hcat
@diffrule hcat(u,v)      u     unhcat(dy, 1, u, v)
@diffrule hcat(u,v)      v     unhcat(dy, 2, u, v)
@diffrule hcat(u,v,w)    u     unhcat(dy, 1, u, v, w)
@diffrule hcat(u,v,w)    v     unhcat(dy, 2, u, v, w)
@diffrule hcat(u,v,w)    w     unhcat(dy, 3, u, v, w)
@diffrule hcat(u,v,w,t)  u     unhcat(dy, 1, u, v, w, t)
@diffrule hcat(u,v,w,t)  v     unhcat(dy, 2, u, v, w, t)
@diffrule hcat(u,v,w,t)  w     unhcat(dy, 3, u, v, w, t)
@diffrule hcat(u,v,w,t)  t     unhcat(dy, 4, u, v, w, t)

#  cat
@diffrule_kw cat(u,v)      u     uncat(dy, 1, u, v)
@diffrule_kw cat(u,v)      v     uncat(dy, 2, u, v)
@diffrule_kw cat(u,v,w)    u     uncat(dy, 1, u, v, w)
@diffrule_kw cat(u,v,w)    v     uncat(dy, 2, u, v, w)
@diffrule_kw cat(u,v,w)    w     uncat(dy, 3, u, v, w)
@diffrule_kw cat(u,v,w,t)  u     uncat(dy, 1, u, v, w, t)
@diffrule_kw cat(u,v,w,t)  v     uncat(dy, 2, u, v, w, t)
@diffrule_kw cat(u,v,w,t)  w     uncat(dy, 3, u, v, w, t)
@diffrule_kw cat(u,v,w,t)  t     uncat(dy, 4, u, v, w, t)

# reshape
@diffrule reshape(u::AbstractArray, _a)             u    reshape(dy, size(u))
@diffrule reshape(u::AbstractArray, _a, _b)         u    reshape(dy, size(u))
@diffrule reshape(u::AbstractArray, _a, _b, _c)     u    reshape(dy, size(u))
@diffrule reshape(u::AbstractArray, _a, _b, _c, _d) u    reshape(dy, size(u))
@nodiff reshape(u::AbstractArray, _a)              _a
@nodiff reshape(u::AbstractArray, _a, _b)          _a
@nodiff reshape(u::AbstractArray, _a, _b)          _b
@nodiff reshape(u::AbstractArray, _a, _b, _c)      _a
@nodiff reshape(u::AbstractArray, _a, _b, _c)      _b
@nodiff reshape(u::AbstractArray, _a, _b, _c)      _c
@nodiff reshape(u::AbstractArray, _a, _b, _c, _d)  _a
@nodiff reshape(u::AbstractArray, _a, _b, _c, _d)  _b
@nodiff reshape(u::AbstractArray, _a, _b, _c, _d)  _c
@nodiff reshape(u::AbstractArray, _a, _b, _c, _d)  _d
@diffrule reshape(u::AbstractArray, _d::Tuple)      u    reshape(dy, size(u))
@nodiff reshape(u::AbstractArray, _d::Tuple)        _d

@diffrule vec(u::AbstractArray)    u    reshape(dy, size(u))

# getindex
@diffrule getindex(u::AbstractArray, i)         u    ungetindex(u, dy, i)
@diffrule getindex(u::AbstractArray, i, j)      u    ungetindex(u, dy, i, j)
@diffrule getindex(u::AbstractArray, i, j, k)   u    ungetindex(u, dy, i, j, k)
@nodiff getindex(u::AbstractArray, i)         i
@nodiff getindex(u::AbstractArray, i, j)      i
@nodiff getindex(u::AbstractArray, i, j)      j
@nodiff getindex(u::AbstractArray, i, j, k)   i
@nodiff getindex(u::AbstractArray, i, j, k)   j
@nodiff getindex(u::AbstractArray, i, j, k)   k

@diffrule getindex(u::Tuple, i) u ungetindex(u, dy, i)
@nodiff getindex(u::Tuple, i) i

# view - in non-mutating world behaves like getting subarrays using getindex
@diffrule view(u::AbstractArray, i)         u    ungetindex(u, dy, i)
@diffrule view(u::AbstractArray, i, j)      u    ungetindex(u, dy, i, j)
@diffrule view(u::AbstractArray, i, j, k)   u    ungetindex(u, dy, i, j, k)
@nodiff view(u::AbstractArray, i)         i
@nodiff view(u::AbstractArray, i, j)      i
@nodiff view(u::AbstractArray, i, j)      j
@nodiff view(u::AbstractArray, i, j, k)   i
@nodiff view(u::AbstractArray, i, j, k)   j
@nodiff view(u::AbstractArray, i, j, k)   k

# square root
@diffrule sqrt(u::Real)              u     one(u) / 2 * u ^ (-one(u) / 2) * dy

# addition
@diffrule +(u::Real         , v::Real )            u     dy
@diffrule +(u::AbstractArray, v::AbstractArray)    u     dy
@diffrule +(u::Real         , v::Real )            v     dy
@diffrule +(u::AbstractArray, v::AbstractArray)    v     dy

@diffrule broadcast(_fn::typeof(+), u, v) u unbroadcast(u, dy)
@diffrule broadcast(_fn::typeof(+), u, v) v unbroadcast(v, dy)

@diffrule broadcast(_fn::typeof(*), u, v) u unbroadcast_prod_x(u, v, dy)
@diffrule broadcast(_fn::typeof(*), u, v) v unbroadcast_prod_y(u, v, dy)

# unary substraction
@diffrule -(u::Real )                               u     -dy
@diffrule -(u::AbstractArray)                       u     -dy

# binary substraction
@diffrule -(u::Real, v::Real)                       u     dy
@diffrule -(u::AbstractArray, v::AbstractArray)     u     dy
@diffrule -(u::Real         , v::Real)              v     -dy
@diffrule -(u::AbstractArray, v::AbstractArray)     v     -dy

# sum() and mean()
@diffrule sum(u::AbstractArray)                     u     sum_grad(u, dy)
@diffrule Base._sum(u::AbstractArray, v::Int)             u     sum_grad(u, dy)
@diffrule Base._sum(u::AbstractArray, v::Int)             v     zero(eltype(u))
@diffrule Core.kwfunc(sum)(_dims, _, u::AbstractArray)     u     sum_grad(u, dy)

# special sums
@diffrule sum(_fn::typeof(log), u::AbstractArray)    u    sum_grad(u, dy) ./ u

@diffrule Statistics.mean(u::AbstractArray)                u     ∇mean(u, dy)
@diffrule Statistics._mean(u::AbstractArray, v::Int)       u     ∇mean(u, dy)
@diffrule Statistics._mean(u::AbstractArray, v::Int)       v     zero(eltype(u))
@diffrule Core.kwfunc(Statistics.mean)(_dims, _, u::AbstractArray) u ∇mean(u, dy, _dims)
@nodiff Core.kwfunc(Statistics.mean)(_dims, _, u::AbstractArray) _dims
@nodiff Core.kwfunc(Statistics.mean)(_dims, _, u::AbstractArray) _

# diag
@diffrule diag(u::AbstractMatrix)    u    Diagonal(dy)

# dot()
@diffrule dot(u::Real, v::Real)                     u     v * dy
@diffrule dot(u::Real, v::Real)                     v     u * dy

# @diffrule dot(u::AbstractArray, v::AbstractArray)   u     v.*dy
# @diffrule dot(u::AbstractArray, v::AbstractArray)   v     u.*dy

# log() and exp()
@diffrule log(u::Real )                            u     dy / u
@diffrule exp(u::Real )                            u     exp(u) * dy
@diffrule log1p(u::Real)                           u     dy  / (one(u) + u)
@diffrule expm1(u::Real)                           u     (one(u) + expm1(u))  * dy
# @diffrule expm1(u::AbstractArray)                  u     (one(u) + expm1(u)) .* dy
# note : derivative uses expm1() and not exp() to reuse the
#   already calculated expm1()

# trig functions
@diffrule sin(u::Real )                            u     cos(u) * dy
@diffrule cos(u::Real )                            u     -sin(u) * dy
@diffrule tan(u::Real )                            u     (one(u) + tan(u)  * tan(u))  * dy
@diffrule sinh(u::Real )                           u     cosh(u) * dy
@diffrule cosh(u::Real )                           u     sinh(u) * dy
@diffrule tanh(u::Real )                           u     (one(u) - tanh(u)  * tanh(u))  * dy
@diffrule asin(u::Real )                           u     dy  / sqrt(one(u) - u*u)
@diffrule acos(u::Real )                           u     -dy  / sqrt(one(u) - u*u)
@diffrule atan(u::Real )                           u     dy  / (one(u) + u *u)

# round, floor, ceil, trunc, mod2pi
@diffrule round(u::Real)                           u     zero(u)

@diffrule floor(u::Real)                           u     zero(u)

@diffrule ceil(u::Real)                            u     zero(u)

@diffrule trunc(u::Real)                           u     zero(u)

@diffrule mod2pi(u::Real)                          u     dy

# abs, max(), min()
@diffrule abs(u::Real)                             u     sign(u) * dy
@diffrule abs2(u::Real)                            u     2 * u * dy

# @diffrule max(u::Real         , v::Real)           u     (u > v) * dy
# @diffrule max(u::Real         , v::AbstractArray)  u     sum((u .> v) .* dy)
# @diffrule max(u::AbstractArray, v::Real)           u     (u .> v) .* dy
# @diffrule max(u::AbstractArray, v::AbstractArray)  u     (u .> v) .* dy

# @diffrule max(u::Real         , v::Real)           v     (u < v) * dy
# @diffrule max(u::Real         , v::AbstractArray)  v     (u .< v) .* dy
# @diffrule max(u::AbstractArray, v::Real)           v     sum((u .< v) .* dy)
# @diffrule max(u::AbstractArray, v::AbstractArray)  v     (u .< v) .* dy

# @diffrule min(u::Real         , v::Real)           u     (u < v) * dy
# @diffrule min(u::Real         , v::AbstractArray)  u     sum((u .< v) .* dy)
# @diffrule min(u::AbstractArray, v::Real)           u     (u .< v) .* dy
# @diffrule min(u::AbstractArray, v::AbstractArray)  u     (u .< v) .* dy

# @diffrule min(u::Real         , v::Real)           v     (u > v) * dy
# @diffrule min(u::Real         , v::AbstractArray)  v     (u .> v) .* dy
# @diffrule min(u::AbstractArray, v::Real)           v     sum((u .> v) .* dy)
# @diffrule min(u::AbstractArray, v::AbstractArray)  v     (u .> v) .* dy

# maximum, minimum
# @diffrule maximum(u::Real         )     u     dy
# @diffrule maximum(u::AbstractArray)     u     (u .== maximum(u)) .* dy

# @diffrule minimum(u::Real         )     u     dy
# @diffrule minimum(u::AbstractArray)     u     (u .== minimum(u)) .* dy

# multiplication
@diffrule *(u::Real         , v::Real )            u     v * dy
# @diffrule *(u::Real         , v::AbstractArray)    u     sum(v .* dy)
# @diffrule *(u::AbstractArray, v::Real )            u     v .* dy
@diffrule *(u::AbstractArray, v::AbstractArray)    u     dy * transpose(v)

@diffrule *(u::Real         , v::Real )            v     u * dy
# @diffrule *(u::Real         , v::AbstractArray)    v     u .* dy
@diffrule *(u::AbstractArray, v::Real )            v     sum(u .* dy)
@diffrule *(u::AbstractArray, v::AbstractArray)    v     transpose(u) * dy

# power  (both args reals)
@diffrule ^(u::Real, v::Real)                      u     v * u ^ (v-(one(u))) * dy
@diffrule ^(u::Real, v::Real)                      v     log(u) * u ^ v * dy

get_val_param(::Val{v}) where v = v
@diffrule Base.literal_pow(_, u::Real, v)         u      get_val_param(v) * u ^ (get_val_param(v)-(one(u))) * dy
@nodiff Base.literal_pow(_, u::Real, v)           _
@nodiff Base.literal_pow(_, u::Real, v)           v

# # division
@diffrule /(u::Real          , v::Real )           u     dy / v
@diffrule /(u::Real          , v::Real )           v     -u * dy / (v * v)

# transpose
@diffrule transpose(u::AbstractVector)             u     untranspose_vec(dy)
@diffrule transpose(u::AbstractArray)              u     transpose(dy)
@diffrule adjoint(u::AbstractVector)             u     untranspose_vec(dy)
@diffrule adjoint(u::AbstractArray)              u     adjoint(dy)

# permutedims
@diffrule permutedims(x, _perm) x ∇permutedims(dy, _perm)
@nodiff permutedims(x, _perm) _perm

# # erf, erfc, gamma, beta, lbeta, lgamma
# @diffrule erf(u::Real)                       u     2.0/sqrt(π) * exp(-u  * u)  * dy
# @diffrule erf(u::AbstractArray)              u     2.0/sqrt(π) .* exp(-u .* u) .* dy

# @diffrule erfc(u::Real)                      u     -2.0/sqrt(π) * exp(-u  * u)  * dy
# @diffrule erfc(u::AbstractArray)             u     -2.0/sqrt(π) .* exp(-u .* u) .* dy

# @diffrule gamma(u::Real)                     u     polygamma(0,u)  * gamma(u)  * dy
# @diffrule gamma(u::AbstractArray)            u     polygamma(0,u) .* gamma(u) .* dy

# @diffrule lgamma(u::Real)                    u     polygamma(0,u)  * dy
# @diffrule lgamma(u::AbstractArray)           u     polygamma(0,u) .* dy

# @diffrule beta(u::Real         , v::Real)            u   beta(u,v)  * (digamma(u)-digamma(u+v))  * dy
# @diffrule beta(u::AbstractArray, v::AbstractArray)   u   beta(u,v) .* (digamma(u)-digamma(u+v)) .* dy
# @diffrule beta(u::Real         , v::Real)            v   beta(u,v)  * (digamma(v)-digamma(u+v))  * dy
# @diffrule beta(u::AbstractArray, v::AbstractArray)   v   beta(u,v) .* (digamma(v)-digamma(u+v)) .* dy

# @diffrule lbeta(u::Real         , v::Real)            u   (polygamma(0,u)-polygamma(0,u+v))  * dy
# @diffrule lbeta(u::AbstractArray, v::AbstractArray)   u   (polygamma(0,u)-polygamma(0,u+v)) .* dy
# @diffrule lbeta(u::Real         , v::Real)            v   (polygamma(0,v)-polygamma(0,u+v))  * dy
# @diffrule lbeta(u::AbstractArray, v::AbstractArray)   v   (polygamma(0,v)-polygamma(0,u+v)) .* dy

@nodiff rand(u) u

@nodiff convert(u, v) u
@diffrule convert(u, v) v dy

@nodiff to_device(_d, u)      _d
@diffrule to_device(_d, u)    u    dy


text:
module SCoBASolver

using StaticArrays
using Parameters
using DataStructures

# Types etc
export
    MODE,
    START,
    FINISH,
    SUCCESS,
    InteractionEvent,
    SearchTree,
    GenericAllocation,
    MRTAEnvironment,
    DecisionNode,
    OutcomeNode,
    TaskUtil

# Solver stuff
export
    generate_search_tree!,
    get_next_attempt_idx,
    SCoBAAlgorithm,
    coordinate_allocation!

include("types.jl")
include("scoba_tree_search.jl")
include("scoba_conflict_resolution.jl")

end


text:
# DispatchAnalyzer
# ================

"""
Every [entry point of dispatch analysis](@ref dispatch-analysis-entry-points) can accept
any of [JET configurations](https://aviatesk.github.io/JET.jl/dev/config/) as well as
the following additional configurations that are specific to dispatch analysis.

---
- `frame_filter = x::$State->true`:\\
  A predicate which takes `InfernceState` or `OptimizationState` and returns `false` to skip analysis on the frame.
  ```julia
  # only checks code within the current module:
  julia> mymodule_filter(x) = x.mod === @__MODULE__;
  julia> @test_nodispatch frame_filter=mymodule_filter f(args...)
  ...
  ```

---
- `function_filter = @nospecialize(ft)->true`:\\
  A predicate which takes a function type and returns `false` to skip analysis on the call.
  ```julia
  # ignores `Core.Compiler.widenconst` calls (since it's designed to be runtime-dispatched):
  julia> myfunction_filter(@nospecialize(ft)) = ft !== typeof(Core.Compiler.widenconst)
  julia> @test_nodispatch function_filter=myfunction_filter f(args...)
  ...
  ```

---
- `skip_nonconcrete_calls::Bool = true`:\\
  Julia's runtime dispatch is "powerful" because it can always compile code with concrete
  runtime arguments so that [a "kernel" function](https://docs.julialang.org/en/v1/manual/performance-tips/#kernel-functions)
  runs very effectively even if it's called from a type-instable call site.
  This means, we (really) often accept that some parts of our code are not inferred statically,
  and rather we want to just rely on information that is only available at runtime.
  To model this programming style, dispatch analyzer does NOT report any optimization failures
  or runtime dispatches detected within non-concrete calls under the default configuration.
  We can turn off this `skip_nonconcrete_calls` configuration to get type-instabilities
  within non-concrete calls.
  ```julia
  # the following examples are adapted from https://docs.julialang.org/en/v1/manual/performance-tips/#kernel-functions
  julia> function fill_twos!(a)
             for i = eachindex(a)
                 a[i] = 2
             end
         end;

  julia> function strange_twos(n)
             a = Vector{rand(Bool) ? Int64 : Float64}(undef, n)
             fill_twos!(a)
             return a
         end;

  # by default, only type-instabilities within concrete call (i.e. `strange_twos(3)`) are reported
  # and those within non-concrete calls (`fill_twos!(a)`) are not reported
  julia> @report_dispatch strange_twos(3)
  ═════ 2 possible errors found ═════
  ┌ @ REPL[2]:2 %45(Main.undef, n)
  │ runtime dispatch detected: %45::Type{Vector{_A}} where _A(Main.undef, n::Int64)
  └─────────────
  ┌ @ REPL[2]:3 Main.fill_twos!(%46)
  │ runtime dispatch detected: Main.fill_twos!(%46::Vector)
  └─────────────
  Vector (alias for Array{_A, 1} where _A)

  # we can get reports from non-concrete calls with `skip_nonconcrete_calls=false`
  julia> @report_dispatch skip_nonconcrete_calls=false strange_twos(3)
  ═════ 4 possible errors found ═════
  ┌ @ REPL[2]:3 Main.fill_twos!(a)
  │┌ @ REPL[1]:3 Base.setindex!(a, 2, %14)
  ││ runtime dispatch detected: Base.setindex!(a::Vector, 2, %14::Int64)
  │└─────────────
  │┌ @ REPL[1]:3 Base.setindex!(a, 2, i)
  ││┌ @ array.jl:877 Base.convert(_, x)
  │││ runtime dispatch detected: Base.convert(_::Any, x::Int64)
  ││└────────────────
  ┌ @ REPL[2]:2 %45(Main.undef, n)
  │ runtime dispatch detected: %45::Type{Vector{_A}} where _A(Main.undef, n::Int64)
  └─────────────
  ┌ @ REPL[2]:3 Main.fill_twos!(%46)
  │ runtime dispatch detected: Main.fill_twos!(%46::Vector)
  └─────────────
  Vector (alias for Array{_A, 1} where _A)
  ```

---
- `skip_unoptimized_throw_blocks::Bool = true`:\\
  By default, Julia's native compilation pipeline intentionally disables inference (and so
  succeeding optimizations too) on "throw blocks", which are code blocks that will eventually
  lead to `throw` calls, in order to ease [the compilation latency problem, a.k.a. "first-time-to-plot"](https://julialang.org/blog/2020/08/invalidations/).
  Accordingly, the dispatch analyzer also ignores runtime dispatches detected within those blocks
  since we _usually_ don't mind if code involved with error handling isn't optimized.
  If `skip_unoptimized_throw_blocks` is set to `false`, it doesn't ignore them and will
  report type instabilities detected within "throw blocks".

  See also <https://github.com/JuliaLang/julia/pull/35982>.

  ```julia
  # by default, unoptimized "throw blocks" are not analyzed
  julia> @test_nodispatch sin(10)
  Test Passed
    Expression: #= none:1 =# JETTest.@test_nodispatch sin(10)

  # we can turn on the analysis on unoptimized "throw blocks" with `skip_unoptimized_throw_blocks=false`
  julia> @test_nodispatch skip_unoptimized_throw_blocks=false sin(10)
  Dispatch Test Failed at none:1
    Expression: #= none:1 =# JETTest.@test_nodispatch skip_unoptimized_throw_blocks = false sin(10)
    ═════ 1 possible error found ═════
    ┌ @ math.jl:1221 Base.Math.sin(xf)
    │┌ @ special/trig.jl:39 Base.Math.sin_domain_error(x)
    ││┌ @ special/trig.jl:28 Base.Math.DomainError(x, "sin(x) is only defined for finite x.")
    │││ runtime dispatch detected: Base.Math.DomainError(x::Float64, "sin(x) is only defined for finite x.")
    ││└──────────────────────

  ERROR: There was an error during testing

  # we can also turns off the heuristic itself
  julia> @test_nodispatch unoptimize_throw_blocks=false skip_unoptimized_throw_blocks=false sin(10)
  Test Passed
    Expression: #= none:1 =# JETTest.@test_nodispatch unoptimize_throw_blocks = false skip_unoptimized_throw_blocks = false sin(10)
  ```

---
"""
mutable struct DispatchAnalyzer{S,T} <: AbstractAnalyzer
    state::AnalyzerState
    opts::BitVector
    frame_filter::S
    function_filter::T
    concrete_frame::Union{Nothing,Bool}
    skip_unoptimized_throw_blocks::Bool
end
function DispatchAnalyzer(;
    frame_filter = x::State->true,
    function_filter = @nospecialize(ft)->true,
    skip_nonconcrete_calls::Bool = true,
    skip_unoptimized_throw_blocks::Bool = true,
    jetconfigs...)
    state = AnalyzerState(; jetconfigs...)
    concrete_frame = skip_nonconcrete_calls ? true : nothing
    return DispatchAnalyzer(state, BitVector(), frame_filter, function_filter, concrete_frame, skip_unoptimized_throw_blocks)
end

# AbstractAnalyzer API requirements
JETInterfaces.AnalyzerState(analyzer::DispatchAnalyzer) = analyzer.state
JETInterfaces.AbstractAnalyzer(analyzer::DispatchAnalyzer, state::AnalyzerState) =
    DispatchAnalyzer(state, analyzer.opts, analyzer.frame_filter, analyzer.function_filter, analyzer.concrete_frame, analyzer.skip_unoptimized_throw_blocks)
JETInterfaces.ReportPass(analyzer::DispatchAnalyzer) = DispatchAnalysisPass()

# we want to run different analysis with a different filter, so include its hash into the cache key
function JET.get_cache_key(analyzer::DispatchAnalyzer)
    h = @invoke get_cache_key(analyzer::AbstractAnalyzer)
    h = @invoke hash(analyzer.frame_filter::Any, h::UInt)    # HACK avoid dynamic dispatch
    h = @invoke hash(analyzer.function_filter::Any, h::UInt) # HACK avoid dynamic dispatch
    h = hash(isnothing(analyzer.concrete_frame), h)
    h = hash(analyzer.skip_unoptimized_throw_blocks, h)
    return h
end

struct DispatchAnalysisPass <: ReportPass end

# ignore most of reports defined by JET.jl
(::DispatchAnalysisPass)(T::Type{<:InferenceErrorReport}, @nospecialize(_...)) = return
(::DispatchAnalysisPass)(T::Type{GeneratorErrorReport}, @nospecialize(args...)) = SoundPass()(T, args...)

@reportdef struct OptimizationFailureReport <: InferenceErrorReport end
JETInterfaces.get_msg(::Type{OptimizationFailureReport}, args...) =
    return "failed to optimize" #: signature of this MethodInstance

function (::DispatchAnalysisPass)(::Type{OptimizationFailureReport}, analyzer::DispatchAnalyzer, frame::InferenceState)
    report!(OptimizationFailureReport, analyzer, frame.linfo)
end

@reportdef struct RuntimeDispatchReport <: InferenceErrorReport end
JETInterfaces.get_msg(::Type{RuntimeDispatchReport}, analyzer, s) =
    return "runtime dispatch detected" #: call signature

function (::DispatchAnalysisPass)(::Type{RuntimeDispatchReport}, analyzer::DispatchAnalyzer, opt::OptimizationState)
    throw_blocks =
        analyzer.skip_unoptimized_throw_blocks && opt.inlining.params.unoptimize_throw_blocks ?
        find_throw_blocks(opt.src.code) : nothing

    sptypes, slottypes = opt.sptypes, opt.slottypes
    for (pc, x) in enumerate(opt.src.code)
        if !isnothing(throw_blocks)
            # optimization is intentionally turned off for this block, let's ignore anything here
            CC.in(pc, throw_blocks) && continue
        end
        if @isexpr(x, :call)
            ft = widenconst(argextype(first(x.args), opt.src, sptypes, slottypes))
            ft <: Builtin && continue # ignore `:call`s of language intrinsics
            if analyzer.function_filter(ft)
                report!(RuntimeDispatchReport, analyzer, (opt, pc))
            end
        end
    end
end

# branch on https://github.com/JuliaLang/julia/pull/39885
@static if isdefined(CC, :finish!)

import .CC:
    _typeinf,
    finish,
    finish!

function CC._typeinf(analyzer::DispatchAnalyzer, frame::InferenceState)
    @assert isempty(analyzer.opts)

    (; concrete_frame) = analyzer
    skip_nonconcrete_calls = !isnothing(concrete_frame)
    if skip_nonconcrete_calls
        analyzer.concrete_frame = isdispatchtuple(frame.linfo.specTypes)
    end

    ret = @invoke _typeinf(analyzer::AbstractAnalyzer, frame::InferenceState)

    if skip_nonconcrete_calls
        analyzer.concrete_frame = concrete_frame
    end

    @assert isempty(analyzer.opts)

    return ret
end

function CC.finish(frame::InferenceState, analyzer::DispatchAnalyzer)
    ret = @invoke finish(frame::InferenceState, analyzer::AbstractAnalyzer)

    (; concrete_frame, opts) = analyzer
    if !(isnothing(concrete_frame) || concrete_frame) || !analyzer.frame_filter(frame)
        push!(opts, false)
    else
        if isa(frame.result.src, OptimizationState)
            push!(opts, true)
        else
            report_pass!(OptimizationFailureReport, analyzer, frame)
            push!(opts, false)
        end
    end

    return ret
end

function CC.finish!(analyzer::DispatchAnalyzer, caller::InferenceResult)
    opt = caller.src

    ret = @invoke finish!(analyzer::AbstractInterpreter, caller::InferenceResult)

    if popfirst!(analyzer.opts) # optimization happened
        if isa(ret, Const) # the optimization was very successful, nothing to report
        elseif isa(ret, OptimizationState) # compiler cached the optimized IR, just analyze it
            report_pass!(RuntimeDispatchReport, analyzer, ret)
        elseif isa(ret, CodeInfo) # compiler didn't cache the optimized IR, but `finish!(::AbstractInterpreter, ::InferenceResult)` transformed it to `opt.src`, so we can analyze it
            @assert isa(opt, OptimizationState) && opt.src === ret
            report_pass!(RuntimeDispatchReport, analyzer, opt)
        else
            Core.eval(@__MODULE__, :(ret = $ret))
            throw("unexpected state happened, inspect $(@__MODULE__).ret") # this pass should never happen
        end
    end

    return ret
end

else # @static if isdefined(CC, :finish!)

include("legacy/dispatch")

end # @static if isdefined(CC, :finish!)

# entries
# =======

"""
    analyze_dispatch(f, types = Tuple{}; jetconfigs...) -> (analyzer::DispatchAnalyzer, frame::Union{InferenceFrame,Nothing})
    analyze_dispatch(tt::Type{<:Tuple}; jetconfigs...) -> (analyzer::DispatchAnalyzer, frame::Union{InferenceFrame,Nothing})

Analyzes the generic function call with the given type signature, and returns:
- `analyzer::DispatchAnalyzer`: contains analyzed optimization failures and runtime dispatch points
- `frame::Union{InferenceFrame,Nothing}`: the final state of the abstract interpretation,
  or `nothing` if `f` is a generator and the code generation has been failed
"""
function analyze_dispatch(@nospecialize(args...);
                          analyzer = DispatchAnalyzer,
                          jetconfigs...)
    @assert analyzer === DispatchAnalyzer "analyzer is fixed to $DispatchAnalyzer"
    return analyze_call(args...; analyzer, jetconfigs...)
end

"""
    report_dispatch(f, types = Tuple{}; jetconfigs...) -> result_type::Any
    report_dispatch(tt::Type{<:Tuple}; jetconfigs...) -> result_type::Any

Analyzes the generic function call with the given type signature, and then prints detected
optimization failures and runtime dispatch points to `stdout`, and finally returns the result
type of the call.
"""
function report_dispatch(@nospecialize(args...);
                         analyzer = DispatchAnalyzer,
                         jetconfigs...)
    @assert analyzer === DispatchAnalyzer "analyzer is fixed to $DispatchAnalyzer"
    return report_call(args...; analyzer, jetconfigs...)
end

"""
    @analyze_dispatch [jetconfigs...] f(args...)

Evaluates the arguments to the function call, determines its types, and then calls
[`analyze_dispatch`](@ref) on the resulting expression.
As with `@code_typed` and its family, any of [JET configurations](https://aviatesk.github.io/JET.jl/dev/config/)
or [dispatch analysis specific configurations](@ref dispatch-analysis-configurations) can be given as the optional arguments like this:
```julia
# reports `rand(::Type{Bool})` with `unoptimize_throw_blocks` configuration turned on
julia> @analyze_dispatch unoptimize_throw_blocks=true rand(Bool)
```
"""
macro analyze_dispatch(ex0...)
    return gen_call_with_extracted_types_and_kwargs(__module__, :analyze_dispatch, ex0)
end

"""
    @report_dispatch [jetconfigs...] f(args...)

Evaluates the arguments to the function call, determines its types, and then calls
[`report_dispatch`](@ref) on the resulting expression.
As with `@code_typed` and its family, any of [JET configurations](https://aviatesk.github.io/JET.jl/dev/config/)
or [dispatch analysis specific configurations](@ref dispatch-analysis-configurations) can be given as the optional arguments like this:
```julia
# reports `rand(::Type{Bool})` with `unoptimize_throw_blocks` configuration turned on
julia> @report_call unoptimize_throw_blocks=true rand(Bool)
```
"""
macro report_dispatch(ex0...)
    return gen_call_with_extracted_types_and_kwargs(__module__, :report_dispatch, ex0)
end

# Test integration
# ================

"""
    @test_nodispatch [jetconfigs...] [broken=false] [skip=false] f(args...)

Tests the generic function call `f(args...)` is free from runtime dispatch.
Returns a `Pass` result if it is, a `Fail` result if if contains any location where runtime
dispatch or optimization failure happens, or an `Error` result if this macro encounters an
unexpected error. When the test `Fail`s, abstract call stack to each problem location will
also be printed to `stdout`.

```julia
julia> @test_nodispatch sincos(10)
Test Passed
  Expression: #= none:1 =# JETTest.@test_nodispatch sincos(10)
```

As with [`@report_dispatch`](@ref) or [`@analyze_dispatch`](@ref), any of [JET configurations](https://aviatesk.github.io/JET.jl/dev/config/)
or [dispatch analysis specific configurations](@ref dispatch-analysis-configurations) can be given as the optional arguments like this:
```julia
julia> function f(n)
            r = sincos(n)
            println(r) # `println` is full of runtime dispatches, but we can ignore the corresponding reports from `Base` by explicit frame filter
            return r
       end;
julia> this_module_filter(x) = x.mod === @__MODULE__;

julia> @test_nodispatch frame_filter=this_module_filter f(10)
Test Passed
  Expression: #= none:1 =# JETTest.@test_nodispatch frame_filter = this_module_filter f(10)
```

`@test_nodispatch` is fully integrated with [`Test` standard library's unit-testing infrastructure](https://docs.julialang.org/en/v1/stdlib/Test/).
It means, the result of `@test_nodispatch` will be included in the final `@testset` summary,
it supports `skip` and `broken` annotations as `@test` macro does, etc.
```julia
julia> using JETTest, Test

julia> f(params) = sin(params.value); # type-stable
julia> params = (; value = 10);       # non-constant global variable
julia> g() = sin(params.value);       # very type-instable

julia> @testset "check optimizations" begin
           @test_nodispatch f((; value = 10)) # pass
           @test_nodispatch g()               # fail
           @test_nodispatch broken=true g()   # annotated as broken, thus still "pass"
       end
check optimizations: Dispatch Test Failed at none:3
  Expression: #= none:3 =# JETTest.@test_nodispatch g()
  ═════ 2 possible errors found ═════
  ┌ @ none:1 Base.getproperty(%1, :value)
  │ runtime dispatch detected: Base.getproperty(%1::Any, :value::Symbol)
  └──────────
  ┌ @ none:1 Main.sin(%2)
  │ runtime dispatch detected: Main.sin(%2::Any)
  └──────────

Test Summary:       | Pass  Fail  Broken  Total
check optimizations |    1     1       1      3
ERROR: Some tests did not pass: 1 passed, 1 failed, 0 errored, 1 broken.
```
"""
macro test_nodispatch(ex0...)
    ex0 = collect(ex0)

    local broken = nothing
    local skip   = nothing
    idx = Int[]
    for (i,x) in enumerate(ex0)
        if iskwarg(x)
            key, val = x.args
            if key === :broken
                if !isnothing(broken)
                    error("invalid test macro call: cannot set `broken` keyword multiple times")
                end
                broken = esc(val)
                push!(idx, i)
            elseif key === :skip
                if !isnothing(skip)
                    error("invalid test macro call: cannot set `skip` keyword multiple times")
                end
                skip = esc(val)
                push!(idx, i)
            end
        end
    end
    if !isnothing(broken) && !isnothing(skip)
        error("invalid test macro call: cannot set both `skip` and `broken` keywords")
    end
    deleteat!(ex0, idx)

    testres, orig_expr = test_dispatch_exs(ex0, __module__, __source__)

    return quote
        if $(!isnothing(skip) && skip)
            $record($get_testset(), $Broken(:skipped, $orig_expr))
        else
            testres = $testres
            if $(!isnothing(broken) && broken)
                if isa(testres, $DispatchTestFailure)
                    testres = $Broken(:test_nodispatch, $orig_expr)
                elseif isa(testres, $Pass)
                    testres = $Error(:test_unbroken, $orig_expr, nothing, nothing, $(QuoteNode(__source__)))
                end
            else
                isa(testres, $Pass) || ccall(:jl_breakpoint, $Cvoid, ($Any,), testres)
            end
            $record($get_testset(), testres)
        end
    end
end

iskwarg(@nospecialize(x)) = @isexpr(x, :(=))

get_exceptions() = @static if isdefined(Base, :current_exceptions)
    Base.current_exceptions()
else
    Base.catch_stack()
end
@static if !hasfield(Pass, :source)
    Pass(test_type::Symbol, orig_expr, data, thrown, source) = Pass(test_type, orig_expr, data, thrown)
end

function test_dispatch_exs(ex0, m, source)
    analyzer_call = gen_call_with_extracted_types_and_kwargs(m, :analyze_dispatch, ex0)
    orig_expr = QuoteNode(
        Expr(:macrocall, GlobalRef(JETTest, Symbol("@test_nodispatch")), source, ex0...))
    source = QuoteNode(source)
    testres = :(try
        analyzer, frame = $analyzer_call
        reports = $get_reports(analyzer)
        if $length(reports) == 0
            $Pass(:test_nodispatch, $orig_expr, nothing, nothing, $source)
        else
            $DispatchTestFailure($orig_expr, $source, reports)
        end
    catch err
        isa(err, $InterruptException) && rethrow()
        $Error(:test_error, $orig_expr, err, $get_exceptions(), $source)
    end) |> Base.remove_linenums!
    return testres, orig_expr
end

"""
    test_nodispatch(f, types = Tuple{}; broken::Bool = false, skip::Bool = false, jetconfigs...)
    test_nodispatch(tt::Type{<:Tuple}; broken::Bool = false, skip::Bool = false, jetconfigs...)

Tests the generic function call with the given type signature is free from runtime dispatch.
Except that it takes a type signature rather than a call expression, this function works
in the same way as [`@test_nodispatch`](@ref).
"""
function test_nodispatch(@nospecialize(args...);
                         broken::Bool = false, skip::Bool = false,
                         jetconfigs...)
    source = LineNumberNode(@__LINE__, @__FILE__)
    kwargs = map(((k,v),)->Expr(:kw, k, v), collect(jetconfigs))
    orig_expr = :(JETTest.test_nodispatch($(args...); $(kwargs...)))

    if skip
        record(get_testset(), Broken(:skipped, orig_expr))
    else
        testres = try
            analyzer, frame = analyze_dispatch(args...; jetconfigs...)
            reports = get_reports(analyzer)
            if length(reports) == 0
                Pass(:test_nodispatch, orig_expr, nothing, nothing, source)
            else
                DispatchTestFailure(orig_expr, source, reports)
            end
        catch err
            isa(err, InterruptException) && rethrow()
            Error(:test_error, orig_expr, err, get_exceptions(), source)
        end

        if broken
            if isa(testres, DispatchTestFailure)
                testres = Broken(:test_nodispatch, orig_expr)
            elseif isa(testres, Pass)
                testres = Error(:test_unbroken, orig_expr, nothing, nothing, source)
            end
        else
            isa(testres, Pass) || ccall(:jl_breakpoint, Cvoid, (Any,), testres)
        end
        record(get_testset(), testres)
    end
end

# NOTE we will just show abstract call strack, and won't show backtrace of actual test executions

struct DispatchTestFailure <: Test.Result
    orig_expr::Expr
    source::LineNumberNode
    reports::Vector{InferenceErrorReport}
end

const TEST_INDENTS = "  "

function Base.show(io::IO, t::DispatchTestFailure)
    printstyled(io, "Dispatch Test Failed"; bold=true, color=Base.error_color())
    print(io, " at ")
    printstyled(io, something(t.source.file, :none), ":", t.source.line, "\n"; bold=true, color=:default)
    println(io, TEST_INDENTS, "Expression: ", t.orig_expr)
    # print abstract call stack, with appropriate indents
    _, ctx = Base.unwrapcontext(io)
    buf = IOBuffer()
    ioctx = IOContext(buf, ctx)
    print_reports(ioctx, t.reports) # TODO kwargs support
    lines = replace(String(take!(buf)), '\n'=>string('\n',TEST_INDENTS))
    print(io, TEST_INDENTS, lines)
end

Base.show(io::IO, ::MIME"application/prs.juno.inline", t::DispatchTestFailure) =
    return t

function Test.record(::Test.FallbackTestSet, t::DispatchTestFailure)
    println(t)
    throw(FallbackTestSetException("There was an error during testing"))
end

function Test.record(ts::Test.DefaultTestSet, t::DispatchTestFailure)
    if Test.TESTSET_PRINT_ENABLE[]
        printstyled(ts.description, ": ", color=:white)
        print(t)
        println()
    end
    # HACK convert to `Fail` so that test summarization works correctly
    push!(ts.results, Fail(:test_nodispatch, t.orig_expr, nothing, nothing, t.source))
    return t
end


text:
struct WebPage <: ShortCode
    url::String
end

function Base.show(io::IO, ::MIME"text/plain", page::WebPage)
    print(io, page.url)
end

function Base.show(io::IO, ::MIME"text/html", page::WebPage)
    print(io, webpage(page.url))
end

function webpage(url="", height=400, width=700, title="")
    htm = """<iframe src=" """* url * """ " height=" """ * string(height) * """ " width=" """ * string(width) * """ " title=" """ * title * """ "></iframe>"""
    return htm
end

struct Vimeo <: ShortCode
    id::Integer
end

function Base.show(io::IO, ::MIME"text/plain", video::Vimeo)
    print(io, "https://vimeo.com/$(video.id)")
end

function Base.show(io::IO, ::MIME"text/html", video::Vimeo)
    write(io, vimeo(video.id))
end

@memoize function vimeo(video_url)
    url = "https://vimeo.com/api/oembed.json?url=$(HTTP.escapeuri(video_url))&maxheight=500&maxwidth=700&width=680&byline=false&portrait=false&title=false"
    response = HTTP.get(url)
    json = JSON3.read(String(response.body))
    return json[:html]
end

function vimeo(video_id::Integer)
    vimeo("https://vimeo.com/$video_id")
end


text:
@testset "Base.:(+) -- SSAValues" begin
    @test (+)(Core.SSAValue(1), 1) == Core.SSAValue(2)
    @test (+)(1, Core.SSAValue(1)) == Core.SSAValue(2)
end

@testset "Pipe -- misc." begin
    ir = code_info(g, Tuple{Int})
    p = CodeInfoTools.Pipe(ir)
    for (v, st) in p
    end
    display(p)
    display(length(p))
    p = CodeInfoTools.Pipe(ir)
    p = identity(p)
    get_slot(p, :x)
end


text:
using IntervalArithmetic

"""
    truncated_supports(udata::AbstractUncertainValueDataset; quantiles = [0.001, 0.999])

Truncate the furnishing distribution of each uncertain value in the dataset
to the provided `quantiles` range. Returns a vector of Interval{Float64}, one
for each value.
"""
function truncated_supports(udata::AbstractUncertainValueDataset;
            quantiles = [0.001, 0.999])
    n_vals = length(udata)

    # Using the provided quantiles, find the parts of the supports of the furnishing
    # distributions from which we're going to sample.
    supports = Vector{Interval{Float64}}(undef, n_vals)

    for i = 1:n_vals
        lowerbound = quantile(udata[i], minimum(quantiles))
        upperbound = quantile(udata[i], maximum(quantiles))

        supports[i] = interval(lowerbound, upperbound)
    end

    return supports
end

"""
    strictly_increasing_sequence_exists(udata::AbstractUncertainValueDataset;
        quantiles = [0.0001, 0.9999])

Does a path through the dataset exist? I.e, check that a strictly
increasing sequence can be found after first
constraining each distribution to the provided quantile range (this
is necessary because some distributions may have infinite support).
"""
function strictly_increasing_sequence_exists(udata, quantiles = [0.0001, 0.9999])
    n_vals = length(udata)
    sample = Vector{Float64}(undef, n_vals)

    # First, constrain all data such that the supports are all finite
    constrained_data = constrain(udata, TruncateQuantiles(quantiles...,))

    # Sample the first value in a way that ensures a strictly
    # increasing sequence from indices 2:end will exist.
    minima = minimum.(constrained_data)
    maxima = maximum.(constrained_data)

    increasing_sequence_exists = true

    for i = 2:n_vals - 1
        # Find the lower and upper bounds of the support from which
        # we can draw values while still ensuring an increasing
        # sequence of values.

        lo = minima[i]
        hi = minimum(maxima[i:end])
        if lo > hi
            increasing_sequence_exists = false
        end
    end

    return increasing_sequence_exists
end

"""
    strictly_decreasing_sequence_exists(udata::AbstractUncertainValueDataset;
        quantiles = [0.0001, 0.9999])

Does a path through the dataset exist? I.e,  check that a strictly
decreasing sequence can be found after first
constraining each distribution to the provided quantile range (this
is necessary because some distributions may have infinite support).
"""
function strictly_decreasing_sequence_exists(udata; quantiles = [0.0001, 0.9999])
    # If a strictly increasing sequence exists for the reversed dataset, then a strictly
    # decreasing sequence exists for the original dataset.
    strictly_increasing_sequence_exists(udata[end:-1:1], quantiles = quantiles)
end

export
strictly_increasing_sequence_exists,
strictly_decreasing_sequence_exists


text:
#applies a function to a vector of trailing variables
function trailingmeasuretest!(F::Function,
  crsp::DataFrame;
  windowsize::Int=error("window size required"),
  Ftarget::Symbol,
  Fmeasurename::Symbol,
  offset::Int=0,
  maxcalendarinterval::DatePeriod = error("max calendar size required"),
  sorted::Bool = false)::Nothing

  N::Int = size(crsp,1)

  #allocate spcae for the new fields
  crsp[!,Fmeasurename] = Vector{MFloat64}(undef, N)

  #need to sort
  (!sorted) && (sort!(crsp, [:permno, :date]))

  #avoid missing data in the returns for simplicity
  scrspnomissing::SubDataFrame = view(crsp, (!ismissing).(crsp[!,Ftarget]),:) #view(crsp, :, :)
  scrsps::GroupedDataFrame = groupby(scrspnomissing, :permno)

  #now compute the trailing returns
  Threads.@threads for i ∈ 1:length(scrsps)
    local scrsp = scrsps[i]

    Nscrsp = size(scrsp, 1)
    if windowsize ≤ Nscrsp #if there are enough  entries to compute the trailing returns
      for r ∈ windowsize:Nscrsp
        r₀::Int = r-windowsize-offset+1 #this si the first row of the period
        rₜ::Int = r - offset
        staledate::Date = scrsp[r,:date] - maxcalendarinterval
        mindate::Date = scrsp[r₀,:date]

        #if the returns are not stale, record them
        if staledate < mindate
          scrsp[r, Fmeasurename] = F(scrsp[r₀:rₜ , Ftarget])
        end
      end
    end
  end

  return nothing
end
#computes the trailing returns from the logged returns
function trailingsigmatest!(crsp::DataFrame; sorted::Bool = false)

  print("computing trailing vol using test routine. Time: ")
  trailingmeasuretest!(std, crsp,
    Ftarget=:ret,
    windowsize=PARAM[:sigmawindow],
    maxcalendarinterval=PARAM[:sigmamaxcalendarinterval],
    sorted=sorted,
    Fmeasurename=PARAM[:sigmaname])::Nothing
  println(" checksum: $(sum(skipmissing(crsp.sigma)))")

  print("computing trailing vol using primary routine. Time: ")
  trailingmeasure!(std, crsp,
    Ftarget=:ret,
    windowsize=PARAM[:sigmawindow],
    maxcalendarinterval=PARAM[:sigmamaxcalendarinterval],
    sorted=sorted,
    Fmeasurename=PARAM[:sigmaname])::Nothing
  println(" checksum: $(sum(skipmissing(crsp.sigma)))")
  return nothing
end

max_stars_repo_name: clintonTE/CCA

max_issues_repo_name: clintonTE/CCA

max_forks_repo_name: clintonTE/CCA


text:
"""
    gcmarray{T, N, AT}

gcmarray data structure. Available constructors:

```
gcmarray{T,N,AT}(grid::gcmgrid,meta::varmeta,f::Array{AT,N},
         fSize::Array{NTuple{N, Int}},fIndex::Array{Int,1},v::String)

gcmarray(grid::gcmgrid,f::Array{Array{T,2},N}) where {T,N}
gcmarray(grid::gcmgrid,f::Array{Array{T,N},1}) where {T,N}

gcmarray(grid::gcmgrid,fSize::Array{NTuple{N, Int}},fIndex::Array{Int,1})
gcmarray(<same as above>,n3::Int)
gcmarray(<same as above>,n3::Int,n4::Int)

gcmarray(grid::gcmgrid)
gcmarray(grid::gcmgrid,::Type{T})
gcmarray(grid::gcmgrid,::Type{T},n3::Int)
gcmarray(grid::gcmgrid,::Type{T},n3::Int,n4::Int)

gcmarray(A::Array{T,N};meta::varmeta=defaultmeta)
```
"""
struct gcmarray{T, N, AT} <: AbstractMeshArray{T, N}
   grid::gcmgrid
   meta::varmeta
   f::OuterArray{AT,N}
   fSize::OuterArray{NTuple{2, Int}}
   fIndex::OuterArray{Int,1}
   version::String
end

function gcmarray(A::Array{T,N};
  meta::varmeta=defaultmeta) where {T, N}
  s=size(A)[1:2]
  fs=Array{NTuple{2, Int},1}(undef,1)
  fs[1]=s
  ios=[s[1] s[2]]
  γ=gcmgrid("","PeriodicDomain",1, fs, ios, T, read, write)
  if N==2
    B=γ.read(A,gcmarray(γ,T;meta=meta));
  else
    B=γ.read(A,gcmarray(γ,T,size(A)[3];meta=meta));
  end
  B
end

function gcmarray(grid::gcmgrid,f::OuterArray{InnerArray{T,2},N};
                  meta::varmeta=defaultmeta) where {T, N}
  gcmarray{T,N,InnerArray{T,2}}(grid,meta,f,grid.fSize,collect(1:grid.nFaces),thisversion)
end

function gcmarray(grid::gcmgrid,f::OuterArray{InnerArray{T,N},1};
                  meta::varmeta=defaultmeta) where {T, N}
  nFaces=grid.nFaces
  if N>2
    n3=size(f[1],3); n4=size(f[1],4);
    g=OuterArray{InnerArray{T,2},3}(undef,nFaces,n3,n4)
    for I in eachindex(view(g,1:nFaces,1:n3,1:n4))
      g[I]=view(f[I[1]],:,:,I[2],I[3])
    end
    n4==1 ? g=dropdims(g,dims=3) : nothing
    gcmarray{T,ndims(g),InnerArray{T,2}}(grid,meta,g,grid.fSize,collect(1:nFaces),thisversion)
  else
    gcmarray{T,1,InnerArray{T,2}}(grid,meta,f,grid.fSize,collect(1:nFaces),thisversion)
  end
end

function gcmarray(grid::gcmgrid,::Type{T},
        fSize::Union{OuterArray{NTuple{2, Int},1},NTuple{2, Int}},
        fIndex::Union{OuterArray{Int,1},Int};
        meta::varmeta=defaultmeta) where {T}
  nFaces=length(fIndex)
  f=OuterArray{InnerArray{T,2},1}(undef,nFaces)
  isa(fSize,NTuple) ? fSize=[fSize] : nothing
  isa(fIndex,Int) ? fIndex=[fIndex] : nothing
  for a=1:nFaces
    f[a]=InnerArray{T}(undef,fSize[a])
  end
  gcmarray{T,1,InnerArray{T,2}}(grid,meta,f,fSize,fIndex,thisversion)
end

function gcmarray(grid::gcmgrid,::Type{T},
        fSize::Union{OuterArray{NTuple{2, Int},1},NTuple{2, Int}},
        fIndex::Union{OuterArray{Int,1},Int},n3::Int;
        meta::varmeta=defaultmeta) where {T}
  nFaces=length(fIndex)
  f=OuterArray{InnerArray{T,2},2}(undef,nFaces,n3)
  isa(fSize,NTuple) ? fSize=[fSize] : nothing
  isa(fIndex,Int) ? fIndex=[fIndex] : nothing
  for a=1:nFaces; for i3=1:n3;
    f[a,i3]=InnerArray{T,2}(undef,fSize[a]...)
  end; end;
  gcmarray{T,2,InnerArray{T,2}}(grid,meta,f,fSize,fIndex,thisversion)
end

function gcmarray(grid::gcmgrid,::Type{T},
        fSize::Union{OuterArray{NTuple{2, Int},1},NTuple{2, Int}},
        fIndex::Union{OuterArray{Int,1},Int},n3::Int,n4::Int;
        meta::varmeta=defaultmeta) where {T}
  nFaces=length(fIndex)
  f=OuterArray{InnerArray{T,2},3}(undef,nFaces,n3,n4)
  isa(fSize,NTuple) ? fSize=[fSize] : nothing
  isa(fIndex,Int) ? fIndex=[fIndex] : nothing
  for a=1:nFaces; for i4=1:n4; for i3=1:n3;
    f[a,i3,i4]=InnerArray{T,2}(undef,fSize[a]...)
  end; end; end;
  gcmarray{T,3,InnerArray{T,2}}(grid,meta,f,fSize,fIndex,thisversion)
end

# +
function gcmarray(grid::gcmgrid; meta::varmeta=defaultmeta)
  nFaces=grid.nFaces
  fSize=grid.fSize
  fIndex=collect(1:grid.nFaces)
  T=grid.ioPrec
  gcmarray(grid,T,fSize,fIndex,meta=meta)
end

function gcmarray(grid::gcmgrid,::Type{T};
                  meta::varmeta=defaultmeta) where {T}
  nFaces=grid.nFaces
  fSize=grid.fSize
  fIndex=collect(1:grid.nFaces)
  gcmarray(grid,T,fSize,fIndex,meta=meta)
end

function gcmarray(grid::gcmgrid,::Type{T},n3::Int;
                  meta::varmeta=defaultmeta) where {T}
  nFaces=grid.nFaces
  fSize=grid.fSize
  fIndex=collect(1:grid.nFaces)
  gcmarray(grid,T,fSize,fIndex,n3,meta=meta)
end

function gcmarray(grid::gcmgrid,::Type{T},n3::Int,n4::Int;
                  meta::varmeta=defaultmeta) where {T}
  nFaces=grid.nFaces
  fSize=grid.fSize
  fIndex=collect(1:grid.nFaces)
  gcmarray(grid,T,fSize,fIndex,n3,n4,meta=meta)
end

# -

# # Interface Methods

# +
Base.size(A::gcmarray) = size(A.f)
Base.size(A::gcmarray, dim::Integer) = size(A)[dim]
#Base.dataids(A::gcmarray) = Base.dataids(A.f)
Base.dataids(A::gcmarray) = (Base.dataids(A.f)..., Base.dataids(A.fSize)..., Base.dataids(A.fIndex)...)

# +
function Base.getindex(A::gcmarray{T, N, Array{T,2}}, I::Vararg{Union{Int,Array{Int},AbstractUnitRange,Colon}, N}) where {T,N}
  J=1:length(A.fIndex)
  !isa(I[1],Colon) ? J=J[I[1]] : nothing
  nFaces=length(J)

  tmpf=A.f[I...]
  if isa(tmpf,Array{eltype(A),2})
    tmp=tmpf
  else
    n3=Int(length(tmpf)/nFaces)
    K=(A.grid,eltype(A),A.fSize[J],A.fIndex[J])
    n3>1 ? tmp=gcmarray(K...,n3) : tmp=gcmarray(K...)
    for I in eachindex(tmpf); tmp.f[I] = view(tmpf[I],:,:); end
  end

  return tmp
end

"""
    getindexetc(A::gcmarray, I::Vararg{_}) where {T,N}

Same as getindex but also returns the face size and index
"""
function getindexetc(A::gcmarray{T, N,InnerArray{T,2}}, I::Vararg{Union{Int,Array{Int},AbstractUnitRange,Colon}, N}) where {T,N}
    f=A[I...]
    fSize=A.fSize[I[1]]
    fIndex=A.fIndex[I[1]]
    return f,fSize,fIndex
end
# -

function Base.setindex!(A::gcmarray{T, N,InnerArray{T,2}}, v, I::Vararg{Int, N}) where {T,N}
  return (A.f[I...] = v)
end

function Base.view(A::gcmarray{T, N,InnerArray{T,2}}, I::Vararg{Union{Int,AbstractUnitRange,Colon}, N}) where {T,N}
  J=1:length(A.fIndex)
  !isa(I[1],Colon) ? J=J[I[1]] : nothing
  nFaces=length(J)

  tmpf=view(A.f,I...)
  n3=Int(length(tmpf)/nFaces) #length(tmpf)>nFaces ? n3=Int(length(tmpf)/nFaces) : n3=1

  K=(A.grid,eltype(A),A.fSize[J],A.fIndex[J])
  n3>1 ? tmp=gcmarray(K...,n3) : tmp=gcmarray(K...)
  for I in eachindex(tmpf); tmp.f[I] = view(tmpf[I],:,:); end

  return tmp
end

# ### Custom pretty-printing, similar, and broadcast

function Base.show(io::IO, z::gcmarray{T, N, Array{T,2}}) where {T,N}
    if ~isa(z.meta.unit,Missing)
      printstyled(io, "  name        = ",color=:normal)
      printstyled(io, "$(z.meta.name)\n",color=:magenta)
      printstyled(io, "  unit        = ",color=:normal)
      printstyled(io, "$(z.meta.unit)\n",color=:magenta)
    end
    printstyled(io, "  data type   = ",color=:normal)
    printstyled(io, "$(eltype(z))\n",color=:magenta)
    printstyled(io, "  cell pos.   = ",color=:normal)
    printstyled(io, "$(z.meta.position)\n",color=:magenta)
    printstyled(io, "  tile array  = ",color=:normal)
    printstyled(io, "$(size(z))\n",color=:cyan)
    printstyled(io, "  tile sizes  = ",color=:normal)
    printstyled(io, "$(size(z[1]))\n",color=:cyan)
    for iFace=2:length(z.fIndex)
      printstyled(io, "                ",color=:normal)
      printstyled(io, "$(size(z[iFace]))\n",color=:cyan)
    end
    printstyled(io, "  grid class  = ",color=:normal)
    printstyled(io, "$(z.grid.class)\n",color=:green)
    printstyled(io, "  MeshArray   = ",color=:normal)
    printstyled(io, "gcmarray \n",color=:green)
    printstyled(io, "  version     = ",color=:normal)
    printstyled(io, "$(z.version) \n",color=:green)
  return
end

function Base.similar(A::gcmarray;m::varmeta=defaultmeta)
    if ndims(A)==1
        B=gcmarray(A.grid,eltype(A),copy(A.fSize),copy(A.fIndex); meta=m)
    else
        B=gcmarray(A.grid,eltype(A),copy(A.fSize),copy(A.fIndex),size(A,2); meta=m)
    end
    return B
end

# ### Customize broadcasting

Base.BroadcastStyle(::Type{<:gcmarray}) = Broadcast.ArrayStyle{gcmarray}()

function Base.similar(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{gcmarray}}, ::Type{ElType}) where ElType
  # Scan the inputs for the gcmarray:
  A = find_gcmarray(bc)
  # Create the gcmarray output:
  if ndims(A)==1
        B=gcmarray(A.grid,ElType,A.fSize,A.fIndex)
  else
        B=gcmarray(A.grid,ElType,A.fSize,A.fIndex,size(A,2))
  end
  return B
end

find_gcmarray(bc::Base.Broadcast.Broadcasted) = find_gcmarray(bc.args)
find_gcmarray(args::Tuple) = find_gcmarray(find_gcmarray(args[1]), Base.tail(args))
find_gcmarray(x) = x
find_gcmarray(a::gcmarray, rest) = a
find_gcmarray(::Any, rest) = find_gcmarray(rest)

####

import Base: copyto!

# Specialize this method if all you want to do is specialize on typeof(dest)
@inline function copyto!(dest::MeshArrays.gcmarray, bc::Broadcast.Broadcasted{Nothing})
    axes(dest) == axes(bc) || throwdm(axes(dest), axes(bc))
    # Performance optimization: broadcast!(identity, dest, A) is equivalent to copyto!(dest, A) if indices match
    if bc.f === identity && bc.args isa Tuple{AbstractArray} # only a single input argument to broadcast!
        A = bc.args[1]
        if axes(dest) == axes(A)
            return copyto!(dest, A)
        end
    end
    bc′ = Broadcast.preprocess(dest, bc)
    @simd for I in eachindex(bc′)
        #@inbounds dest[I] = bc′[I]
        @inbounds dest[I] = gcmarray_getindex_evalf(bc′,I)
    end
    return dest
end

function gcmarray_getindex_evalf(bc,I)
  @boundscheck checkbounds(bc, I)
  args = Broadcast._getindex(bc.args, I)
  return bc.f.(args...)
end

###

"""
    nFacesEtc(a::gcmarray)

Return nFaces, n3 (1 in 2D case; >1 otherwise)
"""
function nFacesEtc(a::gcmarray)
  nFaces=length(a.fIndex)
  ndims(a.f)>1 ? n3=size(a.f,2) : n3=1
  ndims(a.f)>2 ? n4=size(a.f,3) : n4=1
  return nFaces, n3, n4
end


text:
module RTLSDR

import Base.open, Base.close

export
    RtlSdr,
    open,
    close,
    read_bytes,
    read_samples,
    packed_bytes_to_iq

export
    get_rate,
    set_rate,
    get_freq,
    set_freq,
    get_gains,
    get_gain,
    set_gain,
    set_agc_mode,
    set_tuner_gain_mode

include("c_interface.jl")

mutable struct RtlSdr
    valid_ptr::Bool
    dongle_ptr::Ptr{rtlsdr_dev}

    function RtlSdr(index::Int64=0)
        dp = rtlsdr_open(index)

        r = new(true, dp)

        # default sample rate and center freq
        set_rate(r, 2.0e6)
        set_freq(r, 88.5e6)

        return r
    end
end

function open(r::RtlSdr, device_index::Int)
end

function close(r::RtlSdr)
    rtlsdr_close(r.dongle_ptr)
    r.valid_ptr = false
end

"""
`set_rate(r::RtlSdr, sample_rate_Hz)
"""
function set_rate(r::RtlSdr, sample_rate_Hz)
    @assert r.valid_ptr
    rtlsdr_set_sample_rate(r.dongle_ptr, sample_rate_Hz)
end
function get_rate(r::RtlSdr)
    @assert r.valid_ptr
    rate = rtlsdr_get_sample_rate(r.dongle_ptr)
    return Int(rate)
end

"""
`set_freq(r::RtlSdr, freq_Hz)`

Interface for `rtlsdr_set_center_freq`.
"""
function set_freq(r::RtlSdr, freq_Hz)
    @assert r.valid_ptr
    rtlsdr_set_center_freq(r.dongle_ptr, freq_Hz)
end
function get_freq(r::RtlSdr)
    @assert r.valid_ptr
    freq = rtlsdr_get_center_freq(r.dongle_ptr)
    return Int(freq)
end

"""
`get_gains(r::RtlSdr)`

Interface for `rtlsdr_get_tuner_gains`.
Gains are in tenths of a dB.
"""
function get_gains(r::RtlSdr)
    @assert r.valid_ptr
    return rtlsdr_get_tuner_gains(r.dongle_ptr)
end

"""
`set_gain(r::RtlSdr, gain_db)`

Interface for `rtlsdr_set_tuner_gain`.
"""
function set_gain(r::RtlSdr, gain_db)
    @assert r.valid_ptr
    # Implement rounding, taking after roger-'s pyrtlsdr
    avail_gains = get_gains(r)
    errors = [abs((10 * gain_db) - avail_gain) for avail_gain in get_gains(r)]
    closest_gain_idx = findall(errors .== minimum(errors))[1] # broadcast find
    rtlsdr_set_tuner_gain(r.dongle_ptr, avail_gains[closest_gain_idx])
end
function get_gain(r::RtlSdr)
    @assert r.valid_ptr
    return rtlsdr_get_tuner_gain(r.dongle_ptr) / 10
end

"""
Interface for `rtlsdr_set_agc_mode`
"""
function set_agc_mode(r::RtlSdr, mode)
    @assert r.valid_ptr
    rtlsdr_set_agc_mode(r.dongle_ptr, mode)
end

"""
Interface for `rtlsdr_set_tuner_gain_mode`
"""
function set_tuner_gain_mode(r::RtlSdr, manual)
    @assert r.valid_ptr
    rtlsdr_set_tuner_gain_mode(r.dongle_ptr, manual)
end

"""
`read_bytes(r::RtlSdr, num_bytes)`

High-level interface for `rtlsdr_read_sync`.

`num_bytes` must be a multiple of 512.

Returns a vector of length `num_bytes` of Uint8 (bytes).
"""
function read_bytes(r::RtlSdr, num_bytes)
    @assert r.valid_ptr
    return read_bytes(r.dongle_ptr, num_bytes)
end

"""
`read_samples(r::RtlSdr, num_samples)`

Returns a vector of length `num_samples` with complex numbers.
"""
function read_samples(r::RtlSdr, num_samples)
    @assert r.valid_ptr
    num_bytes = 2num_samples
    raw_data = read_bytes(r.dongle_ptr, num_bytes)
    return packed_bytes_to_iq(raw_data)
end

"""
`packed_bytes_to_iq(bytes)`
"""
function packed_bytes_to_iq(bytes)
    num_bytes = length(bytes)
    num_iq = round(Int, num_bytes/2.0, RoundDown)
    iq_vals = zeros(Complex{Float64}, num_iq)
    den = 255.0/2.0

    for i = 1:num_iq
        iq_vals[i] = bytes[2i-1]/den - 1.0 + im*(bytes[2i]/den - 1.0)
    end

    return iq_vals
end

end # module