# This file is part of EndemicModel project, which is licensed under BDS-3-Clause.
# See file LICENSE.md for more information.

"""
A data source is supposed to be imported, exported or both.
Subtypes representing importable data source must implement method `import_data`.
Subtypes representing exportable data source must implement method `export_data`.
The method `pathof` should return the path to the database file or directory.
"""
abstract type AbstractDataSource end
abstract type AbstractDataFile <: AbstractDataSource end

const DataSourceDesignator = Union{AbstractDataSource,Symbol}

"""
    pathof(source::AbstractDataFile)

Return the path to the file or directory of `source`. Defaults to `source.path`.
"""
function pathof(source::AbstractDataFile)
    hasproperty(source, :path) ? source.path : nothing
end

"""
    import_data(source; kwargs...)

Import data from `source` and converts it either into a table or a named collection
thereof. `source` can be a `Symbol` which will be searched in `DATA_SOURCES`.
"""
function import_data end

"""
    import_data(sources::AbstractVector, data; kwargs...)

Import each source in `sources` and return a `Dict`:
```
Dict(source => import_data(source, data; kwargs...) for source ∈ sources)
```
"""
function import_data(
    sources::AbstractVector = option(:sources_to_update);
    cache_data_type = option(:cache_data_type),
    force_update::Bool = option(:force_update),
    kwargs...,
)
    @argcheck !isempty(sources)
    @debug "Importing data." sources
    source_keys = make_symbols(sources)
    try
        if !force_update && cache_data_type != nothing
            data = import_data(
                cache_data_type;
                cache_data_type = nothing,
                force_update = true,
                kwargs...,
            )
            if !(data isa AbstractDict)
                return data
            elseif source_keys ⊆ keys(data)
                return subdict(data, source_keys)
            end
        end
    catch
    end

    result = DataDict()
    for (source, sourcesym) ∈ zip(sources, source_keys)
        data =
            import_data(source; cache_data_type = nothing, force_update = true, kwargs...)
        result[sourcesym] = data
    end
    if cache_data_type != nothing
        export_data(cache_data_type, result; kwargs...)
    end
    @debug "Data imported." sources _debuginfo(result)
    result
end

function import_data(
    sourcesym::Symbol;
    path = nothing,
    cache_data_type = option(:cache_data_type),
    force_update::Bool = option(:force_update),
    kwargs...,
)
    try
        if !force_update && cache_data_type != nothing
            data = import_data(
                cache_data_type;
                cache_data_type = nothing,
                force_update = false,
                kwargs...,
            )
            if !(data isa AbstractDict)
                return data
            elseif haskey(data, sourcesym)
                return data[sourcesym]
            end
        end
    catch
    end

    source = DATA_SOURCES[sourcesym](nothing; path = path, kwargs...)
    data = import_data(source; cache_data_type = nothing, force_update = true, kwargs...)
    if cache_data_type != nothing
        export_data(cache_data_type, Dict(sourcesym => data); kwargs...)
    end
    data
end

"""
    export_data(source, data; kwargs...)

Export `data` to `source`; `data` must be a table or a named collection thereof.
If `source` can be a `Symbol`, search it in `DATA_SOURCES`.
"""
function export_data end

"""
    export_data(sources::AbstractVector, data; kwargs...)

Export each source in `sources` and return an `Dict`:
```
Dict(source => export_data(source, data; kwargs...) for source ∈ sources)
```
"""
function export_data(
    sources::AbstractVector,
    data;
    source_keys::AbstractVector{Symbol} = make_symbols(sources),
    kwargs...,
)
    Dict(
        sourcesym => export_data(source; kwargs...)
        for (source, sourcesym) ∈ zip(sources, source_keys)
    )
end

function export_data(sourcesym::Symbol, data; output = nothing, kwargs...)
    source = DATA_SOURCES[sourcesym](nothing; path = output, kwargs...)
    export_data(source, data; kwargs...)
end

function export_data(data; cache_data_type = option(:cache_data_type), kwargs...)
    export_data(cache_data_type, data; kwargs...)
end

"""
    cache_path(
        ext::Union{Nothing, AbstractString} = databaseExtension;
        directory = cache_directory,
        sub = cache_filename,
        nsamples::Integer = cache_nsamples,
        kwargs...,
    )

Returns the default cache file path with the given `ext`ension, or no extension if `ext`
is `nothing`.
"""
function cache_path(
    ext::Union{Nothing,AbstractString} = data_extension(option(:cache_data_type));
    directory = option(:cache_directory),
    update_period::Period = option(:update_period),
    subpath = option(:cache_filename),
    nsamples::Integer = option(:cache_nsamples),
    kwargs...,
)
    @argcheck nsamples >= 1
    (period_directory, lastdir) = _cache_subpath(Path(directory), update_period)
    subpath = isnothing(ext) ? subpath : subpath * "." * ext
    result = joinpath(period_directory, lastdir, Path(subpath))
    curdir = Path(dirname(result))
    if !exists(curdir)
        mkdir(curdir, recursive = true)
        gc_cache_directory(period_directory, nsamples)
    end
    result
end

function _cache_subpath(directory::AbstractPath, update_period::Period)
    try
        update_period = update_period < Second(1) ? Second(1) : update_period
    catch
    end
    n = Dates.value(update_period)
    timeformat = _dateformat_aux(update_period)
    period_directory = joinpath(directory, Path("Update every " * string(update_period)))
    (perfloor, perceil) = Dates.floorceil(now(), update_period)
    if n == 1
        lastdir = Path(Dates.format(perfloor, timeformat))
    else
        lastdir = Path(
            Dates.format(perfloor, timeformat) * " to " * Dates.format(perceil, timeformat),
        )
    end
    (period_directory, lastdir)
end

function cache_path(sourcetype::Type{T}; kwargs...) where {T<:AbstractDataFile}
    cache_path(data_extension(T); kwargs...)
end
function cache_path(sourcesym::Symbol; kwargs...)
    source = DATA_SOURCES[sourcesym]
    cache_path(data_extension(source); kwargs...)
end

_dateformat_aux(::Year) = dateformat"yyyy"
_dateformat_aux(per::Month) =
    Dates.value(per) >= 12 ? dateformat"yyyy" : dateformat"yyyy-mm"
_dateformat_aux(per::Week) =
    Dates.value(per) >= 5 ? dateformat"yyyy-mm" : dateformat"yyyy-mm-dd"
_dateformat_aux(per::Day) =
    Dates.value(per) >= 31 ? dateformat"yyyy-mm" : dateformat"yyyy-mm-dd"
_dateformat_aux(per::Hour) =
    Dates.value(per) >= 24 ? dateformat"yyyy-mm-dd" : dateformat"yyyy-mm-dd HHh"
_dateformat_aux(per::Minute) =
    Dates.value(per) >= 60 ? dateformat"yyyy-mm-dd HHh" : dateformat"yyyy-mm-dd HHhMM\m"
_dateformat_aux(per::Second) =
    Dates.value(per) >= 60 ? dateformat"yyyy-mm-dd HHhMM\m" :
    dateformat"yyyy-mm-dd HHhMM\mSS\s"

_dateformat_aux(period::Period) = dateformat"yyyy-mm-dd HHhMM\mSS\s"

"""
    gc_cache_directory(directory::AbstractPath, nsamples::Integer)

Garbage collect `directory`. Each subdirectory in cache `directory` is named by
date and/or time it was saved. This function keeps only the oldest `nsamples`
subdirectories.
"""
function gc_cache_directory(directory::AbstractPath, nsamples::Integer)
    @argcheck nsamples >= 1
    dirlist = readpath(directory)
    ndel = length(dirlist) - nsamples
    if ndel > 0
        todelete = sort!(dirlist; by = basename)[1:ndel]
        foreach(d -> rm(d; recursive = true), todelete)
    end
end

tableKeys = [:country, :state, :date]
tableColumns =
    [:country, :state, :latitude, :longitude, :date, :confirmed, :recovered, :deaths]

"""
Auxiliar structure to handle downloads.
"""
struct DownloadSource <: AbstractDataSource
    url::AbstractString
    function DownloadSource(url; kwargs...)
        new(string(url))
    end
end

function import_data(data::DownloadSource; kwargs...)
    url = data.url
    ext = split(url, '.')[end]
    ext = all(isletter, ext) ? ext : nothing
    subpath = joinpath(p"_DOWNLOADS_/", Path(string(hash(url))))
    path = cache_path(ext; subpath = subpath, kwargs...)
    @debug "Downloading data." url path
    download(url, path)
end

# CSV file and directory of files
"""
Path to a CSV file or directory
"""
struct CsvPath <: AbstractDataFile
    path::AbstractPath
end
function CsvPath(::Nothing; path = nothing, kwargs...)
    if path == nothing
        path = cache_path("csv"; kwargs...)
    end
    CsvPath(Path(path))
end

DATA_SOURCES[:csv] = CsvPath
data_extension(::Type{CsvPath}) = "csv"

function import_data(source::CsvPath; simple::Bool = false, kwargs...)
    path = pathof(source)
    @argcheck exists(path)
    if isfile(path)
        @debug "Importing DataFrame from CSV file." path
        data = csv_read(path; copycols = true)
    else
        @debug "Importing Dict from CSV directory." path
        data = DataDict(
            # Parameter `simple` don't need to be passed on
            Symbol(filename(elt)) => import_data(CsvPath(elt); kwargs...)
            for elt ∈ readpath(path)
        )
    end
    if simple
        data = simplify(data; kwargs...)
    end
    data
end

function export_data(
    source::CsvPath,
    data::AbstractDataFrame;
    pretty::Bool = false,
    kwargs...,
)
    destiny = pathof(source)
    @debug "Exporting DataFrame to CSV file." destiny _debuginfo(data)
    if pretty
        data = prettify(data; kwargs...)
    end
    exists(destiny) && rm(destiny; recursive = true)
    csv_write(destiny, data; transform = (col, val) -> something(val, missing))
    destiny
end

function export_data(source::CsvPath, data::AbstractDict; pretty::Bool = false, kwargs...)
    destiny = pathof(source)
    @debug "Exporting Dict to CSV directory." destiny _debuginfo(data)
    if pretty
        data = prettify(data; kwargs...)
    end
    !exists(destiny) && mkdir(destiny; recursive = true)
    for (fname, fdata) ∈ data
        fname = replace(string(fname), '/' => '_')
        fpath = joinpath(destiny, Path(fname * ".csv"))
        export_data(CsvPath(fpath), fdata; pretty = pretty, kwargs...)
    end
    destiny
end

struct OdsPath <: AbstractDataFile
    path::AbstractPath
end
function OdsPath(::Nothing; path = nothing, kwargs...)
    if path == nothing
        path = cache_path("ods"; kwargs...)
    end
    OdsPath(Path(path))
end

DATA_SOURCES[:ods] = OdsPath
data_extension(::Type{OdsPath}) = "ods"

function import_data(source::OdsPath; simple::Bool = false, kwargs...)
    path = pathof(source)
    @argcheck exists(path)
    if isfile(path)
        data = DataDict(Symbol(k) => v for (k, v) ∈ ods_readall(path))
    else
        data = DataDict(
            Symbol(filename(elt)) => import_data(OdsPath(elt); kwargs...)
            for elt ∈ readpath(path)
        )
        dfdict = get(data, :_DEFAULT_, nothing)
        if dfdict != nothing
            delete!(dfdict::Dict, :_DEFAULT_)
            merge!(data, dfdict)
        end
        data
    end
    if simple
        data = simplify(data; kwargs...)
    end
end

function export_data(
    source::OdsPath,
    data::AbstractDataFrame;
    pretty::Bool = false,
    table_name = option(:cache_table_name),
    kwargs...,
)
    if pretty
        data = prettify(data; kwargs...)
    end
    export_data(
        source,
        DataFrameDict(Symbol(table_name) => data);
        pretty = pretty,
        kwargs...,
    )
end

function export_data(source::OdsPath, data::DataFrameDict; pretty::Bool = false, kwargs...)
    destiny = pathof(source)
    rm(destiny; recursive = true, force = true)
    sheetdict = SortedDict{Tuple,DataFrame}()
    for (tname, tbl) ∈ data
        tname = string(tname)
        if pretty
            tname = prettify(tname; kwargs...)
            tbl = prettify(tbl; kwargs...)
        end
        push!(sheetdict, (tname, 1, 1) => tbl)
    end
    ods_write(destiny, sheetdict)
    destiny
end

function export_data(
    source::OdsPath,
    data::AbstractDataDict;
    pretty::Bool = false,
    kwargs...,
)
    if pretty
        data = prettify(data; kwargs...)
    end
    destiny = pathof(source)
    #rm(destiny; recursive = true, force = true)
    subdata = filter(p -> p[2] isa AbstractDict, data)
    dataframes = filter(p -> !(p[2] isa AbstractDict), data)
    dataframes = Dict{Symbol,DataFrame}(k => DataFrame(v) for (k, v) ∈ dataframes)
    if isempty(subdata)
        dfsource = source
    else
        !exists(destiny) && mkdir(destiny; recursive = true)
        dfsource = OdsPath(join(destiny, p"_DEFAULT_.ods"))

        for (fname, fdata) ∈ subdata
            fname = replace(string(fname), '/' => '_')
            fpath = joinpath(destiny, Path(fname * ".ods"))
            export_data(OdsPath(fpath), fdata; pretty = pretty, kwargs...)
        end
    end
    export_data(dfsource::OdsPath, dataframes; kwargs...)
    destiny
