"""
    setline(linewidth)

Set the line width and multiply it with the current multiplier
which is i.e. set by [`appear`](@ref) and [`disappear`](@ref).

Normal behavior without any animation is the same as `Luxor.setline`.

# Example
```
setline(10)
line(O, Point(10, 10))
```

# Arguments:
- `linewidth`: the line width in pixel
"""
function setline(linewidth)
    cs = get_current_setting()
    cs.line_width = linewidth
    current_line_width = cs.line_width * cs.mul_line_width
    Luxor.setline(current_line_width)
end

"""
    setopacity(opacity)

Set the opacity and multiply it with the current multiplier
which is i.e. set by [`appear`](@ref) and [`disappear`](@ref).

Normal behavior without any animation is the same as `Luxor.setopacity`.

# Example
```
setopacity(0.5)
circle(O, 20, :fill)
```

# Arguments:
- `opacity`: the opacity between 0.0 and 1.0
"""
function setopacity(opacity)
    cs = get_current_setting()
    cs.opacity = opacity
    current_opacity = cs.opacity * cs.mul_opacity
    Luxor.setopacity(current_opacity)
end

"""
    fontsize(fsize)

Same as `Luxor.fontsize`: Sets the current font size.

# Example
```
fontsize(12)
text("Hello World!")
```

# Arguments:
- `fsize`: the new font size
"""
function fontsize(fsize)
    cs = get_current_setting()
    cs.fontsize = fsize
    Luxor.fontsize(fsize)
end

"""
    get_fontsize(fsize)

Same as `Luxor.get_fontsize` but works with every version of Luxor that is supported by Javis.

# Example
```
fontsize(12)
fsize = get_fontsize()
text("Hello World! \$fsize")
```

# Returns
- `Float64`: the current font size
"""
function get_fontsize()
    cs = get_current_setting()
    return cs.fontsize
end

"""
    scale(scl)

Set the scale and multiply it with the current multiplier
which is i.e. set by [`appear`](@ref) and [`disappear`](@ref).

Normal behavior without any animation is the same as `Luxor.scale`.

# Example
```
scale(0.5)
circle(O, 20, :fill) # the radius would be 10 because of the scaling
```

# Arguments:
- `scl`: the new default scale
"""
function scale(scl::Number)
    scale(scl, scl)
end

"""
    scale(scl_x, scl_y)

Same as [`scale`](@ref) but the x scale and y scale can be changed independently.

# Arguments:
- `scl_x`: scale in x direction
- `scl_y`: scale in y direction
"""
function scale(scl_x, scl_y)
    cs = get_current_setting()
    cs.desired_scale = (scl_x, scl_y)
    current_scale = cs.desired_scale .* cs.mul_scale
    Luxor.scale(current_scale...)
    cs.current_scale = cs.current_scale .* current_scale
    # println("cs.current_scale: $(cs.current_scale)")
end

"""
    scaleto(x, y)

Scale to a specific scaling instead of multiplying it with the current scale.
For scaling on top of the current scale have a look at [`scale`](@ref).
"""
function scaleto(x, y)
    cs = get_current_setting()
    cs.desired_scale = (x, y)
    scaling = (x, y) ./ cs.current_scale
    # we divided by 0 but clearly we want to scale to 0
    # -> we want scaling to be 0 not Inf
    if x ≈ 0
        scaling = (0.0, scaling[2])
    end
    if y ≈ 0
        scaling = (scaling[1], 0.0)
    end
    Luxor.scale(scaling...)
    cs.current_scale = (x, y)
end


module LFkDTree

abstract type Node end

mutable struct INode <: Node
    i::Float64
    c::Float64
    lt::Ref{Node}
    rt::Ref{Node}
    pr::Ref{Node}
end

mutable struct LNode{K} <: Node where {K}
    k::K
end

end # module


module Neuron

using CUDA

include("lif.jl")
include("srm0.jl")

end


struct NzIterator{T<:AbstractArray}
    filter::Function
    output::Function
    A::T
end

nziterator(A::AbstractMatrix) = NzIterator(A)

NzIterator(A::AbstractMatrix) = NzIterator(_f_all, identity, A)
NzIterator(A::UpperTriangular) = NzIterator(_f_upper, identity, A.data)
NzIterator(A::LowerTriangular) = NzIterator(_f_lower, identity, A.data)
NzIterator(A::Transpose) = NzIterator(_f_all, _out_transpose, A.parent)
NzIterator(A::Adjoint) = NzIterator(_f_all, _out_adjoint, A.parent)
NzIterator(A::Symmetric) = _nziterator(A, identity, _out_transpose)
NzIterator(A::Hermitian) = _nziterator(A, _out_realdiag, _out_adjoint)

_out_transpose((i, j, v)) = (j, i, transpose(v))
_out_adjoint((i, j, v)) = (j, i, adjoint(v))
_out_realdiag((i, j, v)) = (i, j, i == j ? real(v) : adjoint(v))
_f_all(i, j) = true
_f_upper(i,j) = i <= j
_f_upper_strict(i,j) = i < j
_f_lower(i,j) = i >= j
_f_lower_strict(i,j) = i > j

function _nziterator(A::Union{Symmetric,Hermitian}, out1::Function, out2::Function)
    filter1 = A.uplo == 'U' ? _f_upper : _f_lower
    filter2 = A.uplo == 'U' ? _f_upper_strict : _f_lower_strict

    Base.Iterators.flatten((NzIterator(filter1, out1, A.data),
                            NzIterator(filter2, out2, A.data)))
end

import Base: iterate

Base.IteratorSize(::Type{<:NzIterator}) = Base.SizeUnknown()
Base.show(io::IO, itr::NzIterator) = println(io, "Non-zero Iterator for $(typeof(itr.A))")

function iterate(itr::NzIterator{<:SparseMatrixCSC{Tv,Ti}}) where {Tv,Ti}
    A = itr.A
    m, n = size(A)
    m > 0 && n > 0 && nnz(A) > 0 || return nothing
    colptr = A.colptr
    rowval = A.rowval
    j = 1
    while colptr[j+1] <= 1
        j += 1
        j > n && break
    end
    j > n && return nothing
    iterate(itr, (0, colptr[j+1], j, n))
end
function iterate(itr::NzIterator{<:SparseMatrixCSC{Tv,Ti}}, (k, kmax, j, n)::NTuple{4,Ti}) where {Ti,Tv}

    colptr = itr.A.colptr
    rowval = itr.A.rowval
    nzval = itr.A.nzval
    filter = itr.filter
    while j <= n
        k += 1
        while k >= kmax
            j += 1
            j > n && return nothing
            kmax = colptr[j+1]
        end
        i = rowval[k]
        nzv = nzval[k]
        nzv != 0 && filter(i,j) && return itr.output((i, j, nzv)), (k, kmax, j, n)
    end
    nothing
end

function iterate(itr::NzIterator{<:DenseMatrix})
    m, n = size(itr.A)
    m > 0 && n > 0 || return nothing
    iterate(itr, (0, m, 1, n))
end
function iterate(itr::NzIterator{<:DenseMatrix}, (i, m, j, n)::NTuple{4,Int})
    filter = itr.filter
    while j <= n
        i += 1
        if i > m
            j += 1
            i = 0
        else
            nzv = itr.A[i,j]
            nzv != 0 && filter(i, j) && return itr.output((i, j, nzv)), (i, m, j, n)
        end
    end
    nothing
end

function iterate(itr::NzIterator{<:Union{Diagonal,Bidiagonal,Tridiagonal,SymTridiagonal}})
    n = size(itr.A, 1)
    iterate(itr, (0, n))
end
function iterate(itr::NzIterator{<:Diagonal}, (i, n)::Tuple{Int,Int})
    while (i += 1) <= n
        nzv = itr.A.diag[i]
        nzv != 0 && filter(i, i) && return itr.output((i, i, nzv)), (i, n)
    end
end

function iterate(itr::NzIterator{<:Bidiagonal}, (i, n)::Tuple{Int,Int})
    k = itr.A.uplo == 'U' ? 1 : 0
    while (i += 1) <= 2n-1
        nzv, ii, jj = i <= n ? (itr.A.dv[i], i, i) : (itr.A.ev[i-n], i+1-k-n, i+k-n)
        nzv != 0 && itr.filter(ii, jj) && return itr.output((ii, jj, nzv)), (i, n)
    end
end

function iterate(itr::NzIterator{<:SymTridiagonal}, (i, n)::Tuple{Int,Int})
    while (i += 1) <= 3n-2
        nzv, ii, jj =
            i <= n ? (itr.A.dv[i], i, i) :
            i < 2n ? (itr.A.ev[i-n], i-n, i+1-n) :
                     (itr.A.ev[i-2n+1], i+2-2n, i-2n+1)

        nzv != 0 && itr.filter(ii, jj) && return itr.output((ii, jj, nzv)), (i, n)
    end
end

function iterate(itr::NzIterator{<:Tridiagonal}, (i, n)::Tuple{Int,Int})
    while (i += 1) <= 3n-2
        nzv, ii, jj =
            i <= n ? (itr.A.d[i], i, i) :
            i < 2n ? (itr.A.du[i-n], i-n, i+1-n) :
                     (itr.A.dl[i-2n+1], i-2n+2, i-2n+1)

        nzv != 0 && itr.filter(ii, jj) && return itr.output((ii, jj, nzv)), (i, n)
    end
end

function iterate(itr::NzIterator{<:AbstractMatrix})
    stack = NzIterator(itr.A)
    r = iterate(stack)
    _iterate(itr, r, stack)
end
function iterate(itr::NzIterator{<:AbstractMatrix}, (s, stack))
    r = iterate(stack, s)
    _iterate(itr, r, stack)
end
function _iterate(itr::NzIterator, r::Union{Nothing,Tuple}, stack)
    while r !== nothing
        (i, j, nz), s = r
        itr.filter(i, j) && return itr.output((i, j, nz)), (s, stack)
        r = iterate(stack, s)
    end
    nothing
end


using Test
using Tables
using MLJBase
import MLJFlux
using CategoricalArrays
using ColorTypes
using Flux
import Random
import Random.seed!
using Statistics
import StatsBase
using StableRNGs

using ComputationalResources
using ComputationalResources: CPU1, CUDALibs

const RESOURCES = Any[CPU1(), CUDALibs()]
EXCLUDED_RESOURCE_TYPES = Any[]

MLJFlux.gpu_isdead() && push!(EXCLUDED_RESOURCE_TYPES, CUDALibs)

@info "MLJFlux supports these computational resources:\n$RESOURCES"
@info "Current test run to exclude resources with "*
    "these types, as unavailable:\n$EXCLUDED_RESOURCE_TYPES\n"*
    "Excluded tests marked as \"broken\"."

# alternative version of Short builder with no dropout; see
# https://github.com/FluxML/Flux.jl/issues/1372 and
# https://github.com/FluxML/Flux.jl/issues/1372
mutable struct Short2 <: MLJFlux.Builder
    n_hidden::Int     # if zero use geometric mean of input/output
    σ
end
Short2(; n_hidden=0, σ=Flux.sigmoid) = Short2(n_hidden, σ)
function MLJFlux.build(builder::Short2, rng, n, m)
    n_hidden =
        builder.n_hidden == 0 ? round(Int, sqrt(n*m)) : builder.n_hidden
    init = Flux.glorot_uniform(rng)
    return Flux.Chain(
        Flux.Dense(n, n_hidden, builder.σ, init=init),
        Flux.Dense(n_hidden, m, init=init))
end

seed!(123)

include("test_utils.jl")

