# everything related to a DG semidiscretization in 1D,
# currently limited to Lobatto-Legendre nodes

# This method is called when a SemidiscretizationHyperbolic is constructed.
# It constructs the basic `cache` used throughout the simulation to compute
# the RHS etc.
function create_cache(mesh::TreeMesh{1}, equations,
                      dg::DG, RealT, uEltype)
  # Get cells for which an element needs to be created (i.e. all leaf cells)
  leaf_cell_ids = local_leaf_cells(mesh.tree)

  elements = init_elements(leaf_cell_ids, mesh, equations, dg.basis, RealT, uEltype)

  interfaces = init_interfaces(leaf_cell_ids, mesh, elements)

  boundaries = init_boundaries(leaf_cell_ids, mesh, elements)

  cache = (; elements, interfaces, boundaries)

  # Add specialized parts of the cache required to compute the volume integral etc.
  cache = (;cache..., create_cache(mesh, equations, dg.volume_integral, dg, uEltype)...)
  cache = (;cache..., create_cache(mesh, equations, dg.mortar, uEltype)...)

  return cache
end

# The methods below are specialized on the volume integral type
# and called from the basic `create_cache` method at the top.
function create_cache(mesh::TreeMesh{1}, equations,
                      volume_integral::VolumeIntegralFluxDifferencing, dg::DG, uEltype)
  create_cache(mesh, have_nonconservative_terms(equations), equations, volume_integral, dg, uEltype)
end

function create_cache(mesh::TreeMesh{1}, nonconservative_terms::Val{false}, equations,
                      ::VolumeIntegralFluxDifferencing, dg, uEltype)
  NamedTuple()
end

# TODO: MHD in 1D
# function create_cache(mesh::TreeMesh{1}, nonconservative_terms::Val{true}, equations,
#                       ::VolumeIntegralFluxDifferencing, dg, uEltype)
# end

function create_cache(mesh::TreeMesh{1}, equations,
                      volume_integral::VolumeIntegralShockCapturingHG, dg::DG, uEltype)
  element_ids_dg   = Int[]
  element_ids_dgfv = Int[]

  cache = create_cache(mesh, equations,
                       VolumeIntegralFluxDifferencing(volume_integral.volume_flux_dg),
                       dg, uEltype)

  A2dp1_x = Array{uEltype, 2}
  fstar1_threaded = A2dp1_x[A2dp1_x(undef, nvariables(equations), nnodes(dg)+1) for _ in 1:Threads.nthreads()]

  return (; cache..., element_ids_dg, element_ids_dgfv, fstar1_threaded)
end

function create_cache(mesh::TreeMesh{1}, equations,
                      volume_integral::VolumeIntegralPureLGLFiniteVolume, dg::DG, uEltype)

  A2dp1_x = Array{uEltype, 2}
  fstar1_threaded = A2dp1_x[A2dp1_x(undef, nvariables(equations), nnodes(dg)+1) for _ in 1:Threads.nthreads()]

  return (; fstar1_threaded)
end

# The methods below are specialized on the mortar type
# and called from the basic `create_cache` method at the top.
function create_cache(mesh::TreeMesh{1}, equations, mortar_l2::LobattoLegendreMortarL2, uEltype)
  NamedTuple()
end

# TODO: Taal discuss/refactor timer, allowing users to pass a custom timer?

function rhs!(du, u, t,
              mesh::TreeMesh{1}, equations,
              initial_condition, boundary_conditions, source_terms,
              dg::DG, cache)
  # Reset du
  @timeit_debug timer() "reset ∂u/∂t" du .= zero(eltype(du))

  # Calculate volume integral
  @timeit_debug timer() "volume integral" calc_volume_integral!(
    du, u, mesh,
    have_nonconservative_terms(equations), equations,
    dg.volume_integral, dg, cache)

  # Prolong solution to interfaces
  @timeit_debug timer() "prolong2interfaces" prolong2interfaces!(
    cache, u, mesh, equations, dg)

  # Calculate interface fluxes
  @timeit_debug timer() "interface flux" calc_interface_flux!(
    cache.elements.surface_flux_values, mesh,
    have_nonconservative_terms(equations), equations,
    dg, cache)

  # Prolong solution to boundaries
  @timeit_debug timer() "prolong2boundaries" prolong2boundaries!(
    cache, u, mesh, equations, dg)

  # Calculate boundary fluxes
  @timeit_debug timer() "boundary flux" calc_boundary_flux!(
    cache, t, boundary_conditions, mesh, equations, dg)

  # Calculate surface integrals
  @timeit_debug timer() "surface integral" calc_surface_integral!(
    du, mesh, equations, dg, cache)

  # Apply Jacobian from mapping to reference element
  @timeit_debug timer() "Jacobian" apply_jacobian!(
    du, mesh, equations, dg, cache)

  # Calculate source terms
  @timeit_debug timer() "source terms" calc_sources!(
    du, u, t, source_terms, equations, dg, cache)

  return nothing
end

function calc_volume_integral!(du, u,
                               mesh::Union{TreeMesh{1}, CurvedMesh{1}},
                               nonconservative_terms::Val{false}, equations,
                               volume_integral::VolumeIntegralWeakForm,
                               dg::DGSEM, cache)
  @unpack derivative_dhat = dg.basis

  @threaded for element in eachelement(dg, cache)
    for i in eachnode(dg)
      u_node = get_node_vars(u, equations, dg, i, element)

      flux1 = flux(u_node, 1, equations)
      for ii in eachnode(dg)
        integral_contribution = derivative_dhat[ii, i] * flux1
        add_to_node_vars!(du, integral_contribution, equations, dg, ii, element)
      end
    end
  end

  return nothing
end

function calc_volume_integral!(du, u,
                               mesh::TreeMesh{1},
                               nonconservative_terms, equations,
                               volume_integral::VolumeIntegralFluxDifferencing,
                               dg::DGSEM, cache)
  @threaded for element in eachelement(dg, cache)
    split_form_kernel!(du, u, nonconservative_terms, equations, volume_integral.volume_flux, dg, cache, element)
  end
end

@inline function split_form_kernel!(du::AbstractArray{<:Any,3}, u,
                                    nonconservative_terms::Val{false}, equations,
                                    volume_flux, dg::DGSEM, cache,
                                    element, alpha=true)
  # true * [some floating point value] == [exactly the same floating point value]
  # This can (hopefully) be optimized away due to constant propagation.
  @unpack derivative_split = dg.basis

  # Calculate volume integral in one element
  for i in eachnode(dg)
    u_node = get_node_vars(u, equations, dg, i, element)

    # x direction
    # use consistency of the volume flux to make this evaluation cheaper
    flux1 = flux(u_node, 1, equations)
    integral_contribution = alpha * derivative_split[i, i] * flux1
    add_to_node_vars!(du, integral_contribution, equations, dg, i, element)
    # use symmetry of the volume flux for the remaining terms
    for ii in (i+1):nnodes(dg)
      u_node_ii = get_node_vars(u, equations, dg, ii, element)
      flux1 = volume_flux(u_node, u_node_ii, 1, equations)
      integral_contribution = alpha * derivative_split[i, ii] * flux1
      add_to_node_vars!(du, integral_contribution, equations, dg, i,  element)
      integral_contribution = alpha * derivative_split[ii, i] * flux1
      add_to_node_vars!(du, integral_contribution, equations, dg, ii, element)
    end
  end
end

# TODO: Taal dimension agnostic
function calc_volume_integral!(du, u,
                               mesh::TreeMesh{1},
                               nonconservative_terms, equations,
                               volume_integral::VolumeIntegralShockCapturingHG,
                               dg::DGSEM, cache)
  @unpack element_ids_dg, element_ids_dgfv = cache
  @unpack volume_flux_dg, volume_flux_fv, indicator = volume_integral

  # Calculate blending factors α: u = u_DG * (1 - α) + u_FV * α
  alpha = @timeit_debug timer() "blending factors" indicator(u, equations, dg, cache)

  # Determine element ids for DG-only and blended DG-FV volume integral
  pure_and_blended_element_ids!(element_ids_dg, element_ids_dgfv, alpha, dg, cache)

  # Loop over pure DG elements
  @timeit_debug timer() "pure DG" @threaded for element in element_ids_dg
    split_form_kernel!(du, u, nonconservative_terms, equations, volume_flux_dg, dg, cache, element)
  end

  # Loop over blended DG-FV elements
  @timeit_debug timer() "blended DG-FV" @threaded for element in element_ids_dgfv
    alpha_element = alpha[element]

    # Calculate DG volume integral contribution
    split_form_kernel!(du, u, nonconservative_terms, equations, volume_flux_dg, dg, cache, element, 1 - alpha_element)

    # Calculate FV volume integral contribution
    fv_kernel!(du, u, equations, volume_flux_fv, dg, cache, element, alpha_element)
  end

  return nothing
end

# TODO: Taal dimension agnostic
function calc_volume_integral!(du, u,
                               mesh::TreeMesh{1},
                               nonconservative_terms, equations,
                               volume_integral::VolumeIntegralPureLGLFiniteVolume,
                               dg::DGSEM, cache)
  @unpack volume_flux_fv = volume_integral

  # Calculate LGL FV volume integral
  @threaded for element in eachelement(dg, cache)
    fv_kernel!(du, u, equations, volume_flux_fv, dg, cache, element, true)
  end

  return nothing
end

@inline function fv_kernel!(du::AbstractArray{<:Any,3}, u::AbstractArray{<:Any,3},
                            equations, volume_flux_fv, dg::DGSEM, cache, element, alpha=true)
  @unpack fstar1_threaded = cache
  @unpack inverse_weights = dg.basis

  # Calculate FV two-point fluxes
  fstar1 = fstar1_threaded[Threads.threadid()]
  calcflux_fv!(fstar1, u, equations, volume_flux_fv, dg, element)

  # Calculate FV volume integral contribution
  for i in eachnode(dg)
    for v in eachvariable(equations)
      du[v, i, element] += ( alpha *
                             (inverse_weights[i] * (fstar1[v, i+1] - fstar1[v, i])) )

    end
  end

  return nothing
end

@inline function calcflux_fv!(fstar1, u::AbstractArray{<:Any,3},
                              equations, volume_flux_fv, dg::DGSEM, element)

  fstar1[:, 1,           ] .= zero(eltype(fstar1))
  fstar1[:, nnodes(dg)+1,] .= zero(eltype(fstar1))

  for i in 2:nnodes(dg)
    u_ll = get_node_vars(u, equations, dg, i-1, element)
    u_rr = get_node_vars(u, equations, dg, i,   element)
    flux = volume_flux_fv(u_ll, u_rr, 1, equations) # orientation 1: x direction
    set_node_vars!(fstar1, flux, equations, dg, i)
  end

  return nothing
end

