const LOCAL_MEASURE_TYPES = subtypes(MLJBase.Measure)

const LOSSFUNCTIONS_MEASURE_TYPES =
    vcat(subtypes(LossFunctions.MarginLoss),
         subtypes(LossFunctions.DistanceLoss))

const MEASURE_TYPES = vcat(LOCAL_MEASURE_TYPES, LOSSFUNCTIONS_MEASURE_TYPES)

const MeasureProxy = NamedTuple{Tuple(MEASURE_TRAITS)}

function Base.show(stream::IO, p::MeasureProxy)
    instances = "["*join(p.instances, ", ")*"]"
    print(stream, "(name = $(p.name), instances = $instances, ...)")
end

function Base.show(stream::IO, ::MIME"text/plain", p::MeasureProxy)
    printstyled(IOContext(stream, :color=> MLJBase.SHOW_COLOR),
                p.docstring, bold=false, color=:magenta)
    println(stream)
    MLJBase.fancy_nt(stream, p)
end

"""
    measures()

List all measures as named-tuples keyed on measure traits.

    measures(filters...)

List all measures compatible with the target `y`.

    measures(needle::Union{AbstractString,Regex}

List all measures with `needle` in a measure's `name` or `docstring`.

### Example

Find all classification measures supporting sample weights:

    measures(m -> m.target_scitype <: AbstractVector{<:Finite} &&
                  m.supports_weights)

Find all measures in the `rms` family:

    measures("rms")

"""
function measures(conditions...)
    all_measures = map(info, MEASURE_TYPES)
    return filter(all_measures) do measure
        all(c(measure) for c in conditions)
    end
end

function measures(needle::Union{AbstractString,Regex})
    f = m -> occursin(needle, m.name) ||
        occursin(needle, m.docstring)
    return MLJBase.measures(f)
end

measures() = measures(x->true)


using FileIO
using CSV
using DataFrames
using Plots
using Measures      # for better positioning of figures
using PDFmerger

# if you have python installed, you can use pyplot for better formatting
pyplot()
# gr()

include("processing_sim.jl")

function plot_all_forecasts(
    file_date,
    states,
    dates,
    local_case_dict;
    zoom = false,
    confidence_level = "both",
    truncation_days = 7,
)
    """
    Plots all the forecasts and saves the result as a pdf.
    """

    forecast_origin = string(Dates.Date(file_date) - Dates.Day(truncation_days))

    # read in the sim results and kept TP's
    sim_all_states = CSV.read(
        "results/UoA_forecast_output/" * forecast_origin * "/UoA_samples_" * forecast_origin * ".csv",
        DataFrame,
    )

    # read in the case data
    case_dates = collect(dates)

    # initialise the plot
    fig = plot(
        layout = (4, 2),
        dpi = 200,
        size = (700, 1000),
        margin = 3mm,
        rightmargin = 4mm,
        framestyle = :box,
        legend = false
    )

    for (sp, state) in enumerate(states)

        # get relevant columns of the dataframe and merge to matrix
        D = Matrix(sim_all_states[sim_all_states.state .== state, 3:end - 1])
        # if anything in D is missing, it means the sims didn't work so 0 out those entries
        D[ismissing.(D)] .= 0
        D = Matrix{Int}(D)

        onset_dates = sim_all_states[sim_all_states.state .== state, "onset date"]
        # onset_dates = Date.(onset_dates, dateformat"m yyyy")
        df_D_summary = summarise_forecast_for_plotting(D)

        local_cases = local_case_dict[state][dates .>= onset_dates[1]]
        # boolean for correct case dates
        case_dates_ind = [d ∈ onset_dates ? true : false for d in case_dates]

        onset_dates_lims = Dates.value.([onset_dates[1], onset_dates[end]])

        if confidence_level == "50" || confidence_level == "both"
            plot!(
                fig,
                title = state,
                subplot = sp,
                onset_dates,
                df_D_summary[!, "median"],
                linecolor = 1,
                linewidth = 1,
                color = 1,
                fillalpha = 0.2,
                label = false,
            )
            plot!(
                fig,
                subplot = sp,
                onset_dates,
                [df_D_summary[!, "median"] df_D_summary[!, "median"]],
                fillrange = [df_D_summary[!, "lower"] df_D_summary[!, "upper"]],
                fillalpha = 0.4,
                c = 1,
                label = ["50%" nothing]
            )

        end

        if confidence_level == "95" || confidence_level == "both"
            plot!(
                fig,
                title = state,
                subplot = sp,
                onset_dates,
                df_D_summary[!, "median"],
                linecolor = 1,
                linewidth = 1,
                color = 1,
                fillalpha = 0.2,
                label = false,
            )

            plot!(
                fig,
                subplot = sp,
                onset_dates,
                [df_D_summary[!, "median"] df_D_summary[!, "median"]],
                fillrange = [df_D_summary[!, "bottom"] df_D_summary[!, "top"]],
                fillalpha = 0.2,
                c = 1,
                label = ["95%" nothing]
            )
        end

        plot!(
            fig,
            subplot = sp,
            case_dates[case_dates_ind],
            local_cases,
            lc = "black",
            label = false,
        )

        # plot the last date used for conditioning the simulations
        vline!(
            fig,
            subplot = sp,
            [Dates.Date(file_date) - Dates.Day(truncation_days)],
            lc = :gray,
            ls = :dash,
            label = false,
        )

        # adjust limits to make plots a little nicer to look at
        xlims!(fig, subplot = sp, onset_dates_lims...)
        if zoom
            ylims!(fig, subplot = sp, 0, 1.25 * maximum(local_cases))
        end

    end

    dir_name = joinpath("figs", "case_forecasts", forecast_origin)
    if !ispath(dir_name)
        mkpath(dir_name)
    end

    if zoom
        savefig(
            fig,
            dir_name *
            "/UoA_forecast_" *
            forecast_origin *
            "_" *
            "zoomed_" *
            confidence_level *
            "_intervals.pdf"
        )
    else
        savefig(
            fig,
            dir_name *
            "/UoA_forecast_" *
            forecast_origin *
            "_" *
            confidence_level *
            "_intervals.pdf"
        )
    end

    return nothing

end

function plot_all_forecast_intervals(
    file_date,
    states,
    dates,
    local_case_dict;
    truncation_days = 7,
)
    """
    Simple wrapper function to plot the forecasts with various zoom and confidence levels.
    """

    plot_all_forecasts(
        file_date,
        states,
        dates,
        local_case_dict,
        confidence_level = "both",
    )
    plot_all_forecasts(
        file_date,
        states,
        dates,
        local_case_dict,
        zoom = true,
        confidence_level = "both",
    )
    plot_all_forecasts(
        file_date,
        states,
        dates,
        local_case_dict,
        confidence_level = "50",
    )
    plot_all_forecasts(
        file_date,
        states,
        dates,
        local_case_dict,
        confidence_level = "95",
    )

    forecast_origin = string(Dates.Date(file_date) - Dates.Day(truncation_days))
    dir_name = "figs/case_forecasts/" * forecast_origin * "/"
    file_name_tmp = "UoA_forecast_"

    pdf_filenames = [
        dir_name * file_name_tmp * forecast_origin * "_zoomed_both_intervals.pdf",
        dir_name * file_name_tmp * forecast_origin * "_both_intervals.pdf",
        dir_name * file_name_tmp * forecast_origin * "_50_intervals.pdf",
    ]
    # merge the pdfs and delete the files
    merge_pdfs(
        pdf_filenames,
        dir_name * file_name_tmp * forecast_origin * ".pdf",
        cleanup=true,
    )

    return nothing

end


# Project on adaptive evolvability as first described in cce/Adaptive_evolvability9_6_21.docx.

# For each of number_sources starting phenotypes choose number_target random phenotypes within
#     distance_limit of the source phenotype.
# For each starting phenotype and each target phenotype do number_runs epochal evolutions with
#     a maximum of step_limit steps and fitness defined by hamming distance from target.
# Record the number of successes and the averge number of steps for successful runs
# max_steps is the maximum number of steps on the evolution to find the starting chromosome for a run
function compare_evolvabilities( p::Parameters, source_list::GoalList, number_targets::Int64,
    distance_limit::Float64, step_limit::Int64, number_runs::Int64;
    csvfile::String="", evo_df_file::String="", evo_df_column::Symbol=:d_evolvability,
    iteration_limit::Int64=200, max_steps::Int64=200000, distance_function::Function=hamming_distance )
  evo_df = DataFrame()
  if length(evo_df_file) > 0
    try
      evo_df = read_dataframe(evo_df_file)
    catch
      error("reading dataframe from ",evo_df_file," failed")
    end
    @assert size(evo_df)[1] == length(source_list)  # Make sure goal lists are compatible
  end
  df = DataFrame()
  df.source = Goal[]
  #df.numinputs = Int64[]
  #df.numoutputs = Int64[]
  #df.numgates= Int64[]
  #df.numlevelsback = Int64[]
  df.numtargets = Int64[]
  df.dist_limit = Float64[]
  df.step_limit = Int64[]
  df.number_runs = Int64[]
  df.avg_num_successes = Float64[]
  df.avg_steps = Float64[]
  println("size(df): ",size(df))
  # list of parameter values that go into the dataframe
  #df_params_list = Any[p.numinputs, p.numoutputs, p.numinteriors, p.numlevelsback, number_targets, distance_limit, step_limit, number_runs]
  #println("df_params_list: ",df_params_list)
  #for src in source_list
  df_rows = pmap( src->run_compare_evolvabilities( p, src, number_targets, distance_limit, step_limit, number_runs,
        iteration_limit=iteration_limit, max_steps=max_steps, distance_function=distance_function ), source_list )
  for df_row in df_rows
    push!(df,df_row)
  end
  df.evolvability = evo_df[:,evo_df_column]
  if length(csvfile) > 0
    open( csvfile, "w" ) do f
      hostname = chomp(open("/etc/hostname") do f read(f,String) end)
      println(f,"# date and time: ",Dates.now())
      println(f,"# host: ",hostname," with ",nprocs()-1,"  processes: " )
      println(f,"# funcs: ", Main.CGP.default_funcs(p.numinputs))
      print_parameters(f,p,comment=true)
      CSV.write( f, df, append=true, writeheader=true )
    end
  end
  df
end

# Helper function for the above function for a single src goal.
function run_compare_evolvabilities( p::Parameters, src::Goal, number_targets::Int64,
  distance_limit::Float64, step_limit::Int64, number_runs::Int64;
  iteration_limit::Int64=200, max_steps::Int64=200000, distance_function::Function=hamming_distance )
  println("src: ",src)
  df_params_list =
    Any[
      #p.numinputs,
      #p.numoutputs,
      #p.numinteriors,
      #p.numlevelsback,
      number_targets,
      distance_limit,
      step_limit,
      number_runs
    ]
  iteration = 0
  target_list = Goal[]
  while iteration < iteration_limit && length(target_list) < number_targets
    target = randgoal( p.numinputs, p.numoutputs )
    #hd = hamming_distance( src, target, p.numinputs )
    hd = distance_function( src, target, p.numinputs )
    #println("target: ",target,"  hd: ",hd)
    if hd <= distance_limit && target != src
      push!(target_list,target)
      println("target: ",target,"  hd: ",hd,"  iteration: ",iteration)
    end
    iteration += 1
  end
  if iteration == iteration_limit
    error("target list iteration exceeded iteration limit of ",iteration_limit," for source src: ",src)
  end
  println("target_list: ",target_list)
  successes = 0
  sum_steps = 0
  for target in target_list
    ch = random_chromosome( p )
    (src_ch,steps) = neutral_evolution(ch, src, max_steps )
    while iteration < iteration_limit && steps == max_steps
      (src_ch,steps) = neutral_evolution(ch, src, max_steps )
      iteration += 1
    end
    if iteration == iteration_limit
      error("iteration to src exceeded iteration limit of ",iteration_limit," for source src: ",src," and target: ",target)
    end
    @assert  steps < max_steps
    for run = 1:number_runs
      (target_ch,steps) = neutral_evolution(src_ch, target, step_limit )
      successes += steps < step_limit ? 1 : 0
      sum_steps += steps < step_limit ? steps : 0
    end
  end
  avg_successes = successes/length(target_list)/number_runs
  avg_steps = sum_steps/successes
  df_row = vcat([src],df_params_list,[avg_successes,avg_steps])
  return df_row
end

# Doesn't make these assignments global
function test_inputs()
  p = Parameters(3,1,7,4)
   #source_list=randgoallist( 10, p.numinputs, p.numoutputs )
   source_list=map(x->[MyInt(x)],collect(0:255))
   number_targets=3
   distance_limit=0.3
   step_limit=200
   number_runs=10
   evo_df_file="../data/9_3_21/robust_evo_by_walks9_3_21B.csv"
end


# Superflous code
# Contains code ported from Prof. Hicksn's Mapping_Mod.f90

@doc """
### calcJacobian

!!!!!   NOT NEEDED   !!!!
Calculates the metric Jacobian value based on the B-spline explicit mapping.

**Arguments**

*  `map`  : Object of Mapping type
*  `Jac`  : Array of mesh Jacobians in j,k,m format

"""->