end


abstract type Forecast <: Component end

struct Deterministic <: Forecast
    device::Device
    horizon::Int
    resolution::Dates.Period
    interval::Dates.Period
    initialtime::Dates.DateTime
    data::Dict{Any,Dict{Int,TimeSeries.TimeArray}}
end

struct Scenarios <: Forecast
    horizon::Int
    resolution::Dates.Period
    interval::Dates.Period
    initialtime::Dates.DateTime
    scenarioquantity::Int
    data::Dict{Any,Dict{Int,TimeSeries.TimeArray}}
end

struct Probabilistic <: Forecast
    horizon::Int
    resolution::Dates.Period
    interval::Dates.Period
    initialtime::Dates.DateTime
    percentilequantity::Int
    data::Dict{Any,Dict{Int,TimeSeries.TimeArray}}
end


@testset "Read RAWGO" begin
        path = "./test/data/C2FEN02312/scenario_3/"
        raw_path = joinpath(path, "case.raw")
        network = PowerFlowNetwork(path, "RAWGO")
        @test network.baseMVA == 100
        @test size(network.bus, 1) == 2312
        @test nbus(raw_path, "RAWGO") == nbus(network)
        @test nbranch(raw_path, "RAWGO") == nbranch(network)
        @test nbranch(raw_path, "RAWGO", distinct_pair=true) == nbranch(network; distinct_pair=true)
        @test ngen(raw_path, "RAWGO") >= ngen(network)
end


include("point.jl")
include("line.jl")
include("circle.jl")
include("arc.jl")
include("rect.jl")
include("bezier.jl")
include("polygon.jl")
include("text.jl")


"""
    XDR (machine-independent binary) RData format IO stream wrapper.
"""
type XDRIO{T<:IO} <: RDAIO
    sub::T             # underlying IO stream
    buf::Vector{UInt8} # buffer for strings

    XDRIO(io::T) = new(io, Array(UInt8, 1024))
    @compat (::Type{XDRIO}){T <: IO}(io::T) = new{T}(io, Array(UInt8, 1024))
end

readint32(io::XDRIO) = ntoh(read(io.sub, Int32))
readuint32(io::XDRIO) = ntoh(read(io.sub, UInt32))
readfloat64(io::XDRIO) = ntoh(read(io.sub, Float64))

readintorNA(io::XDRIO) = readint32(io)
readintorNA(io::XDRIO, n::RVecLength) = map!(ntoh, read(io.sub, Int32, n))

# this method have Win32 ABI issues, see JuliaStats/RData.jl#5
# R's NA is silently converted to NaN when the value is loaded in the register(?)
#readfloatorNA(io::XDRIO) = readfloat64(io)

readfloatorNA(io::XDRIO, n::RVecLength) = reinterpret(Float64, map!(ntoh, read(io.sub, UInt64, n)))

readuint8(io::XDRIO, n::RVecLength) = readbytes(io.sub, n)

function readnchars(io::XDRIO, n::Int32)  # a single character string
    readbytes!(io.sub, io.buf, n)
    convert(RString, unsafe_string(pointer(io.buf), n))
end


using DelimitedFiles

words1 = readdlm("./examples/forceAlingned.jl/words.txt")

words2 = readdlm("./examples/forceAlingned.jl/aenesa.txt")

words = [words1[:,1] words2[:,3] words2[:,4]]

# set item info
class = "IntervalTier"
name = "Words"
size1 = size(words)[1]
xmin = words[1,:][2]
xmax = words[end,:][3]

# print(" "^4*"item [10]:\n")
# print(" "^8*"name = \""*name*"\" \n")
# print(" "^8*"xmin = "*string(xmin)*" \n")
# print(" "^8*"xmax = "*string(xmax)*" \n")
# print(" "^8*"intervals: size = "*string(size1)*" \n")
# for i in 1:size1
#     print(" "^8*"intervals ["*string(i)*"]:\n")
#     print(" "^12*"xmin = "*string(words[i,:][2])*" \n")
#     print(" "^12*"xmax = "*string(words[i,:][3])*" \n")
#     print(" "^12*"text = "*string(words[i,:][1])*" \n")
# end

open("./examples/forceAlingned.jl/aenesa_results.TextGrid", "a") do file
    write(file, " "^4*"item [10]:\n")
    write(file, " "^8*"class = \""*class*"\" \n")
    write(file, " "^8*"name = \""*name*"\" \n")
    write(file, " "^8*"xmin = "*string(xmin)*" \n")
    write(file, " "^8*"xmax = "*string(xmax)*" \n")
    write(file, " "^8*"intervals: size = "*string(size1)*" \n")
    for i in 1:size1
        write(file, " "^8*"intervals ["*string(i)*"]:\n")
        write(file, " "^12*"xmin = "*string(words[i,:][2])*" \n")
        write(file, " "^12*"xmax = "*string(words[i,:][3])*" \n")
        write(file, " "^12*"text = \""*string(words[i,:][1])*"\" \n")
    end
end


# These are some macros (and supporting functions) to make it easier to define rules.

"""
    @scalar_rule(f(x₁, x₂, ...),
                 @setup(statement₁, statement₂, ...),
                 (∂f₁_∂x₁, ∂f₁_∂x₂, ...),
                 (∂f₂_∂x₁, ∂f₂_∂x₂, ...),
                 ...)

A convenience macro that generates simple scalar forward or reverse rules using
the provided partial derivatives. Specifically, generates the corresponding
methods for `frule` and `rrule`:

    function ChainRulesCore.frule(::typeof(f), x₁::Number, x₂::Number, ...)
        Ω = f(x₁, x₂, ...)
        \$(statement₁, statement₂, ...)
        return Ω, (Rule((Δx₁, Δx₂, ...) -> ∂f₁_∂x₁ * Δx₁ + ∂f₁_∂x₂ * Δx₂ + ...),
                   Rule((Δx₁, Δx₂, ...) -> ∂f₂_∂x₁ * Δx₁ + ∂f₂_∂x₂ * Δx₂ + ...),
                   ...)
    end

    function ChainRulesCore.rrule(::typeof(f), x₁::Number, x₂::Number, ...)
        Ω = f(x₁, x₂, ...)
        \$(statement₁, statement₂, ...)
        return Ω, (Rule((ΔΩ₁, ΔΩ₂, ...) -> ∂f₁_∂x₁ * ΔΩ₁ + ∂f₂_∂x₁ * ΔΩ₂ + ...),
                   Rule((ΔΩ₁, ΔΩ₂, ...) -> ∂f₁_∂x₂ * ΔΩ₁ + ∂f₂_∂x₂ * ΔΩ₂ + ...),
                   ...)
    end

If no type constraints in `f(x₁, x₂, ...)` within the call to `@scalar_rule` are
provided, each parameter in the resulting `frule`/`rrule` definition is given a
type constraint of `Number`.
Constraints may also be explicitly be provided to override the `Number` constraint,
e.g. `f(x₁::Complex, x₂)`, which will constrain `x₁` to `Complex` and `x₂` to
`Number`.

Note that the result of `f(x₁, x₂, ...)` is automatically bound to `Ω`. This
allows the primal result to be conveniently referenced (as `Ω`) within the
derivative/setup expressions.

Note that the `@setup` argument can be elided if no setup code is need. In other
words:

    @scalar_rule(f(x₁, x₂, ...),
                 (∂f₁_∂x₁, ∂f₁_∂x₂, ...),
                 (∂f₂_∂x₁, ∂f₂_∂x₂, ...),
                 ...)

is equivalent to:

    @scalar_rule(f(x₁, x₂, ...),
                 @setup(nothing),
                 (∂f₁_∂x₁, ∂f₁_∂x₂, ...),
                 (∂f₂_∂x₁, ∂f₂_∂x₂, ...),
                 ...)

For examples, see ChainRulesCore' `rules` directory.

See also: [`frule`](@ref), [`rrule`](@ref), [`AbstractRule`](@ref)
"""
macro scalar_rule(call, maybe_setup, partials...)
    if Meta.isexpr(maybe_setup, :macrocall) && maybe_setup.args[1] == Symbol("@setup")
        setup_stmts = map(esc, maybe_setup.args[3:end])
    else
        setup_stmts = (nothing,)
        partials = (maybe_setup, partials...)
    end
    @assert Meta.isexpr(call, :call)
    f = esc(call.args[1])
    # Annotate all arguments in the signature as scalars
    inputs = map(call.args[2:end]) do arg
        esc(Meta.isexpr(arg, :(::)) ? arg : Expr(:(::), arg, :Number))
    end
    # Remove annotations and escape names for the call
    for (i, arg) in enumerate(call.args)
        if Meta.isexpr(arg, :(::))
            call.args[i] = esc(first(arg.args))
        else
            call.args[i] = esc(arg)
        end
    end
    if all(Meta.isexpr(partial, :tuple) for partial in partials)
        input_rep = :(first(promote($(inputs...))))  # stand-in with the right type for an input
        forward_rules = Any[rule_from_partials(input_rep, partial.args...) for partial in partials]
        reverse_rules = Any[]
        for i in 1:length(inputs)
            reverse_partials = [partial.args[i] for partial in partials]
            push!(reverse_rules, rule_from_partials(inputs[i], reverse_partials...))
        end
    else
        @assert length(inputs) == 1 && all(!Meta.isexpr(partial, :tuple) for partial in partials)
        forward_rules = Any[rule_from_partials(inputs[1], partial) for partial in partials]
        reverse_rules = Any[rule_from_partials(inputs[1], partials...)]
    end
    forward_rules = length(forward_rules) == 1 ? forward_rules[1] : Expr(:tuple, forward_rules...)
    reverse_rules = length(reverse_rules) == 1 ? reverse_rules[1] : Expr(:tuple, reverse_rules...)
    return quote
        function ChainRulesCore.frule(::typeof($f), $(inputs...))
            $(esc(:Ω)) = $call
            $(setup_stmts...)
            return $(esc(:Ω)), $forward_rules
        end
        function ChainRulesCore.rrule(::typeof($f), $(inputs...))
            $(esc(:Ω)) = $call
            $(setup_stmts...)
            return $(esc(:Ω)), $reverse_rules
        end
    end
end

function rule_from_partials(input_arg, ∂s...)
    wirtinger_indices = findall(x -> Meta.isexpr(x, :call) && x.args[1] === :Wirtinger,  ∂s)
    ∂s = map(esc, ∂s)
    Δs = [Symbol(string(:Δ, i)) for i in 1:length(∂s)]
    Δs_tuple = Expr(:tuple, Δs...)
    if isempty(wirtinger_indices)
        ∂_mul_Δs = [:(@thunk($(∂s[i])) * $(Δs[i])) for i in 1:length(∂s)]
        return :(Rule($Δs_tuple -> +($(∂_mul_Δs...))))
    else
        ∂_mul_Δs_primal = Any[]
        ∂_mul_Δs_conjugate = Any[]
        ∂_wirtinger_defs = Any[]
        for i in 1:length(∂s)
            if i in wirtinger_indices
                Δi = Δs[i]
                ∂i = Symbol(string(:∂, i))
                push!(∂_wirtinger_defs, :($∂i = $(∂s[i])))
                ∂f∂i_mul_Δ = :(wirtinger_primal($∂i) * wirtinger_primal($Δi))
                ∂f∂ī_mul_Δ̄ = :(conj(wirtinger_conjugate($∂i)) * wirtinger_conjugate($Δi))
                ∂f̄∂i_mul_Δ = :(wirtinger_conjugate($∂i) * wirtinger_primal($Δi))
                ∂f̄∂ī_mul_Δ̄ = :(conj(wirtinger_primal($∂i)) * wirtinger_conjugate($Δi))
                push!(∂_mul_Δs_primal, :($∂f∂i_mul_Δ + $∂f∂ī_mul_Δ̄))
                push!(∂_mul_Δs_conjugate, :($∂f̄∂i_mul_Δ + $∂f̄∂ī_mul_Δ̄))
            else
                ∂_mul_Δ = :(@thunk($(∂s[i])) * $(Δs[i]))
                push!(∂_mul_Δs_primal, ∂_mul_Δ)
                push!(∂_mul_Δs_conjugate, ∂_mul_Δ)
            end
        end
        primal_rule = :(Rule($Δs_tuple -> +($(∂_mul_Δs_primal...))))
        conjugate_rule = :(Rule($Δs_tuple -> +($(∂_mul_Δs_conjugate...))))
        return quote
            $(∂_wirtinger_defs...)
            AbstractRule(typeof($input_arg), $primal_rule, $conjugate_rule)
        end
    end
end


module  TestIterators

using PopGenCore
using Test

x = @nancycats

