module AccuracyBenchmark

using DataFrames
using Distributions
import FITSIO
using StaticArrays
import WCS
import CSV

import ..Config
import ..DeterministicVI
import ..Model
import ..ParallelRun
using ..SDSSIO
using ..Coordinates: angular_separation, match_coordinates

const ARCSEC_PER_DEGREE = 3600
const SDSS_ARCSEC_PER_PIXEL = 0.396
const SDSS_DATA_DIR = joinpath(Pkg.dir("Celeste"), "test", "data")
const STRIPE82_RCF = SDSSIO.RunCamcolField(4263, 5, 119)
const COADD_CATALOG_FITS = joinpath(Pkg.dir("Celeste"), "test", "data", "coadd_for_4263_5_119.fit")
const GALAXY_ONLY_COLUMNS = [:gal_frac_dev, :gal_axis_ratio, :gal_radius_px, :gal_angle_deg]

struct BenchmarkFitsFileNotFound <: Exception
    filename::String
end

struct MatchException <: Exception
    msg::String
end

struct MissingColumnsError <: Exception
    missing_columns::Vector{Symbol}
end

################################################################################
# Load/store a catalog DF from/to a CSV.
################################################################################

CATALOG_COLUMNS = Set([
    :ra,
    :dec,
    :is_star,
    :flux_r_nmgy,
    :color_ug,
    :color_gr,
    :color_ri,
    :color_iz,
    :gal_frac_dev,
    :gal_axis_ratio,
    :gal_radius_px,
    :gal_angle_deg,
])

STDERR_COLUMNS = Set([
    :log_flux_r_stderr,
    :color_ug_stderr,
    :color_gr_stderr,
    :color_ri_stderr,
    :color_iz_stderr,
])

BAD_COADD_OBJID = Set([
    # this object is actually multiple objects,
    # see http://legacysurvey.org/viewer/jpeg-cutout/?ra=0.5636&dec=0.4445&zoom=16&layer=decals-dr3
    8647474692482203853,

    # bright neighbor not accounted for,
    # see http://skyserver.sdss.org/dr10/en/tools/chart/navi.aspx?ra=0.551730166101009&dec=0.48085411481715&scale=0.2
    8647474692482203816,

    #  possibly a quasar; update: decals labels this light source a galaxy!
    # see http://legacysurvey.org/viewer?ra=0.5531&dec=0.4530&zoom=16&layer=decals-dr3&sources
    8647474692482204612,

    # Qusar: http://skyserver.sdss.org/dr10/en/get/SpecById.ashx?id=435863481728657408
    8647474692482204147,
])

function assert_columns_are_present(catalog_df::DataFrame, required_columns::Set{Symbol})
    missing_columns = setdiff(required_columns, Set(names(catalog_df)))
    if !isempty(missing_columns)
        throw(MissingColumnsError([missing_columns...]))
    end
end

function read_catalog(csv_file::String)
    catalog_df = CSV.read(csv_file, rows_for_type_detect=100)
    assert_columns_are_present(catalog_df, CATALOG_COLUMNS)
    catalog_df
end

function write_catalog(filename::String, catalog_df::DataFrame;
                       append_hash=false)
    assert_columns_are_present(catalog_df, CATALOG_COLUMNS)

    if append_hash
        # Serialize the data frame into an array of bytes.
        # (CSV.write(::IO, ...) currently broken, so we use a temp file.
        tmp = tempname()
        CSV.write(tmp, catalog_df)
        data = open(tmp) do f
            read(f)
        end
        rm(tmp)

        # Hash the bytes and add the hash string to the filename.
        hash_string = hex(hash(data))[1:10]
        base, extension = splitext(filename)
        filename = @sprintf("%s_%s%s", base, hash_string, extension)

        # Write out the file
        open(filename, "w") do f
            write(f, data)
        end
    else
        CSV.write(filename, catalog_df)
    end

    return filename
end

################################################################################
# Read various catalogs to a common catalog DF format
################################################################################

"""
Convert between SDSS asinh mags and SDSS flux (nMgy).
See http://www.sdss.org/dr12/algorithms/magnitudes/#asinh.
"""
const ASINH_SOFTENING_PARAMETERS = [
    1.4e-10, # for band 1 = u
    0.9e-10,
    1.2e-10,
    1.8e-10,
    7.4e-10, # for band 5 = z
]

function mag_to_flux(mags::AbstractFloat, band_index::Int)
    b = ASINH_SOFTENING_PARAMETERS[band_index]
     1e9 * 2 * b * sinh(-log(10) / 2.5 * mags - log(b))
end

function flux_to_mag(flux_nmgy::AbstractFloat, band_index::Int)
    b = ASINH_SOFTENING_PARAMETERS[band_index]
    -2.5 / log(10) * (asinh(flux_nmgy * 1e-9 / (2*b)) + log(b))
end

function color_from_fluxes(flux1::AbstractFloat, flux2::AbstractFloat)
    if flux1 <= 0 || flux2 <= 0
        missing
    else
        log(flux2 / flux1)
    end
end

function color_from_mags(mags1::AbstractFloat, band1::Int, mags2::AbstractFloat, band2::Int)
    color_from_fluxes(mag_to_flux(mags1, band1), mag_to_flux(mags2, band2))
end

canonical_angle(gal_angle_deg) = gal_angle_deg - floor(gal_angle_deg / 180) * 180

STRIPE_82_CATALOG_KEYS = [
    :objid, :rerun, :run, :camcol, :field, :flags,
    :ra, :dec, :probpsf,
    :psfmag_u, :psfmag_g, :psfmag_r, :psfmag_i, :psfmag_z,
    :devmag_u, :devmag_g, :devmag_r, :devmag_i, :devmag_z,
    :expmag_u, :expmag_g, :expmag_r, :expmag_i, :expmag_z,
    :fracdev_r,
    :devab_r, :expab_r,
    :devphi_r, :expphi_r,
    :devrad_r, :exprad_r,
    :flags, :is_saturated,
]

function load_fits_table_as_data_frame(filename, extension_index, field_names)
    result_df = DataFrame()
    fits = FITSIO.FITS(filename)
    try
        for key in field_names
            result_df[key] = read(fits[extension_index], string(key))
        end
        return result_df
    finally
        close(fits)
    end
end

function load_stripe82_fits_catalog_as_data_frame(filename, extension_index)
    load_fits_table_as_data_frame(filename, extension_index, STRIPE_82_CATALOG_KEYS)
end

"""
Load Stripe 82 objects into a DataFrame. `fits_filename` should be a FITS file created by running a
CasJobs (skyserver.sdss.org/casjobs/) query on the Stripe82 database. See
https://github.com/jeff-regier/Celeste.jl/wiki/About-SDSS-and-Stripe-82.
"""
function load_coadd_catalog(fits_filename)
    raw_df = load_stripe82_fits_catalog_as_data_frame(fits_filename, 2)

    usedev = raw_df[:fracdev_r] .> 0.5  # true=> use dev, false=> use exp
    dev_or_exp(dev_column, exp_column) = ifelse.(usedev, raw_df[dev_column], raw_df[exp_column])
    is_star = [x != 0 for x in raw_df[:probpsf]]
    function star_or_galaxy(star_column, galaxy_dev_column, galaxy_exp_column)
        ifelse.(is_star, raw_df[star_column], dev_or_exp(galaxy_dev_column, galaxy_exp_column))
    end

    mag_u = star_or_galaxy(:psfmag_u, :devmag_u, :expmag_u)
    mag_g = star_or_galaxy(:psfmag_g, :devmag_g, :expmag_g)
    mag_r = star_or_galaxy(:psfmag_r, :devmag_r, :expmag_r)
    mag_i = star_or_galaxy(:psfmag_i, :devmag_i, :expmag_i)
    mag_z = star_or_galaxy(:psfmag_z, :devmag_z, :expmag_z)

    result = DataFrame()
    result[:objid] = [string(objid) for objid in raw_df[:objid]]
    result[:ra] = raw_df[:ra]
    result[:dec] = raw_df[:dec]
    result[:is_star] = is_star

    flux_r = mag_to_flux.(mag_r, 3)
    result[:flux_r_nmgy] = ifelse.(flux_r .> 0, flux_r, missing)

    result[:color_ug] = color_from_mags.(mag_u, 1, mag_g, 2)
    result[:color_gr] = color_from_mags.(mag_g, 2, mag_r, 3)
    result[:color_ri] = color_from_mags.(mag_r, 3, mag_i, 4)
    result[:color_iz] = color_from_mags.(mag_i, 4, mag_z, 5)

    # gal shape -- fracdev
    result[:gal_frac_dev] = raw_df[:fracdev_r]

    # Note that the SDSS photo pipeline doesn't constrain the de Vaucouleur
    # profile parameters and exponential disk parameters (A/B, angle, scale)
    # to be the same, whereas Celeste does. Here, we pick one or the other
    # from SDSS, based on fracdev - we'll get the parameters corresponding
    # to the dominant component. Later, we limit comparison to objects with
    # fracdev close to 0 or 1 to ensure that we're comparing apples to apples.

    result[:gal_axis_ratio] = dev_or_exp(:devab_r, :expab_r)

    # gal effective radius (re)
    re_arcsec = dev_or_exp(:devrad_r, :exprad_r)
    re_pixel = re_arcsec ./ SDSS_ARCSEC_PER_PIXEL
    result[:gal_radius_px] = convert(Vector{Float64}, re_pixel)

    # gal angle (degrees)
    raw_phi = dev_or_exp(:devphi_r, :expphi_r)
    result[:gal_angle_deg] = canonical_angle.(raw_phi)

    is_saturated = raw_df[:is_saturated] .!= 0
    result = result[.!is_saturated, :]
    bad_rows = [x in BAD_COADD_OBJID for x in result[:, :objid]]
    result = result[.!bad_rows, :]

    # for stars, ensure galaxy-only fields are "missing"
    for col in GALAXY_ONLY_COLUMNS
        result[col] = convert(Vector{Union{Missing, Float64}}, result[col])
        result[result[:is_star], col] = missing
    end

    return result
end

function object_dict_to_data_frame(objects::Dict)
    result_df = DataFrame()
    for (key, values) in objects
        result_df[Symbol(key)] = values
    end
    result_df
end

"""
Load the SDSS photoObj catalog used to initialize celeste, and reformat column
names to match what the rest of the scoring code expects.
"""
function load_primary(rcf::SDSSIO.RunCamcolField, stagedir::String)
    dataset = SDSSDataSet(stagedir)
    raw_df = object_dict_to_data_frame(SDSSIO.read_photoobj(dataset, rcf))

    usedev = raw_df[:frac_dev] .> 0.5  # true=> use dev, false=> use exp
    dev_or_exp(dev_column, exp_column) = ifelse.(usedev, raw_df[dev_column], raw_df[exp_column])
    function star_or_galaxy(star_column, galaxy_dev_column, galaxy_exp_column)
        ifelse.(raw_df[:is_star], raw_df[star_column], dev_or_exp(galaxy_dev_column, galaxy_exp_column))
    end

    flux_u = star_or_galaxy(:psfflux_u, :devflux_u, :expflux_u)
    flux_g = star_or_galaxy(:psfflux_g, :devflux_g, :expflux_g)
    flux_r = star_or_galaxy(:psfflux_r, :devflux_r, :expflux_r)
    flux_i = star_or_galaxy(:psfflux_i, :devflux_i, :expflux_i)
    flux_z = star_or_galaxy(:psfflux_z, :devflux_z, :expflux_z)

    result = DataFrame()
    result[:objid] = raw_df[:objid]
    result[:ra] = raw_df[:ra]
    result[:dec] = raw_df[:dec]
    result[:is_star] = raw_df[:is_star]

    result[:flux_r_nmgy] = flux_r

    result[:color_ug] = color_from_fluxes.(flux_u, flux_g)
    result[:color_gr] = color_from_fluxes.(flux_g, flux_r)
    result[:color_ri] = color_from_fluxes.(flux_r, flux_i)
    result[:color_iz] = color_from_fluxes.(flux_i, flux_z)

    result[:gal_frac_dev] = raw_df[:frac_dev]
    result[:gal_axis_ratio] = dev_or_exp(:ab_dev, :ab_exp)

    # gal effective radius (re)
    re_arcsec = dev_or_exp(:theta_dev, :theta_exp)
    re_pixel = re_arcsec ./ SDSS_ARCSEC_PER_PIXEL
    result[:gal_radius_px] = convert(Vector{Float64}, re_pixel)

    # gal angle (degrees)
    raw_phi = dev_or_exp(:phi_dev, :phi_exp)
    result[:gal_angle_deg] = canonical_angle.(raw_phi)

    # primary is better at indicating oversaturation than coadd
    is_saturated = flux_to_mag.(raw_df[:psfflux_r], 3) .< 16
    result = result[.!is_saturated, :]

    return result
end

