import ModelingToolkit
import ModelingToolkit: build_function, substitute
import Catalyst
import Catalyst: ReactionSystem
using StaticArrays
import Base.show
using Setfield

struct ReactionSet
    rates::Vector{Float64}
    rstoich::Vector{Vector{Pair{Int64,Int64}}}
    nstoich::Vector{Vector{Pair{Int64,Int64}}}
    nspecies::Int
end

abstract type AbstractJumpRateAggregatorAlgorithm end

struct GillespieDirect <: AbstractJumpRateAggregatorAlgorithm end

struct DepGraphDirect <: AbstractJumpRateAggregatorAlgorithm end

abstract type AbstractJumpRateAggregator end

struct DirectAggregator{U} <: AbstractJumpRateAggregator
    "the sum of all reaction rates"
    sumrate::Float64

    "a vector of the current reaction rates"
    rates::Vector{Float64}

    "sum of group propensities"
    gsums::Vector{Float64}

    "maps reaction indices to group indices"
    ridtogroup::U

    "time span for aggregation"
    tspan::Tuple{Float64,Float64}

    "time of last recorded reaction"
    tprev::Float64

    "accumulated log-probability"
    weight::Float64
end

function build_aggregator(alg::GillespieDirect, reactions::ReactionSet, ridtogroup, tspan=(0.0, Inf64))
    ngroups = maximum(ridtogroup)
    nreactions = length(reactions.rates)
    DirectAggregator(0.0, zeros(nreactions), zeros(ngroups), ridtogroup, tspan, tspan[1], 0.0)
end

struct DepGraphAggregator{U,DepGraph} <: AbstractJumpRateAggregator
    "the sum of all reaction rates"
    sumrate::Float64

    "a vector of the current reaction rates"
    rates::Vector{Float64}

    "sum of group propensities"
    gsums::Vector{Float64}

    "maps reaction indices to group indices"
    ridtogroup::U

    "time span for aggregation"
    tspan::Tuple{Float64,Float64}

    "time of last recorded reaction"
    tprev::Float64

    "accumulated log-probability"
    weight::Float64

    "dependency graph"
    depgraph::DepGraph

    "previously executed reaction"
    prev_reaction::Int
end

function build_aggregator(alg::DepGraphDirect, reactions::ReactionSet, ridtogroup, tspan=(0.0, Inf64))
    ngroups = maximum(ridtogroup)
    nreactions = length(reactions.rates)
    depgraph = make_depgraph(reactions)
    DepGraphAggregator(0.0, zeros(nreactions), zeros(ngroups), ridtogroup, tspan, tspan[1], 0.0, depgraph, 0)
end

function initialize_aggregator(agg::AbstractJumpRateAggregator; tspan=(0.0, Inf64))
    agg = @set agg.tspan = tspan
    agg = @set agg.tprev = tspan[1]
    agg = @set agg.weight = 0.0
    agg = @set agg.sumrate = 0.0
    agg.rates .= 0.0
    agg.gsums .= 0.0
    agg
end

function initialize_aggregator(agg::DepGraphAggregator; tspan=(0.0, Inf64))
    agg = @set agg.tspan = tspan
    agg = @set agg.tprev = tspan[1]
    agg = @set agg.weight = 0.0
    agg = @set agg.sumrate = 0.0
    agg = @set agg.prev_reaction = 0
    agg.rates .= 0.0
    agg.gsums .= 0.0
    agg
end

function ReactionSet(js::ModelingToolkit.JumpSystem, p)
    parammap = map(Pair, ModelingToolkit.parameters(js), p)
    statetoid = Dict(ModelingToolkit.value(state) => i for (i, state) in enumerate(ModelingToolkit.states(js)))

    rates = Float64[]
    rstoich_vec = Vector{Pair{Int64,Int64}}[]
    nstoich_vec = Vector{Pair{Int64,Int64}}[]

    for eq in ModelingToolkit.equations(js)
        rate = ModelingToolkit.value(ModelingToolkit.substitute(eq.scaled_rates, parammap))
        rstoich = sort!([statetoid[ModelingToolkit.value(spec)] => stoich for (spec, stoich) in eq.reactant_stoch])
        nstoich = sort!([statetoid[ModelingToolkit.value(spec)] => stoich for (spec, stoich) in eq.net_stoch])

        push!(rates, rate)
        push!(rstoich_vec, rstoich)
        push!(nstoich_vec, nstoich)
    end

    ReactionSet(rates, rstoich_vec, nstoich_vec, length(ModelingToolkit.states(js)))
end

function species_to_dependent_reaction_map(reactions::ReactionSet)
    nspecies = reactions.nspecies
    # map from a species to reactions that depend on it
    spec_to_dep_rxs = [Vector{Int}() for n = 1:nspecies]
    for (rx, complex) in enumerate(reactions.rstoich)
        for (spec, stoch) in complex
            push!(spec_to_dep_rxs[spec], rx)
        end
    end

    foreach(s -> unique!(sort!(s)), spec_to_dep_rxs)
    spec_to_dep_rxs
end

function make_depgraph(reactions::ReactionSet)
    nreactions = length(reactions.rates)
    spec_to_dep_rxs = species_to_dependent_reaction_map(reactions)

    # create map from rx to reactions depending on it
    dep_graph = [Vector{Int}() for n = 1:nreactions]
    for rx in 1:nreactions
        # rx changes spec, hence rxs depending on spec depend on rx
        for (spec, stoch) in reactions.nstoich[rx]
            for dependent_rx in spec_to_dep_rxs[spec]
                push!(dep_graph[rx], dependent_rx)
            end
        end
    end

    foreach(deps -> unique!(sort!(deps)), dep_graph)
    dep_graph
end

struct TrajectoryDistribution{A}
    reactions::ReactionSet
    aggregator::A
end

function TrajectoryDistribution(reactions::ReactionSet, alg::AbstractJumpRateAggregatorAlgorithm, ridtogroup=1:length(reactions.rates))
    agg = build_aggregator(alg, reactions, ridtogroup)
    TrajectoryDistribution(reactions, agg)
end

distribution(rn::ReactionSystem, p; update_map=1:Catalyst.numreactions(rn), alg=GillespieDirect()) = TrajectoryDistribution(ReactionSet(convert(ModelingToolkit.JumpSystem, rn), p), alg, update_map)

@fastmath function Distributions.logpdf(dist::TrajectoryDistribution, trajectory)::Float64
    traj_iter = trajectory_iterator(trajectory)
    tprev = 0.0
    result = 0.0
    agg = initialize_aggregator(dist.aggregator)
    for (u, t, i) in traj_iter
        agg = update_rates(agg, (u, t, i), dist.reactions)

        dt = t - tprev
        result -= dt * agg.sumrate
        if i != 0
            @inbounds gid = agg.ridtogroup[i]
            gid != 0 && @inbounds result += log(agg.gsums[gid])
        end

        tprev = t
    end

    result
end

@fastmath function fold_logpdf(dist::TrajectoryDistribution, agg::AbstractJumpRateAggregator, (u, t, i))
    agg = update_rates(agg, (u, t, i), dist.reactions)
    dt = t - agg.tprev
    log_jump_prob = 0.0
    if i != 0
        @inbounds gid = agg.ridtogroup[i]
        gid != 0 && @inbounds log_jump_prob = log(agg.gsums[gid])
    end
    log_surv_prob = -dt * agg.sumrate
    agg = @set agg.weight = agg.weight + log_surv_prob + log_jump_prob
    agg = @set agg.tprev = t
end

function step_energy(dist::TrajectoryDistribution, agg::AbstractJumpRateAggregator, (u, t, i))
    if t <= agg.tspan[1]
        return agg
    end
    if t > agg.tspan[2]
        if agg.tprev >= agg.tspan[2]
            return agg
        end
        fold_logpdf(dist, agg, (u, agg.tspan[2], 0))
    else
        fold_logpdf(dist, agg, (u, t, i))
    end
end

function trajectory_energy(dist::TrajectoryDistribution, traj; tspan=(0.0, Inf64))
    agg = initialize_aggregator(dist.aggregator, tspan=tspan)
    traj_iter = trajectory_iterator(traj)
    agg = Base.foldl((acc, x) -> step_energy(dist, acc, x), traj_iter; init=agg)
    agg.weight
end

function cumulative_logpdf!(result::AbstractVector, dist::TrajectoryDistribution, traj, dtimes::AbstractVector)
    tspan = (first(dtimes), last(dtimes))
    agg = initialize_aggregator(dist.aggregator, tspan=tspan)
    result[1] = zero(eltype(result))
    traj_iter = trajectory_iterator(traj)
    result_agg, k = Base.foldl(traj_iter; init=(agg, 1)) do (agg, k), (u, t, i)
        if t <= agg.tspan[1]
            return agg, k
        end

        t = min(t, agg.tspan[2])
        agg = update_rates(agg, (u, t, i), dist.reactions)

        tprev = agg.tprev
        while k <= length(dtimes) && dtimes[k] < t
            result[k] -= (dtimes[k] - tprev) * agg.sumrate
            tprev = dtimes[k]
            k += 1
            result[k] = result[k-1]
        end
        result[k] -= (t - tprev) * agg.sumrate

        log_jump_prob = 0.0
        if i != 0
            @inbounds gid = agg.ridtogroup[i]
            gid != 0 && @inbounds log_jump_prob = log(agg.gsums[gid])
        end
        result[k] += log_jump_prob

        agg = @set agg.weight = agg.weight - (t - agg.tprev) * agg.sumrate + log_jump_prob
        agg = @set agg.tprev = t
        agg, k
    end

    result
end

cumulative_logpdf(dist::TrajectoryDistribution, trajectory, dtimes::AbstractVector) = cumulative_logpdf!(zeros(length(dtimes)), dist, trajectory, dtimes)

@inline @fastmath function evalrxrate(speciesvec::AbstractVector{T}, rxidx::Int64, rs::ReactionSet) where {T}
    val = Float64(1.0)
    @inbounds for specstoch in rs.rstoich[rxidx]
        @inbounds specpop = speciesvec[specstoch[1]]
        val *= Float64(specpop)
        @inbounds for k = 2:specstoch[2]
            specpop -= one(specpop)
            val *= Float64(specpop)
        end
    end

    @inbounds val * rs.rates[rxidx]
end

@inline function update_reaction_rate!(aggregator::AbstractJumpRateAggregator, u::AbstractVector, reactions, rx::Int)
    @inbounds gid = aggregator.ridtogroup[rx]
    if gid != 0
        rate = evalrxrate(u, rx, reactions)
        @inbounds aggregator.gsums[gid] += -aggregator.rates[rx] + rate
        @inbounds aggregator.rates[rx] = rate
    end
end

@fastmath function update_rates(aggregator::DirectAggregator, (u, t, i), reactions::ReactionSet)
    for rx in eachindex(reactions.rates)
        update_reaction_rate!(aggregator, u, reactions, rx)
    end
    aggregator = @set aggregator.sumrate = sum(aggregator.gsums)
end

@fastmath function update_rates(aggregator::DepGraphAggregator, (u, t, i), reactions::ReactionSet)
    if aggregator.prev_reaction == 0
        for rx in eachindex(reactions.rates)
            update_reaction_rate!(aggregator, u, reactions, rx)
        end
    else
        @inbounds for rx in aggregator.depgraph[aggregator.prev_reaction]
            update_reaction_rate!(aggregator, u, reactions, rx)
        end
    end
    aggregator = @set aggregator.sumrate = sum(aggregator.gsums)
    aggregator = @set aggregator.prev_reaction = i
end