function calcJacobian(map, Jac)

  # Check if map.jkmmax = size(Jac)
  @assert map.jkmmax[1] == size(Jac,1)
  @assert map.jkmmax[2] == size(Jac,2)
  @assert map.jkmmax[3] == size(Jac,3)

  # Allocate and initialize arrays
  scal = zeros(AbstractFloat, 3)
  xi = zeros(AbstractFloat, 3)
  dx = zeros(AbstractFloat, 3, 3)

  scal[:] = 1 ./ (map.jkmmax[:] - 1)

  for m = 1:map.jkmmax[3]
    for k = 1:map.jkmmax[2]
      for j = 1:map.jkmmax[1]
        xi[:] = map.jkmmax[j,k,m,:]
        calcdXdxi(map, xi, jderiv, dx[:,1])
        calcdXdxi(map, xi, jderiv, dx[:,2])
        calcdXdxi(map, xi, jderiv, dx[:,3])
        dx[:,1] *= scal[1]
        dx[:,2] *= scal[2]
        dx[:,3] *= scal[3]
        Jac[j,k,m] = dx[1,1]*dx[2,2]*dx[3,3] + dx[1,2]*dx[2,3]*dx[3,1] +
                     dx[1,3]*dx[2,1]*dx[3,2] - dx[1,1]*dx[2,3]*dx[3,2] -
                     dx[1,2]*dx[2,1]*dx[3,3] - dx[1,3]*dx[2,2]*dx[3,1]
        Jac[j,k,m] = 1/Jac[j,k,m]
      end  # End for j = 1:map.jkmmax[1]
    end  # End for k = 1:map.jkmmax[2]
  end  # End for m = 1:map.jkmmax[3]

  return nothing
end

@doc """
### resizeMapping

!!!! NOT NEEDED FOR FFD  !!!!

Uses edge spacing parameters to refine/coarsen in each parameter direction.
This might be useful for refinement studies where we want to insert or delete
points in a continous way, or for grid sequencing/multigrid applications.

**Arguments**

*  `map`  : Object of Mapping type
*  `jkmmax` : dersired number of (xi, eta, zeta) parameters

"""->

function resizeMapping(map, jkmmax)

  for i = 1:3
    @assert jkmmax[i] > 0
  end

  # Check if resizing is necessary
  if map.jkmmax[1] != jkmmax[1] || map.jkmmax[2] != jkmmax[2] || map.jkmmax[3]
     != jkmmax[3]
    map.xi = zeros(AbstractFloat, jkmmax[1], jkmmax[2], jkmmax[3], 3)
  end

  uvw = zeros(AbstractFloat, 3)

  # Loop over all nodes, calculate each new (xi, eta, zeta)
  for m = 1:map.jkmmax[3]
    uvw[3] = (m-1)/(map.jkmmax[3] - 1)
    for k = 1:map.jkmmax[2]
      uvw[2] = (k-1)/(map.jkmmax[2]-1)
      for j = 1:map.jkmmax[1]

        uvw[1] = (j-1)/(map.jkmmax[1] - 1)
        calcXi(map, uvw, map.xi[j,k,m,:])
        # need to set lower/upper value explicitly to avoid round-off errors
        if j == 1
          map.xi[j,k,m,1] = 0.0
        end
        if k == 1
          map.xi[j,k,m,2] = 0.0
        end
        if m == 1
          map.xi[j,k,m,3] = 0.0
        end
        if j == map.jkmmax[1]
          map.xi[j,k,m,1] = 1.0
        end
        if k == map.jkmmax[2]
          map.xi[j,k,m,2] = 1.0
        end
        if m == map.jkmmax[3]
          map.xi[j,k,m,3] = 1.0
        end

      end  # End for j = 1:map.jkmmax[1]
    end  # End for k = 1:map.jkmmax[2]
  end  # End for m = 1:map.jkmmax[3]

  return nothing
end  # End function resizeMapping(map, jkmmax)

@doc """
### calcXi

!!!! NOT NEEDED FOR FFD !!!

Calculates the intermediate parameters, (xi,eta,zeta), using the edge spacing
control functions defined by map%edgpar(1:2,:,:)

**Arguments**

*  `map`  : Object of Mapping type
*  `uvw`  : j,k,m indices scaled to fit in [0,1],  i.e. uniform spacing
*  `xi`   : Parametric coordinates (length = 3)

"""->

function calcXi(map, uvw, xi)

  # Allocate and initialize arrays
  # xi = zeros(AbstractFloat, 3)
  xiedg = zeros(AbstractFloat, 4, 3)
  F = zeros(AbstractFloat, 3)
  dF = zeros(AbstractFloat, 3, 3)
  dxi = zeros(AbstractFloat, 3)
  psi = zeros(AbstractFloat, 3)
  dpsi = zeros(AbstractFloat, 3)

  for di = 1:3
    it1 = mod(di,3) + 1
    it2 = mod(di+1,3) + 1
    for edg 1:4
      A = map.edge_param[1, edg, di]
      b = map.edge_param[2, edg, di]
      fac = uvw[di] - 0.5
      if b == 0.0
        # Use taylor series to determine u
        fac = 1.0 + 2.0*fac
      else
        fac = 1.0 + tanh(b*fac)/tanh(0.5*b)
      end  # End if b == 0.0
      xiedg[edg, di] = fac/( 2*A + (1 - A)*fac )
    end  # End for edg 1:4
  end  # End for di = 1:3

  # calculate the intermediate params by find the intersection of 3 trilinear
  # surface (via Newton's method)

  # Get initial estimate
  for di = 1:3
    xi[di] = 0.0
    for edg = 1:4
      xi[di] = 0.25*xiedg[edg,di]
    end
  end

  # loop until sufficiently converged

  for n = 1:100
    # Get the RHS and LHS for Newton's method
    F[:] = -xi[:]
    fill!(dF, 0.0)
    for di = 1:3
      it1 = mod(di,3) + 1
      it2 = mod(di+1,3) + 1
      dF[di,di] = 1.0
      for edg = 1:4
        if mod(edg,2) == 1
          psi[it1] = 1.0 - xi[it1]
          dpsi[it1] = 1.0
        else
          psi[it1] = xi[it1]
          dpsi[it1] = 1.0
        end  # End if mod(edg,2) == 1

        if mod((edg+1)/2,2) == 1
          psi[it2] = 1.0 - xi[it2]
          dpsi[it2] = -1.0
        else
          psi[it2] = xi[it2]
          dpsi[it2] = 1.0
        end  # End if mod((edg+1)/2,2) == 1

        F[di] += psi[it1]*psi[it2]*xiedg[edg,di]
        dF[di,it1] -= -dpsi[it1]*psi[it2]*xiedg[edg,di]
        dF[di,it2] -= -psi[it1]*dpsi[it2]*xiedg[edg,di]
      end  # End for edg = 1:4
    end # End for di = 1:3

    # Solve for the update
    # TODO: Solve for the update

    # update parameter and exit loop
    xi[:] += F[:]
    if norm(F,2) < 1e-13
      return nothing
    end

  end  # End for n = 1:100

  return nothing
end

# parameter functions
@doc """
### calcChordLengthParam

!!! NOT NEEDED FOR FFD !!!

Calculates the mapping parameter values based on the chord lengths in physical
xyz space. It also calculates the edge knot vectors.

**Inputs**

*  `map` : Object of Mapping type
*  `xyz` : Mesh in physical space on which the chord length parameterization is
           based on

**Outputs**

*  None

""" ->

function calcChordLengthParam(map, xyz)

  # Intermediate variables
  jkm = zeros(Int, 3)
  vec = zeros(3)

  # Loop over each direction, Find the parameter values for a curve of constant
  # it1 & it2
  for di = 1:3
    it1 = mod(di,3) + 1
    it2 = mod(di+1,3) + 1

    for jit1 = 1:jkmmax[it1]
      jkm[it1] = jit1
      for jit2 = 1:jkmmax[it2]
        jkm[it2] = jit2
        jkm[di] = 1
        map.xi[jkm[1], jkm[2], jkm[3], di] = 0.0
        for jdi = 1:(map.jkmmax[di]-1)
          jkm[di] = jdi
          j = jkm[1]
          k = jkm[2]
          m = jkm[3]
          jkm[di] = jdi + 1
          jp = jkm[1]
          kp = jkm[2]
          mp = jkm[3]

          # Get the arc-length element
          vec[:] = xyz[jp, kp, mp, :] - xyz[j, k, m, :]
          ds = norm(vec,2)
          map.xi[jp,kp,mp,di] = map.xi[j,k,m,di] + ds
        end  # End for jdi = 1:(map.jkmmax[di]-1)

        # Normalize the arc-length parameters
        jkm[di] = map.jkmmax[di]
        j = jkm[1]
        k = jkm[2]
        m = jkm[3]
        fac = 1 / map.xi[j,k,m,di]
        map.xi[j,k,m,di] = 1.0
        for jdi = 2:(map.jkmmax[di]-1)
          jkm[di] = jdi
          j = jkm[1]
          k = jkm[2]
          m = jkm[3]
          map.xi[j,k,m,di] = fac*map.xi[j,k,m,di]
        end  # End for jdi = 2:(map.jkmmax[di]-1)
      end  # End for jit2 = 1:jkmmax[it2]
    end  # End for jit1 = 1:jkmmax[it1]
  end  # end for di = 1:3

  # At this point, we have arc length parameters, so arc-length-based edge knots
  # can be found

  for di = 1:3
    it1 = mod(di,3) + 1
    it2 = mod(di+1,3) + 1
    order = map.order[di]
    nctl = map.nctl[di]

    for edg = 1:4
      if mod(edg,2) == 1
        jkm[it1] = 1
      else
        jkm[it1] = map.jkmmax[it1]
      end  # End if mod(edg,2) == 1

      if mod((edg+1)/2,2) == 1
        jkm[it2] = 1
      else
        jkm[it2] = map.jkmmax[it2]
      end  # End if mod((edg+1)/2,2) == 1

      # Set the edge knot vector's end knots
      map.edge_knot[1:order, edg, di] = 0.0
      map.edge_knot[nctl+1:nctl+order, edg, di] = 1.0

      # Find point intervals per knot intervals
      ppk = (map.jkmmax[di]-1) / (nctl+1-order)

      for i = order+1:nctl
        knotloc = ppk*(i-order)  # knot location
        for jdi = 1:(map.jkmmax[di]-1)
          jkm[di] = jdi

          if knotloc >= jdi-1 && knotloc < jdi
            jkm[di] = jdi
            j = jkm[1]
            k = jkm[2]
            m = jkm[3]
            jkm[di] = jdi + 1
            jp = jkm[1]
            kp = jkm[2]
            mp = jkm[3]
            map.edge_knot[i,edg,di] = (map.xi[jp,kp,mp,di] - map.xi[j,k,m,di])*
                                      (knotloc - jdi + 1) + map.xi[j,k,m,di]
          end  # End if knotloc >= jdi-1 && knotloc < jdi
        end # End for jdi = = 1:(map.jkmmax[di]-1)
      end  # End for i = order+1:nctl
    end  # End for edg = 1:4
  end # End for di = 1:3

  return nothing
end  # end function calcChordLengthParam

@doc """
### calcEdgeSpacingParam

!!!  NOT NEEDED  !!!
Uses the xi, eta, zeta values to calculate the parameters for hyperbolic tangent
spacing function. These parameters are

A = sqrt(sp2/sp1)
b, where, sinh(b) - b/((N-1) \* sp1 \* sp2) = 0

and sp1 and sp2 are the (xi, eta, zeta) spacing at the ends of the edges.
A is stored in `map.edge_param[1,:,:]`, and b in `map.edge_param[2, :, :]`

**Arguments**

*  `map` : Object of Mapping type

""" ->

function calcEdgeSpacingParam(map)

  # Initialize working array
  jkm = zeros(Int,3)

  for di = 1:3
    it1 = mod(di,3) + 1
    it2 = mod(di+1,3) + 1

    for edg = 1:4
      if mod(edg,2) == 1
        jkm[it1] = 1
      else
        jkm[it1] = map.jkmmax[it1]
      end

      if mod((edg+1)/2,2) == 1
        jkm[it2] = 1
      else
        jkm[it2] = map.jkmmax[it2]
      end

      # Find the spacing at the ends of the edge edg
      jkm[di] = 1
      j = jkm[1]
      k = jkm[2]
      m = jkm[3]
      jkm[di] = 2
      jp = jkm[1]
      kp = jkm[2]
      mp = jkm[3]
      dx2 = map.xi[jp,kp,mp,di] - map.xi[j,k,m,di]

      ### #if 0 satement is supposed to come here
      ###
      ###

      # Set parameter A
      map.edge_param[1, edg,di] = sqrt(dx2/dx1)

      # Set parameter b
      Nm1 = map.jkmmax[di] - 1
      x2 = (map.jkmmax[di] - 1)*sqrt(dx1*dx2)
      if x2 > 1
        # 1/(N-1) is smaller than sqrt(dx1*dx2), so b will be zero
        map.edge_param[2, edg, di] = 0.0
      else
        # need to find b using Newton's method
        x1 = small
        x2 = sqrt(6*(1/x2 - 1))
        map.edge_param[2,edg,di] =  bfuncNewtonSolve(x2, Nm1, dx1, dx2)
      end  # End if x2 > 1
    end    # End for edg = 1:4
  end      # End for di = 1:3

  return nothing
end  # End function calcEdgeSpacingParam(map)

function bfunc(x, Nm1, dx1, dx2)

  tmp = 1 / (Nm1*sqrt(dx1*dx2))
  f = sinh(x) - x*tmp
  df = cosh(x) - tmp

  return f, df
end  # end function bfunc

function bfuncNewtonSolve(guess, Nm1, dx1, dx2)

  b = guess
  b_next = zero(AbstractFloat)

  for i = 1:50
    f, df = bfunc(b, Nm1, dx1, dx2)
    b_next = b - f/df
    if norm(b_next-b, 2) < 1e-12
      break
    else
      b = b_next
    end
  end # End for i = 1:50

  return b_next
end

# Grid.jl

####  -------------------

#  NOT NEEDED !!!!!!!!!!

#### ---------------------