function fluxes_from_colors(flux_r_nmgy::Float64, colors::AbstractVector)
    @assert length(colors) == 4
    color_ratios = exp.(colors)
    fluxes = similar(colors, 5)
    fluxes[3] = flux_r_nmgy
    fluxes[4] = fluxes[3] * color_ratios[3]
    fluxes[5] = fluxes[4] * color_ratios[4]
    fluxes[2] = fluxes[3] / color_ratios[2]
    fluxes[1] = fluxes[2] / color_ratios[1]
    fluxes
end

function get_median_fluxes(variational_params::Vector{Float64}, source_type::Int64)
    fluxes_from_colors(
        exp(variational_params[Model.ids.flux_loc[source_type]]),
        variational_params[Model.ids.color_mean[:, source_type]],
    )
end

function variational_parameters_to_data_frame_row(variational_params::Vector{Float64})
    ids = Model.ids
    result = DataFrame()
    result[:ra] = variational_params[ids.pos[1]]
    result[:dec] = variational_params[ids.pos[2]]
    result[:is_star] = variational_params[ids.is_star[1, 1]]
    result[:gal_frac_dev] = variational_params[ids.gal_frac_dev]
    result[:gal_axis_ratio] = variational_params[ids.gal_axis_ratio]
    result[:gal_radius_px] = (
        variational_params[ids.gal_radius_px] * sqrt(variational_params[ids.gal_axis_ratio])
    )
    result[:gal_angle_deg] = canonical_angle(180 / pi * variational_params[ids.gal_angle])

    star_galaxy_index = (result[1, :is_star] > 0.5 ? 1 : 2)
    fluxes = get_median_fluxes(variational_params, star_galaxy_index)
    result[:flux_r_nmgy] = fluxes[3]
    result[:color_ug] = color_from_fluxes(fluxes[1], fluxes[2])
    result[:color_gr] = color_from_fluxes(fluxes[2], fluxes[3])
    result[:color_ri] = color_from_fluxes(fluxes[3], fluxes[4])
    result[:color_iz] = color_from_fluxes(fluxes[4], fluxes[5])

    result[:log_flux_r_stderr] = sqrt(variational_params[ids.flux_scale[star_galaxy_index]])
    result[:color_ug_stderr] = sqrt(variational_params[ids.color_var[1, star_galaxy_index]])
    result[:color_gr_stderr] = sqrt(variational_params[ids.color_var[2, star_galaxy_index]])
    result[:color_ri_stderr] = sqrt(variational_params[ids.color_var[3, star_galaxy_index]])
    result[:color_iz_stderr] = sqrt(variational_params[ids.color_var[4, star_galaxy_index]])

    result
end

"""
Convert Celeste results to a dataframe.
"""
function celeste_to_df(results::Vector{ParallelRun.OptimizedSource})
    rows = []
    for result in results
        if !result.is_sky_bad
            row = variational_parameters_to_data_frame_row(result.vs)
            push!(rows, row)
        end
    end
    vcat(rows...)
end

################################################################################
# Generate a random catalog from the Celeste prior
################################################################################

# This is the ratio of stars derived from the catalog used the generate the prior; the value 0.99 is
# currently used in inference to account for the extra flexibility of the galaxy model
const PRIOR_PROBABILITY_OF_STAR = 0.28

"""
Draw a random source from the Celeste prior, returning a 1-row data frame.
"""
function draw_source_params(prior)
    is_star = (rand(Bernoulli(PRIOR_PROBABILITY_OF_STAR)) == 1)
    source_type_index = is_star ? 1 : 2
    flux_r_nmgy = exp(rand(
        Normal(prior.flux_mean[source_type_index], sqrt(prior.flux_var[source_type_index]))
    ))

    color_mixture_weights = prior.k[:, source_type_index]
    NUM_COLOR_COMPONENTS = length(color_mixture_weights)
    color_components = MvNormal[
        MvNormal(prior.color_mean[:, k, source_type_index], prior.color_cov[:, :, k, source_type_index])
        for k in 1:NUM_COLOR_COMPONENTS
    ]
    colors = rand(MixtureModel(color_components, color_mixture_weights))

    if !is_star
        gal_radius_px = exp(rand(
            Normal(prior.gal_radius_px_mean, sqrt(prior.gal_radius_px_var))
        ))
        gal_angle_deg = rand(Uniform(0, 180))
        gal_axis_ratio = rand(Beta(2, 2))
        gal_frac_dev = rand(Beta(0.5, 0.5))
    else
        gal_radius_px = -1
        gal_angle_deg = -1
        gal_axis_ratio = -1
        gal_frac_dev = -1
    end

    # We should change ra/dec ranges below if we start using another field
    @assert STRIPE82_RCF.run == 4263 && STRIPE82_RCF.camcol == 5 && STRIPE82_RCF.field == 119

    # Use approximate size of SDSS field in degrees
    ra = rand(Uniform(0.443, 0.606))
    dec = rand(Uniform(0.411, 0.635))

    DataFrame(
        ra=ra,
        dec=dec,
        is_star=is_star,
        flux_r_nmgy=flux_r_nmgy,
        color_ug=colors[1],
        color_gr=colors[2],
        color_ri=colors[3],
        color_iz=colors[4],
        gal_frac_dev=gal_frac_dev,
        gal_axis_ratio=gal_axis_ratio,
        gal_radius_px=gal_radius_px,
        gal_angle_deg=gal_angle_deg,
    )
end

"""
Draw sources at random from Celeste prior, returning a catalog DF.
"""
function generate_catalog_from_celeste_prior(num_sources::Int64, seed::Int64)
    srand(seed)
    prior = Model.load_prior()
    result = vcat(
        [draw_source_params(prior) for index in 1:num_sources]...
    )

    # for stars, ensure galaxy-only fields are "missing"
    for col in GALAXY_ONLY_COLUMNS
        result[col] = convert(Vector{Union{Missing, Float64}}, result[col])
        result[result[:is_star], col] = missing
    end

    return result
end

################################################################################
# Support for running Celeste on test imagery
################################################################################

## Load a multi-extension FITS imagery file

struct FitsImage
    pixels::Matrix{Float32}
    header::FITSIO.FITSHeader
    wcs::WCS.WCSTransform
end

function read_fits(filename::String)
    if !isfile(filename)
        throw(BenchmarkFitsFileNotFound(filename))
    end

    fits = FITSIO.FITS(filename)
    try
        map(fits) do extension
            pixels = read(extension)
            header = FITSIO.read_header(extension)
            wcs = WCS.from_header(FITSIO.read_header(fits[1], String))[1]
            FitsImage(pixels, header, wcs)
        end
    finally
        close(fits)
    end
end

## Create Celeste images from FITS imagery

function make_simple_psf(psf_sigma_px)
    alphaBar = [1.; 0.]
    xiBar = [0.; 0.]
    tauBar = [psf_sigma_px^2 0.; 0. psf_sigma_px^2]
    [
        Model.PsfComponent(
            alphaBar[k],
            StaticArrays.SVector{2, Float64}(xiBar),
            StaticArrays.SMatrix{2, 2, Float64, 4}(tauBar)
        )
        for k in 1:2
    ]
end

function make_psf_from_header(header::FITSIO.FITSHeader)
    if haskey(header, "CLSIGMA")
        return make_simple_psf(header["CLSIGMA"])
    end

    num_components = header["PSFNCOMP"]
    map(1:num_components) do index
        get(field) = header["PSF$(index)_$(field)"]
        alpha_bar = get("A")
        xi_bar = Float64[get("X1"), get("X2")]
        tau_bar = Float64[get("T11") get("T12"); get("T12") get("T22")]
        Model.PsfComponent(
            alpha_bar,
            StaticArrays.SVector{2, Float64}(xi_bar),
            StaticArrays.SMatrix{2, 2, Float64, 4}(tau_bar),
        )
    end
end

function serialize_psf_to_header(psf::Vector{Model.PsfComponent}, header::FITSIO.FITSHeader)
    header["PSFNCOMP"] = length(psf)
    for (index, component) in enumerate(psf)
        @assert component.tauBar[1, 2] == component.tauBar[2, 1]
        function set(field, value)
            header["PSF$(index)_$(field)"] = value
        end
        set("A", component.alphaBar)
        set("X1", component.xiBar[1])
        set("X2", component.xiBar[2])
        set("T11", component.tauBar[1, 1])
        set("T12", component.tauBar[1, 2])
        set("T22", component.tauBar[2, 2])
    end
end

function make_image(pixels::Matrix{Float32}, band_index::Int,
                    wcs::WCS.WCSTransform, psf::Vector{Model.PsfComponent},
                    sky_level_nmgy::Real, nelec_per_nmgy::Real)
    nx, ny = size(pixels)

    # Render the PSF on a grid, to be used as a (spatially constant) PSF map.
    psfstamp = Model.render_psf(psf, (51, 51))

    Model.Image(
        pixels,
        band_index,
        wcs,
        psf,
        fill(Float32(sky_level_nmgy), nx, ny),
        fill(Float32(nelec_per_nmgy), nx),
        Model.ConstantPSFMap(psfstamp)
    )
end

function make_images(band_extensions::Vector{FitsImage})
    map(enumerate(band_extensions)) do pair
        band_index, extension = pair
        make_image(
            extension.pixels,
            band_index,
            extension.wcs,
            make_psf_from_header(extension.header),
            extension.header["CLSKY"]::Float64,
            extension.header["CLIOTA"]::Float64
        )
    end
end

## Create an initialization catalog for Celeste

function typical_band_fluxes(is_star::Bool)
    source_type_index = is_star ? 1 : 2
    prior_parameters::Model.PriorParams = Model.load_prior()
    # this is the mode. brightness is log normal.
    flux_r = exp(
        prior_parameters.flux_mean[source_type_index] - prior_parameters.flux_var[source_type_index]
    )
    # Band relative intensities are a mixture of lognormals. Which mixture component has the most
    # weight?
    dominant_component = indmax(prior_parameters.k[:, source_type_index])
    # What are the most typical log relative intensities for that component?
    colors = (
        prior_parameters.color_mean[:, dominant_component, source_type_index]
        - diag(prior_parameters.color_cov[:, :, dominant_component, source_type_index])
    )
    fluxes_from_colors(flux_r, colors)
end

function make_catalog_entry(x_position_world_coords::Float64,
                            y_position_world_coords::Float64)
    Model.CatalogEntry(
        [x_position_world_coords, y_position_world_coords],
        false, # is_star
        typical_band_fluxes(true),
        typical_band_fluxes(false),
        0.1, # gal_frac_dev
        0.7, # gal_axis_ratio
        pi / 4, # gal_angle
        4., # gal_radius_px
    )
end

ensure_small_flux(value) = (ismissing(value) || value <= 0) ? 1e-6 : value

na_to_default(value, default) = ismissing(value) ? default : value

function make_catalog_entry(row::DataFrameRow)
    colors = Union{Float64, Missing}[row[:color_ug], row[:color_gr],
                                     row[:color_ri], row[:color_iz]]
    fluxes = fluxes_from_colors(row[:flux_r_nmgy], colors)
    fluxes = ensure_small_flux.(fluxes)
    gal_axis_ratio = na_to_default(row[:gal_axis_ratio], 0.8)
    Model.CatalogEntry(
        [row[:ra], row[:dec]],
        row[:is_star] > 0.5,
        fluxes,
        fluxes,
        na_to_default(row[:gal_frac_dev], 0.5),
        gal_axis_ratio,
        na_to_default(row[:gal_angle_deg], 0.) / 180.0 * pi,
        na_to_default(row[:gal_radius_px] / sqrt(gal_axis_ratio), 2.)
        )
end

function make_initialization_catalog(catalog::DataFrame, use_full_initialzation::Bool)
    position_offset_width = SDSS_ARCSEC_PER_PIXEL / ARCSEC_PER_DEGREE # 1 pixel, in degrees
    map(eachrow(catalog)) do row
        if use_full_initialzation
            make_catalog_entry(row)
        else
            make_catalog_entry(
                row[:ra] + position_offset_width,
                row[:dec] - 0.5 * position_offset_width
            )
        end
    end
end

################################################################################
# Support for generating imagery using Synthetic.jl
################################################################################

struct ImageGeometry
    height_px::Int64
    width_px::Int64
    world_coordinate_origin::Tuple{Float64, Float64} # (ra, dec)
    ra_degrees_per_pixel::Float64
    dec_degrees_per_pixel::Float64
end