function prolong2interfaces!(cache, u,
                             mesh::TreeMesh{1}, equations, dg::DG)
  @unpack interfaces = cache

  @threaded for interface in eachinterface(dg, cache)
    left_element  = interfaces.neighbor_ids[1, interface]
    right_element = interfaces.neighbor_ids[2, interface]

    # interface in x-direction
    for v in eachvariable(equations)
      interfaces.u[1, v, interface] = u[v, nnodes(dg), left_element]
      interfaces.u[2, v, interface] = u[v,          1, right_element]
    end
  end

  return nothing
end

function calc_interface_flux!(surface_flux_values,
                              mesh::TreeMesh{1},
                              nonconservative_terms::Val{false}, equations,
                              dg::DG, cache)
  @unpack surface_flux = dg
  @unpack u, neighbor_ids, orientations = cache.interfaces

  @threaded for interface in eachinterface(dg, cache)
    # Get neighboring elements
    left_id  = neighbor_ids[1, interface]
    right_id = neighbor_ids[2, interface]

    # Determine interface direction with respect to elements:
    # orientation = 1: left -> 2, right -> 1
    left_direction  = 2 * orientations[interface]
    right_direction = 2 * orientations[interface] - 1

    # Call pointwise Riemann solver
    u_ll, u_rr = get_surface_node_vars(u, equations, dg, interface)
    flux = surface_flux(u_ll, u_rr, orientations[interface], equations)

    # Copy flux to left and right element storage
    for v in eachvariable(equations)
      surface_flux_values[v, left_direction,  left_id]  = flux[v]
      surface_flux_values[v, right_direction, right_id] = flux[v]
    end
  end
end

# TODO: MHD in 1D
# function calc_interface_flux!(surface_flux_values, mesh::TreeMesh{1},
#                               nonconservative_terms::Val{true}, equations,
#                               dg::DG, cache)
# end

function prolong2boundaries!(cache, u,
                             mesh::TreeMesh{1}, equations, dg::DG)
  @unpack boundaries = cache
  @unpack orientations, neighbor_sides = boundaries

  @threaded for boundary in eachboundary(dg, cache)
    element = boundaries.neighbor_ids[boundary]

    # boundary in x-direction
    if neighbor_sides[boundary] == 1
      # element in -x direction of boundary
      for v in eachvariable(equations)
        boundaries.u[1, v, boundary] = u[v, nnodes(dg), element]
      end
    else # Element in +x direction of boundary
      for v in eachvariable(equations)
        boundaries.u[2, v, boundary] = u[v, 1,          element]
      end
    end
  end

  return nothing
end

# TODO: Taal dimension agnostic
function calc_boundary_flux!(cache, t, boundary_condition::BoundaryConditionPeriodic,
                             mesh::TreeMesh{1}, equations, dg::DG)
  @assert isempty(eachboundary(dg, cache))
end

# TODO: Taal dimension agnostic
function calc_boundary_flux!(cache, t, boundary_condition,
                             mesh::TreeMesh{1}, equations, dg::DG)
  @unpack surface_flux_values = cache.elements
  @unpack n_boundaries_per_direction = cache.boundaries

  # Calculate indices
  lasts = accumulate(+, n_boundaries_per_direction)
  firsts = lasts - n_boundaries_per_direction .+ 1

  # Calc boundary fluxes in each direction
  for direction in eachindex(firsts)
    calc_boundary_flux_by_direction!(surface_flux_values, t, boundary_condition,
                                     equations, dg, cache,
                                     direction, firsts[direction], lasts[direction])
  end
end

function calc_boundary_flux!(cache, t, boundary_conditions::Union{NamedTuple,Tuple},
                             mesh::TreeMesh{1}, equations, dg::DG)
  @unpack surface_flux_values = cache.elements
  @unpack n_boundaries_per_direction = cache.boundaries

  # Calculate indices
  lasts = accumulate(+, n_boundaries_per_direction)
  firsts = lasts - n_boundaries_per_direction .+ 1

  # Calc boundary fluxes in each direction
  calc_boundary_flux_by_direction!(surface_flux_values, t, boundary_conditions[1],
                                   equations, dg, cache, 1, firsts[1], lasts[1])
  calc_boundary_flux_by_direction!(surface_flux_values, t, boundary_conditions[2],
                                   equations, dg, cache, 2, firsts[2], lasts[2])
end

function calc_boundary_flux_by_direction!(surface_flux_values::AbstractArray{<:Any,3}, t,
                                          boundary_condition, equations, dg::DG, cache,
                                          direction, first_boundary, last_boundary)
  @unpack surface_flux = dg
  @unpack u, neighbor_ids, neighbor_sides, node_coordinates, orientations = cache.boundaries

  @threaded for boundary in first_boundary:last_boundary
    # Get neighboring element
    neighbor = neighbor_ids[boundary]

    # Get boundary flux
    u_ll, u_rr = get_surface_node_vars(u, equations, dg, boundary)
    if neighbor_sides[boundary] == 1 # Element is on the left, boundary on the right
      u_inner = u_ll
    else # Element is on the right, boundary on the left
      u_inner = u_rr
    end
    x = get_node_coords(node_coordinates, equations, dg, boundary)
    flux = boundary_condition(u_inner, orientations[boundary], direction, x, t, surface_flux,
                              equations)

    # Copy flux to left and right element storage
    for v in eachvariable(equations)
      surface_flux_values[v, direction, neighbor] = flux[v]
    end
  end

  return nothing
end

function calc_surface_integral!(du, mesh::Union{TreeMesh{1}, CurvedMesh{1}},
                                equations, dg::DGSEM, cache)
  @unpack boundary_interpolation = dg.basis
  @unpack surface_flux_values = cache.elements

  @threaded for element in eachelement(dg, cache)
    for v in eachvariable(equations)
      # surface at -x
      du[v, 1,          element] -= surface_flux_values[v, 1, element] * boundary_interpolation[1,          1]
      # surface at +x
      du[v, nnodes(dg), element] += surface_flux_values[v, 2, element] * boundary_interpolation[nnodes(dg), 2]
    end
  end

  return nothing
end

function apply_jacobian!(du, mesh::Union{TreeMesh{1}, CurvedMesh{1}},
                         equations, dg::DG, cache)

  @threaded for element in eachelement(dg, cache)
    factor = -cache.elements.inverse_jacobian[element]

    for i in eachnode(dg)
      for v in eachvariable(equations)
        du[v, i, element] *= factor
      end
    end
  end

  return nothing
end

# TODO: Taal dimension agnostic
function calc_sources!(du, u, t, source_terms::Nothing,
                       equations::AbstractEquations{1}, dg::DG, cache)
  return nothing
end

function calc_sources!(du, u, t, source_terms,
                       equations::AbstractEquations{1}, dg::DG, cache)

  @threaded for element in eachelement(dg, cache)
    for i in eachnode(dg)
      u_local = get_node_vars(u, equations, dg, i, element)
      x_local = get_node_coords(cache.elements.node_coordinates, equations, dg, i, element)
      du_local = source_terms(u_local, x_local, t, equations)
      add_to_node_vars!(du, du_local, equations, dg, i, element)
    end
  end

  return nothing
end


_hostname = gethostname()
experiment_path = ""
if _hostname == "vit-ThinkPad-E470"
	experiment_path = "/home/vit/vyzkum/bulletproof/pdf_experiment/experiments"
elseif	_hostname in ["gpu-node", "gpu-titan"]
	experiment_path = "/compass/home/skvara/vyzkum/bulletproof/pdf_experiment/experiments"
elseif	_hostname == "tarbik.utia.cas.cz"
	experiment_path = "/home/skvara/work/bulletproof/pdf_experiment/experiments"
end

"""
	elbo(model, x[; β])

Elbo loss function for convolutional models and 4D inputs.
"""
function elbo(m::GenerativeModels.AbstractVAE, x::AbstractArray{T,4}; β=1) where T
    z = rand(m.encoder, x)
    vecdim = reduce(*, size(x)[1:3])
    llh = mean(-loglikelihood(m.decoder, reshape(x, vecdim, :), z))
    kld = mean(kl_divergence(m.encoder, m.prior, x))
    llh + β*kld
end

"""
	reconstruct(m,x)

Reconstructs a sample `x` through the model `m`.
"""
reconstruct(m,x) = reshape(mean(m.decoder, rand(m.encoder, x)), size(x))
reconstruct(x) = reconstruct(model, x)

"""
	history_progress_cb(prog, model, test_data, history, elbo)

Returns a callback function that stores loss values in the history object and also
prints the current progress.
"""
function history_progress_cb(prog::Progress, model::GenerativeModels.AbstractVAE,
							 test_data::AbstractArray, h::MVHistory, elbox::Function)
	iter = 0
	function cb()
		# get losses
		elbo = elbox(test_data) |> cpu
		mse = Flux.mse(test_data, reconstruct(model, test_data)) |> cpu

		# write them to history
		ntuple = GMExtensions.@ntuple elbo mse
		GMExtensions.push_ntuple!(h, ntuple)

		# update the progress bar
		next!(prog, showvalues = [(k, ntuple[k]) for k in keys(ntuple)])

		# save the model
		#iter += 1
		#if save_frequency != nothing && iter%save_frequency == 0
		#	save_checkpoint(filename, model, h)
		#end
	end
end

"""
	load_pdf_data()

Returns a 4D array of pdf representations together with codes of original files and a vector containing
the number of pages.
"""
function load_pdf_data()
	extracted_data = load(PdfFraud.extracted_data)
	representations = extracted_data["representations"]
	sha256 = extracted_data["sha256"]
	labels = extracted_data["labels"]

	# this has to be sone in this way, otherwise stack overflow errs
	lens = map(x->size(x,3), representations)
	inds = cumsum(lens)
	X = zeros(Float32, 64,48,1,sum(lens))
	for (i,representation) in enumerate(representations)
		X[:,:,:,(inds[i]-lens[i]+1):inds[i]] = reshape(representation, 64, 48, 1, :)
	end

	# also expand labels by page
	labels_x = vcat(map(x -> repeat([x[1]], x[2]), zip(labels, lens))...)

	X, labels_x, sha256, lens
end

"""
	save_model(filename, model, history[, experiment_args])
"""
function save_model(filename::String, model::GenerativeModels.AbstractGM,
					history::MVHistory, experiment_args::Dict=Dict())
	model = model |> cpu
    d = @dict model history experiment_args

    @info "Saving checkpoint at $filename"
    DrWatson.tagsave(filename, d, safe=true)
end

"""
	load_model(filename)
"""
function load_model(filename::String)
	@info "Loading checkpoint from $filename"
    res = BSON.load(filename)
    res[:model], res[:history], res[:experiment_args]
end

"""
	encode(model::GenerativeModels.AbstractVAE, X::AbstractArray,
	batchsize::Int, tf::Function)

Encode `X` in batches via encoder of `model` and a given function (e.g. mean/rand).
"""
function encode(model::GenerativeModels.AbstractVAE, X::AbstractArray,
	batchsize::Int, tf::Function)
	s = size(X, 4)
	hcat(map(i->tf(model.encoder, X[:,:,:,(i-1)*batchsize+1:min(i*batchsize, s)]),
		1:ceil(Int, s/batchsize))...)