@doc """
### calcGrid

Calculates the mesh coordinates based on the element node values in the map,
hermite basis functions and the refinement level. The basis function values of
map should be known

**Arguments**

*  `map` : Object of Mapping type
*  `xyz` : Array of mesh node coordinates in j,k,m,1:3 format

"""->

function calcGrid(map, xyz)

  @assert size(map.jkmmax) == size(xyz)

  for m = 1:map.jkmmax[3]
    for k = 1:map.jkmmax[2]
      for j = 1:map.jkmmax[1]
        xi =sview(map.ki, j, k, m, :)
        dXdxi(map, xi, jderiv, xyz[j,k,m,:])
      end  # End for j = 1:map.jkmmax[1]
    end  # End for k = 1:map.jkmmax[2]
  end  # End for m = 1:map.jkmmax[3]

  return nothing
end

@doc """
### contractWithdGdB

Used to contract the matrix dG/dB, the derivative of the mesh node coordinates
with respect to the mapping control points, with a given array, here labelled
dJdGrid and stored in (j,k,m,:) format. This can be used to find
dJ/dB = (dJ/dG)(dG/dB) for some objective J. In practice, the objective used is
actually

             J_practice = J_obj + psi^T dot R

where J_obj is the actual objective, psi are the flow adjoints, and R is the
flow residual.

**Arguments**

*  `map`     : Object of mapping type
*  `dJdGrid` : The derivative of the objective w.r.t the grid coordintes

"""->

function contractWithdGdB(map, dJdGrid)

  @assert map.numnodes[1] == size(dJdGrid, 1)
  @assert map.numnodes[2] == size(dJdGrid, 2)
  @assert map.numnodes[3] == size(dJdGrid, 3)

  # Allocate and initialize arrays
  basis = zeros(AbstractFloat, maximum(map.order), 3)
  fill!(map.work, 0.0)

  # Loop over the nodes of the mapping
  for m = 1:map.numnodes[3]
    for k = 1:map.numnodes[2]
      for j = 1:map.numnodes[1]

        # Store the parameter values and dJdGrid
        xi[:] = map.xi[j,k,m,:]
        dJdG =sview(dJdGrid, j, k, m, 1:3)

        # Evaluate the knot vectors and basis values
        # TODO: Compute the knot vector for subsequent operations
        for di = 1:3
          span[di] = findSpan(xi[di], map, di)
          basisFunctions(map, basis[:,di], di, xi[di], span[di])
        end  # End for di = 1:3

        for r = 1:map.order[3]
          rs = r + span[3] - map.order[3]
          for q = 1:map.order[2]
            qs = q + span[2] - map.order[2]
            for p = 1:map.order[1]

              ps = p + span[1] - map.order[1]
              coeff = basis[p,1]*basis[q,2]*basis[r,3]
              map.work[ps, qs, rs, 1:3] += coeff*dJdG[:]

            end # End for p = 1:map.order[1]
          end   # End for q = 1:map.order[2]
        end     # End for r = 1:map.order[3]

      end  # End for j = 1:map.numnodes[1]
    end    # End for k = 1:map.numnodes[2]
  end      # End for m = 1:map.numnodes[3]

  return nothing
end  # End function contractWithdGdB(map, dJdGrid)

@doc """
### fitGrid

Finds the control point values such that, when evaluated at in computational
space, the mapping produces the mesh xyz as close as possible. Parameter
correction was added by L. Olague and J. Hicken and ported to the language
`Julia`

**Arguments**

*  `map` : Object of Mapping type
*  `xyz` : Mesh to be approximated. On exit, htis is over written with the new
           mesh produced by the map
*  `L2_error` : The L2 error squared in fitting the grid is added to this
                variable at the end of routine
*  `max_error`: The values in this array are updated if the infinity norm in the
                fit of this grid is larger

"""

function fitGrid(map, xyz, L2_error, max_error)

  # Allocate and initialize arrays

  return nothing
end


using CUDA
using CUDNN
if VERSION >= v"0.5-"
    using Base.Test
else
    using BaseTestNext
    const Test = BaseTestNext
end

@testset "activation" for i = 1:5
    x = CuArray(Float32, 10, 5, 3, 2)
    desc = ActivationDesc(CUDNN_ACTIVATION_SIGMOID)
    y = similar(x)
    CUDNN.activation!(desc, x, y)
end

@testset "convolution" for i = 1:5
    x = CuArray(Float32, 10, 5, 3, 2)
    desc = ConvolutionDesc()
    y = similar(x)
    CUDNN.activation!(desc, x, y)
end


###
### parsing
###

const OPTIONS = Parsers.Options(
    sentinel=missing,
    quoted=true,
    openquotechar='\'',
    closequotechar='\'',
    delim=',',
)
# Change delimiter as way to handle end-of-line comments.
const EOL_OPTIONS = Parsers.Options(
    sentinel=missing,
    quoted=true,
    openquotechar='\'',
    closequotechar='\'',
    delim='/',
)

getbytes(source::Vector{UInt8}) = source, 1, length(source)
getbytes(source::IOBuffer) = source.data, source.ptr, source.size
getbytes(source) = getbytes(read(source))

"""
    parse_network(source) -> Network

Read a PSS/E-format `.raw` Power Flow Data file and return a [`Network`](@ref) object.
"""
function parse_network(source)
    bytes, pos, len = getbytes(source)

    caseid, pos = parse_caseid(bytes, pos, len, OPTIONS)
    @debug "caseid" pos

    # Skip the 2 lines of comments
    # TODO: confirm it is always only and exactly 2 lines of comments
    pos = next_line(bytes, pos, len)
    pos = next_line(bytes, pos, len)
    @debug "comments" pos

    nrows = count_nrow(bytes, pos, len, OPTIONS)
    @debug "buses" nrows pos
    buses, pos = parse_records!(Buses(nrows), bytes, pos, len, OPTIONS, EOL_OPTIONS)

    nrows = count_nrow(bytes, pos, len, OPTIONS)
    @debug "loads" nrows pos
    loads, pos = parse_records!(Loads(nrows), bytes, pos, len, OPTIONS, EOL_OPTIONS)

    nrows = count_nrow(bytes, pos, len, OPTIONS)
    @debug "gens" nrows pos
    gens, pos = parse_records!(Generators(nrows), bytes, pos, len, OPTIONS, EOL_OPTIONS)

    nrows = count_nrow(bytes, pos, len, OPTIONS)
    @debug "branches" nrows pos
    branches, pos = parse_records!(Branches(nrows), bytes, pos, len, OPTIONS, EOL_OPTIONS)

    return Network(caseid, buses, loads, gens, branches)
end

function parse_caseid(bytes, pos, len, options)
    ic, pos, code = parse_value(Int, bytes, pos, len, options)
    @debug codes(code) pos newline=newline(code)

    # Support files that have first row like:
    # 0,   100.00          / PSS/E-30.3    WED, SEP 15 2021  21:04
    # and those with first row like:
    # 0,100.0,30 / PSS(tm)E-30 RAW created      Wed, Sep 15 2021 21:04
    # TODO: avoid needing to extract `sbase` value from String?
    sbase = Parsers.tryparse(Float64, bytes, options, pos, len)
    if sbase === nothing
        str, pos, code = parse_value(String, bytes, pos, len, options)
        @debug codes(code) pos newline=newline(code)
        sbase = parse(Float64, first(split(str, '/')))
    end

    # if delimiter after `sbase` value, then won't have reached end of line.
    if !newline(code)
        pos = next_line(bytes, pos, len)
    end
    @debug codes(code) pos newline=newline(code)

    return CaseID(ic, sbase), pos
end

function parse_records!(rec::R, bytes, pos, len, options, eol_options)::Tuple{R, Int} where {R <: Records}
    nrows = length(getfield(rec, 1))
    nrows == 0 && return rec, pos
    for row in 1:nrows
        pos, code = parse_row!(rec, row, bytes, pos, len, options, eol_options)

        # Because we're working around end-of-line comments,
        # rows with comments won't have hit the newline character yet
        if !newline(code)
            pos = next_line(bytes, pos, len)
        end
    end

    # Data input is terminated by specifying a bus number of zero.
    # @assert peekbyte(bytes, pos) == UInt8('0')
    if !(eof(bytes, pos, len) || peekbyte(bytes, pos) == UInt8('0'))
        @warn "Not at end of $(typeof(rec)) records"
    end
    pos = next_line(bytes, pos, len)
    return rec, pos
end

function count_nrow(buf, pos, len, options)
    nlines = 0
    if eof(buf, pos, len) || peekbyte(buf, pos) == UInt8('0')
        return nlines
    end
    while true
        res = xparse(String, buf, pos, len, options)
        pos += res.tlen
        if newline(res.code) || eof(res.code)
            nlines += 1
            if eof(buf, pos, len) || peekbyte(buf, pos) == UInt8('0')
                break
            end
        end
    end
    return nlines
end

# Taken from `Parsers.checkcmtemptylines`
# TODO: move to Parsers.jl?
function next_line(bytes, pos, len)
    eof(bytes, pos, len) && return pos
    b = peekbyte(bytes, pos)
    while b !== UInt8('\n') && b !== UInt8('\r')
        pos += 1
        eof(bytes, pos, len) && break
        b = peekbyte(bytes, pos)
    end
    # Move forward to be past the `\r` or `\n` byte.
    pos += 1
    # if line ends `\r\n`, then we're at `\n`and need to move forward again.
    if b === UInt8('\r') && !eof(bytes, pos, len) && peekbyte(bytes, pos) === UInt8('\n')
        pos += 1
    end
    return pos
end

function parse_row!(rec::Records, row::Int, bytes, pos, len, options, eol_options)
    ncols = nfields(rec)
    local code::Parsers.ReturnCode
    for col in 1:ncols
        eltyp = eltype(fieldtype(typeof(rec), col))
        opts = ifelse(col == ncols, eol_options, options)
        val, pos, code = parse_value(eltyp, bytes, pos, len, opts)
        @inbounds getfield(rec, col)[row] = val

        @debug codes(code) row col pos newline=newline(code)
    end
    return pos, code
end

function parse_value(T, bytes, pos, len, options)
    res = xparse(T, bytes, pos, len, options)

    invalid(res.code) && @warn codes(res.code) pos

    pos += res.tlen
    code = res.code

    val = if T === String
        Parsers.getstring(bytes, res.val, options.e)
    else
        res.val
    end
    return val, pos, code
end


function generate_get_expr(domain::Domain, state::State, term::Const,
                           varmap=Dict{Var,Any}(), state_var=:state)
    return :($state_var.$(term.name))
end

function generate_get_expr(domain::Domain, state::State, term::Var,
                           varmap=Dict{Var,Any}(), state_var=:state)
    return :($(varmap[term]))
end

function generate_get_expr(domain::Domain, state::State, term::Compound,
                           varmap=Dict{Var,Any}(), state_var=:state)
    indices = generate_fluent_ids(domain, state, term,
                                  get_fluent(domain, term.name), varmap, state_var)
    return :($state_var.$(term.name)[$(indices...)])
end

function generate_set_expr(domain::Domain, state::State, term::Const,
                           val, varmap=Dict{Var,Any}(), state_var=:state)
    if domain isa AbstractedDomain && domain.interpreter.autowiden
        prev_val = generate_get_expr(domain, state, term, varmap, :prev_state)
        return :(setfield!($state_var, $(QuoteNode(term.name)), widen($prev_val, $val)))
    else
        return :(setfield!($state_var, $(QuoteNode(term.name)), $val))
    end
end

function generate_set_expr(domain::Domain, state::State, term::Compound,
                           val, varmap=Dict{Var,Any}(), state_var=:state)
    indices = generate_fluent_ids(domain, state, term,
                                  get_fluent(domain, term.name), varmap, state_var)
    if domain isa AbstractedDomain && domain.interpreter.autowiden
        prev_val = generate_get_expr(domain, state, term, varmap, :prev_state)
        return :($state_var.$(term.name)[$(indices...)] = widen($prev_val, $val))
    else
        return :($state_var.$(term.name)[$(indices...)] = $val)
    end
end

function generate_fluent_ids(domain::Domain, state::State, term::Term,
                             sig::Signature, varmap=Dict{Var,Any}(),
                             state_var=:state)
    if get_requirements(domain)[:typing]
        ids = map(zip(term.args, sig.argtypes)) do (arg, type)
            if arg isa Var
                :(objectindex($state_var, $(QuoteNode(type)), $(varmap[arg])))
            else
                objects = sort(get_objects(domain, state, type), by=x->x.name)
                findfirst(isequal(arg), objects)
            end
        end
    else
        objects = sort(collect(get_objects(state)), by=x->x.name)
        ids = map(term.args) do arg
            if arg isa Var
                :(objectindex($state_var, $(varmap[arg])))
            else
                findfirst(isequal(arg), objects)
            end
        end
    end
    return ids
end

function generate_fluent_dims(domain::Domain, state::State, pred::Signature)
    if get_requirements(domain)[:typing]
        dims = [length(get_objects(domain, state, ty)) for ty in pred.argtypes]
    else
        dims = fill(length(get_objects(state)), length(pred.args))
    end
    return dims
end


# Unit tests for judiVector
# Philipp Witte (pwitte.slim@gmail.com)
# May 2018
#
# Mathias Louboutin, mlouboutin3@gatech.edu
# Updated July 2020

datapath = joinpath(dirname(pathof(JUDI)))*"/../data/"

# number of sources/receivers
nrec = 120
ns = 251
ftol = 1e-6

################################################# test constructors ####################################################