function get_image_geometry(catalog_data::DataFrame; field_expand_arcsec=20.0)
    min_ra_deg = minimum(catalog_data[:ra])
    max_ra_deg = maximum(catalog_data[:ra])
    min_dec_deg = minimum(catalog_data[:dec])
    max_dec_deg = maximum(catalog_data[:dec])

    width_arcsec = (max_ra_deg - min_ra_deg) * ARCSEC_PER_DEGREE + 2 * field_expand_arcsec
    height_arcsec = (max_dec_deg - min_dec_deg) * ARCSEC_PER_DEGREE + 2 * field_expand_arcsec
    width_px = convert(Int64, round(width_arcsec / SDSS_ARCSEC_PER_PIXEL))
    height_px = convert(Int64, round(height_arcsec / SDSS_ARCSEC_PER_PIXEL))

    dec_degrees_per_pixel = SDSS_ARCSEC_PER_PIXEL / ARCSEC_PER_DEGREE
    ra_degrees_per_pixel = dec_degrees_per_pixel / cosd(min_dec_deg)

    ImageGeometry(
        height_px,
        width_px,
        (
            min_ra_deg - field_expand_arcsec / ARCSEC_PER_DEGREE,
            min_dec_deg - field_expand_arcsec / ARCSEC_PER_DEGREE,
        ),
        ra_degrees_per_pixel,
        dec_degrees_per_pixel,
    )
end

function make_template_images(
    catalog_data::DataFrame, psf_sigma_px::Float64, band_sky_level_nmgy::Vector{Float64},
    band_nelec_per_nmgy::Vector{Float64}
)
    @assert length(band_sky_level_nmgy) == 5
    @assert length(band_nelec_per_nmgy) == 5
    geometry = get_image_geometry(catalog_data)
    println("  Image dimensions $(geometry.height_px) H x $(geometry.width_px) W px")
    wcs = WCS.WCSTransform(
        2, # dimensions
        # reference pixel coordinates...
        crpix=[1., 1.],
        # ...and corresponding reference world coordinates
        crval=[geometry.world_coordinate_origin[1], geometry.world_coordinate_origin[2]],
        # this WCS is a simple linear transformation
        ctype=["RA---TAN", "DEC--TAN"],
        cunit=["deg", "deg"],
        # these are [du/dx du/dy; dv/dx dv/dy]. (u, v) = world coords, (x, y) = pixel coords.
        pc=[0. geometry.ra_degrees_per_pixel; geometry.dec_degrees_per_pixel 0.],
    )
    map(1:5) do band
        make_image(
            zeros(Float32, (geometry.height_px, geometry.width_px)),
            band,
            wcs,
            make_simple_psf(psf_sigma_px),
            band_sky_level_nmgy[band],
            band_nelec_per_nmgy[band],
        )
    end
end

## Parse a FITS header from a raw string representation

# It's silly that we have to do this, but the Julia FITSIO library doesn't support this (it can only
# read a header from a file), while the Julia WCS library will only generate a header in raw string
# form. (The Julia FITSIO library also won't write a raw string header to a file, so scrap that idea
# :).) Fortunately, the FITS header format is a simple fixed-width textual format with only a few
# data types. See https://fits.gsfc.nasa.gov/fits_primer.html.

function parse_fits_header_value(raw_value::AbstractString)
    if raw_value == "T"
        true
    elseif raw_value == "F"
        false
    elseif raw_value[1] == '\''
        convert(String, raw_value[2:(end - 1)])
    elseif ismatch(r"^[0-9]+$", raw_value)
        parse(Int64, raw_value)
    else
        parse(Float64, raw_value)
    end
end

function parse_fits_header_from_string(header_string::AbstractString)
    header = FITSIO.FITSHeader(String[], [], String[])
    start_position = 1
    while start_position < length(header_string)
        field_string = header_string[start_position:(start_position + 79)]
        field_name = strip(field_string[1:8])
        if length(field_name) != 0 && field_string[9:10] == "= "
            field_parts = split(field_string[11:80], "/", limit=2)
            raw_value = field_parts[1]
            header[field_name] = parse_fits_header_value(strip(raw_value))
            if length(field_parts) == 2
                comment = strip(field_parts[2])
                if length(comment) > 0
                    FITSIO.set_comment!(header, field_name, convert(String, comment))
                end
            end
        end
        start_position += 80
    end
    header
end

function save_images_to_fits(filename::String, images::Vector{<:Model.Image})
    println("Writing images to $filename...")
    fits_file = FITSIO.FITS(filename, "w")
    for band_image in images
        header = parse_fits_header_from_string(WCS.to_header(band_image.wcs))
        serialize_psf_to_header(band_image.psf, header)
        header["CLSKY"] = mean(band_image.sky.sky_small) * mean(band_image.sky.calibration)
        FITSIO.set_comment!(header, "CLSKY", "Mean sky background per pixel, nMgy")
        header["CLIOTA"] = mean(band_image.nelec_per_nmgy)
        FITSIO.set_comment!(header, "CLIOTA", "Gain, nelec per nMgy")
        write(
            fits_file,
            band_image.pixels,
            header=header,
            name="CELESTE_FIELD_$(band_image.b)",
        )
    end
    close(fits_file)
end

################################################################################
# Score a set of predictions against a ground truth catalog
################################################################################

COLOR_COLUMNS = [:color_ug, :color_gr, :color_ri, :color_iz]

ABSOLUTE_ERROR_COLUMNS = vcat(
    [:gal_frac_dev, :gal_axis_ratio, :gal_radius_px],
    COLOR_COLUMNS,
)

function degrees_to_diff(a, b)
    angle_between = abs.(a - b) .% 180
    min.(angle_between, 180 .- angle_between)
end

"""
Given two results data frame, one containing ground truth (i.e Coadd)
and one containing predictions (i.e., either Primary of Celeste),
compute an a data frame containing each prediction's error.
(It's not an average of the errors, it's each error.)
Let's call the return type of this function an \"error data frame\".
"""
function get_error_df(truth::DataFrame, predicted::DataFrame)
    errors = DataFrame()

    predicted_galaxy = predicted[:is_star] .< .5
    true_galaxy = truth[:is_star] .< .5
    errors[:missed_stars] = ifelse.(.!true_galaxy, predicted_galaxy, missing)
    errors[:missed_galaxies] = ifelse.(true_galaxy, .!predicted_galaxy, missing)

    errors[:position] =
        (ARCSEC_PER_DEGREE / SDSS_ARCSEC_PER_PIXEL) .*
        angular_separation.(truth[:ra],
                            truth[:dec],
                            predicted[:ra],
                            predicted[:dec])

    # compare flux in both mags and nMgy for now
    errors[:flux_r_mag] = abs.(
        flux_to_mag.(truth[:flux_r_nmgy], 3)
        .- flux_to_mag.(predicted[:flux_r_nmgy], 3)
    )
    errors[:flux_r_nmgy] = abs.(
        truth[:flux_r_nmgy] .- predicted[:flux_r_nmgy]
    )
    errors[:gal_angle_deg] = degrees_to_diff(truth[:gal_angle_deg], predicted[:gal_angle_deg])

    for column_symbol in ABSOLUTE_ERROR_COLUMNS
        errors[column_symbol] = abs.(truth[column_symbol] .- predicted[column_symbol])
    end
    for color_column in COLOR_COLUMNS
        # to match up with Stripe82Score, which used differences of mags
        errors[color_column] .*= 2.5 / log(10)
    end

    errors
end

function is_good_row(truth_row::DataFrameRow, error_row::DataFrameRow, column_name::Symbol)
    if ismissing(error_row[column_name]) || isnan(error_row[column_name])
        return false
    elseif !ismissing(truth_row[:gal_radius_px]) && truth_row[:gal_radius_px] > 20
        return false
    end

    if column_name in [:gal_axis_ratio, :gal_radius_px, :gal_angle_deg,
                       :gal_frac_dev]
        has_mixture_weight = !ismissing(truth_row[:gal_frac_dev])
        if has_mixture_weight && (0.05 < truth_row[:gal_frac_dev] < 0.95)
            return false
        end
    end
    if column_name == :gal_angle_deg
        if !ismissing(truth_row[:gal_axis_ratio]) && truth_row[:gal_axis_ratio] > .6
            return false
        end
    end
    return true
end

function filter_rows(truth::DataFrame, errors::DataFrame, column_name::Symbol)
    map(zip(eachrow(truth), eachrow(errors))) do rows
        (truth_row, error_row) = rows
        is_good_row(truth_row, error_row, column_name)
    end
end

function score_column(errors::AbstractVector)
    DataFrame(
        N=length(errors),
        first=mean(errors),
    )
end

function score_column(first_errors::AbstractVector,
                      second_errors::AbstractVector)
    @assert length(first_errors) == length(second_errors)
    scores = score_column(first_errors)
    scores[:second] = mean(second_errors)
    diffs = first_errors .- second_errors
    scores[:diff] = mean(diffs)
    scores[:diff_sd] = std(abs.(diffs)) / sqrt(length(diffs))
    scores
end

function get_scores_df(
    truth::DataFrame, first_errors::DataFrame, second_errors::Nullable{DataFrame}
)
    score_rows = DataFrame[]
    for column_name in names(first_errors)
        good_row = filter_rows(truth, first_errors, column_name)
        if !isnull(second_errors)
            good_row .&= filter_rows(truth, get(second_errors), column_name)
        end
        if sum(good_row) <= 1
            continue
        else
            if isnull(second_errors)
                row = score_column(first_errors[good_row, column_name])
            else
                row = score_column(
                    first_errors[good_row, column_name],
                    get(second_errors)[good_row, column_name],
                )
            end

            row[:field] = column_name
            push!(score_rows, row)
        end
    end
    vcat(score_rows...)
end

"""
    match_catalogs(truth, predictions)

Return subset of rows in `truth` and `predictions`, such that each entry in
`truth` has a corresponding match in all `predictions`. A "match" is the
closest object within `tol` degrees. The output data frames all have the same
number of rows, and rows correspond to matching objects.
"""
function match_catalogs(truth::DataFrame, prediction_dfs::Vector{DataFrame};
                        tol::Float64=(0.396 / 3600.))

    matched = trues(nrow(truth))
    idxs = Vector{Int}[]
    for prediction in prediction_dfs
        # remove missings before feeding to match_coordinates
        @assert sum(ismissing.(truth[:ra])) == 0
        @assert sum(ismissing.(truth[:dec])) == 0
        @assert sum(ismissing.(prediction[:ra])) == 0
        @assert sum(ismissing.(prediction[:dec])) == 0
        idx, dists = match_coordinates(Vector{Float64}(truth[:ra]),
                                       Vector{Float64}(truth[:dec]),
                                       Vector{Float64}(prediction[:ra]),
                                       Vector{Float64}(prediction[:dec]))
        matched .&= dists .< tol
        push!(idxs, idx)
    end

    # For each row in `truth`, `matched` is now true if and only if all
    # predictions have a close match.
    #
    # `idxs[i]` is same length as `truth` and contains the index in
    # `prediction_dfs[i]` that is the closest match to each row in
    # `truth`.
    matched_truth = truth[matched, :]
    matched_predictions = [prediction[idx[matched], :]
                           for (prediction, idx) in zip(prediction_dfs, idxs)]

    return matched_truth, matched_predictions
end

function score_predictions(truth::DataFrame, prediction_dfs::Vector{DataFrame})
    matched_truth, matched_prediction_dfs = match_catalogs(truth, prediction_dfs)
    error_dfs = [get_error_df(matched_truth, predictions) for predictions in matched_prediction_dfs]
    @assert length(prediction_dfs) <= 2
    get_scores_df(
        matched_truth,
        error_dfs[1],
        length(error_dfs) > 1 ? Nullable(error_dfs[2]) : Nullable{DataFrame}(),
    )
end

function get_uncertainty_df(truth::DataFrame, predictions::DataFrame)
    assert_columns_are_present(predictions, STDERR_COLUMNS)
    matched_truth, matched_prediction_dfs = match_catalogs(truth, [predictions])
    matched_predictions = matched_prediction_dfs[1]

    valid_rows = (matched_truth[:flux_r_nmgy] .> 0)
    matched_truth = matched_truth[valid_rows, :]
    matched_predictions = matched_predictions[valid_rows, :]

    get_errors(column, map_fn) =
        map_fn.(matched_predictions[column]) .- map_fn.(matched_truth[column])
    get_errors(column) = get_errors(column, x -> x)
    errors = [
        get_errors(:flux_r_nmgy, log),
        get_errors(:color_ug),
        get_errors(:color_gr),
        get_errors(:color_ri),
        get_errors(:color_iz),
    ]
    std_errs = [
        matched_predictions[:log_flux_r_stderr],
        matched_predictions[:color_ug_stderr],
        matched_predictions[:color_gr_stderr],
        matched_predictions[:color_ri_stderr],
        matched_predictions[:color_iz_stderr],
    ]
    names = [:log_flux_r_nmgy, :color_ug, :color_gr,
             :color_ri, :color_iz]

    matched_truth[:log_flux_err] = errors[1]
    matched_truth[:log_flux_stderr] = std_errs[1]
    matched_truth[:flux_r_celeste] = matched_predictions[:flux_r_nmgy]
    CSV.write("stderr.csv", matched_truth)

    mapreduce(vcat, zip(names, errors, std_errs)) do values
        name, error, std_err = values
        DataFrame(name=fill(name, length(error)),
                  error=error,
                  posterior_std_err=std_err)
    end