@with_kw mutable struct Domain{FT}
    # Something CUDA
    max_streams::IndexT
    streams::Nothing

    # Elem-centered
    matElemlist::VD{IndexT} # material indexset
    nodelist::VD{IndexT}    # elemToNode connectivity

    lxim::VD{IndexT}         # element connectivity through face g
    lxip::VD{IndexT}
    letam::VD{IndexT}
    letap::VD{IndexT}
    lzetam::VD{IndexT}
    lzetap::VD{IndexT}

    elemBC::VD{Int}         # elem face symm/free-surf flag g

    e::VD{FT}             # energy g
    d_e::VD{FT}           # change in energy g

    p::VD{FT}             # pressure g

    q::VD{FT}             # q g
    ql::VD{FT}            # linear term for q g
    qq::VD{FT}            # quadratic term for q g

    v::VD{FT}             # relative volume g

    volo::VD{FT}          # reference volume g
    delv::VD{FT}          # m_vnew - m_v g
    vdov::VD{FT}          # volume derivative over volume g

    arealg::VD{FT}        # char length of an element g

    ss::VD{FT}            # "sound speed" g

    elemMass::VD{FT}      # mass g

    # TODO: Not making distinction between pointers to arrays down here

    vnew::VD{FT}          # new relative volume -- temporary g

    delv_xi::VD{FT}       # velocity gradient -- temporary g
    delv_eta::VD{FT}
    delv_zeta::VD{FT}

    delx_xi::VD{FT}       # coordinate gradient -- temporary g
    delx_eta::VD{FT}
    delx_zeta::VD{FT}

    dxx::VD{FT}           # principal strains -- temporary g
    dyy::VD{FT}
    dzz::VD{FT}

    # Node-centered g

    x::VD{FT}             # coordinates g
    y::VD{FT}
    z::VD{FT}

    xd::VD{FT}            # velocities g
    yd::VD{FT}
    zd::VD{FT}

    xdd::VD{FT}           # accelerations g
    ydd::VD{FT}
    zdd::VD{FT}

    fx::VD{FT}            # forces g
    fy::VD{FT}
    fz::VD{FT}

    dfx::VD{FT}          # AD of the forces g
    dfy::VD{FT}
    dfz::VD{FT}

    nodalMass::VD{FT}     # mass g
    h_nodalMass::Vector{FT}     # mass - host g

    # device pointers for comms g
    # TODO: Not sure how to store the pointers or if even necessary
    # Real_t *d_delv_xi       # velocity gradient -- temporary g
    # Real_t *d_delv_eta
    # Real_t *d_delv_zeta

    # Real_t *d_x             # coordinates g
    # Real_t *d_y
    # Real_t *d_z

    # Real_t *d_xd            # velocities g
    # Real_t *d_yd
    # Real_t *d_zd

    # Real_t *d_fx            # forces g
    # Real_t *d_fy
    # Real_t *d_fz
    # Boundary nodesets

    symmX::VD{IndexT}      # symmetry plane nodesets
    symmY::VD{IndexT}
    symmZ::VD{IndexT}

    nodeElemCount::VD{Int}
    nodeElemStart::VD{Int}
    nodeElemCornerList::VD{IndexT}

    # Parameters

    dtfixed::FT                # fixed time increment g
    deltatimemultlb::FT
    deltatimemultub::FT
    stoptime::FT               # end time for simulation g
    dtmax::FT                  # maximum allowable time increment g
    cycle::Int                 # iteration count for simulation g

    dthydro_h::FT             # hydro time constraint g
    d_dthydro_h::FT           # AD change of the hydro time constraint g
    dtcourant_h::FT           # courant time constraint g
    d_dtcourant_h::FT         # AD of the courant time constraint g
    bad_q_h::IndexT              # flag to indicate Q error g
    bad_vol_h::IndexT            # flag to indicate volume error g

    # cuda Events to indicate completion of certain kernels
    # TODO Will check later how this works with KA
    # cudaEvent_t time_constraint_computed;
    time_h::FT                # current time g
    deltatime_h::FT           # variable time increment g

    u_cut::FT                 # velocity tolerance g
    hgcoef::FT                # hourglass control g
    qstop::FT                 # excessive q indicator g
    monoq_max_slope::FT
    monoq_limiter_mult::FT
    e_cut::FT                 # energy tolerance g
    p_cut::FT                 # pressure tolerance g
    ss4o3::FT
    q_cut::FT                 # q tolerance g
    v_cut::FT                 # relative volume tolerance g
    qlc_monoq::FT             # linear term coef for q g
    qqc_monoq::FT             # quadratic term coef for q g
    qqc::FT
    eosvmax::FT
    eosvmin::FT
    pmin::FT                  # pressure floor g
    emin::FT                  # energy floor g
    dvovmax::FT               # maximum allowable volume change g
    refdens::FT               # reference density g

    m_colLoc::IndexT
    m_rowLoc::IndexT
    m_planeLoc::IndexT
    m_tp::IndexT
    sizeX::IndexT
    sizeY::IndexT
    sizeZ::IndexT
    maxPlaneSize::IndexT
    maxEdgeSize::IndexT

    numElem::IndexT
    # padded_numElem::IndexT

    numNode::IndexT
    # padded_numNode::IndexT

    numSymmX::IndexT
    numSymmY::IndexT
    numSymmZ::IndexT

    octantCorner::IndexT

    # Region information
    numReg::Int                     # number of regions (def:11)
    balance::Int                    # Load balance between regions of a domain (def: 1)
    cost::Int                       # imbalance cost (def: 1)
    regElemSize::Vector{Int}        # Size of region sets
    regCSR::VD{Int}                 # records the begining and end of each region
    regReps::VD{Int}                # records the rep number per region
    regNumList::VD{IndexT}         # Region number per domain element
    regElemlist::VD{IndexT}        # region indexset
    regSorted::VD{IndexT}          # keeps index of sorted regions

    # MPI-Related additional data

    # TODO I think we can handle this differently

    # IndexT m_numRanks;
    # IndexT& numRanks() { return m_numRanks ; }

    # # Used in setup
    # IndexT m_rowMin, m_rowMax;
    # IndexT m_colMin, m_colMax;
    # IndexT m_planeMin, m_planeMax ;

    # # Communication Work space
    # Real_t *commDataSend ;
    # Real_t *commDataRecv ;

    # Real_t *d_commDataSend ;
    # Real_t *d_commDataRecv ;

    # # Maximum number of block neighbors
    # MPI_Request recvRequest[26] ; # 6 faces + 12 edges + 8 corners
    # MPI_Request sendRequest[26] ; # 6 faces + 12 edges + 8 corners
end


#!/usr/bin/julia
#
# Plotting functions.
#
using RecipesBase;
using Plots: @layout, grid;
#using IJulia; # Required to prevent plotting window opening and crashing.

export check_path, gen_filename, format_plot_title_variable,
       GeneralPlotOptions, LanczosPlotOptions, SplitStepPlotOptions, SimPlot;

"""
Implementations contain mainly boolean objects specifying what should be
included within plots.

All implementations should at least implement all the properties of
`GeneralPlotOptions`.

See also: [`GeneralPlotOptions`](@ref).
"""
abstract type PlotOptions end;

"Plotting options for a simulation agnostic plot"
struct GeneralPlotOptions <: PlotOptions
  """
  If the real and imaginary parts of the wave function should be plotted in
  addition to the absolute square value.
  """
  components::Bool;

  """
  If the value of the wave function is less than the value of this property at
  the start or end of the plot range, then they will not be plotted.
  """
  extremaϵ::Number;

end
GeneralPlotOptions(;components::Bool=false, extremaϵ::Number=0.1) =
  GeneralPlotOptions(components, extremaϵ);

"Plotting options for split step simulations."
struct SplitStepPlotOptions <: PlotOptions
  components::Bool;
  extremaϵ::Number;
end
SplitStepPlotOptions(;components::Bool=false, extremaϵ::Number=0.1) =
  SplitStepPlotOptions(components, extremaϵ);

"Plotting options for Lanczos simulations."
struct LanczosPlotOptions <: PlotOptions
  components::Bool;
  extremaϵ::Number;
  "Plot basis functions."
  basis::Bool;
  "Plot eigenfunctions of the reduced Hamiltonian."
  eigenfct::Bool;
end

abstract type PlotType end;
struct SimPlot<:PlotType end;

"""
    check_path(path::String)

Check the path given exists and create it if it does not.
"""
function check_path(path::String)
  if (!ispath(path))
    mkpath(path)
  end
end

# Parameters common to all simulation methods.
function gen_filename(p::SimParams, grid::UniformGrid{S,1}) where S
  return string(real(p.Δt) == 0 ? "ti" : "td",
    round(typeof(grid.Δ), grid.Δ, sigdigits=4), "d",
    round(typeof(grid.max), grid.max, sigdigits=4), "mx",
    round(typeof(p.Δt), abs(p.Δt), sigdigits=4), "dt-", p.steps, "s");
end
# Simulation variation method.
function gen_filename(var::Vararg{SimVariation})
  name = string();

  # Iterate through simulation variations and append parameters to filename.
  for i ∈ var
    if i isa AdaptiveΔt
      name *= string("-", round(i.ϵmax, sigdigits=4), "emx",
                     round(i.ϵmin, sigdigits=4), "emn",
                     round(i.Δtdown, sigdigits=4), "tdn",
                     round(i.Δtup, sigdigits=4), "tup",
                     round(i.Δtmax, sigdigits=4), "tmx",
                     round(i.Δtmin, sigdigits=4), "tmn",
                     i.maxeval, "mxev");
      if i.maxevallimit != nothing
        name *= string(i.maxevallimit, "mxevlm");
      end
    elseif i isa ConvergenceVariation
      name *= string("-", i.ϵ, "ecnv");
    end
  end

  return name;
end
# Common plotting options.
function gen_filename(plt::PlotOptions)
  if plt.components
    return string("pcmpnt");
  end
  return string();
end
# Simulation specific methods.
function gen_filename(sim::Lanczos, plt::LanczosPlotOptions)
  name = string("-sil-", sim.basisno, "b", sim.basissteps, "bs-" );

  if plt.basis name *= string("pb"); end
  if plt.eigenfct name *= string("pefct"); end

  name *= gen_filename(plt);

  return name;
end
function gen_filename(sim::SplitStep, plt::SplitStepPlotOptions)
  name = string("-ss-");

  name *= gen_filename(plt);

  return name
end
"""
    gen_filename(p::SimParams, grid::UniformGrid{S,1},
                      sim::SimMethod, plt::PlotOptions,
                      var::Union{Vararg{SimVariation},Nothing}=nothing) where S

Generate an appropriate filename reflecting the parameters used in the
simulation.
"""
function gen_filename(p::SimParams, grid::UniformGrid{S,1},
                      sim::SimMethod, plt::PlotOptions,
                      var::Vararg{Union{Nothing,SimVariation}}=nothing) where S
  name = gen_filename(p, grid);
  name *= gen_filename(sim, plt);
  if (var[1] != nothing) name *= gen_filename(var...); end

  return name;
end

"""
    format_plot_title_variable(variable::Number, decimalplaces::Integer)

Format `variable` with a fixed number of decimal places given by `decimalplaces`
such that it is in a suitable format to be printed in a GIF frame.
"""
function format_plot_title_variable(variable::Unitful.Quantity, decimalplaces::Integer)
  return string(format_plot_title_variable(ustrip(variable), decimalplaces),
                " ", unit(variable));
end
function format_plot_title_variable(variable::Number, decimalplaces::Integer)
  # Determine number of digits before the decimal place.
  predecimal = findfirst(".", string(variable));
  if predecimal === nothing predecimal = length(string(variable));
  else predecimal = predecimal[1]-1 end

  return string(rpad(round(variable, digits=decimalplaces),
                     predecimal+decimalplaces+1, '0'));
end

"""
    get_effective_minmax(v::AbstractVector, ϵ::Float64=0.1)

Return the minimum and maximum indices of `v` below and above which the value of
`v` does not exceed `ϵ`.
"""
function get_effective_minmax(v::AbstractVector, ϵ::Number=0.1)
  len = length(v);
  min = 1;
  max = len;
  # Find the first and last indices of v where v > ϵ.
  for i ∈ 1:len
    if abs(v[i]) > ϵ
      min = i;
      break;
    end;
  end
  for i ∈ len:-1:1
    if abs(v[i]) > ϵ
      max = i;
      break;
    end
  end

  return (min, max);
end
function get_effective_minmax(v::Vararg{AbstractVector}; ϵ=0.1)
  len = length(v);
  mins = Vector{Int}(undef, len);
  maxs = Vector{Int}(undef, len);

  for i ∈ 1:len
    (mins[i], maxs[i]) = get_effective_minmax(v[i], ϵ);
  end

  return (mins, maxs);
end
"""
    get_effective_minmax(m::AbstractMatrix, ϵ::Number=0.1)

Treat the columns of the matrix `m` as vectors.
"""
function get_effective_minmax(m::AbstractMatrix, ϵ::Number=0.1)
  numv = size(m)[2];
  mins = Vector{Int}(undef, numv);
  maxs = Vector{Int}(undef, numv);

  for i ∈ 1:numv
    (mins[i], maxs[i]) = get_effective_minmax(m[:,i], ϵ);
  end

  return (mins, maxs);
end

"""
    get_global_effective_minmax(v::Vararg{AbstractVector}; ϵ=0.1)

Returns the maximum and minimum values found by the respective
`get_effective_minmax()` function.

See also: [`get_effective_minmax`](@ref)
"""
function get_global_effective_minmax(v::Vararg{AbstractVector}; ϵ=0.1)
  (xmins, xmaxs) = get_effective_minmax(v..., ϵ=ϵ);
  return (minimum(xmins), maximum(xmaxs));
end
function get_global_effective_minmax(m::AbstractMatrix, ϵ::Number=0.1)
  (xmins, xmaxs) = get_effective_minmax(m, ϵ);
  return (minimum(xmins), maximum(xmaxs));
end

"""
    get_ranged_minmax(min::Int, max::Int, v::Vararg{AbstractVector})

Find the minimum and maximum of the `v` in the index range `min:max`.

Note: values of the `v` should be real to allow size comparison.
"""
function get_ranged_minmax(min::Int, max::Int, v::Vararg{AbstractVector})
  ymin = 0.0;
  ymax = 0.0;
  for i ∈ 2:numψ
    w = v[i][xmin:xmax];
    t = minimum(w);
    if (t < ymin || i == 1) ymin = t; end
    t = maxmum(w);
    if (t > ymmax || i ==1) ymax = t; end
  end

  return (ymin, ymax);
end

"""
    subscript(i::Integer)

Generate a string with a subscript number `i`. From Stack Overflow user
improbable https://stackoverflow.com/a/46709534.
"""
subscript(i::Integer) = i<0 ? error("$i is negative") : join('₀'+d for d in reverse(digits(i)));

"""
    generate_wavefunction_labels(num::Int)

Generate labels to identify `num` wave functions.
"""
function generate_wavefunction_labels(num::Int)
  labels = Vector{String}(undef, num);
  if (num == 1) labels[1] = "|ψ|²";
  else
    for i ∈ 1:num labels[i] = string("|ψ", subscript(i), "|²"); end
  end
  return reshape(labels, 1, num);
end

"Set appropriate axis labels"
@recipe function f(::Type{O}, x::Vector, y::Array) where
                   {T,N,S<:ConfigurationSpace,O<:WaveObject{T,N,S,1}}
  xguide --> "r";
  yguide --> "Probability density";
  (x, y)
end
@recipe function f(::Type{O}, x::Vector, y::Array) where
                   {T,N,S<:MomentumSpace,O<:WaveObject{T,N,S,1}}
  xguide --> "k";
  yguide --> "Probability density";
  (x, y)
end

"""
    vectors_to_matrix(v::Tuple{Vararg{Vector}}, min::Int=1, max::Int=0)

Construct a matrix from the vectors `v` between the indices `min` and `max`.
"""
function vectors_to_matrix(v::Tuple{Vararg{AbstractVector}}, min::Int=1, max::Int=0)
  num = length(v);
  if (max == 0) max = num; end

  m = Matrix{Number}(undef, max - min + 1, num);
  for i ∈ 1:num m[:,i] = abs2.(v[i][min:max]); end

  return m;
end

"Plot absolute square of wave function across relevant range."
@recipe function f(grid::Grid1D, ψ::Vararg{WaveFunction{T,1}};
                   opts=GeneralPlotOptions()) where T
  (opts, grid, ψ...);
end
# Keyword arguments are not dispatched properly in the recipe chain,
# so define this alternative order of the arguments to implement the function.
# 2020/12/01
@recipe function f(opts::PlotOptions, grid::Grid1D,
                   ψ::Vararg{WaveFunction{T,1}}) where T
  label --> generate_wavefunction_labels(length(ψ));

  (min, max) = get_global_effective_minmax(ψ..., ϵ=opts.extremaϵ);

  (typeof(ψ[1]), grid[min:max], vectors_to_matrix(ψ, min, max));