@testset "judiVector Unit Tests with $(nsrc) sources" for nsrc=[1, 2]
    @timeit TIMEROUTPUT "judiVector (nsrc=$(nsrc))" begin
        # set up judiVector from array
        dsize = (nsrc,)
        rec_geometry = example_rec_geometry(nsrc=nsrc, nrec=nrec)
        data = randn(Float32, ns, nrec)
        d_obs = judiVector(rec_geometry, data)
        @test typeof(d_obs) == judiVector{Float32, Array{Float32, 2}}
        @test isequal(process_input_data(d_obs, rec_geometry), d_obs)
        @test isequal(process_input_data(d_obs), d_obs.data)

        @test isequal(d_obs.nsrc, nsrc)
        @test isequal(typeof(d_obs.data), Array{Array{Float32, 2}, 1})
        @test isequal(typeof(d_obs.geometry), GeometryIC{Float32})
        @test iszero(norm(d_obs.data[1] - d_obs.data[end]))
        @test isequal(size(d_obs), dsize)

        # set up judiVector from cell array
        data = Array{Array{Float32, 2}, 1}(undef, nsrc)
        for j=1:nsrc
            data[j] = randn(Float32, ns, nrec)
        end
        d_obs =  judiVector(rec_geometry, data)

        @test isequal(d_obs.nsrc, nsrc)
        @test isequal(typeof(d_obs.data), Array{Array{Float32, 2}, 1})
        @test isequal(typeof(d_obs.geometry), GeometryIC{Float32})
        @test iszero(norm(d_obs.data - d_obs.data))
        @test isequal(size(d_obs), dsize)
        @test isapprox(convert_to_array(d_obs), vcat([vec(d) for d in data]...); rtol=ftol)

        # contructor for in-core data container
        block = segy_read(datapath*"unit_test_shot_records_$(nsrc).segy"; warn_user=false)
        d_block = judiVector(block; segy_depth_key="RecGroupElevation")
        dsize = (nsrc,)

        @test isequal(d_block.nsrc, nsrc)
        @test isequal(typeof(d_block.data), Array{Array{Float32, 2}, 1})
        @test isequal(typeof(d_block.geometry), GeometryIC{Float32})
        @test isequal(size(d_block), dsize)

        # contructor for in-core data container and given geometry
        rec_geometry = Geometry(block; key="receiver", segy_depth_key="RecGroupElevation")
        d_block = judiVector(rec_geometry, block)

        @test isequal(d_block.nsrc, nsrc)
        @test isequal(typeof(d_block.data), Array{Array{Float32, 2}, 1})
        @test isequal(typeof(d_block.geometry), GeometryIC{Float32})
        @test isequal(rec_geometry, d_block.geometry)
        @test isequal(size(d_block), dsize)

        # contructor for out-of-core data container from single container
        container = segy_scan(datapath, "unit_test_shot_records_$(nsrc)", ["GroupX", "GroupY", "RecGroupElevation", "SourceSurfaceElevation", "dt"])
        d_cont = judiVector(container; segy_depth_key="RecGroupElevation")

        @test typeof(d_cont) == judiVector{Float32, SeisCon}
        @test isequal(d_cont.nsrc, nsrc)
        @test isequal(typeof(d_cont.data), Array{SegyIO.SeisCon, 1})
        @test isequal(typeof(d_cont.geometry), GeometryOOC{Float32})
        @test isequal(size(d_cont), dsize)

        # contructor for single out-of-core data container and given geometry
        d_cont = judiVector(rec_geometry, container)

        @test isequal(d_cont.nsrc, nsrc)
        @test isequal(typeof(d_cont.data), Array{SegyIO.SeisCon, 1})
        @test isequal(typeof(d_cont.geometry), GeometryIC{Float32})
        @test isequal(rec_geometry, d_cont.geometry)
        @test isequal(size(d_cont), dsize)

        # contructor for out-of-core data container from cell array of containers
        container_cell = Array{SegyIO.SeisCon}(undef, nsrc)
        for j=1:nsrc
            container_cell[j] = split(container, j)
        end
        d_cont =  judiVector(container_cell; segy_depth_key="RecGroupElevation")

        @test isequal(d_cont.nsrc, nsrc)
        @test isequal(typeof(d_cont.data), Array{SegyIO.SeisCon, 1})
        @test isequal(typeof(d_cont.geometry), GeometryOOC{Float32})
        @test isequal(size(d_cont), dsize)

        # contructor for out-of-core data container from cell array of containers and given geometry
        d_cont =  judiVector(rec_geometry, container_cell)

        @test isequal(d_cont.nsrc, nsrc)
        @test isequal(typeof(d_cont.data), Array{SegyIO.SeisCon, 1})
        @test isequal(typeof(d_cont.geometry), GeometryIC{Float32})
        @test isequal(rec_geometry, d_cont.geometry)
        @test isequal(size(d_cont), dsize)

    #################################################### test operations ###################################################

        # conj, transpose, adjoint
        @test isequal(size(d_obs), size(conj(d_obs)))
        @test isequal(size(d_block), size(conj(d_block)))
        @test isequal(size(d_cont), size(conj(d_cont)))

        @test isequal(reverse(size(d_obs)), size(transpose(d_obs)))
        @test isequal(reverse(size(d_block)), size(transpose(d_block)))
        @test isequal(reverse(size(d_cont)), size(transpose(d_cont)))

        @test isequal(reverse(size(d_obs)), size(adjoint(d_obs)))
        @test isequal(reverse(size(d_block)), size(adjoint(d_block)))
        @test isequal(reverse(size(d_cont)), size(adjoint(d_cont)))

        # +, -, *, /
        @test iszero(norm(2*d_obs - (d_obs + d_obs)))
        @test iszero(norm(d_obs - (d_obs + d_obs)/2))

        @test iszero(norm(2*d_block - (d_block + d_block)))
        @test iszero(norm(d_block - (d_block + d_block)/2))

        # lmul!, rmul!, ldiv!, rdiv!
        data_ones = Array{Array{Float32, 2}, 1}(undef, nsrc)
        for j=1:nsrc
            data_ones[j] = ones(Float32, ns, nrec)
        end
        d1 =  judiVector(rec_geometry, data_ones)
        lmul!(2f0, d1)
        @test all([all(d1.data[i] .== 2f0) for i = 1:d1.nsrc])
        rmul!(d1, 3f0)
        @test all([all(d1.data[i] .== 6f0) for i = 1:d1.nsrc])
        ldiv!(2f0,d1)
        @test all([all(d1.data[i] .== 3f0) for i = 1:d1.nsrc])
        rdiv!(d1, 3f0)
        @test all([all(d1.data[i] .== 1f0) for i = 1:d1.nsrc])

        # vcat
        d_vcat = [d_block; d_block]
        @test isequal(length(d_vcat), 2*length(d_block))
        @test isequal(d_vcat.nsrc, 2*d_block.nsrc)
        @test isequal(d_vcat.geometry.xloc[1], d_block.geometry.xloc[1])

        # dot, norm, abs
        @test isapprox(norm(d_block), sqrt(dot(d_block, d_block)))
        @test isapprox(norm(d_cont), sqrt(dot(d_cont, d_cont)))
        @test isapprox(abs.(d_block.data[1]), abs(d_block).data[1]) # need to add iterate for JUDI vector

        # vector space axioms
        u = judiVector(rec_geometry, randn(Float32, ns, nrec))
        v = judiVector(rec_geometry, randn(Float32, ns, nrec))
        w = judiVector(rec_geometry, randn(Float32, ns, nrec))
        a = .5f0 + rand(Float32)
        b = .5f0 + rand(Float32)

        @test isapprox(u + (v + w), (u + v) + w; rtol=ftol)
        @test isapprox(u + v, v + u; rtol=ftol)
        @test isapprox(-u, -1f0 * u; rtol=ftol)
        @test iszero(norm(u + u*(-1)))
        @test isapprox(a .* (b .* u), (a * b) .* u; rtol=ftol)
        @test isapprox(u, u .* 1; rtol=ftol)
        @test isapprox(a .* (u + v), a .* u + a .* v; rtol=1f-5)
        @test isapprox((a + b) .* v, a .* v + b .* v; rtol=1f-5)

        # subsamling
        d_block_sub = d_block[1]
        @test isequal(d_block_sub.nsrc, 1)
        @test isequal(typeof(d_block_sub.geometry), GeometryIC{Float32})
        @test isequal(typeof(d_block_sub.data), Array{Array{Float32, 2}, 1})

        inds = nsrc > 1 ? (1:nsrc) : 1
        d_block_sub = d_block[inds]
        @test isequal(d_block_sub.nsrc, nsrc)
        @test isequal(typeof(d_block_sub.geometry), GeometryIC{Float32})
        @test isequal(typeof(d_block_sub.data), Array{Array{Float32, 2}, 1})

        d_cont =  judiVector(container_cell; segy_depth_key="RecGroupElevation")
        d_cont_sub = d_cont[1]
        @test isequal(d_cont_sub.nsrc, 1)
        @test isequal(typeof(d_cont_sub.geometry), GeometryOOC{Float32})
        @test isequal(typeof(d_cont_sub.data), Array{SegyIO.SeisCon, 1})

        d_cont_sub = d_cont[inds]
        @test isequal(d_cont_sub.nsrc, nsrc)
        @test isequal(typeof(d_cont_sub.geometry), GeometryOOC{Float32})
        @test isequal(typeof(d_cont_sub.data), Array{SegyIO.SeisCon, 1})

        # Conversion to SegyIO.Block
        src_geometry = Geometry(block; key="source", segy_depth_key="SourceSurfaceElevation")
        wavelet = randn(Float32, src_geometry.nt[1])
        q = judiVector(src_geometry, wavelet)
        block_out =  judiVector_to_SeisBlock(d_block, q; source_depth_key="SourceSurfaceElevation", receiver_depth_key="RecGroupElevation")

        @test isapprox(block.data, block_out.data)
        @test isapprox(get_header(block, "SourceX"), get_header(block_out, "SourceX"); rtol=1f-6)
        @test isapprox(get_header(block, "GroupX"), get_header(block_out, "GroupX"); rtol=1f-6)
        @test isapprox(get_header(block, "RecGroupElevation"), get_header(block_out, "RecGroupElevation"); rtol=1f-6)
        @test isequal(get_header(block, "ns"), get_header(block_out, "ns"))
        @test isequal(get_header(block, "dt"), get_header(block_out, "dt"))

        block_obs =  judiVector_to_SeisBlock(d_obs, q; source_depth_key="SourceSurfaceElevation", receiver_depth_key="RecGroupElevation")
        d_obs1 = judiVector(block_obs)
        @test isapprox(d_obs1.data, d_obs.data)

        block_q = src_to_SeisBlock(q)
        q_1 = judiVector(block_q)
        for i = 1:nsrc
            vec(q_1.data[i]) == vec(q.data[i])
        end

        # scale
        a = .5f0 + rand(Float32)
        d_scale = deepcopy(d_block)

        # Test norms
        d_ones = judiVector(rec_geometry, 2f0 .* ones(Float32, ns, nrec))
        @test isapprox(norm(d_ones, 2), sqrt(rec_geometry.dt[1]*nsrc*ns*nrec*4))
        @test isapprox(norm(d_ones, 1), rec_geometry.dt[1]*nsrc*ns*nrec*2)
        @test isapprox(norm(d_ones, Inf), 2)

        # Indexing and utilities
        @test isfinite(d_obs)
        @test ndims(judiVector{Float32, Array{Float32, 2}}) == 1
        @test isapprox(sum(d_obs), sum(sum([vec(d) for d in d_obs])))
        d0 = copy(d_obs)
        fill!(d0, 0f0)
        @test iszero(norm(d0))

        @test firstindex(d_obs) == 1
        @test lastindex(d_obs) == nsrc
        @test axes(d_obs) == (Base.OneTo(nsrc),)
        @test ndims(d_obs) == 1

        d0[1] = d_obs.data[1]
        @test isapprox(d0.data[1], d_obs.data[1])

        # broadcast multiplication
        u = judiVector(rec_geometry, randn(Float32, ns, nrec))
        v = judiVector(rec_geometry, randn(Float32, ns, nrec))
        u_scale = deepcopy(u)
        v_scale = deepcopy(v)

        u_scale .*= 2f0
        @test isapprox(u_scale, 2f0 * u; rtol=ftol)
        v_scale .+= 2f0
        @test isapprox(v_scale, 2f0 + v; rtol=ftol)
        u_scale ./= 2f0
        @test isapprox(u_scale, u; rtol=ftol)
        u_scale .= 2f0 .* u_scale .+ v_scale
        @test isapprox(u_scale, 2f0 * u + 2f0 + v; rtol=ftol)
        u_scale .= u .+ v
        @test isapprox(u_scale, u + v)
        u_scale .= u .- v
        @test isapprox(u_scale, u - v)
        u_scale .= u .* v
        @test isapprox(u_scale.data[1], u.data[1].*v.data[1])
        u_scale .= u ./ v
        @test isapprox(u_scale.data[1], u.data[1]./v.data[1])

        # broadcast identity
        u = judiVector(rec_geometry, randn(Float32, ns, nrec))
        v = judiVector(rec_geometry, randn(Float32, ns, nrec))
        u_id = deepcopy(u)
        v_id = deepcopy(v)
        broadcast!(identity, u_id, v_id)    # copy v_id into u_id

        @test isapprox(v, v_id)
        @test isapprox(v, u_id)

        # in-place overwrite
        u = judiVector(rec_geometry, randn(Float32, ns, nrec))
        v = judiVector(rec_geometry, randn(Float32, ns, nrec))
        u_cp = deepcopy(u)
        v_cp = deepcopy(v)
        copy!(v_cp, u_cp)

        @test isapprox(u, u_cp)
        @test isapprox(u, v_cp)

        # similar
        d_zero = similar(d_block, Float32)

        @test isequal(d_zero.geometry, d_block.geometry)
        @test isequal(size(d_zero), size(d_block))

        # retrieve out-of-core data
        d_get = get_data(d_cont)
        @test isapprox(d_block, d_get)

        # Test copies/similar
        w1 = deepcopy(d_obs)
        @test isapprox(w1, d_obs)
        w1 = similar(d_obs)
        @test w1.nsrc == d_obs.nsrc
        w1 .= d_obs
        @test w1.nsrc == d_obs.nsrc
        @test isapprox(w1.data, d_obs.data)

        # Test transducer
        q = judiVector(Geometry(0f0, 0f0, 0f0; dt=2, t=1000), randn(Float32, 251))
        tr = transducer(q, (10, 10), 30, pi/2)
        @test length(tr.geometry.xloc[1]) == 22
        @test tr.geometry.xloc[1][1:11] == range(-30., 30., length=11)
        @test tr.geometry.xloc[1][12:end] == range(-30., 30., length=11)
        @test all(tr.geometry.zloc[1][12:end] .== -10f0)
        @test all(tr.geometry.zloc[1][1:11] .== 0f0)

        q = judiVector(Geometry(0f0, 0f0, 0f0; dt=2, t=1000), randn(Float32, 251))
        tr = transducer(q, (10, 10), 30, pi)
        @test length(tr.geometry.xloc[1]) == 22
        @test isapprox(tr.geometry.zloc[1][1:11], range(30., -30., length=11); atol=1f-14, rtol=1f-14)
        @test isapprox(tr.geometry.zloc[1][12:end], range(30., -30., length=11); atol=1f-14, rtol=1f-14)
        @test isapprox(tr.geometry.xloc[1][12:end], -10f0*ones(11); atol=1f-14, rtol=1f-14)
        @test isapprox(tr.geometry.xloc[1][1:11], zeros(11); atol=1f-14, rtol=1f-14)

        # Test integral & derivative
        refarray = Array{Array{Float32, 2}, 1}(undef, nsrc)
        for j=1:nsrc
            refarray[j] = randn(Float32, ns, nrec)
        end
        d_orig = judiVector(rec_geometry, refarray)

        dt = rec_geometry.dt[1]

        d_cumsum = cumsum(d_orig)
        for i = 1:d_orig.nsrc
            @test isapprox(dt * cumsum(refarray[i],dims=1), d_cumsum.data[i])
        end

        d_diff = diff(d_orig)
        for i = 1:d_orig.nsrc
            @test isapprox(1/dt * refarray[i][1,:], d_diff.data[i][1,:])
            @test isapprox(d_diff.data[i][2:end,:], 1/dt * diff(refarray[i],dims=1))
        end

        @test isapprox(cumsum(d_orig,dims=1),cumsum(d_orig))
        @test isapprox(diff(d_orig,dims=1),diff(d_orig))

        d_cumsum_rec = cumsum(d_orig,dims=2)
        for i = 1:d_orig.nsrc
            @test isapprox(cumsum(refarray[i],dims=2), d_cumsum_rec.data[i])
        end

        d_diff_rec = diff(d_orig,dims=2)
        for i = 1:d_orig.nsrc
            @test isapprox(refarray[i][:,1], d_diff_rec.data[i][:,1])
            @test isapprox(d_diff_rec.data[i][:,2:end], diff(refarray[i],dims=2))
        end
    end