@testset "Iterators.jl" begin
    @testset "partition array" begin
        @test partitionarray(rand(20,5), [10,3,4,3]) .|> size == [(10, 5), (3, 5), (4, 5), (3, 5)]
        @test partitionarray(rand(20), [10,3,4,3]) .|> size == [(10, 1), (3, 1), (4, 1), (3, 1)]
    end

    @testset "pairiwse pairs" begin
        @test length(collect(pairwisepairs(1:3))) == 3
        @test collect(pairwisepairs(1:3)) == [(1, 2), (1, 3), (2, 3)]
        @test length(collect(pairwisepairs(1:10))) == 45
        @test length(collect(pairwisepairs(string.(1:3)))) == 3
        @test collect(pairwisepairs(string.(1:3))) == [("1", "2"), ("1", "3"), ("2", "3")]
        @test length(collect(pairwisepairs(string.(1:10)))) == 45
    end

    #TODO moved to PopGenSims.jl
    @testset "simulated pairs" begin
        @test simpairs(string.(1:4)) == [("1","2"), ("3","4")]
        @test_throws ArgumentError simpairs(string.(1:3))
    end

    @testset "skip iterators" begin
        @test collect(skipinf([1,2,3, Inf, -Inf, 4])) == [1,2,3,4]
        @test collect(skipnan([1,2,3, NaN,4])) == [1,2,3,4]
        @test collect(skipinfnan([1,2,3, Inf, -Inf, NaN, 4])) == [1,2,3,4]
    end
end

end


using AxisArrays

@testset "With AxisArray" begin
    a = AxisArray(reshape(1:6, 3, 2), Axis{:foo}(2:4), Axis{:bar}(3:4))
    # 2-dimensional AxisArray{Int64,2,...} with axes:
    #     :foo, 2:4
    #     :bar, 3:4
    # And data, a 3×2 Base.ReshapedArray{Int64,2,UnitRange{Int64},Tuple{}}:
    #  1  4
    #  2  5
    #  3  6

    v = Axis{:virt}(5:8)

    w = @inferred(EM.WrappedArray(a, (v, AxisArrays.axes(a)...)))

    for i = (5, 10, :a)
        @test @inferred(w[i, :, :]) == a
    end

    w2 = EM.WrappedArray(a, (AxisArrays.axes(a, 2), v, AxisArrays.axes(a, 1)))
    @test isa(@inferred(w2[1, 4, 3]), Number)
    @test w2[1, 4, 3] == a[3, 1]
end

@testset "With constant" begin
    using AxisArrays

    b = 2.0
    v = Axis{:virt}(5:8)

    w3 = EM.WrappedArray(b, v)
    @test @inferred(w3[5]) == b
    @test @inferred(w3[8]) == b
end


function RootsOfUnity()
    return RootsOfUnity(())
end

function compute_roots(obj::RootsOfUnity, arg0::jint)
    return jcall(obj, "computeRoots", void, (jint,), arg0)
end

function get_imaginary(obj::RootsOfUnity, arg0::jint)
    return jcall(obj, "getImaginary", jdouble, (jint,), arg0)
end

function get_number_of_roots(obj::RootsOfUnity)
    return jcall(obj, "getNumberOfRoots", jint, ())
end

function get_real(obj::RootsOfUnity, arg0::jint)
    return jcall(obj, "getReal", jdouble, (jint,), arg0)
end

function is_counter_clock_wise(obj::RootsOfUnity)
    return jcall(obj, "isCounterClockWise", jboolean, ())
end


"""
GenX: An Configurable Capacity Expansion Model
Copyright (C) 2021,  Massachusetts Institute of Technology
This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation; either version 2 of the License, or
(at your option) any later version.
This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.
A complete copy of the GNU General Public License v2 (GPLv2) is available
in LICENSE.txt.  Users uncompressing this from an archive may not have
received this license file.  If not, see <http://www.gnu.org/licenses/>.
"""

@doc raw"""
	thermal(EP::Model, inputs::Dict, UCommit::Int, Reserves::Int, CapacityReserveMargin::Int)
The thermal module creates decision variables, expressions, and constraints related to thermal power plants e.g. coal, oil or natural gas steam plants, natural gas combined cycle and combustion turbine plants, nuclear, hydrogen combustion etc.
This module uses the following 'helper' functions in separate files: ```thermal_commit()``` for thermal resources subject to unit commitment decisions and constraints (if any) and ```thermal_no_commit()``` for thermal resources not subject to unit commitment (if any).
"""
function thermal(EP::Model, inputs::Dict, UCommit::Int, Reserves::Int, CapacityReserveMargin::Int)

	println("Thermal Module")

	dfGen = inputs["dfGen"]

	T = inputs["T"]     # Number of time steps (hours)
	Z = inputs["Z"]     # Number of zones

	THERM_COMMIT = inputs["THERM_COMMIT"]
	THERM_NO_COMMIT = inputs["THERM_NO_COMMIT"]
	T = inputs["T"]
	THERM_ALL = inputs["THERM_ALL"]
	dfGen = inputs["dfGen"]

	if !isempty(THERM_COMMIT)
		EP = thermal_commit(EP::Model, inputs::Dict, Reserves::Int)
	end

	if !isempty(THERM_NO_COMMIT)
		EP = thermal_no_commit(EP::Model, inputs::Dict, Reserves::Int)
	end

	# Capacity Reserves Margin policy
	if CapacityReserveMargin > 0
		@expression(EP, eCapResMarBalanceThermal[res=1:inputs["NCapacityReserveMargin"], t=1:T], sum(dfGen[y,Symbol("CapRes_$res")] * EP[:eTotalCap][y] for y in THERM_ALL))
		EP[:eCapResMarBalance] += eCapResMarBalanceThermal
	end

	## CO2 Policy Module Thermal Generation by zone
	@expression(EP, eGenerationByThermAll[z=1:Z, t=1:T], # the unit is GW
		sum(EP[:vP][y,t] for y in intersect(inputs["THERM_ALL"], dfGen[dfGen[!,:Zone].==z,:R_ID]))
	)
	EP[:eGenerationByZone] += eGenerationByThermAll

	return EP
end


module Stat

import Gadfly
import StatsBase
import Contour
using Colors
using Compat
using Compose
using DataArrays
using DataStructures
using Distributions
using Hexagons
using Loess
using CoupledFields # It is registered in METADATA.jl
using IndirectArrays

import Gadfly: Scale, Coord, input_aesthetics, output_aesthetics,
               default_scales, isconcrete, setfield!, discretize_make_ia
import KernelDensity
# import Distributions: Uniform, Distribution, qqbuild
import IterTools: chain, distinct
import Compat.Iterators: cycle, product

include("bincount.jl")

# Apply a series of statistics.
#
# Args:
#   stats: Statistics to apply in order.
#   scales: Scales used by the plot.
#   aes: A Aesthetics instance.
#
# Returns:
#   Nothing, modifies aes.
#
function apply_statistics(stats::Vector{Gadfly.StatisticElement},
                          scales::Dict{Symbol, Gadfly.ScaleElement},
                          coord::Gadfly.CoordinateElement,
                          aes::Gadfly.Aesthetics)
    for stat in stats
        apply_statistic(stat, scales, coord, aes)
    end
    nothing
end

struct Nil <: Gadfly.StatisticElement end

const nil = Nil

struct Identity <: Gadfly.StatisticElement end

apply_statistic(stat::Identity,
                scales::Dict{Symbol, Gadfly.ScaleElement},
                coord::Gadfly.CoordinateElement,
                aes::Gadfly.Aesthetics) = nothing

const identity = Identity

# Determine bounds of bars positioned at the given values.
function barminmax(vals, iscontinuous::Bool)
    minvalue, maxvalue = extrema(vals)
    span_type = typeof((maxvalue - minvalue) / 1.0)
    barspan = one(span_type)

    if iscontinuous && length(vals) > 1
        sorted_vals = sort(vals)
        T = typeof(sorted_vals[2] - sorted_vals[1])
        z = zero(T)
        minspan = z
        for i in 2:length(vals)
            span = sorted_vals[i] - sorted_vals[i-1]
            if span > z && (span < minspan || minspan == z)
                minspan = span
            end
        end
        barspan = minspan
    end
    position_type = promote_type(typeof(barspan/2.0), eltype(vals))
    minvals = Array{position_type}(length(vals))
    maxvals = Array{position_type}(length(vals))

    for (i, x) in enumerate(vals)
        minvals[i] = x - barspan/2.0
        maxvals[i] = x + barspan/2.0
    end

    return minvals, maxvals
end

struct RectbinStatistic <: Gadfly.StatisticElement end

const rectbin = RectbinStatistic

input_aesthetics(stat::RectbinStatistic) = [:x, :y]
output_aesthetics(stat::RectbinStatistic) = [:xmin, :xmax, :ymin, :ymax]

function apply_statistic(stat::RectbinStatistic,
                         scales::Dict{Symbol, Gadfly.ScaleElement},
                         coord::Gadfly.CoordinateElement,
                         aes::Gadfly.Aesthetics)
    Gadfly.assert_aesthetics_defined("RectbinStatistic", aes, :x, :y)

    isxcontinuous = haskey(scales, :x) && isa(scales[:x], Scale.ContinuousScale)
    isycontinuous = haskey(scales, :y) && isa(scales[:y], Scale.ContinuousScale)

    xminvals, xmaxvals = barminmax(aes.x, isxcontinuous)
    yminvals, ymaxvals = barminmax(aes.y, isycontinuous)

    aes.xmin = xminvals
    aes.xmax = xmaxvals
    aes.ymin = yminvals
    aes.ymax = ymaxvals

    if !isxcontinuous
        aes.pad_categorical_x = Nullable(false)
    end
    if !isycontinuous
        aes.pad_categorical_y = Nullable(false)
    end
end

struct BarStatistic <: Gadfly.StatisticElement
    position::Symbol # :dodge or :stack
    orientation::Symbol # :horizontal or :vertical
end
BarStatistic(; position=:stack, orientation=:vertical) = BarStatistic(position, orientation)

input_aesthetics(stat::BarStatistic) = stat.orientation == :vertical ? [:x] : [:y]
output_aesthetics(stat::BarStatistic) =
    stat.orientation == :vertical ? [:xmin, :xmax] : [:ymin, :ymax]

default_scales(stat::BarStatistic) = stat.orientation == :vertical ?
        [Gadfly.Scale.y_continuous()] : [Gadfly.Scale.x_continuous()]

const bar = BarStatistic

function apply_statistic(stat::BarStatistic,
                         scales::Dict{Symbol, Gadfly.ScaleElement},
                         coord::Gadfly.CoordinateElement,
                         aes::Gadfly.Aesthetics)
    if stat.orientation == :horizontal
        in(:y, Gadfly.defined_aesthetics(aes)) || return
        var = :y
        othervar = :x
        minvar = :ymin
        maxvar = :ymax
        viewminvar = :xviewmin
        viewmaxvar = :xviewmax
        other_viewminvar = :yviewmin
        other_viewmaxvar = :yviewmax
        labelvar = :x_label
    else
        in(:x, Gadfly.defined_aesthetics(aes)) || return
        var = :x
        othervar = :y
        minvar = :xmin
        maxvar = :xmax
        viewminvar = :yviewmin
        viewmaxvar = :yviewmax
        other_viewminvar = :xviewmin
        other_viewmaxvar = :xviewmax
        labelvar = :y_label
    end

    vals = getfield(aes, var)
    if isempty(vals)
      setfield!(aes, minvar, Float64[1.0])
      setfield!(aes, maxvar, Float64[1.0])
      setfield!(aes, var, Float64[1.0])
      setfield!(aes, othervar, Float64[0.0])
      return
    end

    iscontinuous = haskey(scales, var) && isa(scales[var], Scale.ContinuousScale)

    if getfield(aes, minvar) == nothing || getfield(aes, maxvar) == nothing
        minvals, maxvals = barminmax(vals, iscontinuous)

        setfield!(aes, minvar, minvals)
        setfield!(aes, maxvar, maxvals)
    end

    z = zero(eltype(getfield(aes, othervar)))
    if getfield(aes, viewminvar) == nothing && z < minimum(getfield(aes, othervar))
        setfield!(aes, viewminvar, z)
    elseif getfield(aes, viewmaxvar) == nothing && z > maximum(getfield(aes, othervar))
        setfield!(aes, viewmaxvar, z)
    end

    if stat.position == :stack && aes.color !== nothing
        groups = Dict{Any,Float64}()
        for (x, y) in zip(getfield(aes, othervar), vals)
            Gadfly.isconcrete(x) || continue

            if !haskey(groups, y)
                groups[y] = Float64(x)
            else
                groups[y] += Float64(x)
            end
        end

        viewmin, viewmax = extrema(values(groups))
        aes_viewminvar = getfield(aes, viewminvar)
        if aes_viewminvar === nothing || aes_viewminvar > viewmin
            setfield!(aes, viewminvar, viewmin)
        end
        aes_viewmaxvar = getfield(aes, viewmaxvar)
        if aes_viewmaxvar === nothing || aes_viewmaxvar < viewmax
            setfield!(aes, viewmaxvar, viewmax)
        end
    end

    if !iscontinuous
        if stat.orientation == :horizontal
            aes.pad_categorical_y = Nullable(false)
        else
            aes.pad_categorical_x = Nullable(false)
        end
    end
