import Base: getindex, show

@doc raw"""
    BasisSet

Object holding a set of BasisFunction objects associated with an array of atoms.

# Fields

| Name        | Type                               |   Description |
|:------------|:-----------------------------------|:-----------------------------------------------------------|
|`atoms`      | `Vector{Atom}`                     | An array of `Molecules.Atom` objects  |
|`name`       | `String`                           | String holding the basis set name  |
|`basis`      | `Vector{Vector{BasisFunction}}`    | An Array of arrays with BasisFunction          |
|`natoms`     | `Int32`                            | Number of atoms in the BasisSet |
|`nbas`       | `Int32`                            | Number of basis functions (Note, not equal the number of BasisFunction objects) |
|`nshells`    | `Int32`                            | Number of shells, i.e. BasisFunction objects |
|`lc_atoms`   | `Array{Int32,1}`                   | Integer array maping data to libcint |
|`lc_bas`     | `Array{Int32,1}`                   | Integer array maping data to libcint |
|`lc_env`     | `Array{Float64,1}`                 | Float64 array maping data to libcint |

# Example

Build a basis set from default options
```julia
julia> water = Molecules.parse_string(
"O        1.2091536548      1.7664118189     -0.0171613972
 H        2.1984800075      1.7977100627      0.0121161719
 H        0.9197881882      2.4580185570      0.6297938832"
)
julia> bset = BasisSet("sto-3g", water)
sto-3g Basis Set
Number of shells: 5
Number of basis:  7

O: 1s 2s 1p
H: 1s
H: 1s
```
The BasisSet object can be accessed as two-dimensional array.
```julia
julia> bset[1,2] # Show the second basis for the first atom (O 2s)
S shell with 1 basis built from 3 primitive gaussians

χ₀₀  =    0.8486970052⋅Y₀₀⋅exp(-5.033151319⋅r²)
     +    1.1352008076⋅Y₀₀⋅exp(-1.169596125⋅r²)
     +    0.8567529838⋅Y₀₀⋅exp(-0.38038896⋅r²)
```
You can also create your own crazy mix!
Lets create one S and one P basis functions for H
```julia
julia> h2 = Molecules.parse_string(
   "H 0.0 0.0 0.0
    H 0.0 0.0 0.7"
)
julia> s = BasisFunction(0, [0.5215367271], [0.122])
julia> p = BasisFunction(1, [1.9584045349], [0.727])
```
The basis set is constructed with an array of atoms (Vector{Atom}) and a corresponding array of Vector{BasisFunction}
holding all basis functions for that particular atom. In this example, we consider an unequal treatment for the
two atoms in the H₂ molecule.
```
julia> shells = [
    [s],  # A s function on the first hydrogen
    [s,p] # One s and one p function on the second hydrogen
]
julia> BasisSet(h2, "UnequalHydrogens", shells)
UnequalHydrogens Basis Set
Number of shells: 3
Number of basis:  5

H: 1s
H: 1s 1p
```
"""
struct BasisSet
    name::String
    atoms::Vector{Atom}
    basis::Vector{Vector{BasisFunction}}
    natoms::Cint
    nbas::Cint
    nshells::Cint
    lc_atoms::Array{Cint,1}
    lc_bas::Array{Cint, 1}
    lc_env::Array{Cdouble,1}
end

function BasisSet(name::String, atoms::Vector{T}) where T <: Atom

    basis = Vector{BasisFunction}[]

    # Cache entries so repeated atoms don't need to be looked up all the time
    cache = Dict()
    for i in eachindex(atoms)
        S = symbol(atoms[i])
        if S in keys(cache)
            push!(basis, cache[S])
        else
            push!(basis, read_basisset(name, S))
            cache[S] = basis[end]
        end
    end
    BasisSet(name, atoms, basis)
end

function BasisSet(name::String, atoms::Vector{T}, basis::Vector{Vector{BasisFunction}}) where T <: Atom

    length(atoms) != length(basis) ? throw(ArgumentError("Error combining atoms and basis functions")) : nothing

    ATM_SLOTS = 6
    BAS_SLOTS = 8

    natm = length(atoms)
    nbas = 0
    nshells = 0
    nexps = 0
    nprims = 0

    for i in eachindex(atoms)
        for b in basis[i]
            nshells += 1
            nbas += 2*b.l + 1
            nexps += length(b.exp)
            nprims += length(b.coef)
        end
    end

    lc_atm = zeros(Cint, natm*ATM_SLOTS)
    lc_bas = zeros(Cint, nshells*BAS_SLOTS)
    env = zeros(Cdouble, 3*natm+nexps+nprims)

    # Prepare the lc_atom input
    off = 0
    ib = 0
    for i = eachindex(atoms)
        A = atoms[i]
        # lc_atom has ATM_SLOTS (6) "spaces" for each atom
        # The first one (Z_INDEX) is the atomic number
        lc_atm[1 + ATM_SLOTS*(i-1)] = Cint(A.Z)
        # The second one is the env index address for xyz
        lc_atm[2 + ATM_SLOTS*(i-1)] = off
        env[off+1:off+3] .= A.xyz ./ Molecules.bohr_to_angstrom
        off += 3
        # The remaining 4 slots are zero.

        # Prepare the lc_bas input
        for j = eachindex(basis[i])
            B = basis[i][j]
            Ne = length(B.exp)
            Nc = length(B.coef)
            # lc_bas has BAS_SLOTS for each basis set
            # The first one is the index of the atom starting from 0
            lc_bas[1 + BAS_SLOTS*ib] = i-1
            # The second one is the angular momentum
            lc_bas[2 + BAS_SLOTS*ib] = B.l
            # The third is the number of primitive functions
            lc_bas[3 + BAS_SLOTS*ib] = Nc
            # The fourth is the number of contracted functions
            lc_bas[4 + BAS_SLOTS*ib] = 1
            # The fifth is a κ parameter
            lc_bas[5 + BAS_SLOTS*ib] = 0
            # Sixth is the env index address for exponents
            lc_bas[6 + BAS_SLOTS*ib] = off
            env[off+1:off+Ne] .= B.exp
            off += Ne
            # Seventh is the env index address for contraction coeff
            lc_bas[7 + BAS_SLOTS*ib] = off
            env[off+1:off+Nc] .= B.coef
            off += Nc
            # Eigth, nothing
            ib += 1
        end
    end

    return BasisSet(name, atoms, basis, natm, nbas, nshells, lc_atm, lc_bas, env)
end

"""
    GaussianBasis.gto_norm(n::Signed, a::AbstractFloat)

Function that returns the normalization factor for a gaussian basis function rⁿ⋅exp(-a⋅r²)
"""
function gto_norm(n, a)
   # normalization factor of function rⁿ exp(-ar²)
    s = 2^(2n+3) * factorial(n+1) * (2a)^(n+1.5) / (factorial(2n+2) * √π)
    return √s
end

"""
    GaussianBasis.normalize_basisfunction!(B::BasisFunction)

Modify the BasisFunction object by normalizing each primitive.
"""
function normalize_basisfunction!(B::BasisFunction)
    for i = eachindex(B.coef)
        B.coef[i] *= gto_norm(B.l, B.exp[i])
    end
end

function getindex(B::BasisSet, N::Int)
    return B.basis[N]
end

function getindex(B::BasisSet, i::Int, j::Int)
    return B.basis[i][j]
end

function string_repr(B::BasisFunction)
    # Generate Unicode symbol for sub number
    l_sub = Char(0x2080 + B.l)

    # Unicode for superscript is a bit messier, so gotta use control flow
    l_sup = B.l == 1 ? Char(0x00B9) :
            B.l in [2,3] ? Char(0x00B1 + B.l) :
            Char(0x2070 + B.l)

    nbas = 2*B.l + 1
    mvals = collect(-B.l:B.l)
    nprim = length(B.exp)

    # Reverse Dict(symbol=>num) to get Symbols from B.l
    Lsymbol = Dict(value => key for (key, value) in AMDict)[B.l]
    out = "$(Lsymbol) shell with $nbas basis built from $nprim primitive gaussians\n\n"
    for m in mvals
        # Add sub minus sign (0x208B) if necessary
        m_sub = m < 0 ? Char(0x208B)*Char(0x2080 - m) : Char(0x2080 + m)
        out *= format("{:<4s} = ","χ$(l_sub)$m_sub")
        for i in eachindex(B.coef)

            if i > 1
                out *= B.coef[i] > 0 ? "\n     + " : "\n     - "
            end

            #out *= "$(abs(B.coef[i]))⋅Y$(l_sub)$m_sub"
            out *= format("{:>15.10f}⋅Y$(l_sub)$m_sub", abs(B.coef[i]))

            if B.l != 0
                out *= "⋅r$l_sup"
            end

            out *= "⋅exp(-$(B.exp[i])⋅r²)"
        end
        out *="\n\n"
    end
    return strip(out)
end

function string_repr(B::BasisSet)
    out  =  "$(B.name) Basis Set\n"
    out *= "Number of shells: $(B.nshells)\n"
    out *= "Number of basis:  $(B.nbas)\n\n"

    l_to_symbol = Dict(
        0 => "s",
        1 => "p",
        2 => "d",
        3 => "f",
        4 => "g",
        5 => "h",
        6 => "i",
    )

    for i in eachindex(B.atoms)
        A = B.atoms[i]
        # Count how many times s,p,d appears for numbering
        count = zeros(Int16, 7)
        out *= "$(symbol(A)): "
        for bfs in B.basis[i]
            L = bfs.l
            count[L+1] += 1
            out *= "$(count[L+1])$(l_to_symbol[L]) "
        end
        out *="\n"
    end

    return strip(out)
end

# Pretty printing
function show(io::IO, ::MIME"text/plain", X::T) where T<:Union{BasisFunction, BasisSet}
    print(io, string_repr(X))
end


push!(LOAD_PATH, "/Users/czhang/Desktop/Projects/dw_/julialib/")
import DimmWitted
DimmWitted.set_libpath("/Users/czhang/Desktop/Projects/dw_/libdw_julia")

immutable DoublePair
	d1::Cdouble
	d2::Cdouble
end

######################################
# The following piece of code creates a
# synthetic data set:
#    - Data type is Cdouble
#    - Modle type is Array{Cdouble}
#
nexp = 100000
nfeat = 1024
examples = Array(DoublePair, nexp, nfeat+1)
for row = 1:nexp
	for col = 1:nfeat
		examples[row, col] = DoublePair(1.0,1.0)
	end
	if rand() > 0.8
		examples[row, nfeat+1] = DoublePair(1.0,0.0)
	else
		examples[row, nfeat+1] = DoublePair(0.0,1.0)
	end
end
model = DoublePair[DoublePair(0.0,0.0) for i = 1:nfeat]

######################################
# Define the loss function and gradient
# function for logistic regression
#
function loss(row::Array{DoublePair,1}, model::Array{DoublePair,1})
	const label1 = row[length(row)].d1
	const label2 = row[length(row)].d2
	const nfeat = length(model)
	d1 = 0.0
	d2 = 0.0
	for i = 1:nfeat
		d1 = d1 + row[i].d1*model[i].d1
		d2 = d2 + row[i].d2*model[i].d2
	end
	v1 = (-label1 * d1 + log(exp(d1) + 1.0))
	v2 = (-label2 * d2 + log(exp(d2) + 1.0))
	return v1 + v2
end

function grad(row::Array{DoublePair,1}, model::Array{DoublePair,1})
	const label1 = row[length(row)].d1
	const label2 = row[length(row)].d2
	const nfeat = length(model)
	d1 = 0.0
	d2 = 0.0
	for i = 1:nfeat
		d1 = d1 + row[i].d1*model[i].d1
		d2 = d2 + row[i].d2*model[i].d2
	end
	d1 = exp(-d1)
	d2 = exp(-d2)
	Z1 = 0.00001 * (-label1 + 1.0/(1.0+d1))
	Z2 = 0.00001 * (-label2 + 1.0/(1.0+d2))
  	for i = 1:nfeat
  		model[i] = DoublePair(model[i].d1 - row[i].d1 * Z1, model[i].d2 - row[i].d2 * Z2)
  	end
	return 1.0
end

######################################
# Create a DimmWitted object using data
# and model. You do not need to specify
# the type, they are infer'ed by the
# open() function, which is parametric.
#
dw = DimmWitted.open(examples, model,
                DimmWitted.MR_PERMACHINE,
                DimmWitted.DR_SHARDING,
                DimmWitted.AC_ROW)

######################################
# Register functions.
#
handle_loss = DimmWitted.register_row(dw, loss, true)
handle_grad = DimmWitted.register_row(dw, grad, true)
#DimmWitted.register_model_avg(dw, handle_loss, avg, true)
#DimmWitted.register_model_avg(dw, handle_grad, avg, true)

######################################
# Run 10 epoches.
#
for iepoch = 1:10
	rs = DimmWitted.exec(dw, handle_loss)
	println("LOSS: ", rs/nexp)
	rs = DimmWitted.exec(dw, handle_grad)
end

sum1 = 0.0
sum2 = 0.0
for i = 1:length(model)
	sum1 = sum1 + model[i].d1
	sum2 = sum2 + model[i].d2
end
println("SUM OF MODEL1: ", sum1)
println("SUM OF MODEL2: ", sum2)


"""
    AlleeExtinction <: CellRule

    AlleeExtinction(minfounders)
    AlleeExtinction{R}(; minfounders=5.0)
    AlleeExtinction{R,W}(minfounders)

Causes extinction in a cell when a population is below a minimum number of individuals.

- `minfounders`: minimum founding individuals required for a viable population.
    Must be a type that can be compared to the grid values using `isless`.

Pass grid `Symbol`s to `R` or both `R` and `W` type parameters to use to specific grids.
"""
struct AlleeExtinction{R,W,MF} <: CellRule{R,W}
    minfounders::MF
end
AlleeExtinction{R,W}(; minfounders=Param(5.0, bounds=(1.0, 200.0))) where {R,W} =
    AlleeExtinction{R,W}(minfounders)

@inline function applyrule(data, rule::AlleeExtinction, N, I)
    f = get(data, rule.minfounders, I...)

    return N >= f ? N : zero(N)
end


function FieldStopOnIncreasing()
    return FieldStopOnIncreasing(())
end