end


#=
Author : Ng Yeesian
Modified : Guy Farmer / Zander Tedjo / Will Worthington / Michael Hilborn
generates the closest dispatch model
=#
mutable struct ClosestDispatch <: DispatchModel
    drivetime::DataFrame
    candidates::Vector{Vector{Int}}
    available::Vector{Int}
    assignment::Vector{Int} # which location the ambulance is assigned to
    ambulances::Vector{Vector{Int}} # list of ambulances assigned to each location
    status::Vector{Symbol} # the current status of the ambulance
        # possible statuses: :available, :responding, :atscene, :conveying, :returning
    fromtime::Vector{Int} # the time it started the new status
    hospital::Vector{Int}
end

function ClosestDispatch(p::DeploymentProblem, drivetime::DataFrame,distribution::Vector{Int})
    candidates = Any[]
    I = 1:p.nlocations
    for region in 1:p.nregions
        push!(candidates, I[vec(p.coverage[region,:])])
    end

    assignment = zeros(Int, p.nambulances)
    ambulances = [Int[] for i in 1:p.nlocations]
    k = 1
    for i in eachindex(distribution), j in 1:distribution[i]
        assignment[k] = i
        push!(ambulances[i], k)
        k += 1
    end
    status = fill(:available, p.nambulances)
    fromtime = zeros(Int, p.nambulances)
    hospital = zeros(Int, p.nambulances)

    ClosestDispatch(drivetime, candidates, distribution, assignment, ambulances, status, fromtime, hospital)

end

function ClosestDispatch(p::DeploymentProblem, problem::DispatchProblem)
    stn_names = [Symbol("stn$(i)_min") for i in 1:size(p.coverage,2)]
    ClosestDispatch(p, problem.emergency_calls[:, stn_names], problem.available)
end

function update_ambulances!(model::ClosestDispatch, i::Int, delta::Int)
    model.available[i] += delta
    @assert model.available[i] >= 0
end

function closest_available(dispatch::ClosestDispatch, id::Int, problem::DispatchProblem)
    location = 0
    min_time = Inf
    for i in dispatch.candidates[problem.emergency_calls[id, :neighborhood]]
        if problem.available[i] > 0 && dispatch.drivetime[id, i] < min_time
            location = i
            min_time = dispatch.drivetime[id, i]
        end
    end
    location
end

function respond_to!(dispatch::ClosestDispatch, i::Int, t::Int)
    @assert length(dispatch.ambulances[i]) > 0 "$i: $(dispatch.ambulances[i])"
    amb = popfirst!(dispatch.ambulances[i])
    @assert amb != 0
    @assert dispatch.status[amb] == :available "$amb: $(dispatch.status[amb])"
    dispatch.status[amb] = :responding
    amb
end

function arriveatscene!(dispatch::ClosestDispatch, amb::Int, t::Int)
    @assert dispatch.status[amb] == :responding "$amb: $(dispatch.status[amb])"
    @assert dispatch.hospital[amb] == 0
    dispatch.status[amb] = :atscene
    dispatch.fromtime[amb] = t
end

function arriveathospital!(dispatch::DispatchModel, amb::Int, hosp::Int, t::Int)
    @assert dispatch.status[amb] == :atscene "$amb: $(dispatch.status[amb])"
    @assert dispatch.hospital[amb] != 0
    dispatch.status[amb] = :atHospital
    dispatch.fromtime[amb] = t
    dispatch.hospital[amb] = hosp
end

function returning_to!(dispatch::DispatchModel, amb::Int, t::Int)
    @assert dispatch.status[amb] == :atHospital "$amb: $(dispatch.status[amb])"
    @assert dispatch.hospital[amb] != 0
    dispatch.status[amb] = :returning
    dispatch.fromtime[amb] = t
    dispatch.assignment[amb]
end


using MappedArrays
using Test

#ambs = detect_ambiguities(ImageCore,AxisArrays,Base,Core)
#ambs = setdiff(detect_ambiguities(ImageAxes,ImageCore,AxisArrays,Base,Core), ambs)
#=
if !isempty(ambs)
    println("Ambiguities:")
    for a in ambs
        println(a)
    end
end
@test isempty(ambs)
=#

#@traitfn has_time_axis(::AA) where {AA<:AxisArray;  HasTimeAxis{AA}} = true
#@traitfn has_time_axis(::AA) where {AA<:AxisArray; !HasTimeAxis{AA}} = false

@testset "no units, no time" begin
    A = AxisArray(reshape(1:12, 3, 4), x = (1:3), y = (1:4))
    @test @inferred(timeaxis(A)) === nothing
    @test !has_time_axis(A)
    @test timedim(A) == 0
    @test @inferred(pixelspacing(A)) === (1,1)
    @test @inferred(spacedirections(A)) === ((1,0),(0,1))
    @test @inferred(coords_spatial(A)) === (1,2)
    @test spatialorder(A) === (:x, :y)  # TODO: make this inferrable
    @test @inferred(size_spatial(A)) === (3,4)
    @test @inferred(indices_spatial(A)) === (Base.OneTo(3), Base.OneTo(4))
    assert_timedim_last(A)
    @test map(istimeaxis, AxisArrays.axes(A)) == (false,false)

    @test @inferred(timeaxis(rand(3,5))) == nothing

    # deprecate
    @test nimages(A) == 1
end

@testset "units, no time" begin
    mm = u"mm"     # in real use these should be global consts
    m = u"m"
    A = AxisArray(reshape(1:12, 3, 4), x = (1mm:1mm:3mm), y = (1m:2m:7m))
    @test @inferred(timeaxis(A)) === nothing
    @test !has_time_axis(A)
    @test timedim(A) == 0
    @test nimages(A) == 1
    @test @inferred(pixelspacing(A)) === (1mm,2m)
    @test spacedirections(A) === ((1mm,0m),(0mm,2m))   # TODO: make this inferrable
    @test @inferred(coords_spatial(A)) === (1,2)
    @test spatialorder(A) === (:x,:y)
    @test @inferred(size_spatial(A)) === (3,4)
    @test @inferred(indices_spatial(A)) === (Base.OneTo(3),Base.OneTo(4))
    assert_timedim_last(A)
    @test map(istimeaxis, AxisArrays.axes(A)) == (false,false)
end

@testset "units, time" begin
    s = u"s" # again, global const
    axt = Axis(1s:1s:4s)
    A = AxisArray(reshape(1:12, 3, 4), Axis(1:3), axt)
    @test @inferred(timeaxis(A)) === axt
    @test has_time_axis(A)
    @test timedim(A) == 2
    @test nimages(A) == 4
    @test @inferred(pixelspacing(A)) === (1,)
    @test @inferred(spacedirections(A)) === ((1,),)
    @test @inferred(coords_spatial(A)) === (1,)
    @test spatialorder(A) === (:x,)
    @test @inferred(size_spatial(A)) === (3,)
    @test @inferred(indices_spatial(A)) === (Base.OneTo(3),)
    assert_timedim_last(A)
    @test map(istimeaxis, AxisArrays.axes(A)) == (false,true)
end

@testset "units, time first" begin
    s = u"s" # global const
    axt = Axis{:time}(1s:1s:4s)
    A = AxisArray(reshape(1:12, 4, 3), axt, Axis{:x}(1:3))
    @test @inferred(timeaxis(A)) === axt
    @test has_time_axis(A)
    @test timedim(A) == 1
    @test nimages(A) == 4
    @test @inferred(pixelspacing(A)) === (1,)
    @test @inferred(spacedirections(A)) === ((1,),)
    @test @inferred(coords_spatial(A)) === (2,)
    @test spatialorder(A) === (:x,)
    @test @inferred(size_spatial(A)) === (3,)
    @test @inferred(indices_spatial(A)) === (Base.OneTo(3),)
    @test_throws ErrorException assert_timedim_last(A)
    @test map(istimeaxis, AxisArrays.axes(A)) == (true,false)
end

# Possibly-ambiguous functions
@testset "ambig" begin
    A = AxisArray(rand(RGB{N0f8},3,5), :x, :y)
    @test isa(convert(Array{RGB{N0f8},2}, A), Array{RGB{N0f8},2})
    @test isa(convert(Array{Gray{N0f8},2}, A), Array{Gray{N0f8},2})
end

@testset "internal" begin
    A = AxisArray(rand(RGB{N0f8},3,5), :x, :y)
    @test ImageAxes.axtype(A) == Tuple{Axis{:x,Base.OneTo{Int}}, Axis{:y,Base.OneTo{Int}}}
end

# For testing streaming with a non-AxisArray parent
module TestStreaming
using AxisArrays, ImageAxes

struct AVIStream
    dims::NTuple{3,Int}
end
Base.ndims(::AVIStream) = 3
Base.size(A::AVIStream) = A.dims
AxisArrays.axisnames(::Type{AS}) where {AS<:AVIStream} = (:y, :x, :time)
AxisArrays.axes(A::AVIStream) = (Axis{:y}(Base.OneTo(A.dims[1])),
                                 Axis{:x}(Base.OneTo(A.dims[2])),
                                 Axis{:time}(Base.OneTo(A.dims[3])))
ImageAxes.StreamIndexStyle(::Type{AVIStream}, ::Type{typeof(read!)}) =
    IndexIncremental()

end