end

struct HistogramStatistic <: Gadfly.StatisticElement
    minbincount::Int
    maxbincount::Int
    position::Symbol # :dodge or :stack
    orientation::Symbol
    density::Bool
end

function HistogramStatistic(; bincount=nothing,
                              minbincount=3,
                              maxbincount=150,
                              position=:stack,
                              orientation=:vertical,
                              density=false)
    if bincount != nothing
        HistogramStatistic(bincount, bincount, position, orientation, density)
    else
        HistogramStatistic(minbincount, maxbincount, position, orientation, density)
    end
end

input_aesthetics(stat::HistogramStatistic) = stat.orientation == :vertical ? [:x] : [:y]
output_aesthetics(stat::HistogramStatistic) =
    stat.orientation == :vertical ? [:x, :y, :xmin, :xmax] : [:y, :x, :ymin, :ymax]

default_scales(stat::HistogramStatistic) = stat.orientation == :vertical ?
        [Gadfly.Scale.y_continuous()] : [Gadfly.Scale.x_continuous()]

const histogram = HistogramStatistic

function apply_statistic(stat::HistogramStatistic,
                         scales::Dict{Symbol, Gadfly.ScaleElement},
                         coord::Gadfly.CoordinateElement,
                         aes::Gadfly.Aesthetics)
    if stat.orientation == :horizontal
        var = :y
        othervar = :x
        minvar = :ymin
        maxvar = :ymax
        viewminvar = :xviewmin
        viewmaxvar = :xviewmax
        labelvar = :x_label
    else
        var = :x
        othervar = :y
        minvar = :xmin
        maxvar = :xmax
        viewminvar = :yviewmin
        viewmaxvar = :yviewmax
        labelvar = :y_label
    end

    Gadfly.assert_aesthetics_defined("HistogramStatistic", aes, var)

    vals = getfield(aes, var)
    stat.minbincount > stat.maxbincount && error("Histogram minbincount > maxbincount")

    if isempty(getfield(aes, var))
        setfield!(aes, minvar, Float64[1.0])
        setfield!(aes, maxvar, Float64[1.0])
        setfield!(aes, var, Float64[1.0])
        setfield!(aes, othervar, Float64[0.0])
        return
    end

    if haskey(scales, var) && isa(scales[var], Scale.DiscreteScale)
        isdiscrete = true
        x_min, x_max = extrema(vals)
        d = x_max - x_min + 1
        bincounts = zeros(Int, d)
        for x in vals
            bincounts[x - x_min + 1] += 1
        end
        x_min -= 0.5 # adjust the left side of the bar
        binwidth = 1.0
    else
        x_min = Gadfly.concrete_minimum(vals)

        isdiscrete = false
        if estimate_distinct_proportion(vals) <= 0.9
            value_set = sort!(collect(Set(vals[Bool[Gadfly.isconcrete(v) for v in vals]])))
            d, bincounts, x_max = choose_bin_count_1d_discrete(
                        vals, value_set, stat.minbincount, stat.maxbincount)
        else
            d, bincounts, x_max = choose_bin_count_1d(
                        vals, stat.minbincount, stat.maxbincount)
        end

        if stat.density
            x_min = Gadfly.concrete_minimum(vals)
            span = x_max - x_min
            binwidth = span / d
            bincounts = bincounts ./ (sum(bincounts) * binwidth)
        end

        binwidth = (x_max - x_min) / d
    end

    if aes.color === nothing
        T = typeof(x_min + 1*binwidth)
        setfield!(aes, othervar, Array{Float64}(d))
        setfield!(aes, minvar, Array{T}(d))
        setfield!(aes, maxvar, Array{T}(d))
        setfield!(aes, var, Array{T}(d))
        for j in 1:d
            getfield(aes, minvar)[j] = x_min + (j - 1) * binwidth
            getfield(aes, maxvar)[j] = x_min + j * binwidth
            getfield(aes, var)[j] = x_min + (j - 0.5) * binwidth
            getfield(aes, othervar)[j] = bincounts[j]
        end
    else
        groups = Dict()
        for (x, c) in zip(vals, cycle(aes.color))
            Gadfly.isconcrete(x) || continue

            if !haskey(groups, c)
                groups[c] = Float64[x]
            else
                push!(groups[c], x)
            end
        end
        T = typeof(x_min + 1*binwidth)
        setfield!(aes, minvar, Array{T}(d * length(groups)))
        setfield!(aes, maxvar, Array{T}(d * length(groups)))
        setfield!(aes, var, Array{T}(d * length(groups)))

        setfield!(aes, othervar, Array{Float64}(d * length(groups)))
        colors = Array{RGB{Float32}}(d * length(groups))

        x_span = x_max - x_min
        stack_height = zeros(Int, d)
        for (i, (c, xs)) in enumerate(groups)
            fill!(bincounts, 0)
            for x in xs
                Gadfly.isconcrete(x) || continue
                if isdiscrete
                    bincounts[round(Int,x)] += 1
                else
                    bin = max(1, min(d, (ceil(Int, (x - x_min) / binwidth))))
                    bincounts[bin] += 1
                end
            end

            if stat.density
                binwidth = x_span / d
                bincounts = bincounts ./ (sum(bincounts) * binwidth)
            end

            stack_height += bincounts[1:d]

            if isdiscrete
                for j in 1:d
                    idx = (i-1)*d + j
                    getfield(aes, var)[idx] = j
                    getfield(aes, othervar)[idx] = bincounts[j]
                    colors[idx] = c
                end
            else
                for j in 1:d
                    idx = (i-1)*d + j
                    getfield(aes, minvar)[idx] = x_min + (j - 1) * binwidth
                    getfield(aes, maxvar)[idx] = x_min + j * binwidth
                    getfield(aes, var)[idx] = x_min + (j - 0.5) * binwidth
                    getfield(aes, othervar)[idx] = bincounts[j]
                    colors[idx] = c
                end
            end
        end

        if stat.position == :stack
            viewmax = Float64(maximum(stack_height))
            aes_viewmax = getfield(aes, viewmaxvar)
            if aes_viewmax === nothing || aes_viewmax < viewmax
                setfield!(aes, viewmaxvar, viewmax)
            end
        end

        aes.color = discretize_make_ia(colors)
    end

    getfield(aes, viewminvar) === nothing && setfield!(aes, viewminvar, 0.0)

    if haskey(scales, othervar)
        data = Gadfly.Data()
        setfield!(data, othervar, getfield(aes, othervar))
        setfield!(data, viewmaxvar, getfield(aes, viewmaxvar))
        Scale.apply_scale(scales[othervar], [aes], data)

        # See issue #560. Stacked histograms on a non-linear y scale are a strange
        # thing. After some discussion, the least confusing thing is to make the stack
        # partitioned linearly. Here we make that adjustment.
        if stat.position == :stack && aes.color != nothing
            # A little trickery to figure out the scale stack height.
            data = Gadfly.Data()
            setfield!(data, othervar, stack_height)
            scaled_stackheight_aes = Gadfly.Aesthetics()
            Scale.apply_scale(scales[othervar], [scaled_stackheight_aes], data)
            scaled_stackheight = getfield(scaled_stackheight_aes, othervar)

            othervals = getfield(aes, othervar)
            for j in 1:d
                naive_stackheight = 0
                for i in 1:length(groups)
                    idx = (i-1)*d + j
                    naive_stackheight += othervals[idx]
                end

                naive_stackheight == 0 && continue

                for i in 1:length(groups)
                    idx = (i-1)*d + j
                    othervals[idx] = scaled_stackheight[j] * othervals[idx] / naive_stackheight
                end
            end
        end
    else
        setfield!(aes, labelvar, Scale.identity_formatter)
    end
end

struct Density2DStatistic <: Gadfly.StatisticElement
    n::Tuple{Int,Int} # Number of points sampled
    bw::Tuple{Real,Real} # Bandwidth used for the kernel density estimation
    levels::Union{Int,Vector,Function}
end
Density2DStatistic(; n=(256,256), bandwidth=(-Inf,-Inf), levels=15) =
      Density2DStatistic(n, bandwidth, levels)

const density2d = Density2DStatistic

input_aesthetics(stat::Density2DStatistic) = [:x, :y]
output_aesthetics(stat::Density2DStatistic) = [:x, :y, :z]

default_scales(::Density2DStatistic) = [Gadfly.Scale.y_continuous()]

function apply_statistic(stat::Density2DStatistic,
                         scales::Dict{Symbol, Gadfly.ScaleElement},
                         coord::Gadfly.CoordinateElement,
                         aes::Gadfly.Aesthetics)
    Gadfly.assert_aesthetics_defined("Density2DStatistic", aes, :x, :y)

    window = (stat.bw[1] <= 0.0 ? KernelDensity.default_bandwidth(aes.x) : stat.bw[1],
              stat.bw[2] <= 0.0 ? KernelDensity.default_bandwidth(aes.y) : stat.bw[2])
    k = KernelDensity.kde((aes.x,aes.y), bandwidth=window, npoints=stat.n)
    aes.z = k.density
    aes.x = collect(k.x)
    aes.y = collect(k.y)
    apply_statistic(ContourStatistic(levels=stat.levels), scales, coord, aes)
end

struct DensityStatistic <: Gadfly.StatisticElement
    # Number of points sampled
    n::Int
    # Bandwidth used for the kernel density estimation
    bw::Real
end
DensityStatistic(; n=256, bandwidth=-Inf) = DensityStatistic(n, bandwidth)

const density = DensityStatistic

input_aesthetics(stat::DensityStatistic) = [:x]
output_aesthetics(stat::DensityStatistic) = [:x, :y]

default_scales(::DensityStatistic) = [Gadfly.Scale.y_continuous()]

function apply_statistic(stat::DensityStatistic,
                         scales::Dict{Symbol, Gadfly.ScaleElement},
                         coord::Gadfly.CoordinateElement,
                         aes::Gadfly.Aesthetics)
    Gadfly.assert_aesthetics_defined("DensityStatistic", aes, :x)

    if aes.color === nothing
        isa(aes.x[1], Real) || error("Kernel density estimation only works on Real types.")

        x_f64 = collect(Float64, aes.x)

        window = stat.bw <= 0.0 ? KernelDensity.default_bandwidth(x_f64) : stat.bw
        f = KernelDensity.kde(x_f64, bandwidth=window, npoints=stat.n)
        aes.x = collect(Float64, f.x)
        aes.y = f.density
    else
        groups = Dict()
        for (x, c) in zip(aes.x, cycle(aes.color))
            if !haskey(groups, c)
                groups[c] = Float64[x]
            else
                push!(groups[c], x)
            end
        end

        colors = Array{RGB{Float32}}(0)
        aes.x = Array{Float64}(0)
        aes.y = Array{Float64}(0)
        for (c, xs) in groups
            window = stat.bw <= 0.0 ? KernelDensity.default_bandwidth(xs) : stat.bw
            f = KernelDensity.kde(xs, bandwidth=window, npoints=stat.n)
            append!(aes.x, f.x)
            append!(aes.y, f.density)
            for _ in 1:length(f.x)
                push!(colors, c)
            end
        end
        aes.color = discretize_make_ia(colors)
    end
    aes.y_label = Gadfly.Scale.identity_formatter
end

struct Histogram2DStatistic <: Gadfly.StatisticElement
    xminbincount::Int
    xmaxbincount::Int
    yminbincount::Int
    ymaxbincount::Int
end

function Histogram2DStatistic(; xbincount=nothing,
                                xminbincount=3,
                                xmaxbincount=150,
                                ybincount=nothing,
                                yminbincount=3,
                                ymaxbincount=150)
    if xbincount != nothing
        xminbincount = xbincount
        xmaxbincount = xbincount
    end

    if ybincount != nothing
        yminbincount = ybincount
        ymaxbincount = ybincount
    end

    Histogram2DStatistic(xminbincount, xmaxbincount, yminbincount, ymaxbincount)
end

input_aesthetics(stat::Histogram2DStatistic) = [:x, :y]
output_aesthetics(stat::Histogram2DStatistic) = [:xmin, :ymax, :ymin, :ymax, :color]

default_scales(::Histogram2DStatistic, t::Gadfly.Theme=Gadfly.current_theme()) =
    [t.continuous_color_scale]

const histogram2d = Histogram2DStatistic