end

function score_uncertainty(uncertainty_df::DataFrame)
    mapreduce(vcat, groupby(uncertainty_df, :name)) do group_df
        abs_error_sds = abs.(group_df[:error] ./ group_df[:posterior_std_err])
        abs_error_sds = abs_error_sds[.!ismissing.(abs_error_sds)]
        DataFrame(
            field=group_df[1, :name],
            within_half_sd=mean(abs_error_sds .<= 1/2),
            within_1_sd=mean(abs_error_sds .<= 1),
            within_2_sd=mean(abs_error_sds .<= 2),
            within_3_sd=mean(abs_error_sds .<= 3),
        )
    end
end

function plot_image(img)
    xs = img.pixels'

    xs -= 550
    black = minimum(filter(.!isnan, xs))
    xs = map((x)->isnan(x) ? black : x, xs)
    xs = min.(xs, 10_000)
    xs = log.(xs + 100)
    #cutoffs = quantile(xs2[:], 0:0.01:1)
    #xs3 = map(x->findfirst(x .< cutoffs), xs2)
    xs -= log(black + 100)
    xs /= log(10_100)
    return xs
end

end # module AccuracyBenchmark


# Copyright 2018 IBM Corp.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

module Ontology
export Annotation, AnnotationID, ObAnnotation, HomAnnotation

using AutoHashEquals
using Reexport

using ..Doctrine

# Data types
############

""" Semantic annotation of computer program.

This type is agnostic to the programming language of the computer program. All
language-specific information is stored in the `language` dictionary.
"""
abstract type Annotation end

""" Unique identifer of annotation.
"""
@auto_hash_equals struct AnnotationID
  language::String
  package::String
  id::String
end

struct ObAnnotation <: Annotation
  name::AnnotationID
  language::Dict{Symbol,Any}
  definition::Monocl.Ob
  slots::Vector{Monocl.Hom}
end

struct HomAnnotation <: Annotation
  name::AnnotationID
  language::Dict{Symbol,Any}
  definition::Monocl.Hom
end

# Modules
#########

include("OntologyJSON.jl")
include("OntologyDBs.jl")
include("rdf/OntologyRDF.jl")

@reexport using .OntologyJSON
@reexport using .OntologyDBs
@reexport using .OntologyRDF

end


module RemoteS

using GMT, Printf, Statistics, Requires, Dates

const SCENE_HALFW = Dict("AQUA" => 1163479, "TERRA" => 1163479, "LANDSAT8" => 92500)	# half widths

if isdefined(Base, :Experimental) && isdefined(Base.Experimental, Symbol("@optlevel"))
	@eval Base.Experimental.@optlevel 1
end

export
	cutcube, subcube, dn2temperature, dn2radiance, dn2reflectance, reflectance_surf, grid_at_sensor, truecolor,
	clg, clre, evi, evi2, gndvi, mndwi, mtci, mcari, msavi, mbri, ndvi, ndwi, ndwi2, ndrei1,
	ndrei2, satvi, savi, slavi,
	clip_orbits, findscenes, sat_scenes, sat_tracks, reportbands

include("grid_at_sensor.jl")
include("spectral_indices.jl")
include("utils.jl")

function __init__()
	@require SatelliteToolbox="6ac157d9-b43d-51bb-8fab-48bf53814f4a" include("sat_tracks.jl")
end

end # module


# dicom2nifti.jl
# Simple DICOM to NIfTI converter

# This script walks through the current working directory looking for DICOM
# files and sorts them into volumes. It then creates directories corresponding
# to each protocol, and saves nii files named by series number. It has
# (barely) been tested with DICOM files from a Siemens scanner. It may not
# work with other scanners. It may give incorrect information. Test thoroughly
# and use at your own risk.

using DICOM, NIfTI, ArgParse

function parse_commandline()
    s = ArgParseSettings()
    s.description = "DICOM to NIfTI converter"

    @add_arg_table s begin
        "--extension"
            help = "assume files with given extension are DICOM files. Set to blank for no extension."
            default = ".dcm"
        "--sphinx"
            help = "correct for sphinx position (for monkeys)"
            action = :store_true
        "--axes"
            help = """orientation of tensor in the resulting volume. This is a three-character label
            that determines how the data is stored, and how it is displayed in applications that do
            not make use of the orientation matrix in the NIfTI header. Valid characters are:
            L (= left), R (= right), A (= anterior), P (= posterior), I (= inferior), and
            S (= superior). If not specified, the orientation in the DICOM file is preserved."""
            default = nothing
        "dicomdir"
            help = "directory of DICOM files"
            required = true
        "niftidir"
            help = "directory in which to put NIfTI files"
            required = true
    end

    parse_args(s)
end

function walk(fn::Function, path::String)
    for fname in readdir(path)
        if fname == "." || fname == ".."
            continue
        end
        fpath = joinpath(path, fname)
        if isdir(lstat(fpath))
            walk(fn, fpath)
        else
            fn(fpath)
        end
    end
end

function findtag(d, tag, from, what)
    tag = lookup(d, tag)
    if tag == false
        error("$what missing from $from")
    end
    tag.data
end

function main()
    cmd = parse_commandline()

    if cmd["axes"] == nothing
        axes = nothing
    else
        axes = uppercase(cmd["axes"])
        if length(axes) != 3 || (!('R' in axes) && !('L' in axes)) ||
                (!('A' in axes) && !('P' in axes)) ||
                (!('S' in axes) && !('I' in axes))
            error("Invalid axes provided")
        end
    end

    transform = cmd["sphinx"] ? [1 0 0; 0 0 1; 0 -1 0] : [-1 0 0; 0 -1 0; 0 0 1]

    # Load all DICOMs into an array, indexed by series number
    dicoms = Dict{Int, Any}()
    walk(cmd["dicomdir"]) do fpath
        if !endswith(fpath, cmd["extension"])
            return
        end

        local d
        f = open(fpath, "r")
        d = DICOM.dcm_parse(f)

        series_number_index = findfirst((x) -> x.tag == (0x0020, 0x0011), d)
        if series_number_index == 0
            println("No series number found for $fpath; ignoring")
            return
        end
        series_number = d[series_number_index].data[1]

        dicom_arr = get(dicoms, series_number, nothing)
        if dicom_arr == nothing
            dicoms[series_number] = dicom_arr = {}
        end
        push!(dicom_arr, d)
    end

    # Sort DICOMs
    for (series_number, slices) in dicoms
        d = slices[1]

        protocol_name = findtag(d, (0x0018, 0x1030), series_number, "Protocol Name")[1]
        orientation = findtag(d, (0x0020, 0x0037), series_number,
            "Image Orientation (Patient)")
        pixel_spacing = findtag(d, (0x0028, 0x0030), series_number, "Pixel Spacing")
        slice_thickness = findtag(d, (0x0018, 0x0050), series_number, "Slice Thickness")[1]
        time_step = lookup(d, (0x0018, 0x0080)) # This is the TR, but may not be the time step
        phase_encoding_direction = lookup(d, (0x0018, 0x1312))

        # Convert orentation to RAS
        orientation = transform*reshape(orientation, (3, 2))

        # Determine Z as cross product of X and Y
        orientation = hcat(orientation, cross(orientation[:, 1], orientation[:, 2]))

        # Scale by pixel size
        orientation = orientation*[pixel_spacing[1] 0 0; 0 pixel_spacing[2] 0; 0 0 slice_thickness]

        # Find Z coordinates of each slice
        image_positions = transform*hcat(
            [float32(findtag(d, (0x0020, 0x0032), series_number, "Image Position (Patient)"))
            for d in slices]...)
        z_coords = [dot(image_positions[:, i], orientation[:, 3])
            for i = 1:size(image_positions, 2)]

        # Sort Z coordinates
        p = sortperm(z_coords)

        # Concatenate slices along Z axis according to increasing Z coordinate
        raw = cat(3,
            [findtag(slices[i], (0x7FE0, 0x0010), series_number, "Pixel Data")[1] for i in p]...)

        orientation = float32(hcat(orientation, image_positions[:, p[1]]))

        dim_info = phase_encoding_direction.data[1] == "ROW" ? (1, 2, 3) :
                phase_encoding_direction.data[1] == "COL" ? (2, 1, 3) :
                (0, 0, 0)

        voxel_size = tuple(pixel_spacing..., slice_thickness)

        # Permute according to axes specified on command line
        if axes != nothing
            # Determine permutation of current volume to RAS
            ras = Array(Int, 3)
            sign = Array(Bool, 3)
            rg = [1:3;]
            abs_orientation = abs(orientation)
            for i = 1:3
                idx = findmax(abs_orientation[i, rg])[2]
                ras[i] = splice!(rg, idx)
                sign[i] = orientation[i, ras[i]] >= 0
            end

            # Determine permutation of RAS to user-preferred orientation
            specified_axes = Int[(c in "RL" ? 1 :
                                  c in "AP" ? 2 : 3)
                                 for c in axes]
            specified_sign = Bool[c in "RAS" for c in axes]
            output_axes = ras[specified_axes]

            # Permute
            raw = permutedims(raw, output_axes)
            orientation[:, 1:3] = orientation[:, output_axes]
            voxel_size = voxel_size[output_axes]
            dim_info = dim_info[output_axes]

            # Flip signs and dimensions
            flip_sign = find(sign[specified_axes] .!= specified_sign)
            orientation[:, flip_sign] = -orientation[:, flip_sign]
            for dim in flip_sign
                raw = flipdim(raw, dim)
                orientation[dim, 4] += (sign[specified_axes[dim]] ? 1 : -1) *
                    (size(raw, dim) - 1)*voxel_size[dim]
            end
        end

        # Create a directory for each protocol
        protocol_dir = joinpath(cmd["niftidir"], replace(protocol_name, "/", ""))
        if !isdir(protocol_dir)
            mkdir(protocol_dir)
        end

        # Write NIfTI volumes
        ni = NIVolume(raw; voxel_size=voxel_size,
            orientation=orientation, dim_info=dim_info,
            time_step=time_step != false && !isempty(time_step.data) ? time_step.data[1] : 0f0)
        niwrite(joinpath(protocol_dir, "$(series_number).nii"), ni)
    end
end

main()


"""
  - `run(sim::Simulation, until::Event)`
  - `run(sim::Simulation, until::TimeType)`
  - `run(sim::Simulation, period::Period)`
  - `run(sim::Simulation, period::Number)`
  - `run(sim::Simulation)`

Executes [`step`](@ref) until the given criterion `until` is met:

- if it is not specified, the method will return when there are no further events to be processed
- if it is an `Event`, the method will continue stepping until this event has been triggered and will return its value
- if it is a `TimeType`, the method will continue stepping until the simulation’s time reaches until
- if it is a `Period`, the method will continue stepping until the simulation’s time has passed until periods
- if it is a `Number`, the method will continue stepping until the simulation’s time has passed until elementary periods

In the last two cases, the simulation can prematurely stop when there are no further events to be processed.
"""
function run(sim::Simulation, until::Event) :: Any
  append_callback(until, stop_simulation)
  try
    while true
      step(sim)
    end
  catch exc
    if isa(exc, StopSimulation)
      return exc.value
    else
      rethrow(exc)
    end
  end
end

function run(sim::Simulation, period::Period) :: Any
  run(sim, timeout(sim, period))
end

function run(sim::Simulation, period::Number) :: Any
  run(sim, timeout(sim, sim.granularity(period)))
end

function run{T}(sim::Simulation{T}, until::T) :: Any
  run(sim, timeout(sim, now(sim)-until))
end

function run(sim::Simulation) :: Any
  run(sim, timeout(sim, typemax(sim.granularity)))
end

function stop_simulation(sim::Simulation, ev::Event)
  throw(StopSimulation(ev.value))
end

"""
  `now(sim::Simulation) :: TimeType`

Returns the current simulation time.
"""
function now(sim::Simulation) :: TimeType
  return sim.time
end

"""
  - `schedule!(sim::Simulation, ev::Event, delay::Period; priority::Bool=false, value::Any=nothing) :: Event`
  - `schedule!(sim::Simulation, ev::Event, delay::Number=0; priority::Bool=false, value::Any=nothing) :: Event`

Schedules an event at time `sim.time + delay` with a `priority` and a `value`.

If the event is already scheduled, the key is updated with the new `delay` and `priority`. The new `value` is also set.

If the event is being processed, an [`EventProcessing`](@ref) exception is thrown.
"""
function schedule!(sim::Simulation, ev::Event, delay::Period; priority::Bool=false, value::Any=nothing) :: Event
  if ev.state == processed
    throw(EventProcessed())
  end
  ev.value = value
  if ev.state == triggered
    id = sim.heap[ev].id
  else
    id = sim.sid+=one(UInt)
  end
  sim.heap[ev] = EventKey(sim.time + delay, priority, id)
  ev.state = triggered
  return ev