@testset "penalizers" begin
    include("penalizers.jl")
end

@testset "core" begin
    include("core.jl")
end

@testset "builders" begin
    include("builders.jl")
end

@testset "mlj_model_interface" begin
    include("mlj_model_interface.jl")
end

@testset "regressor" begin
    include("regressor.jl")
end

@testset "classifier" begin
    include("classifier.jl")
end

@testset "image" begin
    include("image.jl")
end

@testset "integration" begin
    include("integration.jl")
end


# this script produces Figure 4

using CSV, DataFrames, DelimitedFiles, Glob, GMT, Statistics
stats = expanduser("~/SeisNoise_paper/DATA/gmap-stations.txt")
df = DataFrame!(CSV.File(stats,delim="|"))
lat,lon = df[:,:Latitude], df[:,:Longitude]
latmin, latmax = minimum(lat), maximum(lat)
lonmin, lonmax = minimum(lon), maximum(lon)
latpad = 0.025
lonpad = 0.05

# create figure
coast(
    region=(lonmin-lonpad,lonmax+lonpad,latmin-latpad,latmax+latpad),
    proj=:Mercator,
    res=:full,
    land=:white,
    water=:green,
    # scale=0.03,
    par=(:MAP_FRAME_TYPE,"fancy+"),
    map_scale="jBR+c$(mean(lat))+w10k+f+o2.5/2.5+at+u",
    figsize=15,
)
scatter!(lon,lat,fill=:gold,marker=:triangle,markeredgecolor=:black,markersize=0.2)
basemap!(inset=(anchor=:TL, width=0.25, offset=(1, 5), save="xx000"))
t = readdlm("xx000")
coast!(region=[-107 -91 26 41], proj=:merc,
           land=:lightgray, area=5000, shore=:faint,
           x_off=t[1], y_off=t[2],
           N="a",
           figsize=4,
           frame=:bare,
)
scatter!(
     [mean(lon)],
     [mean(lat)],
     marker=:star,
     markeredgecolor=:black,
     fill=:gold,
     markersize=0.5,
     savefig=expanduser("~/SeisNoise_paper/LASSOmap.png"),
)


Base.minimum(solution::DESolution) = minimum(map(minimum, solution.u))
Base.maximum(solution::DESolution) = maximum(map(maximum, solution.u))

using DiffEqBase: AbstractTimeseriesSolution

struct MissingSolution <: AbstractTimeseriesSolution{Missing,Missing,Missing} end
function Base.show(io::IO, A::MissingSolution)
  print(io,"MissingSolution")
end
function Base.show(io::IO, m::MIME"text/plain", A::MissingSolution)
  print(io,"MissingSolution")
end

struct BareSolution{S,N,U<:Array{<:Array{<:S,N}},X,T} <: AbstractTimeseriesSolution{S,N,U}
    u::U
    x::X
    t::T
    tslocation::Int
end
BareSolution(; u, x, t) = BareSolution(u,x,t)
BareSolution(u::U, x::X, t::T) where {S,N,ELT<:Union{S,Array{<:S,N}},U<:Array{ELT},X,T} = BareSolution{S,N,U,X,T}(u,x,t,0)
BareSolution(u::Missing, x::Missing, t::Missing) = MissingSolution()
function Base.getproperty(bs::BareSolution, sym::Symbol)
    if sym ∈ [:dense, :prob]
        return false
    else
        return Base.getfield(bs, sym)
    end
end
timepoints(bs::BareSolution) = bs.t
coordinates(bs::BareSolution) = bs.x
function Base.show(io::IO, A::BareSolution)
  print(io,"t: ")
  show(io, A.t)
  println(io)
  print(io,"u: ")
  show(io, A.u)
end
function Base.show(io::IO, m::MIME"text/plain", A::BareSolution)
  print(io,"t: ")
  show(io,m,A.t)
  println(io)
  print(io,"u: ")
  show(io,m,A.u)
end
@generated function Simulation73.population_timepoint(solution::BareSolution{T,NP}, pop_dx::Int, time_dx::Int) where {T,NP}
    N = NP - 1 # N + pops
    colons = [:(:) for i in 1:N]
    :(solution[time_dx][$(colons...), pop_dx])
end


# Generic rand and rand! implementations similar to those in Distributions,
# but with an rng argument:

abstract type BATSampler{F<:VariateForm,S<:ValueSupport} <: Sampleable{F,S} end

@inline Random.rand(s::BATSampler, args...) = rand(Random.GLOBAL_RNG, s, args...)
@inline Random.rand!(s::BATSampler, args...) = rand!(Random.GLOBAL_RNG, s, args...)

# To avoid ambiguity with Distributions:
@inline Random.rand(s::BATSampler, dims::Int...) = rand(Random.GLOBAL_RNG, s, dims...)
@inline Random.rand!(s::BATSampler, A::AbstractVector) = rand!(Random.GLOBAL_RNG, s, A)
@inline Random.rand!(s::BATSampler, A::AbstractMatrix) = rand!(Random.GLOBAL_RNG, s, A)

function Random.rand!(rng::AbstractRNG, s::BATSampler{Univariate}, A::AbstractArray)
    @inbounds @simd for i in 1:length(A)
        A[i] = rand(rng, s)
    end
    return A
end

Random.rand(rng::AbstractRNG, s::BATSampler{Univariate}, dims::Dims) =
    rand!(rng, s, Array{eltype(s)}(undef, dims))

Random.rand(rng::AbstractRNG, s::BATSampler{Univariate}, dims::Int...) =
    rand!(s, Array{eltype(s)}(undef, dims))

Random.rand(rng::AbstractRNG, s::BATSampler{Multivariate}) =
    rand!(rng, s, Vector{eltype(s)}(undef, length(s)))

Random.rand(rng::AbstractRNG, s::BATSampler{Multivariate}, n::Integer) =
    rand!(rng, s, Matrix{eltype(s)}(undef, length(s), n))

function Random.rand!(rng::AbstractRNG, s::BATSampler{Multivariate}, A::AbstractMatrix)
    _check_rand_compat(s, A)
    @inbounds for i = 1:size(A,2)
        rand!(rng, s, view(A,:,i))
    end
    return A
end

Random.rand(rng::AbstractRNG, s::BATSampler{Multivariate}, n::Int) = rand!(rng, s, Matrix{eltype(s)}(undef, length(s), n))


for (f, FT, arg) in ((:-, typeof(-), Number),
                     (:+, typeof(+), Real),
                     (:*, typeof(*), Real))
    @eval begin
        function Base.broadcasted(b::DefaultArrayStyle{1}, ::$FT, x::$arg, r::AbstractAxis)
            return maybe_unsafe_reconstruct(r, Base.broadcasted(b, $f, x, parent(r)); keys=keys(r))
        end
        function Base.broadcasted(b::DefaultArrayStyle{1}, ::$FT, r::AbstractAxis, x::$arg)
            return maybe_unsafe_reconstruct(r, Base.broadcasted(b, $f, parent(r), x); keys=keys(r))
        end
    end
end

#=
for (f, FT, arg) in ((:+, typeof(+), Real),)
    @eval begin
        function Base.broadcasted(::DefaultArrayStyle{1}, ::$FT, x::$arg, r::AbstractAxis)
            return unsafe_reconstruct(r, broadcast($f, x, eachindex(r)))
        end
        function Base.broadcasted(::DefaultArrayStyle{1}, ::$FT, r::AbstractAxis, x::$arg)
            return unsafe_reconstruct(r, broadcast($f, eachindex(r), x))
        end
    end
end
=#

function Broadcast.broadcast_shape(
    shape1::Tuple,
    shape2::Tuple{Vararg{<:AbstractAxis}},
    shapes::Tuple...
)

    return Broadcast.broadcast_shape(_bcs(shape1, shape2), shapes...)
end

function Broadcast.broadcast_shape(
    shape1::Tuple{Vararg{<:AbstractAxis}},
    shape2::Tuple,
    shapes::Tuple...
)

    return Broadcast.broadcast_shape(_bcs(shape1, shape2), shapes...)
end
function Broadcast.broadcast_shape(
    shape1::Tuple{Vararg{<:AbstractAxis}},
    shape2::Tuple{Vararg{<:AbstractAxis}},
    shapes::Tuple...
)

    return Broadcast.broadcast_shape(_bcs(shape1, shape2), shapes...)
end

# _bcs consolidates two shapes into a single output shape
_bcs(::Tuple{}, ::Tuple{}) = ()
_bcs(::Tuple{}, newshape::Tuple) = (newshape[1], _bcs((), tail(newshape))...)
_bcs(shape::Tuple, ::Tuple{}) = (shape[1], _bcs(tail(shape), ())...)
function _bcs(shape::Tuple, newshape::Tuple)
    return (combine_axis(first(shape), first(newshape)), _bcs(tail(shape), tail(newshape))...)
end

_parent_if_axis(x::AbstractAxis) = parent(x)

###
### cat
###
cat_indices(x, y) = set_length(indices(x), length(x) + length(y))

cat_keys(x::AbstractVector, y::AbstractRange) = StaticRanges.grow_last(y, length(x))

cat_keys(x::AbstractRange, y::AbstractVector) = StaticRanges.grow_last(x, length(y))

cat_keys(x::AbstractRange, y::AbstractRange) = StaticRanges.grow_last(x, length(y))

cat_keys(x::AbstractVector, y::AbstractVector) = cat_keys!(promote(x, y)...,)

cat_keys(x::AbstractVector{T}, y::AbstractVector{T}) where {T} = cat_keys!(copy(x), y)

function cat_keys!(x::AbstractVector{T}, y::AbstractVector{T}) where {T}
    for x_i in x
        if x_i in y
            error("Element $x_i appears in both collections in call to cat_axis!(collection1, collection2). All elements must be unique.")
        end
    end
    return vcat(x, y)
end

function cat_axis(x::Axis, y::AbstractUnitRange, inds=cat_indices(x, y))
    return maybe_unsafe_reconstruct(x, inds; keys=cat_keys(keys(x), keys(y)))
end

function cat_axis(x::Axis, y::Axis, inds=cat_indices(x, y))
    return maybe_unsafe_reconstruct(x, inds; keys=cat_keys(keys(x), keys(y)))
end

function cat_axis(x::AbstractUnitRange, y::Axis, inds=cat_indices(x, y))
    return maybe_unsafe_reconstruct(y, inds; keys=cat_keys(keys(x), keys(y)))
end

function cat_axis(x::AbstractUnitRange, y::AbstractUnitRange, inds=cat_indices(x, y))
    if x isa AbstractAxis
        return unsafe_reconstruct(x, inds)
    else
        return unsafe_reconstruct(y, inds)
    end
end

#=
    cat_axes(x::AbstractArray, y::AbstractArray, xy::AbstractArray, dims)

Produces the appropriate set of axes where `x` and `y` are the arrays that were
concatenated over `dims` to produce `xy`. The appropriate indices of each axis
are derived from from `xy`.
=#
@inline function cat_axes(x::AbstractArray, y::AbstractArray, xy::AbstractArray{T,N}, dims) where {T,N}
    ntuple(Val(N)) do i
        if i in dims
            cat_axis(axes(x, i), axes(y, i), axes(xy, i))
        else
            combine_axis(axes(x, i), axes(y, i), axes(xy, i))
        end
    end
end
# TODO do these work?
vcat_axes(x::AbstractArray, y::AbstractArray, xy::AbstractArray) = cat_axes(x, y, xy, 1)

hcat_axes(x::AbstractArray, y::AbstractArray, xy::AbstractArray) = cat_axes(x, y, xy, 2)