function apply_statistic(stat::Histogram2DStatistic,
                         scales::Dict{Symbol, Gadfly.ScaleElement},
                         coord::Gadfly.CoordinateElement,
                         aes::Gadfly.Aesthetics)
    Gadfly.assert_aesthetics_defined("Histogram2DStatistic", aes, :x, :y)

    x_min, x_max = Gadfly.concrete_minimum(aes.x), Gadfly.concrete_maximum(aes.x)
    y_min, y_max = Gadfly.concrete_minimum(aes.y), Gadfly.concrete_maximum(aes.y)

    if haskey(scales, :x) && isa(scales[:x], Scale.DiscreteScale)
        x_categorial = true
        xminbincount = x_max - x_min + 1
        xmaxbincount = xminbincount
    else
        x_categorial = false
        xminbincount = stat.xminbincount
        xmaxbincount = stat.xmaxbincount
    end

    if haskey(scales, :y) && isa(scales[:y], Scale.DiscreteScale)
        y_categorial = true
        yminbincount = y_max - y_min + 1
        ymaxbincount = yminbincount
    else
        y_categorial = false
        yminbincount = stat.yminbincount
        ymaxbincount = stat.ymaxbincount
    end

    dy, dx, bincounts = choose_bin_count_2d(aes.x, aes.y,
                                            xminbincount, xmaxbincount,
                                            yminbincount, ymaxbincount)

    wx = x_categorial ? 1 : (x_max - x_min) / dx
    wy = y_categorial ? 1 : (y_max - y_min) / dy

    n = 0
    for cnt in bincounts
        if cnt > 0
            n += 1
        end
    end

    if x_categorial
        aes.x = Array{Int64}(n)
    else
        aes.xmin = Array{Float64}(n)
        aes.xmax = Array{Float64}(n)
    end

    if y_categorial
        aes.y = Array{Int64}(n)
    else
        aes.ymin = Array{Float64}(n)
        aes.ymax = Array{Float64}(n)
    end

    k = 1
    for i in 1:dy, j in 1:dx
        cnt = bincounts[i, j]
        if cnt > 0
            if x_categorial
                aes.x[k] = x_min + (j - 1)
            else
                aes.xmin[k] = x_min + (j - 1) * wx
                aes.xmax[k] = x_min + j * wx
            end

            if y_categorial
                aes.y[k] = y_min + (i - 1)
            else
                aes.ymin[k] = y_min + (i - 1) * wy
                aes.ymax[k] = y_min + i * wy
            end
            k += 1
        end
    end
    @assert k - 1 == n

    haskey(scales, :color) || error("Histogram2DStatistic requires a color scale.")
    color_scale = scales[:color]
    typeof(color_scale) <: Scale.ContinuousColorScale ||
            error("Histogram2DStatistic requires a continuous color scale.")

    aes.color_key_title = "Count"

    data = Gadfly.Data()
    data.color = Array{Int}(n)
    k = 1
    for cnt in transpose(bincounts)
        if cnt > 0
            data.color[k] = cnt
            k += 1
        end
    end

    if x_categorial
        aes.xmin, aes.xmax = barminmax(aes.x, false)
        aes.x = discretize_make_ia(aes.x)
        aes.pad_categorical_x = Nullable(false)
    end

    if y_categorial
        aes.ymin, aes.ymax = barminmax(aes.y, false)
        aes.y = discretize_make_ia(aes.y)
        aes.pad_categorical_y = Nullable(false)
    end

    Scale.apply_scale(color_scale, [aes], data)
    nothing
end

# Find reasonable places to put tick marks and grid lines.
struct TickStatistic <: Gadfly.StatisticElement
    in_vars::Vector{Symbol}
    out_var::AbstractString

    granularity_weight::Float64
    simplicity_weight::Float64
    coverage_weight::Float64
    niceness_weight::Float64

    # fixed ticks, or nothing
    ticks::Union{Symbol, AbstractArray}
end

@deprecate xticks(ticks) xticks(ticks=ticks)

xticks(; ticks=:auto,
         granularity_weight=1/4,
         simplicity_weight=1/6,
         coverage_weight=1/3,
         niceness_weight=1/4) =
    TickStatistic([:x, :xmin, :xmax, :xintercept], "x",
              granularity_weight, simplicity_weight, coverage_weight, niceness_weight, ticks)

@deprecate yticks(ticks) yticks(ticks=ticks)

yticks(; ticks=:auto,
         granularity_weight=1/4,
         simplicity_weight=1/6,
         coverage_weight=1/3,
         niceness_weight=1/4) =
    TickStatistic([:y, :ymin, :ymax, :yintercept, :middle, :lower_hinge, :upper_hinge,
        :lower_fence, :upper_fence], "y",
        granularity_weight, simplicity_weight, coverage_weight, niceness_weight, ticks)

# Apply a tick statistic.

# Args:
#   stat: statistic.
#   aes: aesthetics.

# Returns:
#   nothing

# Modifies:
#   aes
#
function apply_statistic(stat::TickStatistic,
                         scales::Dict{Symbol, Gadfly.ScaleElement},
                         coord::Gadfly.CoordinateElement,
                         aes::Gadfly.Aesthetics)
    isa(stat.ticks, Symbol) && stat.ticks != :auto &&
            error("Invalid value $(stat.ticks) for ticks parameter.")

    isa(coord, Coord.SubplotGrid) &&
            error("TickStatistic cannot be applied to subplot coordinates.")

    # don't clobber existing ticks
    getfield(aes, Symbol(stat.out_var, "tick")) == nothing || return

    in_group_var = Symbol(stat.out_var, "group")
    minval, maxval = nothing, nothing
    in_vals = Any[]
    categorical = (:x in stat.in_vars && Scale.iscategorical(scales, :x)) ||
                  (:y in stat.in_vars && Scale.iscategorical(scales, :y))

    for var in stat.in_vars
        categorical && !in(var,[:x,:y]) && continue
        vals = getfield(aes, var)
        if vals != nothing && eltype(vals) != Function
            if minval == nothing
                minval = first(vals)
            end
            if maxval == nothing
                maxval = first(vals)
            end
            T = promote_type(typeof(minval), typeof(maxval))
            T = promote_type(T, eltype(vals))
            minval = convert(T, minval)
            maxval = convert(T, maxval)

            if stat.out_var == "x"
                dsize = aes.xsize === nothing ? [nothing] : aes.xsize
            elseif stat.out_var == "y"
                dsize = aes.ysize === nothing ? [nothing] : aes.ysize
            else
                dsize = [nothing]
            end

            size = aes.size === nothing ? [nothing] : aes.size

            minval, maxval = apply_statistic_typed(minval, maxval, vals, size, dsize)
            push!(in_vals, vals)
        end
    end

    isempty(in_vals) && return

    in_vals = chain(in_vals...)

    # consider forced tick marks
    if stat.ticks != :auto
        minval = min(minval, minimum(stat.ticks))
        maxval = max(maxval, maximum(stat.ticks))
    end

    # TODO: handle the outliers aesthetic

    n = Gadfly.concrete_length(in_vals)

    # check the x/yviewmin/max pesudo-aesthetics
    if stat.out_var == "x"
        if aes.xviewmin != nothing
            minval = min(minval, aes.xviewmin)
        end
        if aes.xviewmax != nothing
            maxval = max(maxval, aes.xviewmax)
        end
    elseif stat.out_var == "y"
        if aes.yviewmin != nothing
            minval = min(minval, aes.yviewmin)
        end
        if aes.yviewmax != nothing
            maxval = max(maxval, aes.yviewmax)
        end
    end

    # take into account a forced viewport in cartesian coordinates.
    strict_span = false
    if typeof(coord) == Coord.Cartesian
        if stat.out_var == "x"
            if coord.xmin !== nothing
                minval = coord.xmin
                strict_span = true
            end
            if coord.xmax !== nothing
                maxval = coord.xmax
                strict_span = true
            end
        elseif stat.out_var == "y"
            if coord.ymin !== nothing
                minval = coord.ymin
                strict_span = true
            end
            if coord.ymax !== nothing
                maxval = coord.ymax
                strict_span = true
            end
        end
    end

    # all the input values in order.
    if stat.ticks != :auto
        grids = ticks = stat.ticks
        viewmin = minval
        viewmax = maxval
        tickvisible = fill(true, length(ticks))
        tickscale = fill(1.0, length(ticks))
    elseif categorical
        ticks = Set{Int}()
        for val in in_vals
            val>0 && push!(ticks, round(Int, val))
        end
        ticks = Int[t for t in ticks]
        sort!(ticks)
        grids = (ticks .- 0.5)[2:end]
        viewmin = minimum(ticks)
        viewmax = maximum(ticks)
        tickvisible = fill(true, length(ticks))
        tickscale = fill(1.0, length(ticks))
    else
        minval, maxval = promote(minval, maxval)

        ticks, viewmin, viewmax = Gadfly.optimize_ticks(minval, maxval, extend_ticks=true,
                granularity_weight=stat.granularity_weight,
                simplicity_weight=stat.simplicity_weight,
                coverage_weight=stat.coverage_weight,
                niceness_weight=stat.niceness_weight,
                strict_span=strict_span)
        grids = ticks
        multiticks = Gadfly.multilevel_ticks(viewmin - (viewmax - viewmin),
                                             viewmax + (viewmax - viewmin))
        tickcount = length(ticks) + sum([length(ts) for ts in values(multiticks)])
        tickvisible = Array{Bool}(tickcount)
        tickscale = Array{Float64}(tickcount)
        i = 1
        for t in ticks
            tickscale[i] = 1.0
            tickvisible[i] = viewmin <= t <= viewmax
            i += 1
        end

        for (scale, ts) in multiticks, t in ts
            push!(ticks, t)
            tickvisible[i] = false
            tickscale[i] = scale
            i += 1
        end
    end

    # We use the first label function we find for any of the aesthetics. I'm not
    # positive this is the right thing to do, or would would be.
    labeler = getfield(aes, Symbol(stat.out_var, "_label"))

    setfield!(aes, Symbol(stat.out_var, "tick"), ticks)
    setfield!(aes, Symbol(stat.out_var, "grid"), grids)
    setfield!(aes, Symbol(stat.out_var, "tick_label"), labeler)
    setfield!(aes, Symbol(stat.out_var, "tickvisible"), tickvisible)
    setfield!(aes, Symbol(stat.out_var, "tickscale"), tickscale)

    viewmin_var = Symbol(stat.out_var, "viewmin")
    if getfield(aes, viewmin_var) === nothing || getfield(aes, viewmin_var) > viewmin
        setfield!(aes, viewmin_var, viewmin)
    end

    viewmax_var = Symbol(stat.out_var, "viewmax")
    if getfield(aes, viewmax_var) === nothing || getfield(aes, viewmax_var) < viewmax
        setfield!(aes, viewmax_var, viewmax)
    end

    nothing
end

function apply_statistic_typed(minval::T, maxval::T, vals, size, dsize) where T
#     for (val, s, ds) in zip(vals, cycle(size), cycle(dsize))
    lensize  = length(size)
    lendsize = length(dsize)
    for (i, val) in enumerate(vals)
        (!Gadfly.isconcrete(val) || !isfinite(val)) && continue

        s = size[mod1(i, lensize)]
        ds = dsize[mod1(i, lendsize)]

        minval, maxval = minvalmaxval(minval, maxval, convert(T, val), s, ds)
    end
    minval, maxval
end

function apply_statistic_typed(minval::T, maxval::T, vals::DataArray{T}, size, dsize) where T
    lensize  = length(size)
    lendsize = length(dsize)
    for i = 1:length(vals)
        vals.na[i] && continue

        val::T = vals.data[i]
        s = size[mod1(i, lensize)]
        ds = dsize[mod1(i, lendsize)]

        minval, maxval = minvalmaxval(minval, maxval, val, s, ds)
    end
    minval, maxval
end

function minvalmaxval(minval::T, maxval::T, val, s, ds) where T
    if val < minval || !isfinite(minval)
        minval = val
    end

    if val > maxval || !isfinite(maxval)
        maxval = val
    end

    if s != nothing && typeof(s) <: AbstractFloat
        minval = min(minval, val - s)::T
        maxval = max(maxval, val + s)::T
    end

    if ds != nothing
        minval = min(minval, val - ds)::T
        maxval = max(maxval, val + ds)::T
    end

    minval, maxval
end

struct BoxplotStatistic <: Gadfly.StatisticElement
    method::Union{Symbol, Vector}
end

BoxplotStatistic(; method=:tukey) = BoxplotStatistic(method)

input_aesthetics(stat::BoxplotStatistic) = [:x, :y]
output_aesthetics(stat::BoxplotStatistic) =
    [:x, :middle, :lower_hinge, :upper_hinge, :lower_fence, :upper_fence, :outliers]

const boxplot = BoxplotStatistic