@testset "streaming" begin
    P = AxisArray([0 0 0 0;
                   1 2 3 4;
                   0 0 0 0], :x, :time)
    f!(dest, a) = (dest[1] = dest[3] = -0.2*a[2]; dest[2] = 0.6*a[2]; dest)
    # Next inference was special-cased for v0.6
    S = @inferred(StreamingContainer{Float64}(f!, P, Axis{:time}()))
    @test @inferred(axes(S)) === (Base.OneTo(3), Base.OneTo(4))
    @test @inferred(size(S)) == (3,4)
    @test @inferred(axes(S, 2)) === Base.OneTo(4)
    @test @inferred(size(S, 1)) === 3
    @test @inferred(length(S)) == 12
    @test @inferred(axisnames(S)) == (:x, :time)
    @test @inferred(axisvalues(S)) === (Base.OneTo(3), Base.OneTo(4))
    @test axisdim(S, Axis{:x}) == axisdim(S, Axis{:x}(1:2)) == axisdim(S, Axis{:x,UnitRange{Int}}) == 1
    @test axisdim(S, Axis{:time}) == 2
    @test_throws ErrorException axisdim(S, Axis{:y})
    @test axisdim(S, Axis{2}) == 2
    @test_throws ErrorException axisdim(S, Axis{3})
    @test @inferred(timeaxis(S)) === Axis{:time}(Base.OneTo(4))
    @test nimages(S) == 4
    @test @inferred(coords_spatial(S)) == (1,)
    @test @inferred(indices_spatial(S)) == (Base.OneTo(3),)
    @test @inferred(size_spatial(S)) == (3,)
    @test @inferred(spatialorder(S)) == (:x,)
    assert_timedim_last(S)
    for i = 1:4
        @test @inferred(S[:,i]) == [-0.2,0.6,-0.2]*i
        @test @inferred(S[2,i]) === 0.6*i
        @test @inferred(S[Axis{:time}(i)]) == [-0.2,0.6,-0.2]*i
        @test @inferred(S[Axis{:time}(i),Axis{:x}(2)]) === 0.6*i
        @test @inferred(S[Axis{:x}(2),Axis{:time}(i)]) === 0.6*i
    end
    buf = zeros(3)
    @test @inferred(getindex!(buf, S, :, 2)) == [-0.2,0.6,-0.2]*2
    @test StreamIndexStyle(S) === IndexAny()
    @test StreamIndexStyle(zeros(2,2)) === IndexAny()
    # Non-AbstractArray parent
    @test_throws DimensionMismatch StreamingContainer{UInt8}(read!, TestStreaming.AVIStream((1080,1920,10000)), Axis{:foo}())
    V = StreamingContainer{UInt8}(read!, TestStreaming.AVIStream((1080,1920,10000)), Axis{:time}())
    @test size(V) == (1080,1920,10000)
    @test axisnames(V) == (:y, :x, :time)
    @test StreamIndexStyle(V) === IndexIncremental()
    # internal
    @test ImageAxes.streamingaxisnames(S) == (:time,)
    @test ImageAxes.filter_streamed((1,2), S) == (2,)
end

nothing


module TestBatchedVector

include("preamble.jl")

function makedata(V::Type, nbatches::Integer)
    batch_lengths = collect(testdims(nbatches))
    offsets = [0, cumsum(batch_lengths)...]
    nested = Vector{V}[rand(V, bl) for bl in batch_lengths]
    flat = nbatches == 0 ? V[] : reduce(vcat, nested)
    return (batch_lengths = batch_lengths, offsets = offsets, nested=nested, flat=flat)
end

@testset "nbatches = $nbatches, V = $V" for nbatches in (1, 10), V in (Float64, )
    @testset "constructors" begin
        @unpack batch_lengths, offsets, nested, flat = makedata(V, nbatches)
        Expected = BatchedVector{<:AbstractArray{V,inner_ndims(nested)},typeof(flat)}
        @test BatchedVector(flat, offsets) isa Expected
        @test_inferred BatchedVector(flat, offsets)
    end

    @testset "Extra" begin
        @unpack batch_lengths, offsets, nested, flat = makedata(V, nbatches)
        Expected = BatchedVector{<:AbstractArray{V,inner_ndims(nested)},typeof(flat)}
        @test batch(flat, batch_lengths) isa Expected
        @test_inferred batch(flat, batch_lengths)
        let B1 = batch(flat, batch_lengths)
            B2 = similarbatch(copy(flat), B1)
            @test B1 == B2
            @test B1.parent !== B2.parent
        end
    end

    let nbatches = nbatches, V = V
        test_array_AB() do
            @unpack batch_lengths, nested, flat = makedata(V, nbatches)
            A = batch(flat, batch_lengths)
            B = nested
            return A, B
        end
    end
end

@testset "Misc" begin
    d = makedata(Float64, 5)
    B = BatchedVector(d.flat, d.offsets)
    @test Base.dataids(B) === (Base.dataids(d.flat)..., Base.dataids(d.offsets)...)
    @test parent(B) === d.flat
end

@testset "UnsafeArrays" begin
    d = makedata(Float64, 5)
    B = BatchedVector(d.flat, d.offsets)
    U = uview(B)
    @test U == B
    @test U.parent isa UnsafeArray{eltype(d.flat),1}
    @test U.offsets isa Vector{Int}
    @test_inferred uview(B)
end

end # module


function lazy_append(a::Triangulation,b::Triangulation)
  AppendedTriangulation(a,b)
end

struct AppendedTriangulation{Dc,Dp} <: Triangulation{Dc,Dp}
  a::Triangulation{Dc,Dp}
  b::Triangulation{Dc,Dp}
  function AppendedTriangulation(a::Triangulation{Dc,Dp},b::Triangulation{Dc,Dp}) where {Dc,Dp}
    @assert get_background_triangulation(a) === get_background_triangulation(b)
    new{Dc,Dp}(a,b)
  end
end

TriangulationStyle(::Type{<:AppendedTriangulation}) = SubTriangulation()

get_background_triangulation(trian::AppendedTriangulation) = get_background_triangulation(trian.a)

function get_cell_coordinates(trian::AppendedTriangulation)
  a = get_cell_coordinates(trian.a)
  b = get_cell_coordinates(trian.b)
  lazy_append(a,b)
end

function get_cell_ref_coordinates(trian::AppendedTriangulation)
  a = get_cell_ref_coordinates(trian.a)
  b = get_cell_ref_coordinates(trian.b)
  lazy_append(a,b)
end

function get_reffes(trian::AppendedTriangulation)
  vcat(get_reffes(trian.a),get_reffes(trian.b))
end

function get_cell_type(trian::AppendedTriangulation)
  a = get_cell_type(trian.a)
  b = get_cell_type(trian.b) .+ Int8(length(get_reffes(trian.a)))
  lazy_append(a,b)
end

#function reindex(f::AbstractArray,trian::AppendedTriangulation)
#  a = reindex(f,trian.a)
#  b = reindex(f,trian.b)
#  lazy_append(a,b)
#end

function get_cell_to_bgcell(trian::AppendedTriangulation)
  a = get_cell_to_bgcell(trian.a)
  b = get_cell_to_bgcell(trian.b)
  lazy_append(a,b)
end

function get_cell_ref_map(trian::AppendedTriangulation)
  a = get_cell_ref_map(trian.a)
  b = get_cell_ref_map(trian.b)
  lazy_append(a,b)
end

#function restrict(f::AbstractArray,trian::AppendedTriangulation)
#  a = restrict(f,trian.a)
#  b = restrict(f,trian.b)
#  lazy_append(a,b)
#end

# In this case, we do not want a lazy_append since it will become difficult to
# compress / expand the reffes.
#function get_cell_reffe(trian::AppendedTriangulation)
#  a = get_cell_reffe(trian.a)
#  b = get_cell_reffe(trian.b)
#  lazy_append(a,b)
#end

function get_cell_shapefuns(trian::AppendedTriangulation)
  a = get_cell_shapefuns(trian.a)
  b = get_cell_shapefuns(trian.b)
  lazy_append(a,b)
end

function get_cell_map(trian::AppendedTriangulation)
  a = get_cell_map(trian.a)
  b = get_cell_map(trian.b)
  lazy_append(a,b)
end

function get_facet_normal(trian::AppendedTriangulation)
  cm = get_cell_map(trian)
  a = get_facet_normal(trian.a)
  b = get_facet_normal(trian.b)
  lazy_append(a,b)
end

#function CellQuadrature(trian::AppendedTriangulation,degree1::Integer,degree2::Integer)
#  quad1 = CellQuadrature(trian.a,degree1)
#  quad2 = CellQuadrature(trian.b,degree2)
#  lazy_append(quad1,quad2)
#end
#
#function CellQuadrature(trian::AppendedTriangulation,degree::Integer)
#  CellQuadrature(trian,degree,degree)
#end

#function compress_contributions(cell_mat::AppendedArray,trian::AppendedTriangulation)
#  if length(cell_mat.a) == num_cells(trian.a) && length(cell_mat.b) == num_cells(trian.b)
#    a = compress_contributions(cell_mat.a,trian.a)
#    b = compress_contributions(cell_mat.b,trian.b)
#    return lazy_append(a,b)
#  else
#    return cell_mat
#  end
#end
#
#function compress_ids(cell_ids::AppendedArray,trian::AppendedTriangulation)
#  if length(cell_ids.a) == num_cells(trian.a) && length(cell_ids.b) == num_cells(trian.b)
#    a = compress_ids(cell_ids.a,trian.a)
#    b = compress_ids(cell_ids.b,trian.b)
#    return lazy_append(a,b)
#  else
#    return cell_ids
#  end
#end

function compress_contributions(cell_mat,trian::AppendedTriangulation)
  cell_to_bgcell = get_cell_to_bgcell(trian)
  ccell_mat = compress_contributions(cell_mat,cell_to_bgcell)
  ccell_mat
end

function compress_ids(cell_ids,trian::AppendedTriangulation)
  cell_to_bgcell = get_cell_to_bgcell(trian)
  compress_ids(cell_ids,cell_to_bgcell)
end


module TimeToContact

using Plots
using Images
using FastConv
using ImageView
import VideoIO

include("./utils.jl");

end # module


# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Description
#
#   Private functions.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

################################################################################
#                                    Macros
################################################################################

"""
    @parse_value(T, str, line_num)

Parse the string `str` using the type `T`. If it is not succeeded, then throw an
error indicating the line `line_num` with the problem.

"""
macro parse_value(T, str, line_num)
    quote
        local _T = $(esc(T))
        local _str = $(esc(str))
        local _line_num = $(esc(line_num))

        try
            parse(_T, _str)
        catch e
            throw(ErrorException("The TLE file is not valid (error in line $_line_num): $(e.msg)."))
        end
    end
end

################################################################################
#                                  Functions
################################################################################

function _get_mant_exp(n)
    if abs(n) < 1e-20
        return 0., 0
    end

    exp  = floor(Int, log10(n)) + 1
    mant = n/10.0^exp

    return mant,exp
end

function _parse_tle(io::IO, verify_checksum::Bool = true)
    # State machine to read the TLE. It has three possible states:
    #
    #   :name -> Satellite name.
    #   :l1   -> Line 1.
    #   :l2   -> Line 2.
    #
    # The transitions are:
    #
    #   :name --> :l1 --> :l2
    #     ^                |
    #     |                |
    #      ----------------
    #
    state = :name

    # Auxiliary variables that will store the TLE.
    name = ""
    sat_num = 0
    classification = Char(0)
    int_designator = ""
    epoch_year = 0
    epoch_day = 0.0
    epoch = 0.0
    dn_o2 = 0.0
    ddn_o6 = 0.0
    bstar = 0.0
    elem_set_number = 0
    checksum_l1 = 0
    i = 0.0
    Ω = 0.0
    e = 0.0
    ω = 0.0
    M = 0.0
    n = 0.0
    rev_num = 0
    checksum_l2 = 0

    # Output array with the TLEs found in the file.
    tles = Array{TLE}(undef,0)

    line_num = 0

    skip_line_read = false

    line = nothing

    while !eof(io)
        # Read the current line, strip white spaces, and skip if it is blank or
        # is a comment.
        if !skip_line_read
            line = strip(readline(io))
            line_num += 1
            (isempty(line))  && continue
            (line[1] == '#') && continue
        else
            skip_line_read = false
        end

        # Check the state of the reading.
        if state == :name
            # Check if the line seems to be the first line of the TLE. In this
            # case, maybe the user has not provided a name. Then, change the
            # state to `:l1`, and skip the line reading in the next loop.
            if (line[1:2] == "1 ") && (length(line) == 69)
                name = "UNDEFINED"
                state = :l1
                skip_line_read  = true
                continue
            end

            # Otherwise, if the line is not blank, then it must be the name of
            # the satellite.
            #
            # NOTE: The name should not be bigger than 24 characters. However,
            # we will not check this here.

            # Copy the name and change to state to wait for the 1st line.
            name  = line
            state = :l1

        # TLE Line 1
        # ==========
        elseif state == :l1
            # The next non-blank line must be the first line of the TLE.
            # Otherwise, the file is not valid.

            # The first line must start with "1 " and have 69 characters.
            if (line[1:2] != "1 ") || (length(line) != 69)
                throw(ErrorException("The TLE file is not valid (error in line $line_num): This is not a valid 1st line."))
            end

            # Verify the checksum.
            # ====================

            checksum_l1   = @parse_value(Int, line[69], line_num)
            checksum_line = compute_checksum(line[1:end-1])

            if (verify_checksum) && (checksum_l1 != checksum_line)
                throw(ErrorException("The TLE file is not valid (error in line $line_num): Expected checksum: $checksum_line, line checksum: $checksum_l1."))
            end

            # Satellite number and classification
            # ===================================
            sat_num        = @parse_value(Int, line[3:7], line_num)
            classification = Char(line[8])

            # International designator
            # ========================
            int_designator = line[10:17]

            # Epoch
            # =====
            epoch_year = @parse_value(Int,     line[19:20], line_num)
            epoch_day  = @parse_value(Float64, line[21:32], line_num)

            # Convert the TLE year and date to JD.
            # ------------------------------------

            # If the two-number year is higher than 75, then consider it in the
            # past (e.g. 50 = 2050, but 81 = 1981).
            if epoch_year > 75
                epoch = datetime2julian(DateTime(1900 + epoch_year, 1, 1, 0, 0, 0)) - 1 +
                        epoch_day
            else
                epoch = datetime2julian(DateTime(2000 + epoch_year, 1, 1, 0, 0, 0)) - 1 +
                        epoch_day
            end

            # Mean motion derivatives
            # =======================
            dn_o2  = @parse_value(Float64, line[34:43], line_num)

            aux = ( (line[45] == ' ') ? "+." : line[45:45] * "." )*line[46:50]

            ddn_o6_dec = @parse_value(Float64, aux, line_num)
            ddn_o6_exp = @parse_value(Float64, line[51:52], line_num)
            ddn_o6 = ddn_o6_dec*10^ddn_o6_exp

            # BSTAR
            # =====

            aux = ( (line[54] == ' ') ? "+." : line[54:54] * "." )*line[55:59]

            bstar_dec = @parse_value(Float64, aux, line_num)
            bstar_exp = @parse_value(Float64, line[60:61], line_num)
            bstar = bstar_dec*10^bstar_exp

            # Ephemeris type
            # ==============
            (line[63] != '0' && line[63] != ' ') &&
                warn("Warning in TLE file (line $line_num): Ephemeris type should be 0!")

            # Element number
            # ==============
            elem_set_number = @parse_value(Int, line[65:68], line_num)

            # Change the state to wait for the line 2.
            state = :l2

        # TLE Line 2
        # ==========
        elseif state == :l2
            # The next non-blank line must be the second line of the TLE.
            # Otherwise, the file is not valid.

            # The second line must start with "2 " and have 69 characters.
            if (line[1:2] != "2 ") || (length(line) != 69)
                throw(ErrorException("The TLE file is not valid (error in line $line_num): This is not a valid 2st line."))
            end

            # Verify the checksum.
            # ====================

            checksum_l2   = @parse_value(Int, line[69], line_num)
            checksum_line = compute_checksum(line[1:end-1])

            if (verify_checksum) && (checksum_l2 != checksum_line)
                throw(ErrorException("The TLE file is not valid (error in line $line_num): Expected checksum: $checksum_line, line checksum: $checksum_l2."))
            end

            # Check satellite number with the one in the first line.
            # ======================================================
            sat_num_l2 = @parse_value(Int, line[3:7], line_num)

            if sat_num_l2 != sat_num
                throw(ErrorException("The TLE file is not valid (error in line $line_num): Satellite number in line 2 is not equal to that in line 1."))
            end

            # Inclination
            # ===========
            i = @parse_value(Float64, line[9:16], line_num)

            # RAAN
            # ====
            Ω = @parse_value(Float64, line[18:25], line_num)

            # Eccentricity
            # ============
            e = @parse_value(Float64, "." * line[27:33], line_num)

            # Argument of perigee
            # ===================
            ω = @parse_value(Float64, line[35:42], line_num)

            # Mean anomaly
            # ============
            M = @parse_value(Float64, line[44:51], line_num)

            # Mean motion
            # ===========
            n = @parse_value(Float64, line[53:63], line_num)

            # Revolution number at epoch
            # ==========================
            rev_num = @parse_value(Int, line[64:68], line_num)

            # Now that we have all the information, we can store our TLE in the
            # output array.
            push!(tles,
                  TLE(name,
                      sat_num,
                      classification,
                      int_designator,
                      epoch_year,
                      epoch_day,
                      epoch,
                      dn_o2,
                      ddn_o6,
                      bstar,
                      elem_set_number,
                      checksum_l1,
                      i,
                      Ω,
                      e,
                      ω,
                      M,
                      n,
                      rev_num,
                      checksum_l2))

            # Change the state to wait for another TLE.
            state = :name
        end
    end

    # If the final state is not :name, then we have an incomplete TLE. Thus,
    # throw and exception because the file is not valid.
    if state != :name
        throw(ErrorException("The TLE file is not valid.\nThere is an incomplete TLE."))
    end

    tles