end

"""
Plot potential function and absolute square of wave function across interesting
wave function range.
"""
@recipe function f(grid::Grid1D, H::SeparableHamiltonian,
                   ψ::Vararg{WaveFunction{T,1}};
                   opts=GeneralPlotOptions()) where T
  (opts, grid, H, ψ...);
end
# Workaround for same issue as above.
@recipe function f(opts::PlotOptions, grid::Grid1D, H::SeparableHamiltonian,
                   ψ::Vararg{WaveFunction{T,1}}) where T
  label --> hcat(["Potential"], generate_wavefunction_labels(length(ψ)));

  (min, max) = get_global_effective_minmax(ψ..., ϵ=opts.extremaϵ);

  (typeof(ψ[1]), grid[min:max], hcat(real.(H.V[min:max]),
                                     vectors_to_matrix(ψ, min, max)));
end

"""
Plot absolute square of wave function components of wave matrix across relevant
range.
"""
@recipe function f(grid::G, ϕ::WaveMatrix{T,S,G},
                   opts::PlotOptions=GeneralPlotOptions()) where {T,S,G}
  label --> generate_wavefunction_labels(size(ϕ)[2]);

  (min, max) = get_global_effective_minmax(ϕ, opts.extremaϵ);

  (typeof(ϕ), grid[min:max], abs2.(ϕ[min:max,:]))
end

"""
Plot potential function and absolute square of wave function components of wave
matrix across relevant range.
"""
@recipe function f(grid::G, H::SeparableHamiltonian, ϕ::WaveMatrix{T,S,G},
                   opts::PlotOptions=GeneralPlotOptions()) where {T,S,G}
  label --> hcat(["Potential"], generate_wavefunction_labels(size(ϕ)[2]));

  (min, max) = get_global_effective_minmax(ϕ, opts.extremaϵ);

  (typeof(ϕ), grid[min:max], hcat(real.(H.V[min:max]),
                                  abs2.(ϕ[min:max,:])))
end

"Plot potential functions included in the passed Hamiltonians."
@recipe function f(grid::Grid1D, H::Vararg{SeparableHamiltonian})
  potentials = Vector{Vector{real(eltype(H[1]))}}(undef, length(H));
  for i ∈ 1:length(H) potentials[i] = real.(H[i].V); end
  (grid, potentials)
end

"Plot real and imaginary parts of wave functions alongside main plot."
@recipe function f(p::SimPlot, opts::PlotOptions, g::G, ψ::WaveFunction{T,1,S,G},
                   H::Union{SeparableHamiltonian,Nothing}=nothing) where {T,S,G}
  size := (800, 450);
  dpi := 300;

  if opts.components
    #layout := @layout [ a{0.6w} grid(2,1) ]
    layout := @layout [ a{0.5w} b ]

    @series begin
      subplot := 2;
      title := "Wave function decomposition";
      ylims := (-1,1);
      realψ = real.(ψ);
      (remin, remax) = get_effective_minmax(realψ, opts.extremaϵ);
      imagψ = imag.(ψ);
      (immin, immax) = get_effective_minmax(imagψ, opts.extremaϵ);
      min = Base.min(immin, remin);
      max = Base.max(immax, remax);
      label := hcat(["Re(ψ)"], ["Im(ψ)"]);
      (g[min:max], hcat(realψ[min:max], imagψ[min:max]))
    end
    # Allow negative y values on subplots.
    #=
    try
      subylims = ylims();
      if subylims[1] >= 0
        subylims = (-subylims[2], subylims[2]);
      end
    catch
    end=#

    #=@series begin
      subplot := 2;
      title := "Re(ψ)";
      ylims := (-1,1);
      realψ = real.(ψ);
      (min, max) = get_effective_minmax(realψ, opts.extremaϵ);
      (g[min:max], realψ[min:max])
    end

    @series begin
      subplot := 3;
      title := "Im(ψ)";
      ylims := (-1,1);
      imagψ = imag.(ψ);
      (min, max) = get_effective_minmax(imagψ, opts.extremaϵ);
      (g[min:max], imagψ[min:max])
    end=#
  end

  @series begin
    subplot := 1;
    if (H === nothing) return (opts, g, ψ);
    else return (opts, g, H, ψ); end
  end
end

"Plot Krylov basis and reduced Hamiltonian eigenfunctions from Lanczos method."
@recipe function f(p::SimPlot, opts::LanczosPlotOptions, g::G, ψ::WaveFunction{T,1,S,G},
                   Φ::KrylovBasis{T,S,G}, H::Union{SeparableHamiltonian,Nothing}=nothing) where {T,S,G}
  if opts.components && opts.basis && opts.eigenfct
    layout := @layout [ a{0.7w} grid(2,1); grid(1,2){0.4h}]
    size := (1400, 750);
  elseif !opts.components && opts.basis && opts.eigenfct
    layout := @layout [ a; grid(1,2)]
    size := (1400, 750);
  elseif !opts.components && !opts.basis && !opts.eigenfct
    size := (1400, 750);
  elseif opts.components && !opts.basis && !opts.eigenfct
    layout := @layout [ a{0.6w} grid(2,1) ]
    size := (1400, 750);
  else
    error("the requested plotting format has not yet been implemented.");
  end

  if opts.components
    @series begin
      subplot := 2;
      title := "Re(ψ)";
      ylims := (-1,1);
      realψ = real.(ψ);
      (min, max) = get_effective_minmax(realψ, opts.extremaϵ);
      (g[min:max], realψ[min:max])
    end
    @series begin
      subplot := 3;
      title := "Im(ψ)";
      ylims := (-1,1);
      imagψ = imag.(ψ);
      (min, max) = get_effective_minmax(imagψ, opts.extremaϵ);
      (g[min:max], imagψ[min:max])
    end
  end
  if opts.basis && opts.eigenfct
    # Plot basis functions.
    @series begin
      if !opts.components subplot := 2;
      else subplot := 4; end
      title := "Krylov Basis"
      (g, Φ, opts)
    end

    # Plot eigenfunctions.
    @series begin
      if !opts.components subplot := 3;
      else subplot := 5; end

      # Label the eigenstates by their energies.
      len = length(Φ.λ);
      labels = Vector{String}(undef, len)
      for i ∈ 1:len labels[i] = string("|ψ|² where E = ", round(Φ.λ[i], sigdigits=4)); end
      label := reshape(labels, 1, len);

      title := "Reduced Hamiltonian Eigenstates"
      (g, KrylovBasis{T,S,G}(Φ*Φ.v, Φ.λ, Φ.v, Φ.Δ), opts)
    end
  end

  @series begin
    subplot := 1;
    if (H === nothing) return (opts, g, ψ);
    else return (opts, g, H, ψ); end
  end
end


#=
This file is auto-generated. Do not edit.
=#
"""
    mutable struct AVRTypeI <: AVR
        Ka::Float64
        Ke::Float64
        Kf::Float64
        Ta::Float64
        Te::Float64
        Tf::Float64
        Tr::Float64
        Vr_max::Float64
        Vr_min::Float64
        Ae::Float64
        Be::Float64
        V_ref::Float64
        ext::Dict{String, Any}
        states::Vector{Symbol}
        n_states::Int64
        internal::InfrastructureSystemsInternal
    end

Parameters of an Automatic Voltage Regulator Type I - Resembles IEEE Type DC1

# Arguments
- `Ka::Float64`: Amplifier Gain, validation range: (0, nothing)
- `Ke::Float64`: Field circuit integral deviation, validation range: (0, nothing)
- `Kf::Float64`: Stabilizer Gain in s * pu/pu, validation range: (0, nothing)
- `Ta::Float64`: Amplifier Time Constant in s, validation range: (0, nothing)
- `Te::Float64`: Field Circuit Time Constant in s, validation range: (0, nothing)
- `Tf::Float64`: Stabilizer Time Constant in s, validation range: (0, nothing)
- `Tr::Float64`: Voltage Measurement Time Constant in s, validation range: (0, nothing)
- `Vr_max::Float64`: Maximum regulator voltage in pu, validation range: (0, nothing)
- `Vr_min::Float64`: Minimum regulator voltage in pu, validation range: (0, nothing)
- `Ae::Float64`: 1st ceiling coefficient, validation range: (0, nothing)
- `Be::Float64`: 2nd ceiling coefficient, validation range: (0, nothing)
- `V_ref::Float64`: Reference Voltage Set-point, validation range: (0, nothing)
- `ext::Dict{String, Any}`
- `states::Vector{Symbol}`
- `n_states::Int64`
- `internal::InfrastructureSystemsInternal`: power system internal reference, do not modify
"""
mutable struct AVRTypeI <: AVR
    "Amplifier Gain"
    Ka::Float64
    "Field circuit integral deviation"
    Ke::Float64
    "Stabilizer Gain in s * pu/pu"
    Kf::Float64
    "Amplifier Time Constant in s"
    Ta::Float64
    "Field Circuit Time Constant in s"
    Te::Float64
    "Stabilizer Time Constant in s"
    Tf::Float64
    "Voltage Measurement Time Constant in s"
    Tr::Float64
    "Maximum regulator voltage in pu"
    Vr_max::Float64
    "Minimum regulator voltage in pu"
    Vr_min::Float64
    "1st ceiling coefficient"
    Ae::Float64
    "2nd ceiling coefficient"
    Be::Float64
    "Reference Voltage Set-point"
    V_ref::Float64
    ext::Dict{String, Any}
    states::Vector{Symbol}
    n_states::Int64
    "power system internal reference, do not modify"
    internal::InfrastructureSystemsInternal
end

function AVRTypeI(Ka, Ke, Kf, Ta, Te, Tf, Tr, Vr_max, Vr_min, Ae, Be, V_ref=1.0, ext=Dict{String, Any}(), )
    AVRTypeI(Ka, Ke, Kf, Ta, Te, Tf, Tr, Vr_max, Vr_min, Ae, Be, V_ref, ext, [:Vf, :Vr1, :Vr2, :Vm], 4, InfrastructureSystemsInternal(), )
end

function AVRTypeI(; Ka, Ke, Kf, Ta, Te, Tf, Tr, Vr_max, Vr_min, Ae, Be, V_ref=1.0, ext=Dict{String, Any}(), )
    AVRTypeI(Ka, Ke, Kf, Ta, Te, Tf, Tr, Vr_max, Vr_min, Ae, Be, V_ref, ext, )
end

# Constructor for demo purposes; non-functional.
function AVRTypeI(::Nothing)
    AVRTypeI(;
        Ka=0,
        Ke=0,
        Kf=0,
        Ta=0,
        Te=0,
        Tf=0,
        Tr=0,
        Vr_max=0,
        Vr_min=0,
        Ae=0,
        Be=0,
        V_ref=0,
        ext=Dict{String, Any}(),
    )
end

"""Get AVRTypeI Ka."""
get_Ka(value::AVRTypeI) = value.Ka
"""Get AVRTypeI Ke."""
get_Ke(value::AVRTypeI) = value.Ke
"""Get AVRTypeI Kf."""
get_Kf(value::AVRTypeI) = value.Kf
"""Get AVRTypeI Ta."""
get_Ta(value::AVRTypeI) = value.Ta
"""Get AVRTypeI Te."""
get_Te(value::AVRTypeI) = value.Te
"""Get AVRTypeI Tf."""
get_Tf(value::AVRTypeI) = value.Tf
"""Get AVRTypeI Tr."""
get_Tr(value::AVRTypeI) = value.Tr
"""Get AVRTypeI Vr_max."""
get_Vr_max(value::AVRTypeI) = value.Vr_max
"""Get AVRTypeI Vr_min."""
get_Vr_min(value::AVRTypeI) = value.Vr_min
"""Get AVRTypeI Ae."""
get_Ae(value::AVRTypeI) = value.Ae
"""Get AVRTypeI Be."""
get_Be(value::AVRTypeI) = value.Be
"""Get AVRTypeI V_ref."""
get_V_ref(value::AVRTypeI) = value.V_ref
"""Get AVRTypeI ext."""
get_ext(value::AVRTypeI) = value.ext
"""Get AVRTypeI states."""
get_states(value::AVRTypeI) = value.states
"""Get AVRTypeI n_states."""
get_n_states(value::AVRTypeI) = value.n_states
"""Get AVRTypeI internal."""
get_internal(value::AVRTypeI) = value.internal

"""Set AVRTypeI Ka."""
set_Ka!(value::AVRTypeI, val::Float64) = value.Ka = val
"""Set AVRTypeI Ke."""
set_Ke!(value::AVRTypeI, val::Float64) = value.Ke = val
"""Set AVRTypeI Kf."""
set_Kf!(value::AVRTypeI, val::Float64) = value.Kf = val
"""Set AVRTypeI Ta."""
set_Ta!(value::AVRTypeI, val::Float64) = value.Ta = val
"""Set AVRTypeI Te."""
set_Te!(value::AVRTypeI, val::Float64) = value.Te = val
"""Set AVRTypeI Tf."""
set_Tf!(value::AVRTypeI, val::Float64) = value.Tf = val
"""Set AVRTypeI Tr."""
set_Tr!(value::AVRTypeI, val::Float64) = value.Tr = val
"""Set AVRTypeI Vr_max."""
set_Vr_max!(value::AVRTypeI, val::Float64) = value.Vr_max = val
"""Set AVRTypeI Vr_min."""
set_Vr_min!(value::AVRTypeI, val::Float64) = value.Vr_min = val
"""Set AVRTypeI Ae."""
set_Ae!(value::AVRTypeI, val::Float64) = value.Ae = val
"""Set AVRTypeI Be."""
set_Be!(value::AVRTypeI, val::Float64) = value.Be = val
"""Set AVRTypeI V_ref."""
set_V_ref!(value::AVRTypeI, val::Float64) = value.V_ref = val
"""Set AVRTypeI ext."""
set_ext!(value::AVRTypeI, val::Dict{String, Any}) = value.ext = val
"""Set AVRTypeI states."""
set_states!(value::AVRTypeI, val::Vector{Symbol}) = value.states = val
"""Set AVRTypeI n_states."""
set_n_states!(value::AVRTypeI, val::Int64) = value.n_states = val
"""Set AVRTypeI internal."""
set_internal!(value::AVRTypeI, val::InfrastructureSystemsInternal) = value.internal = val