function event_occurred(obj::FieldStopOnIncreasing, arg0::FieldSpacecraftState, arg1::FieldEventDetector, arg2::jboolean)
    return jcall(obj, "eventOccurred", Action, (FieldSpacecraftState, FieldEventDetector, jboolean), arg0, arg1, arg2)
end


abstract type AbstractIntegratorVPRK{DT,TT,D,S} <: AbstractIntegratorIRK{DT,TT} end
abstract type AbstractIntegratorVPRKwProjection{DT,TT,D,S} <: AbstractIntegratorVPRK{DT,TT,D,S} end

@inline parameters(integrator::AbstractIntegratorVPRK) = integrator.params

@inline nstages(integrator::AbstractIntegratorVPRK{DT,TT,D,S}) where {DT,TT,D,S} = S
@inline Base.ndims(integrator::AbstractIntegratorVPRK{DT,TT,D,S}) where {DT,TT,D,S} = D

@inline eachstage(integrator::AbstractIntegratorVPRK) = 1:nstages(integrator)
@inline eachdim(integrator::AbstractIntegratorVPRK) = 1:ndims(integrator)


fname = :ARCLDLt
c = Combi(hessian_dense,PDataLDLt,solve_modelARCDiag,preprocessLDLt,decreaseFact,Tparam())
include("Template.jl")


################################################################################
# Custom combined place-and-route
################################################################################
function lowtemp_pnr(m::Map, iterations; place_kwargs...)
    # Do a more low-level treatment of the place and route routine - acquire
    # a SAStruct and use it directly. If routing fails, do a low-temp placement
    # to perturb the location of processors and try routing again.
    #
    # Repeat for the specified number of iterations.
    pstruct = SAStruct(m)
    # First placement - capture the final state.
    state = SA.place(pstruct; place_kwargs...)
    # Record the placement back to the Map.
    Mapper2.SA.record(m, pstruct)

    # Try routing
    m = route(m)
    # If it passes a routing check, return.
    check_routing(m, true) && return m

    @info """
        Initial attempt at routing failed. Trying again.
        """

    for i in 1:iterations
        # Low temperature perturbation.
        SA.place(pstruct;
                 place_kwargs...,
                 supplied_state = state,
                )
        SA.record(m, pstruct)
        # Try routing again.
        m = route(m);
        check_routing(m, true) && return m
    end
    @error "Mapping Failed!"

    return m
end

function basic_place(arch::TopLevel, taskgraph::Taskgraph)
    m = NewMap(arch, taskgraph)
    place(m)
    MapperCore.save(m, "temp.jls")
end

function basic_route(arch::TopLevel, taskgraph::Taskgraph)
    m = NewMap(arch, taskgraph)
    MapperCore.load(m, "temp.jls")
    route(m)
    return m.mapping
end

################################################################################
# Custom Placement Functions
################################################################################
function shotgun_placement(arch     ::TopLevel,
                           taskgraph::Taskgraph;
                           nplacements::Int = 1,
                           nsamples::Int    = 1,
                           kwargs...)

    @info "Total Placements: $(nplacements)"
    # closure for parallel placement
    p(x) = low_temp_placement(arch, taskgraph, x;
                              nsamples = nsamples,
                              kwargs...)

    # Parallelize placements
    pmap(x -> p(x), 1:nplacements)
end

function low_temp_placement(arch            ::TopLevel,
                            taskgraph       ::Taskgraph,
                            iteration_number::Int;
                            nsamples = 1,
                            place_kwargs...)

    @info "Sublacement iteration $iteration_number"
    # Construct a new Map object
    m = NewMap(arch, taskgraph)
    # Get the placement structure
    pstruct = SAStruct(m)
    # Scope "state" out of the loop to avoid renaming.
    local state
    for i = 1:nsamples
        if i == 1
            state = SA.place(pstruct; place_kwargs...)
        else
            SA.place(pstruct;
                    place_kwargs...,
                    supplied_state = state,
                    warmer = SA.TrueSAWarm()
                )
        end
        Mapper2.SA.record(m, pstruct)

        # serialize to temp/ dir so routing can find it.
        savename    = "$(iteration_number)_$(i)"
        save_path   = joinpath(PKGDIR, "temp", savename)
        Mapper2.MapperCore.save(m, save_path)
    end
end

################################################################################
# Custom Routing Functions
################################################################################
function low_temp_route(arch, taskgraph;
                        nplacements::Int    = 1,
                        nsamples::Int       = 1,
                        route_kwargs...)

    # create closure for parallelism
    r(x) = multisample_routing(arch, taskgraph, x, nsamples)

    results = pmap(x -> r(x), 1:nplacements)
    return results
end

function multisample_routing(arch, taskgraph, iteration_number, nsamples)
    results = Mapping[]

    m = NewMap(arch, taskgraph)
    for i in 1:nsamples
        # load serialized placement mapping
        load_path = joinpath(PKGDIR, "temp", "$(iteration_number)_$(i)")
        Mapper2.MapperCore.load(m, load_path)
        try
            m = route(m)
        catch
        end
        push!(results, m.mapping)
    end
    return results
end


function _JBox(
    cornerpoint1::Point,
    cornerpoint2::Point,
    color,
    action::Symbol,
    vertices::Bool,
)
    sethue(color)
    verts = box(cornerpoint1, cornerpoint2, action, vertices = vertices)
    return verts[2]
end
function _JBox(points::Array, color, action::Symbol, vertices::Bool)
    sethue(color)
    verts = box(points, action, vertices = vertices)
    return verts[2]
end
function _JBox(pt::Point, width::Real, height::Real, color, action::Symbol, vertices::Bool)
    sethue(color)
    box(pt, width, height, action, vertices = vertices)
    return Point(pt.x - width / 2, pt.y + height / 2)
end
function _JBox(
    pt::Point,
    width::Real,
    height::Real,
    cornerradius::Float64,
    color,
    action::Symbol,
)
    sethue(color)
    box(pt, width, height, cornerradius, action)
    return Point(pt.x - width / 2, pt.y + height / 2)
end

"""
    JBox(cornerpoint1::Point, cornerpoint2::Point; color="black", action=:stroke, vertices=false)

Create a box (rectangle) between two points and do an action.
Returns the top left corner point of the box.
"""
JBox(
    cornerpoint1::Point,
    cornerpoint2::Point;
    color = "black",
    action = :stroke,
    vertices = false,
) =
    (
        args...;
        cornerpoint1 = cornerpoint1,
        cornerpoint2 = cornerpoint2,
        color = color,
        action = action,
        vertices = vertices,
    ) -> _JBox(cornerpoint1, cornerpoint2, color, action, vertices)

"""
    JBox(points::Array; color="black", action=:stroke, vertices=false)

Create a box/rectangle using the first two points of an array of Points to defined opposite corners.
Returns the top left corner point of the box.
"""
JBox(points::Array; color = "black", action = :stroke, vertices = false) =
    (args...; points = points, color = color, action = action, vertices = vertices) ->
        _JBox(points, color, action, vertices)

"""
    JBox(pt::Point, width::Real, height::Real; color="black", action=:stroke, vertices=false)

Create a box/rectangle centered at point pt with width and height. Use vertices=true to return an array of the four corner points rather than draw the box.
"""
JBox(
    pt::Point,
    width::Real,
    height::Real;
    color = "black",
    action = :stroke,
    vertices = false,
) =
    (
        args...;
        pt = pt,
        width = width,
        height = height,
        color = color,
        action = action,
        vertices = vertices,
    ) -> _JBox(pt, width, height, color, action, vertices)

"""
    JBox(x::Int64, y::Int64, width::Real, height::Real; color="black", action=:stroke)

Create a box/rectangle centered at point x/y with width and height.
"""
JBox(x::Int64, y::Int64, width::Real, height::Real; color = "black", action = :stroke) =
    JBox(Point(x, y), width, height, color = color, action = action)

"""
    JBox(pt::Point, width::Real, height::Real, cornerradius::Float64; color="black", action=:stroke)

Draw a box/rectangle centered at point pt with width and height and round each corner by cornerradius.
"""
JBox(
    pt::Point,
    width::Real,
    height::Real,
    cornerradius::Float64;
    color = "black",
    action = :stroke,
) =
    (
        args...;
        pt = pt,
        width = width,
        height = height,
        cornerradius = cornerradius,
        color = color,
        action = action,
    ) -> _JBox(pt, width, height, cornerradius, color, action)


####################
# value extraction #
####################

@inline extract_value!(::Type{T}, out::DiffResult, ydual) where {T} =
    DiffResults.value!(d -> value(T,d), out, ydual)
@inline extract_value!(::Type{T}, out, ydual) where {T} = out # ???

@inline function extract_value!(::Type{T}, out, y, ydual) where {T}
    map!(d -> value(T,d), y, ydual)
    copy_value!(out, y)
end

@inline copy_value!(out::DiffResult, y) = DiffResults.value!(out, y)
@inline copy_value!(out, y) = out

###################################
# vector mode function evaluation #
###################################

@generated function dualize(::Type{T}, x::SArray{S,V,D,N}) where {T,S,V,D,N}
    dx = Expr(:tuple, [:(Dual{T}(x[$i], chunk, Val{$i}())) for i in 1:N]...)
    return quote
        chunk = Chunk{N}()
        $(Expr(:meta, :inline))
        return SArray{S}($(dx))
    end
end

@inline static_dual_eval(::Type{T}, f, x::SArray) where {T} = f(dualize(T, x))

function vector_mode_dual_eval(f, x, cfg::Union{JacobianConfig,GradientConfig})
    xdual = cfg.duals
    seed!(xdual, x, cfg.seeds)
    return f(xdual)
end

function vector_mode_dual_eval(f!, y, x, cfg::JacobianConfig)
    ydual, xdual = cfg.duals
    seed!(xdual, x, cfg.seeds)
    seed!(ydual, y)
    f!(ydual, xdual)
    return ydual
end

##################################
# seed construction/manipulation #
##################################

@generated function construct_seeds(::Type{Partials{N,V}}) where {N,V}
    return Expr(:tuple, [:(single_seed(Partials{N,V}, Val{$i}())) for i in 1:N]...)
end

function seed!(duals::AbstractArray{Dual{T,V,N}}, x,
               seed::Partials{N,V} = zero(Partials{N,V})) where {T,V,N}
    for i in eachindex(duals)
        duals[i] = Dual{T,V,N}(x[i], seed)
    end
    return duals
end

function seed!(duals::AbstractArray{Dual{T,V,N}}, x,
               seeds::NTuple{N,Partials{N,V}}) where {T,V,N}
    for i in 1:N
        duals[i] = Dual{T,V,N}(x[i], seeds[i])
    end
    return duals
end

function seed!(duals::AbstractArray{Dual{T,V,N}}, x, index,
               seed::Partials{N,V} = zero(Partials{N,V})) where {T,V,N}
    offset = index - 1
    for i in 1:N
        j = i + offset
        duals[j] = Dual{T,V,N}(x[j], seed)
    end
    return duals
end

function seed!(duals::AbstractArray{Dual{T,V,N}}, x, index,
               seeds::NTuple{N,Partials{N,V}}, chunksize = N) where {T,V,N}
    offset = index - 1
    for i in 1:chunksize
        j = i + offset
        duals[j] = Dual{T,V,N}(x[j], seeds[i])
    end
    return duals
end


using Printf
using Grids

struct TrainingSample
    # @todo make training sample 3 fields: game, position, steps
    image::AbstractImage
    to_play::Integer # This information is encoded in the image, and is only needed for logging purposes.
    actions::Vector{<:Integer}
    targets::Vector{Target}
    image_label::String
    # @todo change image_state to initial_state
    image_state::AbstractEnvState
    states::Vector{AbstractEnvState}
    target_states::Vector{AbstractEnvState}
    function TrainingSample(image, to_play, actions, targets, states, target_states; image_label, image_state)
        @assert length(targets) == length(actions) + 1 "Num targets ($(length(targets))) must be 1 more than num actions ($(length(actions)))"
        @assert size(image)[end] == 3 "Image/Observation must have 3 channels (has $(size(image)[end])"
        new(image, to_play, actions, targets, image_label, image_state, states, target_states)
    end
end

struct TrainingSampleSVG
    sample::TrainingSample
end

svg(x::TrainingSample) = TrainingSampleSVG(x)

actions(x::TrainingSample) = x.actions
targets(x::TrainingSample) = x.targets
image(x::TrainingSample) = x.image
image_channels(x::TrainingSample) = x |> image |> channels

function TrainingSample(game::MuZeroGame, position::Integer, num_unroll_steps::Integer, td_steps::Integer)
    image = make_image(game, position)
    image_state = get_state(game, position)
    image_label = make_image_label(game, position)
    to = min(length(history(game)), position + num_unroll_steps - 1)
    actions = map(x->x.idx, history(game)[position:to])
    targets = make_target(game, position, num_unroll_steps, td_steps)
    num_targets = length(actions)+1
    player = to_play(position)
    states = map(idx->get_state(game, idx), position:(to+1))
    game_copy = copy_at(game.game, position)
    target_states = map(target->begin
        greedy_action = argmax(policy(target))
        step!(game_copy, greedy_action)
        get_state(game_copy)
    end, targets)
    TrainingSample(image, player, actions, targets[1:num_targets], states, target_states; image_label=image_label, image_state=image_state)
end

function TrainingSample(config::MuZeroConfig, game::MuZeroGame)
    TrainingSample(game, sample_position(game), config.num_unroll_steps, config.td_steps)
end

tuple(x::TrainingSample) = x.image, x.actions, x.targets

get_states(x::TrainingSample) = x.states
get_target_states(x::TrainingSample) = x.target_states