end

"""
    compute_checksum(str::AbstractString)

Compute the checksum of the line `str` modulo 10.

The algorithm is simple: add all the numbers in the line, ignoring letters,
spaces, periods, and plus signs, but assigning +1 to the minus signs. The
checksum is the remainder of the division by 10.

"""
function compute_checksum(str::AbstractString)
    checksum = 0

    for c in str
        # Check if `c` is a number.
        if isnumeric(c)
            checksum += parse(Int, c)

        # Check if `c` is a minus sign, which has value 1.
        elseif c == '-'
            checksum += 1
        end

        # Otherwise, the value of the character is 0.
    end

    # Return the checksum modulo 10.
    checksum % 10
end

"""
    _show_tle(io::IO, tle::TLE, color::Bool = true)

Show the TLE `tle` in the IO `io`.

If `color` is `true`, then the text will be printed using colors. If `color` is
omitted, then it defaults to `true`.

"""
function _show_tle(io::IO, tle::TLE, color::Bool = true)
    # Colors will be printed only for STDOUT.
    b = (color) ? _b : ""
    d = (color) ? _d : ""
    g = (color) ? _g : ""
    y = (color) ? _y : ""

    # Print the TLE information.
    print(io, "                             $(g)TLE$(d)\n")
    print(io, "$(y)    ==========================================================$(d)\n")
    print(io, "$(b)                            Name: $(d)"); @printf(io, "%s\n", tle.name)
    print(io, "$(b)                Satellite number: $(d)"); @printf(io, "%d\n", tle.sat_num)
    print(io, "$(b)        International designator: $(d)"); @printf(io, "%s\n", tle.int_designator)
    print(io, "$(b)                    Epoch (Year): $(d)"); @printf(io, "%d\n", tle.epoch_year)
    print(io, "$(b)                     Epoch (Day): $(d)"); @printf(io, "%12.8f\n", tle.epoch_day)
    print(io, "$(b)              Epoch (Julian Day): $(d)"); @printf(io, "%12.5f\n", tle.epoch)
    print(io, "$(b)              Element set number: $(d)"); @printf(io, "%d\n", tle.elem_set_number)
    print(io, "$(b)                    Eccentricity: $(d)"); @printf(io, "%12.8f deg\n", tle.e)
    print(io, "$(b)                     Inclination: $(d)"); @printf(io, "%12.8f deg\n", tle.i)
    print(io, "$(b)                            RAAN: $(d)"); @printf(io, "%12.8f deg\n", tle.Ω)
    print(io, "$(b)             Argument of perigee: $(d)"); @printf(io, "%12.8f deg\n", tle.ω)
    print(io, "$(b)                    Mean anomaly: $(d)"); @printf(io, "%12.8f deg\n", tle.M)
    print(io, "$(b)                 Mean motion (n): $(d)"); @printf(io, "%12.8f revs/day\n", tle.n)
    print(io, "$(b)               Revolution number: $(d)"); @printf(io, "%d\n", tle.rev_num)
    print(io, "\n")
    print(io, "$(b)                              B*: $(d)"); @printf(io, "%f 1/[er]\n", tle.bstar)
    print(io, "\n")
    print(io, "$(b)                        1   d\n$(d)")
    print(io, "$(b)                       ---.--- n: $(d)"); @printf(io, "%f rev/day²\n", tle.dn_o2)
    print(io, "$(b)                        2  dt\n$(d)")
    print(io, "\n")
    print(io, "$(b)                        1   d²\n$(d)")
    print(io, "$(b)                       ---.--- n: $(d)"); @printf(io, "%f rev/day³\n", tle.ddn_o6)
    print(io, "$(b)                        6  dt²\n$(d)")
    print(io, "$(y)    ==========================================================$(d)")

    nothing
end


function runfig4()
    p = Params(Ln = 2000, W = 2700, Ls =8.649,  Ws = 8.649, scale = 40, λ= 5, α = 0,
        EZ = SA[0, 0.6, 0], μN = 1.192, Δ = 0.3, d = 0, τ = 1);

    presets_fig = Fig4_presets(0.0, 0.0, 0, 4, false)
    data = ldosonlattice_averaged_sc(p, presets_fig)
    savepsi("fig4", p, data[1])

    presets_fig = Fig4_presets(0.0, 0.0, 0, 4, true)
    data = ldosonlattice_averaged_sc(p, presets_fig)
    savepsi("fig4", p, data[1])

    sp = spectrumsweepb(p, 0:.1:2, true)
    savespectrum("fig4", p, sp[1], sp[2])

    sp = spectrumsweepb(p, 0:.1:2, false)
    savespectrum("fig4", p, sp[1], sp[2])
end


module Systems

export PHPSystem,Tube,Evaporator,Condenser,Liquid,Vapor,Wall,Mapping,PHPResult

# using ..Tools

"""
PHPSystem is a struct containing
    γ
    Hc
    He
    θc
    θe
    ω0
    ζ
    L dimensionless pipe total length
    Xc dimensionless condenser range
    Xe dimensionless evaporater range
"""

mutable struct Tube
    d::Float64
    L::Float64
    L2D::Float64
    angle::Float64
    closedornot::Bool
end

"""
PHPSystem is a struct containing
    γ
    Hc
    He
    θc
    θe
    ω0
    ζ
    L dimensionless pipe total length
    Xc dimensionless condenser range
    Xe dimensionless evaporater range
"""

mutable struct Evaporator
    He::Float64
    θe::Float64
    Xe::Array{Tuple{Float64,Float64},1}
    We::Array{Float64,1}
end

"""
PHPSystem is a struct containing
    γ
    Hc
    He
    θc
    θe
    ω0
    ζ
    L dimensionless pipe total length
    Xc dimensionless condenser range
    Xe dimensionless evaporater range
"""

mutable struct Condenser
    Hc::Float64
    θc::Float64
    Xc::Array{Tuple{Float64,Float64},1}
    Hwc::Float64
end

"""
PHPSystem is a struct containing
    γ
    Hc
    He
    θc
    θe
    ω0
    ζ
    L dimensionless pipe total length
    Xc dimensionless condenser range
    Xe dimensionless evaporater range
"""

mutable struct Liquid
    γ::Float64
    Hₗ::Float64
    ρ::Float64
    ω0::Array{Float64,1}
    ℘::Array{Float64,1}
    Xp::Array{Tuple{Float64,Float64},1}
    dXdt::Array{Tuple{Float64,Float64},1}
    Xarrays::Array{Array{Float64,1},1}
    θarrays::Array{Array{Float64,1},1}

end

"""
PHPSystem is a struct containing
    γ
    Hc
    He
    θc
    θe
    ω0
    ζ
    L dimensionless pipe total length
    Xc dimensionless condenser range
    Xe dimensionless evaporater range
"""

mutable struct Vapor
    γ::Float64
    P::Array{Float64,1}
    δ::Array{Float64,1}
end

"""
PHPSystem is a struct containing
    γ
    Hc
    He
    θc
    θe
    ω0
    ζ
    L dimensionless pipe total length
    Xc dimensionless condenser range
    Xe dimensionless evaporater range
"""

mutable struct Wall
    α::Float64
    Δt::Float64
    Xstations::Array{Float64,1}
    Xarray::Array{Float64,1}
    θarray::Array{Float64,1}
end

"""
PHPSystem is a struct containing
    γ
    Hc
    He
    θc
    θe
    ω0
    ζ
    L dimensionless pipe total length
    Xc dimensionless condenser range
    Xe dimensionless evaporater range
"""

mutable struct Mapping
    walltoliquid::Array{Tuple{Int64,Int64},1}
    liquidtowall::Array{Array{Int64,1},1}
end

"""
PHPSystem is a struct containing
    γ
    Hc
    He
    θc
    θe
    ω0
    ζ
    L dimensionless pipe total length
    Xc dimensionless condenser range
    Xe dimensionless evaporater range
"""
# add type names!
mutable struct PHPSystem
    tube::Tube
    evaporator::Evaporator
    condenser::Condenser
    liquid::Liquid
    vapor::Vapor
    wall::Wall
    mapping::Mapping
end

"""
PHPSystem is a struct containing
    γ
    Hc
    He
    θc
    θe
    ω0
    ζ
    L dimensionless pipe total length
    Xc dimensionless condenser range
    Xe dimensionless evaporater range
"""

struct PHPResult
    t
    Xp
    dXdt
    P
    θ
    M
end

end


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

using Base.Threads: threadid, nthreads
using ..Model
import ..Model: BivariateNormalDerivatives, BvnComponent, GalaxyCacheComponent,
                GalaxySigmaDerivs, SkyPatch,
                get_bvn_cov, eval_bvn_pdf!, get_bvn_derivs!,
                transform_bvn_derivs!, init_source
import ..Infer
using ..SensitiveFloats
import ..SensitiveFloats.clear!
import ..Log
using ..Transform
import DataFrames
import Optim
using StaticArrays
using ForwardDiff
using ReverseDiff

export ElboArgs

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

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

Arguments:
    images: a collection of astronomical images
    neighbors: the other light sources near `entry`
    entry: the source to infer
"""
function infer_source(images::Vector{Image},
                      neighbors::Vector{CatalogEntry},
                      entry::CatalogEntry)
    if length(neighbors) > 100
        Log.warn("Excessive number ($(length(neighbors))) of neighbors")
    end

    # It's a bit inefficient to call the next 5 lines every time we optimize_f.
    # But, as long as runtime is dominated by the call to maximize_f, that
    # isn't a big deal.
    cat_local = vcat([entry], neighbors)
    vp = Vector{Float64}[init_source(ce) for ce in cat_local]
    patches = Infer.get_sky_patches(images, cat_local)
    Infer.load_active_pixels!(images, patches)

    ea = ElboArgs(images, vp, patches, [1])
    f_evals, max_f, max_x, nm_result = maximize_f(elbo, ea)
    vp[1]
end

end


using AbstractPlotting
using MakieRipserer
using Ripserer
using Suppressor
using Test

@testset "Nothing throws" begin
    data = [(sin(t), cos(t)) for t in range(0, 2π; length=13)[1:12]]
    time = Observable(0.0)
    @suppress begin
        @test begin
            MakieRipserer.app(data, Rips; time)
            time[] = 1
            true
        end
    end
end


__precompile__(true)

"The module for [PCG Family](@ref)."
module PCG

# PCG
export PCGStateMCG, PCGStateOneseq, PCGStateUnique, PCGStateSetseq

# PCG Methods
export PCG_XSH_RS, PCG_XSH_RR, PCG_RXS_M_XS, PCG_XSL_RR, PCG_XSL_RR_RR

export PCGUInt, PCGMethod, PCG_LIST

export bounded_rand, advance!

const pcg_uints = (UInt8, UInt16, UInt32, UInt64, UInt128)
const PCGUInt = Union{pcg_uints...}

"""
One of PCG output method: high xorshift, followed by a random shift.