###
### combine
###
# LinearIndices indicates that keys are not formally defined so the collection
# that isn't LinearIndices is used. If both are LinearIndices then take the underlying
# OneTo as the new keys.
combine_keys(x, y) = _combine_keys(keys(x), keys(y))
_combine_keys(x, y) = promote_axis_collections(x, y)
_combine_keys(x,                y::LinearIndices) = x
_combine_keys(x::LinearIndices, y               ) = y
_combine_keys(x::LinearIndices, y::LinearIndices) = first(y.indices)

# TODO I still really don't like this solution but the result seems better than any
# alternative I've seen out there
#=
Problem: String, Symbol, Second, etc. don't promote neatly with Int (default key value)
Solution: Symbol(element), Second(element) for promotion
Exception: String(element) doesn't work so string(element) has to be used
=#

promote_axis_collections(x::X, y::X) where {X} = x

function promote_axis_collections(x::LinearIndices{1}, y::Y) where {Y}
    return promote_axis_collections(x.indices[1], y)
end
function promote_axis_collections(x::X, y::LinearIndices{1}) where {X}
    return promote_axis_collections(x, y.indices[1])
end
function promote_axis_collections(x::LinearIndices{1}, y::LinearIndices{1})
    return promote_axis_collections(x.indices[1], y.indices[1])
end
function promote_axis_collections(x::X, y::Y) where {X,Y}
    if promote_rule(X, Y) <: Union{}
        Z = promote_rule(Y, X)
    else
        Z = promote_rule(X, Y)
    end
    if Z <: Union{}
        Tx = eltype(X)
        Ty = eltype(Y)
        Tnew = promote_type(Tx, Ty)
        if Tnew == Any
            if is_key(Tx)
                if Tx <: AbstractString
                    return promote_axis_collections(x, string.(y))
                else
                    return promote_axis_collections(x, Tx.(y))
                end
            else
                if is_key(Ty)
                    if Ty <: AbstractString
                        return promote_axis_collections(string.(x), y)
                    else
                        return promote_axis_collections(Ty.(x), y)
                    end
                else
                    error("No method available for promoting keys of type $Tx and $Ty.")
                end
            end
        else
            return promote_axis_collections(Tnew.(x), y)
        end
    else
        return Z(x)
    end
end

###
### combine_axis
###
combine_axis(x, y) = combine_axis(x, y, _combine_indices(x, y))
combine_axis(x, y, inds) = SimpleAxis(inds)
combine_axis(x, y::AbstractAxis, inds) = combine_axis(y, x, inds)
combine_axis(x::AbstractAxis, y, inds) = unsafe_reconstruct(x, inds)
function combine_axis(x::AbstractAxis, y::AbstractAxis, inds)
    if y isa AbstractOffsetAxis
        return unsafe_reconstruct(x, unsafe_reconstruct(parent(y), inds))
    else
        return unsafe_reconstruct(x, unsafe_reconstruct(y, inds))
    end
end

combine_axis(x::SimpleAxis, y::AbstractAxis, inds) = combine_axis(parent(x), y, inds)
combine_axis(x::AbstractAxis, y::SimpleAxis, inds) = combine_axis(x, parent(y), inds)
combine_axis(x::SimpleAxis, y::SimpleAxis, inds) = combine_axis(parent(x), parent(y), inds)

# Axis
combine_axis(x::Axis, y, inds) = resize_last(x, inds)
combine_axis(x::Axis, y::SimpleAxis, inds) = combine_axis(x, parent(y), inds)
function combine_axis(x::Axis, y::AbstractAxis, inds)
    return resize_last(x, combine_axis(parent(x), y, inds))
end
# TODO check this stuff
function combine_axis(x::Axis, y::Axis, inds)
    return Axis(
        combine_keys(x, y),
        combine_axis(parent(x), parent(y), inds)
    )
end

_combine_indices(x, y) = One():_combine_length(static_length(x), static_length(y))
function _combine_length(::StaticInt{X}, ::StaticInt{Y}) where {X,Y}
    return StaticInt(_combine_length(X, Y))
end
_combine_length(::StaticInt{X}, y) where {X} = _combine_length(X, y)
_combine_length(x, ::StaticInt{Y}) where {Y} = _combine_length(x, Y)
function _combine_length(x, y)
    if x == y || y == 1
        return x
    elseif x == 1
        return y
    else
        throw(DimensionMismatch("arrays could not be broadcast to a common size;" *
                                " got a dimension with lengths $(x) and $(y)"))
    end
end


function FieldEquinoctialOrbit(arg0::FieldOrbit)
    return FieldEquinoctialOrbit((FieldOrbit,), arg0)
end

function FieldEquinoctialOrbit(arg0::FieldPVCoordinates, arg1::Frame, arg2::FieldAbsoluteDate, arg3::RealFieldElement)
    return FieldEquinoctialOrbit((FieldPVCoordinates, Frame, FieldAbsoluteDate, RealFieldElement), arg0, arg1, arg2, arg3)
end

function FieldEquinoctialOrbit(arg0::RealFieldElement, arg1::RealFieldElement, arg10::RealFieldElement, arg11::RealFieldElement, arg12::PositionAngle, arg13::Frame, arg14::FieldAbsoluteDate, arg15::RealFieldElement, arg2::RealFieldElement, arg3::RealFieldElement, arg4::RealFieldElement, arg5::RealFieldElement, arg6::RealFieldElement, arg7::RealFieldElement, arg8::RealFieldElement, arg9::RealFieldElement)
    return FieldEquinoctialOrbit((RealFieldElement, RealFieldElement, RealFieldElement, RealFieldElement, PositionAngle, Frame, FieldAbsoluteDate, RealFieldElement, RealFieldElement, RealFieldElement, RealFieldElement, RealFieldElement, RealFieldElement, RealFieldElement, RealFieldElement, RealFieldElement), arg0, arg1, arg10, arg11, arg12, arg13, arg14, arg15, arg2, arg3, arg4, arg5, arg6, arg7, arg8, arg9)
end

function FieldEquinoctialOrbit(arg0::RealFieldElement, arg1::RealFieldElement, arg2::RealFieldElement, arg3::RealFieldElement, arg4::RealFieldElement, arg5::RealFieldElement, arg6::PositionAngle, arg7::Frame, arg8::FieldAbsoluteDate, arg9::RealFieldElement)
    return FieldEquinoctialOrbit((RealFieldElement, RealFieldElement, RealFieldElement, RealFieldElement, RealFieldElement, RealFieldElement, PositionAngle, Frame, FieldAbsoluteDate, RealFieldElement), arg0, arg1, arg2, arg3, arg4, arg5, arg6, arg7, arg8, arg9)
end

function FieldEquinoctialOrbit(arg0::TimeStampedFieldPVCoordinates, arg1::Frame, arg2::RealFieldElement)
    return FieldEquinoctialOrbit((TimeStampedFieldPVCoordinates, Frame, RealFieldElement), arg0, arg1, arg2)
end

function add_kepler_contribution(obj::FieldEquinoctialOrbit, arg0::PositionAngle, arg1::RealFieldElement, arg2::Vector{RealFieldElement})
    return jcall(obj, "addKeplerContribution", void, (PositionAngle, RealFieldElement, Vector{RealFieldElement}), arg0, arg1, arg2)
end

function eccentric_to_mean(::Type{FieldEquinoctialOrbit}, arg0::RealFieldElement, arg1::RealFieldElement, arg2::RealFieldElement)
    return jcall(FieldEquinoctialOrbit, "eccentricToMean", RealFieldElement, (RealFieldElement, RealFieldElement, RealFieldElement), arg0, arg1, arg2)
end

function eccentric_to_true(::Type{FieldEquinoctialOrbit}, arg0::RealFieldElement, arg1::RealFieldElement, arg2::RealFieldElement)
    return jcall(FieldEquinoctialOrbit, "eccentricToTrue", RealFieldElement, (RealFieldElement, RealFieldElement, RealFieldElement), arg0, arg1, arg2)
end

function get_a(obj::FieldEquinoctialOrbit)
    return jcall(obj, "getA", RealFieldElement, ())
end

function get_a_dot(obj::FieldEquinoctialOrbit)
    return jcall(obj, "getADot", RealFieldElement, ())
end

function get_e(obj::FieldEquinoctialOrbit)
    return jcall(obj, "getE", RealFieldElement, ())
end

function get_e_dot(obj::FieldEquinoctialOrbit)
    return jcall(obj, "getEDot", RealFieldElement, ())
end

function get_equinoctial_ex(obj::FieldEquinoctialOrbit)
    return jcall(obj, "getEquinoctialEx", RealFieldElement, ())
end

function get_equinoctial_ex_dot(obj::FieldEquinoctialOrbit)
    return jcall(obj, "getEquinoctialExDot", RealFieldElement, ())
end

function get_equinoctial_ey(obj::FieldEquinoctialOrbit)
    return jcall(obj, "getEquinoctialEy", RealFieldElement, ())
end

function get_equinoctial_ey_dot(obj::FieldEquinoctialOrbit)
    return jcall(obj, "getEquinoctialEyDot", RealFieldElement, ())
end

function get_hx(obj::FieldEquinoctialOrbit)
    return jcall(obj, "getHx", RealFieldElement, ())
end

function get_hx_dot(obj::FieldEquinoctialOrbit)
    return jcall(obj, "getHxDot", RealFieldElement, ())
end

function get_hy(obj::FieldEquinoctialOrbit)
    return jcall(obj, "getHy", RealFieldElement, ())
end

function get_hy_dot(obj::FieldEquinoctialOrbit)
    return jcall(obj, "getHyDot", RealFieldElement, ())
end

function get_i(obj::FieldEquinoctialOrbit)
    return jcall(obj, "getI", RealFieldElement, ())
end

function get_i_dot(obj::FieldEquinoctialOrbit)
    return jcall(obj, "getIDot", RealFieldElement, ())
end

function get_l(obj::FieldEquinoctialOrbit, arg0::PositionAngle)
    return jcall(obj, "getL", RealFieldElement, (PositionAngle,), arg0)
end

function get_l_dot(obj::FieldEquinoctialOrbit, arg0::PositionAngle)
    return jcall(obj, "getLDot", RealFieldElement, (PositionAngle,), arg0)
end

function get_le(obj::FieldEquinoctialOrbit)
    return jcall(obj, "getLE", RealFieldElement, ())
end

function get_le_dot(obj::FieldEquinoctialOrbit)
    return jcall(obj, "getLEDot", RealFieldElement, ())
end

function get_lm(obj::FieldEquinoctialOrbit)
    return jcall(obj, "getLM", RealFieldElement, ())
end

function get_lm_dot(obj::FieldEquinoctialOrbit)
    return jcall(obj, "getLMDot", RealFieldElement, ())
end

function get_lv(obj::FieldEquinoctialOrbit)
    return jcall(obj, "getLv", RealFieldElement, ())
end

function get_lv_dot(obj::FieldEquinoctialOrbit)
    return jcall(obj, "getLvDot", RealFieldElement, ())
end

function get_type(obj::FieldEquinoctialOrbit)
    return jcall(obj, "getType", OrbitType, ())
end

function has_derivatives(obj::FieldEquinoctialOrbit)
    return jcall(obj, "hasDerivatives", jboolean, ())
end

function interpolate(obj::FieldEquinoctialOrbit, arg0::FieldAbsoluteDate, arg1::Stream)
    return jcall(obj, "interpolate", FieldEquinoctialOrbit, (FieldAbsoluteDate, Stream), arg0, arg1)
end