function getdf(x::TrainingSample)
    inftype_col = "Inf Type"
    action_col = "Action"
    reward_col = "Target Reward"
    value_col = "Target Value"
    policy_col = "Target Policy"
    d = Dict(
        inftype_col=>String[],
        action_col=>Union{Int,Nothing}[],
        reward_col=>Float32[],
        value_col=>Float32[],
        policy_col=>String[],
    )
    # @todo change this to x.action_space_size
    num_policy_actions = x.targets |> first |> policy |> length
    for i in 1:num_policy_actions
        d["π[$i]"] = Float32[]
    end

    states = get_states(x)
    actions = [nothing, x.actions...]
    targets = x.targets
    for (action_idx, zipped) in enumerate(zip(states, actions, targets))
        (s,a,target) = zipped
        push!(d[inftype_col], action_idx == 1 ? "Initial" : "Recurrent")
        push!(d[action_col], a)
        push!(d[reward_col], reward(target))
        push!(d[value_col], value(target))

        pred = policy(target)
        limit = 3
        from = 1
        to = min(limit, length(pred))
        actions = reverse(sortperm(pred))[from:to]
        prs = pred[actions]
        action_prs = map(x->begin
            action, pr = x
            pr_s = @sprintf "%.2f" pr
            "$action ($(pr_s[1:end]))"
        end, zip(actions, prs))
        result_policy = join(action_prs, ", ")
        push!(d[policy_col], result_policy)
        # π = policy(target)
        # for i in 1:num_policy_actions
        #     push!(d["π[$i]"], round(π[i], sigdigits=2))
        # end
    end

    DataFrame(
        inftype_col=>d[inftype_col],
        action_col=>d[action_col],
        reward_col=>d[reward_col],
        value_col=>d[value_col],
        policy_col=>d[policy_col],
    )
end

function getvrgrid(target)
    value = @sprintf "%.1f" target.value
    reward = @sprintf "%.1f" target.reward
    vr = Grid(
        [
            LabelRect("value: $value", width=200, halign=:left),
            LabelRect("reward: $reward", width=200, halign=:left),
        ],
        :,
        1,
        margin_left=0,
        line_thickness=0
    )
    Grid(
        [vr],
        :,
        1,
        header="TARGET V/R",
        margin_left=10,
        line_thickness=0
    )
end

function getinputgrid(s, header)
    player_name = s.to_play.name
    state_grid = getgrid(s; footer="to play: $player_name", margin_left=0)
    Grid(
        [state_grid],
        :,
        1,
        header=header,
        margin_left=0,
        line_thickness=0
    )
end

function getpolicygrid(s, target)
    pred = target.policy
    pred_states = Rect[]
    limit = 3
    from = 1
    to = min(limit, length(pred))
    pred_actions = reverse(sortperm(pred))[from:to]
    prs = pred[pred_actions]
    for (a,pr) in zip(pred_actions, prs)
        pr_str = @sprintf "%.2f" pr
        push!(
            pred_states,
            getgrid(
                s,
                a,
                0f0,
                getnextstate(s, a);
                margin_left=0,
                footer="action: $a",
                subfooter=(text=pr_str, color="#666"),
            )
        )
    end
    Grid(
        pred_states,
        :,
        1;
        header="TARGET π",
        margin_left=0,
        line_thickness=0,
    )
end

function getinferencegrid(s, target, header, input_header)
    input_grid = getinputgrid(s, input_header)
    vr_grid = getvrgrid(target)
    policy_grid = getpolicygrid(s, target)
    policy_col = Grid([Grid(3,3), Grid(3,3)], :, 1, header=(text="TARGET π", offset=(10,0)))
    Grid(
        [input_grid, vr_grid, policy_grid],
        header=header,
        margin_left=0,
        line_thickness=0,
    )
end

function getrecurrentinfgrid(sample)
    state_targets = zip(sample.states[2:end],sample.targets[2:end])
    map(st->begin
        s, target = st
        getinferencegrid(
            s,
            target,
            "RECURRENT INFERENCE",
            "HIDDEN STATE"
        )
    end, state_targets)
end

function Base.show(io::IO, m::MIME"image/svg+xml", x::TrainingSampleSVG)
    sample = x.sample
    initial_grid = getinferencegrid(
        sample.image_state,
        sample.targets[1],
        "INITIAL INFERENCE",
        "INPUT"
    )
    recurrent_grids = getrecurrentinfgrid(sample)
    g = Grid(
        [initial_grid, recurrent_grids...],
        :,
        1,
        header=(text="TRAINING SAMPLE"),
        line_thickness=0,
        margin=15,
    )
    show(io, m, g)
end

function Base.show(io::IO, x::TrainingSample)
    action_hl = Highlighter((data, i, j)->(j==2) && data[i,j] == nothing , crayon"#666666");
    pretty_table(
        io,
        getdf(x);
        show_row_number=true,
        nosubheader=true,
        title="Training Sample",
        highlighters=(action_hl),
    )
end

function gridsvg(x::TrainingSample)
    first_input_state, remaining_states... = get_states(x)
    first_target_state, remaining_target_states... = x.target_states
    first_target, remaining_targets... = x.targets
    state_targets = zip(get_states(x), x.targets)
    state_type = typeof(x.image_state)
    initial_grid = GridSVG([
        gridsvg(first_input_state),
        targetsvg(first_input_state, first_target_state, first_target)
    ], dim=2)
    grids = [
        initial_grid,
    ]
    for (input_state, action, target, target_state) in zip(remaining_states, x.actions, remaining_targets, remaining_target_states)
        push!(grids, GridSVG([
            gridsvg(input_state, title="hiddenstate + action")
            targetsvg(input_state, target_state, target)
        ], dim=2))
    end
    GridSVG(grids, main_title="Training Sample")
end

function get_player_label(x::TrainingSample, position)
    if x.to_play == 1
        return isodd(position) ? "A" : "B"
    end
    isodd(position) ? "B" : "A"
end


mutable struct SummaryStat
    step::Int64
    totalStep::Int64
    visited::Vector{Float64}
    reweight::Vector{Float64} # only record the reweight factor of a typical configuration
    propose::Array{Float64,3}
    accept::Array{Float64,3}
end

function addStat(config, summary = nothing)
    if isnothing(summary)
        return SummaryStat(config.step, config.totalStep, config.visited, config.reweight, config.propose, config.accept)
    else
        summary = deepcopy(summary)
        summary.step += config.step
        summary.totalStep += config.totalStep
        summary.visited += config.visited
        summary.reweight = config.reweight # only record the newest reweight factor
        summary.propose += config.propose
        summary.accept += config.accept
        return summary
    end
end

# sum SummaryStat from all workers to the root
function reduceStat(summary, root, comm)
    step = MPI.Reduce(summary.step, MPI.SUM, root, comm)
    totalStep = MPI.Reduce(summary.totalStep, MPI.SUM, root, comm)
    visited = MPI.Reduce(summary.visited, MPI.SUM, root, comm)
    reweight = MPI.Reduce(summary.reweight, MPI.SUM, root, comm)
    propose = MPI.Reduce(summary.propose, MPI.SUM, root, comm)
    accept = MPI.Reduce(summary.accept, MPI.SUM, root, comm)
    if MPI.Comm_rank(comm) == root
        reweight ./= MPI.Comm_size(comm)
        return SummaryStat(step, totalStep, visited, reweight, propose, accept)
    else
        return summary
    end
end

function MPIreduce(data)
    comm = MPI.COMM_WORLD
    Nworker = MPI.Comm_size(comm)  # number of MPI workers
    rank = MPI.Comm_rank(comm)  # rank of current MPI worker
    root = 0 # rank of the root worker

    if Nworker == 1 #no parallelization
        return data
    end
    if typeof(data) <: AbstractArray
        MPI.Reduce!(data, MPI.SUM, root, comm) # root node gets the sum of observables from all blocks
        return data
    else
        result = [data,]  # MPI.Reduce works for array only
        MPI.Reduce!(result, MPI.SUM, root, comm) # root node gets the sum of observables from all blocks
        return result[1]
    end
end


# Conversions
# -----------

# no-op conversions

for CV in CVparametric
    @eval begin
        convert{T}(::Type{$CV{T}}, c::$CV{T}) = c
        convert(::Type{$CV}, c::$CV) = c
    end
end
convert(::RGB24, c::RGB24) = c
convert(::ARGB32, c::ARGB32) = c

# changing the datatype without changing the underlying colorspace
# Note: this assumes that the underlying type constructor doesn't swap the layout!
for CV in CVparametric
    @eval begin
        convert{T}(::Type{$CV{T}}, c::$CV) = $CV{T}(convert(T, getfield(c, 1)), convert(T, getfield(c, 2)), convert(T, getfield(c, 3)))
    end
end

# Conversions where the datatype is not specified
for (ElementClass, Colorspace) in [(Fractional, CVfractional), (FloatingPoint, CVfloatingpoint)]
    for CV in Colorspace
        @eval begin
            # preserves the datatype of the original space
            convert{T<:$ElementClass}(::Type{$CV}, c::ColorValue{T}) = convert($CV{T}, c)
            # fallback is Float64 (needed for RGB24)
            convert{T}(::Type{$CV}, c::ColorValue{T}) = convert($CV{Float64}, c)
        end
    end
end
convert{T}(::Type{RGB}, c::ColorValue{T}) = _convert(RGB{T}, c)
convert{T}(::Type{RGB{T}}, c::ColorValue) = _convert(RGB{T}, c)
convert(::Type{RGB}, c::RGB24) = _convert(RGB{Float64}, c)

# Everything to RGB
# -----------------

function _convert{CV<:AbstractRGB}(::Type{CV}, c::HSV)
    h = c.h / 60
    i = floor(h)
    f = h - i
    if round(Int, i) & 1 == 0
        f = 1 - f
    end
    m = c.v * (1 - c.s)
    n = c.v * (1 - c.s * f)
    i = round(Int, i)
    if i == 6 || i == 0; CV(c.v, n, m)
    elseif i == 1;       CV(n, c.v, m)
    elseif i == 2;       CV(m, c.v, n)
    elseif i == 3;       CV(m, n, c.v)
    elseif i == 4;       CV(n, m, c.v)
    else;                CV(c.v, m, n)
    end
end

function _convert{CV<:AbstractRGB}(::Type{CV}, c::HSL)
    function qtrans(u, v, hue)
        if     hue > 360; hue -= 360
        elseif hue < 0;   hue += 360
        end

        if     hue < 60;  u + (v - u) * hue / 60
        elseif hue < 180; v
        elseif hue < 240; u + (v - u) * (240 - hue) / 60
        else;             u
        end
    end

    v = c.l <= 0.5 ? c.l * (1 + c.s) : c.l + c.s - (c.l * c.s)
    u = 2 * c.l - v

    if c.s == 0; CV(c.l, c.l, c.l)
    else;        CV(qtrans(u, v, c.h + 120),
                        qtrans(u, v, c.h),
                        qtrans(u, v, c.h - 120))
    end
end

const M_XYZ_RGB = [ 3.2404542 -1.5371385 -0.4985314
                   -0.9692660  1.8760108  0.0415560
                    0.0556434 -0.2040259  1.0572252 ]

correct_gamut{CV<:AbstractRGB}(c::CV) = CV(clamp01(c.r), clamp01(c.g), clamp01(c.b))
clamp01{T<:Fractional}(v::T) = ifelse(v < zero(T), zero(T), ifelse(v > one(T), one(T), v))

function srgb_compand(v)
    v <= 0.0031308 ? 12.92v : 1.055v^(1/2.4) - 0.055
end

function _convert{CV<:AbstractRGB}(::Type{CV}, c::XYZ)
    ans = M_XYZ_RGB * [c.x, c.y, c.z]
    correct_gamut(CV(srgb_compand(ans[1]),
                     srgb_compand(ans[2]),
                     srgb_compand(ans[3])))
end

_convert{CV<:AbstractRGB}(::Type{CV}, c::xyY)    = _convert(CV, convert(XYZ{eltype(c)}, c))
_convert{CV<:AbstractRGB}(::Type{CV}, c::Lab)    = _convert(CV, convert(XYZ{eltype(c)}, c))
_convert{CV<:AbstractRGB}(::Type{CV}, c::LCHab)  = _convert(CV, convert(Lab{eltype(c)}, c))
_convert{CV<:AbstractRGB}(::Type{CV}, c::Luv)    = _convert(CV, convert(XYZ{eltype(c)}, c))
_convert{CV<:AbstractRGB}(::Type{CV}, c::LCHuv)  = _convert(CV, convert(Luv{eltype(c)}, c))
_convert{CV<:AbstractRGB}(::Type{CV}, c::DIN99)  = _convert(CV, convert(XYZ{eltype(c)}, c))
_convert{CV<:AbstractRGB}(::Type{CV}, c::DIN99o) = _convert(CV, convert(XYZ{eltype(c)}, c))
_convert{CV<:AbstractRGB}(::Type{CV}, c::DIN99d) = _convert(CV, convert(XYZ{eltype(c)}, c))
_convert{CV<:AbstractRGB}(::Type{CV}, c::LMS)    = _convert(CV, convert(XYZ{eltype(c)}, c))

_convert{CV<:AbstractRGB{Ufixed8}}(::Type{CV}, c::RGB24) = CV(Ufixed8(c.color&0x00ff0000>>>16,0), Ufixed8(c.color&0x0000ff00>>>8,0), Ufixed8(c.color&0x000000ff,0))
_convert{CV<:AbstractRGB}(::Type{CV}, c::RGB24) = CV((c.color&0x00ff0000>>>16)/255, ((c.color&0x0000ff00)>>>8)/255, (c.color&0x000000ff)/255)

# Everything to HSV
# -----------------

function convert{T}(::Type{HSV{T}}, c::AbstractRGB)
    c_min = min(c.r, c.g, c.b)
    c_max = max(c.r, c.g, c.b)
    if c_min == c_max
        return HSV(zero(T), zero(T), c_max)
    end

    if c_min == c.r
        f = c.g - c.b
        i = convert(T, 3)
    elseif c_min == c.g
        f = c.b - c.r
        i = convert(T, 5)
    else
        f = c.r - c.g
        i = convert(T, 1)
    end

    HSV{T}(60 * (i - f / (c_max - c_min)),
        (c_max - c_min) / c_max,
        c_max)
end

convert{T}(::Type{HSV{T}}, c::ColorValue) = convert(HSV{T}, convert(RGB{T}, c))

# Everything to HSL
# -----------------