end


module ITensorChainRules

using ITensors.NDTensors
using ITensors: Indices

using ChainRulesCore
using ..ITensors

import ChainRulesCore: rrule

ITensors.dag(z::AbstractZero) = z

broadcast_notangent(a) = broadcast(_ -> NoTangent(), a)

include(joinpath("NDTensors", "tensor.jl"))
include(joinpath("NDTensors", "dense.jl"))
include("indexset.jl")
include("itensor.jl")
include(joinpath("physics", "sitetype.jl"))
include(joinpath("LazyApply", "LazyApply.jl"))
include(joinpath("mps", "abstractmps.jl"))
include(joinpath("mps", "mpo.jl"))
include(joinpath("Ops", "Ops.jl"))
include("zygoterules.jl")

@non_differentiable broadcast_notangent(::Any)
@non_differentiable Index(::Any...)
@non_differentiable delta(::Any...)
@non_differentiable dag(::Index)
@non_differentiable inds(::Any...)
@non_differentiable commoninds(::Any...)
@non_differentiable noncommoninds(::Any...)
@non_differentiable uniqueinds(::Any...)
@non_differentiable SiteType(::Any)
@non_differentiable ITensors._sitetypes(::Any)
@non_differentiable addtags(::TagSet, ::Any)
@non_differentiable ITensors.filter_inds_set_function(::Function, ::Function, ::Any...)
@non_differentiable ITensors.filter_inds_set_function(::Function, ::Any...)

end


module CTMLEs

using Logging
if isdefined(Main, :CTMLE_LOG_LEVEL)
    @eval @Logging.configure(level=$(Main.CTMLE_LOG_LEVEL))
else
    @Logging.configure(level=WARNING)
end

export ctmle,

SearchStrategy,
ForwardStepwise,
PreOrdered,
OrderingStrategy,
LogisticOrdering,
PartialCorrOrdering,
HDPSOrdering,

StratifiedKfold,
StratifiedRandomSub,

flucinfo

using StatsBase, MLBase
using ..LReg, ..Common, ..Parameters, ..Qmodels

import ..Qmodels: fluctuate!, defluctuate!, Fluctuation

include("pcor.jl")
include("strategies.jl")

subset{P<:Parameter}(param::P, idx::AbstractVector{Int}) =
    P(map(d -> subset(d, idx), regimens(param))...)

subset(d::StaticRegimen, idx) = d
subset(d::DynamicRegimen, idx) = DynamicRegimen(d.a[idx])

"""
A `Qchunk` holds (a subset of) the data set and corresponding indexes along with a `Qmodel` which can compute predictions for
each observation, a `Parameter` (whos definition may depend on each observation) and the value of the emperical risk
for the current `Qmodel`.
"""
type Qchunk{T<:AbstractFloat}
    q::Qmodel{T}
    W::Matrix{T}
    A::Vector{T}
    Y::Vector{T}
    param::Parameter{T}
    idx::AbstractVector{Int}
    gbounds::Vector{T}
    penalize::Bool
    risk::T
end

function make_chunk_subset{T<:AbstractFloat}(logitQnA1::Vector{T}, logitQnA0::Vector{T},
                           W::Matrix{T}, A::Vector{T}, Y::Vector{T},
                           param::Parameter{T}, idx::AbstractVector{Int}, gbounds::Vector{T},
                           penalize::Bool)
    length(logitQnA1) == length(logitQnA0) == size(W, 1) == length(A) == length(Y) ||
        error(ArgumentError("Input sizes do not match"))
    q = Qmodel(logitQnA1[idx], logitQnA0[idx])
    param = subset(param, idx)
    W = W[idx, :]
    A = A[idx]
    Y = Y[idx]
    #Can't compute penalized risk if q isn't fluctuated yet,
    #so set to Inf
    q_risk = penalize? convert(T, Inf) : risk(q, A, Y, param, false)
    Qchunk(q, W, A, Y, param, idx, gbounds, penalize, q_risk)
end

StatsBase.nobs(chunk::Qchunk) = nobs(chunk.q)

"""
A `QCV` holds a training `Qchunk` and a validation `Qchunk`, a sequence of estimated gs corresponding to
each fluctuation of the `Qchunk`s, the `SearchStrategy` in use, and indexes of used and unused covariates.
"""
type QCV{T<:AbstractFloat}
    chunk_train::Qchunk{T}
    chunk_val::Qchunk{T}
    gseq::Vector{LR{T}}
    searchstrategy::SearchStrategy
    used_covars::IntSet
    unused_covars::IntSet
end

"""
Removes the last fluctuation from a `QCV`

** Arguments **

* `qcv` - A `QCV`

Returns last estimated g and last fluctuation from the training and validation chunks.
"""
function defluctuate!(qcv::QCV)
    (pop!(qcv.gseq),
     defluctuate!(qcv.chunk_train.q),
     defluctuate!(qcv.chunk_val.q))
end

"""
Adds a previously calculated fluctuation to a `QCV`

** Arguments **

* `qcv` - A `QCV`
* `g` - A `LR`, estimate of g
* `fluc_train` - fluctuation for the training set
* `fluc_val` - fluctuation for the validation set

Returns training and validation fluctuations.
"""
function fluctuate!(qcv::QCV, g::LR, fluc_train::Fluctuation, fluc_val::Fluctuation)
    push!(qcv.gseq, g)
    fluctuate!(qcv.chunk_train, fluc_train)
    fluctuate!(qcv.chunk_val, fluc_val)
    (fluc_train, fluc_val)
end

"""
Adds a fluctuation to a `QCV` given an estimated g

** Arguments **

* `qcv` - A `QCV`
* `g` - A `LR`, estimate of g

** Details **

A new fluctuation for the training chunk is calculated along with a corresponding flucutation for the validation
chunk, based on the estiamted `g`, then the `QCV` is updated.

Returns training and validation fluctuations.
"""
function fluctuate!(qcv, g::LR)
    train = qcv.train_chunk
    val = qcv.val_chunk
    gn1_train = bound!(predict(g, train.W), train.gbounds)
    gn1_val = bound!(predict(g, val.W), val.gbounds)
    fluc_train = computefluc(train.q, train.param, gn1_train, train.A, train.Y)
    fluc_val = valfluc(fluc_train, val.param, gn1_val, val.A)
    fluctuate!(qcv, g, fluc_train, fluc_val)
end

function makeQCV{T<:AbstractFloat}(logitQnA1::Vector{T}, logitQnA0::Vector{T},
                                   W::Matrix{T}, A::Vector{T}, Y::Vector{T},
                                   param::Parameter{T}, searchstrategy::SearchStrategy,
                                   idx_train::AbstractVector{Int}, gbounds::Vector{T},
                                   penalize::Bool)
    n, p = size(W)
    idx_train = sort(idx_train)
    idx_val = setdiff(1:n, idx_train)
    chunk_train = make_chunk_subset(logitQnA1, logitQnA0, W, A, Y, param, idx_train, gbounds, penalize)
    chunk_val= make_chunk_subset(logitQnA1, logitQnA0, W, A, Y, param, idx_val, gbounds, false)
    used_covars = IntSet()
    unused_covars = setdiff(IntSet(1:p), used_covars)
    gseq = LR{T}[]
    QCV(chunk_train, chunk_val, gseq, deepcopy(searchstrategy), used_covars, unused_covars)
end

immutable FluctuationInfo
    steps::Int
    covar_order::Vector{Vector{Int}}
    new_flucseq::Vector{Bool}
    added_each_fluc::Vector{Vector{Int}}
    function FluctuationInfo(steps::Int, covar_order::Vector{Vector{Int}},
                             new_flucseq::Vector{Bool},
                             added_each_fluc::Vector{Vector{Int}})
        steps == length(covar_order) == length(new_flucseq) || throw(ArgumentError("steps should equal length of covar_order and new_flucseq"))
        length(added_each_fluc) == sum(new_flucseq) || throw(ArgumentError("length(added_each_fluc) not equal to number of new flucs"))
        new(steps, covar_order, new_flucseq, added_each_fluc)
    end
end

function FluctuationInfo{T}(steps::Int, gseq::Vector{LR{T}}, new_flucseq::Vector{Bool})
    new_flucseq[1] || throw(ArgumentError("new_flucseq[1] should be true"))
    first_covars = Vector{Int}[convert(Vector{Int}, gseq[1].idx)]::Vector{Vector{Int}}
    added_covar_order = [setdiff(a.idx, b.idx)::Vector{Int} for (a, b) in zip(gseq[2:end], gseq[1:end-1])]::Vector{Vector{Int}}
    covar_order = [first_covars; added_covar_order]::Vector{Vector{Int}}
    start_new_fluc = find(new_flucseq)
    end_fluc = [find(new_flucseq)[2:end] - 1; length(new_flucseq)]
    added_each_fluc = [vcat(covar_order[st:en]...)::Vector{Int}
                       for (st, en) in zip(start_new_fluc, end_fluc)]
    FluctuationInfo(steps, covar_order, new_flucseq, added_each_fluc)
end

function Base.show(io::IO, flucinfo::FluctuationInfo)
    println(io, "Fluctuation info:")
    println(io, "Covariates added in $(flucinfo.steps) steps to $(length(flucinfo.added_each_fluc)) fluctuations.")
    for (i, new_covars) in enumerate(flucinfo.added_each_fluc)
        println(io, "Fluc $i covars added: $new_covars")
    end
end

type CTMLE{T<:AbstractFloat} <: AbstractScalarEstimate
    psi::T
    ic::Vector{T}
    n::Int
    estimand::AbstractString
    seasrchstrategy::SearchStrategy
    flucinfo::FluctuationInfo
    function CTMLE(psi, ic, n, estimand, searchstrategy, flucinfo)
        est = ScalarEstimate(psi, ic)
        new(est.psi, est.ic, n, estimand, searchstrategy, flucinfo)
    end
end

CTMLE{T<:AbstractFloat}(psi::T, ic::Vector{T}, n::Int, estimand::AbstractString,
                        searchstrategy::SearchStrategy, flucinfo::FluctuationInfo) =
    CTMLE{T}(psi, ic, n, estimand, searchstrategy, flucinfo)