function mean_to_eccentric(::Type{FieldEquinoctialOrbit}, arg0::RealFieldElement, arg1::RealFieldElement, arg2::RealFieldElement)
    return jcall(FieldEquinoctialOrbit, "meanToEccentric", RealFieldElement, (RealFieldElement, RealFieldElement, RealFieldElement), arg0, arg1, arg2)
end

function shifted_by(obj::FieldEquinoctialOrbit, arg0::RealFieldElement)
    return jcall(obj, "shiftedBy", FieldEquinoctialOrbit, (RealFieldElement,), arg0)
end

function shifted_by(obj::FieldEquinoctialOrbit, arg0::jdouble)
    return jcall(obj, "shiftedBy", FieldEquinoctialOrbit, (jdouble,), arg0)
end

function to_orbit(obj::FieldEquinoctialOrbit)
    return jcall(obj, "toOrbit", EquinoctialOrbit, ())
end

function to_string(obj::FieldEquinoctialOrbit)
    return jcall(obj, "toString", JString, ())
end

function true_to_eccentric(::Type{FieldEquinoctialOrbit}, arg0::RealFieldElement, arg1::RealFieldElement, arg2::RealFieldElement)
    return jcall(FieldEquinoctialOrbit, "trueToEccentric", RealFieldElement, (RealFieldElement, RealFieldElement, RealFieldElement), arg0, arg1, arg2)
end


#!/usr/bin/env bash
# -*- mode: julia -*-
#=
curdir="$(dirname "$(readlink -f "$0")")"
parentdir=$(dirname "$curdir")
echo "$parentdir"
exec julia-1.1.1 --project="$parentdir" "${BASH_SOURCE[0]}"
=#
module ROSWrapper
#TODO strongly and consistently type (esp w/ gpu to avoid Float64/32 conv)
#TODO get rid of branching in cb's
using RoboLib.Util: binarize, takeN, rangebearing2point
using CoordinateTransformations, Rotations
using RoboLib.Geom: Pose2D, project2D, Scale2D
using StaticArrays
using LinearAlgebra
using Statistics
using Random
using mushr_pf.ParticleFilter: setup_pf, update_act!, update_obs!, reset!, resample!, OccMap

# ROS Imports
using RobotOS
@rosimport sensor_msgs.msg: LaserScan
@rosimport nav_msgs.msg: Odometry
@rosimport nav_msgs.srv: GetMap
@rosimport geometry_msgs.msg: PoseStamped, PoseWithCovarianceStamped, PoseArray, Quaternion
@rosimport std_msgs.msg: Header, Float64
@rosimport ackermann_msgs.msg: AckermannDriveStamped
@rosimport vesc_msgs.msg: VescStateStamped

rostypegen(@__MODULE__)
using .sensor_msgs.msg
using .nav_msgs.msg
using .nav_msgs.srv
using .geometry_msgs.msg
using .std_msgs.msg: Header, Float64Msg
using .ackermann_msgs.msg
using .vesc_msgs.msg: VescStateStamped
# -----

include("settings.jl")
include("util.jl")

function laser_cb(msg::LaserScan, rospf::ROSPF)
    @debugtask begin
        h = msg.header
        trecv = time()
        if rospf.last_laser == -1
            rospf.last_laser = trecv

            # a wise man once said "buffers are the way to happiness"
            angles = LinRange(msg.angle_min, msg.angle_max, length(msg.ranges))
            resize!(rospf.angles, length(angles))
            resize!(rospf.laseridxs, length(angles))
            rospf.angles .= angles
        else
            dt = trecv - rospf.last_laser
            rospf.last_laser = trecv

            #println((1/dt, size())

            idxs = validindices(msg.ranges, rospf)
            angles = takeN(view(rospf.angles, idxs), rospf.nrays)
            ranges = takeN(view(msg.ranges, idxs), rospf.nrays)

            @assert length(angles) == length(ranges) == rospf.nrays

            update_obs!(rospf.pf, (angles, ranges), dt)
            resample!(rospf.pf)

            meanpose = mean(rospf.pf.belief.particles)
            publish(rospf.odom_pub, rosodometrystamped(meanpose, "map"))
            publish(rospf.pose_pub, rosposestamped(meanpose, "map"))

            # Publish a subset of the particles
            particles = view(rospf.pf.belief.particles, rand!(rospf.rand_part_idxs, eachindex(rospf.rand_part_idxs)))
            publish(rospf.particle_pub, rosposearraystamped(particles, "map"))

            # Publish the range measurements used in the reweighting step
            laserposes!(rospf, angles, ranges, meanpose)
            laserposes = rosposearraystamped(rospf.laserposes, "map")
            publish(rospf.laserpose_pub, laserposes)
        end
    end
end

function clicked_pose_cb(msg::PoseWithCovarianceStamped, rospf::ROSPF)
    @debugtask begin
        p = msg.pose.pose.position
        q = msg.pose.pose.orientation

        r2d = project2D(Quat(q.w, q.x, q.y, q.z))
        t2d = SVector{2}(p.x, p.y)
        reset!(rospf.pf, Pose2D(r2d, t2d))

        rospf.initialized = true
    end
end

function servo_cb(msg, rospf::ROSPF)
    rospf.last_servo = msg.data
end

function vesc_cb(msg, rospf::ROSPF)
    @debugtask begin
        trecv = to_sec(msg.header.stamp)
        if rospf.last_servo == -1
            return
        elseif rospf.last_vesc == -1
            rospf.last_vesc = trecv
        else
            dt = trecv - rospf.last_vesc
            rospf.last_vesc = trecv

            v = (msg.state.speed - rospf.speed2erpm_offset) / rospf.speed2erpm_gain
            delta = (rospf.last_servo - rospf.steering2servo_offset) / rospf.steering2servo_gain

            update_act!(rospf.pf, SVector(v, delta), dt)
        end
    end
end

function ackermann_cb(msg, rospf::ROSPF)
    @debugtask begin
        h = msg.header
        d = msg.drive
        trecv = time()
        if rospf.last_acker == -1
            rospf.last_acker = trecv
        else
            dt = trecv - rospf.last_acker
            rospf.last_acker = trecv

            ctrl = SVector(d.speed, d.steering_angle)
            update_act!(rospf.pf, ctrl, dt)
        end
    end
end

#TODO(cxs): nested params
#TODO(cxs): handle array/dict params
#TODO(cxs): handle unknown
function get_params(d::Dict{Symbol, Any})
    for (param_name, value) in d
        if isa(value, Dict)
            get_params(value)
        elseif isa(value, Union{Number, String})
            if has_param('~' * string(param_name))
                new_value = oftype(value, get_param('~' * string(param_name), value))
                @assert typeof(value) == typeof(new_value)
                loginfo("Found param \"$param_name\". Default value: \"$value\", new value: \"$new_value\"")
                d[param_name] = new_value
            else
                loginfo("Param \"$param_name\" not set. Using default value: \"$value\"")
            end
        end
    end
end