"""
Calculate value, gradient, and hessian of the variational ELBO.
"""
module DeterministicVI

using Base.Threads: threadid, nthreads

import ..Config
using ..BivariateNormals: BivariateNormalDerivatives, BvnComponent,
                          GalaxySigmaDerivs, get_bvn_cov, eval_bvn_pdf!,
                          get_bvn_derivs!, transform_bvn_derivs!
using ..Model
using ..Model: SkyPatch, BvnBundle, populate_fsm!
import ..Celeste: Const, @aliasscope, @unroll_loop
import ..Infer
using ..SensitiveFloats
import ..Log
using ..Transform
import DataFrames
import Optim
import ForwardDiff.Dual
using StaticArrays
import Base: convert

export ElboArgs, generic_init_source, catalog_init_source, init_sources,
       VariationalParams, elbo, ElboIntermediateVariables

function init_thread_pool!(pool::Vector, create)
    if length(pool) != Base.Threads.nthreads()
        empty!(pool)
        for i = 1:Base.Threads.nthreads()
            push!(pool, create())
        end
    end
end

"""
Return a default-initialized VariationalParams instance.
"""
function generic_init_source(init_pos::Vector{Float64})
    ret = Vector{Float64}(length(CanonicalParams))
    ret[ids.is_star] = 0.5
    ret[ids.pos] = init_pos
    ret[ids.flux_loc] = log(2.0)
    ret[ids.flux_scale] = 1e-3
    ret[ids.gal_fracdev] = 0.5
    ret[ids.gal_ab] = 0.5
    ret[ids.gal_angle] = 0.0
    ret[ids.gal_scale] = 1.0
    ret[ids.k] = 1.0 / size(ids.k, 1)
    ret[ids.color_mean] = 0.0
    ret[ids.color_var] =  1e-2
    ret
end

"""
Return VariationalParams instance initialized form a catalog entry
"""
function catalog_init_source(ce::CatalogEntry; max_gal_scale=Inf)
    # TODO: sync this up with the transform bounds
    ret = generic_init_source(ce.pos)

    # TODO: don't do this thresholding for background sources,
    # just for sources that are being optimized
    ret[ids.is_star[1]] = ce.is_star ? 0.8: 0.2
    ret[ids.is_star[2]] = ce.is_star ? 0.2: 0.8

    ret[ids.flux_loc[1]] = log(max(0.1, ce.star_fluxes[3]))
    ret[ids.flux_loc[2]] = log(max(0.1, ce.gal_fluxes[3]))

    function get_color(color_var, color_mean)
        color_var > 0 && color_mean > 0 ? min(max(log(color_var / color_mean), -9.), 9.) :
            color_var > 0 && color_mean <= 0 ? 3.0 :
                color_var <= 0 && color_mean > 0 ? -3.0 : 0.0
    end

    function get_colors(raw_fluxes)
        [get_color(raw_fluxes[c+1], raw_fluxes[c]) for c in 1:4]
    end

    ret[ids.color_mean[:, 1]] = get_colors(ce.star_fluxes)
    ret[ids.color_mean[:, 2]] = get_colors(ce.gal_fluxes)

    ret[ids.gal_fracdev] = min(max(ce.gal_frac_dev, 0.015), 0.985)

    ret[ids.gal_ab] = ce.is_star ? .8 : min(max(ce.gal_ab, 0.015), 0.985)
    ret[ids.gal_angle] = ce.gal_angle
    ret[ids.gal_scale] = ce.is_star ? 0.2 : min(max_gal_scale, max(ce.gal_scale, 0.2))

    ret
end

function init_sources(target_sources::Vector{Int}, catalog::Vector{CatalogEntry})
    ret = Vector{Vector{Float64}}(length(catalog))
    for s in 1:length(catalog)
        ret[s] = catalog_init_source(catalog[s])
    end
    for s in target_sources
        ret[s][:] = generic_init_source(catalog[s].pos)
    end
    ret
end

include("deterministic_vi/elbo_args.jl")
include("deterministic_vi/elbo_kl.jl")
include("deterministic_vi/source_brightness.jl")
include("deterministic_vi/elbo_objective.jl")
include("deterministic_vi/ConstraintTransforms.jl")
include("deterministic_vi/ElboMaximize.jl")

"""
Infers one light source. This routine is intended to be called in parallel,
once per target light source.

Arguments:
    images: a collection of astronomical images
    neighbors: the other light sources near `entry`
    entry: the source to infer
"""
function infer_source(config::Config,
                      images::Vector{Image},
                      neighbors::Vector{CatalogEntry},
                      entry::CatalogEntry)
    if length(neighbors) > 100
        msg = string("object at RA, Dec = $(entry.pos) has an excessive",
                     "number ($(length(neighbors))) of neighbors")
        Log.warn(msg)
    end

    # It's a bit inefficient to call the next 5 lines every time we optimize_f.
    # But, as long as runtime is dominated by the call to maximize!, that
    # isn't a big deal.
    cat_local = vcat([entry], neighbors)
    vp = init_sources([1], cat_local)
    patches = Infer.get_sky_patches(images, cat_local)
    Infer.load_active_pixels!(config, images, patches)

    ea = ElboArgs(images, patches, [1])
    f_evals, max_f, max_x, nm_result = ElboMaximize.maximize!(ea, vp)
    return vp[1]
end

# legacy wrapper
function infer_source(images::Vector{Image},
                      neighbors::Vector{CatalogEntry},
                      entry::CatalogEntry)
    infer_source(Config(), images, neighbors, entry)
end

end


# low-level interface, meaning works with vectors and matrices, not data frames and estimation objects

"Structure to store test results"
struct BootTestResult{T}
  stat::T; stattype::String
  p::T; padj::T
  reps::Int64; repsfeas::Int64
  nbootclust::Int64
  dof::Int64; dof_r::T
  plot::Union{Nothing, NamedTuple{(:X, :p), Tuple{Tuple{Vararg{Vector{T}, N} where N},Vector{T}}}}
  peak::Union{Nothing, NamedTuple{(:X, :p), Tuple{Vector{T}, T}}}
  ci::Union{Nothing, Matrix{T}}
  dist::Matrix{T}
  b::Vector{T}
  V::Matrix{T}
  auxweights::Union{Nothing,Matrix{T}}
  # M::StrBootTest
end

"Return test statistic"
teststat(o::BootTestResult) = o.stat

"Return numerator of test statistic"
statnumer(o::BootTestResult) = o.b

"Return denominator of test statistic"
statvar(o::BootTestResult) = o.V

"""Return type of test statistic: "t", "z", "F", or "χ²" """
stattype(o::BootTestResult) = o.stattype

"Return p value"
p(o::BootTestResult) = o.p

"Return p value after multiple-hypothesis adjustment, if any"
padj(o::BootTestResult) = o.padj

"Return requested number of replications"
reps(o::BootTestResult) = o.reps

"Return actual number of replications, subject to enumeration of Rademacher draws"
repsfeas(o::BootTestResult) = o.repsfeas

"Return number of bootstrapping clusters in test"
nbootclust(o::BootTestResult) = o.nbootclust

"Return degrees of freedom of test"
dof(o::BootTestResult) = o.dof

"Return residual degrees of freedom of test"
dof_r(o::BootTestResult) = o.dof_r

"""
Return data for confidence plot of test.
Return value is a 2-tuple with named entries `X` and `p` holding
the confidence sampling locations and p values respectively. `X` is in turn
a 1- or 2-tuple of vectors of sampling coordinates for each
dimension of the tested hypothesis.
"""
plotpoints(o::BootTestResult) = o.plot

"Return parameter value with peak p value in test"
peak(o::BootTestResult) = o.peak

"Return confidence interval matrix from test, one row per disjoint piece"
ci(o::BootTestResult) = o.ci

"Return bootstrap distribution of statistic or statistic numerator in bootstrap test"
dist(o::BootTestResult) = o.dist

"Return auxilliary weight matrix for wild bootstrap"
auxweights(o::BootTestResult) = o.auxweights

strint(x) = iszero(mod(x,1)) ? "$(Int64(x))" : "$x"
function Base.show(io::IO, o::BootTestResult{T}) where T
	s = stattype(o) * ( iszero(dof_r(o)) ? isone(dof(o)) ? " " : "(" * strint(dof(o)) * ")" :                                                     # z, χ²
	                                       isone(dof(o)) ? "(" * strint(dof_r(o)) * ")" : "(" * strint(dof(o)) * ", " * strint(dof_r(o)) * ")" )  # t, F
	Printf.@printf(io, "%s = %6.4f\n", s, teststat(o))
	Printf.@printf(io, "p%s = %6.4f\n", repeat(" ", length(s)-1), p(o))
	isdefined(o, :ci) && !isnothing(o.ci) && length(o.ci)>0 && print(io, "CI" * repeat(" ", length(s)-2) * " = $(round.(ci(o); sigdigits=4))\n")
end