function convert{T}(::Type{HSL{T}}, c::AbstractRGB)
    c_min = min(c.r, c.g, c.b)
    c_max = max(c.r, c.g, c.b)
    l = (c_max + c_min) / 2

    if c_max == c_min
        return HSL(zero(T), zero(T), l)
    end

    if l < 0.5; s = (c_max - c_min) / (c_max + c_min)
    else;       s = (c_max - c_min) / (convert(T, 2) - c_max - c_min)
    end

    if c_max == c.r
        h = (c.g - c.b) / (c_max - c_min)
    elseif c_max == c.g
        h = convert(T, 2) + (c.b - c.r) / (c_max - c_min)
    else
        h = convert(T, 4) + (c.r - c.g) / (c_max - c_min)
    end

    h *= 60
    if h < 0
        h += 360
    elseif h > 360
        h -= 360
    end

    HSL{T}(h,s,l)
end

convert{T}(::Type{HSL{T}}, c::ColorValue) = convert(HSL{T}, convert(RGB{T}, c))

# Everything to XYZ
# -----------------

function invert_rgb_compand(v)
    v <= 0.04045 ? v/12.92 : ((v+0.055) /1.055)^2.4
end

const M_RGB_XYZ =
    [ 0.4124564  0.3575761  0.1804375
      0.2126729  0.7151522  0.0721750
      0.0193339  0.1191920  0.9503041 ]

function convert{T}(::Type{XYZ{T}}, c::AbstractRGB)
    v = [invert_rgb_compand(c.r),
         invert_rgb_compand(c.g),
         invert_rgb_compand(c.b)]
    ans = M_RGB_XYZ * v
    XYZ{T}(ans[1], ans[2], ans[3])
end

convert{T}(::Type{XYZ{T}}, c::HSV) = convert(XYZ{T}, convert(RGB{T}, c))
convert{T}(::Type{XYZ{T}}, c::HSL) = convert(XYZ{T}, convert(RGB{T}, c))

function convert{T}(::Type{XYZ{T}}, c::xyY)
    X = c.Y*c.x/c.y
    Z = c.Y*(1-c.x-c.y)/c.y
    XYZ{T}(X, c.Y, Z)
end

const xyz_epsilon = 216 // 24389
const xyz_kappa   = 24389 // 27

function convert{T}(::Type{XYZ{T}}, c::Lab, wp::XYZ)
    fy = (c.l + 16) / 116
    fx = c.a / 500 + fy
    fz = fy - c.b / 200

    fx3 = fx^3
    fz3 = fz^3

    x = fx3 > xyz_epsilon ? fx3 : (116fx - 16) / xyz_kappa
    y = c.l > xyz_kappa * xyz_epsilon ? ((c. l+ 16) / 116)^3 : c.l / xyz_kappa
    z = fz3 > xyz_epsilon ? fz3 : (116fz - 16) / xyz_kappa

    XYZ{T}(x*wp.x, y*wp.y, z*wp.z)
end

convert{T}(::Type{XYZ{T}}, c::Lab)   = convert(XYZ{T}, c, WP_DEFAULT)
convert{T}(::Type{XYZ{T}}, c::LCHab) = convert(XYZ{T}, convert(Lab{T}, c))
convert{T}(::Type{XYZ{T}}, c::DIN99) = convert(XYZ{T}, convert(Lab{T}, c))
convert{T}(::Type{XYZ{T}}, c::DIN99o) = convert(XYZ{T}, convert(Lab{T}, c))

convert{T<:Ufixed}(::Type{XYZ{T}}, c::LCHab) = convert(XYZ{T}, convert(Lab{eltype(c)}, c))
convert{T<:Ufixed}(::Type{XYZ{T}}, c::DIN99) = convert(XYZ{T}, convert(Lab{eltype(c)}, c))
convert{T<:Ufixed}(::Type{XYZ{T}}, c::DIN99o) = convert(XYZ{T}, convert(Lab{eltype(c)}, c))

function xyz_to_uv(c::XYZ)
    d = c.x + 15c.y + 3c.z
    d==0 && return (d, d)
    u = 4c.x / d
    v = 9c.y / d
    return (u, v)
end