end

function schedule!(sim::Simulation, ev::Event, delay::Number=0; priority::Bool=false, value::Any=nothing) :: Event
  schedule!(sim, ev, sim.granularity(delay), priority=priority, value=value)
end

"""
  - `schedule(sim::Simulation, ev::Event, delay::Period; priority::Bool=false, value::Any=nothing) :: Event`
  - `schedule(sim::Simulation, ev::Event, delay::Number=0; priority::Bool=false, value::Any=nothing) :: Event`

Schedules an event at time `sim.time + delay` with a `priority` and a `value`.

If the event is already scheduled or is being processed, an [`EventNotIdle`](@ref) exception is thrown.
"""
function schedule(sim::Simulation, ev::Event, delay::Period; priority::Bool=false, value::Any=nothing) :: Event
  if !(ev.state == idle)
    throw(EventNotIdle())
  end
  ev.value = value
  sim.heap[ev] = EventKey(sim.time + delay, priority, sim.sid+=one(UInt))
  ev.state = triggered
  return ev
end

function schedule(sim::Simulation, ev::Event, delay::Number=0; priority::Bool=false, value::Any=nothing) :: Event
  schedule(sim, ev, sim.granularity(delay), priority=priority, value=value)
end

"""
  `step(sim::Simulation) :: Bool`

Does a simulation step and processes the next event.

Only used internally.
"""
function step(sim::Simulation)
  if isempty(sim.heap)
    throw(EmptySchedule())
  end
  (ev, key) = peek(sim.heap)
  dequeue!(sim.heap)
  sim.time = key.time
  ev.state = processed
  while !isempty(ev.callbacks)
    cb = dequeue!(ev.callbacks)
    cb(sim, ev)
  end
end


module SLEEF

export xsin,  xcos,  xtan,  xasin,  xacos,  xatan,  xlog,
	xexp,  xsinh,  xcosh,  xtanh,  xasinh,  xacosh,  xatanh,
	xcbrt,  xexp2,  xexp10,  xexpm1,  xlog10,  xlog1p,
	xsin_u1, xcos_u1,  xtan_u1,  xasin_u1,
	xacos_u1,  xatan_u1,  xlog_u1,  xcbrt_u1,
	xsincos, xsincos_u1, xatan2, xatan2_u1, xpow,
	xldexp, xilogb

if isfile(joinpath(dirname(@__FILE__),"..","deps","deps.jl"))
    include("../deps/deps.jl")
else
    error("SLEEF not properly installed. Please run Pkg.build(\"SLEEF\")")
end

include("sleef_inner.jl")

##############################

end # SLEEF Module


function DiffEqBase.__solve(prob::Union{DiffEqBase.AbstractODEProblem,DiffEqBase.AbstractDAEProblem},
                            alg::Union{OrdinaryDiffEqAlgorithm,DAEAlgorithm}, args...;
                            kwargs...)
  integrator = DiffEqBase.__init(prob, alg, args...; kwargs...)
  solve!(integrator)
  integrator.sol
end