function apply_statistic(stat::BoxplotStatistic,
                         scales::Dict{Symbol, Gadfly.ScaleElement},
                         coord::Gadfly.CoordinateElement,
                         aes::Gadfly.Aesthetics)
    if aes.y === nothing
        Gadfly.assert_aesthetics_defined("BoxplotStatistic", aes,
            :x, :lower_hinge, :upper_hinge, :lower_fence, :upper_fence)

        aes_color = aes.color === nothing ? [nothing] : aes.color
        groups = Any[]
        for (x, c) in zip(aes.x, cycle(aes_color))
            push!(groups, (x, c))
        end

        if aes.color !== nothing
            aes.color = discretize_make_ia([c for (x, c) in groups],
                filter(!ismissing, aes.color.values))
        end

        return
    end

    if aes.x === nothing
        aes_x = [1]
        aes.x_label = x -> fill("", length(x))
    else
        aes_x = aes.x
    end
    aes_color = aes.color === nothing ? [nothing] : aes.color

    T = isempty(aes.y) ? eltype(aes.y) : typeof(aes.y[1] / 1)
    groups = DefaultOrderedDict(() -> T[])

    for (x, y, c) in zip(cycle(aes_x), aes.y, cycle(aes_color))
        push!(groups[(x, c)], y)
    end

    if aes.y != nothing
        m = length(groups)
        aes.x = Array{eltype(aes.x)}(m)
        aes.middle = Array{T}(m)
        aes.lower_hinge = Array{T}(m)
        aes.upper_hinge = Array{T}(m)
        aes.lower_fence = Array{T}(m)
        aes.upper_fence = Array{T}(m)
        aes.outliers = Vector{T}[]

        for (i, ((x, c), ys)) in enumerate(groups)
            sort!(ys)

            aes.x[i] = x

            if stat.method == :tukey
                aes.lower_hinge[i], aes.middle[i], aes.upper_hinge[i] =
                        quantile(ys, [0.25, 0.5, 0.75])
                iqr = aes.upper_hinge[i] - aes.lower_hinge[i]

                idx = searchsortedfirst(ys, aes.lower_hinge[i] - 1.5iqr)
                aes.lower_fence[i] = ys[idx]

                idx = searchsortedlast(ys, aes.upper_hinge[i] + 1.5iqr)
                aes.upper_fence[i] = ys[idx]
            elseif isa(stat.method, Vector)
                qs = stat.method
                if length(qs) != 5
                    error("Stat.boxplot requires exactly five quantiles.")
                end

                aes.lower_fence[i], aes.lower_hinge[i], aes.middle[i],
                aes.upper_hinge[i], aes.upper_fence[i] = quantile!(ys, qs)
            else
                error("Invalid method specified for State.boxplot")
            end

            push!(aes.outliers,
                 filter(y -> y < aes.lower_fence[i] || y > aes.upper_fence[i], ys))
        end
    end

    if length(aes.x) > 1 && (haskey(scales, :x) && isa(scales[:x], Scale.ContinuousScale))
        xmin, xmax = minimum(aes.x), maximum(aes.x)
        minspan = minimum([xj - xi for (xi, xj) in zip(aes.x[1:end-1], aes.x[2:end])])

        xviewmin = xmin - minspan / 2
        xviewmax = xmax + minspan / 2

        if aes.xviewmin === nothing || aes.xviewmin > xviewmin
            aes.xviewmin = xviewmin
        end

        if aes.xviewmax === nothing || aes.xviewmax < xviewmax
            aes.xviewmax = xviewmax
        end
    end

    if isa(aes_x, IndirectArray)
        aes.x = discretize_make_ia(aes.x, aes_x.values)
    end

    if aes.color !== nothing
        aes.color = discretize_make_ia(RGB{Float32}[c for (x, c) in keys(groups)],
            aes.color.values)
    end

    nothing
end

struct SmoothStatistic <: Gadfly.StatisticElement
    method::Symbol
    smoothing::Float64
end
SmoothStatistic(; method=:loess, smoothing=0.75) = SmoothStatistic(method, smoothing)

const smooth = SmoothStatistic

input_aesthetics(::SmoothStatistic) = [:x, :y]
output_aesthetics(::SmoothStatistic) = [:x, :y]

function apply_statistic(stat::SmoothStatistic,
                         scales::Dict{Symbol, Gadfly.ScaleElement},
                         coord::Gadfly.CoordinateElement,
                         aes::Gadfly.Aesthetics)

    Gadfly.assert_aesthetics_defined("Stat.smooth", aes, :x, :y)
    Gadfly.assert_aesthetics_equal_length("Stat.smooth", aes, :x, :y)

    stat.method in [:loess,:lm] ||
            error("The only Stat.smooth methods currently supported are loess and lm.")

    max_num_steps = 750
    aes_color = aes.color === nothing ? [nothing] : aes.color

    groups = Dict(c => (eltype(aes.x)[], eltype(aes.y)[]) for c in unique(aes_color))
    for (x, y, c) in zip(aes.x, aes.y, cycle(aes_color))
        push!(groups[c][1], x)
        push!(groups[c][2], y)
    end

    local xs, ys, xsp
    aes.x = eltype(aes.x)[]
    # For aes.y returning a Float is ok if `y` is an Int or a Float
    # There does not seem to be strong demand for other types of `y`
    aes.y = Float64[]
    colors = eltype(aes_color)[]

    for (c, (xv, yv)) in groups
        x_min, x_max = minimum(xv), maximum(xv)
        x_min == x_max && error("Stat.smooth requires more than one distinct x value")
        try
            xs = Float64.( eltype(xv) <: Dates.TimeType ? Dates.value.(xv) : xv )
            ys = Float64.( eltype(yv) <: Dates.TimeType ? Dates.value.(yv) : yv )
        catch e
            error("Stat.loess and Stat.lm require that x and y be bound to arrays of plain numbers.")
        end

        nudge = 1e-5 * (x_max - x_min)

        dx = (x_max-x_min)*(1/max_num_steps)
        # For a Date, dx might be 0 days, so correct
        # For Ints, correct dx
        if isa(xv[1], Date)
            dx = max(dx, Dates.Day(1))
        elseif isa(xv[1], Int)
            dx = ceil(Int, dx)
            nudge = 0
        end

        steps = collect((x_min + nudge):dx:(x_max - nudge))
        xsp = Float64.( eltype(steps) <: Dates.TimeType ? Dates.value.(steps) : steps )
        if stat.method == :loess
            smoothys = Loess.predict(loess(xs, ys, span=stat.smoothing), xsp)
        elseif stat.method == :lm
            lmcoeff = linreg(xs,ys)
            smoothys = lmcoeff[2].*xsp .+ lmcoeff[1]
        end

    # New aes
        append!(aes.x, steps)
        append!(aes.y, smoothys)
        append!(colors, fill(c, length(steps)))
    end

    if !(aes.color===nothing)
        aes.color = discretize_make_ia(colors)
    end
end

struct HexBinStatistic <: Gadfly.StatisticElement
    xbincount::Int
    ybincount::Int
end
HexBinStatistic(; xbincount=50, ybincount=50) = HexBinStatistic(xbincount, ybincount)

const hexbin = HexBinStatistic

function apply_statistic(stat::HexBinStatistic,
                         scales::Dict{Symbol, Gadfly.ScaleElement},
                         coord::Gadfly.CoordinateElement,
                         aes::Gadfly.Aesthetics)
    xmin, xmax = minimum(aes.x), maximum(aes.x)
    ymin, ymax = minimum(aes.y), maximum(aes.y)
    xspan, yspan = xmax - xmin, ymax - ymin

    xsize = xspan / stat.xbincount
    ysize = yspan / stat.ybincount

    counts = Dict{(Tuple{Int, Int}), Int}()
    for (x, y) in zip(aes.x, aes.y)
        h = convert(HexagonOffsetOddR, cube_round(x - xmin + xspan/2,
                                                  y - ymin + yspan/2,
                                                  xsize, ysize))
        idx = (h.q, h.r)
        if !haskey(counts, idx)
            counts[idx] = 1
        else
            counts[idx] += 1
        end
    end

    N = length(counts)
    aes.x = Array{Float64}(N)
    aes.y = Array{Float64}(N)
    data = Gadfly.Data()
    data.color = Array{Int}(N)
    k = 1
    for (idx, cnt) in counts
        x, y = center(HexagonOffsetOddR(idx[1], idx[2]), xsize, ysize,
                      xmin - xspan/2, ymin - yspan/2)
        aes.x[k] = x
        aes.y[k] = y
        data.color[k] = cnt
        k += 1
    end
    aes.xsize = [xsize]
    aes.ysize = [ysize]

    color_scale = scales[:color]
    typeof(color_scale) <: Scale.ContinuousColorScale ||
            error("HexBinGeometry requires a continuous color scale.")

    aes.color_key_title = "Count"

    Scale.apply_scale(color_scale, [aes], data)
end

default_scales(::HexBinStatistic, t::Gadfly.Theme) = [t.continuous_color_scale]

struct StepStatistic <: Gadfly.StatisticElement
    direction::Symbol
end
StepStatistic(; direction=:hv) = StepStatistic(direction)

const step = StepStatistic

input_aesthetics(::StepStatistic) = [:x, :y]
output_aesthetics(::StepStatistic) = [:x, :y]

function apply_statistic(stat::StepStatistic,
                         scales::Dict{Symbol, Gadfly.ScaleElement},
                         coord::Gadfly.CoordinateElement,
                         aes::Gadfly.Aesthetics)
    Gadfly.assert_aesthetics_defined("StepStatistic", aes, :x)
    Gadfly.assert_aesthetics_defined("StepStatistic", aes, :y)
    Gadfly.assert_aesthetics_equal_length("StepStatistic", aes, :x, :y)

    p = sortperm(aes.x, alg=MergeSort)
    permute!(aes.x, p)
    permute!(aes.y, p)
    aes.group != nothing && permute!(aes.group, p)
    aes.color != nothing && permute!(aes.color, p)

    if aes.group != nothing
        Gadfly.assert_aesthetics_equal_length("StepStatistic", aes, :x, :group)
        permute!(aes.x, p)
        permute!(aes.y, p)
        permute!(aes.group, p)
        aes.color != nothing && permute!(aes.color, p)
    end

    if aes.color != nothing
        Gadfly.assert_aesthetics_equal_length("StepStatistic", aes, :x, :color)
        # TODO: use this when we switch to 0.4
        # sortperm!(p, aes.color, alg=MergeSort, lt=Gadfly.color_isless)
        p = sortperm(aes.color, alg=MergeSort, lt=Gadfly.color_isless)
        permute!(aes.x, p)
        permute!(aes.y, p)
        permute!(aes.color, p)
        aes.group != nothing && permute!(aes.group, p)
    end

    x_step = Array{eltype(aes.x)}(0)
    y_step = Array{eltype(aes.y)}(0)
    color_step = aes.color == nothing ? nothing : Array{eltype(aes.color)}(0)
    group_step = aes.group == nothing ? nothing : Array{eltype(aes.group)}(0)

    i = 1
    i_offset = 1
    while true
        u = i_offset + div(i - 1, 2) + (isodd(i) || stat.direction != :hv ? 0 : 1)
        v = i_offset + div(i - 1, 2) + (isodd(i) || stat.direction != :vh ? 0 : 1)

        (u > length(aes.x) || v > length(aes.y)) && break

        if (aes.color != nothing &&
             (aes.color[u] != aes.color[i_offset] || aes.color[v] != aes.color[i_offset])) ||
           (aes.group != nothing &&
             (aes.group[u] != aes.color[i_offset] || aes.color[v] != aes.group[i_offset]))
            i_offset = max(u, v)
            i = 1
        else
            push!(x_step, aes.x[u])
            push!(y_step, aes.y[v])
            aes.color != nothing && push!(color_step, aes.color[i_offset])
            aes.group != nothing && push!(group_step, aes.group[i_offset])
            i += 1
        end
    end

    aes.x = x_step
    aes.y = y_step
    aes.color = color_step
    aes.group = group_step
end

struct FunctionStatistic <: Gadfly.StatisticElement
    # Number of points to evaluate the function at
    num_samples::Int
end
FunctionStatistic(; num_samples=250) = FunctionStatistic(num_samples)

const func = FunctionStatistic

default_scales(::FunctionStatistic) = [Gadfly.Scale.x_continuous(), Gadfly.Scale.y_continuous()]

input_aesthetics(::FunctionStatistic) = [:y, :xmin, :xmax]
output_aesthetics(::FunctionStatistic) = [:x, :y, :group]

function apply_statistic(stat::FunctionStatistic,
                         scales::Dict{Symbol, Gadfly.ScaleElement},
                         coord::Gadfly.CoordinateElement,
                         aes::Gadfly.Aesthetics)
    Gadfly.assert_aesthetics_defined("FunctionStatistic", aes, :y)
    Gadfly.assert_aesthetics_defined("FunctionStatistic", aes, :xmin)
    Gadfly.assert_aesthetics_defined("FunctionStatistic", aes, :xmax)
    Gadfly.assert_aesthetics_equal_length("FunctionStatistic", aes, :xmin, :xmax)

    aes.x = Array{Float64}(length(aes.y) * stat.num_samples)
    ys = Array{Float64}(length(aes.y) * stat.num_samples)

    i = 1
    for (f, xmin, xmax) in zip(aes.y, cycle(aes.xmin), cycle(aes.xmax))
        for x in linspace(xmin, xmax, stat.num_samples)
            aes.x[i] = x
            ys[i] = f(x)
            i += 1
        end
    end

    # color was bound explicitly
    if aes.color != nothing
        func_color = aes.color
        aes.color = Array{eltype(aes.color)}(length(aes.y) * stat.num_samples)
        groups = DataArray(Int, length(aes.y) * stat.num_samples)
        for i in 1:length(aes.y)
            aes.color[1+(i-1)*stat.num_samples:i*stat.num_samples] = func_color[i]
            groups[1+(i-1)*stat.num_samples:i*stat.num_samples] = i
        end
        aes.group = discretize_make_ia(groups)
    elseif length(aes.y) > 1 && haskey(scales, :color)
        data = Gadfly.Data()
        data.color = Array{AbstractString}(length(aes.y) * stat.num_samples)
        groups = DataArray(Int, length(aes.y) * stat.num_samples)
        for i in 1:length(aes.y)
            fname = "f<sub>$(i)</sub>"
            data.color[1+(i-1)*stat.num_samples:i*stat.num_samples] = fname
            groups[1+(i-1)*stat.num_samples:i*stat.num_samples] = i
        end
        Scale.apply_scale(scales[:color], [aes], data)
        aes.group = discretize_make_ia(groups)
    end

    data = Gadfly.Data()
    data.y = ys
    Scale.apply_scale(scales[:y], [aes], data)