"""
Performs colaborative targeted minimum loss-based estimation

** Arguments **

* `logitQnA1` - Vector of length n of initial estimates logit(\bar{Q}_n(1, W_i))
* `logitQnA0` - Vector of length n of initial estimates logit(\bar{Q}_n(0, W_i))
* `W` - Matrix of covariates to be potentially used to estimate g. n rows.
        The first column should be all ones for an intercept.
* `A` - Vector of length n of treatments, 0.0 or 1.0
* `Y` - Vector of length n of outcomes, bounded between 0 and 1.

** Keyword Arguments **

* `param` - Target parameter. See ?ATE or ?Mean. Default is ATE()
* `gbounds` - values for bounding g. Estimates of g will be bounded by `extrema(gbounds)`.
   To only bound away from 0 (or 1), set to e.g. `[0.01, 1.0]` (or `[0.0, 0.99]`).
   To turn off bounding all together, set to `[0.0, 1.0]`. Defaults to `[0.01, 0.99]`.
* `penalize_risk` - should a penalized risk be used when choosing covariates to estimate g?
   Defaults to false.
* `searchstrategy` - Strategy for adding covariates to estimates of g. See ?SearchStrategy.
* `cvplan` - An iterator of vectors of Int indexes of observations in each training set for
   cross validaiton. StratifiedKfold and StratifiedRandomSub from the MLBase package are useful
   here. Defaults to 10 fold CV stratifying by treatment for outcomes with more than 2 levels,
   or by treatment and outcome for binary outcomes.
* `patience` - For how many steps should CV continue after finding a local optimum? Defaults to typemax(Int)

"""
function ctmle{T<:AbstractFloat}(logitQnA1::Vector{T}, logitQnA0::Vector{T},
                                 W::Matrix{T}, A::Vector{T}, Y::Vector{T};
                                 param::Parameter{T}=ATE(),
                                 gbounds::Vector{T}=[0.01, 0.99],
                                 penalize_risk::Bool=false,
                                 searchstrategy::SearchStrategy = ForwardStepwise(),
                                 cvplan = StratifiedKfold(length(unique(Y))<3? zip(A, Y) : A, 10),
                                 patience::Int=typemax(Int)
                                )
    n,p = size(W)
    n == length(logitQnA1) ==
        length(logitQnA0) ==
        length(A) ==
        length(Y) ||
        throw(ArgumentError("logitQnA1, logitQnA0, A, and Y should have length equal to the number of rows of W."))

    all(W[:, 1] .== 1) || throw(ArgumentError("The first column of W should be all ones."))

    #create vector of QCV objects
    cvqs = [makeQCV(logitQnA1, logitQnA0, W, A, Y, param, searchstrategy,
                    idx_train, gbounds, penalize_risk)::QCV{T}
            for idx_train in cvplan]

    #using cross-validation, find best number of steps to take
    best_steps = find_steps(cvqs, patience=patience)

    #build a chunk with the full data set
    fullchunk = Qchunk(Qmodel(logitQnA1, logitQnA0), W, A, Y, param, 1:n, gbounds, penalize_risk, convert(T, Inf))
    gseq = LR{T}[]
    new_flucseq = Bool[]
    used_covars = IntSet()
    unused_covars = IntSet(1:p)
    #and add covariates steps # of times, saving g fits
    @info("Fitting final estimate of Q")
    for step in 1:best_steps
#         @debug("used_covars: $used_covars")
#         @debug("unused_covars: $unused_covars")
        gfit, new_fluc = add_covars!(fullchunk, searchstrategy, used_covars, unused_covars)
        push!(gseq, gfit)
        push!(new_flucseq, new_fluc)
        @info(if length(gseq) == 1
                  "Step $step of $best_steps complete, covariate [1] added in new fluctuation."
              else
                  whichfluc = new_fluc? "in new" : "to current"
                  whichcovars = setdiff(gseq[end].idx, gseq[end-1].idx)
                  "Step $step of $best_steps complete, covariates $whichcovars added $whichfluc fluctuation."
              end)
    end

    #create flucinfo
    flucinfo = FluctuationInfo(best_steps, gseq, new_flucseq)

    q=fullchunk.q
    CTMLE(applyparam(param, q, A, Y)..., nobs(q), estimand(param), searchstrategy, flucinfo)

end

flucinfo(est::CTMLE) = est.flucinfo

"""
Determines the number of times covariates should be added using cross validation.

** Arguments **

* `qcvs` - A vector of one or more `QCV`s

** Keyword Arguments **

* `patience` - An integer. How many more steps should we try after finding a local minimum?
Defaults to a big number.

** Details **

Adds covariates to each `qcv` and computes the mean risk on validation chunks.
If minimum validation risk occured more than `patience` steps ago, we stop looking.
So for example if `patience` is `2`, we try for two more steps after finding a local minimum to see
if we can do better.
If `patience` is `0`, we stop as soon as risk starts increasing.
If `pateince` is negative, we stop after 1 step, which is probably not a good idea.

Returns best risk and number of steps.

"""
function find_steps{T<:AbstractFloat}(qcvs::Vector{QCV{T}}; patience::Int=typemax(Int))
    notdone = true
    steps = 0
    best_steps = 0
    best_risk = Inf

    @info("Choosing number of steps via cross validation")
    while notdone
        steps += 1
        map(add_covars!, qcvs)
        val_risk = mean(map(qcv -> qcv.chunk_val.risk, qcvs))
        if val_risk < best_risk
            best_risk = val_risk
            best_steps = steps
        end
        @info("Step $steps validation risk: $val_risk. Best risk so far: $best_risk at step $best_steps")
        if steps - best_steps > patience
            @info("No improvement in $patience steps, stopping at step $steps.")
            notdone = false
        end
        if all(map(qcv -> isempty(qcv.unused_covars), qcvs))
            @info("Crossvalidation completed after $steps steps.")
            notdone = false
        end
    end
    best_steps
end

"""
Adds next covariate(s) to a Qchunk based on a search strategy

** Arguments **

* `chunk` - a `Qchunk`
* `searchstrategy` - a `SearchStrategy`
* `used_covars` - indexes of covariates already included in g
* `unused_covars` - indexes of covariates not yet included in g

** Details **

If `used_covars` is emtpy (because this is the first time `add_covars!` has been called on this chunk,
a fluctuation based on an intercept only model is added to Q.
Otherwise, next covariates are added where the order is determined by `searchstrategy`

Returns most recent `gfit` object and a boolean which is `true` if the most recent covariates were added to a new fluctuation
or `false` if the previous fluctuation was updated with new covariates.

"""
function add_covars!(chunk::Qchunk, searchstrategy::SearchStrategy, used_covars, unused_covars)
    if isempty(used_covars)
#         @debug("Initial fluctuation")
        push!(used_covars, 1)
        delete!(unused_covars, 1)
        gfit = lreg(chunk.W, chunk.A, subset=1:1)
        fluc = computefluc(chunk.q, chunk.param, bound!(predict(gfit, chunk.W), chunk.gbounds), chunk.A, chunk.Y)
        fluctuate!(chunk.q, fluc)
        chunk.risk = risk(chunk.q,chunk.A, chunk.Y, chunk.param, chunk.penalize)
        return gfit, true
    else
#         @debug("Adding covariate(s)")
        abc = add_covars!(searchstrategy, chunk.q, chunk.param, chunk.W, chunk.A, chunk.Y,
                          used_covars, unused_covars, chunk.risk, chunk.gbounds, chunk.penalize)
        q_risk, gfit, new_fluc = abc
        chunk.risk=q_risk
        return gfit, new_fluc
    end
end

"""
Adds next covariate(s) to a QCV.

** Arguments **

* `qcv` - a `QCV`

** Details **
If `qcv` has no unused covariates, none are added.  New covariates are first added to the training chunk,
`qcv.train`. Then the validation chunk fluctuated using the newest `gfit`.
Returns updated `qcv`.

"""
function add_covars!(qcv::QCV)
    if isempty(qcv.unused_covars)
        @debug("No covariates to add")
        return qcv
    end
    train = qcv.chunk_train
    val = qcv.chunk_val
    #find next fluctuation for q_train
    # update q_val, which dpeends on g fit on q_train, and fluc

    #add next covar to training chunk and store gfit
    gfit, new_fluc = add_covars!(train, qcv.searchstrategy, qcv.used_covars, qcv.unused_covars)
    push!(qcv.gseq, gfit)

    #add newest fluctuation to val chunk and compute risk
    val = qcv.chunk_val
    gn1_val = bound!(predict(gfit, val.W), val.gbounds)
    if !new_fluc
        defluctuate!(val.q)
    end
    fluctuate!(val.q, valfluc(lastfluc(train.q), val.param, gn1_val, val.A))
    val.risk = risk(val.q, val.A, val.Y, val.param, false)

    qcv
end

function bound!{T<:AbstractFloat}(gn1::Vector{T}, gbounds::Vector{T})
    mn, mx = extrema(gbounds)
    mn < 0.5 || ArgumentError("minimum(gbounds) should be < 0.5")
    mx > 0.5 || ArgumentError("maximum(gbounds) should be > 0.5")

    mn > 0.0 || mx < 1.0 || return gn1

    for i in 1:length(gn1)
        @inbounds gn1[i] = max(mn, min(mx,gn1[i]))
    end
    gn1
end

function risk{T<:AbstractFloat}(q::Qmodel{T}, A::Vector{T}, Y::Vector{T}, param::Parameter{T}, penalize::Bool)
    total_loss = 0.0
    length(Y) == nobs(q) || error()
    lp = linpred(q, A)
    for i in 1:nobs(q)
      total_loss += LReg.loss(Y[i], lp[i])
    end
    q_risk = total_loss/nobs(q)

    if penalize
        ic = applyparam(param, q, A, Y)[2]
        pen = var(ic) / length(A)
    else
        pen = 0.0
    end
    q_risk + pen
end

end # module


import QuantumControlBase
using QuantumControlBase:
    getcontrols, getcontrolderivs, discretize_on_midpoints, evalcontrols
using QuantumControlBase: GradVector, TimeDependentGradGenerator
using QuantumControlBase.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

    fg_count::Vector{Int64}

    # 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

    # backward-propagated grad-vectors
    bw_grad_states

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

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

    # dynamical generator for grad-bw-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

    fw_storage # backward storage array (per objective)

    fw_prop_wrk # for normal forward propagation

    grad_prop_wrk # for gradient forward propagation

    use_threads::Bool

end