# single entry point with arguments already converted to standardized types, to allow a smaller set of precompile() calls(?)
function __wildboottest(
	R::Matrix{T},
	r::Vector{T};
	resp::VecOrMat{T},
	predexog::Matrix{T},
	predendog::Matrix{T},
	inst::Matrix{T},
	R1::Matrix{T},
	r1::Vector{T},
	clustid::Matrix{Int64},
	nbootclustvar::Int64,
	nerrclustvar::Int64,
	issorted::Bool,
	hetrobust::Bool,
	nfe::Int64,
	feid::VecOrMat{Int64},
	fedfadj::Int64,
	obswt::VecOrMat{T},
	fweights::Bool,
	maxmatsize::Float16,
	ptype::Symbol,
	bootstrapc::Bool,
	liml::Bool,
	fuller::T,
	kappa::T,
	arubin::Bool,
	small::Bool,
	clusteradj::Bool,
	clustermin::Bool,
	scorebs::Bool,
	reps::Int64,
	imposenull::Bool,
	auxwttype::Symbol,
	rng::AbstractRNG,
	level::T,
	rtol::T,
	madjtype::Symbol,
	nH0::Int16,
	ml::Bool,
	scores::Matrix{T},
	beta::Vector{T},
	A::Symmetric{T,Matrix{T}},
	gridmin::VecOrMat{T},
	gridmax::VecOrMat{T},
	gridpoints::VecOrMat{T},
	diststat::Symbol,
	getci::Bool,
	getplot::Bool,
	getauxweights::Bool) where T

	M = StrBootTest{T}(R, r, R1, r1, resp, predexog, predendog, inst, obswt, fweights, liml, fuller, kappa, arubin,
	                   reps, auxwttype, rng, maxmatsize, ptype, imposenull, scorebs, !bootstrapc, clustid, nbootclustvar, nerrclustvar, issorted, hetrobust, small, clusteradj, clustermin,
	                   nfe, feid, fedfadj, level, rtol, madjtype, nH0, ml, beta, A, scores, getplot,
	                   gridmin, gridmax, gridpoints)

	if getplot || (level<1 && getci)
		plot!(M)
		plot = getplot & isdefined(M, :plotX) ? (X=Tuple(M.plotX), p=M.plotY) : nothing
		peak = M.peak
		ci = level<1 & getci ? M.ci : nothing
	else
		ci = plot = peak = nothing
	end

	padj = getp(M)  # trigger main (re)computation

	BootTestResult{T}(getstat(M),
	                  isone(nrows(R)) ? (small ? "t" : "z") : (small ? "F" : "χ²"),
	                  M.p, padj, M.B, M.BFeas, M.N✻, M.dof, M.dof_r, plot, peak, ci,
	                  getdist(M, diststat),
	                  getb(M), getV(M),
	                  getauxweights && reps>0 ? getv(M) : nothing #=, M=#)
end

vecconvert(T::DataType, X) = Vector(isa(X, AbstractArray) ? vec(    eltype(X)==T ? X : T.(X)) : X)
matconvert(T::DataType, X) = Matrix(isa(X, AbstractArray) ? reshape(eltype(X)==T ? X : T.(X), size(X,1), size(X,2)) : X)

function _wildboottest(T::DataType,
					  R::AbstractVecOrMat,
						r::AbstractVecOrMat;
					  resp::AbstractVecOrMat{<:Real},
					  predexog::AbstractVecOrMat{<:Real}=zeros(T,0,0),
					  predendog::AbstractVecOrMat{<:Real}=zeros(T,0,0),
					  inst::AbstractVecOrMat{<:Real}=zeros(T,0,0),
					  R1::AbstractVecOrMat=zeros(T,0,0),
						r1::AbstractVecOrMat=zeros(T,0),
					  hetrobust::Bool=true,
					  clustid::AbstractVecOrMat{<:Integer}=zeros(Int,0,0),  # bootstrap-only clust vars, then boot&err clust vars, then err-only clust vars
					  nbootclustvar::Integer=ncols(clustid),
					  nerrclustvar::Integer=nbootclustvar,
						issorted::Bool=false,
						nfe::Integer=0,
					  feid::AbstractVecOrMat{<:Integer}=Int8[],
					  fedfadj::Integer=length(feid)>0 ? -1 : 0,
					  obswt::AbstractVecOrMat{<:Real}=T[],
					  fweights::Bool=false,
					  maxmatsize::Number=0,
					  ptype::Symbol=:symmetric,
					  bootstrapc::Bool=false,
					  liml::Bool=false,
					  fuller::Number=0,
					  kappa::Number=NaN,
					  arubin::Bool=false,
					  small::Bool=true,
						clusteradj::Bool=small,
						clustermin::Bool=false,
					  scorebs::Bool=false,
					  reps::Integer=999,
					  imposenull::Bool=true,
					  auxwttype::Symbol=:rademacher,
					  rng::AbstractRNG=MersenneTwister(),
					  level::Number=.95,
					  rtol::Number=1e-6,
					  madjtype::Symbol=:none,
					  nH0::Integer=1,
					  ml::Bool=false,
					  scores::AbstractVecOrMat=Matrix{Float32}(undef,0,0),
					  beta::AbstractVecOrMat=T[],
					  A::AbstractMatrix=zeros(T,0,0),
					  gridmin::Union{VecOrMat{S},VecOrMat{Union{S,Missing}}} where S<:Number = T[],
					  gridmax::Union{VecOrMat{S},VecOrMat{Union{S,Missing}}} where S<:Number = T[],
					  gridpoints::Union{VecOrMat{S},VecOrMat{Union{S,Missing}}} where S<:Number = Int64[],
					  diststat::Symbol=:none,
					  getci::Bool=true,
					  getplot::Bool=getci,
					  getauxweights::Bool=false)

	nrows(R)>2 && (getplot = getci = false)

	@assert any(auxwttype .== (:rademacher, :mammen, :webb, :gamma, :normal)) "auxwttype shoud be :rademacher, :mammen, :webb, :gamma, or :normal"
	@assert any(ptype .==(:symmetric, :equaltail, :lower, :upper)) "ptype should be :symmetric, :equaltail, :lower, or :upper"
	@assert any(madjtype .== (:none, :bonferroni, :sidak)) "madjtype should be :none, :bonferroni, or :sidak"
	@assert any(diststat .== (:none, :t, :numer))
	@assert ml || ncols(resp)==1 "resp should have one column"
  @assert (length(predexog)==0 || nrows(predexog)==nrows(resp)) &&
	        (length(predendog)==0 || nrows(predendog)==nrows(resp)) &&
					(length(inst)==0 || nrows(inst)==nrows(resp)) "All data vectors/matrices must have same height"
	@assert ncols(inst) >= ncols(predendog) "Model has fewer instruments than instrumented variables"
  @assert length(feid)==0 || nrows(feid)==nrows(resp) "feid vector must have same height as data matrices"
	@assert ncols(feid)≤1 "feid should have one column"
  @assert length(clustid)==0 || nrows(clustid)==nrows(resp) "clustid must have same height as data matrices"
  @assert nrows(obswt)==0 || nrows(obswt)==nrows(resp) "obswt must have same height as data matrices"
	@assert ncols(obswt)≤1 "obswt must have one column"
  @assert nrows(R)==nrows(r) "R and r must have same height"
  @assert (ncols(R) == (ml ? nrows(beta) : ncols(predexog)+ncols(predendog)) && isone(ncols(r))) "Wrong number of columns in null specification"
  @assert nrows(R1)==nrows(r1) "R₁ and r₁ must have same height"
  @assert length(R1)==0 || ncols(R1)==ncols(predexog)+ncols(predendog) "Wrong number of columns in model constraint specification"
	@assert ncols(r)==1 "r should have one column"
	@assert length(R1)==0 || ncols(r1)==1 "r1 should have one column"
  @assert nbootclustvar ≤ ncols(clustid) "nbootclustvar > width of clustid"
  @assert nerrclustvar ≤ ncols(clustid) "nerrclustvar > width of clustid"
  @assert reps ≥ 0 "reps < 0"
  @assert level ≥ 0. && level≤1. "level must be in the range [0,1]"
  @assert rtol > 0. "rtol ≤ 0"
  @assert nH0 > 0 "nH0 ≤ 0"
	@assert !liml || (ncols(predendog)>0 && ncols(inst)>0) "For liml, non-empty predendog and inst arguments are needed"
	@assert fuller==0 || (ncols(predendog)>0 && ncols(inst)>0) "For Fuller liml, non-empty predendog and inst arguments are needed"
	@assert iszero(ncols(predendog)) || ncols(inst)>0 "predendog provided without inst"
	@assert !arubin || ncols(predendog)>0 "Anderson-Rubin test requested but predendog not provided"

	if getplot || getci
		@assert iszero(length(gridmin   )) || length(gridmin   )==nrows(R) "Length of gridmin doesn't match number of hypotheses being jointly tested"
		@assert iszero(length(gridmax   )) || length(gridmax   )==nrows(R) "Length of gridmax doesn't match number of hypotheses being jointly tested"
		@assert iszero(length(gridpoints)) || length(gridpoints)==nrows(R) "Length of gridpoints doesn't match number of hypotheses being jointly tested"
		@assert iszero(length(gridmin   )) || ncols(gridmin)   ==1 "gridmin should have one column"
		@assert iszero(length(gridmax   )) || ncols(gridmax)   ==1 "gridmax should have one column"
		@assert iszero(length(gridpoints)) || ncols(gridpoints)==1 "gridpoints should have one column"
	end

	_gridmin = Vector{T}(undef, length(gridmin))
	_gridmax = Vector{T}(undef, length(gridmax))
	_gridpoints = Vector{T}(undef, length(gridpoints))
	for i ∈ 1:length(gridmin)  # cumbersome loops because map() and list comprehensions mess up type inference(?!)
		_gridmin[i] = T(ismissing(gridmin[i]) ? NaN : gridmin[i])
	end
	for i ∈ 1:length(gridmax)
		_gridmax[i] = T(ismissing(gridmax[i]) ? NaN : gridmax[i])
	end
	for i ∈ 1:length(gridpoints)
		_gridpoints[i] = T(ismissing(gridpoints[i]) ? NaN : gridpoints[i])
	end

	__wildboottest(
		matconvert(T,R),
		vecconvert(T,r);
		resp=vecconvert(T,resp),
		predexog=matconvert(T,predexog),
		predendog=matconvert(T,predendog),
		inst=matconvert(T,inst),
		R1=matconvert(T,R1),
		r1=vecconvert(T,r1),
		clustid=matconvert(Int64,clustid),
		nbootclustvar=Int64(nbootclustvar),
		nerrclustvar=Int64(nerrclustvar),
		issorted,
		hetrobust,
		nfe,
		feid=vecconvert(Int64,feid),
		fedfadj,
		obswt=vecconvert(T,obswt),
		fweights,
		maxmatsize=Float16(maxmatsize),
		ptype,
		bootstrapc,
		liml,
		fuller=T(fuller),
		kappa=T(kappa),
		arubin,
		small,
		clusteradj,
		clustermin,
		scorebs,
		reps=Int64(reps),
		imposenull,
		auxwttype,
		rng,
		level=T(level),
		rtol=T(rtol),
		madjtype,
		nH0=Int16(nH0),
		ml,
		scores=matconvert(T,scores),
		beta=vecconvert(T,beta),
		A=Symmetric(matconvert(T,A)),
		gridmin=_gridmin,
		gridmax=_gridmax,
		gridpoints=_gridpoints,
		diststat,
		getci,
		getplot,
		getauxweights)
end

_wildboottest(T::DataType, R, r::Number; kwargs...) = _wildboottest(T, R, [r]; kwargs...)
_wildboottest(T::DataType, R::UniformScaling{Bool}, r; kwargs...) = _wildboottest(T, diagm(fill(T(R.λ),nrows(r))), r; kwargs...)
_wildboottest(R::Union{UniformScaling{Bool},AbstractVecOrMat}, r::Union{Number,AbstractVecOrMat}; kwargs...) = _wildboottest(Float64, R, r; kwargs...)

"""
wildboottest([T::DataType=Float64,] R::AbstractMatrix, r::AbstractVector;
             resp, <optional keyword arguments>) -> WildBootTests.BootTestResult

Function to perform wild-bootstrap-based hypothesis test

# Positional arguments
* `T::DataType`: data type for inputs, results, and computations: Float32 or Float64 (default)
* `R::AbstractMatrix` and `r::AbstractVector`: required matrix and vector expressing the null Rβ=r; see notes below

# Required keyword argument
* `resp::AbstractVector`: response/dependent variable (y or y₁ in Roodman et al. (2019))

# Optional keyword arguments
* `predexog::AbstractVecOrMat`: exogenous predictors, including constant term, if any (X/X₁)
* `predendog::AbstractVecOrMat`: endogenous predictors (Y₂)
* `inst::AbstractVecOrMat`: instruments (X₂)
* `R1::AbstractMatrix` and `r1::AbstractVector`: model constraints; same format as for `R` and `r`
* `clustid::AbstractVecOrMat{<:Integer}`: data vector/matrix of error and bootstrapping cluster identifiers; see Notes
* `nbootclustvar::Integer=1`: number of bootstrap-clustering variables
* `nerrclustvar::Integer=nbootclustvar`: number of error-clustering variables
* `issorted:Bool=false`: time-saving flag: data matrices are already sort by column types 2, then 3, then 1 (see notes)
* `hetrobust::Bool=true`: true unless errors are treated as iid
* `nfe::Integer=0`: number of fixed-effect groups; if 0 yet `feid` is provided, will be computed
* `feid::AbstractVector{<:Integer}`: data vector for one-way fixed effect group identifier
* `fedfadj::Integer=nfe`: degrees of freedom that fixed effects (if any) consume
* `obswt::AbstractVector=[]`: observation weight vector; default is equal weighting
* `fweights::Bool=false`: true for frequency weights
* `maxmatsize::Number`: maximum size of auxilliary weight matrix (v), in gigabytes
* `ptype::Symbol=:symmetric`: p value type (`:symmetric`, `:equaltail`, `:lower`, `:upper`)
* `bootstrapc::Bool=false`: true to request bootstrap-c instead of bootstrap-t
* `liml::Bool=false`: true for LIML or Fuller LIML
* `fuller::Number`: Fuller LIML factor
* `kappa::Number`: fixed κ for _k_-class estimation
* `arubin::Bool=false`: true for Anderson-Rubin test
* `small::Bool=true`: true to multiply test statistics by G/(G-1) × N/(N-k), where G, N, k are number of clusters, observations, and predictors
* `clusteradj::Bool=true`: false to drop G/(G-1) factor
* `clustermin::Bool=false``: for multiway clustering, true to base G/(G-1) factor for all clusterings ]on the smallest G across clusterings
* `scorebs::Bool=false`: true for score bootstrap instead of wild bootstrap
* `reps::Integer=999`: number of bootstrap replications; `reps` = 0 requests classical Rao (or Wald) test if `imposenull` = `true` (or `false`)
* `imposenull::Bool=true`: true to impose null
* `auxwttype::Symbol=:rademacher`: auxilliary weight type (`:rademacher`, `:mammen`, `:webb`, `:normal`, `:gamma`)
* `rng::AbstractRNG=MersenneTwister()`: randon number generator
* `level::Number=.95`: significance level (0-1)
* `rtol::Number=1e-6`: tolerance for ci bound determination
* `madjtype::Symbol=:none`: multiple hypothesis adjustment (`none`, `:bonferroni`, `:sidak`)
* `nH0::Integer=1`: number of hypotheses tested, including one being tested now
* `ml::Bool=false`: true for (nonlinear) ML estimation
* `scores::AbstractVecOrMat`: for ML, pre-computed scores
* `beta::AbstractVector`: for ML, parameter estimates
* `A::AbstractMatrix`: for ML, covariance estimates
* `gridmin`: vector of graph lower bounds; max length 2, `missing`/`NaN` entries ask wildboottest() to choose
* `gridmax`: vector of graph upper bounds; `missing`/`NaN` entries ask wildboottest() to choose
* `gridpoints`: vector of number of sampling points; `missing`/`NaN` entries ask wildboottest() to choose
* `diststat::Symbole=:none`: `:t` to save bootstrap distribution of t/z/F/χ² statistics; `:numer` to save numerators thereof
* `getci::Bool=true`: whether to return confidence interval
* `getplot::Bool=getci`: whether to generate plot data
* `getauxweights::Bool=false`: whether to save auxilliary weight matrix (v)

# Notes
`T`, `ptype`, `auxwttype`, `madjtype`, and `diststat` may also be strings. Examples: `"Float32"` and `"webb"`.

The columns of `R` in the statement of the null should correspond to those of the matrix [`predexog` `predendog`],
where `predendog` is non-empty only in regressions with instruments.

Order the columns of `clustid` this way:
1. Variables only used to define bootstrapping clusters, as in the subcluster bootstrap.
2. Variables used to define both bootstrapping and error clusters.
3. Variables only used to define error clusters.
`nbootclustvar` is then the number of columns of type 1 or 2; `nerrclustvar` is the number of columns of type 2 or 3. Typically `clustid` is a single column of type 2.

`wildboottest()` does not handle missing data values: all data and identifier matrices must
be restricted to the estimation sample.

"""
wildboottest(   R, r; kwargs...) = _wildboottest(                                                  R, r; Dict(a.first => isa(a.second, AbstractString) ? Meta.parse(a.second) : a.second for a ∈ kwargs)...)
wildboottest(T, R, r; kwargs...) = _wildboottest(isa(T, AbstractString) ? eval(Meta.parse(T)) : T, R, r; Dict(a.first => isa(a.second, AbstractString) ? Meta.parse(a.second) : a.second for a ∈ kwargs)...)


using AtomsBase
# Key functionality to integrate DFTK and AtomsBase

function construct_system(lattice::AbstractMatrix, atoms::Vector, positions::Vector, magnetic_moments::Vector)
    @assert length(atoms) == length(positions)
    @assert isempty(magnetic_moments) || length(magnetic_moments) == length(atoms)
    atomsbase_atoms = map(enumerate(atoms)) do (i, element)
        kwargs = Dict{Symbol, Any}(:potential => element, )
        if !isempty(magnetic_moments)
            kwargs[:magnetic_moment] = normalize_magnetic_moment(magnetic_moments[i])
        end
        if element isa ElementPsp
            kwargs[:pseudopotential] = element.psp.identifier
        end

        position = lattice * positions[i] * u"bohr"
        if atomic_symbol(element) == :X  # dummy element ... should solve this upstream
            Atom(:X, position; atomic_symbol=:X, atomic_number=0, atomic_mass=0u"u", kwargs...)
        else
            Atom(atomic_symbol(element), position; kwargs...)
        end
    end
    periodic_system(atomsbase_atoms, collect(eachcol(lattice)) * u"bohr")
end

function parse_system(system::AbstractSystem{D}) where {D}
    if !all(periodicity(system))
        error("DFTK only supports calculations with periodic boundary conditions.")
    end

    # Parse abstract system and return data required to construct model
    mtx = austrip.(reduce(hcat, bounding_box(system)))
    T = eltype(mtx)
    lattice = zeros(T, 3, 3)
    lattice[1:D, 1:D] .= mtx

    # Cache for instantiated pseudopotentials. This is done to ensure that identical
    # atoms are indistinguishable in memory, which is used in the Model constructor
    # to deduce the atom_groups.
    cached_pspelements = Dict{String, ElementPsp}()
    atoms = map(system) do atom
        if hasproperty(atom, :potential)
            atom.potential
        elseif hasproperty(atom, :pseudopotential)
            get!(cached_pspelements, atom.pseudopotential) do
                ElementPsp(atomic_symbol(atom); psp=load_psp(atom.pseudopotential))
            end
        else
            ElementCoulomb(atomic_symbol(atom))
        end
    end

    positions = map(system) do atom
        coordinate = zeros(T, 3)
        coordinate[1:D] = lattice[1:D, 1:D] \ T.(austrip.(position(atom)))
        Vec3{T}(coordinate)
    end

    magnetic_moments = map(system) do atom
        hasproperty(atom, :magnetic_moment) || return nothing
        getproperty(atom, :magnetic_moment)
    end
    if all(m -> isnothing(m) || iszero(m) || isempty(m), magnetic_moments)
        empty!(magnetic_moments)
    else
        magnetic_moments = normalize_magnetic_moment.(magnetic_moments)
    end

    # TODO Use system to determine n_electrons

    (; lattice, atoms, positions, magnetic_moments)
end

function _call_with_system(f, system::AbstractSystem, args...; kwargs...)
    @assert !(:magnetic_moments in keys(kwargs))
    parsed = parse_system(system)
    f(parsed.lattice, parsed.atoms, parsed.positions, args...;
      parsed.magnetic_moments, kwargs...)
end

"""
    Model(system::AbstractSystem; kwargs...)

AtomsBase-compatible Model constructor. Sets structural information (`atoms`, `positions`,
`lattice`, `n_electrons` etc.) from the passed `system`.
"""
Model(system::AbstractSystem; kwargs...) = _call_with_system(Model, system; kwargs...)

# Generate equivalent functions for AtomsBase
for fun in (:model_atomic, :model_DFT, :model_LDA, :model_PBE, :model_SCAN)
    @eval function $fun(system::AbstractSystem, args...; kwargs...)
        _call_with_system($fun, system, args...; kwargs...)
    end
end

# Extra methods to AtomsBase functions for DFTK data structures
"""
    atomic_system(model::DFTK.Model, magnetic_moments=[])

Construct an AtomsBase atomic system from a DFTK model and associated magnetic moments.
"""
function AtomsBase.atomic_system(model::Model, magnetic_moments=[])
    construct_system(model.lattice, model.atoms, model.positions, magnetic_moments)
end

"""
    periodic_system(model::DFTK.Model, magnetic_moments=[])

Construct an AtomsBase atomic system from a DFTK model and associated magnetic moments.
"""
function AtomsBase.periodic_system(model::Model, magnetic_moments=[])
    atomic_system(model, magnetic_moments)
end

AtomsBase.chemical_formula(model::Model) = chemical_formula(atomic_symbol.(model.atoms))


__precompile__(true)

module ColorVectorSpace

using ColorTypes, FixedPointNumbers, Compat

import Base: ==, +, -, *, /, .+, .-, .*, ./, ^, .^, <, ~
import Base: abs, abs2, clamp, convert, copy, div, eps, isfinite, isinf,
    isnan, isless, length, mapreduce, norm, one, promote_array_type,
    promote_op, promote_rule, zero, trunc, floor, round, ceil, bswap,
    mod, rem, atan2, hypot, max, min, varm, real, histrange

if VERSION < v"0.5.0-dev"
    import Base.nan
else
    export nan
end

# The unaryOps
import Base:      conj, sin, cos, tan, sinh, cosh, tanh,
                  asin, acos, atan, asinh, acosh, atanh,
                  sec, csc, cot, asec, acsc, acot,
                  sech, csch, coth, asech, acsch, acoth,
                  sinc, cosc, cosd, cotd, cscd, secd,
                  sind, tand, acosd, acotd, acscd, asecd,
                  asind, atand, rad2deg, deg2rad,
                  log, log2, log10, log1p, exponent, exp,
                  exp2, expm1, cbrt, sqrt, erf,
                  erfc, erfcx, erfi, dawson,
                  significand, lgamma,
                  gamma, lfact, frexp, modf, airy, airyai,
                  airyprime, airyaiprime, airybi, airybiprime,
                  besselj0, besselj1, bessely0, bessely1,
                  eta, zeta, digamma

typealias AbstractGray{T} Color{T,1}
typealias TransparentRGB{C<:AbstractRGB,T}   TransparentColor{C,T,4}
typealias TransparentGray{C<:AbstractGray,T} TransparentColor{C,T,2}
typealias TransparentRGBFloat{C<:AbstractRGB,T<:AbstractFloat} TransparentColor{C,T,4}
typealias TransparentGrayFloat{C<:AbstractGray,T<:AbstractFloat} TransparentColor{C,T,2}
typealias TransparentRGBUFixed{C<:AbstractRGB,T<:UFixed} TransparentColor{C,T,4}
typealias TransparentGrayUFixed{C<:AbstractGray,T<:UFixed} TransparentColor{C,T,2}

typealias MathTypes{T,C} Union{AbstractRGB{T},TransparentRGB{C,T},AbstractGray{T},TransparentGray{C,T}}

# convert(RGB{Float32}, NaN) doesn't and shouldn't work, so we need to reintroduce nan
if VERSION >= v"0.5.0-dev"
    nan{T<:AbstractFloat}(::Type{T}) = convert(T, NaN)
end
nan{C<:MathTypes}(::Type{C}) = _nan(eltype(C), C)
_nan{T<:AbstractFloat,C<:AbstractGray}(::Type{T}, ::Type{C}) = (x = convert(T, NaN); C(x))
_nan{T<:AbstractFloat,C<:TransparentGray}(::Type{T}, ::Type{C}) = (x = convert(T, NaN); C(x,x))
_nan{T<:AbstractFloat,C<:AbstractRGB}(::Type{T}, ::Type{C}) = (x = convert(T, NaN); C(x,x,x))
_nan{T<:AbstractFloat,C<:TransparentRGB}(::Type{T}, ::Type{C}) = (x = convert(T, NaN); C(x,x,x,x))

## Generic algorithms
mapreduce(f, op::Base.ShortCircuiting, a::MathTypes) = f(a)  # ambiguity
mapreduce(f, op, a::MathTypes) = f(a)

for f in (:trunc, :floor, :round, :ceil, :eps, :bswap)
    @eval $f{T}(g::Gray{T}) = Gray{T}($f(gray(g)))
    @eval Compat.@dep_vectorize_1arg Gray $f
end
eps{T}(::Type{Gray{T}}) = Gray(eps(T))
Compat.@dep_vectorize_1arg AbstractGray isfinite
Compat.@dep_vectorize_1arg AbstractGray isinf
Compat.@dep_vectorize_1arg AbstractGray isnan
Compat.@dep_vectorize_1arg AbstractGray abs
Compat.@dep_vectorize_1arg AbstractGray abs2
for f in (:trunc, :floor, :round, :ceil)
    @eval $f{T<:Integer}(::Type{T}, g::Gray) = Gray{T}($f(T, gray(g)))
end

for f in (:mod, :rem, :mod1)
    @eval $f(x::Gray, m::Gray) = Gray($f(gray(x), gray(m)))
end

# Real values are treated like grays
ColorTypes.gray(x::Real) = x

# Return types for arithmetic operations
multype{A,B}(::Type{A}, ::Type{B}) = coltype(typeof(zero(A)*zero(B)))
sumtype{A,B}(::Type{A}, ::Type{B}) = coltype(typeof(zero(A)+zero(B)))
divtype{A,B}(::Type{A}, ::Type{B}) = coltype(typeof(zero(A)/zero(B)))
powtype{A,B}(::Type{A}, ::Type{B}) = coltype(typeof(zero(A)^zero(B)))
multype(a::Colorant, b::Colorant) = multype(eltype(a),eltype(b))
sumtype(a::Colorant, b::Colorant) = sumtype(eltype(a),eltype(b))
divtype(a::Colorant, b::Colorant) = divtype(eltype(a),eltype(b))
powtype(a::Colorant, b::Colorant) = powtype(eltype(a),eltype(b))

coltype{T<:Fractional}(::Type{T}) = T
coltype{T}(::Type{T})             = Float64

# Scalar binary RGB operations require the same RGB type for each element,
# otherwise we don't know which to return
color_rettype{A<:AbstractRGB,B<:AbstractRGB}(::Type{A}, ::Type{B}) = _color_rettype(base_colorant_type(A), base_colorant_type(B))
color_rettype{A<:AbstractGray,B<:AbstractGray}(::Type{A}, ::Type{B}) = _color_rettype(base_colorant_type(A), base_colorant_type(B))
color_rettype{A<:TransparentRGB,B<:TransparentRGB}(::Type{A}, ::Type{B}) = _color_rettype(base_colorant_type(A), base_colorant_type(B))
color_rettype{A<:TransparentGray,B<:TransparentGray}(::Type{A}, ::Type{B}) = _color_rettype(base_colorant_type(A), base_colorant_type(B))
_color_rettype{A<:Colorant,B<:Colorant}(::Type{A}, ::Type{B}) = error("binary operation with $A and $B, return type is ambiguous")
_color_rettype{C<:Colorant}(::Type{C}, ::Type{C}) = C

color_rettype(c1::Colorant, c2::Colorant) = color_rettype(typeof(c1), typeof(c2))

## Math on Colors. These implementations encourage inlining and,
## for the case of UFixed types, nearly halve the number of multiplications (for RGB)

# Scalar RGB
copy(c::AbstractRGB) = c
(*)(f::Real, c::AbstractRGB) = base_colorant_type(c){multype(typeof(f),eltype(c))}(f*red(c), f*green(c), f*blue(c))
(*)(f::Real, c::TransparentRGB) = base_colorant_type(c){multype(typeof(f),eltype(c))}(f*red(c), f*green(c), f*blue(c), f*alpha(c))
function (*){T<:UFixed}(f::Real, c::AbstractRGB{T})
    fs = f*(1/reinterpret(one(T)))
    base_colorant_type(c){multype(typeof(f),T)}(fs*reinterpret(red(c)), fs*reinterpret(green(c)), fs*reinterpret(blue(c)))
end
function (*){T<:UFixed}(f::UFixed, c::AbstractRGB{T})
    fs = reinterpret(f)*(1/widen(reinterpret(one(T)))^2)
    base_colorant_type(c){multype(typeof(f),T)}(fs*reinterpret(red(c)), fs*reinterpret(green(c)), fs*reinterpret(blue(c)))
end
function (/){T<:UFixed}(c::AbstractRGB{T}, f::Real)
    fs = (one(f)/reinterpret(one(T)))/f
    base_colorant_type(c){divtype(typeof(f),T)}(fs*reinterpret(red(c)), fs*reinterpret(green(c)), fs*reinterpret(blue(c)))
end
function (/){T<:UFixed}(c::AbstractRGB{T}, f::Integer)
    fs = (1/reinterpret(one(T)))/f
    base_colorant_type(c){divtype(typeof(f),T)}(fs*reinterpret(red(c)), fs*reinterpret(green(c)), fs*reinterpret(blue(c)))
end
(+){S,T}(a::AbstractRGB{S}, b::AbstractRGB{T}) = color_rettype(a, b){sumtype(S,T)}(red(a)+red(b), green(a)+green(b), blue(a)+blue(b))
(-){S,T}(a::AbstractRGB{S}, b::AbstractRGB{T}) = color_rettype(a, b){sumtype(S,T)}(red(a)-red(b), green(a)-green(b), blue(a)-blue(b))
(+)(a::TransparentRGB, b::TransparentRGB) =
    color_rettype(a, b){sumtype(a,b)}(red(a)+red(b), green(a)+green(b), blue(a)+blue(b), alpha(a)+alpha(b))
(-)(a::TransparentRGB, b::TransparentRGB) =
    color_rettype(a, b){sumtype(a,b)}(red(a)-red(b), green(a)-green(b), blue(a)-blue(b), alpha(a)-alpha(b))
(*)(c::AbstractRGB, f::Real) = (*)(f, c)
(*)(c::TransparentRGB, f::Real) = (*)(f, c)
(.*)(f::Real, c::AbstractRGB) = (*)(f, c)
(.*)(f::Real, c::TransparentRGB) = (*)(f, c)
(.*)(c::AbstractRGB, f::Real) = (*)(f, c)
(.*)(c::TransparentRGB, f::Real) = (*)(f, c)
(/)(c::AbstractRGB, f::Real) = (one(f)/f)*c
(/)(c::TransparentRGB, f::Real) = (one(f)/f)*c
(/)(c::AbstractRGB, f::Integer) = (one(eltype(c))/f)*c
(/)(c::TransparentRGB, f::Integer) = (one(eltype(c))/f)*c
(./)(c::AbstractRGB, f::Real) = (/)(c, f)
(./)(c::TransparentRGB, f::Real) = (/)(c, f)

isfinite{T<:UFixed}(c::Colorant{T}) = true
isfinite{T<:AbstractFloat}(c::AbstractRGB{T}) = isfinite(red(c)) && isfinite(green(c)) && isfinite(blue(c))
isfinite(c::TransparentRGBFloat) = isfinite(red(c)) && isfinite(green(c)) && isfinite(blue(c)) && isfinite(alpha(c))
isnan{T<:UFixed}(c::Colorant{T}) = false
isnan{T<:AbstractFloat}(c::AbstractRGB{T}) = isnan(red(c)) || isnan(green(c)) || isnan(blue(c))
isnan(c::TransparentRGBFloat) = isnan(red(c)) || isnan(green(c)) || isnan(blue(c)) || isnan(alpha(c))
isinf{T<:UFixed}(c::Colorant{T}) = false
isinf{T<:AbstractFloat}(c::AbstractRGB{T}) = isinf(red(c)) || isinf(green(c)) || isinf(blue(c))
isinf(c::TransparentRGBFloat) = isinf(red(c)) || isinf(green(c)) || isinf(blue(c)) || isinf(alpha(c))
abs(c::AbstractRGB) = abs(red(c))+abs(green(c))+abs(blue(c)) # should this have a different name?
abs{T<:UFixed}(c::AbstractRGB{T}) = Float32(red(c))+Float32(green(c))+Float32(blue(c)) # should this have a different name?
abs(c::TransparentRGB) = abs(red(c))+abs(green(c))+abs(blue(c))+abs(alpha(c)) # should this have a different name?
abs{T<:UFixed}(c::TransparentRGB{T}) = Float32(red(c))+Float32(green(c))+Float32(blue(c))+Float32(alpha(c)) # should this have a different name?
abs2(c::AbstractRGB) = red(c)^2+green(c)^2+blue(c)^2
abs2{T<:UFixed}(c::AbstractRGB{T}) = Float32(red(c))^2+Float32(green(c))^2+Float32(blue(c))^2
abs2(c::TransparentRGB) = (ret = abs2(color(c)); ret + convert(typeof(ret), alpha(c))^2)
norm(c::AbstractRGB) = sqrt(abs2(c))
norm(c::TransparentRGB) = sqrt(abs2(c))

one{C<:AbstractRGB}(::Type{C})     = C(1,1,1)
one{C<:TransparentRGB}(::Type{C})  = C(1,1,1,1)
zero{C<:AbstractRGB}(::Type{C})    = C(0,0,0)
zero{C<:TransparentRGB}(::Type{C}) = C(0,0,0,0)
zero{C<:YCbCr}(::Type{C}) = C(0,0,0)
zero{C<:HSV}(::Type{C}) = C(0,0,0)
one(p::Colorant) = one(typeof(p))
zero(p::Colorant) = zero(typeof(p))

# Arrays
(+){CV<:AbstractRGB}(A::AbstractArray{CV}, b::AbstractRGB) = (.+)(A, b)
(+){CV<:AbstractRGB}(b::AbstractRGB, A::AbstractArray{CV}) = (.+)(b, A)
(-){CV<:AbstractRGB}(A::AbstractArray{CV}, b::AbstractRGB) = (.-)(A, b)
(-){CV<:AbstractRGB}(b::AbstractRGB, A::AbstractArray{CV}) = (.-)(b, A)
(*){T<:Number}(A::AbstractArray{T}, b::AbstractRGB) = A.*b
(*){T<:Number}(b::AbstractRGB, A::AbstractArray{T}) = A.*b
(.+){C<:AbstractRGB}(A::AbstractArray{C}, b::AbstractRGB) = plus(A, b)
(.+){C<:AbstractRGB}(b::AbstractRGB, A::AbstractArray{C}) = plus(b, A)
(.-){C<:AbstractRGB}(A::AbstractArray{C}, b::AbstractRGB) = minus(A, b)
(.-){C<:AbstractRGB}(b::AbstractRGB, A::AbstractArray{C}) = minus(b, A)
(.*){T<:Number}(A::AbstractArray{T}, b::AbstractRGB) = mul(A, b)
(.*){T<:Number}(b::AbstractRGB, A::AbstractArray{T}) = mul(b, A)

(+){CV<:TransparentRGB}(A::AbstractArray{CV}, b::TransparentRGB) = (.+)(A, b)
(+){CV<:TransparentRGB}(b::TransparentRGB, A::AbstractArray{CV}) = (.+)(b, A)
(-){CV<:TransparentRGB}(A::AbstractArray{CV}, b::TransparentRGB) = (.-)(A, b)
(-){CV<:TransparentRGB}(b::TransparentRGB, A::AbstractArray{CV}) = (.-)(b, A)
(*){T<:Number}(A::AbstractArray{T}, b::TransparentRGB) = A.*b
(*){T<:Number}(b::TransparentRGB, A::AbstractArray{T}) = A.*b
(.+){C<:TransparentRGB}(A::AbstractArray{C}, b::TransparentRGB) = plus(A, b)
(.+){C<:TransparentRGB}(b::TransparentRGB, A::AbstractArray{C}) = plus(b, A)
(.-){C<:TransparentRGB}(A::AbstractArray{C}, b::TransparentRGB) = minus(A, b)
(.-){C<:TransparentRGB}(b::TransparentRGB, A::AbstractArray{C}) = minus(b, A)
(.*){T<:Number}(A::AbstractArray{T}, b::TransparentRGB) = mul(A, b)
(.*){T<:Number}(b::TransparentRGB, A::AbstractArray{T}) = mul(b, A)

# Scalar Gray
copy(c::AbstractGray) = c
const unaryOps = (:~, :conj, :abs,
                  :sin, :cos, :tan, :sinh, :cosh, :tanh,
                  :asin, :acos, :atan, :asinh, :acosh, :atanh,
                  :sec, :csc, :cot, :asec, :acsc, :acot,
                  :sech, :csch, :coth, :asech, :acsch, :acoth,
                  :sinc, :cosc, :cosd, :cotd, :cscd, :secd,
                  :sind, :tand, :acosd, :acotd, :acscd, :asecd,
                  :asind, :atand, :rad2deg, :deg2rad,
                  :log, :log2, :log10, :log1p, :exponent, :exp,
                  :exp2, :expm1, :cbrt, :sqrt, :erf,
                  :erfc, :erfcx, :erfi, :dawson,
                  :significand, :lgamma,
                  :gamma, :lfact, :frexp, :modf, :airy, :airyai,
                  :airyprime, :airyaiprime, :airybi, :airybiprime,
                  :besselj0, :besselj1, :bessely0, :bessely1,
                  :eta, :zeta, :digamma)
for op in unaryOps
    @eval ($op)(c::AbstractGray) = $op(gray(c))
end

(*)(f::Real, c::AbstractGray) = base_colorant_type(c){multype(typeof(f),eltype(c))}(f*gray(c))
(*)(f::Real, c::TransparentGray) = base_colorant_type(c){multype(typeof(f),eltype(c))}(f*gray(c), f*alpha(c))
(*)(c::AbstractGray, f::Real) = (*)(f, c)
(.*)(f::Real, c::AbstractGray) = (*)(f, c)
(.*)(c::AbstractGray, f::Real) = (*)(f, c)
(*)(c::TransparentGray, f::Real) = (*)(f, c)
(.*)(f::Real, c::TransparentGray) = (*)(f, c)
(.*)(c::TransparentGray, f::Real) = (*)(f, c)
(/)(c::AbstractGray, f::Real) = (one(f)/f)*c
(/)(n::Number, c::AbstractGray) = n/gray(c)
(/)(c::TransparentGray, f::Real) = (one(f)/f)*c
(/)(c::AbstractGray, f::Integer) = (one(eltype(c))/f)*c
(/)(c::TransparentGray, f::Integer) = (one(eltype(c))/f)*c
(./)(c::AbstractGray, f::Real) = c/f
(./)(n::Number, c::AbstractGray) = n/gray(c)
(./)(c::TransparentGray, f::Real) = c/f
(+){S,T}(a::AbstractGray{S}, b::AbstractGray{T}) = color_rettype(a,b){sumtype(S,T)}(gray(a)+gray(b))
(+)(a::TransparentGray, b::TransparentGray) = color_rettype(a,b){sumtype(eltype(a),eltype(b))}(gray(a)+gray(b),alpha(a)+alpha(b))
(-){S,T}(a::AbstractGray{S}, b::AbstractGray{T}) = color_rettype(a,b){sumtype(S,T)}(gray(a)-gray(b))
(-)(a::TransparentGray, b::TransparentGray) = color_rettype(a,b){sumtype(eltype(a),eltype(b))}(gray(a)-gray(b),alpha(a)-alpha(b))
(*){S,T}(a::AbstractGray{S}, b::AbstractGray{T}) = color_rettype(a,b){multype(S,T)}(gray(a)*gray(b))
(^){S}(a::AbstractGray{S}, b::Integer) = base_colorant_type(a){powtype(S,Int)}(gray(a)^convert(Int,b))
(^){S}(a::AbstractGray{S}, b::Real) = base_colorant_type(a){powtype(S,typeof(b))}(gray(a)^b)
(.^){S}(a::AbstractGray{S}, b) = a^b
(+)(c::AbstractGray) = c
(+)(c::TransparentGray) = c
(-)(c::AbstractGray) = typeof(c)(-gray(c))
(-)(c::TransparentGray) = typeof(c)(-gray(c),-alpha(c))
(/)(a::AbstractGray, b::AbstractGray) = gray(a)/gray(b)
div(a::AbstractGray, b::AbstractGray) = div(gray(a), gray(b))
(+)(a::AbstractGray, b::Number) = gray(a)+b
(-)(a::AbstractGray, b::Number) = gray(a)-b
(+)(a::Number, b::AbstractGray) = a+gray(b)
(-)(a::Number, b::AbstractGray) = a-gray(b)
(.+)(a::AbstractGray, b::Number) = gray(a)+b
(.-)(a::AbstractGray, b::Number) = gray(a)-b
(.+)(a::Number, b::AbstractGray) = a+gray(b)
(.-)(a::Number, b::AbstractGray) = a-gray(b)
max{T<:AbstractGray}(a::T, b::T) = T(max(gray(a),gray(b)))
max(a::AbstractGray, b::AbstractGray) = max(promote(a,b)...)
max(a::Number, b::AbstractGray) = max(promote(a,b)...)
max(a::AbstractGray, b::Number) = max(promote(a,b)...)
min{T<:AbstractGray}(a::T, b::T) = T(min(gray(a),gray(b)))
min(a::AbstractGray, b::AbstractGray) = min(promote(a,b)...)
min(a::Number, b::AbstractGray) = min(promote(a,b)...)
min(a::AbstractGray, b::Number) = min(promote(a,b)...)

isfinite{T<:AbstractFloat}(c::AbstractGray{T}) = isfinite(gray(c))
isfinite(c::TransparentGrayFloat) = isfinite(gray(c)) && isfinite(alpha(c))
isnan{T<:AbstractFloat}(c::AbstractGray{T}) = isnan(gray(c))
isnan(c::TransparentGrayFloat) = isnan(gray(c)) && isnan(alpha(c))
isinf{T<:AbstractFloat}(c::AbstractGray{T}) = isinf(gray(c))
isinf(c::TransparentGrayFloat) = isinf(gray(c)) && isnan(alpha(c))
norm(c::AbstractGray) = abs(gray(c))
abs(c::TransparentGray) = abs(gray(c))+abs(alpha(c)) # should this have a different name?
abs(c::TransparentGrayUFixed) = Float32(gray(c)) + Float32(alpha(c)) # should this have a different name?
abs2(c::AbstractGray) = gray(c)^2
abs2{T<:UFixed}(c::AbstractGray{T}) = Float32(gray(c))^2
abs2(c::TransparentGray) = gray(c)^2+alpha(c)^2
abs2(c::TransparentGrayUFixed) = Float32(gray(c))^2 + Float32(alpha(c))^2
atan2(x::Gray, y::Gray) = atan2(convert(Real, x), convert(Real, y))
hypot(x::Gray, y::Gray) = hypot(convert(Real, x), convert(Real, y))
norm(c::TransparentGray) = sqrt(abs2(c))

(<)(g1::AbstractGray, g2::AbstractGray) = gray(g1) < gray(g2)
(<)(c::AbstractGray, r::Real) = gray(c) < r
(<)(r::Real, c::AbstractGray) = r < gray(c)
isless(g1::AbstractGray, g2::AbstractGray) = isless(gray(g1), gray(g2))
isless(c::AbstractGray, r::Real) = isless(gray(c), r)
isless(r::Real, c::AbstractGray) = isless(r, gray(c))
Base.isapprox(x::AbstractGray, y::AbstractGray; kwargs...) = isapprox(gray(x), gray(y); kwargs...)
Base.isapprox(x::TransparentGray, y::TransparentGray; kwargs...) = isapprox(gray(x), gray(y); kwargs...) && isapprox(alpha(x), alpha(y); kwargs...)
Base.isapprox(x::AbstractRGB, y::AbstractRGB; kwargs...) = isapprox(red(x), red(y); kwargs...) && isapprox(green(x), green(y); kwargs...) && isapprox(blue(x), blue(y); kwargs...)
Base.isapprox(x::TransparentRGB, y::TransparentRGB; kwargs...) = isapprox(alpha(x), alpha(y); kwargs...) && isapprox(red(x), red(y); kwargs...) && isapprox(green(x), green(y); kwargs...) && isapprox(blue(x), blue(y); kwargs...)

function Base.isapprox{Cx<:MathTypes,Cy<:MathTypes}(x::AbstractArray{Cx},
                                                    y::AbstractArray{Cy};
                                                    rtol::Real=Base.rtoldefault(eltype(Cx),eltype(Cy)),
                                                    atol::Real=0,
                                                    norm::Function=vecnorm)
    d = norm(x - y)
    if isfinite(d)
        return d <= atol + rtol*max(norm(x), norm(y))
    else
        # Fall back to a component-wise approximate comparison
        return all(ab -> isapprox(ab[1], ab[2]; rtol=rtol, atol=atol), zip(x, y))
    end
end

zero{C<:TransparentGray}(::Type{C}) = C(0,0)
zero{C<:Gray}(::Type{C}) = C(0)
one{C<:TransparentGray}(::Type{C}) = C(1,1)
one{C<:Gray}(::Type{C}) = C(1)

 # Arrays
(+){CV<:AbstractGray}(A::AbstractArray{CV}, b::AbstractGray) = (.+)(A, b)
(+){CV<:AbstractGray}(b::AbstractGray, A::AbstractArray{CV}) = (.+)(b, A)
(-){CV<:AbstractGray}(A::AbstractArray{CV}, b::AbstractGray) = (.-)(A, b)
(-){CV<:AbstractGray}(b::AbstractGray, A::AbstractArray{CV}) = (.-)(b, A)
(*){T<:Number}(A::AbstractArray{T}, b::AbstractGray) = A.*b
(*){T<:Number}(b::AbstractGray, A::AbstractArray{T}) = A.*b
(/){C<:AbstractGray}(A::AbstractArray{C}, b::AbstractGray) = A./b
(.+){C<:AbstractGray}(A::AbstractArray{C}, b::AbstractGray) = plus(A, b)
(.+){C<:AbstractGray}(b::AbstractGray, A::AbstractArray{C}) = plus(b, A)
(.-){C<:AbstractGray}(A::AbstractArray{C}, b::AbstractGray) = minus(A, b)
(.-){C<:AbstractGray}(b::AbstractGray, A::AbstractArray{C}) = minus(b, A)
(.*){T<:Number}(A::AbstractArray{T}, b::AbstractGray) = mul(A, b)
(.*){T<:Number}(b::AbstractGray, A::AbstractArray{T}) = mul(b, A)
(./){C<:AbstractGray}(A::AbstractArray{C}, b::AbstractGray) = divd(A, b)

Compat.@dep_vectorize_2arg Gray max
Compat.@dep_vectorize_2arg Gray min
for f in (:min, :max)
    @eval begin
        @deprecate($f{T<:Gray}(x::Number, y::AbstractArray{T}),
                   @compat $f.(x, y))
        @deprecate($f{T<:Gray}(x::AbstractArray{T}, y::Number),
                   @compat $f.(x, y))
    end
end

(+){CV<:TransparentGray}(A::AbstractArray{CV}, b::TransparentGray) = (.+)(A, b)
(+){CV<:TransparentGray}(b::TransparentGray, A::AbstractArray{CV}) = (.+)(b, A)
(-){CV<:TransparentGray}(A::AbstractArray{CV}, b::TransparentGray) = (.-)(A, b)
(-){CV<:TransparentGray}(b::TransparentGray, A::AbstractArray{CV}) = (.-)(b, A)
(*){T<:Number}(A::AbstractArray{T}, b::TransparentGray) = A.*b
(*){T<:Number}(b::TransparentGray, A::AbstractArray{T}) = A.*b
(.+){C<:TransparentGray}(A::AbstractArray{C}, b::TransparentGray) = plus(A, b)
(.+){C<:TransparentGray}(b::TransparentGray, A::AbstractArray{C}) = plus(b, A)
(.-){C<:TransparentGray}(A::AbstractArray{C}, b::TransparentGray) = minus(A, b)
(.-){C<:TransparentGray}(b::TransparentGray, A::AbstractArray{C}) = minus(b, A)
(.*){T<:Number}(A::AbstractArray{T}, b::TransparentGray) = mul(A, b)
(.*){T<:Number}(b::TransparentGray, A::AbstractArray{T}) = mul(b, A)

varm{C<:AbstractGray}(v::AbstractArray{C}, s::AbstractGray; corrected::Bool=true) =
        varm(map(gray,v),gray(s); corrected=corrected)
real{C<:AbstractGray}(::Type{C}) = real(eltype(C))

# Called plus/minus instead of plus/sub because `sub` already has a meaning!
function plus(A::AbstractArray, b::Colorant)
    bT = convert(eltype(A), b)
    out = similar(A)
    plus!(out, A, bT)
end
plus(b::Colorant, A::AbstractArray) = plus(A, b)
function minus(A::AbstractArray, b::Colorant)
    bT = convert(eltype(A), b)
    out = similar(A)
    minus!(out, A, bT)
end
function minus(b::Colorant, A::AbstractArray)
    bT = convert(eltype(A), b)
    out = similar(A)
    minus!(out, bT, A)
end
function mul{T<:Number}(A::AbstractArray{T}, b::Colorant)
    bT = typeof(b*one(T))
    out = similar(A, bT)
    mul!(out, A, b)
end
mul{T<:Number}(b::Colorant, A::AbstractArray{T}) = mul(A, b)
function divd{C<:AbstractGray}(A::AbstractArray{C}, b::AbstractGray)
    bT = typeof(zero(C)/b)
    out = similar(A, bT)
    div!(out, A, b)
end

for (func, op) in ((:plus!, :+),
                   (:minus!, :-),
                   (:mul!, :*),
                   (:div!, :/))
    @eval begin
        function $func{T,N}(out, A::AbstractArray{T,N}, b)
            Rout, RA = eachindex(out), eachindex(A)
            if Rout == RA
                for I in RA
                    @inbounds out[I] = $op(A[I], b)
                end
            else
                for (Iout, IA) in zip(Rout, RA)
                    @inbounds out[Iout] = $op(A[IA], b)
                end
            end
            out
        end
    end
end

# This needs separate implementation because we can take -b of unsigned types
function minus!{T,N}(out, b::Colorant, A::AbstractArray{T,N})
    Rout, RA = eachindex(out), eachindex(A)
    if Rout == RA
        for I in RA
            @inbounds out[I] = b - A[I]
        end
    else
        for (Iout, IA) in zip(Rout, RA)
            @inbounds out[Iout] = b - A[IA]
        end
    end
    out
end

#histrange for Gray type
Base.histrange{T}(v::AbstractArray{Gray{T}}, n::Integer) = histrange(convert(Array{Float32}, map(gray, v)), n)

# Promotions for reductions
if VERSION < v"0.5.0-dev+3701"
    Base.r_promote{T<:FixedPoint}(::Base.AddFun, c::MathTypes{T}) = convert(base_colorant_type(typeof(c)){Float64}, c)
    Base.r_promote{T<:FixedPoint}(::Base.MulFun, c::MathTypes{T}) = convert(base_colorant_type(typeof(c)){Float64}, c)
else
    Base.r_promote{T<:FixedPoint}(::typeof(+), c::MathTypes{T}) = convert(base_colorant_type(typeof(c)){Float64}, c)
    Base.r_promote{T<:FixedPoint}(::typeof(*), c::MathTypes{T}) = convert(base_colorant_type(typeof(c)){Float64}, c)
end

# To help type inference
promote_array_type{T<:Real,C<:MathTypes}(F, ::Type{T}, ::Type{C}) = base_colorant_type(C){Base.promote_array_type(F, T, eltype(C))}
promote_op{C<:MathTypes,T<:Real}(F, ::Type{C}, ::Type{T}) = typeof(F(zero(C), zero(T)))
promote_op{C<:MathTypes,T<:Real}(F, ::Type{T}, ::Type{C}) = typeof(F(zero(T), zero(C)))
promote_rule{C1<:Colorant,C2<:Colorant}(::Type{C1}, ::Type{C2}) = color_rettype(C1,C2){promote_type(eltype(C1), eltype(C2))}
promote_rule{T<:Real,C<:AbstractGray}(::Type{T}, ::Type{C}) = promote_type(T, eltype(C))

@deprecate sumsq abs2

end


############################################################################################
"""
$(SIGNATURES)
Return kernel that is used for jittering.

# Examples
```julia
```

"""
function jitterkernel(particles::SMCParticles, iter::Int64)
    return particles.kernel[iter]
end

############################################################################################
"""
$(SIGNATURES)
Resample particle ancestors, shuffle current particles and rejuvenate θ.

# Examples
```julia
```

"""
function resample!(
    _rng::Random.AbstractRNG,
    particles::SMCParticles,
    tune::SMCTune,
    data::D,
    temperature::F
) where {D, F<:AbstractFloat}
    ## Set resample back to false from previous iteration
    init!(tune.resample, false)
    ## Compute ESS
    ESS = BaytesCore.computeESS(particles.weights)
    resampled = BaytesCore.issmaller(ESS, tune.chains.Nchains * tune.chains.threshold)
    if resampled
        ## Record in tuning for diagnostics and propagation step
        init!(tune.resample, true)
        ## Resample ancestors ~ buffer already contain normalized weights from computeESS(weights) call
        resample!(
            _rng,
            tune.resample.method,
            particles.buffer.parameter.index,
            tune.iter.current,
            particles.weights.buffer,
            tune.chains.Nchains,
        )
        ## Equal weight normalized weights for next iteration memory
        Base.fill!(particles.weights.ℓweightsₙ, log(1.0 / tune.chains.Nchains))
        ## Reshuffle θ and rejuvenation particles - can be done inplace via buffer
        #!NOTE: Also set cumulative particle weights smc.buffer.weights back to correct position
        BaytesCore.shuffle!(particles.buffer.parameter, particles.kernel, particles.model, particles.buffer.cumweights)
        ## Rejuvenate Particles
        jitter!(_rng, particles, tune, data, temperature)
        ## Reweight ℓweights used for tempering so have correct index and temperature for next iteration
        reweight!(_rng, particles, tune, data, temperature)
    end
    return ESS, resampled
end

############################################################################################
"""
$(SIGNATURES)
Jitter θ particles with given kernels. This is performed in 2 stages:
1. kernel is updated with new data and shuffled particles, then one proposal step is performed.
2. If more than one jitterstep has to be performed, previous results might be captured depending on the kernel.

# Examples
```julia
```

"""
function jitter!(_rng::Random.AbstractRNG, particles::SMCParticles, tune::SMCTune, data::D, temperature::F) where {D, F<:AbstractFloat}
    ## Make first jitter step, update for new data and shuffled parameter
    Base.Threads.@threads for iter in eachindex(particles.kernel)
        ## Propose new parameter
        _, particles.buffer.jitterdiagnostics[iter] = propose!(
            _rng,
            jitterkernel(particles, iter),
            particles.model[iter],
            data,
            temperature,
            BaytesCore.UpdateTrue()
        )
        #!NOTE: has to be in loop for scoping rules
        particles.buffer.predictions[iter] = BaytesCore.get_prediction(
            particles.buffer.jitterdiagnostics[iter]
        )
    end
    # Calculate Correlation between old and new θ particles
    compute_ρ!(particles.buffer.correlation, particles.buffer.parameter.result, particles.kernel)
    # Check if jittering can be stopped
    jitter = BaytesCore.jitter!(tune.jitter, tune.jitterfun(particles.buffer.correlation.ρ))
    ## If more than one jitterstep has to be performed, previous results might be captured depending on the kernel.
    #!NOTE buffer.parameter.result.θᵤ at correct place because jitter! only happens after shuffle!
    while jitter
        ## Jitter Particles
        Base.Threads.@threads for iter in eachindex(particles.kernel)
            ## Propose new parameter
            #!TODO: This is only true in the first call though, if more than 1 step jittered, MCMC only kernels could be set to UpdateFalse()
            _, particles.buffer.jitterdiagnostics[iter] = propose!(
                _rng,
                jitterkernel(particles, iter),
                particles.model[iter],
                data,
                temperature,
                tune.capture
            )
            #!NOTE: has to be in loop for scoping rules
            particles.buffer.predictions[iter] = BaytesCore.get_prediction(
                particles.buffer.jitterdiagnostics[iter]
            )
        end
        ## Calculate Correlation between old and new θ particles
        compute_ρ!(particles.buffer.correlation, particles.buffer.parameter.result, particles.kernel)
        ## Check if jittering can be stopped
        jitter = jitter!(tune.jitter, tune.jitterfun(particles.buffer.correlation.ρ))
    end
    return nothing
end

############################################################################################
"""
$(SIGNATURES)
Propagate data forward over time. If (latent) data has to be extended, need to overload this function for specific smc kernel.

# Examples
```julia
```

"""
function propagate!(_rng::Random.AbstractRNG, particles::SMCParticles, tune::SMCTune, data::D, temperature::F) where {D, F<:AbstractFloat}
    return nothing
end

############################################################################################
"""
$(SIGNATURES)
Predict new data for each smc particle.

# Examples
```julia
```

"""
function predict!(_rng::Random.AbstractRNG, particles::SMCParticles, tune::SMCTune, data::D, temperature::F) where {D, F<:AbstractFloat}
    ## Propagate series forward with recent particle
    for iter in eachindex(particles.model)
        ## Predict new data point
        particles.buffer.predictions[iter] = predict(
            _rng,
            particles.kernel[iter],
            Objective(particles.model[iter], data, tune.tagged, temperature)
        )
    end
    #!NOTE: Here, particles.buffer.jitterdiagnostics will not be updated, and for SMCDiagnostics, last available jitterdiagnostics are provided.
    return nothing
end

############################################################################################
"""
$(SIGNATURES)
Compute new cumulative and incremental particle weights.

# Examples
```julia
```

"""
function weight!(
    _rng::Random.AbstractRNG,
    particles::SMCParticles,
    tune::SMCTune,
    data::D,
    temperature::F
) where {D, F<:AbstractFloat}
    ## Compute incremental weight for resampling step, and cumulative weight for i) next iteration and ii) temperature adjustment
    @inbounds for iter in eachindex(particles.weights.ℓweights)
        objective = Objective(particles.model[iter], data, tune.tagged, temperature)
        #!NOTE: particles.buffer.cumweights is cumulative weight at previous iteration, accounting for jittering step.
        particles.buffer.cumweights[iter], particles.weights.ℓweights[iter] = SMCweight(
            _rng,
            objective, particles.kernel[iter],
            particles.buffer.cumweights[iter]
        )
    end
    BaytesFilters.normalize!(particles.weights)
    return nothing
end

############################################################################################
"""
$(SIGNATURES)
Compute cumulative and incremental weight of objective at time/iteration t, given weight at t-1.
Incremental weight will be used as particle weight for resampling.
Cumulative weight will be used to adapt temperature.

If temperature is constant, this function can be overloaded with your ModelName if incremental weight can be computed independent of previous weight,
which speeds up computation. `cumweightsₜ` is not needed in this case.

# Examples
```julia
```

"""
function SMCweight(_rng::Random.AbstractRNG, objective::Objective, algorithm, cumweightsₜ₋₁)
    cumweightsₜ = objective.temperature * objective(objective.model.val)
    #!NOTE: cumweightsₜ₋₁ already has correct temperature at t-1
    return cumweightsₜ, cumweightsₜ - cumweightsₜ₋₁
end

############################################################################################
"""
$(SIGNATURES)
Compute new cumulative particle weights, accounting for jittering steps. This is only needed for
next iteration's weight calculation, and will not adjust current incremental and normalized weight.

# Examples
```julia
```

"""
function reweight!(
    _rng::Random.AbstractRNG,
    particles::SMCParticles,
    tune::SMCTune,
    data::D,
    temperature::F
) where {D, F<:AbstractFloat}
    ## Compute incremental weight for resampling step, and cumulative weight for i) next iteration and ii) temperature adjustment
    @inbounds for iter in eachindex(particles.weights.ℓweights)
        objective = Objective(particles.model[iter], data, tune.tagged, temperature)
        #!NOTE: particles.buffer.cumweights is cumulative weight at previous iteration, accounting for jittering step.
        particles.buffer.cumweights[iter], _ = SMCreweight(
            _rng,
            objective,
            particles.kernel[iter],
            particles.buffer.cumweights[iter]
        )
    end
    return nothing
end

############################################################################################
"""
$(SIGNATURES)
Computes particle weights after jiterring step. Defaults to `SMCweight` function.

# Examples
```julia
```

"""
function SMCreweight(_rng::Random.AbstractRNG, objective::Objective, algorithm, cumweightsₜ₋₁)
    return SMCweight(_rng, objective, algorithm, cumweightsₜ₋₁)
end

############################################################################################
# Export
export
    resample!,
    jitter!,
    propagate!,
    weight!,
    reweight!,
    SMCweight,
    SMCreweight