end

struct ContourStatistic <: Gadfly.StatisticElement
    levels::Union{Int,Vector,Function}
    samples::Int
end
ContourStatistic(; levels=15, samples=150) = ContourStatistic(levels, samples)

input_aesthetics(::ContourStatistic) = [:z, :xmin, :xmax, :ymin, :ymax]
output_aesthetics(::ContourStatistic) = [:x, :y, :color, :group]

const contour = ContourStatistic

default_scales(::ContourStatistic, t::Gadfly.Theme=Gadfly.current_theme()) =
        [Gadfly.Scale.z_func(), Gadfly.Scale.x_continuous(), Gadfly.Scale.y_continuous(),
            t.continuous_color_scale]

function apply_statistic(stat::ContourStatistic,
                         scales::Dict{Symbol, Gadfly.ScaleElement},
                         coord::Gadfly.CoordinateElement,
                         aes::Gadfly.Aesthetics)
    xs = aes.x === nothing ? nothing : convert(Vector{Float64}, aes.x)
    ys = aes.y === nothing ? nothing : convert(Vector{Float64}, aes.y)

    if typeof(aes.z) <: Function
        if xs == nothing && aes.xmin != nothing && aes.xmax != nothing
            xs = linspace(aes.xmin[1], aes.xmax[1], stat.samples)
        end

        if ys == nothing && aes.ymin != nothing && aes.ymax != nothing
            ys = linspace(aes.ymin[1], aes.ymax[1], stat.samples)
        end

        zs = Float64[aes.z(x, y) for x in xs, y in ys]

    elseif typeof(aes.z) <: Matrix
        zs = convert(Matrix{Float64}, aes.z)

        if xs == nothing
            xs = collect(Float64, 1:size(zs)[1])
        end
        if ys == nothing
            ys = collect(Float64, 1:size(zs)[2])
        end
        size(zs) != (length(xs), length(ys)) &&
                error("Stat.contour requires dimension of z to be length(x) by length(y)")
    else
        error("Stat.contour requires either a matrix or a function")
    end

    levels = Float64[]
    contour_xs = eltype(xs)[]
    contour_ys = eltype(ys)[]

    stat_levels = typeof(stat.levels) <: Function ? stat.levels(zs) : stat.levels

    groups = discretize_make_ia(Int[])
    group = 0
    for level in Contour.levels(Contour.contours(xs, ys, zs, stat_levels))
        for line in Contour.lines(level)
            xc, yc = Contour.coordinates(line)
            append!(contour_xs, xc)
            append!(contour_ys, yc)
            for _ in 1:length(xc)
                push!(groups, group)
                push!(levels, Contour.level(level))
            end
            group += 1
        end
    end

    aes.group = groups
    color_scale = get(scales, :color, Gadfly.Scale.color_continuous_gradient())
    Scale.apply_scale(color_scale, [aes], Gadfly.Data(color=levels))
    Scale.apply_scale(scales[:x], [aes],  Gadfly.Data(x=contour_xs))
    Scale.apply_scale(scales[:y], [aes], Gadfly.Data(y=contour_ys))
end

struct QQStatistic <: Gadfly.StatisticElement end

input_aesthetics(::QQStatistic) = [:x, :y]
output_aesthetics(::QQStatistic) = [:x, :y]

const qq = QQStatistic

default_scales(::QQStatistic) =
        [Gadfly.Scale.x_continuous(), Gadfly.Scale.y_continuous]

function apply_statistic(stat::QQStatistic,
                         scales::Dict{Symbol, Gadfly.ScaleElement},
                         coord::Gadfly.CoordinateElement,
                         aes::Gadfly.Aesthetics)

    Gadfly.assert_aesthetics_defined("Stat.qq", aes, :x, :y)
    Gadfly.assert_aesthetics_undefined("State.qq", aes, :color)

    # NOTES:
    #
    # apply_scales happens before apply_statistics, so we need to handle in
    # apply_scales the Distributions that might be bound to x and y... By
    # analogy with Stat.func, we can add a check in apply_statistic which defers
    # application.  Stat.func though requires an ARRAY of Functions, and doesn't
    # work on naked functions bound to aes.y.  If we want to bind Distributions,
    # we'd need to extend the types that are allowed for aes.y/.x (e.g. change
    # type of Aesthetics fields x and y).  Right now these are of type
    # NumericalOrCategoricalAesthetic.  The .x and .y fields are the _only_
    # place where this type is used, but I'm not sure if there's a reason that
    # changing this typealias would be a bad idea...for now I've just used a
    # direct `@compat(Union{NumericalOrCategoricalAesthetic, Distribution})`.
    #
    # TODO:
    #
    # Grouping by color etc.?

    # a little helper function to convert either numeric or distribution
    # variables to a format suitable to input to qqbuild.
    toVecOrDist = v -> typeof(v) <: Distribution ? v : convert(Vector{Float64}, v)

    # check and convert :x and :y to proper types for input to qqbuild
    local xs, ys
    try
        (xs, ys) = map(toVecOrDist, (aes.x, aes.y))
    catch e
        error("Stat.qq requires that x and y be bound to either a Distribution or to arrays of plain numbers.")
    end

    qqq = qqbuild(xs, ys)

    aes.x = qqq.qx
    aes.y = qqq.qy

    # apply_scale to Distribution-bound aesthetics is deferred, so re-apply here
    # (but only for Distribution, numeric data is already scaled).  Only one of
    # :x or :y can be a Distribution since qqbuild will throw an error for two
    # Distributions.
    data = Gadfly.Data()
    if typeof(xs) <: Distribution
        data.x = aes.x
        Scale.apply_scale(scales[:x], [aes], data)
    elseif typeof(ys) <: Distribution
        data.y = aes.y
        Scale.apply_scale(scales[:y], [aes], data)
    end
end

struct ViolinStatistic <: Gadfly.StatisticElement
    # Number of points sampled
    n::Int
end
ViolinStatistic() = ViolinStatistic(300)

input_aesthetics(::ViolinStatistic) = [:x, :y, :width]

const violin = ViolinStatistic

function apply_statistic(stat::ViolinStatistic,
                         scales::Dict{Symbol, Gadfly.ScaleElement},
                         coord::Gadfly.CoordinateElement,
                         aes::Gadfly.Aesthetics)

    isa(aes.y[1], Real) || error("Kernel density estimation only works on Real types.")

    grouped_y = Dict(1=>aes.y)
    grouped_color = Dict{Int, Gadfly.ColorOrNothing}(1=>nothing)
    ux = unique(aes.x)
    uxflag = length(ux) < length(aes.x)
    colorflag = aes.color != nothing

    uxflag && (grouped_y = Dict(x=>aes.y[aes.x.==x] for x in ux))

    grouped_color = (colorflag ? Dict(x=>first(aes.color[aes.x.==x]) for x in ux) :
        uxflag && Dict(x=>nothing for x in ux) )

    aes.x     = Array{Float64}(0)
    aes.y     = Array{Float64}(0)
    aes.width = Array{Float64}(0)
    colors = eltype(aes.color)[]

    for (x, ys) in grouped_y
        window = stat.n > 1 ? KernelDensity.default_bandwidth(ys) : 0.1
        f = KernelDensity.kde(ys, bandwidth=window, npoints=stat.n)
        append!(aes.x, fill(x, length(f.x)))
        append!(aes.y, f.x)
        append!(aes.width, f.density)
        append!(colors, fill(grouped_color[x], length(f.x)))
    end

    colorflag  && (aes.color = colors)

    pad = 0.1
    maxwidth = maximum(aes.width)
    broadcast!(*, aes.width, aes.width, 1 - pad)
    broadcast!(/, aes.width, aes.width, maxwidth)
end

struct JitterStatistic <: Gadfly.StatisticElement
    vars::Vector{Symbol}
    range::Float64
    seed::UInt32
end
JitterStatistic(vars; range=0.8, seed=0x0af5a1f7) = JitterStatistic(vars, range, seed)

x_jitter(; range=0.8, seed=0x0af5a1f7) = JitterStatistic([:x], range=range, seed=seed)
y_jitter(; range=0.8, seed=0x0af5a1f7) = JitterStatistic([:y], range=range, seed=seed)

input_aesthetics(stat::JitterStatistic) = stat.vars
output_aesthetics(stat::JitterStatistic) = stat.vars

function minimum_span(vars::Vector{Symbol}, aes::Gadfly.Aesthetics)
    span = nothing
    for var in vars
        data = getfield(aes, var)
        length(data) < 2 && continue
        dataspan = data[2] - data[1]
        T = eltype(data)
        z = convert(T, zero(T))
        sorteddata = sort(data)
        for δ in diff(sorteddata)
            if δ != z && (δ < dataspan || dataspan == z)
                dataspan = δ
            end
        end

        if span == nothing || (dataspan != nothing && dataspan < span)
            span = dataspan
        end
    end

    return span
end

function apply_statistic(stat::JitterStatistic,
                         scales::Dict{Symbol, Gadfly.ScaleElement},
                         coord::Gadfly.CoordinateElement,
                         aes::Gadfly.Aesthetics)
    span = minimum_span(stat.vars, aes)
    span == nothing && return

    rng = MersenneTwister(stat.seed)
    for var in stat.vars
        data = getfield(aes, var)
        outdata = Array{Float64}(size(data))
        broadcast!(+, outdata, data, stat.range * (rand(rng, length(data)) - 0.5) .* span)
        setfield!(aes, var, outdata)
    end
end

# Bin mean returns the mean of x and y in n bins of x
struct BinMeanStatistic <: Gadfly.StatisticElement
    n::Int
end
BinMeanStatistic(; n=20) = BinMeanStatistic(n)

const binmean = BinMeanStatistic

input_aesthetics(::BinMeanStatistic) = [:x, :y]
output_aesthetics(::BinMeanStatistic) = [:x, :y]

function apply_statistic(stat::BinMeanStatistic,
                         scales::Dict{Symbol, Gadfly.ScaleElement},
                         coord::Gadfly.CoordinateElement,
                         aes::Gadfly.Aesthetics)

    Gadfly.assert_aesthetics_defined("Stat.binmean", aes, :x, :y)

    breaks = quantile(aes.x, collect(1:stat.n)/stat.n)

    Tx = eltype(aes.x)
    Ty = eltype(aes.y)

    if aes.color === nothing
        (aes.x, aes.y) = mean_by_group(aes.x, aes.y, breaks)
    else
        groups = Dict()
        for (x, y, c) in zip(aes.x, aes.y, cycle(aes.color))
            if !haskey(groups, c)
                xs = append!(Tx[], collect(Tx, aes.x))
                ys = append!(Ty[], collect(Ty, aes.y))
                groups[c] = Array[xs, ys]
            else
                push!(groups[c][1], x)
                push!(groups[c][2], y)
            end
        end
        colors = Array{RGB{Float32}}(0)
        aes.x = Array{Tx}(0)
        aes.y = Array{Ty}(0)
        for (c, v) in groups
            (fx, fy) = mean_by_group(v[1], v[2], breaks)
            append!(aes.x, fx)
            append!(aes.y, fy)
            for _ in 1:length(fx)
                push!(colors, c)
            end
        end
        aes.color = discretize_make_ia(colors)
    end
end

function mean_by_group(x::Vector{Tx}, y::Vector{Ty}, breaks::Vector{Float64}) where {Tx, Ty}
    count = zeros(Int64, length(breaks))
    totalx = zeros(Tx, length(breaks))
    totaly = zeros(Ty, length(breaks))
    for i in 1:length(x)
        refs = searchsortedfirst(breaks, x[i])
        count[refs] += 1
        totalx[refs] += x[i]
        totaly[refs] += y[i]
    end
    subset = count .> 0
    count = count[subset]
    return (totalx[subset] ./ count, totaly[subset] ./ count)
end

struct EnumerateStatistic <: Gadfly.StatisticElement
    var::Symbol
end

input_aesthetics(stat::EnumerateStatistic) = [stat.var]
output_aesthetics(stat::EnumerateStatistic) = [stat.var]

function default_scales(stat::EnumerateStatistic)
    if stat.var == :y
        return [Gadfly.Scale.y_continuous()]
    elseif stat.var == :x
        return [Gadfly.Scale.x_continuous()]
    else
        return Gadfly.ScaleElement[]
    end
end

const x_enumerate = EnumerateStatistic(:x)
const y_enumerate = EnumerateStatistic(:y)