function convert{T}(::Type{XYZ{T}}, c::Luv, wp::XYZ = WP_DEFAULT)
    (u_wp, v_wp) = xyz_to_uv(wp)

    a = (52 * (c.l==0 ? zero(T) : c.l / (c.u + 13 * c.l * u_wp)) - 1) / 3
    y = c.l > xyz_kappa * xyz_epsilon ? wp.y * ((c.l + 16) / 116)^3 : wp.y * c.l / xyz_kappa
    b = -5y
    d = y * (39 * (c.l==0 ? zero(T) : c.l / (c.v + 13 * c.l * v_wp)) - 5)
    x = d==b ? zero(T) : (d - b) / (a + 1//3)
    z = a * x + b + zero(T)

    XYZ{T}(x, y, z)
end

convert{T}(::Type{XYZ{T}}, c::LCHuv) = convert(XYZ{T}, convert(Luv{T}, c))
convert{T<:Ufixed}(::Type{XYZ{T}}, c::LCHuv) = convert(XYZ{T}, convert(Luv{eltype(c)}, c))

function convert{T}(::Type{XYZ{T}}, c::DIN99d)

    # Go back to C-h space
    # FIXME: Clean this up (why is there no atan2d?)
    h = rad2deg(atan2(c.b,c.a)) + 50
    while h > 360; h -= 360; end
    while h < 0;   h += 360; end

    C = sqrt(c.a^2 + c.b^2)

    # Intermediate terms
    G = (exp(C/22.5)-1)/0.06
    f = G*sind(h - 50)
    ee = G*cosd(h - 50)

    l = (exp(c.l/325.22)-1)/0.0036
    # a = ee*cosd(50) - f/1.14*sind(50)
    a = ee*0.6427876096865393 - f/1.14*0.766044443118978
    # b = ee*sind(50) - f/1.14*cosd(50)
    b = ee*0.766044443118978 - f/1.14*0.6427876096865393

    adj = convert(XYZ, Lab(l, a, b))

    XYZ{T}((adj.x + 0.12*adj.z)/1.12, adj.y, adj.z)

end

function convert{T}(::Type{XYZ{T}}, c::LMS)
    ans = CAT02_INV * [c.l, c.m, c.s]
    XYZ{T}(ans[1], ans[2], ans[3])
end

convert{T}(::Type{XYZ{T}}, c::RGB24) = convert(XYZ{T}, convert(RGB{T}, c))

# Everything to xyY
# -----------------

function convert{T}(::Type{xyY{T}}, c::XYZ)

    x = c.x/(c.x + c.y + c.z)
    y = c.y/(c.x + c.y + c.z)

    xyY{T}(x, y, convert(typeof(x), c.y))

end

convert{T}(::Type{xyY{T}}, c::ColorValue) = convert(xyY{T}, convert(XYZ{T}, c))

# Everything to Lab
# -----------------

convert{T}(::Type{Lab{T}}, c::AbstractRGB) = convert(Lab{T}, convert(XYZ{T}, c))
convert{T}(::Type{Lab{T}}, c::HSV) = convert(Lab{T}, convert(RGB{T}, c))
convert{T}(::Type{Lab{T}}, c::HSL) = convert(Lab{T}, convert(RGB{T}, c))

function convert{T}(::Type{Lab{T}}, c::XYZ, wp::XYZ)
    function f(v)
        v > xyz_epsilon ? cbrt(v) : (xyz_kappa * v + 16) / 116
    end

    fx, fy, fz = f(c.x / wp.x), f(c.y / wp.y), f(c.z / wp.z)
    Lab{T}(116fy - 16, 500(fx - fy), 200(fy - fz))
end

convert{T}(::Type{Lab{T}}, c::XYZ{T}) = convert(Lab{T}, c, WP_DEFAULT)

function convert{T}(::Type{Lab{T}}, c::LCHab)
    hr = deg2rad(c.h)
    Lab{T}(c.l, c.c * cos(hr), c.c * sin(hr))
end

function convert{T}(::Type{Lab{T}}, c::DIN99)

    # We assume the adjustment parameters are always 1; the standard recommends
    # that they not be changed from these values.
    kch = 1
    ke = 1

    # Calculate Chroma (C99) in the DIN99 space
    cc = sqrt(c.a^2 + c.b^2)

    # NOTE: This is calculated in degrees, against the standard, to save
    # computation steps later.
    if (c.a > 0 && c.b >= 0)
        h = atand(c.b/c.a)
    elseif (c.a == 0 && c.b > 0)
        h = 90
    elseif (c.a < 0)
        h = 180+atand(c.b/c.a)
    elseif (c.a == 0 && c.b < 0)
        h = 270
    elseif (c.a > 0 && c.b <= 0)
        h = 360 + atand(c.b/c.a)
    else
        h = 0
    end

    # Temporary variable for chroma
    g = (e^(0.045*cc*kch*ke)-1)/0.045

    # Temporary redness
    ee = g*cosd(h)

    # Temporary yellowness
    f = g*sind(h)

    # CIELAB a*b*
    # ciea = ee*cosd(16) - (f/0.7)*sind(16)
    ciea = ee*0.9612616959383189 - (f/0.7)*0.27563735581699916
    # cieb = ee*sind(16) + (f/0.7)*cosd(16)
    cieb = ee*0.27563735581699916 + (f/0.7)*0.9612616959383189

    # CIELAB L*
    ciel = (e^(c.l*ke/105.51)-1)/0.0158

    Lab{T}(ciel, ciea, cieb)
end

function convert{T}(::Type{Lab{T}}, c::DIN99o)

    # We assume the adjustment parameters are always 1; the standard recommends
    # that they not be changed from these values.
    kch = 1
    ke = 1

    # Calculate Chroma (C99) in the DIN99o space
    co = sqrt(c.a^2 + c.b^2)

    # hue angle h99o
    h = atan2(c.b, c.a)

    # revert rotation by 26°
    ho= rad2deg(h)-26

    # revert logarithmic chroma compression
    g = (e^(co*kch*ke/23.0)-1)/0.075

    # Temporary redness
    eo = g*cosd(ho)

    # Temporary yellowness
    fo = g*sind(ho)

    # CIELAB a*b* (revert b* axis compression)
    # ciea = eo*cosd(26) - (fo/0.83)*sind(26)
    ciea = eo*0.898794046299167 - (fo/0.83)*0.4383711467890774
    # cieb = eo*sind(26) + (fo/0.83)*cosd(26)
    cieb = eo*0.4383711467890774 + (fo/0.83)*0.898794046299167

    # CIELAB L* (revert logarithmic lightness compression)
    ciel = (e^(c.l*ke/303.67)-1)/0.0039

    Lab{T}(ciel, ciea, cieb)
end

convert{T}(::Type{Lab{T}}, c::ColorValue) = convert(Lab{T}, convert(XYZ{T}, c))

# Everything to Luv
# -----------------

convert{T}(::Type{Luv{T}}, c::AbstractRGB) = convert(Luv{T}, convert(XYZ{T}, c))
convert{T}(::Type{Luv{T}}, c::HSV) = convert(Luv{T}, convert(RGB{T}, c))
convert{T}(::Type{Luv{T}}, c::HSL) = convert(Luv{T}, convert(RGB{T}, c))

function convert{T}(::Type{Luv{T}}, c::XYZ, wp::XYZ = WP_DEFAULT)
    (u_wp, v_wp) = xyz_to_uv(wp)
    (u_, v_) = xyz_to_uv(c)

    y = c.y / wp.y

    l = y > xyz_epsilon ? 116 * cbrt(y) - 16 : xyz_kappa * y
    u = 13 * l * (u_ - u_wp) + zero(T)
    v = 13 * l * (v_ - v_wp) + zero(T)

    Luv{T}(l, u, v)
end

function convert{T}(::Type{Luv{T}}, c::LCHuv)
    hr = deg2rad(c.h)
    Luv{T}(c.l, c.c * cos(hr), c.c * sin(hr))
end

convert{T}(::Type{Luv{T}}, c::ColorValue) = convert(Luv{T}, convert(XYZ{T}, c))

# Everything to LCHuv
# -------------------

function convert{T}(::Type{LCHuv{T}}, c::Luv)
    h = rad2deg(atan2(c.v, c.u))
    while h > 360; h -= 360; end
    while h < 0;   h += 360; end
    LCHuv{T}(c.l, sqrt(c.u^2 + c.v^2), h)
end

convert{T}(::Type{LCHuv{T}}, c::ColorValue) = convert(LCHuv{T}, convert(Luv{T}, c))

# Everything to LCHab
# -------------------

function convert{T}(::Type{LCHab{T}}, c::Lab)
    h = rad2deg(atan2(c.b, c.a))
    while h > 360; h -= 360; end
    while h < 0;   h += 360; end
    LCHab{T}(c.l, sqrt(c.a^2 + c.b^2), h)
end

convert{T}(::Type{LCHab{T}}, c::ColorValue) = convert(LCHab{T}, convert(Lab{T}, c))

# Everything to DIN99
# -------------------

function convert{T}(::Type{DIN99{T}}, c::Lab)

    # We assume the adjustment parameters are always 1; the standard recommends
    # that they not be changed from these values.
    kch = 1
    ke = 1

    # Calculate DIN99 L
    l99 = (1/ke)*105.51*log(1+0.0158*c.l)

    # Temporary value for redness and yellowness
    # ee = c.a*cosd(16) + c.b*sind(16)
    ee = c.a*0.9612616959383189 + c.b*0.27563735581699916
    # f = -0.7*c.a*sind(16) + 0.7*c.b*cosd(16)
    f = -0.7*c.a*0.27563735581699916 + 0.7*c.b*0.9612616959383189

    # Temporary value for chroma
    g = sqrt(ee^2 + f^2)

    # Hue angle
    # Calculated in degrees, against the specification.
    if (ee > 0 && f >= 0)
        h = atand(f/ee)
    elseif (ee == 0 && f > 0)
        h = 90
    elseif (ee < 0)
        h = 180+atand(f/ee)
    elseif (ee == 0 && f < 0)
        h = 270
    elseif (ee > 0 && f <= 0)
        h = 360 + atand(f/ee)
    else
        h = 0
    end

    # DIN99 chroma
    cc = log(1+0.045*g)/(0.045*kch*ke)

    # DIN99 chromaticities
    a99, b99 = cc*cosd(h), cc*sind(h)

    DIN99{T}(l99, a99, b99)

end

convert{T}(::Type{DIN99{T}}, c::ColorValue) = convert(DIN99{T}, convert(Lab{T}, c))

# Everything to DIN99d
# --------------------

function convert{T}(::Type{DIN99d{T}}, c::XYZ{T})

    # Apply tristimulus-space correction term
    adj_c = XYZ(1.12*c.x - 0.12*c.z, c.y, c.z)

    # Apply L*a*b*-space correction
    lab = convert(Lab, adj_c)
    adj_L = 325.22*log(1+0.0036*lab.l)

    # Calculate intermediate parameters
    # ee = lab.a*cosd(50) + lab.b*sind(50)
    ee = lab.a*0.6427876096865393 + lab.b*0.766044443118978
    # f = 1.14*(lab.b*cosd(50) - lab.a*sind(50))
    f = 1.14*(lab.b*0.6427876096865393 - lab.a*0.766044443118978)
    G = sqrt(ee^2+f^2)

    # Calculate hue/chroma
    C = 22.5*log(1+0.06*G)
    # FIXME: Clean this up (why is there no atan2d?)
    h = rad2deg(atan2(f,ee)) + 50
    while h > 360; h -= 360; end
    while h < 0;   h += 360; end

    DIN99d{T}(adj_L, C*cosd(h), C*sind(h))

end

convert{T}(::Type{DIN99d{T}}, c::ColorValue) = convert(DIN99d{T}, convert(XYZ{T}, c))

# Everything to DIN99o
# -------------------

function convert{T}(::Type{DIN99o{T}}, c::Lab)

    # We assume the adjustment parameters are always 1; the standard recommends
    # that they not be changed from these values.
    kch = 1
    ke = 1

    # Calculate DIN99o L (logarithmic compression)
    l99 = 303.67/ke*log(1+0.0039*c.l)

    # Temporary value for redness and yellowness
    # including rotation by 26°
    # eo = c.a*cosd(26) + c.b*sind(26)
    eo = c.a*0.898794046299167 + c.b*0.4383711467890774
    # compression along the yellowness (blue-yellow) axis
    # fo = 0.83 * (c.b*cosd(26) - c.a*sind(26))
    fo = 0.83 * (c.b*0.898794046299167 - c.a*0.4383711467890774)

    # Temporary value for chroma
    go = sqrt(eo^2 + fo^2)
    ho = atan2(fo,eo)
    # rotation of the color space by 26°
    h  = rad2deg(ho) + 26

    # DIN99o chroma (logarithmic compression)
    cc = 23.0*log(1+0.075*go)/(kch*ke)

    # DIN99o chromaticities
    a99, b99 = cc*cosd(h), cc*sind(h)

    DIN99o{T}(l99, a99, b99)

end

convert{T}(::Type{DIN99o{T}}, c::ColorValue) = convert(DIN99o{T}, convert(Lab{T}, c))

# Everything to LMS
# -----------------

# Chromatic adaptation from CIECAM97s
const CAT97s = [ 0.8562  0.3372 -0.1934
                -0.8360  1.8327  0.0033
                 0.0357 -0.0469  1.0112 ]

const CAT97s_INV = inv(CAT97s)

# Chromatic adaptation from CIECAM02
const CAT02 = [ 0.7328 0.4296 -0.1624
               -0.7036 1.6975  0.0061
                0.0030 0.0136  0.9834 ]

const CAT02_INV = inv(CAT02)

function convert{T}(::Type{LMS{T}}, c::XYZ{T})
    ans = CAT02 * [c.x, c.y, c.z]
    LMS{T}(ans[1], ans[2], ans[3])
end

convert{T}(::Type{LMS{T}}, c::ColorValue) = convert(LMS{T}, convert(XYZ{T}, c))

# Everything to RGB24
# -------------------

convert(::Type{RGB24}, c::AbstractRGB{Ufixed8}) = RGB24(c.r, c.g, c.b)
convert(::Type{RGB24}, c::AbstractRGB) = RGB24(round(Uint32, 255*c.r)<<16 +
                                               round(Uint32, 255*c.g)<<8 +
                                               round(Uint32, 255*c.b))
to32(x::Ufixed8) = convert(Uint32, reinterpret(x))
convert(::Type{RGB24}, val::Uint32) = RGB24(val & 0x00ffffff)

convert(::Type{RGB24}, c::ColorValue) = convert(RGB24, convert(RGB{Ufixed8}, c))

# To Uint32
# ----------------

convert(::Type{Uint32}, c::RGB24) = c.color

convert(::Type{Uint32}, ac::ARGB32) = ac.color

#### AlphaColorValue

convert{C,T}(::Type{AlphaColorValue{C,T}}, c::AlphaColorValue{C,T}) = c
function convert{C,T,D,U}(::Type{AlphaColorValue{C,T}}, c::AbstractAlphaColorValue{D,U})
    AlphaColorValue{C,T}(convert(C, c.c), c.alpha)
end
for CV in CVAlpha
    @eval begin
        function convert{D,T}(AC::Type{$CV}, c::AbstractAlphaColorValue{D,T})
            AlphaColorValue(convert(colortype($CV), c.c), c.alpha)
        end
        convert(AC::Type{$CV}, c::ColorValue) = AlphaColorValue(convert(colortype($CV), c))
    end
end

convert{C,T}(::Type{AlphaColorValue{C,T}}, c::ColorValue) =
    AlphaColorValue{C,T}(convert(C, c), one(T))
convert{C<:ColorValue,D,T}(::Type{C}, c::AlphaColorValue{D, T}) = convert(C, c.c)

convert(::Type{ARGB32}, c::ARGB32) = c
convert{CV<:AbstractRGB{Ufixed8}}(::Type{ARGB32}, c::AbstractAlphaColorValue{CV,Ufixed8}) =
    ARGB32(c.c.r, c.c.g, c.c.b, c.alpha)
convert(::Type{ARGB32}, c::AbstractAlphaColorValue) =
    ARGB32(convert(RGB24, c.c).color | round(Uint32, 255*c.alpha)<<24)
convert(::Type{ARGB32}, c::ColorValue) = ARGB32(convert(RGB24, c).color | 0xff000000)
convert(::Type{ARGB32}, val::Uint32) = ARGB32(val)

function convert(::Type{RGBA{Ufixed8}}, ac::ARGB32)
    RGBA{Ufixed8}(Ufixed8(ac.color&0x00ff0000>>>16,0),
                  Ufixed8(ac.color&0x0000ff00>>>8,0),
                  Ufixed8(ac.color&0x000000ff,0),
                  Ufixed8(ac.color>>>24,0))
end
function convert{T}(::Type{RGBA{T}}, c::ARGB32)
    RGBA{T}((c.color&0x00ff0000>>>16)/255,
            (c.color&0x0000ff00>>>8)/255,
            (c.color&0x000000ff)/255,
            (c.color>>>24)/255)
end
function convert{C,T}(::Type{AlphaColorValue{C,T}}, c::ARGB32)
    convert(AlphaColorValue{C,T}, convert(RGBA{Ufixed8}, c))
end

### Equality
==(c1::AbstractRGB, c2::AbstractRGB) = c1.r == c2.r && c1.g == c2.g && c1.b == c2.b


@testset "geometricflux.jl" begin

    @testset "GraphConv" begin
        gnn = SeaPearl.GraphConv(3 => 3)

        graphs = Matrix.(adjacency_matrix.([random_regular_graph(10, 4) for i = 1:4]))
        nodeFeatures = [rand(Float32, 3, 10) for i = 1:4]
        featuredGraphs = [SeaPearl.FeaturedGraph(g; nf=nf) for (g, nf) in zip(graphs, nodeFeatures)]
        batchedFG = SeaPearl.BatchedFeaturedGraph{Float32}(featuredGraphs)

        classicOutput = (cat(map(fg->fg.nf,gnn.(featuredGraphs))...; dims=3))
        batchedOutput = gnn(batchedFG).nf

        @test all(abs.(classicOutput .- batchedOutput) .< 1e-5)

        grad = gradient(batchedFG) do inp
            gnn(inp).nf[1, 3, 2]
        end
        grad = grad[1]

        @test all(grad.nf[:, :, [1, 3, 4]] .== 0)
        @test !all(grad.nf[:, :, 2] .== 0)
        @test all(grad.graph[:, :, [1, 3, 4]] .== 0)
        @test !all(grad.graph[:, :, 2] .== 0)
    end
end


# Implemment the mutable arithmetics interface for UnitAffExpr

_MA.mutability(::Type{UnitAffExpr}) = _MA.IsMutable()

#  Minimal support of operate!! to allow @rewrite macro to work on linear constraints

UnitAffOrVar = Union{UnitVariableRef, UnitAffExpr}
AddSub = Union{typeof(_MA.add_mul), typeof(_MA.sub_mul)}
NumQuant = Union{Number, Quantity}

function Base.convert(::Type{UnitAffExpr}, uv::UnitVariableRef)
    return UnitAffExpr(convert(AffExpr, uv.vref), uv.u)
end

function update_expression!(ua::UnitAffExpr, a::NumQuant, x::UnitVariableRef)
    aval = ustrip(uconvert(ua.u, a * Quantity(1, x.u)))
    return UnitAffExpr(add_to_expression!(ua.expr, aval, x.vref), ua.u)
end

function update_expression!(ua::UnitAffExpr, a::Quantity, x::VariableRef)
    aval = ustrip(uconvert(ua.u, a))
    return UnitAffExpr(add_to_expression!(ua.expr, aval, x), ua.u)
end

function update_expression!(ua::UnitAffExpr, a::Quantity)
    aval = ustrip(uconvert(ua.u, a))
    return UnitAffExpr(add_to_expression!(ua.expr, aval), ua.u)
end

function update_expression!(ua::UnitAffExpr, x::UnitAffExpr)
    factor = ustrip(uconvert(ua.u, Quantity(1, x.u)))
    return UnitAffExpr(add_to_expression!(ua.expr, factor, x.expr), ua.u)
end

function update_expression!(ua::UnitAffExpr, a::NumQuant, b::NumQuant)
    aval = ustrip(uconvert(ua.u, a * b))
    return UnitAffExpr(add_to_expression!(ua.expr, aval), ua.u)
end

function create_expression(t::AddSub, z::typeof(_MA.Zero()), a::NumQuant, x::UnitVariableRef)
    val = a * Quantity(1, x.u)
    return UnitAffExpr(_MA.operate!!(t, z, ustrip(val), x.vref), Unitful.unit(val))
end

function create_expression(t::AddSub, z::typeof(_MA.Zero()), a::Quantity, x::VariableRef)
    return UnitAffExpr(_MA.operate!!(t, z, ustrip(a), x), Unitful.unit(a))
end

function create_expression(t::AddSub, z::typeof(_MA.Zero()), a::Quantity)
    return UnitAffExpr(_MA.operate!!(t, z, ustrip(a)), Unitful.unit(a))
end

# Two arguments
_MA.operate!!(::typeof(_MA.add_mul), uav::UnitAffOrVar, x::UnitVariableRef) = update_expression!(convert(UnitAffExpr, uav), 1, x)
_MA.operate!!(::typeof(_MA.sub_mul), uav::UnitAffOrVar, x::UnitVariableRef) = update_expression!(convert(UnitAffExpr, uav), -1, x)

_MA.operate!!(::typeof(_MA.add_mul), uav::UnitAffOrVar, x::UnitAffExpr) = update_expression!(convert(UnitAffExpr, uav), x)
_MA.operate!!(::typeof(_MA.sub_mul), uav::UnitAffOrVar, x::UnitAffExpr) = update_expression!(convert(UnitAffExpr, uav), -1, x)

_MA.operate!!(::typeof(_MA.add_mul), uav::UnitAffOrVar, a::Quantity) = update_expression!(convert(UnitAffExpr, uav), a)
_MA.operate!!(::typeof(_MA.sub_mul), uav::UnitAffOrVar, a::Quantity) = update_expression!(convert(UnitAffExpr, uav), -a)

# Three arguments
_MA.operate!!(::typeof(_MA.add_mul), uav::UnitAffOrVar, a::NumQuant, x::UnitVariableRef) = update_expression!(convert(UnitAffExpr, uav), a, x)
_MA.operate!!(::typeof(_MA.sub_mul), uav::UnitAffOrVar, a::NumQuant, x::UnitVariableRef) = update_expression!(convert(UnitAffExpr, uav), -a, x)

_MA.operate!!(::typeof(_MA.add_mul), uav::UnitAffOrVar, a::Quantity, x::VariableRef) = update_expression!(convert(UnitAffExpr, uav), a, x)
_MA.operate!!(::typeof(_MA.sub_mul), uav::UnitAffOrVar, a::Quantity, x::VariableRef) = update_expression!(convert(UnitAffExpr, uav), -a, x)

_MA.operate!!(::typeof(_MA.add_mul), uav::UnitAffOrVar, a::NumQuant, b::NumQuant) = update_expression!(convert(UnitAffExpr, uav), a, b)
_MA.operate!!(::typeof(_MA.sub_mul), uav::UnitAffOrVar, a::NumQuant, b::NumQuant) = update_expression!(convert(UnitAffExpr, uav), -a, b)

_MA.operate!!(::typeof(_MA.add_mul), uav::UnitAffOrVar, a::Number, b::Unitful.Units) = update_expression!(convert(UnitAffExpr, uav), a * Quantity(1.0, b))
_MA.operate!!(::typeof(_MA.sub_mul), uav::UnitAffOrVar, a::Number, b::Unitful.Units) = update_expression!(convert(UnitAffExpr, uav), - a * Quantity(1.0, b))

# Multiple arguments
function _MA.operate!!(t::AddSub, uav::UnitAffOrVar, x, y, z, other_args...)
    args = [x, y, z, other_args...]
    n = length(args)
    varidx = findall(t -> (typeof(t) <: UnitAffExpr) || (typeof(t) <: UnitVariableRef) || (typeof(t) <: VariableRef), args)
    var = args[varidx[1]]
    val = prod(args[i] for i in setdiff(1:n, varidx))
    return _MA.operate!!(t, uav, val, var)
end

# Zero handling
_MA.operate!!(t::AddSub, z::_MA.Zero, a::Quantity) = create_expression(t, z, a)
_MA.operate!!(t::AddSub, z::_MA.Zero, x::UnitVariableRef) = create_expression(t, z, 1, x)
_MA.operate!!(t::AddSub, z::_MA.Zero, ua::UnitAffExpr) = ua
_MA.operate!!(t::AddSub, z::_MA.Zero, a::NumQuant, x::UnitVariableRef) = create_expression(t, z, a, x)
_MA.operate!!(t::AddSub, z::_MA.Zero, x::UnitVariableRef, a::NumQuant) = create_expression(t, z, a, x)
_MA.operate!!(t::AddSub, z::_MA.Zero, a::Quantity, x::VariableRef) = create_expression(t, z, a, x)
_MA.operate!!(t::AddSub, z::_MA.Zero, x::VariableRef, a::Quantity) = create_expression(t, z, a, x)

function _MA.operate!!(t::AddSub, zero::_MA.Zero, x, y, z, other_args...)
    args = [x, y, z, other_args...]
    n = length(args)
    varidx = findall(t -> (typeof(t) <: UnitAffExpr) || (typeof(t) <: UnitVariableRef) || (typeof(t) <: VariableRef), args)
    var = args[varidx[1]]
    val = prod(args[i] for i in setdiff(1:n, varidx))
    return _MA.operate!!(t, zero, val, var)
end


using CrystalNets
using Test, Random
using PeriodicGraphs
using StaticArrays
import Base.Threads

const cifs = joinpath(@__DIR__, "cif")
const crystalnetsdir = normpath(@__DIR__, "..")

@testset "Archive" begin
    Threads.@threads for (genome, id) in collect(CrystalNets.CRYSTAL_NETS_ARCHIVE)
        (id == "sxt" || id == "llw-z") && continue # special case for these known unstable nets
        @test reckognize_topology(topological_genome(PeriodicGraph(genome))) == id
    end
end

@testset "Module" begin
    Threads.@threads for target in
        ["pcu", "afy, AFY", "apc, APC", "bam", "bcf", "cdp", "cnd", "ecb", "fiv",
         "ftd", "ftj", "ins", "kgt", "mot", "moz", "muh", "pbz", "qom", "sig",
         "sma", "sod-f", "sod-h", "utj", "utp"]
       @info "Testing $target"
       graph = PeriodicGraph(CrystalNets.REVERSE_CRYSTAL_NETS_ARCHIVE[target])
       n = PeriodicGraphs.nv(graph)
       for k in 1:50
           r = randperm(n)
           offsets = [SVector{3,Int}([rand(-3:3) for _ in 1:3]) for _ in 1:n]
           graph = swap_axes!(offset_representatives!(graph[r], offsets), randperm(3))
           @test reckognize_topology(topological_genome(graph)) == target
       end
   end

   @test_broken reckognize_topology(topological_genome(PeriodicGraph(CrystalNets.REVERSE_CRYSTAL_NETS_ARCHIVE["sxt"]))) == "sxt"
   @test_broken reckognize_topology(topological_genome(PeriodicGraph(CrystalNets.REVERSE_CRYSTAL_NETS_ARCHIVE["llw-z"]))) == "llw-z"

   @info """The following warnings about guessing bonds are expected."""
   @test reckognize_topology(topological_genome(CrystalNet(parse_chemfile(joinpath(cifs, "Moganite.cif"))))) == "mog"
end

@testset "Executable" begin
    safestdout = deepcopy(stdout)
    safeARGS = deepcopy(ARGS)
    piperead, _ = redirect_stdout()
    empty!(ARGS)

    push!(ARGS, "-g", "3   1 2  0 0 0   1 2  0 0 1   1 2  0 1 0   1 2  1 0 0")
    @test CrystalNets.julia_main() == 0
    @test readuntil(piperead, '\n') == "dia"

    empty!(ARGS)
    path = joinpath(cifs, "ABW.cif")
    push!(ARGS, path)
    @test CrystalNets.julia_main() == 0
    @test readuntil(piperead, '\n') == "sra, ABW"

    empty!(ARGS)
    path = joinpath(cifs, "ABW.cif")
    push!(ARGS, "-a", CrystalNets.arc_location*"rcsr.arc", path)
    @test CrystalNets.julia_main() == 0
    @test readuntil(piperead, '\n') == "sra"

    empty!(ARGS)
    path = joinpath(cifs, "RRO.cif")
    push!(ARGS, path)
    @test CrystalNets.julia_main() == 0
    @test readuntil(piperead, '\n') == "RRO"

    empty!(ARGS)
    path = joinpath(cifs, "RRO.cif")
    push!(ARGS, "-a", CrystalNets.arc_location*"rcsr.arc", path)
    @test CrystalNets.julia_main() == 1
    @test readuntil(piperead, '\n') == "UNKNOWN"

    empty!(ARGS)
    path = joinpath(cifs, "HKUST-1.cif")
    push!(ARGS, "-c", "mof", path)
    @test CrystalNets.julia_main() == 0
    @test readuntil(piperead, '\n') == "tbo"

    empty!(ARGS)
    path = joinpath(cifs, "Diamond.cif")
    push!(ARGS, "-c", "atom", path)
    @test CrystalNets.julia_main() == 0
    @test readuntil(piperead, '\n') == "dia"

    empty!(ARGS)
    push!(ARGS, "--help")
    @test CrystalNets.julia_main() == 0
    @test startswith(readuntil(piperead, '\n'), "usage: CrystalNets")
    @test occursin("CRYSTAL_FILE", readuntil(piperead, '\n'))
    @test occursin("Form B", readuntil(piperead, '\n'))
    @test occursin("Form C", readuntil(piperead, '\n'))
    @test isempty(readuntil(piperead, '\n'))
    @test isempty(readuntil(piperead, '\n'))
    @test readuntil(piperead, '\n') == "Automatic reckognition of crystal net topologies."

    redirect_stdout(safestdout)
    empty!(ARGS)
    append!(ARGS, safeARGS)
end


using Oceananigans.Architectures: device_event
using Oceananigans.Fields: location
using Oceananigans.TimeSteppers: ab2_step_field!
using Oceananigans.TurbulenceClosures: implicit_step!

using KernelAbstractions: NoneEvent

import Oceananigans.TimeSteppers: ab2_step!

#####
##### Step everything
#####

function ab2_step!(model::HydrostaticFreeSurfaceModel, Δt, χ)

    workgroup, worksize = work_layout(model.grid, :xyz)

    explicit_velocity_step_events = ab2_step_velocities!(model.velocities, model, Δt, χ)
    explicit_tracer_step_events = ab2_step_tracers!(model.tracers, model, Δt, χ)
    free_surface_event = ab2_step_free_surface!(model.free_surface, model, Δt, χ,
                                                MultiEvent(Tuple(explicit_velocity_step_events)))

    prognostic_field_events = MultiEvent(tuple(free_surface_event,
                                               explicit_velocity_step_events...,
                                               explicit_tracer_step_events...))

    wait(device(model.architecture), prognostic_field_events)

    return nothing
end

#####
##### Step velocities
#####

function ab2_step_velocities!(velocities, model, Δt, χ)

    barrier = device_event(model)

    # Launch velocity update kernels
    explicit_velocity_step_events = []

    for name in (:u, :v)
        Gⁿ = model.timestepper.Gⁿ[name]
        G⁻ = model.timestepper.G⁻[name]
        velocity_field = model.velocities[name]

        event = launch!(model.architecture, model.grid, :xyz,
                        ab2_step_field!, velocity_field, Δt, χ, Gⁿ, G⁻,
                        dependencies = device_event(model))

        push!(explicit_velocity_step_events, event)
    end

    for (i, name) in enumerate((:u, :v))
        velocity_field = model.velocities[name]

        # TODO: let next implicit solve depend on previous solve + explicit velocity step
        # Need to distinguish between solver events and tendency calculation events.
        # Note that BatchedTridiagonalSolver has a hard `wait`; this must be solved first.
        implicit_step!(velocity_field,
                       model.timestepper.implicit_solver,
                       model.clock,
                       Δt,
                       model.closure,
                       nothing,
                       model.diffusivities,
                       model.velocities,
                       model.tracers,
                       model.buoyancy,
                       dependencies = explicit_velocity_step_events[i])
    end

    return explicit_velocity_step_events
end

#####
##### Step velocities
#####

const EmptyNamedTuple = NamedTuple{(),Tuple{}}

ab2_step_tracers!(::EmptyNamedTuple, model, Δt, χ) = [NoneEvent()]

function ab2_step_tracers!(tracers, model, Δt, χ)

    # Tracer update kernels
    explicit_tracer_step_events = []

    for (tracer_index, tracer_name) in enumerate(propertynames(tracers))
        Gⁿ = model.timestepper.Gⁿ[tracer_name]
        G⁻ = model.timestepper.G⁻[tracer_name]
        tracer_field = tracers[tracer_name]

        event = launch!(model.architecture, model.grid, :xyz,
                        ab2_step_field!, tracer_field, Δt, χ, Gⁿ, G⁻,
                        dependencies = device_event(model))

        push!(explicit_tracer_step_events, event)
    end

    for (tracer_index, tracer_name) in enumerate(propertynames(tracers))
        tracer_field = tracers[tracer_name]

        implicit_step!(tracer_field,
                       model.timestepper.implicit_solver,
                       model.clock,
                       Δt,
                       model.closure,
                       tracer_index,
                       model.diffusivities,
                       model.velocities,
                       model.tracers,
                       model.buoyancy,
                       dependencies = explicit_tracer_step_events[tracer_index])
    end

    return explicit_tracer_step_events
end


module WildBootTests
export BootTestResult, wildboottest, teststat, stattype, p, padj, reps, repsfeas, nbootclust, dof, dof_r, plotpoints, peak, ci, dist, statnumer, statvar, auxweights

using LinearAlgebra, Random, Distributions, SortingAlgorithms, Printf

include("StrBootTest.jl")
include("utilities.jl")
include("estimators.jl")
include("init.jl")
include("WRE.jl")
include("nonWRE.jl")
include("plot-CI.jl")
include("interface.jl")

# top-level computation routine for OLS/arubin (and score BS on IV/2SLS); split off to reduce latency when just doing WRE
function boottestOLSARubin!(o::StrBootTest{T}) where T
	if !o.initialized
		Init!(o)
	elseif !o.null
		NoNullUpdate!(o)
		return
	end

	if o.Nw > 1  # if more than one weight group to save memory, make on every call to boottest(), not just once in Init!()
		Random.seed!(o.rng,o.seed)
		MakeWildWeights!(o, last(o.WeightGrp[1])-1, first=true)
	end

	MakeInterpolables!(o)  # make stuff that depends linearly on r, possibly by interpolating, for first weight group

	MakeNonWREStats!(o, 1)
	for w ∈ 2:o.Nw  # do group 1 first because it includes col 1, which is all that might need updating in constructing CI in WCU
		MakeWildWeights!(o, length(o.WeightGrp[w]), first=false)
		MakeNonWREStats!(o, w)
	end
	!o.bootstrapt && UpdateBootstrapcDenom!(o)

	o.BFeas = isnan(o.dist[1]) ? 0 : sum(.!(isnan.(o.dist) .| isinf.(o.dist))) - 1
	o.distCDR = zeros(T,0,0)
	nothing
end

# top-level computation routine for non-arubin WRE; split off to reduce latency when just doing other tests
function boottestWRE!(o::StrBootTest{T}) where T
	if !o.initialized
		Init!(o)
	elseif !o.null
		NoNullUpdate!(o)
		return
	end

	if o.Nw > 1  # if more than one weight group to save memory, make on every call to boottest(), not just once in Init!()
		Random.seed!(o.rng,o.seed)
		MakeWildWeights!(o, last(o.WeightGrp[1])-1, first=true)
	end

	PrepWRE!(o)

	MakeWREStats!(o, 1)
	for w ∈ 2:o.Nw  # do group 1 first because it includes col 1, which is all that might need updating in constructing CI in WCU
		MakeWildWeights!(o, length(o.WeightGrp[w]), first=false)
		MakeWREStats!(o, w)
	end
	!o.bootstrapt && UpdateBootstrapcDenom!(o)

	o.BFeas = isnan(o.dist[1]) ? 0 : sum(.!(isnan.(o.dist) .| isinf.(o.dist))) - 1
	o.distCDR = zeros(T,0,0)
	nothing
end

# if not imposing null and we have returned to boottest!(), then dof=1 or 2; we're plotting or finding CI, and only test stat, not distribution, changes with r
function NoNullUpdate!(o::StrBootTest{T} where T)
	if o.WREnonARubin
		o.numer[:,1] = o.R * o.DGP.Rpar * o.β̈s[1] - o.r
	elseif o.arubin
		EstimateARubin!(o.DGP, o, o.r)
		o.numer[:,1] = o.v_sd * @view o.DGP.β̈[o.kX₁+1:end,:]  # coefficients on excluded instruments in arubin OLS
	else
		o.numer[:,1] = o.v_sd * (o.R * (o.ml ? o.β̈ : iszero(o.κ) ? o.M.β̈ : o.M.Rpar * o.M.β̈) - o.r)  # Analytical Wald numerator; if imposing null then numer[:,1] already equals this. If not, then it's 0 before this
	end
	o.dist[1] = isone(o.dof) ? o.numer[1] / sqrtNaN(o.statDenom[1]) : o.numer[:,1]'invsym(o.statDenom)*o.numer[:,1]
	nothing
end

# compute bootstrap-c denominator from all bootstrap numerators
function UpdateBootstrapcDenom!(o::StrBootTest{T} where T)
	numer1 = o.numer[:,1]
	o.numersum = rowsum(o.numer) - numer1
	o.statDenom = (o.numer * o.numer' - numer1 * numer1' - o.numersum * o.numersum' / o.B) / o.B
	if o.sqrt
		o.dist .= o.numer ./ sqrtNaN.(o.statDenom)
	else
		negcolquadform!(o, o.dist, -invsym(o.statDenom), o.numer)  # to reduce latency by minimizing #=@tturbo=# instances, work with negative of colquadform in order to fuse code with colquadformminus!
	end
	nothing
end

include("precompile_WildBootTests.jl")  # https://timholy.github.io/SnoopCompile.jl/stable/snoopi_deep_parcel/#SnoopCompile.write
_precompile_()

end


using Distributions

@enum DogState begin
    ACTIVE = 1
    OCCUPIED
    SLEEPY
end # enum
Base.length( ::DogState ) = 1
Base.iterate( ds::DogState, state = 1 ) = nothing

Base.@kwdef struct Dog{F<:Number}
    state::DogState = ACTIVE
    sleepyness::UInt32 = 0
    position::Point2{F}
    velocity::Point2{F}
end # struct

function Dog( position, velocity )
    T = eltype(position)
    @assert T == eltype(velocity)
    Dog{T}(;
        position = convert(Point2{T}, position),
        velocity = convert(Point2{T}, velocity),
        )
end # function

#TODO: this is slow and allocating
function Dog( v::D; kwargs... ) where D<:Dog
    nt = NamedTuple()
    for property in propertynames(v)
        if property in keys(kwargs)
            nt = merge( nt, NamedTuple{(property,)}((kwargs[property],)) )
            continue
        end
        value = getproperty( v, property )
        # if value isa Vector
        #     nt = merge( nt, NamedTuple{(property,)}(Ref(copy(value))) )
        #     continue
        # end # if
        nt = merge( nt, NamedTuple{(property,)}((value,)) )
    end
    return D(;nt...)
end # function


export
    SWR_CH_MAX,
    SWR_FLAG_RESAMPLE,
    SwrDitherType,
    SWR_DITHER_NONE,
    SWR_DITHER_RECTANGULAR,
    SWR_DITHER_TRIANGULAR,
    SWR_DITHER_TRIANGULAR_HIGHPASS,
    SWR_DITHER_NS,
    SWR_DITHER_NS_LIPSHITZ,
    SWR_DITHER_NS_F_WEIGHTED,
    SWR_DITHER_NS_MODIFIED_E_WEIGHTED,
    SWR_DITHER_NS_IMPROVED_E_WEIGHTED,
    SWR_DITHER_NS_SHIBATA,
    SWR_DITHER_NS_LOW_SHIBATA,
    SWR_DITHER_NS_HIGH_SHIBATA,
    SWR_DITHER_NB,
    SwrEngine,
    SWR_ENGINE_SWR,
    SWR_ENGINE_SOXR,
    SWR_ENGINE_NB,
    SwrFilterType,
    SWR_FILTER_TYPE_CUBIC,
    SWR_FILTER_TYPE_BLACKMAN_NUTTALL,
    SWR_FILTER_TYPE_KAISER,
    SwrContext

const SWR_CH_MAX  =  32
const SWR_FLAG_RESAMPLE  =  1

# begin enum SwrDitherType
const SwrDitherType = UInt32
const SWR_DITHER_NONE  =  UInt32(0)
const SWR_DITHER_RECTANGULAR  =  UInt32(1)
const SWR_DITHER_TRIANGULAR  =  UInt32(2)
const SWR_DITHER_TRIANGULAR_HIGHPASS  =  UInt32(3)
const SWR_DITHER_NS  =  UInt32(64)
const SWR_DITHER_NS_LIPSHITZ  =  UInt32(65)
const SWR_DITHER_NS_F_WEIGHTED  =  UInt32(66)
const SWR_DITHER_NS_MODIFIED_E_WEIGHTED  =  UInt32(67)
const SWR_DITHER_NS_IMPROVED_E_WEIGHTED  =  UInt32(68)
const SWR_DITHER_NS_SHIBATA  =  UInt32(69)
const SWR_DITHER_NS_LOW_SHIBATA  =  UInt32(70)
const SWR_DITHER_NS_HIGH_SHIBATA  =  UInt32(71)
const SWR_DITHER_NB  =  UInt32(72)
# end enum SwrDitherType

# begin enum SwrEngine
const SwrEngine = UInt32
const SWR_ENGINE_SWR  =  UInt32(0)
const SWR_ENGINE_SOXR  =  UInt32(1)
const SWR_ENGINE_NB  =  UInt32(2)
# end enum SwrEngine

# begin enum SwrFilterType
const SwrFilterType = UInt32
const SWR_FILTER_TYPE_CUBIC  =  UInt32(0)
const SWR_FILTER_TYPE_BLACKMAN_NUTTALL  =  UInt32(1)
const SWR_FILTER_TYPE_KAISER  =  UInt32(2)
# end enum SwrFilterType

const SwrContext = Void


#=========
Requires
=========#

using GLM
using DataFrames
using StatsBase
using StatsModels
#import StatsBase: residuals

const INNERMOD = Union{GLM.GeneralizedLinearModel, GLM.LinearModel}
#const UNIONC = Union{Type{Nothing}, Type{LinearAlgebra.Cholesky}, Type{PositiveFactorizations.Positive}}
#============
General GLM methods
=============#
StatsModels.modelmatrix(m::StatsModels.DataFrameRegressionModel{T}) where T<:INNERMOD = m.mm.m
StatsModels.modelmatrix(m::T) where T<:INNERMOD = m.pp.X

residuals(m::StatsModels.DataFrameRegressionModel{T}) where T<:INNERMOD = m.model.rr.wrkresid
residuals(m::T) where T<:INNERMOD = m.rr.y .- m.rr.mu

modelweights(m::StatsModels.DataFrameRegressionModel{T}) where T<:INNERMOD = m.model.rr.wrkwt
modelweights(m::GLM.GeneralizedLinearModel) = m.rr.wrkwt
modelweights(m::GLM.LinearModel) = m.rr.wts

smplweights(m::StatsModels.DataFrameRegressionModel{T}) where T<:INNERMOD = m.model.rr.wts
smplweights(m::GLM.GeneralizedLinearModel) = m.rr.wts
smplweights(m::GLM.LinearModel) = m.rr.wts

choleskyfactor(m::StatsModels.DataFrameRegressionModel{T}) where T<:INNERMOD = m.model.pp.chol.UL
choleskyfactor(m::T) where T<:INNERMOD = m.pp.chol.UL

unweighted_nobs(m::StatsModels.DataFrameRegressionModel{T}) where T<:INNERMOD = size(modelmatrix(m), 1)
unweighted_nobs(m::T) where T<:INNERMOD = size(modelmatrix(m), 1)

unweighted_dof_residual(m::StatsModels.DataFrameRegressionModel{T}) where T<:INNERMOD = unweighted_nobs(m) - length(coef(m))
unweighted_dof_residual(m::T) where T<:INNERMOD = unweighted_nobs(m) - length(coef(m))

function installxuw!(cache, m::T) where T<:INNERMOD
    copyto!(cache.X, copy(modelmatrix(m)))
    copyto!(cache.u, copy(residuals(m)))
    if !isempty(m.rr.wts)
        broadcast!(*, cache.u, cache.u, smplweights(m))
    end
end

function esteq!(cache, m::StatsModels.RegressionModel, k::T) where T<:Union{HC, CRHC}
    broadcast!(*, cache.q, cache.X, cache.u)
    return cache.q
end

function esteq!(cache, m::RegressionModel, k::HAC)
    X = copy(modelmatrix(m))
    u = copy(residuals(m))
    if !isempty(smplweights(m))
        broadcast!(*, u, u, smplweights(m))
    end
    broadcast!(*, X, X, u)
end

pseudohessian(m::StatsModels.DataFrameRegressionModel{T}) where T<:INNERMOD = GLM.invchol(m.model.pp)
pseudohessian(m::T) where T<:INNERMOD = GLM.invchol(m.pp)

#==============
HAC GLM Methods
===============#

function HACCache(m::StatsModels.DataFrameRegressionModel{T}; prewhiten = false) where T<:INNERMOD
    HACCache(m.model, prewhiten = prewhiten)
end

function HACCache(m::T; prewhiten = false) where T<:INNERMOD
    HACCache(copy(modelmatrix(m)), prewhiten = prewhiten)
end

function StatsBase.vcov(m::StatsModels.DataFrameRegressionModel{T}, k::HAC, args...; kwargs...) where T<:INNERMOD
    set_bw_weights!(k, m)
    vcov(m.model, k, args...; kwargs...)
end

function StatsBase.vcov(m::T, k::K, returntype::Type{T1} = CovarianceMatrix, factortype::Type{T2} = Cholesky; kwargs...) where {T<:INNERMOD, K<:HAC, T1<:Union{CovarianceMatrix, Matrix}, T2<:Union{Nothing, Factorization}}
    cache = HACCache(m, prewhiten = k.prewhiten)
    vcov(m, k, cache, returntype, factortype; kwargs...)
end

function StatsBase.vcov(m::T, k::K, cache::HACCache, returntype::Type{T1} = CovarianceMatrix, factortype::Type{T2} = Cholesky; demean::Bool = false, dof_adjustment::Bool = true) where {T<:INNERMOD, K<:HAC, T1<:Union{CovarianceMatrix, Matrix}, T2<:Union{Nothing, Factorization}}
    mf = esteq!(cache, m, k)
    br = pseudohessian(m)
    set_bw_weights!(k, m)
    Ω = covariance(mf, k, cache, Matrix, demean = demean, scale = 1)
    V = br*Ω*br'
    rmul!(V, size(mf, 1)^2)
    scale = dof_adjustment ? unweighted_dof_residual(m) : unweighted_nobs(m)
    finalize(V, T2, T1, k, scale)
end

# setcholesky!(cache, cholesky::Type{Nothing}, V) = nothing

# function setcholesky!(cache, cholesky::Type{Cholesky}, V)
#     chol = LinearAlgebra.cholesky(Symmetric(V))
#     copyto!(cache.chol.UL.data, chol.UL.data)
#     copyto!(cache.chol.U.data, chol.U.data)
#     copyto!(cache.chol.L.data, chol.L.data)
# end

# function setcholesky!(cache, cholesky::Type{PositiveFactorizations.Positive}, V)
#     chol = LinearAlgebra.cholesky(Positive, V)
#     copyto!(cache.chol.UL.data, chol.UL.data)
#     copyto!(cache.chol.U.data, chol.U.data)
#     copyto!(cache.chol.L.data, chol.L.data)
# end

function set_bw_weights!(k::HAC, m::StatsModels.DataFrameRegressionModel{T}) where T<:INNERMOD
    if isempty(k.weights)
        for j in eachindex(coef(m))
            push!(k.weights, 1.0)
        end
    end
    if length(k.weights) == length(coef(m))
        "(Intercept)" ∈ coefnames(m) ? (k.weights .= 1.0; k.weights[1] = 0.0) : k.weights .= 1.0
    else
        error("Bandwidth weights have wrong dimension for the problem")
    end
end

function set_bw_weights!(k::HAC, m::T) where T<:INNERMOD
    if isempty(k.weights)
        for j in eachindex(coef(m))
            push!(k.weights, 1.0)
        end
    end
end

#==============
HC GLM Methods
===============#
function HCCache(m::StatsModels.DataFrameRegressionModel{T}) where T<:INNERMOD
    HCCache(m.model)
end

function HCCache(m::T) where {T<:INNERMOD}
    n, p = unweighted_nobs(m), length(coef(m))
    HCCache(similar(Array{eltype(modelmatrix(m)), 2}(undef, n, p)))
end

function StatsBase.vcov(m::StatsModels.DataFrameRegressionModel{T},
                        k::K, returntype::Type{T1} = CovarianceMatrix,
                        factortype::Type{T2} = Cholesky;
                        kwargs...) where {T<:INNERMOD,
                                          K<:HC,
                                          T1<:Union{CovarianceMatrix, Matrix},
                                          T2<:Union{Nothing, Factorization}}
    vcov(m.model, k, HCCache(m.model), returntype, factortype; kwargs...)
end

function StatsBase.vcov(m::T,
    k::K, returntype::Type{T1} = CovarianceMatrix,
    factortype::Type{T2} = Cholesky;
    kwargs...) where {T<:INNERMOD,
                      K<:HC,
                      T1<:Union{CovarianceMatrix, Matrix},
                      T2<:Union{Nothing, Factorization}}
    vcov(m, k, HCCache(m), returntype, factortype; kwargs...)
end

# function StatsBase.vcov(m::T,
#                         k::K,
#                         cache::HCCache,
#                         returntype::Type{T1} = CovarianceMatrix,
#                         factortype::Type{T2} = Cholesky; kwargs...) where {T<:INNERMOD,
#                                                                            K<:HC,
#                                                                            T1<:Union{CovarianceMatrix, Matrix},
#                                                                            T2<:Union{Nothing, Factorization}}

#     vcov(m, k, HCCache(m), returntype, factortype; kwargs...)
# end

function StatsBase.vcov(m::T,
                        k::K,
                        cache::HCCache,
                        returntype::Type{T1} = CovarianceMatrix,
                        factortype::Type{T2} = Cholesky;
                        demean::Bool = false) where {T<:INNERMOD,
                                                            K<:HC,
                                                            T1<:Union{CovarianceMatrix, Matrix},
                                                            T2<:Union{Nothing, Factorization}}
    CovarianceMatrices.installxuw!(cache, m)
    mf = CovarianceMatrices.esteq!(cache, m, k)
    br = CovarianceMatrices.pseudohessian(m)
    CovarianceMatrices.adjfactor!(cache, m, k)
    mul!(cache.x, mf', broadcast!(*, cache.X, mf, cache.η))
    V = br*cache.x*br'
    finalize(V, T2, T1, k, 1.0)
end

## -----
## DataFramesRegressionModel/AbstractGLM methods
## -----

function hatmatrix(cache, m::T) where T<:INNERMOD
      z = cache.X  ## THIS ASSUME THAT X IS WEIGHTED BY SQRT(W)
      if !isempty(modelweights(m))
          z .= z.*sqrt.(modelweights(m))
      end
      cf = choleskyfactor(m)::UpperTriangular
      rdiv!(z, cf)
      sum!(cache.v, z.^2)
 end

  adjfactor!(cache, m::StatsModels.RegressionModel, k::HC0) = cache.η .= one(eltype(cache.u))
  adjfactor!(cache, m::StatsModels.RegressionModel, k::HC1) = cache.η .= unweighted_nobs(m)./unweighted_dof_residual(m)
  adjfactor!(cache, m::StatsModels.RegressionModel, k::HC2) = cache.η .= one(eltype(cache.u))./(one(eltype(cache.u)).-hatmatrix(cache, m))
  adjfactor!(cache, m::StatsModels.RegressionModel, k::HC3) = cache.η .= one(eltype(cache.u))./(one(eltype(cache.u)).-hatmatrix(cache, m)).^2

  function adjfactor!(cache, m::StatsModels.RegressionModel, k::HC4)
      n, p = size(cache.X)
      tone = one(eltype(cache.u))
      h = hatmatrix(cache, m)
      @inbounds for j in eachindex(h)
          delta = min(4, n*h[j]/p)
          cache.η[j] = tone/(tone-h[j])^delta
      end
      cache.η
  end

  function adjfactor!(cache, m::StatsModels.RegressionModel, k::HC4m)
      n, p = size(cache.X)
      tone = one(eltype(cache.u))
      h = hatmatrix(cache, m)
      @inbounds for j in eachindex(h)
          delta = min(tone, n*h[j]/p) + min(1.5, n*h[j]/p)
          cache.η[j] = tone/(tone-h[j])^delta
      end
      cache.η
  end

  function adjfactor!(cache, m::StatsModels.RegressionModel, k::HC5)
      n, p = size(cache.X)
      tone = one(eltype(cache.u))
      h     = hatmatrix(cache, m)
      mx    = max(n*0.7*maximum(h)/p, 4)
      @inbounds for j in eachindex(h)
          alpha =  min(n*h[j]/p, mx)
          cache.η[j] = sqrt(tone/(tone-h[j])^alpha)
      end
      cache.η
  end

#==========
Cluster GLM
=========#

function CRHCCache(m::StatsModels.DataFrameRegressionModel{T}, cl::AbstractArray{F}) where {T<:INNERMOD, F}
    CRHCCache(m.model, cl, returntype=returntype)
end

function CRHCCache(m::T, cl::AbstractVector{F}) where {T<:INNERMOD, F}
    n, p = unweighted_nobs(m), length(coef(m))
    CRHCCache(similar(modelmatrix(m)), cl)
end

function StatsBase.vcov(m::StatsModels.DataFrameRegressionModel{T},
                        k::K,
                        returntype::Type{T1} = CovarianceMatrix,
                        factortype::Type{T2} = Cholesky; kwargs...) where {T<:INNERMOD,
                                                                           K<:CRHC,
                                                                           T1<:Union{CovarianceMatrix, Matrix},
                                                                           T2<:Union{Nothing, Factorization}}
    vcov(m.model, k, CRHCCache(m.model, k.cl), returntype, factortype; kwargs...)
end

function StatsBase.vcov(m::T,
                        k::K,
                        returntype::Type{T1} = CovarianceMatrix,
                        factortype::Type{T2} = Cholesky;
                        kwargs...) where {T<:INNERMOD,
                                          K<:CRHC,
                                          T1<:Union{CovarianceMatrix, Matrix},
                                          T2<:Union{Nothing, Factorization}}
    vcov(m, k, CRHCCache(m, k.cl), returntype, factortype; kwargs...)
end

function StatsBase.vcov(m::T,
                        k::CRHC,
                        cache::CRHCCache,
                        returntype::Type{T1} = CovarianceMatrix,
                        factortype::Type{T2} = Cholesky;
                        sorted::Bool = false) where {T<:INNERMOD,
                                                     K<:HC,
                                                     T1<:Union{CovarianceMatrix, Matrix},
                                                     T2<:Union{Nothing, Factorization}}
    B = CovarianceMatrices.pseudohessian(m)
    CovarianceMatrices.installsortedxuw!(cache, m, k, Val{sorted})
    bstarts = (searchsorted(cache.clus, j[2]) for j in enumerate(unique(cache.clus)))
    CovarianceMatrices.adjresid!(k, cache, B, bstarts)
    CovarianceMatrices.esteq!(cache, m, k)
    CovarianceMatrices.clusterize!(cache, bstarts)
    V = dof_adjustment(cache, k, bstarts).*(B*cache.M*B)
    finalize(V, T2, T1, k, 1.0)
end

StatsBase.stderror(m::StatsModels.DataFrameRegressionModel{T}, k::RobustVariance; kwargs...) where T<:INNERMOD = sqrt.(diag(vcov(m, k, Matrix; kwargs...)))


"""
A `Worker` object holds the necessary local states for a learner worker:

 - dset_tr: the local partition of the training dataset
 - dset_tt: the local partition of the test dataset
 - hp: hyper parameters specific to each algorithm (e.g. batch size)
 - lp: local parameters specific to each algorithm (e.g. algorithms in the
   dual might store dual variables locally)
"""
mutable struct Worker{HPType, LPType}
  dset_tr   :: Dataset # training set
  dset_tt   :: Dataset # test / validation set

  hp        :: HPType  # hyper parameters
  lp        :: LPType  # worker local parameters

  Worker{HPType, LPType}() where {HPType, LPType} = new()
end

"""
    spanw_workers(ctor, args...)

Create `Worker` objects on each worker process. Return a list
of (pid, worker future ref) tuples.
"""
function spawn_workers(ctor :: Function, args...)
  return map(1:nworkers()) do wid
    pid = wid + 1  # head node is 1, all other nodes are workers
    (pid, remotecall(ctor, pid, wid, nworkers(), args...))
  end
end

"""
    invoke_on_workers(f, workers, args...)

Invoke `f` on each workers in parallel, with `f(worker_ref, args...)`.
This function blocks until all workers finishes, and return a list
of results from each worker.
"""
function invoke_on_workers(f :: Function, workers, args...)
  rets = Array{Any}(length(workers))
  @sync begin
    for (i, (pid, worker_ref)) in enumerate(workers)
      @async begin
        rets[i] = remotecall_fetch(f, pid, worker_ref, args...)
      end
    end
  end

  # propagate remote exceptions
  for obj in rets
    if obj isa RemoteException
      throw(obj)
    end
  end
  return rets
end

"""
    accumulate_metric(workers, Type{MetricType}, w)

Collect and accumulate metrics from workers. This function assumes linear
predictors and assume only one metric object is needed. One can define this
in more general ways for general predictors and multiple metrics. Also it
might be more efficient to collect the metrics during the forward stage of
the training to avoid duplicated computation.
"""
function accumulate_metric{T <: Metric}(workers, :: Type{T}, w :: Vector{Float64})
  function eval_metric{T <: Metric}(worker_ref :: Future, :: Type{T}, w :: Vector{Float64})
    worker = fetch(worker_ref)
    return map((worker.dset_tr, worker.dset_tt)) do dset
      pred = dset.data' * w
      evaluate(T, pred, dset.labels)
    end
  end
  # rets is [(m_tr1, m_tt1), (m_tr2, m_tt2), ...]
  rets = invoke_on_workers(eval_metric, workers, T, w)
  return merge([x[1] for x in rets]), merge([x[2] for x in rets])
end


# Greetings! from path ~/.julia/config/startup.jl
############################################
## 1.Set parameters
############################################
# process and threads parameters
cores = 1                                                           # Int64; total number of processes. cores number should be less than "N_SR_pairs" number.
threads = 1                                                         # Int64; total number of threads. threads number should be less than N piece of data when "stack_all=false". please set to be "1" when "stack_all=true"
flag = true                                                         # true of false; print computing time for debugging purpose.
precompile = false                                                  # true or false; precompile or not.

# absolute path parameters
CCFDIR = expanduser("./CORR_fft")                              # dir where CC data are stored.
STACKDIR = expanduser("./STACK_d")                               # dir to store STACK data.

# select time-chunck folder for stack
select_time = false                                                  # true or false; whether specify time-chunck folder path.
TIMETXT = expanduser("./time_stack.txt")                            # time-chunck folder path.

# select station for stack
select_sta = "sr"                                                   # "sta" or "sr"; "sta" cross-correlate only works for selected net-sta pairs base on STATXT; "sr" cross-correlate only works for specified source and receiver based on SRTXT.
STATXT = expanduser("./sta_stack.txt")                              # station infomation including: network-name and station-name.
SRTXT = expanduser("./sr_stack.txt")                                # source and receiver information including: source-1 receiver-1 r-2 r-3...\nsource-2 r-1 r-2 r-3...
channel_regular = "*.CHN.*.CHN.*"                                   # channel(regular expression) of CCF-data; we will use 'glob(comp,CCFDIR)' to select data in 'CCFDIR' path.
comp = "NN"                                                         # component name corresponds to "channel_regular"; different from NoiseCC.CC, we will use 'load_corr(file,comp)' to read CCFData in 'CCFDIR' path, and "comp" is the group name of jld2 file(HDF5 in julia).

# stack time interval
stack_all = true                                                    # true or false; "true" for stack corr-data over all the selected time chunck.
time_interval = 1 * 24 * 60 * 60                                    # Int64 (sec); stack every "time_interval" time, and it works when stack_all = false || 中文: 每间隔多少秒叠加一次，要根据corr-data文件之间的时间间隔决定，例如corr-data文件是6个小时一个，那么time_interval=12*60*60，表示每12个小时叠加一次，即需要2个corr文件叠加在一起。该参数工作时，先搜索指定台站对的全部时间的文件，找到其中最小的开始时间与最大的结束时间，利用range函数生成间隔为""的叠加时窗，叠加每个时窗内的corr-data文件，叠加后的文件名称为该时窗的起始与终止时刻。
min_stack_chunck = 1                                                # Int64; each stack window needs at least "min_stack_chunck" to be able to stack, and it works for both stack_all = true and stack_all = false. || 中文:每个时间窗内叠加，至少需要 "min_stack_chunck" 个corr-data文件才会叠加,我们在读取corr-data文件后，会进行质量评估，剔除异常文件，导致corr-data文件数量减少。

# stack method
substack_method = "pws"                                             # "mean", "pws", and "robust"; substack method corresponding with NoiseCC.CC process, but can only choose one type.
stack_method = "mean"                                               # "mean", "pws", and "robust"; stack method, but can only choose one type.
auto_stack = true                                                   # true or false; perform auto-correlation or not.
cross_stack = true                                                  # true or false; perform cross-correlation or not.

# filter
select_filter = true                                                # true or false; whether filter before stacking.
freqmin = 0.5                                                       # Float64;
freqmax = 20.0                                                      # Float64;
corners = 4                                                         # Int64; number of corners in Butterworth filter.
zerophase = false                                                   # true or false; if true, apply filter twice for no phase shifting.

# criteria for corr-data selection
select_corr = true                                                  # true or false; whether to perform quality selecting based on median value of amp of corr-data.
median_high = 10.0                                                  # Float64; max median value (default value).
median_low = 0.0                                                    # Float64; min median value (default value).

############################################
# 2.run NoiseCC.FFT_CC
############################################
# a.using and import
# using Distributed
# NoiseCC_path = "/Users/yf/3.Project/2-1.Velocity_Change/SeisNoise_Code_YF/code_pmap_thread/"
# push!(LOAD_PATH, NoiseCC_path)
# @everywhere using NoiseCC # or @everywhere import NoiseCC

# b.include
using Distributed
addprocs(cores, topology = :master_worker, exeflags = "-t $(threads)")
@everywhere include(
    "/Users/yf/3.Project/2-1.Velocity_Change/SeisNoise_Code_YF/code_pmap_thread/NoiseCC.jl",
)

NoiseCC.STACK(
    cores,
    threads,
    flag,
    precompile,
    CCFDIR,
    STACKDIR,
    select_time,
    TIMETXT,
    select_sta,
    STATXT,
    SRTXT,
    channel_regular,
    comp,
    stack_all,
    time_interval,
    min_stack_chunck,
    substack_method,
    stack_method,
    auto_stack,
    cross_stack,
    select_filter,
    freqmin,
    freqmax,
    corners,
    zerophase,
    select_corr,
    median_high,
    median_low,
)


using NNlib: PoolDims

# descriptor

mutable struct PoolDesc
    ptr::cudnnPoolingDescriptor_t
end

unsafe_free!(pd::PoolDesc)=cudnnDestroyPoolingDescriptor(pd.ptr)

Base.unsafe_convert(::Type{cudnnPoolingDescriptor_t}, pd::PoolDesc)=pd.ptr

function PoolDesc(nd, window, padding, stride, mode, maxpoolingNanOpt=CUDNN_NOT_PROPAGATE_NAN)
    pd = Ref{cudnnPoolingDescriptor_t}()
    cudnnCreatePoolingDescriptor(pd)
    cudnnSetPoolingNdDescriptor(pd[],cudnnPoolingMode_t(mode),maxpoolingNanOpt,nd,pdsize(window,nd),pdsize(padding,nd),pdsize(stride,nd))
    this = PoolDesc(pd[])
    finalizer(unsafe_free!, this)
    return this
end

function PoolDesc(pdims::PoolDims, mode, maxpoolingNanOpt=CUDNN_NOT_PROPAGATE_NAN)
    pd = NNlib.padding(pdims)
    if !all(pd[1:2:end] .== pd[2:2:end])
        @warn("CuDNN does not support asymmetric padding; defaulting to symmetric choice")
    end
    return PoolDesc(NNlib.spatial_dims(pdims), NNlib.kernel_size(pdims), pd[1:2:end],
                    NNlib.stride(pdims), mode, maxpoolingNanOpt)
end

# wrappers

function cudnnPoolingForward(y::CuArray{T,N}, x::CuArray{T,N}, pdims::PoolDims;
                             alpha=1, mode=0) where {T,N}
    beta = 0
    cudnnPoolingForward(handle(), PoolDesc(pdims, mode), Ref(T(alpha)), TensorDesc(x), x, Ref(T(beta)), TensorDesc(y), y)
    return y
end

function cudnnPoolingBackward(dx::CuArray{T,N}, dy::CuArray{T,N}, x::CuArray{T,N}, y::CuArray{T,N},
                              pdims::PoolDims; alpha=1, mode=0) where {T,N}
    if alpha!=1 && mode==0; error("Gradient of pool(alpha!=1,mode=0) broken in CUDNN"); end
    beta = 0
    cudnnPoolingBackward(handle(),
        PoolDesc(pdims, mode), Ref(T(alpha)), TensorDesc(y), y,
        TensorDesc(dy), dy, TensorDesc(x), x, Ref(T(beta)), TensorDesc(dx), dx
    )
    return dx
end


"""
    britter_puff_factory(scenario)

Generates a Britter-McQuaid dispersion model on the given scenario and returns a
function giving the concentration of the form
c(x, y, z, t)

"""
function britter_puff_factory(scenario)
    # TODO
    error("Not Implemented")
end