function process_ros_map(map_msg)
    info = map_msg.map.info

    s = info.resolution
    p = info.origin.position
    q = info.origin.orientation

    # something something row major vs column major
    occmap = reshape(map_msg.map.data', (info.width, info.height))'
    occmap = OccMap((el)->el!=0, occmap, SVector(p.x, p.y), project2D(Quat(q.w, q.x, q.y, q.z)), s)

    return occmap
end

function main()
    init_node("particle_filter")

    get_params(conf)

    loginfo("Waiting for map service...")
    wait_for_service(rosconf[:static_map_service])
    map_msg = ServiceProxy(rosconf[:static_map_service], GetMap)(GetMapRequest())
    occmap = process_ros_map(map_msg)

    pose_pub      = Publisher(rosconf[:pose_topic], PoseStamped, queue_size = 1)
    particle_pub  = Publisher(rosconf[:particle_topic], PoseArray, queue_size = 1)
    laserpose_pub = Publisher(rosconf[:laserpose_topic], PoseArray, queue_size = 1)
    odom_pub = Publisher(rosconf[:odom_topic], Odometry, queue_size = 1)

    # setup PF
    pfconf[:rng] = Random.MersenneTwister(pfconf[:rng_seed])
    pf = setup_pf(occmap, pfconf)
    rospf = ROSPF(conf, pf, pose_pub, particle_pub, odom_pub, occmap, laserpose_pub)

    pose_sub  = Subscriber(rosconf[:initialpose_topic], PoseWithCovarianceStamped, clicked_pose_cb, (rospf,), queue_size=1)

    loginfo("Waiting for initial pose...")
    while !rospf.initialized
        rossleep(Duration(0.001))
    end

    laser_sub = Subscriber(rosconf[:scan_topic], LaserScan, laser_cb, (rospf,), queue_size=1)
    servo_sub = Subscriber(rosconf[:servo_pos_topic], Float64Msg, servo_cb, (rospf,), queue_size=1)
    vesc_sub = Subscriber(rosconf[:vesc_topic], VescStateStamped, vesc_cb, (rospf,), queue_size=1)
    #ackermann_sub = Subscriber(rosconf[:ackermann_topic], AckermannDriveStamped, ackermann_cb, (rospf,), queue_size=1)

    loginfo("Initialization complete")
    d = Duration(0.001)
    while !is_shutdown() rossleep(d) end
end

end # module

using .ROSWrapper: main
if !isinteractive()
    main()
end


#using Distributions
#include("constants.jl")
#include("orbit.jl")
#include("planet.jl")

if !isdefined(:PlanetarySystemAbstract)
  @compat abstract type PlanetarySystemAbstract end

  immutable PlanetarySystem{StarT<:StarAbstract} <: PlanetarySystemAbstract
    star::StarT
    planet::Vector{Planet}         # TODO OPT: If want to optimize, consider using  something like ImmmutableArrays?
    orbit::Vector{Orbit}           # TODO OPT: If want to optimize, consider using  something like ImmmutableArrays?
    # TODO DETAIL: Setup inner constructor to enforce equal number of planets & orbits
    #PlanetarySystem{StarT}(s::StarAbstract, p::Vector{Planet}, o::Vector{Orbit})
    #  @assert(length(p)==length(o)) # else error(string("Number of planets must match number of orbits: Np= ",length(p)," No= ",length(o)))
    #  new(s,p,o)
    #end
  end
  PlanetarySystemSingleStar = PlanetarySystem{SingleStar}

end
PlanetarySystem{StarT<:StarAbstract}(s::StarT) = PlanetarySystem(s,Array{Planet}(0),Array{Orbit}(0))  # Constructor for a Planetary System with no planets

PlanetarySystem{StarT<:StarAbstract}(s::StarT, p::Planet, o::Orbit) = PlanetarySystem(s,[p],[o])  # Constructor for a single Planet System

function PlanetarySystem{StarT}(ps::PlanetarySystem{StarT}, keep::Vector{Int64}) # Why doesn't this work?
   PlanetarySystem{StarT}(ps.star,ps.planet[keep],ps.orbit[keep])
end

#function PlanetarySystemSingleStar(ps::PlanetarySystemSingleStar, keep::Vector{Int64})
#function PlanetarySystem{StarT<:StarAbstract}(ps::PlanetarySystem{StarT}, keep::Vector{Int64})
#   PlanetarySystem(ps.star,ps.planet[keep],ps.orbit[keep])
#end

flux(ps::PlanetarySystem{StarAbstract}) = flux(ps.star)
flux(ps::PlanetarySystem{Star}) = flux(ps.star)
flux(ps::PlanetarySystem{BinaryStar}) = flux(ps.star)
flux(ps::PlanetarySystem{MultipleStar}) = flux(ps.star)

function num_planets(s::PlanetarySystem{Star})
  @assert( length(s.planet) == length(s.orbit) )    # TODO OPT: Deactivate inner assert's like this for speed once tested
  return length(s.planet)
end

function generate_planet_mass_from_radius_powerlaw(r::Float64, sim_param::SimParam)
  const mr_power_index::Float64 = get_real(sim_param,"mr_power_index")
  const mr_const::Float64 = get_real(sim_param,"mr_const")
  const mr_max_mass::Float64 = get_real(sim_param,"mr_max_mass")
  m = mr_const*earth_mass*(r/earth_radius)^mr_power_index
  if m > mr_max_mass
     m = mr_max_mass
  end
  return m
end

function generate_planet_mass_from_radius_powerlaw(r::Float64, s::Star, o::Orbit, sim_param::SimParam)  # TODO SCI: IMPORTANT once have stability criteria, replace w/ better M-R relationship
  generate_planet_mass_from_radius_powerlaw(r,sim_param)
end

function generate_num_planets_poisson(lambda::Real, max_planets::Integer; min_planets::Integer = 0)
  if lambda < min_planets*1e-3
      return min_planets
  end
  bug_fixed = false # TODO OPT: true case should work, but Danley found bug in Distributions package.  Revert once fixed for speed.
  local n
  if bug_fixed
     d = Distributions.Truncated(Distributions.Poisson(lambda),min_planets,max_planets)
     n = rand(d)
  else
     if min_planets == 0
        min_planets = -1
     end
     d = Distributions.Truncated(Distributions.Poisson(lambda),min_planets,max_planets)
     n = rand(d)
     #=
     n = -1
     while !(min_planets<=n<=max_planets)
        n = rand(Distributions.Poisson(lambda))
     end
     =#
  end
  return n
end

function generate_num_planets_poisson(s::Star, sim_param::SimParam)
  const lambda::Float64 = exp(get_real(sim_param,"log_eta_pl"))
  const max_tranets_in_sys::Int64 = get_int(sim_param,"max_tranets_in_sys")
  generate_num_planets_poisson(lambda,max_tranets_in_sys)
end

function generate_period_and_sizes_log_normal(s::Star, sim_param::SimParam; num_pl::Integer = 1)  # TODO USER SCI:  User should make sure planetary properties are being drawn appropriately for their scientific purposes
    const mu_log_r::Float64 = get_real(sim_param,"mean_log_planet_radius")
    const sigma_log_r::Float64 = get_real(sim_param,"sigma_log_planet_radius")
    const mu_log_P::Float64 = get_real(sim_param,"mean_log_planet_period")
    const sigma_log_P::Float64 = get_real(sim_param,"sigma_log_planet_period")
    const min_period::Float64 = get_real(sim_param,"min_period")
    const max_period::Float64 = get_real(sim_param,"max_period")
    const min_radius::Float64 = get_real(sim_param,"min_radius")
    const max_radius::Float64 = get_real(sim_param,"max_radius")
    const max_draws::Int64 = 100

    if   sigma_log_r <= 0. || sigma_log_P<=0.
     println("# mu_log_r= ", mu_log_r, " sigma_log_r= ", sigma_log_r, " mu_log_P= ", mu_log_P, " sigma_log_P= ", sigma_log_P)
    end
    const rdist = LogNormal(mu_log_r,sigma_log_r)
    const Pdist = LogNormal(mu_log_P,sigma_log_P)
    #Rlist = rand(rdist,num_pl)
    #Plist = rand(Pdist,num_pl)
    #idx_keep = find(i->(min_radius<=Rlist[i]<=max_radius) && (min_period<=Plist[i]<=max_period), 1:num_pl )
    #return Plist[idx_keep], Rlist[idx_keep]  # replaced because want to return exactly num_pl.  Could use Truncated to restore above code.
    Rlist = zeros(num_pl)
    Plist = zeros(num_pl)
    for i in 1:num_pl
      j = 0
      while ! (min_radius<Rlist[i]<=max_radius) && j<max_draws
            Rlist[i] = rand(rdist)
            j+=1
      end
      if j>=max_draws
         println("# Struggled to draw size for: ",mu_log_r, " ", sigma_log_r)
      end
      j = 0
      while ! (min_period<Plist[i]<=max_period) && j<max_draws
            Plist[i] = rand(Pdist)
            j+=1
      end
      if j>=max_draws
         println("# Struggled to draw period for: ",mu_log_P, " ", sigma_log_P)
      end
    end
    return Plist, Rlist
end

function draw_power_law(n::Real, x0::Real, x1::Real, num_pl::Integer)
     ((x1^(n+1) - x0^(n+1)).*rand(num_pl) .+ x0^(n+1)).^(1/(n+1))
end

function draw_broken_power_law(n1::Real, n2::Real, x0::Real, x1::Real, xb::Real, num_pl::Integer)
    #x0 and x1 are the lower and upper truncation limits, and xb is the break point, i.e. x0 <= xb <= x1 (all must be positive)
    #n1 and n2 are the power law indices between x0 and xb, and xb and x1, respectively (can be positive or negative)
    @assert(x0 <= xb <= x1)
    @assert(num_pl >= 1)

    u_draws = rand(num_pl) #'num_pl' draws from the uniform distribution between 0 and 1
    x_draws = zeros(num_pl)

    if (n1 != -1) & (n2 != -1)
        C1 = 1./(((xb^(n1+1) - x0^(n1+1))/(n1+1)) + ((xb^(n1-n2)*(x1^(n2+1) - xb^(n2+1)))/(n2+1))) #normalization constant
        ub = (C1*(xb^(n1+1) - x0^(n1+1)))/(n1+1) #break point in u, between 0 and 1
        for (i,u) in enumerate(u_draws)
            if u <= ub
                x_draws[i] = (((n1+1)*u)/C1 + x0^(n1+1))^(1/(n1+1))
            else #if u > ub
                x_draws[i] = (((n2+1)/(C1*xb^(n1-n2)))*(u - (C1*(xb^(n1+1) - x0^(n1+1)))/(n1+1)) + xb^(n2+1))^(1/(n2+1))
            end
        end
    elseif (n1 == -1) & (n2 != -1)
        C1 = 1./(log(xb/x0) + ((xb^(-1-n2))*(x1^(n2+1)) - 1)/(n2+1)) #normalization constant
        ub = C1*log(xb/x0) #break point in u, between 0 and 1
        for (i,u) in enumerate(u_draws)
            if u <= ub
                x_draws[i] = x0*exp(u/C1)
            else #if u > ub
                x_draws[i] = (((n2+1)/(C1*xb^(-1-n2)))*(u - C1*log(xb/x0)) + xb^(n2+1))^(1/(n2+1))
            end
        end
    elseif (n1 != -1) & (n2 == -1)
        C1 = 1./(((xb^(n1+1) - x0^(n1+1))/(n1+1)) + (xb^(n1+1))*log(x1/xb)) #normalization constant
        ub = (C1*(xb^(n1+1) - x0^(n1+1)))/(n1+1) #break point in u, between 0 and 1
        for (i,u) in enumerate(u_draws)
            if u <= ub
                x_draws[i] = (((n1+1)*u)/C1 + x0^(n1+1))^(1/(n1+1))
            else #if u > ub
                x_draws[i] = xb*exp((1/(C1*xb^(n1+1)))*(u - (C1*(xb^(n1+1) - x0^(n1+1)))/(n1+1)))
            end
        end
    else #if n1 == -1 and n2 == -1 (i.e. it is a single power-law with index of -1)
        for (i,u) in enumerate(u_draws)
            x_draws[i] = x0*exp(u*log(x1/x0))
        end
    end

    return x_draws
end

function generate_periods_power_law(s::Star, sim_param::SimParam; num_pl::Integer = 1)
    const power_law_P::Float64 = get_real(sim_param,"power_law_P")
    const min_period::Float64 = get_real(sim_param,"min_period")
    const max_period::Float64 = get_real(sim_param,"max_period")
    Plist = power_law_P!=-1.0 ? draw_power_law(power_law_P,min_period,max_period, num_pl) : exp(log(min_period).+rand()*log(max_period/min_period))
    return Plist
end

function generate_sizes_power_law(s::Star, sim_param::SimParam; num_pl::Integer = 1)
    const power_law_r::Float64 = get_real(sim_param,"power_law_r")
    const min_radius::Float64 = get_real(sim_param,"min_radius")
    const max_radius::Float64 = get_real(sim_param,"max_radius")
    Rlist = power_law_r!=-1.0 ? draw_power_law(power_law_r,min_radius,max_radius, num_pl) : exp(log(min_radius).+rand()*log(max_radius/min_radius))
    return Rlist
end

function generate_sizes_broken_power_law(s::Star, sim_param::SimParam; num_pl::Integer = 1)
    const power_law_r1::Float64 = get_real(sim_param,"power_law_r1")
    const power_law_r2::Float64 = get_real(sim_param,"power_law_r2")
    const min_radius::Float64 = get_real(sim_param,"min_radius")
    const max_radius::Float64 = get_real(sim_param,"max_radius")
    const break_radius::Float64 = get_real(sim_param,"break_radius")
    Rlist = draw_broken_power_law(power_law_r1,power_law_r2,min_radius,max_radius,break_radius, num_pl)
    return Rlist
end

function generate_period_and_sizes_power_law(s::Star, sim_param::SimParam; num_pl::Integer = 1)
    return (generate_periods_power_law(s, sim_param, num_pl=num_pl), generate_sizes_power_law(s, sim_param, num_pl=num_pl))
end

function TruncatedUpper(d::Distributions.UnivariateDistribution, u::Float64)
    zero(u) < u || error("lower bound should be less than upper bound.")
    lcdf = zero(u)
    ucdf = isinf(u) ? one(u) : cdf(d, u)
    tp = ucdf - lcdf
    Distributions.Truncated{typeof(d),Distributions.value_support(typeof(d))}(d, zero(u), u, lcdf, ucdf, tp, log(tp))
end

function generate_e_omega_rayleigh_direct(sigma_hk::Float64; max_e::Float64 = 1.0)
  @assert(0<max_e<=1.0)
  ecc::Float64 = rand( TruncatedUpper(Rayleigh(sigma_hk),max_e) )
  w::Float64 = 2pi*rand()
  return ecc, w
end

function generate_e_omega_rayleigh_two_gaussians(sigma_hk::Float64; max_e::Float64 = 1.0)
  @assert(0<max_e<=1.0)
  h = k = 1.0
  while h*h+k*k >= max_e*max_e
    h = sigma_hk*randn()
    k = sigma_hk*randn()
  end
  ecc::Float64 = sqrt(h*h+k*k)
  w::Float64 = atan2(k,h)
  return ecc, w
end

function generate_e_omega_rayleigh(sigma_hk::Float64; max_e::Float64 = 1.0)
   if max_e > sigma_hk
      return generate_e_omega_rayleigh_two_gaussians(sigma_hk,max_e=max_e)
   else
      return generate_e_omega_rayleigh_direct(sigma_hk,max_e=max_e)
   end
end

function generate_e_omega_rayleigh(sim_param::SimParam; max_e::Float64 = 1.0)
  sigma_hk::Float64 = get_real(sim_param,"sigma_hk")
  generate_e_omega_rayleigh(sigma_hk, max_e=max_e)
end

function generate_planetary_system_hardcoded_example(star::StarAbstract, sim_param::SimParam; verbose::Bool = false)
  # in this version we specify fixed functions that are known at compile time, allowing for additional optimizations (~0.6 second faster per Kepler catalog out of ~3.6 sec on my laptop w/ 1 core)
  const  generate_planet_mass_from_radius = generate_planet_mass_from_radius_powerlaw
  const  generate_num_planets = generate_num_planets_poisson
  const  generate_period_and_sizes = generate_period_and_sizes_log_normal

  const  generate_e_omega =  generate_e_omega_rayleigh

  # const  generate_star = get_function(sim_param,"generate_star")
  # const star::StarAbstract = generate_star(sim_param)
  const num_pl::Int64 = generate_num_planets(star, sim_param)

  if( num_pl==0 )
    return PlanetarySystem(star)
  else
    (Plist::Vector{Float64}, Rlist::Vector{Float64}) = generate_period_and_sizes(star, sim_param, num_pl=num_pl)
    idx = sortperm(Plist)                   # TODO OPT: Check to see if sorting is significant time sink.  If so, it might could be deferred

    min_a_in_rstar = 2.0
    min_P_orbit = day_in_year*sqrt((min_a_in_rstar*star.radius*rsol_in_au)^3 / star.mass) # minimum semi-major axis of two stellar radii
    idx = idx[Plist[idx] .> min_P_orbit]
    if( length(idx)==0 )
        return PlanetarySystem(star)
    end

    pl = Array{Planet}(length(idx))
    orbit = Array{Orbit}(length(idx))
    a = map(i->semimajor_axis(Plist[i],star.mass),idx)
    max_e = ones(length(idx))
    max_e_factor = 0.999 # A factor just less than 1 to prevent numerical issues with near-crossing orbits
    if length(a)>=2
       for i in 1:length(a)
         if i==1
           max_e[i] = max_e_factor*(1-a[i]/a[i+1])/(1+a[i]/a[i+1])
         elseif i==length(a)
           max_e[i] = max_e_factor*(1-a[i-1]/a[i])/(1+a[i-1]/a[i])
         else
           max_e[i] = max_e_factor*min( (1-a[i]/a[i+1])/(1+a[i]/a[i+1]), (1-a[i-1]/a[i])/(1+a[i-1]/a[i]) )
         end
       end
    end
    for i in 1:length(idx)
      # if verbose   println("i=",i," idx=",idx," Plist=",Plist[idx] );     end
      P = Plist[idx[i]]
      Rpl = Rlist[idx[i]]
      (ecc::Float64,  omega::Float64) = generate_e_omega(sim_param, max_e=max_e[i])
      incl::Float64 = acos(rand())
      orbit[i] = Orbit(P,ecc,incl,omega,2pi*rand(),2pi*rand())
      mass::Float64 = generate_planet_mass_from_radius(Rpl, star, orbit[i], sim_param)
      pl[i] = Planet( Rpl,  mass )
    end
  return PlanetarySystem(star,pl,orbit)
  end
end

function generate_planetary_system_uncorrelated_incl(star::StarAbstract, sim_param::SimParam; verbose::Bool = false)
  # load functions to use for drawing parameters
  const  generate_planet_mass_from_radius = get_function(sim_param,"generate_planet_mass_from_radius")
  const  generate_num_planets = get_function(sim_param,"generate_num_planets")
  const  generate_period_and_sizes = get_function(sim_param,"generate_period_and_sizes")
  const  generate_e_omega = get_function(sim_param,"generate_e_omega")

  # const  generate_star = get_function(sim_param,"generate_star")
  # const star::StarAbstract = generate_star(sim_param)
  const num_pl::Int64 = generate_num_planets(star, sim_param)::Int64
  const sigma_ecc::Float64 = haskey(sim_param,"sigma_hk") ? get_real(sim_param,"sigma_hk") : 0.0

  if( num_pl==0 )
    return PlanetarySystem(star)::PlanetarySystem
  else
    (Plist::Vector{Float64}, Rlist::Vector{Float64}) = generate_period_and_sizes(star, sim_param, num_pl=num_pl)
    idx = sortperm(Plist)                   # TODO OPT: Check to see if sorting is significant time sink.  If so, it might could be deferred

    min_a_in_rstar = 2.0
    min_P_orbit = day_in_year*sqrt((min_a_in_rstar*star.radius*rsol_in_au)^3 / star.mass) # minimum semi-major axis of two stellar radii
    idx = idx[Plist[idx] .> min_P_orbit]
    if( length(idx)==0 )
        return PlanetarySystem(star)
    end

    pl = Array{Planet}(length(idx))
    orbit = Array{Orbit}(length(idx))
    a = map(i->semimajor_axis(Plist[i],star.mass),idx)
    max_e = ones(length(idx))
    max_e_factor = 0.999 # A factor just less than 1 to prevent numerical issues with near-crossing orbits
    if length(a)>=2
       for i in 1:length(a)
          if i==1
            max_e[i] = max_e_factor*(1-a[i]/a[i+1])/(1+a[i]/a[i+1])
          elseif i==length(a)
            max_e[i] = max_e_factor*(1-a[i-1]/a[i])/(1+a[i-1]/a[i])
          else
            max_e[i] = max_e_factor*min( (1-a[i]/a[i+1])/(1+a[i]/a[i+1]), (1-a[i-1]/a[i])/(1+a[i-1]/a[i]) )
          end
        end
    end
    for i in 1:length(idx)
      # if verbose   println("i=",i," idx=",idx," Plist=",Plist[idx] );     end
      P = Plist[idx[i]]
      Rpl = Rlist[idx[i]]
      if haskey(sim_param,"sigma_hk_one") && haskey(sim_param,"sigma_hk_multi")
         sigma_ecc = num_pl == 1 ? get_real(sim_param,"sigma_hk_one") : get_real(sim_param,"sigma_hk_multi")
      end
      (ecc::Float64,  omega::Float64) = generate_e_omega(sim_param, max_e=max_e[i])
      incl::Float64 = acos(rand())
      orbit[i] = Orbit(P,ecc,incl,omega,2pi*rand(),2pi*rand())
      # set!(orbit[idx[i]],P,ecc,incl,omega,2pi*rand(),2pi*rand())
      mass::Float64 = generate_planet_mass_from_radius(Rpl, star, orbit[i], sim_param)
      pl[i] = Planet( Rpl,  mass )
    end
  return PlanetarySystem(star,pl,orbit)
  end
end

# This version generates more systems roughly near a common plane, but until incorporate CORBITS data, ABC can't match input param
function generate_planetary_system_simple(star::StarAbstract, sim_param::SimParam; verbose::Bool = false)
  # load functions to use for drawing parameters
  const  generate_planet_mass_from_radius = get_function(sim_param,"generate_planet_mass_from_radius")
  const  generate_num_planets = get_function(sim_param,"generate_num_planets")
  #const  generate_num_planets = generate_num_planets_christiansen
  #const  generate_period_and_sizes = generate_period_and_sizes_christiansen
  const  generate_period_and_sizes = get_function(sim_param,"generate_period_and_sizes")
  const  generate_e_omega = get_function(sim_param,"generate_e_omega")
  const  sigma_incl = deg2rad(get_real(sim_param,"sigma_incl"))

  # const  generate_star = get_function(sim_param,"generate_star")
  # const star::StarAbstract = generate_star(sim_param)
  const num_pl = generate_num_planets(star, sim_param)::Int64
  const sigma_ecc::Float64 = haskey(sim_param,"sigma_hk") ? get_real(sim_param,"sigma_hk") : 0.0

  if( num_pl==0 )
    return PlanetarySystem(star)
  else
    (Plist::Vector{Float64}, Rlist::Vector{Float64}) = generate_period_and_sizes(star, sim_param, num_pl=num_pl)
    idx = sortperm(Plist)                   # TODO OPT: Check to see if sorting is significant time sink.  If so, it might could be deferred
    incl_sys = acos(rand())

    min_a_in_rstar = 2.0
    min_P_orbit = day_in_year*sqrt((min_a_in_rstar*star.radius*rsol_in_au)^3 / star.mass) # minimum semi-major axis of two stellar radii
    idx = idx[Plist[idx] .> min_P_orbit]
    if( length(idx)==0 )
        return PlanetarySystem(star)
    end

    pl = Array{Planet}(length(idx))
    orbit = Array{Orbit}(length(idx))
    for i in 1:length(idx)
      # if verbose   println("i=",i," idx=",idx," Plist=",Plist[idx] );     end
      P = Plist[idx[i]]
      Rpl = Rlist[idx[i]]
      if haskey(sim_param,"sigma_hk_one") && haskey(sim_param,"sigma_hk_multi")
         sigma_ecc = num_pl == 1 ? get_real(sim_param,"sigma_hk_one") : get_real(sim_param,"sigma_hk_multi")
      end
      (ecc,  omega) = generate_e_omega(sim_param)::Tuple{Float64,Float64}
      incl_mut = sigma_incl*sqrt(randn()^2+randn()^2) # rand(Distributions.Rayleigh(sigma_incl)) # sigma_incl*randn()
      asc_node = 2pi*rand()
      mean_anom = 2pi*rand()
      #incl = incl_sys + sigma_incl*randn()
      incl =  incl_mut!=zero(incl_mut) ? acos( cos(incl_sys)*cos(incl_mut) + sin(incl_sys)*sin(incl_mut)*cos(asc_node) ) : incl_sys
      orbit[i] = Orbit(P,ecc,incl,omega,asc_node,mean_anom)
      mass = generate_planet_mass_from_radius(Rpl, star, orbit[i], sim_param)::Float64
      pl[i] = Planet( Rpl,  mass )
    end
  return PlanetarySystem(star,pl,orbit)
  end
end

function test_planetary_system_constructors(sim_param::SimParam)
  generate_star = get_function(sim_param,"generate_star")
  star = generate_star(sim_param)
  empty_sys = PlanetarySystem(star)
  earth = Planet(earth_radius,earth_mass)
  earth_orbit = Orbit(365.2425,0.0167,0.5*pi,0.0,0.0,0.0)
  solar_sys = PlanetarySystem(star, earth,earth_orbit)
  m = generate_planet_mass_from_radius_powerlaw(0.02,star,earth_orbit,sim_param)/earth_mass
  generate_planetary_system_simple(star,sim_param,verbose=true)
end


using NeXLUncertainties
using NeXLCore
using Colors
using ImageCore: colorview
using Statistics

"""
The members in common between KRatio and KRatios
"""
abstract type KRatioBase
    # element::Element
    # xrays::Vector{CharXRay} # Which CharXRays were measured?
    # unkProps::Dict{Symbol,Any} # Beam energy, take-off angle, coating, ???
    # stdProps::Dict{Symbol,Any} # Beam energy, take-off angle, coating, ???
    # standard::Material
end

"""
    Base.findfirst(krs::AbstractVector{<:KRatioBase}, cxr::CharXRay)

Find the first KRatio or KRatios in which the .xrays field contains the cxr::CharXRay.
"""
Base.findfirst(krs::AbstractVector{<:KRatioBase}, cxr::CharXRay) =
    krs[findfirst(kr -> cxr in kr.xrays, krs)]
element(kr::KRatioBase) = kr.element
xrays(krs::KRatioBase) = krs.xrays
standard(krs::KRatioBase) = krs.standard
"""
    elms(krs::Vector{KRatio})::Set{Element}

Return a set containing the elements present in krs.
"""
function elms(krs::Vector{<:KRatioBase})::Set{Element}
    res = Set{Element}()
    for kr in krs
        push!(res, kr.element)
    end
    return res
end

"""
    KRatio

The k-ratio is the ratio of two similar intensity measurements - one on
a material of unknown composition and one on a standard with known
composition. Each measurement has properties like :BeamEnergy (req),
:TakeOffAngle (req), :Coating (opt) that characterize the measurement.
A minimal set of properties includes:

Properties: (These `Symbol`s are intentionally the same used in NeXLSpectrum)

    :BeamEnergy incident beam energy in eV
    :TakeOffAngle in radians
    :Coating A NeXLCore.Film object or Film[] detailing a conductive coating

Some algorithms may require additional properties.

k-ratios are associated with characteristic X-rays (`CharXRay`) from a single element.
WDS k-ratios are typically associated with a single `CharXRay` while EDS
measurements may be associated with many `CharXRay` that are similar in energy.
k-ratios are always relative to another material.  Usually the composition of this
`Material` is well-known.  However, when a k-ratio is restandardized, it is possible
for the intermediate material to be less well-known.
"""
struct KRatio <: KRatioBase
    element::Element
    xrays::Vector{CharXRay} # Which CharXRays were measured?
    unkProps::Dict{Symbol,Any} # Beam energy, take-off angle, coating, ???
    stdProps::Dict{Symbol,Any} # Beam energy, take-off angle, coating, ???
    standard::Material
    kratio::UncertainValue

    function KRatio(
        xrays::AbstractVector{CharXRay},
        unkProps::Dict{Symbol,<:Any},
        stdProps::Dict{Symbol,<:Any},
        standard::Material,
        kratio::AbstractFloat,
    )
        if length(xrays) < 1
            error("A k-ratio must specify at least one characteristic X-ray.")
        end
        elm = element(xrays[1])
        if !all(element(l) == elm for l in xrays)
            error(
                "The characteristic X-rays in a k-ratio must all be from the same element.",
            )
        end
        if value(standard[elm]) <= 1.0e-6
            error("The standard $standard does not contain the element $(elm).")
        end
        if haskey(unkProps, :TakeOffAngle) &&
           haskey(stdProps, :TakeOffAngle) &&
           (
               !isapprox(
                   unkProps[:TakeOffAngle],
                   stdProps[:TakeOffAngle],
                   atol = deg2rad(0.1),
               )
           )
            @warn "The unknown and standard take-off angles do not match for $elm in $standard and $xrays."
        end
        return new(
            elm,
            sort(xrays, rev = true),
            copy(unkProps),
            copy(stdProps),
            standard,
            convert(UncertainValue, kratio),
        )
    end
end

NeXLUncertainties.value(kr::KRatio) = value(kr.kratio)
NeXLUncertainties.σ(kr::KRatio) = σ(kr.kratio)
nonnegk(kr::KRatio) = value(kr.kratio) < 0.0 ? uv(0.0, σ(kr.kratio)) : kr.kratio
Statistics.mean(krs::AbstractVector{KRatio})::UncertainValue =
    mean((kr.kratio for kr in krs)...)

function Base.getindex(krs::AbstractVector{KRatio}, cxr::CharXRay)
    i = findfirst(kr->cxr in kr.xrays, krs)
    return isnothing(i) ? zero(UncertainValue) : krs[i].kratio
end

"""
    strip(krs::AbstractVector{KRatio}, els::Element...)::Vector{KRatio}

Creates a new Vector{KRatio} containing all the KRatio objects in `krs` except those associated with the specified elements.
"""
Base.strip(krs::AbstractVector{KRatio}, els::Element...) =
    collect(filter(k -> !(element(k) in els), krs))

Base.show(io::IO, kr::KRatio) =
    print(io, "k[$(name(kr.xrays)), $(name(kr.standard))] = $(round(kr.kratio))")

function NeXLUncertainties.asa(::Type{DataFrame}, krs::AbstractVector{KRatio})::DataFrame
    return DataFrame(
        Symbol("X-rays") => [ repr(kr.xrays) for kr in krs ],
        Symbol("Standard") => [ name(kr.standard) for kr in krs ],
        Symbol("C[std]") => [ value(kr.standard[kr.element]) for kr in krs ],
        Symbol("ΔC[std]") => [ σ(kr.standard[kr.element]) for kr in krs ],
        Symbol("E₀[unk]") => [ get(kr.unkProps, :BeamEnergy, missing) for kr in krs ],
        Symbol("θ[unk]") => [ get(kr.unkProps, :TakeOffAngle, missing) for kr in krs ],
        Symbol("E₀[std]") => [ get(kr.stdProps, :BeamEnergy, missing) for kr in krs ],
        Symbol("θ[std]") => [ get(kr.stdProps, :TakeOffAngle, missing) for kr in krs ],
        Symbol("k") => [ value(kr.kratio) for kr in krs ],
        Symbol("Δk") => [ σ(kr.kratio) for kr in krs ],
    )
end

"""
`KRatios` represents the hyper-spectral equivalent of the KRatio type.  Each pixel in the `KRatios` object
must be characterized by the same unknown and standard properties, the same X-ray lines and the other
properties.
"""
struct KRatios{T} <: KRatioBase where {T <: AbstractFloat}
    element::Element
    xrays::Vector{CharXRay} # Which CharXRays were measured?
    unkProps::Dict{Symbol,Any} # Beam energy, take-off angle, coating, ???
    stdProps::Dict{Symbol,Any} # Beam energy, take-off angle, coating, ???
    standard::Material
    kratios::Array{T}

    function KRatios(
        xrays::Vector{CharXRay},
        unkProps::Dict{Symbol,<:Any},
        stdProps::Dict{Symbol,<:Any},
        standard::Material,
        kratios::Array{T},
    ) where { T <: AbstractFloat }
        if length(xrays) < 1
            error("A k-ratio must specify at least one characteristic X-ray.")
        end
        elm = element(xrays[1])
        if !all(element(l) == elm for l in xrays)
            error(
                "The characteristic X-rays in a k-ratio must all be from the same element.",
            )
        end
        if value(standard[elm]) <= 1.0e-6
            error("The standard $standard does not contain the element $(elm).")
        end
        if haskey(unkProps, :TakeOffAngle) &&
           haskey(stdProps, :TakeOffAngle) &&
           (
               !isapprox(
                   unkProps[:TakeOffAngle],
                   stdProps[:TakeOffAngle],
                   atol = deg2rad(0.1),
               )
           )
            @warn "The unknown and standard take-off angles do not match for $elm in $standard and $xrays."
        end
        return new{T}(elm, xrays, copy(unkProps), copy(stdProps), standard, kratios)
    end
end

Base.show(io::IO, kr::KRatios) = print(
    io,
    "k[$(name(kr.standard)), $(name(kr.xrays))] = $(eltype(kr.kratios))[ $(size(kr.kratios)) ]",
)

Base.getindex(krs::KRatios, idx::Int...) = KRatio(
    krs.xrays,
    krs.unkProps,
    krs.stdProps,
    krs.standard,
    getindex(krs.kratios, idx...),
)

Base.getindex(krs::KRatios, ci::CartesianIndex) = KRatio(
    krs.xrays,
    krs.unkProps,
    krs.stdProps,
    krs.standard,
    getindex(krs.kratios, ci),
)

Base.size(krs::KRatios) = size(krs.kratios)
Base.size(krs::KRatios, idx::Int) = size(krs.kratios, idx)
Base.CartesianIndices(krs::KRatios) = CartesianIndices(krs.kratios)

"""
    LinearAlgebra.normalize(krs::AbstractVector{<:KRatios}; norm::Float32=1.0f)::Vector{KRatios}
    LinearAlgebra.normalize(krs::AbstractVector{KRatio}; norm::Float32=1.0f)::Vector{KRatio}

Computes the pixel-by-pixel normalized k-ratio for each point in the KRatios data array. `norm` specifies normalization
constants other than 1.0.
"""
function LinearAlgebra.normalize(
    krs::AbstractVector{<:KRatios};
    norm::Float64 = 1.0,
    minsum::Float64 = 0.0
)
    sz = size(krs[1].kratios)
    @assert all((sz == size(kr.kratios) for kr in krs[2:end]))
    s = map(CartesianIndices(krs[1].kratios)) do ci
        ss = sum( max(0.0, kr.kratios[ci]) for kr in krs)
        ss<=minsum ? 1.0e100 : ss
    end
    return map(krs) do kr
        KRatios(kr.xrays, copy(kr.unkProps), copy(kr.stdProps), kr.standard, norm .* (kr.kratios ./ s))
    end
end
function LinearAlgebra.normalize(
    krs::AbstractVector{KRatio};
    norm::Float64 = 1.0,
    minsum::Float64 = 0.0
)
    s = sum(kr->max(0.0, kr.kratio), krs)
    s = s<=minsum ? 1.0e100 : s
    map(krs) do kr
        KRatio(kr.xrays, copy(kr.unkProps), copy(kr.stdProps), kr.standard, (norm/s) * kr.kratio)
    end
end

"""
    brightest(krs::Union{KRatios, KRatio})

Returns a new KRatios (referencing same basic data as krs) but with a single CharXRay in the `lines` field.
"""
brightest(krs::KRatios) = KRatios([brightest(krs.xrays)], krs.unkProps, krs.stdProps, krs.standard, krs.kratios )
brightest(krs::KRatio) = KRatio([brightest(krs.xrays)], krs.unkProps, krs.stdProps, krs.standard, krs.kratio )

"""
    colorize(krs::AbstractVector{<:KRatios}, red::Element, green::Element, blue::Element, normalize=:All[|:Each])
    colorize(krs::AbstractVector{<:KRatios}, elms::AbstractVector{Element}, normalize=:All)

Create RGB colorized images from up to three `Element`s.  The elements
are normalized relative to all `KRatios` in `krs`. The resulting images are scaled by the factor
`scale` to allow visualization of trace elements.
"""
function colorize(krs::AbstractVector{<:KRatios}, red::Element, green::Element, blue::Element, normalize=:All)
    colorize(krs, [red, green, blue], normalize)
end
function colorize(krs::AbstractVector{<:KRatios}, elms::AbstractVector{Element}, normalize=:All)
    idx  = collect( findfirst(kr->isequal(kr.element, elm), krs) for elm in elms )
    if normalize==:All
        # Normalize relative to sum of KRatios at each point
        s = map(CartesianIndices(krs[1])) do ci
            ss = sum( max(0.0, kr.kratios[ci]) for kr in krs)
            ss<=0.0 ? 1.0 : ss
        end
        clip(skr) = min(skr, 1.0)
        norm(kr) = clip.(kr.kratios ./ s)
        colorview(RGB,
            norm(krs[idx[1]]),
            length(idx)>1 ? norm(krs[idx[2]]) : zeroarray,
            length(idx)>2 ? norm(krs[idx[3]]) : zeroarray)
    else
        # Normalize relative to max of each KRatios independently
        each(kr) = kr.kratios / maximum(kr.kratios)
        colorview(RGB,
            each(krs[idx[1]]),
            length(idx)>1 ? each(krs[idx[2]]) : zeroarray,
            length(idx)>2 ? each(krs[idx[3]]) : zeroarray)
    end
end

function Base.getindex(krs::AbstractVector{<:KRatios}, elm::Element)
    colorize(krs, [ elm, elm, elm ])
end
function Base.getindex(krs::AbstractVector{<:KRatios}, red::Element, green::Element)
    colorize(krs, [ red, green ])
end
function Base.getindex(krs::AbstractVector{<:KRatios}, red::Element, green::Element, blue::Element)
    colorize(krs, [ red, green, blue ])
end


# Prediction
# ==========
#
# Sequence type predictor.
#
# This file is a part of BioJulia.
# License is MIT: https://github.com/BioJulia/Bio.jl/blob/master/LICENSE.md

# predict sequence type based on character frequencies in `seq[start:stop]`
function predict(seq::Vector{UInt8}, start, stop)
    # count characters
    a = c = g = t = u = n = alpha = 0
    for i in start:stop
        x = seq[i]
        if x == 0x41 || x == 0x61
            a += 1
        elseif x == 0x43 || x == 0x63
            c += 1
        elseif x == 0x47 || x == 0x67
            g += 1
        elseif x == 0x54 || x == 0x74
            t += 1
        elseif x == 0x55 || x == 0x75
            u += 1
        elseif x == 0x4e || x == 0x6e
            n += 1
        end
        if 0x41 ≤ x ≤ 0x5a || 0x61 ≤ x ≤ 0x7a
            alpha += 1
        end
    end

    # the threshold (= 0.95) is somewhat arbitrary
    if (a + c + g + t + u + n) / alpha > 0.95
        if t ≥ u
            return DNASequence
        else
            return RNASequence
        end
    else
        return AminoAcidSequence
    end
end


using Base.Test
using DataFrames
import WCS

const rcf = RunCamcolField(3900, 6, 269)

"""
Get the PSF located at a particular world location in an image.

Args:
 - world_loc: A location in world coordinates.
 - img: An Image

Returns:
 - An array of PsfComponent objects that represents the PSF as a mixture
     of Gaussians.
"""
function get_source_psf(world_loc::Vector{Float64}, img::Image, psf_K::Int)
    # Some stamps or simulated data have no raw psf information. In that case,
    # just use the psf from the image.
    if size(img.raw_psf_comp.rrows) == (0, 0)
        # Also return a vector of empty psf params
        return img.psf, fill(fill(NaN, length(PsfParams)), psf_K)
    else
        pixel_loc = WCS.world_to_pix(img.wcs, world_loc)
        psfstamp = Model.eval_psf(img.raw_psf_comp, pixel_loc[1], pixel_loc[2])
        return PSF.fit_raw_psf_for_celeste(psfstamp, psf_K)
    end
end

function test_images()
    # A lot of tests are in a single function to avoid having to reload
    # the full image multiple times.
    images = SDSSIO.load_field_images(rcf, datadir)

    dir = "$datadir/$(rcf.run)/$(rcf.camcol)/$(rcf.field)"
    fname = @sprintf "%s/photoObj-%06d-%d-%04d.fits" dir rcf.run rcf.camcol rcf.field
    cat_entries = SDSSIO.read_photoobj_celeste(fname)

    ea = make_elbo_args(images, cat_entries, patch_radius_pix=1e-6)

    # Just check some basic facts about the catalog.
    @test length(cat_entries) == 805
    @test 0 < sum([ce.is_star for ce in cat_entries]) < 805

    # Find an object near the middle of the image.
    ctr = WCS.pix_to_world(images[3].wcs, [images[3].H / 2, images[3].W / 2])
    dist = [sqrt((ce.pos[1] - ctr[1])^2 + (ce.pos[2] - ctr[2])^2)
                    for ce in cat_entries]
    obj_index = findmin(dist)[2]    # index of closest object
    obj_loc = cat_entries[obj_index].pos    # location of closest object

    # Test get_source_psf at point while we have the images loaded.
    test_b = 3
    img = ea.images[test_b]
    ea_obj = make_elbo_args(images, cat_entries[obj_index:obj_index])
    pixel_loc = WCS.world_to_pix(img.wcs, obj_loc)
    original_psf_val = Model.eval_psf(img.raw_psf_comp, pixel_loc[1], pixel_loc[2])

    original_psf_celeste =
        PSF.fit_raw_psf_for_celeste(original_psf_val, ea.psf_K)[1]
    fit_original_psf_val = PSF.get_psf_at_point(original_psf_celeste)

    obj_psf = get_source_psf(ea_obj.vp[1][ids.u], img, ea.psf_K)[1]
    obj_psf_val = PSF.get_psf_at_point(obj_psf)

    # The fits should match exactly.
    @test_approx_eq_eps(obj_psf_val, fit_original_psf_val, 1e-6)

    # The raw psf will not be as good.
    @test_approx_eq_eps(obj_psf_val, original_psf_val, 1e-2)

    cat_several = [cat_entries[1]; cat_entries[obj_index]]
    ea_several = make_elbo_args(ea_obj.images, cat_several)

    # The second set of vp is the object of interest
    point_patch_psf = PSF.get_psf_at_point(ea_several.patches[2, test_b].psf)
    # The threshold for the test below was formerly 1e-6.
    # Is it a problem I needed to increase it?
    @test_approx_eq_eps(obj_psf_val, point_patch_psf, 1e-4)
end

function test_stamp_get_object_psf()
    stamp_blob, stamp_mp, body = gen_sample_star_dataset()
    img = stamp_blob[3]
    obj_index =    stamp_mp.vp[1][ids.u]
    pixel_loc = WCS.world_to_pix(img.wcs, obj_index)
    original_psf_val = PSF.get_psf_at_point(img.psf)

    obj_psf_val = PSF.get_psf_at_point(
      get_source_psf(stamp_mp.vp[1][ids.u], img, 2)[1])
    @test_approx_eq_eps(obj_psf_val, original_psf_val, 1e-6)
end

function test_set_patch_size()
    # Test that the patch size gets most of the light from a variety of
    # galaxy shapes.
    # This shows that the current patch size is actually far too conservative.

    function gal_catalog_from_scale(gal_scale::Float64, flux_scale::Float64)
        CatalogEntry[CatalogEntry(world_location, false,
                                   flux_scale * fluxes, flux_scale * fluxes,
                                   0.1, .01, pi/4, gal_scale, "sample", 0) ]
    end

    srand(1)
    images0 = SampleData.load_stamp_blob(datadir, "164.4311-39.0359_2kpsf")
    img_size = 150
    for b in 1:5
            images0[b].H, images0[b].W = img_size, img_size
    end
    fluxes = [4.451805E+03,1.491065E+03,2.264545E+03,2.027004E+03,1.846822E+04]

    world_location = WCS.pix_to_world(images0[3].wcs,
                                      Float64[img_size / 2, img_size / 2])

    for gal_scale in [1.0, 10.0], flux_scale in [0.1, 10.0]
        cat = gal_catalog_from_scale(gal_scale, flux_scale)
        images = Synthetic.gen_blob(images0, cat)
        ea = make_elbo_args(images, cat)

        for b=1:length(images)
            @assert size(ea.images[b].tiles) == (1, 1)
            tile_image = DeterministicVI.tile_predicted_image(
                ea.images[b].tiles[1,1], ea, ea.tile_source_map[b][1,1])

            pixel_center = WCS.world_to_pix(images[b].wcs, cat[1].pos)
            radius = Model.choose_patch_radius(
                pixel_center, cat[1], images[b].psf, ea.images[b])

            circle_pts = fill(false, images[b].H, images[b].W)
            in_circle = 0.0
            for x=1:size(tile_image)[1], y=1:size(tile_image)[2]
                if ((x - pixel_center[1]) ^ 2 + (y - pixel_center[2]) ^ 2) < radius ^ 2
                    in_circle += tile_image[x, y]
                    circle_pts[x, y] = true
                end
            end
            @test in_circle / sum(tile_image) > 0.95
        end
    end
end

test_images()
#test_stamp_get_object_psf()
#test_set_patch_size()


# Suppress warnings during testing.
# required for reducing logging during tests
using Memento
setlevel!(getlogger(PowerModels), "error")

# is this the best way to find a file in a package?
base_dir = string(dirname(dirname(pathof(PowerSystems))))
case5_data = PM.parse_file(joinpath(base_dir, "data/matpower/case5.m"))
case5_data = PM.replicate(case5_data, 2)

case5_dc_data = PM.parse_file(joinpath(base_dir, "data/matpower/case5_dc.m"))
case5_dc_data = PM.replicate(case5_dc_data, 2)

# TODO: currently JuMP.num_variables is the best we can do to introspect the JuMP model.
#  Ideally this would also test the number of constraints generated

@testset "PowerModels Model Build" begin
    pm = PowerSimulations.instantiate_nip_model(case5_data, DCPPowerModel)
    @test JuMP.num_variables(pm.model) == 34
    pm = PowerSimulations.instantiate_nip_model(case5_data, PM.ACPPowerModel)
    @test JuMP.num_variables(pm.model) == 96
    pm = PowerSimulations.instantiate_nip_model(case5_data, PM.SOCWRPowerModel)
    @test JuMP.num_variables(pm.model) == 110
end

@testset "PM with type extensions" begin
    pm = PowerSimulations.instantiate_nip_model(case5_dc_data, DCPPowerModel)
    JuMP.num_variables(pm.model) == 34 # Repetitive?
    pm = PowerSimulations.instantiate_nip_model(case5_dc_data, ACPPowerModel)
    JuMP.num_variables(pm.model) == 96
    pm = PowerSimulations.instantiate_nip_model(case5_dc_data, DCPPowerModel)
    @test JuMP.num_variables(pm.model) == 36
    pm = PowerSimulations.instantiate_nip_model(case5_dc_data, DCPPowerModel)
    @test JuMP.num_variables(pm.model) == 36 # 48 was originally here, verify??
end


const COND_DATA_PATH = joinpath(@__DIR__,"data","boolean_table_RTG13.csv")
const DEL_CONC = 1e-4

"""
Steady state solver.

A Stiff solver is chosen `TRBDF2` (stand for Tranpezoidal Backward Differeital Formula) from **DifferentialEquations.jl** for robust simulation of stability.

`TRBF2` is a one-step method based on trapezoidal rule and the backward differentiaion formula of order 2, and it is strongly L-stable that means `TRBFS` is good at integrating stiff equations [1].

References
----------
1. Hosea, M. E., & Shampine, L. F. (1996). Analysis and implementation of TR-BDF2. Applied Numerical Mathematics, 20(1-2), 21-37. URL: https://www.sciencedirect.com/science/article/pii/0168927495001158?via%3Dihub
2. DifferentialEqiations.jl - Stiff Problems. [[link](https://diffeq.sciml.ai/stable/solvers/ode_solve/#Stiff-Problems)]
"""
const SSMETHOD = DynamicSS(TRBDF2())

const DataFiles = (;
    RNAseq= joinpath(@__DIR__,"data","RNAseq_RTG_expression.csv"),
    BoolCond = joinpath(@__DIR__,"data","boolean_table_RTG13.csv"),
    solution_rtgM4 = joinpath(@__DIR__,"data","solution_rtgM4.csv")
)

const default_sol_i = 1

"""
Parameter searching
"""
const K_dist = Exponential(1000)
const K_N_dist = Geometric(0.3)
const S_SPAN = (1e-6,1.0)
const Max_expr = 10000.
const TRANS_THRESHOLD = 1.5
const DataTables = getTables(DataFiles)
const DefaultCondition = "stressed"


module SimpleDamage

using Base: @kwdef
using StaticArrays
import DamagedShearBand as DSB
using IfElse

using Reexport
@reexport using DamagedShearBand: Rheology

include("types.jl")
include("damage_functions.jl")
include("functions.jl")
include("data_simulation.jl")

export Model, NumericalSetup, TriaxialSetup, DeformationControl, ConstantStrainRate, ConstantStress, ConstitutiveModel, LinearWeakening, AsymptoticWeakening, EnergyBasedWeakening
export ConstantYieldStress, CoulombYieldStress, SimpleCharlesLaw, InvariantsKICharlesLaw, PrincipalKICharlesLaw, MicroMechanicalCharlesLaw, Elasticity, IncompressibleElasticity, Geom2D, Geom3D
export Plasticity, MinViscosityThreshold, DamageThreshold
export update_derivatives, update_derivatives!
export init_variables, init_params

end#module