Fast.
"""
const PCG_XSH_RS = Val{:XSH_RS}
"""
One of PCG output method: high xorshift, followed by a random rotate.

Fast. Slightly better statistically than `PCG_XSH_RS`.
"""
const PCG_XSH_RR = Val{:XSH_RR}
"""
One of PCG output method: random xorshift, mcg multiply, fixed xorshift.

The most statistically powerful generator, but slower than some of the others.
"""
const PCG_RXS_M_XS = Val{:RXS_M_XS}
"""
One of PCG output method: fixed xorshift (to low bits), random rotate.

Useful for 128-bit types that are split across two CPU registers.
"""
const PCG_XSL_RR = Val{:XSL_RR}
"""
One of PCG output method: fixed xorshift (to low bits), random rotate (both parts).

Useful for 128-bit types that are split across two CPU registers. Use this in need of an invertable 128-bit
RNG.
"""
const PCG_XSL_RR_RR = Val{:XSL_RR_RR}

const pcg_methods = (PCG_XSH_RS, PCG_XSH_RR, PCG_RXS_M_XS, PCG_XSL_RR, PCG_XSL_RR_RR)

"""
The `Union` of all the PCG method types: `PCG_XSH_RS`, `PCG_XSH_RR`, `PCG_RXS_M_XS`, `PCG_XSL_RR`, and `PCG_XSL_RR_RR`.
"""
const PCGMethod = Union{pcg_methods...}

include("pcg_list.jl")

include("bases.jl")

include("main.jl")

end


##### Prognostic variable

export Mass, Momentum, Energy, ρθ_liq_ice
export AbstractMoistureVariable, TotalMoisture, LiquidMoisture, IceMoisture
export AbstractPrecipitationVariable, Rain, Snow
export Tracers
export AbstractSGSTKEVariable, SGSTKE

struct Mass <: AbstractPrognosticVariable end
struct Momentum <: AbstractMomentumVariable end

struct Energy <: AbstractEnergyVariable end
struct ρθ_liq_ice <: AbstractEnergyVariable end

struct TotalMoisture <: AbstractMoistureVariable end
struct LiquidMoisture <: AbstractMoistureVariable end
struct IceMoisture <: AbstractMoistureVariable end

struct Rain <: AbstractPrecipitationVariable end
struct Snow <: AbstractPrecipitationVariable end

struct Tracers{N} <: AbstractTracersVariable{N} end

abstract type AbstractSGSTKEVariable <: AbstractPrognosticVariable end
struct SGSTKE <: AbstractSGSTKEVariable end


##################################################
# Helper functions
##################################################

_elementtuple(t::HomoNuclearTable) = (t.element_symbol, t.element_symbol)
_elementtuple(t::HeteroNuclearTable) = (t.first_element_symbol, t.second_element_symbol)

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

"""
    loaddir([T], dirpath::String)

Parse a directory at `dirpath` containing SKF files and return a
`SlaterKosterTable` derived from its contents.

See also [`load`](@ref SlaterKoster.load).

"""
loaddir(dirpath) = loaddir(Float64, dirpath)
function loaddir(T::Type, dirpath)
    data = Dict{Tuple{Symbol,Symbol},PrimitiveTable{T}}()
    for file in filter(f -> occursin(r"(skf|SKF)$", f), Base.readdir(dirpath))
        table = load(T, joinpath(dirpath, file))
        data[_elementtuple(table)] = table
    end
    return SlaterKosterTable(data)
end

"""
    load([T], filepath::String; [element_symbols::NTuple{2,String}])

Parse an SKF located at `filepath` and return a primitive table
derived from its contents.  The element symbols will be inferred from
the file name if it follows the standard `X-Y.skf` format, where `X`,
and `Y` are element symbols.  Alternatively, if a tuple of symbols is
provided to the keyword argument `element_symbols`, the element
symbols listed in the tuple will be used instead.  If `T` is supplied,
floating-point data will be parsed using the type `T` instead of the
default `Float64`.

See also [`loaddir`](@ref SlaterKoster.loaddir).

"""
load(filepath; element_symbols=(nothing, nothing)) = load(Float64, filepath, element_symbols=element_symbols)
function load(T::Type, filepath::String; element_symbols=(nothing, nothing))
    dirname, filename = splitdir(filepath)
    if element_symbols === (nothing, nothing)
        # Try to derive the element symbols from the file name by
        # matching file name patterns of known datasets
        m = match(r"^([A-Z][a-z]*)[-]*([A-Z][a-z]*)\.(skf|SKF)$", filename)
        m === nothing && error("could not determine element symbols from SKF file ($filename); please specify them manually using the keyword argument `element_symbols`")
        first_element_symbol, second_element_symbol, _ = m.captures
    else
        first_element_symbol, second_element_symbol = element_symbols
    end
    lines = Base.readlines(filepath)
    if first_element_symbol == second_element_symbol
        if startswith(lines[1], "@")
            return HomoNuclearTable(
                Symbol(first_element_symbol),
                _readextendedformat_homonuclear(T, lines)...)
        else
            return HomoNuclearTable(
                Symbol(first_element_symbol),
                _readsimpleformat_homonuclear(T, lines)...)
        end
    else
        if startswith(lines[1], "@")
            return HeteroNuclearTable(
                Symbol(first_element_symbol),
                Symbol(second_element_symbol),
                _readextendedformat_heteronuclear(T, lines)...)
        else
            return HeteroNuclearTable(
                Symbol(first_element_symbol),
                Symbol(second_element_symbol),
                _readsimpleformat_heteronuclear(T, lines)...)
        end
    end
end

function parseheader(T::Type, line)
    words = filter(word -> !all(isspace, word), split(line, x -> isspace(x) || x === ','))
    return parse(T, words[1]), parse(Int, words[2])
end

function parseline(T::Type, line)
    words = filter(word -> !all(isspace, word), split(line, x -> isspace(x) || x === ','))
    i = 1
    while i <= length(words)
        m = match(r"([0-9]+)\*([0-9\.]+)", words[i])
        if m !== nothing
            deleteat!(words, i)
            for _ in 1:parse(Int, m[1])
                insert!(words, i, m[2])
            end
        end
        i += 1
    end
    return parse.(T, words)
end

function _readsimpleformat_homonuclear(T, lines)
    grid_dist, n_grid_points = parseheader(T, lines[1])

    energy_table = DataFrame(
        Ed = T[], Ep = T[], Es = T[],
        SPE = T[],
        Ud = T[], Up = T[], Us = T[],
        fd = T[], fp = T[], fs = T[])
    push!(energy_table, parseline(T, lines[2]))

    mass_table = DataFrame(
        mass = T[],
        c2 = T[], c3 = T[], c4 = T[], c5 = T[],
        c6 = T[], c7 = T[], c8 = T[], c9 = T[],
        rcut = T[],
        d1 = T[], d2 = T[], d3 = T[], d4 = T[], d5 = T[],
        d6 = T[], d7 = T[], d8 = T[], d9 = T[], d10 = T[])
    push!(mass_table, parseline(T, lines[3]))

    integral_table = DataFrame(
        dist = T[],
        Hdd0 = T[], Hdd1 = T[], Hdd2 = T[], Hpd0 = T[], Hpd1 = T[], Hpp0 = T[],
        Hpp1 = T[], Hsd0 = T[], Hsp0 = T[], Hss0 = T[],
        Sdd0 = T[], Sdd1 = T[], Sdd2 = T[], Spd0 = T[], Spd1 = T[], Spp0 = T[],
        Spp1 = T[], Ssd0 = T[], Ssp0 = T[], Sss0 = T[])

    for (i, line) in enumerate(lines[4:4+n_grid_points-1])
        push!(integral_table, vcat([i*grid_dist], parseline(T, line)))
    end
    return grid_dist, n_grid_points, energy_table, mass_table, integral_table
end

function _readsimpleformat_heteronuclear(T, lines)
    grid_dist, n_grid_points = parseheader(T, lines[1])

    mass_table = DataFrame(
        mass = T[],
        c2 = T[], c3 = T[], c4 = T[], c5 = T[],
        c6 = T[], c7 = T[], c8 = T[], c9 = T[],
        rcut = T[],
        d1 = T[], d2 = T[], d3 = T[], d4 = T[], d5 = T[],
        d6 = T[], d7 = T[], d8 = T[], d9 = T[], d10 = T[])
    push!(mass_table, parseline(T, lines[2]))

    integral_table = DataFrame(
        dist = T[],
        Hdd0 = T[], Hdd1 = T[], Hdd2 = T[], Hpd0 = T[], Hpd1 = T[], Hpp0 = T[],
        Hpp1 = T[], Hsd0 = T[], Hsp0 = T[], Hss0 = T[],
        Sdd0 = T[], Sdd1 = T[], Sdd2 = T[], Spd0 = T[], Spd1 = T[], Spp0 = T[],
        Spp1 = T[], Ssd0 = T[], Ssp0 = T[], Sss0 = T[])

    for (i, line) in enumerate(lines[3:3+n_grid_points-1])
        push!(integral_table, vcat([i*grid_dist], parseline(T, line)))
    end
    return grid_dist, n_grid_points, mass_table, integral_table
end

function _readextendedformat_homonuclear(T, lines)
    grid_dist, n_grid_points = parseheader(T, lines[1])

    energy_table = DataFrame(
        Ef = T[], Ed = T[], Ep = T[], Es = T[],
        SPE = T[],
        Uf = T[], Ud = T[], Up = T[], Us = T[],
        ff = T[], fd = T[], fp = T[], fs = T[])
    push!(energy_table, parseline(T, lines[3]))

    mass_table = DataFrame(
        mass = T[],
        c2 = T[], c3 = T[], c4 = T[], c5 = T[],
        c6 = T[], c7 = T[], c8 = T[], c9 = T[],
        rcut = T[],
        d1 = T[], d2 = T[], d3 = T[], d4 = T[], d5 = T[],
        d6 = T[], d7 = T[], d8 = T[], d9 = T[], d10 = T[])
    push!(mass_table, parseline(T, lines[4]))

    integral_table = DataFrame(
        dist = T[],
        Hff0 = T[], Hff1 = T[], Hff2 = T[], Hff3 = T[],
        Hdf0 = T[], Hdf1 = T[], Hdf2 = T[], Hdd0 = T[],
        Hdd1 = T[], Hdd2 = T[], Hpf0 = T[], Hpf1 = T[],
        Hpd0 = T[], Hpd1 = T[], Hpp0 = T[], Hpp1 = T[],
        Hsf0 = T[], Hsd0 = T[], Hsp0 = T[], Hss0 = T[],
        Sdd1 = T[], Sdd2 = T[], Sff0 = T[], Sff1 = T[],
        Sff2 = T[], Sff3 = T[], Sdf0 = T[], Sdf1 = T[],
        Sdf2 = T[], Sdd0 = T[], Spf0 = T[], Spf1 = T[],
        Spd0 = T[], Spd1 = T[], Spp0 = T[], Spp1 = T[],
        Ssf0 = T[], Ssd0 = T[], Ssp0 = T[], Sss0 = T[])

    for (i, line) in enumerate(lines[5:5+n_grid_points-1])
        push!(integral_table, vcat([i*grid_dist], parseline(T, line)))
    end
    return grid_dist, n_grid_points, energy_table, mass_table, integral_table
end

function _readextendedformat_heteronuclear(T, lines)
    grid_dist, n_grid_points = parseheader(T, lines[1])

    mass_table = DataFrame(
        mass = T[],
        c2 = T[], c3 = T[], c4 = T[], c5 = T[],
        c6 = T[], c7 = T[], c8 = T[], c9 = T[],
        rcut = T[],
        d1 = T[], d2 = T[], d3 = T[], d4 = T[], d5 = T[],
        d6 = T[], d7 = T[], d8 = T[], d9 = T[], d10 = T[])
    push!(mass_table, parseline(T, lines[3]))

    integral_table = DataFrame(
        dist = T[],
        Hff0 = T[], Hff1 = T[], Hff2 = T[], Hff3 = T[],
        Hdf0 = T[], Hdf1 = T[], Hdf2 = T[], Hdd0 = T[],
        Hdd1 = T[], Hdd2 = T[], Hpf0 = T[], Hpf1 = T[],
        Hpd0 = T[], Hpd1 = T[], Hpp0 = T[], Hpp1 = T[],
        Hsf0 = T[], Hsd0 = T[], Hsp0 = T[], Hss0 = T[],
        Sdd1 = T[], Sdd2 = T[], Sff0 = T[], Sff1 = T[],
        Sff2 = T[], Sff3 = T[], Sdf0 = T[], Sdf1 = T[],
        Sdf2 = T[], Sdd0 = T[], Spf0 = T[], Spf1 = T[],
        Spd0 = T[], Spd1 = T[], Spp0 = T[], Spp1 = T[],
        Ssf0 = T[], Ssd0 = T[], Ssp0 = T[], Sss0 = T[])

    for (i, line) in enumerate(lines[4:4+n_grid_points-1])
        push!(integral_table, vcat([i*grid_dist], parseline(T, line)))
    end
    return grid_dist, n_grid_points, mass_table, integral_table
end