function DiffEqBase.__init(prob::Union{DiffEqBase.AbstractODEProblem,DiffEqBase.AbstractDAEProblem},
                           alg::Union{OrdinaryDiffEqAlgorithm,DAEAlgorithm},
                           timeseries_init = (),
                           ts_init = (),
                           ks_init = (),
                           recompile::Type{Val{recompile_flag}} = Val{true};
                           saveat = (),
                           tstops = (),
                           d_discontinuities = (),
                           save_idxs = nothing,
                           save_everystep = isempty(saveat),
                           save_on = true,
                           save_start = save_everystep || isempty(saveat) || saveat isa Number || prob.tspan[1] in saveat,
                           save_end = nothing,
                           callback = nothing,
                           dense = save_everystep && !(typeof(alg) <: Union{DAEAlgorithm,FunctionMap}) && isempty(saveat),
                           calck = (callback !== nothing && callback != CallbackSet()) || (dense), # and no dense output
                           dt = alg isa FunctionMap && isempty(tstops) ? eltype(prob.tspan)(1) : eltype(prob.tspan)(0),
                           dtmin = nothing,
                           dtmax = eltype(prob.tspan)((prob.tspan[end]-prob.tspan[1])),
                           force_dtmin = false,
                           adaptive = isadaptive(alg),
                           gamma = gamma_default(alg),
                           abstol = nothing,
                           reltol = nothing,
                           qmin = qmin_default(alg),
                           qmax = qmax_default(alg),
                           qsteady_min = qsteady_min_default(alg),
                           qsteady_max = qsteady_max_default(alg),
                           qoldinit = isadaptive(alg) ? 1//10^4 : 0,
                           fullnormalize = true,
                           failfactor = 2,
                           beta1 = nothing,
                           beta2 = nothing,
                           maxiters = adaptive ? 1000000 : typemax(Int),
                           internalnorm = ODE_DEFAULT_NORM,
                           internalopnorm = LinearAlgebra.opnorm,
                           isoutofdomain = ODE_DEFAULT_ISOUTOFDOMAIN,
                           unstable_check = ODE_DEFAULT_UNSTABLE_CHECK,
                           verbose = true,
                           timeseries_errors = true,
                           dense_errors = false,
                           advance_to_tstop = false,
                           stop_at_next_tstop = false,
                           initialize_save = true,
                           progress = false,
                           progress_steps = 1000,
                           progress_name = "ODE",
                           progress_message = ODE_DEFAULT_PROG_MESSAGE,
                           userdata = nothing,
                           allow_extrapolation = alg_extrapolates(alg),
                           initialize_integrator = true,
                           alias_u0 = false,
                           alias_du0 = false,
                           initializealg = DefaultInit(),
                           kwargs...) where recompile_flag

  if prob isa DiffEqBase.AbstractDAEProblem && alg isa OrdinaryDiffEqAlgorithm
    error("You cannot use an ODE Algorithm with a DAEProblem")
  end

  if prob isa DiffEqBase.AbstractODEProblem && alg isa DAEAlgorithm
    error("You cannot use an DAE Algorithm with a ODEProblem")
  end

  if typeof(prob.f)<:DynamicalODEFunction && typeof(prob.f.mass_matrix)<:Tuple
    if any(mm != I for mm in prob.f.mass_matrix)
      error("This solver is not able to use mass matrices.")
    end
  elseif !(typeof(prob)<:DiscreteProblem) &&
         !(typeof(prob)<:DiffEqBase.AbstractDAEProblem) &&
         !is_mass_matrix_alg(alg) &&
         prob.f.mass_matrix != I
    error("This solver is not able to use mass matrices.")
  end

  if !isempty(saveat) && dense
    @warn("Dense output is incompatible with saveat. Please use the SavingCallback from the Callback Library to mix the two behaviors.")
  end

  progress && @logmsg(LogLevel(-1),progress_name,_id=_id = :OrdinaryDiffEq,progress=0)

  tType = eltype(prob.tspan)
  tspan = prob.tspan
  tdir = sign(tspan[end]-tspan[1])

  t = tspan[1]

  if (((!(typeof(alg) <: OrdinaryDiffEqAdaptiveAlgorithm) && !(typeof(alg) <: OrdinaryDiffEqCompositeAlgorithm) && !(typeof(alg) <: DAEAlgorithm)) || !adaptive) && dt == tType(0) && isempty(tstops)) && !(typeof(alg) <: Union{FunctionMap,LinearExponential})
      error("Fixed timestep methods require a choice of dt or choosing the tstops")
  end

  isdae = alg isa DAEAlgorithm || (!(typeof(prob)<:DiscreteProblem) &&
                                     prob.f.mass_matrix != I &&
                                     !(typeof(prob.f.mass_matrix)<:Tuple) &&
                                     ArrayInterface.issingular(prob.f.mass_matrix))
  if alg isa CompositeAlgorithm && alg.choice_function isa AutoSwitch
    auto = alg.choice_function
    _alg = CompositeAlgorithm(alg.algs,
                             AutoSwitchCache(
                                             0,0,
                                             auto.nonstiffalg,
                                             auto.stiffalg,
                                             auto.stiffalgfirst,
                                             auto.maxstiffstep,
                                             auto.maxnonstiffstep,
                                             auto.nonstifftol,
                                             auto.stifftol,
                                             auto.dtfac,
                                             auto.stiffalgfirst,
                                             auto.switch_max
                                            ))
  else
    _alg = alg
  end
  f = prob.f
  p = prob.p

  # Get the control variables

  if alias_u0
    u = prob.u0
  else
    u = recursivecopy(prob.u0)
  end

  if _alg isa DAEAlgorithm
    if alias_du0
      du = prob.du0
    else
      du = recursivecopy(prob.u0)
    end
    duprev = recursivecopy(du)
  else
    du = nothing
    duprev = nothing
  end

  uType = typeof(u)
  uBottomEltype = recursive_bottom_eltype(u)
  uBottomEltypeNoUnits = recursive_unitless_bottom_eltype(u)

  uEltypeNoUnits = recursive_unitless_eltype(u)
  tTypeNoUnits   = typeof(one(tType))

  if typeof(_alg) <: FunctionMap
    abstol_internal = false
  elseif abstol === nothing
    if uBottomEltypeNoUnits == uBottomEltype
      abstol_internal = real(convert(uBottomEltype,oneunit(uBottomEltype)*1//10^6))
    else
      abstol_internal = real.(oneunit.(u).*1//10^6)
    end
  else
    abstol_internal = real.(abstol)
  end

  if typeof(_alg) <: FunctionMap
    reltol_internal = false
  elseif reltol === nothing
    if uBottomEltypeNoUnits == uBottomEltype
      reltol_internal = real(convert(uBottomEltype,oneunit(uBottomEltype)*1//10^3))
    else
      reltol_internal = real.(oneunit.(u).*1//10^3)
    end
  else
    reltol_internal = real.(reltol)
  end

  dtmax > zero(dtmax) && tdir < 0 && (dtmax *= tdir) # Allow positive dtmax, but auto-convert
  # dtmin is all abs => does not care about sign already.

  if !isdae && isinplace(prob) && typeof(u) <: AbstractArray && eltype(u) <: Number && uBottomEltypeNoUnits == uBottomEltype && tType == tTypeNoUnits # Could this be more efficient for other arrays?
    rate_prototype = recursivecopy(u)
  elseif prob isa DAEProblem
    rate_prototype = prob.du0
  else
    if (uBottomEltypeNoUnits == uBottomEltype && tType == tTypeNoUnits) || eltype(u) <: Enum
      rate_prototype = u
    else # has units!
      rate_prototype = u/oneunit(tType)
    end
  end
  rateType = typeof(rate_prototype) ## Can be different if united

  if isdae
    if uBottomEltype == uBottomEltypeNoUnits
      res_prototype = u
    else
      res_prototype = one(u)
    end
    resType = typeof(res_prototype)
  end

  tstops_internal = initialize_tstops(tType, tstops, d_discontinuities, tspan)
  saveat_internal = initialize_saveat(tType, saveat, tspan)
  d_discontinuities_internal = initialize_d_discontinuities(tType, d_discontinuities, tspan)

  callbacks_internal = CallbackSet(callback)

  max_len_cb = DiffEqBase.max_vector_callback_length(callbacks_internal)
  if max_len_cb isa VectorContinuousCallback
    uBottomEltypeReal = real(uBottomEltype)
    if isinplace(prob)
      callback_cache = DiffEqBase.CallbackCache(u,max_len_cb.len,uBottomEltypeReal,uBottomEltypeReal)
    else
      callback_cache = DiffEqBase.CallbackCache(max_len_cb.len,uBottomEltypeReal,uBottomEltypeReal)
    end
  else
    callback_cache = nothing
  end

  ### Algorithm-specific defaults ###
  if save_idxs === nothing
    ksEltype = Vector{rateType}
  else
    ks_prototype = rate_prototype[save_idxs]
    ksEltype = Vector{typeof(ks_prototype)}
  end

  # Have to convert incase passed in wrong.
  if save_idxs === nothing
    timeseries = timeseries_init === () ? uType[] : convert(Vector{uType},timeseries_init)
  else
    u_initial = u[save_idxs]
    timeseries = timeseries_init === () ? typeof(u_initial)[] :
                                          convert(Vector{uType},timeseries_init)
  end

  ts = ts_init === () ? tType[]    : convert(Vector{tType},ts_init)
  ks = ks_init === () ? ksEltype[] : convert(Vector{ksEltype},ks_init)
  alg_choice = typeof(_alg) <: CompositeAlgorithm ? Int[] : ()

  if !adaptive && save_everystep && tspan[2]-tspan[1] != Inf
    if dt == 0
      steps = length(tstops)
    else
      dtmin === nothing && (dtmin = DiffEqBase.prob2dtmin(prob; use_end_time=true))
      abs(dt) < dtmin && throw(ArgumentError("Supplied dt is smaller than dtmin"))
      steps = ceil(Int,internalnorm((tspan[2]-tspan[1])/dt,tspan[1]))
    end
    sizehint!(timeseries,steps+1)
    sizehint!(ts,steps+1)
    sizehint!(ks,steps+1)
  elseif save_everystep
    sizehint!(timeseries,50)
    sizehint!(ts,50)
    sizehint!(ks,50)
  elseif !isempty(saveat_internal)
    sizehint!(timeseries,length(saveat_internal)+1)
    sizehint!(ts,length(saveat_internal)+1)
    sizehint!(ks,length(saveat_internal)+1)
  else
    sizehint!(timeseries,2)
    sizehint!(ts,2)
    sizehint!(ks,2)
  end

  QT = if tTypeNoUnits <: Integer
    typeof(qmin)
  elseif prob isa DiscreteProblem
    # The QT fields are not used for DiscreteProblems
    constvalue(tTypeNoUnits)
  else
    typeof(internalnorm(u, t))
  end

  k = rateType[]

  if uses_uprev(_alg, adaptive) || calck
    uprev = recursivecopy(u)
  else
    # Some algorithms do not use `uprev` explicitly. In that case, we can save
    # some memory by aliasing `uprev = u`, e.g. for "2N" low storage methods.
    uprev = u
  end
  if allow_extrapolation
    uprev2 = recursivecopy(u)
  else
    uprev2 = uprev
  end

  if prob isa DAEProblem
    cache = alg_cache(_alg,du,u,res_prototype,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol_internal,p,calck,Val(isinplace(prob)))
  else
    cache = alg_cache(_alg,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol_internal,p,calck,Val(isinplace(prob)))
  end

  if typeof(_alg) <: OrdinaryDiffEqCompositeAlgorithm
    id = CompositeInterpolationData(f,timeseries,ts,ks,alg_choice,dense,cache)
    beta2 === nothing && ( beta2=_composite_beta2_default(_alg.algs, cache.current, QT) )
    beta1 === nothing && ( beta1=_composite_beta1_default(_alg.algs, cache.current, QT, beta2) )
  else
    id = InterpolationData(f,timeseries,ts,ks,dense,cache)
    beta2 === nothing && ( beta2=beta2_default(_alg) )
    beta1 === nothing && ( beta1=beta1_default(_alg,beta2) )
  end

  dtmin === nothing && (dtmin = DiffEqBase.prob2dtmin(prob; use_end_time=false))

  save_end_user = save_end
  save_end = save_end === nothing ? save_everystep || isempty(saveat) || saveat isa Number || prob.tspan[2] in saveat : save_end

  opts = DEOptions{typeof(abstol_internal),typeof(reltol_internal),QT,tType,
                   typeof(internalnorm),typeof(internalopnorm),typeof(save_end_user),
                   typeof(callbacks_internal),
                   typeof(isoutofdomain),
                   typeof(progress_message),typeof(unstable_check),typeof(tstops_internal),
                   typeof(d_discontinuities_internal),typeof(userdata),typeof(save_idxs),
                   typeof(maxiters),typeof(tstops),typeof(saveat),
                   typeof(d_discontinuities)}(
                       maxiters,save_everystep,adaptive,abstol_internal,
                       reltol_internal,QT(gamma),QT(qmax),
                       QT(qmin),QT(qsteady_max),
                       QT(qsteady_min),QT(failfactor),tType(dtmax),
                       tType(dtmin),internalnorm,internalopnorm,save_idxs,tstops_internal,saveat_internal,
                       d_discontinuities_internal,
                       tstops,saveat,d_discontinuities,
                       userdata,progress,progress_steps,
                       progress_name,progress_message,timeseries_errors,dense_errors,
                       QT(beta1),QT(beta2),QT(qoldinit),dense,
                       save_on,save_start,save_end,save_end_user,
                       callbacks_internal,isoutofdomain,
                       unstable_check,verbose,
                       calck,force_dtmin,advance_to_tstop,stop_at_next_tstop)

  destats = DiffEqBase.DEStats(0)

  if typeof(_alg) <: OrdinaryDiffEqCompositeAlgorithm
    sol = DiffEqBase.build_solution(prob,_alg,ts,timeseries,
                      dense=dense,k=ks,interp=id,
                      alg_choice=alg_choice,
                      calculate_error = false, destats=destats)
  else
    sol = DiffEqBase.build_solution(prob,_alg,ts,timeseries,
                      dense=dense,k=ks,interp=id,
                      calculate_error = false, destats=destats)
  end

  if recompile_flag == true
    FType = typeof(f)
    SolType = typeof(sol)
    cacheType = typeof(cache)
  else
    FType = Function
    if _alg isa OrdinaryDiffEqAlgorithm
      SolType = DiffEqBase.AbstractODESolution
      cacheType =  OrdinaryDiffEqCache
    else
      SolType = DiffEqBase.AbstractDAESolution
      cacheType =  DAECache
    end
  end

  # rate/state = (state/time)/state = 1/t units, internalnorm drops units
  # we don't want to differentiate through eigenvalue estimation
  eigen_est = inv(one(tType))
  tprev = t
  dtcache = tType(dt)
  dtpropose = tType(dt)
  iter = 0
  kshortsize = 0
  reeval_fsal = false
  u_modified = false
  EEst = QT(1)
  just_hit_tstop = false
  isout = false
  accept_step = false
  force_stepfail = false
  last_stepfail = false
  do_error_check = true
  event_last_time = 0
  vector_event_last_time = 1
  last_event_error = typeof(_alg) <: FunctionMap ? false : zero(uBottomEltypeNoUnits)
  dtchangeable = isdtchangeable(_alg)
  q11 = QT(1)
  success_iter = 0
  erracc = QT(1)
  dtacc = tType(1)
  reinitiailize = true
  saveiter = 0 # Starts at 0 so first save is at 1
  saveiter_dense = 0

  integrator = ODEIntegrator{typeof(_alg),isinplace(prob),uType,typeof(du),
                             tType,typeof(p),
                             typeof(eigen_est),
                             QT,typeof(tdir),typeof(k),SolType,
                             FType,cacheType,
                             typeof(opts),fsal_typeof(_alg,rate_prototype),
                             typeof(last_event_error),typeof(callback_cache),
                             typeof(initializealg)}(
                             sol,u,du,k,t,tType(dt),f,p,uprev,uprev2,duprev,tprev,
                             _alg,dtcache,dtchangeable,
                             dtpropose,tdir,eigen_est,EEst,QT(qoldinit),q11,
                             erracc,dtacc,success_iter,
                             iter,saveiter,saveiter_dense,cache,callback_cache,
                             kshortsize,force_stepfail,last_stepfail,
                             just_hit_tstop,do_error_check,
                             event_last_time,vector_event_last_time,
                             last_event_error,accept_step,
                             isout,reeval_fsal,
                             u_modified,reinitiailize,isdae,
                             opts,destats,initializealg)
  if initialize_integrator
    if isdae
      DiffEqBase.initialize_dae!(integrator)
    end

    if save_start
      integrator.saveiter += 1 # Starts at 1 so first save is at 2
      integrator.saveiter_dense += 1
      copyat_or_push!(ts,1,t)
      if save_idxs === nothing
        copyat_or_push!(timeseries,1,integrator.u)
        copyat_or_push!(ks,1,[rate_prototype])
      else
        copyat_or_push!(timeseries,1,u_initial,Val{false})
        copyat_or_push!(ks,1,[ks_prototype])
      end
    else
      saveiter = 0 # Starts at 0 so first save is at 1
      saveiter_dense = 0
    end

    initialize_callbacks!(integrator, initialize_save)
    initialize!(integrator,integrator.cache)

    if typeof(_alg) <: OrdinaryDiffEqCompositeAlgorithm && save_start
      # Loop to get all of the extra possible saves in callback initialization
      for i in 1:integrator.saveiter
        copyat_or_push!(alg_choice,i,integrator.cache.current)
      end
    end
  end

  handle_dt!(integrator)

  integrator
end

function DiffEqBase.solve!(integrator::ODEIntegrator)
  @inbounds while !isempty(integrator.opts.tstops)
    while integrator.tdir * integrator.t < first(integrator.opts.tstops)
      loopheader!(integrator)
      if integrator.do_error_check && check_error!(integrator) != :Success
        return integrator.sol
      end
      perform_step!(integrator,integrator.cache)
      loopfooter!(integrator)
      if isempty(integrator.opts.tstops)
        break
      end
    end
    handle_tstop!(integrator)
  end
  postamble!(integrator)

  f = integrator.sol.prob.f

  if DiffEqBase.has_analytic(f)
    DiffEqBase.calculate_solution_errors!(integrator.sol;timeseries_errors=integrator.opts.timeseries_errors,dense_errors=integrator.opts.dense_errors)
  end
  if integrator.sol.retcode != :Default
    return integrator.sol
  end
  integrator.sol = DiffEqBase.solution_new_retcode(integrator.sol,:Success)
end

# Helpers

function handle_dt!(integrator)
  if iszero(integrator.dt) && integrator.opts.adaptive
    auto_dt_reset!(integrator)
    if sign(integrator.dt)!=integrator.tdir && !iszero(integrator.dt) && !isnan(integrator.dt)
      error("Automatic dt setting has the wrong sign. Exiting. Please report this error.")
    end
    if isnan(integrator.dt)
      if integrator.opts.verbose
        @warn("Automatic dt set the starting dt as NaN, causing instability.")
      end
    end
  elseif integrator.opts.adaptive && integrator.dt > zero(integrator.dt) && integrator.tdir < 0
    integrator.dt *= integrator.tdir # Allow positive dt, but auto-convert
  end
end

# time stops
function initialize_tstops(::Type{T}, tstops, d_discontinuities, tspan) where T
  tstops_internal = BinaryMinHeap{T}()

  t0, tf = tspan
  tdir = sign(tf - t0)
  tdir_t0 = tdir * t0
  tdir_tf = tdir * tf

  if isempty(d_discontinuities) && isempty(tstops) # TODO: Specialize more
    push!(tstops_internal, tdir_tf)
  else
    for t in tstops
      tdir_t = tdir * t
      tdir_t0 < tdir_t ≤ tdir_tf && push!(tstops_internal, tdir_t)
    end

    for t in d_discontinuities
      tdir_t = tdir * t
      tdir_t0 < tdir_t ≤ tdir_tf && push!(tstops_internal, tdir_t)
    end

    push!(tstops_internal, tdir_tf)
  end

  return tstops_internal
end

# saving time points
function initialize_saveat(::Type{T}, saveat, tspan) where T
  saveat_internal = BinaryMinHeap{T}()

  t0, tf = tspan
  tdir = sign(tf - t0)
  tdir_t0 = tdir * t0
  tdir_tf = tdir * tf

  if typeof(saveat) <: Number
    directional_saveat = tdir * abs(saveat)
    for t in (t0 + directional_saveat):directional_saveat:tf
      push!(saveat_internal, tdir * t)
    end
  elseif !isempty(saveat)
    for t in saveat
      tdir_t = tdir * t
      tdir_t0 < tdir_t ≤ tdir_tf && push!(saveat_internal, tdir_t)
    end
  end

  return saveat_internal
end

# discontinuities
function initialize_d_discontinuities(::Type{T}, d_discontinuities, tspan) where T
  d_discontinuities_internal = BinaryMinHeap{T}()
  sizehint!(d_discontinuities_internal, length(d_discontinuities))

  t0, tf = tspan
  tdir = sign(tf - t0)

  for t in d_discontinuities
    push!(d_discontinuities_internal, tdir * t)
  end

  return d_discontinuities_internal
end

function initialize_callbacks!(integrator, initialize_save = true)
  t = integrator.t
  u = integrator.u
  callbacks = integrator.opts.callback
  integrator.u_modified = true

  u_modified = initialize!(callbacks,u,t,integrator)

  # if the user modifies u, we need to fix previous values before initializing
  # FSAL in order for the starting derivatives to be correct
  if u_modified

    if isinplace(integrator.sol.prob)
      recursivecopy!(integrator.uprev,integrator.u)
    else
      integrator.uprev = integrator.u
    end

    if alg_extrapolates(integrator.alg)
      if isinplace(integrator.sol.prob)
        recursivecopy!(integrator.uprev2,integrator.uprev)
      else
        integrator.uprev2 = integrator.uprev
      end
    end

    if initialize_save &&
      (any((c)->c.save_positions[2],callbacks.discrete_callbacks) ||
      any((c)->c.save_positions[2],callbacks.continuous_callbacks))
      savevalues!(integrator,true)
    end
  end

  # reset this as it is now handled so the integrators should proceed as normal
  integrator.u_modified = false
end


# binary.jl: Elementwise broadcasting binary functions for arrays and scalars.
# uses binary_ops from broadcast.jl.

import Base.Broadcast: broadcasted

# binary_op defines the broadcast_func of a Julia function for KnetArrays.
# The corresponding kernel is defined in libknet8.
function binary_op(f, j=f, o...)
    J=Symbol(j)
    M = which(@__MODULE__, J)
    for S in (32,64)
        T = Symbol("Float$S")

        F01 = "$(f)_$(S)_01" # Scalar,Array->Array
        F11 = "$(f)_$(S)_11" # Array,Array->Array (same size) (not broadcast)
        F12 = "$(f)_$(S)_12" # Array,Array->Array (one have to be vector)
        F13_x_y = "$(f)_$(S)_13_x_y" # e.g. (A(x,y,z,w,t...), B(1,1,1,w,1...))
        F13_y_x = "$(f)_$(S)_13_y_x" # different versions for efficiency
        # F14_x_y = "$(f)_$(S)_14_x_y" # e.g. (M(x,y,z,w,t...), N(w,1,1,1...)
        # F14_y_x = "$(f)_$(S)_14_y_x" # different versions for efficiency
        # F15 reserved for another kernel, eliminated later and combined with F16
        F16_3 = "$(f)_$(S)_16_3" # multi-dimensional bcast unrolled up to 5 dims
        F16_4 = "$(f)_$(S)_16_4" # multi-dimensional bcast unrolled up to 5 dims
        F16_5 = "$(f)_$(S)_16_5" # multi-dimensional bcast unrolled up to 5 dims
        F17 = "$(f)_$(S)_17" # multi-dimensional bcast with loops

        @eval begin
            # Scalar,Array->Array
            function broadcasted(::typeof($J),x::$T,y::KnetArray{$T})
                z = similar(y)
                @knet8($F01,(Cint,$T,Ptr{$T},Ptr{$T}),length(z),x,y,z)
                return z
            end
            # Array,Array->Array
            function broadcasted(::typeof($J),x::KnetArray{$T},y::KnetArray{$T})
                if size(x)==size(y)
                    z = similar(x)
                    @knet8($F11,(Cint,Ptr{$T},Ptr{$T},Ptr{$T}),length(z),x,y,z)
                    return z
                else
                    bs = vbroadcast_shape(x,y)
                    z = similar(x,bs[1])
                    _broadcasted($J,x,y,z,bs)
                end
            end
            # Helpers
            function _broadcasted(::typeof($J),x::KnetArray{$T},y::KnetArray{$T},z::KnetArray{$T,1},bs)
                if length(x) == 1
                    broadcasted($J,x[1],y)
                elseif length(y) == 1
                    broadcasted($J,x,y[1])
                else # length(x) == length(y) was handled above
                    throw(DimensionMismatch("$(map(size,(x,y,z)))"))
                end
            end
            function _broadcasted(::typeof($J),x::KnetArray{$T},y::KnetArray{$T},z::KnetArray{$T,2},bs)
                # xlast or ylast will be broadcasting dimension
                (dz,sx,nx,sy,ny,xlast,ylast,xdims,ydims,multi) = bs
                if (nx == 1
                    || ny == 1
                    || ((xdims == 1 && (xlast==1 || 512 < sx )) ||
                        (ydims == 1 && (ylast==1 || sy < 512 )))
                    || (xdims==1 && nx<704)
                    || (ydims==1 && (ny<704)))
                    @knet8($F12,
                           (Cint,Ptr{$T},Cint,Cint,Ptr{$T},Cint,Cint,Ptr{$T}),
                           length(z),x,sx,nx,y,sy,ny,z)
                elseif xdims == 1
                    dim_stride = strides(y)[xlast]
                    next_stride = (xlast+1) > ndims(y) ?
                        0 : strides(y)[xlast+1]
                    dim_size = prod(size(y)[xlast+1:end])
                    @knet8($F13_y_x,
                           (Ptr{$T},Ptr{$T},Ptr{$T},Cint,Cint,Cint,Cint,Cint),
                           y,x,z,dim_stride,next_stride,dim_size,
                           length(y),length(x))
                elseif ydims == 1
                    dim_stride = strides(x)[ylast]
                    next_stride = (ylast+1) > ndims(x) ?
                        0 : strides(x)[ylast+1]
                    dim_size = prod(size(x)[ylast+1:end])
                    @knet8($F13_x_y,
                           (Ptr{$T},Ptr{$T},Ptr{$T},Cint,Cint,Cint,Cint,Cint),
                           x,y,z,dim_stride,next_stride,dim_size,
                           length(x), length(y))
                else
                    throw(DimensionMismatch("$(map(size,(x,y,z)))"))
                end
                return z
            end
            function _broadcasted(::typeof($J),x::KnetArray{$T},y::KnetArray{$T},z::KnetArray{$T,3},bs)
                sx,sy,sz = get_strides(x,y,z)
                @knet8($F16_3,
                       (Ptr{$T},Ptr{$T},Ptr{$T},Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint),
                       x,y,z, sx[1],sx[2],sx[3], sy[1],sy[2],sy[3], sz[1],sz[2],sz[3], length(z))
                return z
            end
            function _broadcasted(::typeof($J),x::KnetArray{$T},y::KnetArray{$T},z::KnetArray{$T,4},bs)
                sx,sy,sz = get_strides(x,y,z)
                @knet8($F16_4,
                       (Ptr{$T},Ptr{$T},Ptr{$T},Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint),
                       x,y,z, sx[1],sx[2],sx[3],sx[4], sy[1],sy[2],sy[3],sy[4], sz[1],sz[2],sz[3],sz[4], length(z))
                return z
            end
            function _broadcasted(::typeof($J),x::KnetArray{$T},y::KnetArray{$T},z::KnetArray{$T,5},bs)
                sx,sy,sz = get_strides(x,y,z)
                @knet8($F16_5,
                       (Ptr{$T},Ptr{$T},Ptr{$T},Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint),
                       x,y,z, sx[1],sx[2],sx[3],sx[4],sx[5], sy[1],sy[2],sy[3],sy[4],sy[5], sz[1],sz[2],sz[3],sz[4],sz[5], length(z))
                return z
            end
            function _broadcasted(::typeof($J),x::KnetArray{$T},y::KnetArray{$T},z::KnetArray{$T},bs)
                # ndims(z) <= 5 handled above, this is for > 5
                sx,sy,sz = map(s->convert(KnetArray, s), get_strides(x,y,z))
                @knet8($F17,
                       (Ptr{$T},Ptr{$T},Ptr{$T},Ptr{Cint},Ptr{Cint},Ptr{Cint},Cint,Cint),
                       x,y,z, sx, sy, sz, length(z), ndims(z))
                return z
            end

            # Bcasted methods
            ($M).$J(x::Bcasted{<:KnetArray{$T}}, y::Bcasted{<:KnetArray{$T}}) = broadcasted($J, x.value, y.value) |> Bcasted
            ($M).$J(x::KnetArray{$T}, y::Bcasted{<:KnetArray{$T}}) = broadcasted($J, x, y.value) |> Bcasted
            ($M).$J(x::Bcasted{<:KnetArray{$T}}, y::KnetArray{$T}) = broadcasted($J, x.value, y) |> Bcasted
            ($M).$J(x::Bcasted{$T}, y::Bcasted{<:KnetArray{$T}}) = broadcasted($J, x.value, y.value) |> Bcasted
            ($M).$J(x::$T, y::Bcasted{<:KnetArray{$T}}) = broadcasted($J, x, y.value) |> Bcasted
            ($M).$J(x::Bcasted{$T}, y::KnetArray{$T}) = broadcasted($J, x.value, y) |> Bcasted
            broadcasted(::typeof($J),x::Bcasted{<:KnetArray{$T}}, y::Bcasted{<:KnetArray{$T}}) = broadcasted($J, x.value, y.value) |> Bcasted
            broadcasted(::typeof($J),x::KnetArray{$T}, y::Bcasted{<:KnetArray{$T}}) = broadcasted($J, x, y.value) |> Bcasted
            broadcasted(::typeof($J),x::Bcasted{<:KnetArray{$T}}, y::KnetArray{$T}) = broadcasted($J, x.value, y) |> Bcasted
            broadcasted(::typeof($J),x::Bcasted{$T}, y::Bcasted{<:KnetArray{$T}}) = broadcasted($J, x.value, y.value) |> Bcasted
            broadcasted(::typeof($J),x::$T, y::Bcasted{<:KnetArray{$T}}) = broadcasted($J, x, y.value) |> Bcasted
            broadcasted(::typeof($J),x::Bcasted{$T}, y::KnetArray{$T}) = broadcasted($J, x.value, y) |> Bcasted
        end # @eval
    end # for
    @eval begin # so we do not trigger some default Base implementation
        ($M).$J(x::Bcasted, y::Bcasted) = throw(MethodError($J,(x,y)))
        ($M).$J(x, y::Bcasted) = throw(MethodError($J,(x,y)))
        ($M).$J(x::Bcasted, y) = throw(MethodError($J,(x,y)))
        broadcasted(::typeof($J),x::Bcasted, y::Bcasted) = throw(MethodError(broadcasted, ($J, x, y)))
        broadcasted(::typeof($J),x, y::Bcasted) = throw(MethodError(broadcasted, ($J, x, y)))
        broadcasted(::typeof($J),x::Bcasted, y) = throw(MethodError(broadcasted, ($J, x, y)))
    end
end # function binary_op

# vbroadcast_shape computes index/offset arguments for a broadcasting kernel call.
function vbroadcast_shape(x,y)
    nz = max(ndims(x),ndims(y))
    dz = ones(Int,nz)
    xdims = ydims = xsame = ysame = xlast = ylast = 0; zlen = 1;
    sx = sy = -1;
    nx = ny = -1;

    for i = 1:nz
        # xdims: number of xdims != 1
        # xlast: last index != 1

        if size(x,i) > 1
            xdims += 1; xlast = i
            dz[i] = size(x,i)
        end
        if size(y,i) > 1
            ydims += 1; ylast = i
            if dz[i] == 1
                dz[i] = size(y,i)
            elseif dz[i] != size(y,i)
                throw(DimensionMismatch(
                    "arrays could not be broadcast to a common size"))
            end
        end

        xsame += (dz[i] == size(x,i))
        ysame += (dz[i] == size(y,i))
        zlen *= dz[i]
    end

    multi = false
    if (xsame != nz && xdims > 1) || (ysame != nz && ydims > 1)
        multi = true
    end

    if xdims == 0
        sx = zlen; nx = 1
    elseif xdims == 1
        sx = prod(dz[1:xlast-1]); nx=dz[xlast]
    elseif xsame == nz
        sx = 1
        nx = zlen
    elseif !multi
        error("Broadcasting error")
    end

    if ydims == 0
        sy = zlen
        ny = 1
    elseif ydims == 1
        sy = prod(dz[1:ylast-1])
        ny = dz[ylast]
    elseif ysame == nz
        sy = 1
        ny = zlen
    elseif !multi
        error("Broadcasting error")
    end

    return (tuple(dz...), sx, nx, sy, ny,xlast,ylast,xdims,ydims,multi)
end

function get_strides(x,y,z)
    # x,y,z may have different ndims, work with the max
    n = max(ndims(x), ndims(y), ndims(z))
    stride_x = Int32[ stride(x,i) for i=1:n ]
    stride_y = Int32[ stride(y,i) for i=1:n ]
    stride_z = Int32[ stride(z,i) for i=1:n ]
    dims_x = Int32[ size(x,i) for i=1:n ]
    dims_y = Int32[ size(y,i) for i=1:n ]
    for i in 1:n
        dims_x[i] == dims_y[i] && continue
        if dims_x[i]==1
            stride_x[i]=0
        else
            stride_y[i]=0
        end
    end
    return stride_x, stride_y, stride_z
end

# Additional imports: fns in binary_ops are defined using broadcasted.
import Base: +, -, *, /, \

# Here we'll just define some functions that specifically do not have broadcasting.
(+)(x::KnetArray{T},y::KnetArray{T}) where {T} = (size(x)==size(y)||throw(DimensionMismatch("$(map(size,(x,y)))"));(.+)(x,y))
(-)(x::KnetArray{T},y::KnetArray{T}) where {T} = (size(x)==size(y)||throw(DimensionMismatch("$(map(size,(x,y)))"));(.-)(x,y))
#(*){T}(x::KnetArray{T},y::KnetArray{T})=(.*)(x,y) # This is matmul
#(/){T}(x::KnetArray{T},y::KnetArray{T})=(./)(x,y) # This is another linalg op

# Broadcast max/min haven't been defined in Base:
# max(a::Array,b::Array)=broadcast(max,a,b)
# min(a::Array,b::Array)=broadcast(min,a,b)
# tkelman: These two methods aren't necessary, and overwrite Base. You can get this behavior via max.(a,b), with @compat needed on 0.4.

# import Base: broadcast

# Scalar kernels are defined for scalar,array order only.
# For array,scalar we can get most for free.
@eval begin
    broadcasted(::typeof(+),a::KnetArray{T},s::Number) where {T} = (.+)(T(s),a)
    broadcasted(::typeof(+),s::Number,a::KnetArray{T}) where {T} = (.+)(T(s),a)
    broadcasted(::typeof(-),a::KnetArray{T},s::Number) where {T} = (.+)(T(-s),a)
    broadcasted(::typeof(-),s::Number,a::KnetArray{T}) where {T} = (.-)(T(s),a)
    broadcasted(::typeof(*),a::KnetArray{T},s::Number) where {T} = (.*)(T(s),a)
    broadcasted(::typeof(*),s::Number,a::KnetArray{T}) where {T} = (.*)(T(s),a)
    broadcasted(::typeof(/),a::KnetArray{T},s::Number) where {T} = (.*)(T(1/s),a)
    broadcasted(::typeof(/),s::Number,a::KnetArray{T}) where {T} = (./)(T(s),a)
    broadcasted(::typeof(max),a::KnetArray{T},s::Number) where {T} = max.(T(s),a)
    broadcasted(::typeof(max),s::Number,a::KnetArray{T}) where {T} = max.(T(s),a)
    broadcasted(::typeof(min),a::KnetArray{T},s::Number) where {T} = min.(T(s),a)
    broadcasted(::typeof(min),s::Number,a::KnetArray{T}) where {T} = min.(T(s),a)
    broadcasted(::typeof(^),s::Number,a::KnetArray{T}) where {T} = (.^)(T(s),a)
    # Pow is the one exception, we need to define a separate kernel:
    rpow(s,a)=a^s # only broadcast#rpow is defined above, we need rpow defined
    broadcasted(::typeof(^),a::KnetArray{T},s::Number) where {T} = rpow.(T(s),a)
end

@eval begin
    broadcasted(::typeof(==),a::KnetArray{T},s::Number) where {T} = (T(s).==a)
    broadcasted(::typeof(==),s::Number,a::KnetArray{T}) where {T} = (T(s).==a)
    broadcasted(::typeof(!=),a::KnetArray{T},s::Number) where {T} = (T(s).!=a)
    broadcasted(::typeof(!=),s::Number,a::KnetArray{T}) where {T} = (T(s).!=a)
    broadcasted(::typeof(>),a::KnetArray{T},s::Number) where {T} = (T(s).<a)
    broadcasted(::typeof(>),s::Number,a::KnetArray{T}) where {T} = (T(s).>a)
    broadcasted(::typeof(>=),a::KnetArray{T},s::Number) where {T} = (T(s).<=a)
    broadcasted(::typeof(>=),s::Number,a::KnetArray{T}) where {T} = (T(s).>=a)
    broadcasted(::typeof(<),a::KnetArray{T},s::Number) where {T} = (T(s).>a)
    broadcasted(::typeof(<),s::Number,a::KnetArray{T}) where {T} = (T(s).<a)
    broadcasted(::typeof(<=),a::KnetArray{T},s::Number) where {T} = (T(s).>=a)
    broadcasted(::typeof(<=),s::Number,a::KnetArray{T}) where {T} = (T(s).<=a)
end

# Bcasted methods

for f in Symbol.((+,-,*,/,max,min,^,==,!=,>,>=,<,<=))
    M = which(@__MODULE__, f)
    @eval begin
        broadcasted(::typeof($f),s::Bcasted{<:Number},a::Bcasted{<:KnetArray}) = broadcasted($f, s.value, a.value) |> Bcasted
        broadcasted(::typeof($f),s::Bcasted{<:Number},a::KnetArray) = broadcasted($f, s.value, a) |> Bcasted
        broadcasted(::typeof($f),s::Number,a::Bcasted{<:KnetArray}) = broadcasted($f, s, a.value) |> Bcasted
        broadcasted(::typeof($f),a::Bcasted{<:KnetArray},s::Bcasted{<:Number}) = broadcasted($f, a.value, s.value) |> Bcasted
        broadcasted(::typeof($f),a::KnetArray,s::Bcasted{<:Number}) = broadcasted($f, a, s.value) |> Bcasted
        broadcasted(::typeof($f),a::Bcasted{<:KnetArray},s::Number) = broadcasted($f, a.value, s) |> Bcasted
        ($M).$f(s::Bcasted{<:Number},a::Bcasted{<:KnetArray}) = broadcasted($f, s.value, a.value) |> Bcasted
        ($M).$f(s::Bcasted{<:Number},a::KnetArray) = broadcasted($f, s.value, a) |> Bcasted
        ($M).$f(s::Number,a::Bcasted{<:KnetArray}) = broadcasted($f, s, a.value) |> Bcasted
        ($M).$f(a::Bcasted{<:KnetArray},s::Bcasted{<:Number}) = broadcasted($f, a.value, s.value) |> Bcasted
        ($M).$f(a::KnetArray,s::Bcasted{<:Number}) = broadcasted($f, a, s.value) |> Bcasted
        ($M).$f(a::Bcasted{<:KnetArray},s::Number) = broadcasted($f, a.value, s) |> Bcasted
    end
end

# familiar aliases for broadcasting operations of array & scalar (#7226):
# (+)(a::KnetArray{T},s::Number) where {T} = (.+)(T(s),a)  -- deprecated
# (+)(s::Number,a::KnetArray{T}) where {T} = (.+)(T(s),a)  -- deprecated
# (-)(a::KnetArray{T},s::Number) where {T} = (.+)(T(-s),a) -- deprecated
# (-)(s::Number,a::KnetArray{T}) where {T} = (.-)(T(s),a)  -- deprecated
(*)(a::KnetArray{T},s::Number) where {T} = (.*)(T(s),a)
(*)(s::Number,a::KnetArray{T}) where {T} = (.*)(T(s),a)
(/)(a::KnetArray{T},s::Number) where {T} = (.*)(T(1/s),a)
(\)(s::Number,a::KnetArray{T}) where {T} = (.*)(T(1/s),a)

#(/)(s::Number,a::KnetArray{T}) where {T} = (.*)(T(1/s),a) # TODO: non-elementwise definition in linalg
#(^)(a::KnetArray{T},s::Number) where {T} = (.^)(a,T(s)) # non-elementwise definition in linalg
#(^)(s::Number,a::KnetArray{T}) where {T} = (.^)(T(s),a) # non-elementwise definition in linalg

# Define all overloaded Julia functions for KnetArrays:

for f in binary_ops
    if !isa(f,Tuple); f=(f,); end
    binary_op(f...)
end

# Fix #412 where KnetArray(randn(Float64,4,4,4,4)).^2 gives a 1-D result
broadcasted(::typeof(Base.literal_pow), ::typeof(^), k::KnetArray{T}, n::Val{N}) where {T,N} = broadcasted(^, k, N)


import QuantumControlBase
using QuantumControlBase: getcontrols, getcontrolderivs, discretize_on_midpoints, evalcontrols
using QuantumControlBase: GradVector, TimeDependentGradGenerator
using QuantumPropagators: init_storage, initpropwrk
using ConcreteStructs

# GRAPE workspace (for internal use)
@concrete terse struct GrapeWrk

    # a copy of the objectives
    objectives

    # the adjoint objectives, containing the adjoint generators for the
    # backward propagation
    adjoint_objectives

    # The kwargs from the control problem
    kwargs

    # Tuple of the original controls (probably functions)
    controls

    # TODO: pulsevals0 and pulsevals1
    pulsevals :: Vector{Float64}

    gradient :: Vector{Float64}  # gradient for guess in iterations

    searchdirection :: Vector{Float64}  # search-direction for guess in iterations

    # Result object
    result

    #################################
    # scratch objects, per objective:

    # backward-propagated states
    # note: storage for fw-propagated states is in result.states
    bw_states

    # forward-propagated states (functional evaluation only)
    fw_states

    # foward-propagated grad-vectors
    fw_grad_states

    # gradients ∂τₖ/ϵₗ(tₙ)
    tau_grads :: Vector{Matrix{ComplexF64}}

    # dynamical generator (normal propagation) at a particular point in time
    G

    # dynamica generator for grad-propagation, time-dependent
    TDgradG

    # dynamical generator for grad-propagation at a particular point in time
    gradG

    control_derivs :: Vector{Vector{Union{Function, Nothing}}}

    vals_dict

    bw_storage # backward storage array (per objective)

    fw_prop_wrk # for normal forward propagation

    bw_prop_wrk # for normal propagation

    grad_prop_wrk # for gradient forward propagation

    use_threads :: Bool

end

function GrapeWrk(problem::QuantumControlBase.ControlProblem)
    use_threads = get(problem.kwargs, :use_threads, false)
    objectives = [obj for obj in problem.objectives]
    adjoint_objectives = [adjoint(obj) for obj in problem.objectives]
    controls = getcontrols(objectives)
    control_derivs = [
        getcontrolderivs(obj.generator, controls) for obj in objectives
    ]
    tlist = problem.tlist
    # interleave the pulse values as [ϵ₁(t̃₁), ϵ₂(t̃₁), ..., ϵ₁(t̃₂), ϵ₂(t̃₂), ...]
    # to allow access as reshape(pulsevals0, L :)[l, n] where l is the control
    # index and n is the time index
    pulsevals = convert(
        Vector{Float64},
        reshape(transpose(hcat(
            [discretize_on_midpoints(control, tlist)
            for control in controls]...
        )), :)
    )
    kwargs = Dict(problem.kwargs)
    pulse_options = problem.pulse_options
    if haskey(kwargs, :continue_from)
        @info "Continuing previous optimization"
        result = kwargs[:continue_from]
        if !(result isa GrapeResult)
            # account for continuing from a different optimization method
            result = convert(GrapeResult, result)
        end
        result.iter_stop = get(problem.kwargs, :iter_stop, 5000)
        result.converged = false
        result.start_local_time = now()
        result.message = "in progress"
        pulsevals = convert(
            Vector{Float64},
            reshape(transpose(hcat(
                [discretize_on_midpoints(control, wrk.result.tlist)
                for control in result.optimized_controls]...
            )), :)
        )
    else
        result = GrapeResult(problem)
    end
    gradient = zeros(length(pulsevals))
    searchdirection = zeros(length(pulsevals))
    bw_states = [similar(obj.initial_state) for obj in objectives]
    fw_states = [similar(obj.initial_state) for obj in objectives]
    fw_grad_states = [
        GradVector(obj.initial_state, length(controls))
        for obj in objectives
    ]
    dummy_vals = IdDict(control => 1.0 for (i, control) in enumerate(controls))
    G = [evalcontrols(obj.generator, dummy_vals) for obj in objectives]
    vals_dict = [copy(dummy_vals) for _ in objectives]
    bw_storage = [init_storage(obj.initial_state, tlist) for obj in objectives]
    fw_prop_method = [
        Val(
            QuantumControlBase.get_objective_prop_method(
                obj, :fw_prop_method, :prop_method; problem.kwargs...
            )
        )
        for obj in objectives
    ]
    bw_prop_method = [
        Val(
            QuantumControlBase.get_objective_prop_method(
                obj, :bw_prop_method, :prop_method; problem.kwargs...
            )
        )
        for obj in objectives
    ]
    grad_prop_method = [
        Val(
            QuantumControlBase.get_objective_prop_method(
                obj, :grad_prop_method, :fw_prop_method, :prop_method;
                problem.kwargs...
            )
        )
        for obj in objectives
    ]

    fw_prop_wrk = [
        QuantumControlBase.initobjpropwrk(obj, tlist, fw_prop_method[k];
                                          initial_state=obj.initial_state,
                                          kwargs...)
        for (k, obj) in enumerate(objectives)
    ]
    bw_prop_wrk = [
        QuantumControlBase.initobjpropwrk(obj, tlist, bw_prop_method[k];
                                          initial_state=obj.initial_state,
                                          kwargs...)
        for (k, obj) in enumerate(objectives)
    ]
    TDgradG = [TimeDependentGradGenerator(obj.generator) for obj in objectives]
    gradG = [evalcontrols(G̃_of_t, dummy_vals) for G̃_of_t ∈ TDgradG]
    grad_prop_wrk = [
        initpropwrk(Ψ̃, tlist, grad_prop_method[k], gradG[k]; kwargs...)
        for (k, Ψ̃) in enumerate(fw_grad_states)
    ]
    tau_grads :: Vector{Matrix{ComplexF64}} = [
        zeros(ComplexF64, length(controls), length(tlist)-1)
        for _ in objectives
    ]
    GrapeWrk(
        objectives,
        adjoint_objectives,
        kwargs,
        controls,
        pulsevals,
        gradient,
        searchdirection,
        result,
        bw_states,
        fw_states,
        fw_grad_states,
        tau_grads,
        G,
        TDgradG,
        gradG,
        control_derivs,
        vals_dict,
        bw_storage,
        fw_prop_wrk,
        bw_prop_wrk,
        grad_prop_wrk,
        use_threads
    )
end