function apply_statistic(stat::EnumerateStatistic,
                         scales::Dict{Symbol, Gadfly.ScaleElement},
                         coord::Gadfly.CoordinateElement,
                         aes::Gadfly.Aesthetics)
    has_x = aes.x != nothing
    has_y = aes.y != nothing

    if stat.var == :x && !has_x && has_y
        aes.x = collect(1:length(aes.y))
    elseif stat.var == :y && !has_y && has_x
        aes.y = collect(1:length(aes.x))
    end
end

### Vector Field Statistic

struct VecFieldStatistic <: Gadfly.StatisticElement
    smoothness::Float64
    scale::Float64
    samples::Int64
end
VecFieldStatistic(; smoothness=1.0, scale=1.0, samples=20) =
        VecFieldStatistic(smoothness, scale, samples)

input_aesthetics(stat::VecFieldStatistic) = [:z, :x, :y, :color]
output_aesthetics(stat::VecFieldStatistic) = [:x, :y, :xend, :yend, :color]
default_scales(stat::VecFieldStatistic, t::Gadfly.Theme=Gadfly.current_theme()) =
        [Gadfly.Scale.z_func(), Gadfly.Scale.x_continuous(), Gadfly.Scale.y_continuous(),
            t.continuous_color_scale ]

const vectorfield = VecFieldStatistic

function apply_statistic(stat::VecFieldStatistic,
                         scales::Dict{Symbol, Gadfly.ScaleElement},
                         coord::Gadfly.CoordinateElement,
                         aes::Gadfly.Aesthetics)
    xs = aes.x === nothing ? nothing : Float64.(aes.x)
    ys = aes.y === nothing ? nothing : Float64.(aes.y)

    if isa(aes.z, Function)
        if xs == nothing && aes.xmin != nothing && aes.xmax != nothing
            xs = linspace(aes.xmin[1], aes.xmax[1], stat.samples)
        end
        if ys == nothing && aes.ymin != nothing && aes.ymax != nothing
            ys = linspace(aes.ymin[1], aes.ymax[1], stat.samples)
        end

        zs = Float64[aes.z(x, y) for x in xs, y in ys]

    elseif isa(aes.z, Matrix)
        zs = Float64.(aes.z)

        if xs == nothing
            xs = collect(Float64, 1:size(zs)[1])
        end
        if ys == nothing
            ys = collect(Float64, 1:size(zs)[2])
        end
        if size(zs) != (length(xs), length(ys))
            error("Stat.vectorfield requires dimension of z to be length(x) by length(y)")
        end
    else
        error("Stat.vectorfield requires either a matrix or a function")
    end

    X = vcat([[x y] for x in xs, y in ys]...)
    Z = vec(zs)
    # The next two lines make use of the package CoupledFields.jl
    kpars = GaussianKP(X)
    ∇g = hcat(gradvecfield([stat.smoothness -7.0], X, Z[:,1:1], kpars)...)'
    vecf = [X-∇g*stat.scale X+∇g*stat.scale]

    aes.z = nothing
    aes.x = vecf[:,1]
    aes.y = vecf[:,2]
    aes.xend = vecf[:,3]
    aes.yend = vecf[:,4]
    color_scale = get(scales, :color, Gadfly.Scale.color_continuous_gradient())
    Scale.apply_scale(color_scale, [aes], Gadfly.Data(color=Z))

end

### Hair Statistic

struct HairStatistic <: Gadfly.StatisticElement
    intercept
    orientation::Symbol # :horizontal or :vertical like BarStatistic
end
HairStatistic(;intercept=0.0, orientation=:vertical) = HairStatistic(intercept, orientation)

input_aesthetics(stat::HairStatistic) = [:x, :y]

output_aesthetics(stat::HairStatistic) = [:x, :y, :xend, :yend]

default_scales(stat::HairStatistic) = [Gadfly.Scale.x_continuous(), Gadfly.Scale.y_continuous()]

const hair = HairStatistic

function apply_statistic(stat::HairStatistic,
                         scales::Dict{Symbol, Gadfly.ScaleElement},
                         coord::Gadfly.CoordinateElement,
                         aes::Gadfly.Aesthetics)
    if stat.orientation == :vertical
        aes.xend = aes.x
        aes.yend = fill(stat.intercept, length(aes.y))
    else
        aes.yend = aes.y
        aes.xend = fill(stat.intercept, length(aes.x))
    end
end

### Ellipse Statistic

struct EllipseStatistic <: Gadfly.StatisticElement
    distribution::Type{<:ContinuousMultivariateDistribution}
    levels::Vector{<:AbstractFloat}
    nsegments::Int
end

function EllipseStatistic(;
        distribution::(Type{<:ContinuousMultivariateDistribution})=MvNormal,
        levels::Vector{Float64}=[0.95],
        nsegments::Int=51 )
    return EllipseStatistic(distribution, levels, nsegments)
end

Gadfly.input_aesthetics(stat::EllipseStatistic) = [:x, :y]
Gadfly.output_aesthetics(stat::EllipseStatistic) = [:x, :y]
Gadfly.default_scales(stat::EllipseStatistic) = [Gadfly.Scale.x_continuous(), Gadfly.Scale.y_continuous()]

const ellipse = EllipseStatistic

function Gadfly.Stat.apply_statistic(stat::EllipseStatistic,
                         scales::Dict{Symbol, Gadfly.ScaleElement},
                         coord::Gadfly.CoordinateElement,
                         aes::Gadfly.Aesthetics)

    Dat = [aes.x aes.y]
    grouped_xy = Dict(1=>Dat)
    grouped_color = Dict{Int, Gadfly.ColorOrNothing}(1=>nothing)
    colorflag = aes.color != nothing
    aes.group = (colorflag ? aes.color : aes.group)

    if aes.group != nothing
        ug = unique(aes.group)
        grouped_xy = Dict(g=>Dat[aes.group.==g,:] for g in ug)
        grouped_color = Dict(g=>first(aes.group[aes.group.==g]) for g in ug)
    end

    levels = Float64[]
    colors = eltype(aes.color)[]
    ellipse_x = eltype(Dat)[]
    ellipse_y = eltype(Dat)[]

    dfn = 2
    θ = 2π*(0:stat.nsegments)/stat.nsegments
    n = length(θ)
    for (g, data) in grouped_xy
        dfd = size(data,1)-1
        dhat = fit(stat.distribution, data')
        Σ½ = chol(cov(dhat))
        rv = sqrt.(dfn*[quantile(FDist(dfn,dfd), p) for p in stat.levels])
        ellxy =  [cos.(θ) sin.(θ)] * Σ½
        μ = mean(dhat)
        for r in rv
            append!(ellipse_x, r*ellxy[:,1].+μ[1])
            append!(ellipse_y, r*ellxy[:,2].+μ[2])
            append!(colors, fill(grouped_color[g], n))
            append!(levels, fill(r, n))
        end
    end

    aes.group = discretize_make_ia(levels)
    colorflag && (aes.color = colors)
    aes.x = ellipse_x
    aes.y = ellipse_y
end

end # module Stat


"""
    rstar(rng=Random.GLOBAL_RNG, classifier, chains::Chains; kwargs...)

Compute the ``R^*`` convergence diagnostic of the MCMC `chains` with the `classifier`.

The keyword arguments supported here are the same as those in `rstar` for arrays of samples
and chain indices.

# Examples

```jldoctest rstar; setup = :(using Random; Random.seed!(100))
julia> using MLJBase, MLJXGBoostInterface, Statistics

julia> chains = Chains(fill(4.0, 100, 2, 3));
```

One can compute the distribution of the ``R^*`` statistic with the probabilistic classifier.

```jldoctest rstar
julia> distribution = rstar(XGBoostClassifier(), chains);

julia> isapprox(mean(distribution), 1; atol=0.1)
true
```

For deterministic classifiers, a single ``R^*`` statistic is returned.

```jldoctest rstar
julia> @pipeline XGBoostClassifier name = XGBoostDeterministic operation = predict_mode;

julia> value = rstar(XGBoostDeterministic(), chains);

julia> isapprox(value, 1; atol=0.2)
true
```
"""
function MCMCDiagnosticTools.rstar(
    classif::MLJModelInterface.Supervised, chn::Chains; kwargs...
)
    return MCMCDiagnosticTools.rstar(Random.GLOBAL_RNG, classif, chn; kwargs...)
end

function MCMCDiagnosticTools.rstar(
    rng::Random.AbstractRNG, classif::MLJModelInterface.Supervised, chn::Chains; kwargs...
)
    nchains = size(chn, 3)
    nchains <= 1 && throw(DimensionMismatch())

    # collect data
    x = Array(chn)
    y = repeat(chains(chn); inner = size(chn,1))

    return MCMCDiagnosticTools.rstar(rng, classif, x, y; kwargs...)
end


module debug

# Home to all of the routines that are useful for debugging purposes

using ..constants
using ..model
using ..image

import PyPlot
import PyPlot.plt
using LaTeXStrings

export format_grid, synthesize_image, plot_chmaps, plot_slices

"""Format grid values to a standard string"""
function format_grid(grid::Grid, npix::Int)
    #nr, ntheta, r_in, r_out, npix
    r_in = @sprintf("%.1f", grid.rs[1]/AU)
    r_out = @sprintf("%.1f", grid.rs[end]/AU)
    return "$(grid.nr)_$(grid.ntheta)_$(r_in)_$(r_out)_$npix"
end

"""Synthesize images for these parameters, grid, and npix values. dlams provides the frequencies at which to synthesize images. Assumes that you have run DJInitialize.jl inside of homedir. When done, copies the image to outfile."""
function synthesize_image(dlams::Array{Float64, 1}, pars::Parameters, grid::Grid, npix::Int, outfile::AbstractString, species::AbstractString)

    # Assumes that all of the necessary files are already in the current working directory.
    write_grid("", grid)

    # Compute parameter file using model.jl, write to disk in current directory
    write_model(pars, "", grid, species)

    # Doppler shift the dataset wavelengths to rest-frame wavelength
    nchan = length(dlams)
    beta = pars.vel/c_kms # relativistic Doppler formula
    lams = Array(Float64, nchan)
    for i=1:nchan
        lams[i] =  dlams[i] * sqrt((1. - beta) / (1. + beta)) # [microns]
    end

    write_lambda(lams, "") # write into current directory

    # Run RADMC-3D, redirect output to /dev/null
    run(pipeline(`radmc3d image incl $(pars.incl) posang $(pars.PA) npix $npix loadlambda`, DevNull))

    # Copy the output image to the new filename
    run(`mv image.out image_$(outfile).out`)
end

# Plot the channel maps using sky convention
function plot_chmaps(img::image.SkyImage, lam0, fname="channel_maps_sky.png")

    vels = c_kms * (img.lams .- lam0)/lam0

    (im_ny, im_nx) = size(img.data)[1:2] # y and x dimensions of the image

    # Image needs to be flipped along RA dimension
    ext = (img.ra[end], img.ra[1], img.dec[1], img.dec[end])

    nlam = length(img.lams)

    # Figure out how many plots we'll have.
    ncols = 8
    nrows = ceil(Int, nlam/ncols)

    fig, ax = plt[:subplots](nrows=nrows, ncols=ncols, figsize=(12, 1.5 * nrows))

    # max = maximum(log10(img.data))

    vvmax = maxabs(img.data)
    norm = PyPlot.matplotlib[:colors][:Normalize](0, vvmax)

    for row=1:nrows
        for col=1:ncols
            iframe = col + (row - 1) * ncols

            if col != 1 || row != nrows
                ax[row, col][:xaxis][:set_ticklabels]([])
                ax[row, col][:yaxis][:set_ticklabels]([])
            else
                ax[row, col][:set_xlabel](L"$\Delta \alpha$ ('')")
                ax[row, col][:set_ylabel](L"$\Delta \delta$ ('')")
            end

            if iframe > nlam
                # Plot a blank square if we run out of channels
                ax[row, col][:imshow](zeros((im_ny, im_nx)), cmap=plt[:get_cmap]("PuBu"), vmin=0, vmax=20, extent=ext, origin="lower")

            else
                #Flip the frame for Sky convention
                frame = flipdim(img.data[:,:,iframe], 2)

                # frame += 1e-99 #Add a tiny bit so that we don't have log10(0)
                # lframe = log10(frame)
                #
                # ax[row, col][:imshow](lframe, extent=ext, interpolation="none", vmin=max - 6, vmax=max, origin="lower", cmap=plt[:get_cmap]("PuBu"))

                ax[row, col][:imshow](frame, extent=ext, interpolation="none", origin="lower", cmap=plt[:get_cmap]("PuBu"), norm=norm)

                ax[row, col][:annotate](@sprintf("%.1f", vels[iframe]), (0.1, 0.8), xycoords="axes fraction", size=8)
            end
        end
    end
    fig[:subplots_adjust](hspace=0.06, wspace=0.01, top=0.9, bottom=0.1, left=0.05, right=0.95)
    plt[:savefig](fname)
end

function plot_slices(img::image.SkyImage, index, fname="")
end

end #module