function GrapeWrk(problem::QuantumControlBase.ControlProblem; verbose=false)
    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
    fg_count = zeros(Int64, 2)
    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]
    bw_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]
    fw_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,
                :bw_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,
            info_msg="Initializing fw-prop of objective $k",
            verbose=verbose,
            kwargs...
        ) for (k, obj) in enumerate(objectives)
    ]
    TDgradG = [TimeDependentGradGenerator(obj.generator) for obj in adjoint_objectives]
    gradG = [evalcontrols(G̃_of_t, dummy_vals) for G̃_of_t ∈ TDgradG]
    grad_prop_wrk = [
        _init_objgradpropwrk(
            Ψ̃,
            tlist,
            grad_prop_method[k],
            gradG[k];
            verbose=verbose,
            info_msg="Initializing gradient bw-prop of objective $k",
            kwargs...
        ) for (k, Ψ̃) in enumerate(bw_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,
        fg_count,
        result,
        bw_states,
        fw_states,
        bw_grad_states,
        tau_grads,
        G,
        TDgradG,
        gradG,
        control_derivs,
        vals_dict,
        fw_storage,
        fw_prop_wrk,
        grad_prop_wrk,
        use_threads
    )
end

function _init_objgradpropwrk(
    Ψ̃,
    tlist,
    method,
    generator;
    verbose=false,
    info_msg=nothing,
    kwargs...
)
    if verbose && !isnothing(info_msg)
        @info info_msg
    end
    initpropwrk(Ψ̃, tlist, method, generator; verbose=verbose, kwargs...)
end


using ArgParse
s = ArgParseSettings()
@add_arg_table s begin
    "output_file"
        help = "Output location to save the new data bundle."
        arg_type = ASCIIString
        required = true
    "input_config"
        help = "File listing which external data files to integrate into the network. These can be"*
               "either BAM files or BED files."*
               "For each line trailing fields may be omitted and follow the tab separated format: "*
               "FILE_NAME SHORT_TITLE LONG_TITLE CELL_TYPE LAB EXPERIMENT_ID ANTIBODY_ID TREATMENTS ORGANISM LIFE_STAGE LINK"
        arg_type = ASCIIString
        required = true
    "--assembly"
        help = "Which assembly genome coordinates are in (currently GRCh38, GRCh38_alt, hg19, mm10, mm9)."
        arg_type = ASCIIString
        default = "GRCh38_alt"
    "--quiet", "-q"
        help = "Don't print anything"
        action = :store_true
end
args = parse_args(s)

using ChromNet
using JSON
using GZip
using JLD
using SamIO
using ProgressMeter

# parse the input arguments
inputConfig = args["input_config"]
outputFile = args["output_file"]
if !ismatch(r"\.ChromNet\.jld$", outputFile)
    outputFile = "$outputFile.ChromNet.jld"
end
quiet = args["quiet"]

# load the metadata
inStream = STDIN
dataFileRoot = "."
if inputConfig != "-"
    inStream = open(inputConfig)
    dataFileRoot = dirname(inputConfig)
end
metadata = parse_config(inStream, dataFileRoot)
inputConfig != "-" && close(inStream)

# extract all the data files and custom dataset ids
fileTypes = ASCIIString[]
fileNames = ASCIIString[]
ids = ASCIIString[]
for (k,v) in metadata
    if v["fileType"] == "unknown"
        quiet || println(STDERR, "Ignoring file of unknown type: ", v["dataFile"])
        continue
    end
    push!(fileTypes, v["fileType"])
    push!(fileNames, v["dataFile"])
    push!(ids, k)
end

referenceContigs = SamIO.assembly[args["assembly"]]

# build the custom data matrix
numBins = ceil(Int64, sum(referenceContigs.sizes) / 1000)
data = zeros(Float32, length(ids), numBins)
@showprogress "Binning BAM/BED files: " for i in 1:length(ids)
    if fileTypes[i] == "bed" || fileTypes[i] == "narrowPeak"
        reader = BinningMap(BedReader(fileNames[i], referenceContigs), 1000)
        while !eof(reader)
            data[i,position(reader)] = value(reader)
            advance!(reader)
        end
        close(reader)

    elseif fileTypes[i] == "bed.gz" || fileTypes[i] == "narrowPeak.gz"
        f = GZip.open(fileNames[i])
        reader = BinningMap(BedReader(f, referenceContigs), 1000)
        while !eof(reader)
            data[i,position(reader)] = value(reader)
            advance!(reader)
        end
        close(reader)

    elseif fileTypes[i] == "bam"
        reader = BinningMap(BamReader(fileNames[i], :any, referenceContigs), 1000)
        while !eof(reader)
            data[i,position(reader)] = value(reader)
            advance!(reader)
        end
        close(reader)
    end
end

# build summary statistics
!quiet && println(STDERR, "Computing means and norms...")
means = vec(mean(data, 2))
norms = Float64[norm(data[i,:] - means[i]) for i in 1:size(data)[1]]

# write the JLD file
!quiet && println(STDERR, "Writing results to $outputFile...")
jldopen(outputFile, "w") do f
    f["data"] = data
    f["means"] = vec(mean(data, 2))
    f["norms"] = Float64[norm(data[i,:] - means[i]) for i in 1:size(data)[1]]
    f["ids"] = ids
    f["P"] = length(ids)
    f["N"] = size(data)[2]
    f["metadata"] = metadata
    f["assembly"] = args["assembly"]
end


# Transmission & distribution coupling functions

## Functions replacing PowerModels or InfrastructureModels functions

""
function run_model(
        t_data::Dict{String,<:Any},
        d_data::Dict{String,<:Any},
        t_model_type::Type,
        d_model_type::Type{BF},
        optimizer::Union{MathOptInterface.AbstractOptimizer, MathOptInterface.OptimizerWithAttributes},
        build_method::Function;
        t_ref_extensions::Vector{<:Function} = Function[],
        d_ref_extensions::Vector{<:Function} = Function[],
        t_solution_processors::Vector{<:Function} = Function[],
        d_solution_processors::Vector{<:Function} = Function[],
        kwargs...
    ) where BF <: _PM.AbstractBFModel

    # Check that transmission and distribution network ids are different.
    if !isempty(intersect(Set(id for (id,nw) in t_data["nw"]), Set(id for (id,nw) in d_data["nw"])))
        Memento.error(_LOGGER, "Transmission and distribution data contain networks having the same IDs.")
    end

    # Instantiate models.
    start_time = time()
    t_pm, d_pm = instantiate_model(t_data, d_data, t_model_type, d_model_type, build_method; t_ref_extensions, d_ref_extensions, kwargs...)
    Memento.debug(_LOGGER, "combined T&D model build time: $(time() - start_time)")

    start_time = time()

    # Solve the optimization model and store the transmission result.
    t_result = _IM.optimize_model!(t_pm; optimizer, solution_processors=t_solution_processors)

    # Build the distribution result using the same model as above.
    d_result = _IM.build_result(d_pm, t_result["solve_time"]; solution_processors=d_solution_processors)

    # The asymmetric code above for building results produces inaccurate debugging messages;
    # this behavior can be fixed by writing a custom optimize_model!() that takes 2 models.

    Memento.debug(_LOGGER, "combined T&D model solution time: $(time() - start_time)")

    # Combine the result objects.
    result = t_result # All fields are pairwise equal, except for "solution".
    solution = result["solution"]
    for (k,v) in d_result["solution"]
        if k == "nw"
            # Merge transmission and distribution "nw" fields.
            # Network ids do not clash (checked before model instantiation).
            solution["nw"] = merge(solution["nw"], v)
        else
            if solution[k] != v
                Memento.warning(_LOGGER, "Transmission and distribution solutions differ on key \"$k\"; only transmission value is kept.")
            end
        end
    end

    return result
end

""
function instantiate_model(
        t_data::Dict{String,<:Any},
        d_data::Dict{String,<:Any},
        t_model_type::Type,
        d_model_type::Type{BF},
        build_method::Function;
        t_ref_extensions::Vector{<:Function} = Function[],
        d_ref_extensions::Vector{<:Function} = Function[],
        kwargs...
    ) where BF <: _PM.AbstractBFModel

    # Instantiate the transmission PowerModels struct, without building the model.
    t_pm = _PM.instantiate_model(t_data, t_model_type, method->nothing; ref_extensions=t_ref_extensions, kwargs...)

    # Instantiate the distribution PowerModels struct, without building the model.
    # Distribution and transmission structs share the same JuMP model. The `jump_model` parameter is used by _IM.InitializeInfrastructureModel().
    d_pm = _PM.instantiate_model(d_data, d_model_type, method->nothing; ref_extensions=d_ref_extensions, jump_model=t_pm.model, kwargs...)

    # Build the combined model.
    build_method(t_pm, d_pm)

    return t_pm, d_pm
end

## Functions that manipulate data structures

"""
    add_td_coupling_data!(d_data; sub_nw, qs_ratio_bound=0.48)

Add to distribution single-network data structures the data needed for T&D coupling.

- Add to distribution network data structure a coupling generator connected to the reference
  bus; if it already exists, overwrite all its parameters except cost.
- Define a bound on reactive power exchanged between transmission and distribution networks:
  it is computed as a fraction `qs_ratio_bound` of the rated power of the distribution
  network (based on the rated power of existing and candidate branches connected to the
  reference bus); default value is 0.48 as per “Network Code on Demand Connection” –
  Commission Regulation (EU) 2016/1388.
- Add to `dim_prop(d_data, :sub_nw, sub_nw)` the following entries:
  - `d_gen`: the id of the coupling generator of distribution network;
  - `qs_ratio_bound`: the aforementioned allowable fraction of the rated power.

Return `d_gen`.

This function is intended to be the last that edits distribution single-network data
structures: it should be called just before `make_multinetwork`.
"""
function add_td_coupling_data!(d_data::Dict{String,Any}; sub_nw::Int, qs_ratio_bound::Float64=0.48)

    # Get the reference bus id
    d_ref_buses = [b for (b,bus) in d_data["bus"] if bus["bus_type"] == 3]
    if length(d_ref_buses) != 1
        Memento.error(_LOGGER, "Distribution network data must have 1 ref bus, but $(length(d_ref_buses)) are present.")
    end
    d_ref_bus = parse(Int, first(d_ref_buses))

    # Get a list of the ids of generators connected to the reference bus
    d_ref_gens = [g for (g,gen) in d_data["gen"] if gen["gen_bus"] == d_ref_bus]
    if length(d_ref_gens) > 1
        Memento.error(_LOGGER, "Distribution network data must have 0 or 1 generators connected to ref bus, but $(length(d_ref_gens)) are present.")
    end

    # Define an upper bound on the rated apparent power based on the rated power of existing and candidate branches connected to the reference bus
    d_s_rate = (
        # `Float64[...]` is to prevent `ArgumentError: reducing over an empty collection is not allowed` in case at least one of the following collections is empty.
        # TODO: when dropping support for Julia < 1.6, remove `Float64[...]` and use `init=0.0` keyword argument of `sum`.
          sum(Float64[branch["rate_a"] for (b,branch) in d_data["branch"]    if (branch["f_bus"]==d_ref_bus || branch["t_bus"]==d_ref_bus) && branch["br_status"]==1]) # In t_bus here, the t stands for "to" (not for "transmission" as in the rest of the function)
        + sum(Float64[branch["rate_a"] for (b,branch) in d_data["ne_branch"] if (branch["f_bus"]==d_ref_bus || branch["t_bus"]==d_ref_bus) && branch["br_status"]==1]) # In t_bus here, the t stands for "to" (not for "transmission" as in the rest of the function)
    )

    # Add a coupling generator connected to d_ref_bus (or use existing one) to model the transmission network
    if isempty(d_ref_gens)
        d_gen_idx = length(d_data["gen"]) + 1 # Assumes that gens have contiguous indices starting from 1, as should be
        d_data["gen"]["$d_gen_idx"] = Dict{String,Any}(
            "gen_bus"      => d_ref_bus,
            "index"        => d_gen_idx,
            "dispatchable" => true,
            "model"        => 2, # Cost model (2 => polynomial cost)
            "ncost"        => 0, # Number of cost coefficients
            "cost"         => Any[]
        )
    else
        d_gen_idx = parse(Int, first(d_ref_gens))
    end

    # Set coupling generator parameters
    d_gen = d_data["gen"]["$d_gen_idx"]
    d_gen["mbase"]      = d_data["baseMVA"]
    d_gen["pmin"]       = -d_s_rate
    d_gen["pmax"]       =  d_s_rate
    d_gen["qmin"]       = -d_s_rate
    d_gen["qmax"]       =  d_s_rate
    d_gen["gen_status"] = 1

    # Store the T&D coupling parameters
    dim_prop(d_data, :sub_nw, sub_nw)["qs_ratio_bound"] = qs_ratio_bound
    dim_prop(d_data, :sub_nw, sub_nw)["d_gen"] = d_gen_idx
end

"""
    add_td_coupling_data!(t_data, d_data; t_bus, sub_nw, qs_ratio_bound=0.48)

Add to transmission and distribution single-network data structures the data needed for T&D coupling.

In addition to `add_td_coupling_data!(d_data; sub_nw, qs_ratio_bound)`, do the following:
- Add to transmission network data structure a coupling generator connected to `t_bus`, that
  is the bus to which the distribution network is to be connected.
- Add to `dim_prop(d_data, :sub_nw, sub_nw)` an entry `t_gen` for the id of the coupling
  generator of transmission network.

Return `t_gen`.

This function is intended to be the last that edits transmission and distribution
single-network data structures: it should be called just before `make_multinetwork`.
"""
function add_td_coupling_data!(t_data::Dict{String,Any}, d_data::Dict{String,Any}; t_bus::Int, sub_nw::Int, qs_ratio_bound::Float64=0.48)

    d_gen_idx = add_td_coupling_data!(d_data; sub_nw, qs_ratio_bound)

    # Check that t_bus exists
    if !haskey(t_data["bus"], "$t_bus")
        Memento.error(_LOGGER, "Bus $t_bus does not exist in transmission network data.")
    end

    # Add a coupling generator connected to t_bus, to model the distribution network
    t_s_rate = (d_data["baseMVA"]/t_data["baseMVA"]) * d_data["gen"]["$d_gen_idx"]["pmax"]
    t_gen_idx = length(t_data["gen"]) + 1 # Assumes that gens have contiguous indices starting from 1, as should be
    t_data["gen"]["$t_gen_idx"] = Dict{String,Any}(
        "gen_bus"      => t_bus,
        "index"        => t_gen_idx,
        "dispatchable" => true,
        "mbase"        => t_data["baseMVA"],
        "pmin"         => -t_s_rate,
        "pmax"         =>  t_s_rate,
        "qmin"         => -t_s_rate,
        "qmax"         =>  t_s_rate,
        "gen_status"   => 1,
        "model"        => 2, # Cost model (2 => polynomial cost)
        "ncost"        => 0, # Number of cost coefficients
        "cost"         => Any[]
    )

    # Add generator id to T&D coupling parameters
    dim_prop(d_data, :sub_nw, sub_nw)["t_gen"] = t_gen_idx
end

## Solution processors

"""
    sol_td_coupling!(pm, solution)

Add T&D coupling data to solution.

Report in `solution` the active and reactive power that distribution network `pm` exchanges with the
transmission network (positive if from transmission to distribution) in units of `baseMVA` of `pm`.
"""
function sol_td_coupling!(pm::_PM.AbstractBFModel, solution::Dict{String,Any})
    solution = _PM.get_pm_data(solution)
    if haskey(solution, "nw")
        nws_sol = solution["nw"]
    else
        nws_sol = Dict("0" => solution)
    end

    for (nw, nw_sol) in nws_sol
        n = parse(Int, nw)
        if !(haskey(dim_prop(pm), :sub_nw) && haskey(dim_prop(pm, n, :sub_nw), "d_gen"))
                Memento.error(_LOGGER, "T&D coupling data is missing from the model of nw $nw.")
        end
        d_gen = dim_prop(pm, n, :sub_nw, "d_gen")
        nw_sol["td_coupling"] = Dict{String,Any}()
        nw_sol["td_coupling"]["p"] = nw_sol["gen"]["$d_gen"]["pg"]
        nw_sol["td_coupling"]["q"] = nw_sol["gen"]["$d_gen"]["qg"]
    end
end

## Functions that group constraint templates, provided for convenience

"""
Connect each distribution nw to the corresponding transmission nw and apply coupling constraints.

The coupling constraint is applied to the two generators that each distribution nw indicates.
"""
function constraint_td_coupling(t_pm::_PM.AbstractPowerModel, d_pm::_PM.AbstractBFModel)
    t_nws = nw_ids(t_pm)

    for s in keys(dim_prop(d_pm, :sub_nw))
        d_nws = nw_ids(d_pm; sub_nw = s)
        for i in 1:length(t_nws)
            t_nw = t_nws[i]
            d_nw = d_nws[i]

            constraint_td_coupling_power_balance(t_pm, d_pm, t_nw, d_nw)
        end
    end
end

## Constraint templates

"""
State the power conservation between a distribution nw and the corresponding transmission nw.
"""
function constraint_td_coupling_power_balance(t_pm::_PM.AbstractPowerModel, d_pm::_PM.AbstractBFModel, t_nw::Int, d_nw::Int)
    sub_nw = dim_prop(d_pm, d_nw, :sub_nw)
    t_gen = sub_nw["t_gen"]
    d_gen = sub_nw["d_gen"]
    t_mbase = _PM.ref(t_pm, t_nw, :gen, t_gen, "mbase")
    d_mbase = _PM.ref(d_pm, d_nw, :gen, d_gen, "mbase")

    constraint_td_coupling_power_balance_active(t_pm, d_pm, t_nw, d_nw, t_gen, d_gen, t_mbase, d_mbase)
    constraint_td_coupling_power_balance_reactive(t_pm, d_pm, t_nw, d_nw, t_gen, d_gen, t_mbase, d_mbase)
end

"""
Apply bounds on reactive power exchange at the point of common coupling (PCC) of a distribution nw.
"""
function constraint_td_coupling_power_reactive_bounds(d_pm::_PM.AbstractBFModel; nw::Int=_PM.nw_id_default)
    sub_nw = dim_prop(d_pm, nw, :sub_nw)
    d_gen = sub_nw["d_gen"]
    qs_ratio_bound = sub_nw["qs_ratio_bound"] # Allowable fraction of rated apparent power

    constraint_td_coupling_power_reactive_bounds(d_pm, nw, d_gen, qs_ratio_bound)
end

## Constraint implementations

""
function constraint_td_coupling_power_balance_active(t_pm::_PM.AbstractPowerModel, d_pm::_PM.AbstractBFModel, t_nw::Int, d_nw::Int, t_gen::Int, d_gen::Int, t_mbase::Float64, d_mbase::Float64)
    t_p_in = _PM.var(t_pm, t_nw, :pg, t_gen)
    d_p_in = _PM.var(d_pm, d_nw, :pg, d_gen)
    JuMP.@constraint(t_pm.model, t_mbase*t_p_in + d_mbase*d_p_in == 0.0) # t_pm.model == d_pm.model
end

""
function constraint_td_coupling_power_balance_reactive(t_pm::_PM.AbstractPowerModel, d_pm::_PM.AbstractBFModel, t_nw::Int, d_nw::Int, t_gen::Int, d_gen::Int, t_mbase::Float64, d_mbase::Float64)
    t_q_in = _PM.var(t_pm, t_nw, :qg, t_gen)
    d_q_in = _PM.var(d_pm, d_nw, :qg, d_gen)
    JuMP.@constraint(t_pm.model, t_mbase*t_q_in + d_mbase*d_q_in == 0.0) # t_pm.model == d_pm.model
end

"Nothing to do because the transmission network model does not support reactive power."
function constraint_td_coupling_power_balance_reactive(t_pm::_PM.AbstractActivePowerModel, d_pm::_PM.AbstractBFModel, t_nw::Int, d_nw::Int, t_gen::Int, d_gen::Int, t_mbase::Float64, d_mbase::Float64)
end

""
function constraint_td_coupling_power_reactive_bounds(d_pm::_PM.AbstractBFModel, d_nw::Int, d_gen::Int, qs_ratio_bound)

    # Compute the rated apparent power of the distribution network, based on the rated power of
    # existing and candidate branches connected to its reference bus. This value depends on the
    # indicator variables of both existing (if applicable) and candidate branches (i.e. whether they
    # are built or not).
    if haskey(_PM.var(d_pm, d_nw), :z_branch) # Some `branch`es can be replaced by `ne_branch`es
        z_branch = _PM.var(d_pm, d_nw, :z_branch)
        s_rate = (
            sum(branch["rate_a"] * get(z_branch, b, 1.0) for (b,branch) in _PM.ref(d_pm, d_nw, :frb_branch))
            + sum(branch["rate_a"] * _PM.var(d_pm, d_nw, :branch_ne, b) for (b,branch) in _PM.ref(d_pm, d_nw, :frb_ne_branch))
        )
    else # No `ne_branch`es at all
        s_rate = sum(branch["rate_a"] for (b,branch) in _PM.ref(d_pm, d_nw, :frb_branch))
    end

    q = _PM.var(d_pm, d_nw, :qg, d_gen) # Exchanged reactive power (positive if from T to D)

    JuMP.@constraint(d_pm.model, q <=  qs_ratio_bound*s_rate)
    JuMP.@constraint(d_pm.model, q >= -qs_ratio_bound*s_rate)
end


"""
    defaultSettings()

Defines the list of configuration variables and returns their default values
in a Dict.
"""
function defaultSettings()
    # Return the default settings. All parameters must be registered here.
    Dict(
        "biggenelength" => 200, # length of the compatibility gene's sequence (if `usebiggenes`)
        "bodytemp"=> 312.0, # body temperature in K for homeothermic species
        "borders" => "absorbing", # border behaviour: absorbing/reflecting/mainland
        #"bufferoutput" => false, # use output buffering (faster execution, more RAM usage) #TODO
        "burn-in" => 1000, # timesteps before invasion starts
        "capgrowth" => false, # strictly limit individuals' size to `repsize`
        "cellsize" => 20e6, # maximum biomass per hectare in gramm (based on Clark et al. 2001)
        "compressgenes" => true, # save gene sequences as integers to reduce memory usage
        "config" => "", # configuration file name
        "debug" => false, # write out debug statements
        "degpleiotropy" => 0.1, # How frequent are pleitropy and polygenic inheritance? 0 <= degpleiotropy < 1
        "dest" => string(Dates.today()), # output folder name
        "dispmean" => 1.0, # maximum value of randomly drawn mean dispersal distance in cells
        "dispshape" => 1.0, # maximum value of randomly drawn shape parameter for dispersal kernel. determines tail fatness/long distance dispersal
        "dispmortality" => false, #boolean toggle for switching mortality dependent on dispersal distance, depends on dispfactor(see below), only in zdisperse()
        "dispfactor"=> 2, #: when dispmean multiplied by dispfactor is larger than  evolved dispersaldistance -> individual dies with 100 % chance, when smaller, chances of survival higher
        "disturbance" => 0, # percentage of individuals killed per update per cell
        "dumpindforfasta"=> false, #should dumpinds be called simultaneously with the fasta output?
        "fasta" => "off", # record fasta data? "off", "compat", "all" (high detail output)
        "fastaoutfreq" => 10, #frequency with which to write fasta data
        "fertility" => exp(28.0), # global base reproduction rate 23.8 from Brown et al. 2004, alternatively 25.0, default 30.0
        "fixtol" => true, # fix mating tolerance globally to `tolerance`
        "global-species-pool" => 0, # size of the global species pool (invasion source)
        "globalmating" => false, # global pollen "dispersal"
        "growthrate" => exp(25.2), # global base growth/biomass production from Brown et al. 2004
        "heterozygosity" => false, # keep track of heterozygosity when studying hybridisation
        "hybridisation" => "off", #toggle hybridisation on and off to allow simultaneaous hybridisation and speciation modes "ecological" and "neutral"
        "indsize" => "seed", # initialize organisms as seed, adult or mixed
        "isolationweight" => 3.0, # additional distance to be crossed when dispersing from or to isolated patches
        "lineages" => false, # record lineage and diversity data (low detail output)
        "linkage" => "random", # gene linkage type (random/full/none)
        "logging" => false, # write output to logfile
        "maps" => "", # comma-separated list of map files
        "maxbreadth" => 5.0, # maximum niche breadth
        "maxloci" => 1, # maximum number of loci/copies per gene
        "maxprec" => 10.0, # max optimum precipitation
        "maxrepsize" => 14.0, # maximal repsize in grams calculated as exp(maxrepsize) -> 1.2 t
        "maxseedsize" => 10.0, # maximal seedsize in grams calculated as exp(maxseedsize) -> 22 kg
        "maxtemp" => 313.0, # max optimum temp in K
        "mortalitytype" => "metabolic", # metabolic/habitat/global
        "minprec" => 0.0, # min optimum precipitation
        "minrepsize" => 3.0, # minimal repsize in grams calculated as exp(minrepsize) -> 20 g
        "minseedsize" => -2.0, # minimal seedsize in grams calculated as exp(minseedsize) -> 0.14 g
        "mintemp" => 273.0, # min optimum temp in K
        "mode" => "default", # experiment type ("default", "invasion", or "zosterops")
        "mortality" => exp(22), # global base mortality from Brown et al. 2004 is 26.3, but competition and dispersal introduce add. mort.
        "mutate" => true, # mutations occur
        "mutationrate" => 3.6e10, # one mutation per generation/individual, corrected for metabolic function
        "nniches" => 2, # number of environmental niches (max. 2) #XXX currently, only 2 is sensible
        "outfreq" => 10, # output frequency
        "overfill" => 1.0, # how much to overfill grid cells beyond their capacity
        "perfecttol" => 2.0, #if mortality is habitat dependent, how tolerant should a perfectly adapted individual be? scales for specialist selection pressure
        "phylconstr" => 0.1, # phylogenetic constraint during mutation and inter-loci variation. scales trait value as sd.
        "popsize" => "metabolic", # initialisation algorithm: metabolic/bodysize/minimal/single
        "precrange" => 0, # max optimum precipitation - deprecated! (use `maxprec`)
        "propagule-pressure" => 0, # number of non-native individuals introduced per invasion event
        "quiet" => false, # don't write output to screen
        "raw" => true, # record raw data
        "scaleadaptation" => 0.0, #either 0.0 or 1.0, needed parameter for the gausscurve funciton. For phylogeny experiment set to 1.0!
        "sdprec" => 0.0, # SD of precipitation change per time step
        "sdtemp" => 0.0, # SD of temperature change per time step
        "seed" => 0, # for the RNG, seed = 0 -> random seed
        "smallgenelength" => 20, # standard gene sequence length (max. 21)
        "speciation" => "neutral", # allow lineage differentiation? off/neutral/ecological
        "species" => Dict{String,Any}[], # define trait values for each Zosterops species
        "static" => false, # whether mainland sites undergo eco-evolutionary processes (implies "mainland" borders)
        "stats" => true, # record population statistics (medium detail output)
        "tolerance" => 0.8, # sequence similarity threshold for reproduction if `fixtol` == true
        # minimal required traitnames
        "traitnames" => ["compat", # place-holder for the neutral ("compatibility") gene
                         "dispmean", # mean dispersal distance
                         "dispshape", # dispersal shape parameter
                         "numpollen", # number of offspring
                         "precopt", # precipitation/resource optimum
                         "prectol", # precipitation/resource tolerance
                         "repsize", # reproductive size
                         "seqsimilarity", # required genetic sequence similarity for mating
                         "selfing", # probability of self-fertilisation
                         "seedsize", # seed size
                         "tempopt", # temperature optimum
                         "temptol"], # temperature tolerance
        "traitsforecospec" => ["dispmean", # mean dispersal distance
                               "dispshape", # dispersal shape parameter
                               "numpollen", # number of offspring
                               "precopt", # precipitation/resource optimum
                               "prectol", # precipitation/resource tolerance
                               "repsize", # reproductive size
                               "seqsimilarity", # required genetic sequence similarity for mating
                               "selfing", # probability of self-fertilisation
                               "seedsize", # seed size
                               "tempopt", # temperature optimum
                               "temptol"], # temperature tolerance
        "usebiggenes" => true, # use a longer sequence for the compatibility gene
    )
end


using Interpolations: LinearInterpolation
import ..ContinuumAbsorption: total_continuum_absorption

"""
    synthesize(atm, linelist, λ_start, λ_stop, [λ_step=0.01]; metallicity=0, abundances=Dict(), vmic=0, ... )

Solve the transfer equation in the model atmosphere `atm` with the transitions in `linelist` at the
from `λ_start` to `λ_stop` in steps of `λ_step` (all Å) to get the resultant astrophysical flux at
each wavelength.  An `AbstractRange` can also be probvided directly in place of `λ_start`, `λ_stop`,
and `λ_step`.

Returns a named tuple with keys:
- `flux`: the output spectrum
- `alpha`: the linear absorption coefficient at each wavelenth and atmospheric layer a Matrix of
   size (layers x wavelengths)
- `number_densities`: A dictionary mapping `Species` to vectors of number densities at each
   atmospheric layer
- `wavelengths`: The vacuum wavelenths (in Å) over which the synthesis was performed.  If
  `air_wavelengths=true` this will not be the same as the input wavelenths.

Optional arguments:
- `metallicity`, i.e. [metals/H] is log_10 solar relative
- `abundances` is a `Dict` mapping atomic symbols to ``A(X)`` format abundances, i.e.
   ``A(x) = \\log_{10}(n_X/n_\\mathrm{H}) + 12``, where ``n_X`` is the number density of ``X``.
   These override `metallicity`.
- `vmic` (default: 0) is the microturbulent velocity, ``\\xi``, in km/s.
- `air_wavelengths` (default: `false`): Whether or not the input wavelengths are air wavelenths to
   be converted to vacuum wavelengths by Korg.  The conversion will not be exact, so that the
   wavelenth range can internally be represented by an evenly-spaced range.  If the approximation
   error is greater than `wavelength_conversion_warn_threshold`, an error will be thrown.
- `wavelength_conversion_warn_threshold` (default: 1e-4): see `air_wavelengths`.
- `line_buffer` (default: 10): the farthest (in Å) any line can be from the provided wavelenth range
   before it is discarded.  If the edge of your window is near a strong line, you may have to turn
   this up.
- `cntm_step` (default 1): the distance (in Å) between point at which the continuum opacity is
  calculated.
- `hydrogen_lines` (default: `true`): whether or not to include H lines in the synthesis.
- `mu_grid`: the range of (surface) μ values at which to calculate the surface flux when doing
   transfer in spherical geometry (when `atm` is a `ShellAtmosphere`).
- `line_cutoff_threshold` (default: `1e-3`): the fraction of the continuum absorption coefficient
   at which line profiles are truncated.  This has major performance impacts, since line absorption
   calculations dominate more syntheses.  Turn it down for more precision at the expense of runtime.
   The default value should effect final spectra below the 10^-3 level.
- `ionization_energies`, a `Dict` mapping `Species` to their first three ionization energies,
   defaults to `Korg.ionization_energies`.
- `partition_funcs`, a `Dict` mapping `Species` to partition functions. Defaults to data from
   Barklem & Collet 2016, `Korg.partition_funcs`.
- `equilibrium_constants`, a `Dict` mapping `Species` representing diatomic molecules to their
   molecular equilbrium constants in partial pressure form.  Defaults to data from
   Barklem and Collet 2016, `Korg.equilibrium_constants`.
"""
function synthesize(atm::ModelAtmosphere, linelist, λ_start, λ_stop, λ_step=0.01
                    ; air_wavelengths=false, wavelength_conversion_warn_threshold=1e-4, kwargs...)
    wls = if air_wavelengths
        len = Int(round((λ_stop - λ_start)/λ_step))+1
        vac_start, vac_stop = air_to_vacuum.((λ_start, λ_stop))
        vac_step = (vac_stop - vac_start) / (len-1)
        wls = StepRangeLen(vac_start, vac_step, len)
        max_diff = maximum(abs.(wls .- air_to_vacuum.(λ_start:λ_step:λ_stop)))
        if max_diff > wavelength_conversion_warn_threshold
            throw(ArgumentError("A linear air wavelength range can't be approximated exactly with a"
                                *"linear vacuum wavelength range. This solution differs by up to " *
                                "$max_diff Å.  Adjust wavelength_conversion_warn_threshold if you" *
                                "want to suppress this error."))
        end
        wls
    else
        StepRangeLen(λ_start, λ_step, Int(round((λ_stop - λ_start)/λ_step))+1)
    end
    synthesize(atm, linelist, wls; kwargs...)
end
function synthesize(atm::ModelAtmosphere, linelist, λs::AbstractRange; metallicity::Real=0.0,
                    vmic::Real=1.0, abundances=Dict(), line_buffer::Real=10.0, cntm_step::Real=1.0,
                    hydrogen_lines=true, mu_grid=0:0.05:1, line_cutoff_threshold=1e-3,
                    ionization_energies=ionization_energies,
                    partition_funcs=partition_funcs, equilibrium_constants=equilibrium_constants)
    #work in cm
    λs = λs * 1e-8
    cntm_step *= 1e-8
    line_buffer *= 1e-8
    cntmλs = (λs[1] - line_buffer - cntm_step) : cntm_step : (λs[end] + line_buffer + cntm_step)

    #sort the lines if necessary and check that λs is sorted
    issorted(linelist; by=l->l.wl) || sort!(linelist, by=l->l.wl)
    if step(λs) < 0
        throw(ArgumentError("λs must be in increasing order."))
    end

    #discard lines far from the wavelength range being synthesized
    linelist = filter(l-> λs[1] - line_buffer*1e-8 <= l.wl <= λs[end] + line_buffer*1e-8, linelist)

    abundances = get_absolute_abundances(metallicity, abundances)
    MEQs = molecular_equilibrium_equations(abundances, ionization_energies, partition_funcs,
                                           equilibrium_constants)

    sorted_cntmνs = c_cgs ./ reverse(cntmλs) #frequencies at which to calculate the continuum

    #float-like type general to handle dual numbers
    α_type = typeof(promote(atm.layers[1].temp, length(linelist) > 0 ? linelist[1].wl : 1.0, λs[1],
                            metallicity, vmic, abundances[1])[1])
    #the absorption coefficient, α, for each wavelength and atmospheric layer
    α = Matrix{α_type}(undef, length(atm.layers), length(λs))
    α_cntm = Vector(undef, length(atm.layers))     #vector of continuum-absorption interpolators
    α5 = Vector{α_type}(undef, length(atm.layers)) #each layer's absorption at reference λ (5000 Å)
    n_dicts = Vector(undef, length(atm.layers))    #vector of (species -> number density) Dicts
    for (i, layer) in enumerate(atm.layers)
        n_dicts[i] = molecular_equilibrium(MEQs, layer.temp, layer.number_density,
                                           layer.electron_number_density)

        α_cntm_vals = reverse(total_continuum_absorption(sorted_cntmνs, layer.temp,
                                                         layer.electron_number_density,
                                                         n_dicts[i], partition_funcs))
        α_cntm[i] = LinearInterpolation(cntmλs, α_cntm_vals)
        α[i, :] .= α_cntm[i].(λs)

        α5[i] = total_continuum_absorption([c_cgs/5e-5], layer.temp, layer.electron_number_density,
                                           n_dicts[i], partition_funcs)[1]

        if hydrogen_lines
            α[i, :] .+= hydrogen_line_absorption(λs, layer.temp, layer.electron_number_density,
                                                 n_dicts[i][species"H_I"],
                                                 partition_funcs[species"H_I"],
                                                 hline_stark_profiles, vmic*1e5)
        end
    end

    #put number densities in a dict of vectors, rather than a vector of dicts.
    number_densities = Dict([spec=>[n[spec] for n in n_dicts] for spec in keys(n_dicts[1])])

    #add contribution of line absorption to α
    line_absorption!(α, linelist, λs, [layer.temp for layer in atm.layers],
                     [layer.electron_number_density for layer in atm.layers], number_densities,
                     partition_funcs, vmic*1e5; α_cntm=α_cntm,
                     cutoff_threshold=line_cutoff_threshold)

    source_fn = blackbody.((l->l.temp).(atm.layers), λs')
    flux = radiative_transfer(atm, α, source_fn, α5, mu_grid)

    (flux=flux, alpha=α, number_densities=number_densities, wavelengths=λs.*1e8)
end

"""
    get_absolute_abundances(metallicity, A_X)

Calculate ``n_X/n_\\mathrm{total}`` for each element X given some specified abundances, ``A(X)``.  Use the
metallicity [``X``/H] to calculate those remaining from the solar values (except He).
"""
function get_absolute_abundances(metallicity, A_X::Dict) :: Vector{Number}
    if "H" in keys(A_X)
        throw(ArgumentError("A(H) set, but A(H) = 12 by definition. Adjust \"metallicity\" and "
                           * "\"abundances\" to implicitly set the amount of H"))
    end

    #populate dictionary of absolute abundaces
    abundances = map(0x01:Natoms) do elem
        if elem == 0x01 #hydrogen
            1.0
        elseif atomic_symbols[elem] in keys(A_X)
            10^(A_X[atomic_symbols[elem]] - 12.0)
        else
            #I'm accessing the module global solar_abundances here, but it doesn't make sense to
            #make this an optional argument because this behavior can be completely overridden by
            #specifying all abundances explicitely.
            Δ = (elem == 0x02 #= helium =#) ? 0.0 : metallicity
            10^(solar_abundances[elem] + Δ - 12.0)
        end
    end
    #now normalize so that sum(N_x/N_total) = 1
    abundances ./= sum(abundances)
    abundances
end

"""
    blackbody(T, λ)

The value of the Planck blackbody function for temperature `T` at wavelength `λ`.
"""
function blackbody(T, λ)
    h = hplanck_cgs
    c = c_cgs
    k = kboltz_cgs

    2*h*c^2/λ^5 * 1/(exp(h*c/λ/k/T) - 1)
end


"""
Common interface for all modes and key lengths
"""
function encrypt end

encrypt(plain, cipher::Cipher{CBC,T,U}; iv=rand(UInt8, 16)) where {T,U} = encrypt(transcode(UInt8, plain), cipher, iv)
encrypt(plain, cipher::Cipher{CTR,T,U}; iv=rand(UInt8, 16)) where {T,U} = encrypt(transcode(UInt8, plain), cipher, iv)
encrypt(plain, cipher::Cipher) = encrypt(transcode(UInt8, plain), cipher)

encrypt(plaintext::AbstractArray{UInt8}, cipher::Cipher{CBC,T,U}, iv) where {T,U} = CipherText(AESCBC(plaintext, iv, cipher.key, cipher.cache), iv, get_key_length(cipher), CBC)
encrypt(plaintext::AbstractArray{UInt8}, cipher::Cipher{CTR,T,U}, iv) where {T,U} = CipherText(AESCTR(plaintext, iv, cipher.key, cipher.cache), iv, get_key_length(cipher), CTR)
encrypt(plaintext::AbstractArray{UInt8}, cipher::Cipher{ECB,T,U}) where {T,U} = CipherText(AESECB(plaintext, cipher.key, cipher.cache), nothing, get_key_length(cipher), ECB)

encrypt!(ciphertext, plaintext, cipher::Cipher{CBC,T,U}; iv=rand(UInt8, 16)) where {T,U} = encrypt!(ciphertext, transcode(UInt8, plaintext), cipher, iv)
encrypt!(ciphertext, plaintext, cipher::Cipher{CTR,T,U}; iv=rand(UInt8, 16)) where {T,U} = encrypt!(ciphertext, transcode(UInt8, plaintext), cipher, iv)
encrypt!(ciphertext, plaintext, cipher::Cipher) = encrypt!(ciphertext, transcode(UInt8, plaintext), cipher)

encrypt!(ciphertext, plaintext::AbstractArray{UInt8}, cipher::Cipher{CBC,T,U}, iv) where {T,U} = CipherText(AESCBC!(ciphertext, plaintext, iv, cipher.key, cipher.cache), iv, get_key_length(cipher), CBC)
encrypt!(ciphertext, plaintext::AbstractArray{UInt8}, cipher::Cipher{CTR,T,U}, iv) where {T,U} = CipherText(AESCTR!(ciphertext, plaintext, iv, cipher.key, cipher.cache), iv, get_key_length(cipher), CTR)
encrypt!(ciphertext, plaintext::AbstractArray{UInt8}, cipher::Cipher{ECB,T,U}) where {T,U} = CipherText(AESECB!(ciphertext, plaintext, cipher.key, cipher.cache), nothing, get_key_length(cipher), ECB)


module Coulter

    __precompile__()

    using KernelDensity
    using Distributions
    using Gadfly
    using StatsBase
    using DataStructures
    import Base.-, Base.deepcopy
    using Dates: Second, DateTime

    export CoulterCounterRun, loadZ2, -, volume, diameter,
        extract_peak, extract_peak!, extract_peak_interval

    include("utils.jl")

    """
    A simplified representation of a coulter counter run
    """
    mutable struct CoulterCounterRun
        filename::String
        sample::String
        timepoint::DateTime
        reltime::Union{Second, Nothing}
        binlims::Vector{Float64}
        binvols::Vector{Float64}
        binheights::Vector{Number}
        """Whether the `binheights` is the raw count or total volume"""
        yvariable::Symbol
        data::Vector{Float64}
        livepeak::Union{Float64, Nothing}
        allpeaks::Vector{Float64}
        params::Dict{String, Any}
    end

    include("analysis.jl")

    # copy constructor
    deepcopy(a::CoulterCounterRun) = CoulterCounterRun([deepcopy(getfield(a, field)) for field in fieldnames(CoulterCounterRun)]...)

    """
    loadZ2(filename::String, sample::String)

    Loads `filename` and assigns it a sample, returns a
    `CoulterCounterRun` object
    """
    function loadZ2(filepath::String, sample::String; yvariable=:count)
        open(filepath) do s
            # split windows newlines if present
            filebody = replace(read(s, String), "\r\n"=>"\n")
            # extract start time and date from body
            datetime = match(r"^StartTime= \d*\s*(?<time>\d*:\d*:\d*)\s*(?<date>\d*\s\w{3}\s\d{4})$"m, filebody)
            timepoint = DateTime("$(datetime[:date]) $(datetime[:time])", "dd uuu yyy HH:MM:SS")

            params = Dict{String, Any}()
            params["Current"] = parse(Float64, match(r"^Cur=([+-]?[0-9]*[.]?[0-9]+)$"m, filebody)[1])
            params["Pre-Amp Gain"] = parse(Float64, match(r"^PAGn=([+-]?[0-9]*[.]?[0-9]+)$"m, filebody)[1])
            params["Volume metered"] = parse(Float64, match(r"^Vol=([+-]?[0-9]*[.]?[0-9]+)$"m, filebody)[1])

            # extract data
            matcheddata = match(r"^\[#Bindiam\]\n(?<binvols>.*?)\n^\[Binunits\].*?\[#Binheight\]\n(?<binheight>.*?)\n^\[end\]"sm, filebody)
            binlims = [volume(parse(Float64, x)) for x in split(matcheddata[:binvols], "\n ")]
            # the Coulter software takes the mean the upper and lower limits to
            # compute the volume of everything that falls in the bin
            binvols = [(binlims[i] + binlims[i+1])/2 for i in 1:length(binlims)-1]
            binheights = Int[parse(Int, x) for x in split(matcheddata[:binheight], "\n ")]
            pop!(binheights) # remove the last extraneous value

            # unbin data, i.e. the inverse of the hist function
            data = repvec(binvols, binheights)

            if yvariable == :volume
                binheights = binvols .* binheights
            end

            CoulterCounterRun(basename(filepath), sample, timepoint, nothing,
                              binlims, binvols, binheights, yvariable, data,
                              nothing, Float64[], params)
        end
    end

    loadZ2(filepath::String) = loadZ2(filepath, "N/A")

    -(a::CoulterCounterRun, b::CoulterCounterRun) = a.timepoint - b.timepoint
    -(a::CoulterCounterRun, b::DateTime) = a.timepoint - b
end