abstract type AbstractTree{NDIMS} <: AbstractContainer end

# Type traits to obtain dimension
@inline Base.ndims(::AbstractTree{NDIMS}) where NDIMS = NDIMS

# Auxiliary methods to allow semantic queries on the tree
# Check whether cell has parent cell
has_parent(t::AbstractTree, cell_id::Int) = t.parent_ids[cell_id] > 0

# Count number of children for a given cell
n_children(t::AbstractTree, cell_id::Int) = count(x -> (x > 0), @view t.child_ids[:, cell_id])

# Check whether cell has any child cell
has_children(t::AbstractTree, cell_id::Int) = n_children(t, cell_id) > 0

# Check whether cell is leaf cell
is_leaf(t::AbstractTree, cell_id::Int) = !has_children(t, cell_id)

# Check whether cell has specific child cell
has_child(t::AbstractTree, cell_id::Int, child::Int) = t.child_ids[child, cell_id] > 0

# Check if cell has a neighbor at the same refinement level in the given direction
has_neighbor(t::AbstractTree, cell_id::Int, direction::Int) = t.neighbor_ids[direction, cell_id] > 0

# Check if cell has a coarse neighbor, i.e., with one refinement level lower
function has_coarse_neighbor(t::AbstractTree, cell_id::Int, direction::Int)
  return has_parent(t, cell_id) && has_neighbor(t, t.parent_ids[cell_id], direction)
end

# Check if cell has any neighbor (same-level or lower-level)
function has_any_neighbor(t::AbstractTree, cell_id::Int, direction::Int)
  return has_neighbor(t, cell_id, direction) || has_coarse_neighbor(t, cell_id, direction)
end

# Check if cell is own cell, i.e., belongs to this MPI rank
is_own_cell(t::AbstractTree, cell_id) = true

# Return cell length for a given level
length_at_level(t::AbstractTree, level::Int) = t.length_level_0 / 2^level

# Return cell length for a given cell
length_at_cell(t::AbstractTree, cell_id::Int) = length_at_level(t, t.levels[cell_id])

# Return minimum level of any leaf cell
minimum_level(t::AbstractTree) = minimum(t.levels[leaf_cells(t)])

# Return maximum level of any leaf cell
maximum_level(t::AbstractTree) = maximum(t.levels[leaf_cells(t)])

# Check if tree is periodic
isperiodic(t::AbstractTree) = all(t.periodicity)
isperiodic(t::AbstractTree, dimension) = t.periodicity[dimension]

# Auxiliary methods for often-required calculations
# Number of potential child cells
n_children_per_cell(::AbstractTree{NDIMS}) where NDIMS = 2^NDIMS
n_children_per_cell(dims::Integer) = 2^dims

# Number of directions
#
# Directions are indicated by numbers from 1 to 2*ndims:
# 1 -> -x
# 2 -> +x
# 3 -> -y
# 4 -> +y
# 5 -> -z
# 6 -> +z
@inline n_directions(::AbstractTree{NDIMS}) where NDIMS = 2 * NDIMS
# TODO: Taal performance, 1:n_directions(tree) vs. Base.OneTo(n_directions(tree)) vs. SOneTo(n_directions(tree))
@inline eachdirection(tree::AbstractTree) = Base.OneTo(n_directions(tree))

# For a given direction, return its opposite direction
#
# dir -> opp
#  1  ->  2
#  2  ->  1
#  3  ->  4
#  4  ->  3
#  5  ->  6
#  6  ->  5
opposite_direction(direction::Int) = direction + 1 - 2 * ((direction + 1) % 2)

# For a given child position (from 1 to 8) and dimension (from 1 to 3),
# calculate a child cell's position relative to its parent cell.
#
# Essentially calculates the following
#         dim=1 dim=2 dim=3
# child     x     y     z
#   1       -     -     -
#   2       +     -     -
#   3       -     +     -
#   4       +     +     -
#   5       -     -     +
#   6       +     -     +
#   7       -     +     +
#   8       +     +     +
child_sign(child::Int, dim::Int) = 1 - 2 * (div(child + 2^(dim - 1) - 1, 2^(dim-1)) % 2)

# For each child position (1 to 8) and a given direction (from 1 to 6), return
# neighboring child position.
adjacent_child(child::Int, direction::Int) = [2 2 3 3 5 5;
                                              1 1 4 4 6 6;
                                              4 4 1 1 7 7;
                                              3 3 2 2 8 8;
                                              6 6 7 7 1 1;
                                              5 5 8 8 2 2;
                                              8 8 5 5 3 3;
                                              7 7 6 6 4 4][child, direction]

# For each child position (1 to 8) and a given direction (from 1 to 6), return
# if neighbor is a sibling
function has_sibling(child::Int, direction::Int)
  return (child_sign(child, div(direction + 1, 2)) * (-1)^(direction - 1)) > 0
end

# Obtain leaf cells that fulfill a given criterion.
#
# The function `f` is passed the cell id of each leaf cell
# as an argument.
function filter_leaf_cells(f, t::AbstractTree)
  filtered = Vector{Int}(undef, length(t))
  count = 0
  for cell_id in 1:length(t)
    if is_leaf(t, cell_id) && f(cell_id)
      count += 1
      filtered[count] = cell_id
    end
  end

  return filtered[1:count]
end

# Return an array with the ids of all leaf cells
leaf_cells(t::AbstractTree) = filter_leaf_cells((cell_id)->true, t)

# Return an array with the ids of all leaf cells for a given rank
leaf_cells_by_rank(t::AbstractTree, rank) = leaf_cells(t)

# Return an array with the ids of all local leaf cells
local_leaf_cells(t::AbstractTree) = leaf_cells(t)

# Count the number of leaf cells.
count_leaf_cells(t::AbstractTree) = length(leaf_cells(t))

# Store cell id in each cell to use for post-AMR analysis
function reset_original_cell_ids!(t::AbstractTree)
  t.original_cell_ids[1:length(t)] .= 1:length(t)
end

# Refine entire tree by one level
refine!(t::AbstractTree) = refine!(t, leaf_cells(t))

# Refine given cells and rebalance tree.
#
# Note 1: Rebalancing is iterative, i.e., neighboring cells are refined if
#         otherwise the 2:1 rule would be violated, which can cause more
#         refinements.
# Note 2: Rebalancing currently only considers *Cartesian* neighbors, not diagonal neighbors!
function refine!(t::AbstractTree, cell_ids)
  # Reset original cell ids such that each cell knows its current id
  reset_original_cell_ids!(t)

  # Refine all requested cells
  refined = refine_unbalanced!(t, cell_ids)
  refinement_count = length(refined)

  # Iteratively rebalance the tree until it does not change anymore
  while length(refined) > 0
    refined = rebalance!(t, refined)
    refinement_count += length(refined)
  end

  # Determine list of *original* cell ids that were refined
  # Note: original_cell_ids contains the cell_id *before* refinement. At
  # refinement, the refined cell's original_cell_ids value has its sign flipped
  # to easily find it now.
  @views refined_original_cells = (
      -t.original_cell_ids[1:length(t)][t.original_cell_ids[1:length(t)] .< 0])

  # Check if count of refinement cells matches information in original_cell_ids
  @assert refinement_count == length(refined_original_cells) (
      "Mismatch in number of refined cells")

  return refined_original_cells
end

# Refine all leaf cells with coordinates in a given rectangular box
function refine_box!(t::AbstractTree{NDIMS}, coordinates_min, coordinates_max) where NDIMS
  for dim in 1:NDIMS
    @assert coordinates_min[dim] < coordinates_max[dim] "Minimum coordinates are not minimum."
  end

  # Find all leaf cells within box
  cells = filter_leaf_cells(t) do cell_id
    return (all(coordinates_min .< t.coordinates[:, cell_id]) &&
            all(coordinates_max .> t.coordinates[:, cell_id]))
  end

  # Refine cells
  refine!(t, cells)
end

# Convenience method for 1D
function refine_box!(t::AbstractTree{1}, coordinates_min::Real, coordinates_max::Real)
  return refine_box!(t, [convert(Float64, coordinates_min)], [convert(Float64, coordinates_max)])
end

# For the given cell ids, check if neighbors need to be refined to restore a rebalanced tree.
#
# Note 1: Rebalancing currently only considers *Cartesian* neighbors, not diagonal neighbors!
# Note 2: The current algorithm assumes that a previous refinement step has
#         created level differences of at most 2. That is, before the previous
#         refinement step, the tree was balanced.
function rebalance!(t::AbstractTree, refined_cell_ids)
  # Create buffer for newly refined cells
  to_refine = zeros(Int, n_directions(t) * length(refined_cell_ids))
  count = 0

  # Iterate over cell ids that have previously been refined
  for cell_id in refined_cell_ids
    # Go over all potential neighbors of child cell
    for direction in eachdirection(t)
      # Continue if refined cell has a neighbor in that direction
      if has_neighbor(t, cell_id, direction)
        continue
      end

      # Continue if refined cell has no coarse neighbor, since that would
      # mean it there is no neighbor in that direction at all (domain
      # boundary)
      if !has_coarse_neighbor(t, cell_id, direction)
        continue
      end

      # Otherwise, the coarse neighbor exists and is not refined, thus it must
      # be marked for refinement
      coarse_neighbor_id = t.neighbor_ids[direction, t.parent_ids[cell_id]]
      count += 1
      to_refine[count] = coarse_neighbor_id
    end
  end

  # Finally, refine all marked cells...
  refined = refine_unbalanced!(t, unique(to_refine[1:count]))

  # ...and return list of refined cells
  return refined
end

# Refine given cells without rebalancing tree.
#
# Note: After a call to this method the tree may be unbalanced!
function refine_unbalanced!(t::AbstractTree, cell_ids) end

# Wrap single-cell refinements such that `sort(...)` does not complain
refine_unbalanced!(t::AbstractTree, cell_id::Int) = refine_unbalanced!(t, [cell_id])

# Coarsen entire tree by one level
function coarsen!(t::AbstractTree)
  # Special case: if there is only one cell (root), there is nothing to do
  if length(t) == 1
    return Int[]
  end

  # Get list of unique parent ids for all leaf cells
  parent_ids = unique(t.parent_ids[leaf_cells(t)])
  coarsen!(t, parent_ids)
end

# Coarsen given *parent* cells (= these cells must have children who are all
# leaf cells) while retaining a balanced tree.
#
# A cell to be coarsened might cause an unbalanced tree if the neighboring cell
# was already refined. Since it is generally not desired that cells are
# coarsened without specifically asking for it, these cells will then *not* be
# coarsened.
function coarsen!(t::AbstractTree, cell_ids::AbstractArray{Int})
  # Return early if array is empty
  if length(cell_ids) == 0
    return Int[]
  end

  # Reset original cell ids such that each cell knows its current id
  reset_original_cell_ids!(t)

  # To maximize the number of cells that may be coarsened, start with the cells at the highest level
  sorted_by_level = sort(cell_ids, by = i -> t.levels[i])

  # Keep track of number of cells that were actually coarsened
  n_coarsened = 0

  # Local function to adjust cell ids after some cells have been removed
  function adjust_cell_ids!(cell_ids, coarsened_cell_id, count)
    for (id, cell_id) in enumerate(cell_ids)
      if cell_id > coarsened_cell_id
        cell_ids[id] = cell_id - count
      end
    end
  end

  # Iterate backwards over cells to coarsen
  while true
    # Retrieve next cell or quit
    if length(sorted_by_level) > 0
      coarse_cell_id = pop!(sorted_by_level)
    else
      break
    end

    # Ensure that cell has children (violation is an error)
    if !has_children(t, coarse_cell_id)
      error("cell is leaf and cannot be coarsened to: $coarse_cell_id")
    end

    # Ensure that all child cells are leaf cells (violation is an error)
    for child in 1:n_children_per_cell(t)
      if has_child(t, coarse_cell_id, child)
        if !is_leaf(t, t.child_ids[child, coarse_cell_id])
          error("cell $coarse_cell_id has child cell at position $child that is not a leaf cell")
        end
      end
    end

    # Check if coarse cell has refined neighbors that would prevent coarsening
    skip = false
    # Iterate over all children (which are to be removed)
    for child in 1:n_children_per_cell(t)
      # Continue if child does not exist
      if !has_child(t, coarse_cell_id, child)
        continue
      end
      child_id = t.child_ids[child, coarse_cell_id]

      # Go over all neighbors of child cell. If it has a neighbor that is *not*
      # a sibling and that is not a leaf cell, we cannot coarsen its parent
      # without creating an unbalanced tree.
      for direction in eachdirection(t)
        # Continue if neighbor would be a sibling
        if has_sibling(child, direction)
          continue
        end

        # Continue if child cell has no neighbor in that direction
        if !has_neighbor(t, child_id, direction)
          continue
        end
        neighbor_id = t.neighbor_ids[direction, child_id]

        if !has_children(t, neighbor_id)
          continue
        end

        # If neighbor is not a sibling, is existing, and has children, do not coarsen
        skip = true
        break
      end
    end
    # Skip if a neighboring cell prevents coarsening
    if skip
      continue
    end

    # Flip sign of cell to be coarsened to such that we can easily find it
    t.original_cell_ids[coarse_cell_id] = -t.original_cell_ids[coarse_cell_id]

    # If a coarse cell has children that are all leaf cells, they must follow
    # immediately due to depth-first ordering of the tree
    count = n_children(t, coarse_cell_id)
    @assert count == n_children_per_cell(t) "cell $coarse_cell_id does not have all child cells"
    remove_shift!(t, coarse_cell_id + 1, coarse_cell_id + count)

    # Take into account shifts in tree that alters cell ids
    adjust_cell_ids!(sorted_by_level, coarse_cell_id, count)

    # Keep track of number of coarsened cells
    n_coarsened += 1
  end

  # Determine list of *original* cell ids that were coarsened to
  # Note: original_cell_ids contains the cell_id *before* coarsening. At
  # coarsening, the coarsened parent cell's original_cell_ids value has its sign flipped
  # to easily find it now.
  @views coarsened_original_cells = (
      -t.original_cell_ids[1:length(t)][t.original_cell_ids[1:length(t)] .< 0])

  # Check if count of coarsened cells matches information in original_cell_ids
  @assert n_coarsened == length(coarsened_original_cells) (
      "Mismatch in number of coarsened cells")

  return coarsened_original_cells
end

# Wrap single-cell coarsening such that `sort(...)` does not complain
coarsen!(t::AbstractTree, cell_id::Int) = coarsen!(t::AbstractTree, [cell_id])

# Coarsen all viable parent cells with coordinates in a given rectangular box
function coarsen_box!(t::AbstractTree{NDIMS}, coordinates_min::AbstractArray{Float64},
                     coordinates_max::AbstractArray{Float64}) where NDIMS
  for dim in 1:NDIMS
    @assert coordinates_min[dim] < coordinates_max[dim] "Minimum coordinates are not minimum."
  end

  # Find all leaf cells within box
  leaves = filter_leaf_cells(t) do cell_id
    return (all(coordinates_min .< t.coordinates[:, cell_id]) &&
            all(coordinates_max .> t.coordinates[:, cell_id]))
  end

  # Get list of unique parent ids for all leaf cells
  parent_ids = unique(t.parent_ids[leaves])

  # Filter parent ids to be within box
  parents = filter(parent_ids) do cell_id
    return (all(coordinates_min .< t.coordinates[:, cell_id]) &&
            all(coordinates_max .> t.coordinates[:, cell_id]))
  end

  # Coarsen cells
  coarsen!(t, parents)
end

# Convenience method for 1D
function coarsen_box!(t::AbstractTree{1}, coordinates_min::Real, coordinates_max::Real)
  return coarsen_box!(t, [convert(Float64, coordinates_min)], [convert(Float64, coordinates_max)])
end

# Return coordinates of a child cell based on its relative position to the parent.
function child_coordinates(::AbstractTree{NDIMS}, parent_coordinates, parent_length::Number, child::Int) where NDIMS
  # Calculate length of child cells and set up data structure
  child_length = parent_length / 2
  coordinates = MVector{NDIMS, Float64}(undef)

  # For each dimension, calculate coordinate as parent coordinate + relative position x length/2
  for d in 1:NDIMS
    coordinates[d] = parent_coordinates[d] + child_sign(child, d) * child_length / 2
  end

  return coordinates
end

# Reset range of cells to values that are prone to cause errors as soon as they are used.
#
# Rationale: If an invalid cell is accidentally used, we want to know it as soon as possible.
function invalidate!(t::AbstractTree, first::Int, last::Int) end
invalidate!(t::AbstractTree, id::Int) = invalidate!(t, id, id)
invalidate!(t::AbstractTree) = invalidate!(t, 1, length(t))

# Delete connectivity with parents/children/neighbors before cells are erased
function delete_connectivity!(t::AbstractTree, first::Int, last::Int)
  @assert first > 0
  @assert first <= last
  @assert last <= t.capacity + 1

  # Iterate over all cells
  for cell_id in first:last
    # Delete connectivity from parent cell
    if has_parent(t, cell_id)
      parent_id = t.parent_ids[cell_id]
      for child in 1:n_children_per_cell(t)
        if t.child_ids[child, parent_id] == cell_id
          t.child_ids[child, parent_id] = 0
          break
        end
      end
    end

    # Delete connectivity from child cells
    for child in 1:n_children_per_cell(t)
      if has_child(t, cell_id, child)
        t.parent_ids[t._child_ids[child, cell_id]] = 0
      end
    end

    # Delete connectivity from neighboring cells
    for direction in eachdirection(t)
      if has_neighbor(t, cell_id, direction)
        t.neighbor_ids[opposite_direction(direction), t.neighbor_ids[direction, cell_id]] = 0
      end
    end
  end
end

# Move connectivity with parents/children/neighbors after cells have been moved
function move_connectivity!(t::AbstractTree, first::Int, last::Int, destination::Int)
  @assert first > 0
  @assert first <= last
  @assert last <= t.capacity + 1
  @assert destination > 0
  @assert destination <= t.capacity + 1

  # Strategy
  # 1) Loop over moved cells (at target location)
  # 2) Check if parent/children/neighbors connections are to a cell that was moved
  #    a) if cell was moved: apply offset to current cell
  #    b) if cell was not moved: go to connected cell and update connectivity there

  offset = destination - first
  has_moved(n) = (first <= n <= last)

  for source in first:last
    target = source + offset

    # Update parent
    if has_parent(t, target)
      # Get parent cell
      parent_id = t.parent_ids[target]
      if has_moved(parent_id)
        # If parent itself was moved, just update parent id accordingly
        t.parent_ids[target] += offset
      else
        # If parent was not moved, update its corresponding child id
        for child in 1:n_children_per_cell(t)
          if t.child_ids[child, parent_id] == source
            t.child_ids[child, parent_id] = target
          end
        end
      end
    end

    # Update children
    for child in 1:n_children_per_cell(t)
      if has_child(t, target, child)
        # Get child cell
        child_id = t.child_ids[child, target]
        if has_moved(child_id)
          # If child itself was moved, just update child id accordingly
          t.child_ids[child, target] += offset
        else
          # If child was not moved, update its parent id
          t.parent_ids[child_id] = target
        end
      end
    end

    # Update neighbors
    for direction in eachdirection(t)
      if has_neighbor(t, target, direction)
        # Get neighbor cell
        neighbor_id = t.neighbor_ids[direction, target]
        if has_moved(neighbor_id)
          # If neighbor itself was moved, just update neighbor id accordingly
          t.neighbor_ids[direction, target] += offset
        else
          # If neighbor was not moved, update its opposing neighbor id
          t.neighbor_ids[opposite_direction(direction), neighbor_id] = target
        end
      end
    end
  end
end

# Raw copy operation for ranges of cells.
#
# This method is used by the higher-level copy operations for AbstractContainer
function raw_copy!(target::AbstractTree, source::AbstractTree, first::Int, last::Int, destination::Int) end

# Reset data structures by recreating all internal storage containers and invalidating all elements
function reset_data_structures!(t::AbstractTree{NDIMS}) where NDIMS end


" Outputs the unstructured mesh with the solution fields "
type VTKOutput <: Output
end

"""
    Create the lists of Cells and Points for Outputting VTK

    Returns a tuple of `cells, points`
"""
function createVTKGeometry(mesh::Mesh)
    # Create the cells
    elems = mesh.elements
    n_elem = length(elems)

    cells = Array{MeshCell}(n_elem)

    for i in 1:n_elem
        elem = mesh.elements[i]
        cells[i] = MeshCell(VTKCellTypes.VTK_QUAD, [ node.id for node in elem.nodes ])
    end

    # Create the Points array
    nodes = mesh.nodes
    n_nodes = length(nodes)

    points = Matrix{Float64}((2, n_nodes))

    for i in 1:n_nodes
        node = mesh.nodes[i]
        points[1,i] = node.coords[1]
        points[2,i] = node.coords[2]
    end

    return cells, points
end

"""
    Creates the point data suitable for outputing to VTK

    returns a Dict{VariableName, Array{Float64}}
"""
function createPointData(solver::Solver, solution::Array)
    # One entry per variable
    point_data = Dict{String, Array{Float64}}()

    sys = solver.system
    mesh = sys.mesh
    nodes = mesh.nodes
    n_nodes = length(nodes)

    for var in sys.variables
        id = var.id
        data = Array{Float64}(n_nodes)

        for i in 1:n_nodes
            data[i] = solution[mesh.nodes[i].dofs[id]]
        end

        point_data[var.name] = data
    end

    return point_data
end

" Write out a VTK file "
function output(out::VTKOutput, solver::Solver, filebase::String)
    startLog(main_perf_log, "output(VTKOutput)")

    serialized_solution = serializeToZero(solver.solution)

    if MPI.Comm_rank(MPI.COMM_WORLD) == 0
        cells, points = createVTKGeometry(solver.system.mesh)

        vtkfile = vtk_grid(filebase, points, cells)

        point_data = createPointData(solver, serialized_solution)

        for name_data in point_data
            vtk_point_data(vtkfile, name_data[2], name_data[1])
        end

        outfiles = vtk_save(vtkfile)
    end

    stopLog(main_perf_log, "output(VTKOutput)")
end


using BenchmarkTools
using Arrow
using CategoricalArrays

data = convert(Vector{UInt8}, reinterpret(UInt8, [2,3,5,7]))
data = vcat(0x0f, data)

# v = NullablePrimitive([2,missing,5,7])

v = NullablePrimitive{Int64}(data, 1, 2, 4)


"""
    gcmarray{T, N}

gcmarray data structure. Available constructors:

```
gcmarray{T,N}(grid::gcmgrid,f::Array{Array{T,2},N},
         fSize::Array{NTuple{N, Int}},fIndex::Array{Int,1})

gcmarray(grid::gcmgrid,f::Array{Array{T,2},N}) where {T,N}
gcmarray(grid::gcmgrid,f::Array{Array{T,N},1}) where {T,N}

gcmarray(grid::gcmgrid,fSize::Array{NTuple{N, Int}},fIndex::Array{Int,1})
gcmarray(<same as above>,n3::Int)
gcmarray(<same as above>,n3::Int,n4::Int)

gcmarray(grid::gcmgrid)
gcmarray(grid::gcmgrid,::Type{T})
gcmarray(grid::gcmgrid,::Type{T},n3::Int)
gcmarray(grid::gcmgrid,::Type{T},n3::Int,n4::Int)
```
"""
struct gcmarray{T, N} <: AbstractMeshArray{T, N}
   grid::gcmgrid
   f::Array{Array{T,2},N}
   fSize::Array{NTuple{2, Int}}
   fIndex::Array{Int,1}
end

function gcmarray(grid::gcmgrid,f::Array{Array{T,2},N}) where {T, N}
  gcmarray{T,N}(grid,f,grid.fSize,collect(1:grid.nFaces))
end

function gcmarray(grid::gcmgrid,f::Array{Array{T,N},1}) where {T, N}
  nFaces=grid.nFaces
  if N>2
    n3=size(f[1],3); n4=size(f[1],4);
    g=Array{Array{T,2},3}(undef,nFaces,n3,n4)
    for I in eachindex(view(g,1:nFaces,1:n3,1:n4))
      g[I]=view(f[I[1]],:,:,I[2],I[3])
    end
    n4==1 ? g=dropdims(g,dims=3) : nothing
    gcmarray{T,ndims(g)}(grid,g,grid.fSize,collect(1:nFaces))
  else
    gcmarray{T,1}(grid,f,grid.fSize,collect(1:nFaces))
  end
end

function gcmarray(grid::gcmgrid,::Type{T},
        fSize::Union{Array{NTuple{2, Int},1},NTuple{2, Int}},
        fIndex::Union{Array{Int,1},Int}) where {T}
  nFaces=length(fIndex)
  f=Array{Array{T,2},1}(undef,nFaces)
  isa(fSize,NTuple) ? fSize=[fSize] : nothing
  isa(fIndex,Int) ? fIndex=[fIndex] : nothing
  for a=1:nFaces
    f[a]=Array{T}(undef,fSize[a])
  end
  gcmarray{T,1}(grid,f,fSize,fIndex)
end

function gcmarray(grid::gcmgrid,::Type{T},
        fSize::Union{Array{NTuple{2, Int},1},NTuple{2, Int}},
        fIndex::Union{Array{Int,1},Int},n3::Int) where {T}
  nFaces=length(fIndex)
  f=Array{Array{T,2},2}(undef,nFaces,n3)
  isa(fSize,NTuple) ? fSize=[fSize] : nothing
  isa(fIndex,Int) ? fIndex=[fIndex] : nothing
  for a=1:nFaces; for i3=1:n3;
    f[a,i3]=Array{T}(undef,fSize[a]...)
  end; end;
  gcmarray{T,2}(grid,f,fSize,fIndex)
end

function gcmarray(grid::gcmgrid,::Type{T},
        fSize::Union{Array{NTuple{2, Int},1},NTuple{2, Int}},
        fIndex::Union{Array{Int,1},Int},n3::Int,n4::Int) where {T}
  nFaces=length(fIndex)
  f=Array{Array{T,2},3}(undef,nFaces,n3,n4)
  isa(fSize,NTuple) ? fSize=[fSize] : nothing
  isa(fIndex,Int) ? fIndex=[fIndex] : nothing
  for a=1:nFaces; for i4=1:n4; for i3=1:n3;
    f[a,i3,i4]=Array{T}(undef,fSize[a]...)
  end; end; end;
  gcmarray{T,3}(grid,f,fSize,fIndex)
end

# +
function gcmarray(grid::gcmgrid)
  nFaces=grid.nFaces
  fSize=grid.fSize
  fIndex=collect(1:grid.nFaces)
  T=grid.ioPrec
  gcmarray(grid,T,fSize,fIndex)
end

function gcmarray(grid::gcmgrid,::Type{T}) where {T}
  nFaces=grid.nFaces
  fSize=grid.fSize
  fIndex=collect(1:grid.nFaces)
  gcmarray(grid,T,fSize,fIndex)
end

function gcmarray(grid::gcmgrid,::Type{T},n3::Int) where {T}
  nFaces=grid.nFaces
  fSize=grid.fSize
  fIndex=collect(1:grid.nFaces)
  gcmarray(grid,T,fSize,fIndex,n3)
end

function gcmarray(grid::gcmgrid,::Type{T},n3::Int,n4::Int) where {T}
  nFaces=grid.nFaces
  fSize=grid.fSize
  fIndex=collect(1:grid.nFaces)
  gcmarray(grid,T,fSize,fIndex,n3,n4)
end

# -

# # Interface Methods

# +
Base.size(A::gcmarray) = size(A.f)
Base.size(A::gcmarray, dim::Integer) = size(A)[dim]

# +
function Base.getindex(A::gcmarray{T, N}, I::Vararg{Union{Int,Array{Int},AbstractUnitRange,Colon}, N}) where {T,N}
  J=1:length(A.fIndex)
  !isa(I[1],Colon) ? J=J[I[1]] : nothing
  nFaces=length(J)

  tmpf=A.f[I...]
  if isa(tmpf,Array{eltype(A),2})
    tmp=tmpf
  else
    n3=Int(length(tmpf)/nFaces)
    K=(A.grid,eltype(A),A.fSize[J],A.fIndex[J])
    n3>1 ? tmp=gcmarray(K...,n3) : tmp=gcmarray(K...)
    for I in eachindex(tmpf); tmp.f[I] = view(tmpf[I],:,:); end
  end

  return tmp
end

"""
    getindexetc(A::gcmarray, I::Vararg{_}) where {T,N}

Same as getindex but also returns the face size and index
"""
function getindexetc(A::gcmarray{T, N}, I::Vararg{Union{Int,Array{Int},AbstractUnitRange,Colon}, N}) where {T,N}
    f=A[I...]
    fSize=A.fSize[I[1]]
    fIndex=A.fIndex[I[1]]
    return f,fSize,fIndex
end
# -

function Base.setindex!(A::gcmarray{T, N}, v, I::Vararg{Int, N}) where {T,N}
  return (A.f[I...] = v)
end

function Base.view(A::gcmarray{T, N}, I::Vararg{Union{Int,AbstractUnitRange,Colon}, N}) where {T,N}
  J=1:length(A.fIndex)
  !isa(I[1],Colon) ? J=J[I[1]] : nothing
  nFaces=length(J)

  tmpf=view(A.f,I...)
  n3=Int(length(tmpf)/nFaces) #length(tmpf)>nFaces ? n3=Int(length(tmpf)/nFaces) : n3=1

  K=(A.grid,eltype(A),A.fSize[J],A.fIndex[J])
  n3>1 ? tmp=gcmarray(K...,n3) : tmp=gcmarray(K...)
  for I in eachindex(tmpf); tmp.f[I] = view(tmpf[I],:,:); end

  return tmp
end

# ### Custom pretty-printing, similar, and broadcast

function Base.show(io::IO, z::gcmarray{T, N}) where {T,N}
    printstyled(io, " gcmarray \n",color=:normal)
    printstyled(io, "  grid type   = ",color=:normal)
    printstyled(io, "$(z.grid.class)\n",color=:blue)
    printstyled(io, "  data type   = ",color=:normal)
    printstyled(io, "$(eltype(z))\n",color=:blue)
    printstyled(io, "  tile array  = ",color=:normal)
    printstyled(io, "$(size(z))\n",color=:blue)
    printstyled(io, "  tile sizes  = ",color=:normal)
    printstyled(io, "$(size(z[1]))\n",color=:blue)
    for iFace=2:length(z.fIndex)
      printstyled(io, "                ",color=:normal)
      printstyled(io, "$(size(z[iFace]))\n",color=:blue)
    end
  return
end

function Base.similar(A::gcmarray)
    if ndims(A)==1
        B=gcmarray(A.grid,eltype(A),A.fSize,A.fIndex)
    else
        B=gcmarray(A.grid,eltype(A),A.fSize,A.fIndex,size(A,2))
    end
    return B
end

# ### Customize broadcasting

Base.BroadcastStyle(::Type{<:gcmarray}) = Broadcast.ArrayStyle{gcmarray}()

function Base.similar(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{gcmarray}}, ::Type{ElType}) where ElType
  # Scan the inputs for the gcmarray:
  A = find_gcmarray(bc)
  # Create the gcmarray output:
  if ndims(A)==1
        B=gcmarray(A.grid,ElType,A.fSize,A.fIndex)
  else
        B=gcmarray(A.grid,ElType,A.fSize,A.fIndex,size(A,2))
  end
  return B
end

find_gcmarray(bc::Base.Broadcast.Broadcasted) = find_gcmarray(bc.args)
find_gcmarray(args::Tuple) = find_gcmarray(find_gcmarray(args[1]), Base.tail(args))
find_gcmarray(x) = x
find_gcmarray(a::gcmarray, rest) = a
find_gcmarray(::Any, rest) = find_gcmarray(rest)

####

import Base: copyto!

# Specialize this method if all you want to do is specialize on typeof(dest)
@inline function copyto!(dest::MeshArrays.gcmarray, bc::Broadcast.Broadcasted{Nothing})
    axes(dest) == axes(bc) || throwdm(axes(dest), axes(bc))
    # Performance optimization: broadcast!(identity, dest, A) is equivalent to copyto!(dest, A) if indices match
    if bc.f === identity && bc.args isa Tuple{AbstractArray} # only a single input argument to broadcast!
        A = bc.args[1]
        if axes(dest) == axes(A)
            return copyto!(dest, A)
        end
    end
    bc′ = Broadcast.preprocess(dest, bc)
    @simd for I in eachindex(bc′)
        #@inbounds dest[I] = bc′[I]
        @inbounds dest[I] = gcmarray_getindex_evalf(bc′,I)
    end
    return dest
end

function gcmarray_getindex_evalf(bc,I)
  @boundscheck checkbounds(bc, I)
  args = Broadcast._getindex(bc.args, I)
  return bc.f.(args...)
end

###

"""
    nFacesEtc(a::gcmarray)

Return nFaces, n3 (1 in 2D case; >1 otherwise)
"""
function nFacesEtc(a::gcmarray)
  nFaces=length(a.fIndex)
  ndims(a.f)>1 ? n3=size(a.f,2) : n3=1
  ndims(a.f)>2 ? n4=size(a.f,3) : n4=1
  return nFaces, n3, n4
end


export LinearBackwardEulerSolver, AbstractBackwardEulerSolver
export NonLinearBackwardEulerSolver

abstract type AbstractImplicitOperator end

"""
    op! = EulerOperator(f!, ϵ)

Construct a linear operator which performs an explicit Euler step ``Q + α
f(Q)``, where `f!` and `op!` both operate inplace, with extra arguments passed
through, i.e.
```
op!(LQ, Q, args...)
```
is equivalent to
```
f!(dQ, Q, args...)
LQ .= Q .+ ϵ .* dQ
```
"""
mutable struct EulerOperator{F, FT} <: AbstractImplicitOperator
    f!::F
    ϵ::FT
end

function (op::EulerOperator)(LQ, Q, args...)
    op.f!(LQ, Q, args..., increment = false)
    @. LQ = Q + op.ϵ * LQ
end

"""
    AbstractBackwardEulerSolver

An abstract backward Euler method
"""
abstract type AbstractBackwardEulerSolver end

"""
    (be::AbstractBackwardEulerSolver)(Q, Qhat, α, param, time)

Each concrete implementations of `AbstractBackwardEulerSolver` should provide a
callable version which solves the following system for `Q`
```
    Q = Qhat + α f(Q, param, time)
```
where `f` is the ODE tendency function, `param` are the ODE parameters, and
`time` is the current ODE time. The arguments `Q` should be modified in place
and should not be assumed to be initialized to any value.
"""
(be::AbstractBackwardEulerSolver)(Q, Qhat, α, p, t) =
    throw(MethodError(be, (Q, Qhat, α, p, t)))

"""
    Δt_is_adjustable(::AbstractBackwardEulerSolver)

Return `Bool` for whether this backward Euler solver can be updated. default is
`false`.
"""
Δt_is_adjustable(::AbstractBackwardEulerSolver) = false

"""
    update_backward_Euler_solver!(::AbstractBackwardEulerSolver, α)

Update the given backward Euler solver for the parameter `α`; see
['AbstractBackwardEulerSolver'](@ref). Default behavior is no change to the
solver.
"""
update_backward_Euler_solver!(::AbstractBackwardEulerSolver, Q, α) = nothing

"""
    setup_backward_Euler_solver(solver, Q, α, tendency!)

Returns a concrete implementation of an `AbstractBackwardEulerSolver` that will
solve for `Q` in systems of the form of
```
    Q = Qhat + α f(Q, param, time)
```
where `tendency!` is the in-place tendency function. Not the array `Q` is just
passed in for type information, e.g., `Q` the same `Q` will not be used for all
calls to the solver.
"""
setup_backward_Euler_solver(solver::AbstractBackwardEulerSolver, _...) = solver

"""
    LinearBackwardEulerSolver(::AbstractSystemSolver; isadjustable = false)

Helper type for specifying building a backward Euler solver with a linear
solver.  If `isadjustable == true` then the solver can be updated with a new
time step size.
"""
struct LinearBackwardEulerSolver{LS}
    solver::LS
    isadjustable::Bool
    preconditioner_update_freq::Int
    LinearBackwardEulerSolver(
        solver;
        isadjustable = false,
        preconditioner_update_freq = -1,
    ) = new{typeof(solver)}(solver, isadjustable, preconditioner_update_freq)
end

"""
    LinBESolver

Concrete implementation of an `AbstractBackwardEulerSolver` to use linear
solvers of type `AbstractSystemSolver`. See helper type
[`LinearBackwardEulerSolver`](@ref)
```
    Q = Qhat + α f(Q, param, time)
```
"""
mutable struct LinBESolver{FT, F, LS} <: AbstractBackwardEulerSolver
    α::FT
    f_imp!::F
    solver::LS
    isadjustable::Bool
    # used only for iterative solver
    preconditioner::AbstractPreconditioner
    # used only for direct solver
    factors::Any
end

Δt_is_adjustable(lin::LinBESolver) = lin.isadjustable

function setup_backward_Euler_solver(
    lin::LinearBackwardEulerSolver,
    Q,
    α,
    f_imp!,
)
    FT = eltype(α)
    rhs! = EulerOperator(f_imp!, -α)

    factors = prefactorize(rhs!, lin.solver, Q, nothing, FT(NaN))

    # when direct solver is applied preconditioner_update_freq <= 0
    @assert(
        typeof(lin.solver) <: AbstractIterativeSystemSolver ||
        lin.preconditioner_update_freq <= 0
    )

    preconditioner_update_freq = lin.preconditioner_update_freq
    # construct an empty preconditioner
    preconditioner = (
        preconditioner_update_freq > 0 ?
        ColumnwiseLUPreconditioner(f_imp!, Q, preconditioner_update_freq) :
        NoPreconditioner()
    )

    LinBESolver(
        α,
        f_imp!,
        lin.solver,
        lin.isadjustable,
        preconditioner,
        factors,
    )
end

function update_backward_Euler_solver!(lin::LinBESolver, Q, α)
    lin.α = α
    FT = eltype(Q)
    # for direct solver, update factors
    # for iterative solver, set factors to Nothing (TODO optimize)
    lin.factors = prefactorize(
        EulerOperator(lin.f_imp!, -α),
        lin.solver,
        Q,
        nothing,
        FT(NaN),
    )
end

function (lin::LinBESolver)(Q, Qhat, α, p, t)

    # If α is not the same as the already assembled operator
    # with its previous value of α (lin.α), then we need to
    # first check that the solver CAN be rebuilt (lin.isadjustable)
    # followed by recalling the set up routine by updating the
    # coefficient with the new version of α
    if lin.α != α
        @assert lin.isadjustable
        update_backward_Euler_solver!(lin, Q, α)
    end

    rhs! = EulerOperator(lin.f_imp!, -α)

    if typeof(lin.solver) <: AbstractIterativeSystemSolver
        FT = eltype(α)
        preconditioner_update!(rhs!, rhs!.f!, lin.preconditioner, p, t)
        linearsolve!(rhs!, lin.preconditioner, lin.solver, Q, Qhat, p, t)
        preconditioner_counter_update!(lin.preconditioner)
    else
        linearsolve!(rhs!, lin.factors, lin.solver, Q, Qhat, p, t)
    end
end

"""
    struct NonLinearBackwardEulerSolver{NLS}
        nlsolver::NLS
        isadjustable::Bool
        preconditioner_update_freq::Int64
    end

Helper type for specifying building a nonlinear backward Euler solver with a nonlinear
solver.

# Arguments
- `nlsolver`: iterative nonlinear solver, i.e., JacobianFreeNewtonKrylovSolver
- `isadjustable`: TODO not used, might use for updating preconditioner
- `preconditioner_update_freq`:  relavent to Jacobian free -1: no preconditioner;
                             positive number, update every freq times
"""
struct NonLinearBackwardEulerSolver{NLS}
    nlsolver::NLS
    isadjustable::Bool
    # preconditioner_update_freq, -1: no preconditioner;
    # positive number, update every freq times
    preconditioner_update_freq::Int
    function NonLinearBackwardEulerSolver(
        nlsolver;
        isadjustable = false,
        preconditioner_update_freq = -1,
    )
        NLS = typeof(nlsolver)
        return new{NLS}(nlsolver, isadjustable, preconditioner_update_freq)
    end
end

"""
    LinBESolver

Concrete implementation of an `AbstractBackwardEulerSolver` to use nonlinear
solvers of type `NLS`. See helper type
[`NonLinearBackwardEulerSolver`](@ref)
```
    Q = Qhat + α f_imp(Q, param, time)
```
"""
mutable struct NonLinBESolver{FT, F, NLS} <: AbstractBackwardEulerSolver
    # Solve Q - α f_imp(Q) = Qrhs
    α::FT
    # implcit operator
    f_imp!::F
    # jacobian action, which approximates drhs!/dQ⋅ΔQ , here rhs!(Q) = Q - α f_imp(Q)
    jvp!::JacobianAction
    # nonlinear solver
    nlsolver::NLS
    # whether adjust the time step or not
    isadjustable::Bool
    # preconditioner, approximation of drhs!/dQ
    preconditioner::AbstractPreconditioner

end

Δt_is_adjustable(nlsolver::NonLinBESolver) = nlsolver.isadjustable

"""
    setup_backward_Euler_solver(solver::NonLinearBackwardEulerSolver, Q, α, tendency!)

Returns a concrete implementation of an `AbstractBackwardEulerSolver` that will
solve for `Q` in nonlinear systems of the form of
```
    Q = Qhat + α f(Q, param, time)
```
Create an empty JacobianAction

Create an empty preconditioner if preconditioner_update_freq > 0
"""
function setup_backward_Euler_solver(
    nlbesolver::NonLinearBackwardEulerSolver,
    Q,
    α,
    f_imp!,
)
    # Create an empty JacobianAction (without operator)
    jvp! = JacobianAction(nothing, Q, nlbesolver.nlsolver.ϵ)

    # Create an empty preconditioner if preconditioner_update_freq > 0
    preconditioner_update_freq = nlbesolver.preconditioner_update_freq
    # construct an empty preconditioner
    preconditioner = (
        preconditioner_update_freq > 0 ?
        ColumnwiseLUPreconditioner(f_imp!, Q, preconditioner_update_freq) :
        NoPreconditioner()
    )
    NonLinBESolver(
        α,
        f_imp!,
        jvp!,
        nlbesolver.nlsolver,
        nlbesolver.isadjustable,
        preconditioner,
    )
end

"""
Nonlinear solve

Update rhs! with α

Update the rhs! in the jacobian action jvp!
"""
function (nlbesolver::NonLinBESolver)(Q, Qhat, α, p, t)

    rhs! = EulerOperator(nlbesolver.f_imp!, -α)

    nlbesolver.jvp!.rhs! = rhs!

    nonlinearsolve!(
        rhs!,
        nlbesolver.jvp!,
        nlbesolver.preconditioner,
        nlbesolver.nlsolver,
        Q,
        Qhat,
        p,
        t;
        max_newton_iters = nlbesolver.nlsolver.M,
    )

end


############################################################################################
# Nested

############################################################################################
@testset "Nested - AbstractArray" begin
    for output in outputtypes
        for flattentype in flattentypes
            flatdefault = FlattenDefault(; output = output, flattentype = flattentype)
            val = [1., [2, 3], [4. 5. ; 6. 7.], 8., [9., 10.]]
            ReConstructor(val)
            reconstruct = ReConstructor(flatdefault, val)
# Flatten
            x_flat = flatten(reconstruct, val)
            @test x_flat isa AbstractVector
            @test eltype(x_flat) == output
# Flatten AD
            x_flatAD = flattenAD(reconstruct, val)
            @test x_flatAD isa AbstractVector
            @test eltype(x_flatAD) == Float64 #Always most common type
# Unflatten
            x_unflat = unflatten(reconstruct, x_flat)
            @test x_unflat isa typeof(val)
# Unflatten AD
            x_unflatAD = unflattenAD(reconstruct, x_flat)
#            @test x_unflatAD isa eltype(x_flat)
            x_unflatAD2 = unflattenAD(reconstruct, x_flatAD)
#            @test x_unflatAD2 isa eltype(x_flatAD)
################################################################################
# Constraint ~ Make Distribution as standard constraints are separately tested
            val = [1., [2, 3], [4. 5. ; 6. 7.], 8., [9., 10.]]
            constraint = [Normal(), MvNormal([1., 1.]), Fixed(), Gamma(), Unconstrained()]
            ReConstructor(constraint, val)
            reconstruct = ReConstructor(flatdefault, constraint, val)
# Flatten
            x_flat = flatten(reconstruct, val)
            @test x_flat isa AbstractVector
            @test eltype(x_flat) == output
# Flatten AD
            x_flatAD = flattenAD(reconstruct, val)
            @test x_flatAD isa AbstractVector
            @test eltype(x_flatAD) == Float64
# Unflatten
            x_unflat = unflatten(reconstruct, x_flat)
            @test x_unflat isa typeof(val)
# Unflatten AD
            x_unflatAD = unflattenAD(reconstruct, x_flat)
#            @test x_unflatAD isa eltype(x_flat)
            x_unflatAD2 = unflattenAD(reconstruct, x_flatAD)
#            @test x_unflatAD2 isa eltype(x_flatAD)
################################################################################
# Transforms
            _transform, _inversetransform = construct_transform(constraint, val)
            x_transformed = unconstrain(_transform, val)
            x_inversetransformed = constrain(_inversetransform, x_transformed)

            transformer = TransformConstructor(constraint, val)
            val_uncon = unconstrain(transformer, val)
            val_con = constrain(transformer, val_uncon)

            @test transformer.unconstrain == _transform
            @test transformer.constrain == _inversetransform
            @test x_transformed == val_uncon
            @test x_inversetransformed == val_con

            @test x_inversetransformed isa typeof(val)
            val_unconstrained = unconstrain(_transform, val)
            val_constrained = constrain(_inversetransform, val_unconstrained)
            @test val_constrained isa typeof(val)
            logabs = log_abs_det_jac(_transform, val)
            @test logabs isa AbstractFloat
            #!NOTE - Second constraint not correct -> would not be a ModelWrapper
            @test !_check(_RNG, constraint, val)
        end
    end
end

@testset "Nested - AbstractArray - Automatic Differentiation" begin
    val = [1., [2., 3.], [4. 5. ; 6. 7.], 8., [9., 10.]]
    constraint = [Normal(), MvNormal([1., 1.]), Fixed(), Gamma(), Unconstrained()]
    reconstruct = ReConstructor(constraint, val)
    val_flat = flatten(reconstruct, val)

    check_AD = check_AD_closure(constraint, val)
    check_AD(val_flat)
    grad_mod_fd = ForwardDiff.gradient(check_AD, val_flat)
    grad_mod_rd = ReverseDiff.gradient(check_AD, val_flat)
    grad_mod_zy = Zygote.gradient(check_AD, val_flat)[1]
    @test sum(abs.(grad_mod_fd - grad_mod_rd)) ≈ 0 atol = _TOL
    @test sum(abs.(grad_mod_fd - grad_mod_zy)) ≈ 0 atol = _TOL

    grad_mod_fd = ForwardDiff.gradient(check_AD, flatten(reconstruct, val))
    grad_mod_rd = ReverseDiff.gradient(check_AD, flatten(reconstruct, val))
    grad_mod_zy = Zygote.gradient(check_AD, flatten(reconstruct, val))[1]
end

############################################################################################
@testset "Nested - Tuple" begin
    for output in outputtypes
        for flattentype in flattentypes
            flatdefault = FlattenDefault(; output = output, flattentype = flattentype)
            val = (1., [2, 3], [4. 5. ; 6. 7.], 8, [9., 10.])
            ReConstructor(val)
            reconstruct = ReConstructor(flatdefault, val)
# Flatten
            x_flat = flatten(reconstruct, val)
            @test x_flat isa AbstractVector
            @test eltype(x_flat) == output
# Flatten AD
            x_flatAD = flattenAD(reconstruct, val)
            @test x_flatAD isa AbstractVector
            @test eltype(x_flatAD) == Float64 #Always most common type
# Unflatten
            x_unflat = unflatten(reconstruct, x_flat)
            @test x_unflat isa typeof(val)
# Unflatten AD
            x_unflatAD = unflattenAD(reconstruct, x_flat)
#            @test x_unflatAD isa eltype(x_flat)
            x_unflatAD2 = unflattenAD(reconstruct, x_flatAD)
#            @test x_unflatAD2 isa eltype(x_flatAD)
################################################################################
# Constraint ~ Make Distribution as standard constraints are separately tested
            val = (1., [2, 3], [4. 5. ; 6. 7.], 8, [9., 10.])
            constraint = (Normal(), MvNormal([1., 1.]), Fixed(), Fixed(), Unconstrained())
            ReConstructor(constraint, val)
            reconstruct = ReConstructor(flatdefault, constraint, val)
# Flatten
            x_flat = flatten(reconstruct, val)
            @test x_flat isa AbstractVector
            @test eltype(x_flat) == output
# Flatten AD
            x_flatAD = flattenAD(reconstruct, val)
            @test x_flatAD isa AbstractVector
            @test eltype(x_flatAD) == Float64
# Unflatten
            x_unflat = unflatten(reconstruct, x_flat)
            @test x_unflat isa typeof(val)
# Unflatten AD
            x_unflatAD = unflattenAD(reconstruct, x_flat)
#            @test x_unflatAD isa eltype(x_flat)
            x_unflatAD2 = unflattenAD(reconstruct, x_flatAD)
#            @test x_unflatAD2 isa eltype(x_flatAD)
################################################################################
# Transforms
            _transform, _inversetransform = construct_transform(constraint, val)
            x_transformed = unconstrain(_transform, val)
            x_inversetransformed = constrain(_inversetransform, x_transformed)

            transformer = TransformConstructor(constraint, val)
            val_uncon = unconstrain(transformer, val)
            val_con = constrain(transformer, val_uncon)

            @test transformer.unconstrain == _transform
            @test transformer.constrain == _inversetransform
            @test x_transformed == val_uncon
            @test x_inversetransformed == val_con

            @test x_inversetransformed isa typeof(val)
            val_unconstrained = unconstrain(_transform, val)
            val_constrained = constrain(_inversetransform, val_unconstrained)
            @test val_constrained isa typeof(val)
            logabs = log_abs_det_jac(_transform, val)
            @test logabs isa AbstractFloat
            #!NOTE - Second constraint not correct -> would not be a ModelWrapper
            @test !_check(_RNG, constraint, val)
        end
    end
end

@testset "Nested - Tuple - Automatic Differentiation" begin
    val = (1., [2., 3.], [4. 5. ; 6. 7.], 8., [9., 10.])
    constraint = (Normal(), MvNormal([1., 1.]), Fixed(), Gamma(), Unconstrained())
    reconstruct = ReConstructor(constraint, val)
    val_flat = flatten(reconstruct, val)

    check_AD = check_AD_closure(constraint, val)
    check_AD(val_flat)
    grad_mod_fd = ForwardDiff.gradient(check_AD, val_flat)
    grad_mod_rd = ReverseDiff.gradient(check_AD, val_flat)
    grad_mod_zy = Zygote.gradient(check_AD, val_flat)[1]
    @test sum(abs.(grad_mod_fd - grad_mod_rd)) ≈ 0 atol = _TOL
    @test sum(abs.(grad_mod_fd - grad_mod_zy)) ≈ 0 atol = _TOL

    grad_mod_fd = ForwardDiff.gradient(check_AD, flatten(reconstruct, val))
    grad_mod_rd = ReverseDiff.gradient(check_AD, flatten(reconstruct, val))
    grad_mod_zy = Zygote.gradient(check_AD, flatten(reconstruct, val))[1]
end

############################################################################################
@testset "Nested - NamedTuple" begin
    for output in outputtypes
        for flattentype in flattentypes
            flatdefault = FlattenDefault(; output = output, flattentype = flattentype)
            val = (a = Float16(1.0), b = [2, 3], c = [4. 5. ; 6. 7.], d = 8, e = [9., 10.], f = (g = (h = 3.)))
            ReConstructor(val)
            reconstruct = ReConstructor(flatdefault, val)
# Flatten
            x_flat = flatten(reconstruct, val)
            @test x_flat isa AbstractVector
            @test eltype(x_flat) == output
# Flatten AD
            x_flatAD = flattenAD(reconstruct, val)
            @test x_flatAD isa AbstractVector
            @test eltype(x_flatAD) == Float64 #Always most common type
# Unflatten
            x_unflat = unflatten(reconstruct, x_flat)
            @test x_unflat isa typeof(val)
# Unflatten AD
            x_unflatAD = unflattenAD(reconstruct, x_flat)
#            @test x_unflatAD isa eltype(x_flat)
            x_unflatAD2 = unflattenAD(reconstruct, x_flatAD)
#            @test x_unflatAD2 isa eltype(x_flatAD)
################################################################################
# Constraint ~ Make Distribution as standard constraints are separately tested
            val = (a = Float16(1.0), b = [2, 3], c = [4. 5. ; 6. 7.], d = 8, e = [9., 10.], f = (g = (h = 3.)))
            constraint = (a = Normal(), b = MvNormal([1., 1.]), c = Fixed(), d = Fixed(), e = Unconstrained(), f = (g = (h = Gamma())))
            ReConstructor(constraint, val)
            reconstruct = ReConstructor(flatdefault, constraint, val)
# Flatten
            x_flat = flatten(reconstruct, val)
            @test x_flat isa AbstractVector
            @test eltype(x_flat) == output
# Flatten AD
            x_flatAD = flattenAD(reconstruct, val)
            @test x_flatAD isa AbstractVector
            @test eltype(x_flatAD) == Float64
# Unflatten
            x_unflat = unflatten(reconstruct, x_flat)
            @test x_unflat isa typeof(val)
# Unflatten AD
            x_unflatAD = unflattenAD(reconstruct, x_flat)
#            @test x_unflatAD isa eltype(x_flat)
            x_unflatAD2 = unflattenAD(reconstruct, x_flatAD)
#            @test x_unflatAD2 isa eltype(x_flatAD)
################################################################################
# Transforms
            _transform, _inversetransform = construct_transform(constraint, val)
            x_transformed = unconstrain(_transform, val)
            x_inversetransformed = constrain(_inversetransform, x_transformed)

            transformer = TransformConstructor(constraint, val)
            val_uncon = unconstrain(transformer, val)
            val_con = constrain(transformer, val_uncon)

            @test transformer.unconstrain == _transform
            @test transformer.constrain == _inversetransform
            @test x_transformed == val_uncon
            @test x_inversetransformed == val_con

            @test x_inversetransformed isa typeof(val)
            val_unconstrained = unconstrain(_transform, val)
            val_constrained = constrain(_inversetransform, val_unconstrained)
            @test val_constrained isa typeof(val)
            logabs = log_abs_det_jac(_transform, val)
            @test logabs isa AbstractFloat
            #!NOTE - Second constraint not correct -> would not be a ModelWrapper
            @test !_check(_RNG, constraint, val)
        end
    end
end

@testset "Nested - NamedTuple - Automatic Differentiation" begin
    val = (a = Float32(1.0), b = [2., 3.], c = [4. 5. ; 6. 7.], d = 8., e = [9., 10.], f = (g = (h = 3.)))
    constraint = (a = Normal(), b = MvNormal([1., 1.]), c = Fixed(), d = Gamma(), e = Unconstrained(), f = (g = (h = Gamma())))
    reconstruct = ReConstructor(constraint, val)
    val_flat = flatten(reconstruct, val)

    check_AD = check_AD_closure(constraint, val)
    check_AD(val_flat)
    grad_mod_fd = ForwardDiff.gradient(check_AD, val_flat)
    grad_mod_rd = ReverseDiff.gradient(check_AD, val_flat)
    grad_mod_zy = Zygote.gradient(check_AD, val_flat)[1]
    @test sum(abs.(grad_mod_fd - grad_mod_rd)) ≈ 0 atol = _TOL
    @test sum(abs.(grad_mod_fd - grad_mod_zy)) ≈ 0 atol = _TOL

    grad_mod_fd = ForwardDiff.gradient(check_AD, flatten(reconstruct, val))
    grad_mod_rd = ReverseDiff.gradient(check_AD, flatten(reconstruct, val))
    grad_mod_zy = Zygote.gradient(check_AD, flatten(reconstruct, val))[1]
end


# Author: Youngjun Kim, youngjun@stanford.edu
# Date: 08/19/2014
#
# Originally written in MATLAB by Mykel Kochenderfer, mykel@stanford.edu
# Converted by Youngjun Kim, youngjun@stanford.edu

#   SAVE_SCRIPTS(FILENAME,SCRIPTS) saves a structure holding the encounter
#   scripts stored to the specified file.
#
#   SCRIPTS FILE:
#   The scripts file contains a set of scripted encounters. Each
#   encounter is defined by a set of encounter scripts associated with a
#   fixed number of aircraft. The file is organized as follows:
#
#   [Header]
#   uint32 (number of encounters)
#   uint32 (number of aircraft)
#       [Encounter 1]
#           [Initial]
#               [Aircraft 1]
#               double (initial airspeed in ft/s)
#               double (initial north position in ft)
#               double (initial east position in ft)
#               double (initial altitude in ft)
#               double (initial heading angle in radians)
#               double (initial flight path angle in radians)
#               double (initial roll angle in radians)
#               double (initial airspeed acceleration in ft/s^2)
#               ...
#               [Aircraft n]
#               double (initial airspeed in ft/s)
#               double (initial north position in ft)
#               double (initial east position in ft)
#               double (initial altitude in ft)
#               double (initial heading angle in radians)
#               double (initial flight path angle in radians)
#               double (initial roll angle in radians)
#               double (initial airspeed acceleration in ft/s^2)
#           [Updates]
#               [Aircraft 1]
#               uint8 (number of updates)
#                   [Update 1]
#                   double (time in s)
#                   double (commanded vertical rate in ft/s)
#                   double (commanded turn rate in rad/s)
#                   double (commanded airspeed acceleration in ft/s^2)
#                   ...
#                   [Update m]
#                   double (time in s)
#                   double (commanded vertical rate in ft/s)
#                   double (commanded turn rate in rad/s)
#                   double (commanded airspeed acceleration in ft/s^2)
#               ...
#               [Aircraft n]
#                   ...
#       ...
#       [Encounter k]
#           ...
#
#   SCRIPTS STRUCTURE:
#   The scripts structure is an m x n structure matrix, where m is the
#   number of aircraft and n is the number of encounters. This structure
#   matrix has two fields: initial and update. The initial field is an 8
#   element array specifying the airspeed, north position, east position,
#   altitude, heading, flight path angle, roll angle, and airspeed
#   acceleration. The update field is a 4 x n matrix, where n is the number
#   of updates. The rows correspond to the time, vertical rate, turn rate
#   and airspeed acceleration.

# Inputs:
#       filename: name of output file
#
#       encounters: structure of encounters to save to file specified by
#       "filename".  For a guide to the format of this structure, see
#       the help documentation in load_scripts.m.
#
#       varargin: parameter/value pairs
#
#           +   parameter: numupdatetype
#
#               value (character string): data type of "num_update"
#               variable: the number of updates scripted for some aircraft
#               in some encounter (default is 'uint32')
#
#           +   parameter: append
#
#               value (logical): set to true to append encounters to the
#               end of the file specified by "filename".  Set to false
#               (default) to write a new file with the name specified by
#               "filename".
#
#           +   parameter: floattype
#
#               value (character string): data type of encounter initial
#               conditions and updates (default is 'double')

function save_encounters(filename, encounters; numupdatetype = UInt32, append = false, floattype = Float64)

    num_encounters::UInt32 = size(encounters, 2)
    num_ac = size(encounters, 1)

    if append
        if isreadable(filename) == false
            fio = open(filename, "w")

            write(fio, UInt32(num_encounters))
            write(fio, UInt32(num_ac))
        else
            fio = open(filename, "a+")
            seekstart(fio)

            file_num_encounters = read(fio, UInt32)
            file_num_ac = read(fio, UInt32)

            if file_num_ac != num_ac
                error("Must have the same number of aircraft when appending encounter files")
            end

            seekstart(fio)
            write(fio, UInt32(file_num_encounters + num_encounters))

            seekend(fio)
        end
    else
        fio = open(filename, "w")

        write(fio, UInt32(num_encounters))
        write(fio, UInt32(num_ac))
    end

    for i = 1:num_encounters
        for j = 1:num_ac
            write(fio, convert(Array{floattype}, encounters[j, i]["initial"]))
        end

        for j = 1:num_ac
            num_update = size(encounters[j, i]["update"], 2)

            write(fio, convert(numupdatetype, num_update))
            write(fio, convert(Array{floattype}, (encounters[j, i]["update"])))
        end
    end

    close(fio)
end

function save_scripts(filename, scripts; numupdatetype = UInt32, append = false, floattype = Float64)

    save_encounters(filename, scripts; numupdatetype = UInt8, append = append, floattype = floattype)
end

if false
    save_scripts("encounters.dat", scripts; append = false)
end


using DataFrames, Dates

# This contains code which makes up the pipeline of csv -> interaction seqs

export add_page_column!, get_path_sequences, my_filter

# Some intial helper functions
# Take row (vector) and output a label which merges the entries with seperator
# '-', e.g. [1,2,1] -> '1-2-1'

function make_label(X::AbstractVector; nan="NaN")
    return join(X[findall(X.!=nan)], "-")
end

function make_labels(X::AbstractArray; nan="NaN")
    Y = Vector{String}()
    for row in eachrow(X)
        push!(Y, make_label(row))
    end
    return Y
end

# tmp = rand(1:10, 100,10)
# make_labels(tmp)
# tmp1

# Now a function for dataframes

function make_labels(X::AbstractDataFrame, cols::Vector{Int}; nan="NaN")
    X = Matrix(X[:,cols])
    return make_labels(X, nan=nan)
end

function make_labels(X::AbstractDataFrame, cols::Int; nan="NaN")
    return make_labels(X, [cols], nan=nan)
end

function add_page_column!(
    X::AbstractDataFrame,
    cols::Vector{Int};
    level::Int=length(cols), nan="NaN")
    # @show level
    X.page = make_labels(X, cols[1:level], nan=nan)
end

## Splitting into sessions

function get_paths(
    X::AbstractDataFrame;
    T::TimePeriod=Minute(15),
    vertex_col_name::String="vertex",
    time_col_name::String="time_gmt")

    @assert(issorted(X, time_col_name), "Must be sorted according to time.")

    function recurse_split!(
        X::AbstractDataFrame,
        Y::AbstractVector,
        ind::BitArray{1},
        i::Int;
        vertex_col_name::String)

        if findnext(ind, i) == nothing
            push!(Y, X[i:end, vertex_col_name])
            return
        else
            j = findnext(ind, i)
            # println("Path from $(i) to $(j)")
            # println(X[i:j])
            push!(Y, X[i:j, vertex_col_name])
            recurse_split!(X, Y, ind, j+1, vertex_col_name=vertex_col_name)
        end
    end
    Y = Vector{Vector{eltype(X[:,vertex_col_name])}}() # Initialise storage
    ind = diff(X[:,time_col_name]).>T
    recurse_split!(X, Y, ind, 1, vertex_col_name=vertex_col_name)
    return Y
end

## Now do on a grouped data frame (will return a vector of vector of vectors)
"""
Given a `DataFrames.GroupedDataFrame` object and some time threshold of type `Dates.TimePeriod` this will convert the data to interaction sequences.

Column names for vertex labels and timestamps are passed names arguments.

**Note** the time column must be of type `Dates.DateTime`.
"""
function get_path_sequences(
    X::GroupedDataFrame;
    T::TimePeriod=Minute(15),
    vertex_col_name::String="vertex",
    time_col_name::String="time_gmt")

    data = Vector{Vector{Vector{String}}}(undef, length(X))
    for (i, key) in enumerate(keys(X))
        data[i] = get_paths(X[key], T=T, vertex_col_name=vertex_col_name, time_col_name=time_col_name)
    end
    return data
end

## Filtering functions
# On raw dataframe

function max_length(
    X::AbstractDataFrame;
    time_col_name::String="time_gmt",
    T::TimePeriod=Minute(15))

    tmp = diff(X[:,time_col_name]) .> T
    push!(tmp, true)
    pushfirst!(tmp, true)
    return maximum(diff(findall(tmp)))
end

function get_num_paths(
    X::AbstractDataFrame;
    time_col_name::String="time_gmt",
    T::TimePeriod=Minute(15)
    )
    return sum(diff(X[:, time_col_name]) .> T) + 1
end

function my_filter(
    X::GroupedDataFrame;
    T::TimePeriod=Minute(15),
    time_col_name::String="time_gmt",
    max_clicks=7000,
    max_path_length=750,
    max_num_paths=450
    )
    clicks = [size(x)[1] for x in X]
    path_lengths = [max_length(group, time_col_name=time_col_name, T=T) for group in X]
    num_paths = [get_num_paths(group, time_col_name=time_col_name, T=T) for group in X]

    ind = (clicks .< max_clicks) .& (path_lengths .< max_path_length) .& (num_paths .< max_num_paths)

    X[ind]
end

# On session data

function max_length(X::Vector{Vector{T}}) where T
    return maximum(map(length, X))
end

function collapsed_length(X::Vector{Vector{T}}) where T
    z = 0
    for x in X
        z += length(x)
    end
    return z
end

function my_filter(
    X::Vector{Vector{Vector{String}}};
    max_clicks=7000,
    max_path_length=750,
    max_num_paths=450
    )
    ind = (collapsed_length.(X) .< max_clicks) .& (max_length.(X) .< max_path_length) .& (length.(X) .< max_num_paths)
    return X[ind]
end


using VCFTools
using Test
using VariantCallFormat
using CodecZlib
using DelimitedFiles

# packages needed only for testing
using Random
using CSV
using DataFrames
using StatsBase
using LinearAlgebra

include("gtstats_test.jl")
include("conformgt_test.jl")
include("convert_test.jl")
include("filter_test.jl")
include("aim_test.jl")
include("grm_test.jl")


import MITgcmTools: read_namelist

function read_namelist(x:: SPEEDY_config)

    fil=joinpath(pathof(x),"rundir/namelist.nml")
    nml=readlines(fil)
    i=findall([nml[i][1]!=='!' for i in 1:length(nml)])
    nml=nml[i]

    i0=findall([nml[i][1]=='&' for i in 1:length(nml)])
    i1=findall([nml[i][1]=='/' for i in 1:length(nml)])

    tmp0=OrderedDict()
    for i in 1:length(i0)
        tmp1=Symbol(nml[i0[i]][2:end])
        tmp2=OrderedDict()
        for j in i0[i]+1:i1[i]-1
            tmp3=split(nml[j],'=')
            tmp2[tmp3[1]]=parse(Int,tmp3[2])
        end
        tmp0[tmp1]=tmp2
    end

    return tmp0
end


# This file was generated, do not modify it. # hide
DataFrames.describe(features)


# LibTOAST.jl: interface to the TOAST++ library
# Copyright (C) 2019 Samuel Powell

@testset "Meshes" begin

  # Load 2D .msh file from disk and create from vertices, elements
  mesh1 = Mesh(joinpath(meshdir_2D, "circle25_32.msh"));
  vtx, ele, elt = data(mesh1)
  mesh2 = Mesh(vtx,ele,elt)

  for m in [mesh1 mesh2]
    @test nodecount(m) == 3511
    @test elemcount(m) == 6840
    @test dimensions(m) == 2
    @test boundingbox(m) == [-25.0 -25.0; 25.0 25.0]
    @test fullsize(m) == 1963.0959349692218
    @test maxnodes(m) == 3
    vtx, ele = surface(m)
    @test vtx[10,:] == [19.7003,15.3915]
    @test ele[1:5] == [90,135,89,164,88]
  end

  # Load 3D .msh file from disk and create from vertices, elements
  mesh1 = Mesh(joinpath(meshdir_3D, "cyl3.msh"));
  vtx, ele, elt = data(mesh1)
  mesh2 = Mesh(vtx,ele,elt)

  for m in [mesh1 mesh2]
    @test nodecount(m) == 27084
    @test elemcount(m) == 141702
    @test dimensions(m) == 3
    @test boundingbox(m) == [-25.0 -25.0 -25.0; 25.0 25.0 25.0]
    @test fullsize(m) == 98056.71142861326
    @test maxnodes(m) == 4
    vtx, ele = surface(m)
    @test vtx[10,:] == [-1.29504,6.49008,25.0]
    @test ele[1:5] == [ 2791,1429,3014,2552,124]
  end

end


import GLPK

include("LPBenchmark.jl")

function run_glpk(fname::String)
    glpk = GLPK.Optimizer(method=GLPK.INTERIOR)
    bridged = MOIB.full_bridge_optimizer(glpk, Float64)

    load_problem!(bridged, fname)

    MOI.set(bridged, MOI.TimeLimitSec(), 10_000)
    MOI.set(bridged, MOI.RawParameter("msg_lev"), 3)

    # Solve the problem
    run(bridged)

    # Display timing info
    return nothing
end

if abspath(PROGRAM_FILE) == @__FILE__
    run_glpk(ARGS[1])
end


using Test

for testscen in 1:2
    if testscen == 2
        continue # this validation data is missing
    end

    valdir, scenario, use_permafrost, use_seaice = get_scenario(testscen)
    println(scenario)

    m = page_model()
    include("../src/components/RCPSSPScenario.jl")
    include("../src/components/EquityWeighting.jl")

    rcpsspscenario = addrcpsspscenario(m, scenario)
    equityweighting = add_comp!(m, EquityWeighting)

    equityweighting[:tct_percap_totalcosts_total] = readpagedata(m, "test/validationdata/$valdir/tct_per_cap_totalcostspercap.csv")
    equityweighting[:act_adaptationcosts_total] = readpagedata(m, "test/validationdata/$valdir/act_adaptationcosts_tot.csv")
    equityweighting[:act_percap_adaptationcosts] = readpagedata(m, "test/validationdata/$valdir/act_percap_adaptationcosts.csv")
    equityweighting[:cons_percap_aftercosts] = readpagedata(m, "test/validationdata/$valdir/cons_percap_aftercosts.csv")
    equityweighting[:cons_percap_consumption] = readpagedata(m, "test/validationdata/$valdir/cons_percap_consumption.csv")
    equityweighting[:cons_percap_consumption_0] = readpagedata(m, "test/validationdata/cons_percap_consumption_0.csv")
    equityweighting[:rcons_percap_dis] = readpagedata(m, "test/validationdata/$valdir/rcons_per_cap_DiscRemainConsumption.csv")
    equityweighting[:yagg_periodspan] = readpagedata(m, "test/validationdata/yagg_periodspan.csv")
    equityweighting[:pop_population] = readpagedata(m, "test/validationdata/$valdir/pop_population.csv")
    equityweighting[:y_year_0] = 2015.
    equityweighting[:y_year] = Mimi.dim_keys(m.md, :time)
    equityweighting[:grw_gdpgrowthrate] = rcpsspscenario[:grw_gdpgrowthrate]
    equityweighting[:popgrw_populationgrowth] = rcpsspscenario[:popgrw_populationgrowth]

    p = load_parameters(m)
    set_leftover_params!(m, p)

    run(m)

    # Generated data
    df = m[:EquityWeighting, :df_utilitydiscountfactor]
    wtct_percap = m[:EquityWeighting, :wtct_percap_weightedcosts]
    pct_percap = m[:EquityWeighting, :pct_percap_partiallyweighted]
    dr = m[:EquityWeighting, :dr_discountrate]
    dfc = m[:EquityWeighting, :dfc_consumptiondiscountrate]
    pct = m[:EquityWeighting, :pct_partiallyweighted]
    pcdt = m[:EquityWeighting, :pcdt_partiallyweighted_discounted]
    wacdt = m[:EquityWeighting, :wacdt_partiallyweighted_discounted]
    aact = m[:EquityWeighting, :aact_equityweightedadaptation_discountedaggregated]

    wit = m[:EquityWeighting, :wit_equityweightedimpact]
    addt = m[:EquityWeighting, :addt_equityweightedimpact_discountedaggregated]
    addt_gt = m[:EquityWeighting, :addt_gt_equityweightedimpact_discountedglobal]
    te = m[:EquityWeighting, :te_totaleffect]

    # Recorded data
    df_compare = readpagedata(m, "test/validationdata/df_utilitydiscountfactor.csv")
    wtct_percap_compare = readpagedata(m, "test/validationdata/$valdir/wtct_percap_weightedcosts.csv")
    pct_percap_compare = readpagedata(m, "test/validationdata/$valdir/pct_percap_partiallyweighted.csv")
    dr_compare = readpagedata(m, "test/validationdata/$valdir/dr_discountrate.csv")
    dfc_compare = readpagedata(m, "test/validationdata/$valdir/dfc_consumptiondiscountrate.csv")
    pct_compare = readpagedata(m, "test/validationdata/$valdir/pct_partiallyweighted.csv")
    pcdt_compare = readpagedata(m, "test/validationdata/$valdir/pcdt_partiallyweighted_discounted.csv")
    wacdt_compare = readpagedata(m, "test/validationdata/$valdir/wacdt_partiallyweighted_discounted.csv")
    aact_compare = readpagedata(m, "test/validationdata/$valdir/aact_equityweightedadaptation_discountedaggregated.csv")

    wit_compare = readpagedata(m, "test/validationdata/$valdir/wit_equityweightedimpact.csv")
    addt_compare = readpagedata(m, "test/validationdata/$valdir/addt_equityweightedimpact_discountedaggregated.csv")
    addt_gt_compare = (scenario == "NDCs" ? 9.960706559386551e8 : 3.445096140740376e8)
    te_compare = (scenario == "NDCs" ? 1.0320923880568126e9 : 5.0496576104580307e8)

    @test df ≈ df_compare rtol=1e-8
    @test wtct_percap ≈ wtct_percap_compare rtol=1e-7
    @test pct_percap ≈ pct_percap_compare rtol=1e-7
    @test dr ≈ dr_compare rtol=1e-5
    @test dfc ≈ dfc_compare rtol=1e-7
    @test pct ≈ pct_compare rtol=1e-3
    @test pcdt ≈ pcdt_compare rtol=1e-3
    @test wacdt ≈ wacdt_compare rtol=1e-4

    @test aact ≈ aact_compare rtol=1e-3

    @test wit ≈ wit_compare rtol=1e-3
    @test addt ≈ addt_compare rtol=1e-2
    @test addt_gt ≈ addt_gt_compare rtol=1e-2

    @test te ≈ te_compare rtol=1e-2
end


twosort(array) = join(minimum(array),"***")


"""
    IdentityAxis(start, stop) -> axis
    IdentityAxis(keys::AbstractUnitRange) -> axis
    IdentityAxis(keys::AbstractUnitRange, indices::AbstractUnitRange) -> axis

`AbstractAxis` subtype that preserves indices after indexing.

# Examples
```julia
julia> using AxisIndices

julia> axis = AxisIndices.IdentityAxis(3:5)
idaxis(3:5)(SimpleAxis(1:3))

julia> axis[4:5]
idaxis(4:5)(SimpleAxis(2:3))

```
"""
struct IdentityAxis{I,Inds<:AbstractAxis,F} <: AbstractOffsetAxis{I,Inds,F}
    offset::F
    parent::Inds

    @inline function IdentityAxis{I,Inds,F}(f::Integer, inds::AbstractUnitRange) where {I,Inds,F}
        if inds isa Inds && f isa F
            return new{I,Inds,F}(f, inds)
        else
            return IdentityAxis(f, Inds(inds))
        end
    end

    function IdentityAxis{I,Inds}(f::Integer, inds::AbstractAxis) where {I,Inds}
        return IdentityAxis{I,Inds,typeof(f)}(f, inds)
    end
    function IdentityAxis{I,Inds}(f::Integer, inds::AbstractArray) where {I,Inds}
        return IdentityAxis{I,Inds}(f, compose_axis(inds))
    end
    function IdentityAxis{I,Inds}(f::AbstractUnitRange, inds::AbstractAxis) where {I,Inds}
        return IdentityAxis{I,Inds}(static_first(ks) - static_first(inds), inds)
    end

    function IdentityAxis{I,Inds}(f::AbstractUnitRange, inds::AbstractArray) where {I,Inds}
        return IdentityAxis{I,Inds}(f, compose_axis(inds))
    end
    function IdentityAxis{I,Inds}(ks::AbstractUnitRange) where {I,Inds}
        if can_change_size(ks)
            inds = OneToMRange(length(ks))
        else
            inds = OneTo(static_length(ks))
        end
        f = static_first(ks) - one(F)
        return IdentityAxis{eltype(inds),typeof(f),typeof(inds)}(f, inds)
    end

    # IdentityAxis{I}
    function IdentityAxis{I}(f::Integer, inds::AbstractOffsetAxis) where {I}
        return IdentityAxis(f + offsets(inds, 1), parent(inds))
    end
    function IdentityAxis{I}(f::Integer, inds::AbstractArray) where {I}
        return IdentityAxis{I}(f, compose_axis(inds))
    end
    function IdentityAxis{I}(f::Integer, inds::AbstractAxis) where {I}
        return IdentityAxis{I,typeof(inds)}(f, inds)
    end
    function IdentityAxis{I}(ks::AbstractRange{<:Integer}) where {I}
        return IdentityAxis{I}(ks, OneTo{I}(length(ks)))
    end
    function IdentityAxis{I}(start::Integer, stop::Integer) where {I}
        return IdentityAxis{I}(start:stop)
    end
    function IdentityAxis{I}(ks::AbstractRange, inds::AbstractAxis) where {I}
        return IdentityAxis{I}(static_first(ks) - static_first(inds), inds)
    end
    function IdentityAxis{I}(ks::AbstractRange, inds::AbstractOffsetAxis) where {I}
        return IdentityAxis{I}(ks, parent(inds))
    end
    function IdentityAxis{I}(ks::AbstractRange, inds::AbstractArray) where {I}
        return IdentityAxis{I}(ks, compose_axis(inds))
    end

    # IdentityAxis
    function IdentityAxis(f::Integer, inds::AbstractAxis)
        return IdentityAxis{eltype(inds)}(f, inds)
    end
    function IdentityAxis(ks::AbstractRange, inds::AbstractAxis)
        return IdentityAxis{eltype(inds)}(ks, compose_axis(inds))
    end
    function IdentityAxis(ks::AbstractRange, inds::AbstractArray)
        return IdentityAxis(ks, compose_axis(inds))
    end
    IdentityAxis(start::Integer, stop::Integer) = IdentityAxis(start:stop)
    function IdentityAxis(ks::Ks) where {Ks}
        if can_change_size(ks)
            inds = SimpleAxis(OneToMRange(length(ks)))
        else
            inds = SimpleAxis(StaticInt(1):static_length(ks))
        end
        f = static_first(ks)
        f = f - one(f)
        return new{eltype(inds),typeof(inds),typeof(f)}(f, inds)
    end
end

@inline Base.getproperty(axis::IdentityAxis, k::Symbol) = getproperty(parent(axis), k)

ArrayInterface.known_first(::Type{T}) where {T<:IdentityAxis{<:Any,<:Any,<:Any}} = nothing
function ArrayInterface.known_first(::Type{T}) where {Inds,F,T<:IdentityAxis{<:Any,Inds,StaticInt{F}}}
    if known_first(Inds) === nothing
        return nothing
    else
        return known_first(Inds) + F
    end
end
Base.first(axis::IdentityAxis) = first(parent(axis)) + getfield(axis, :offset)

ArrayInterface.known_last(::Type{T}) where {T<:IdentityAxis{<:Any,<:Any,<:Any}} = nothing
function ArrayInterface.known_last(::Type{T}) where {Inds,F,T<:IdentityAxis{<:Any,Inds,StaticInt{F}}}
    if known_last(Inds) === nothing
        return nothing
    else
        return known_last(Inds) + F
    end
end
Base.last(axis::IdentityAxis) = last(parent(axis)) + getfield(axis, :offset)

function ArrayInterface.unsafe_reconstruct(axis::IdentityAxis, inds; keys=nothing, kwargs...)
    if keys === nothing
        if inds isa AbstractOffsetAxis
            f_axis = offsets(axis, 1)
            f_inds = offsets(inds, 1)
            if f_axis === f_inds
                return OffsetAxis(offsets(axis, 1), unsafe_reconstruct(parent(axis), parent(inds); kwargs...))
            else
                return OffsetAxis(f_axis + f_inds, unsafe_reconstruct(parent(axis), parent(inds); kwargs...))
            end
        else
            return OffsetAxis(offsets(axis, 1), unsafe_reconstruct(parent(axis), inds; kwargs...))
        end
    else
        return IdentityAxis(keys, inds)
    end
end

"""
    idaxis(inds::AbstractUnitRange{<:Integer})

Shortcut for creating [`IdentityAxis`](@ref).

# Examples

```jldoctest
julia> using AxisIndices

julia> AxisArray(ones(3), idaxis)[2:3]
2-element AxisArray(::Vector{Float64}
  • axes:
     1 = 2:3
)
     1
  2  1.0
  3  1.0

```
"""
struct IdAxis <: AxisInitializer end
axis_method(::IdAxis, x, inds) = IdentityAxis(x, inds)
const idaxis = IdAxis()
function idaxis(collection)
    if known_step(collection) === 1
        return IdentityAxis(collection)
    else
        return idaxis(collection, axes(collection))
    end
end

"""
    IdentityArray(A::AbstractArray)

Provides [`IdentityAxis`](@ref)s for indexing `A`.

## Examples
```jldoctest
julia> using AxisIndices

julia> AxisIndices.IdentityArray(ones(3,3))[2:3, 2:3]
2×2 AxisArray(::Matrix{Float64}
  • axes:
     1 = 2:3
     2 = 2:3
)
     2    3
  2  1.0  1.0
  3  1.0  1.0

```
"""
const IdentityArray{T,N,P,A<:Tuple{Vararg{<:IdentityAxis}}} = AxisArray{T,N,P,A}

IdentityArray(A::AbstractArray{T,N}) where {T,N} = IdentityArray{T,N}(A)

IdentityArray(A::AbstractArray{T,0}) where {T} = AxisArray(A)

IdentityArray{T}(A::AbstractArray) where {T} = IdentityArray{T,ndims(A)}(A)

IdentityArray{T,N}(A::AbstractArray) where {T,N} = IdentityArray{T,N,typeof(A)}(A)

function IdentityArray{T,N,P}(x::P) where {T,N,P<:AbstractArray{T,N}}
    return initialize_axis_array(x, map(IdentityAxis, axes(x)))
end
function AxisArray{T,N,P}(x::P, axs::Tuple) where {T,N,P}
    axs = map(IdentityAxis, axs)
    return AxisArray{T,N,P}(convert(P, x), axs)
end

IdentityArray{T,N,P}(A::IdentityArray{T,N,P}) where {T,N,P} = A

"""
    IdentityVector(v::AbstractVector)

Provides an [`IdentityAxis`](@ref) for indexing `v`.

## Examples
```jldoctest
julia> using AxisIndices

julia> AxisIndices.IdentityVector(ones(4))[3:4]
2-element AxisArray(::Vector{Float64}
  • axes:
     1 = 3:4
)
     1
  3  1.0
  4  1.0

```
"""
const IdentityVector{T,P<:AbstractVector{T},Ax1<:IdentityAxis} = IdentityArray{T,1,P,Tuple{Ax1}}

IdentityVector(A::AbstractVector) = IdentityArray{eltype(A)}(A)
IdentityVector{T}(A::AbstractVector) where {T} = IdentityArray{T,1}(A)

"""
    IdentityVector{T}(init::ArrayInitializer, sz::Integer)

Creates a vector with elements of type `T` of size `sz` an [`IdentityAxis`](@ref).

## Examples
```jldoctest
julia> using AxisIndices

julia> AxisIndices.IdentityVector{Union{Missing, Int}}(missing, 3)[2:3]
2-element AxisArray(::Vector{Union{Missing, Int64}}
  • axes:
     1 = 2:3
)
     1
  2   missing
  3   missing

```
"""
function IdentityVector{T}(init::ArrayInitializer, arg) where {T}
    return IdentityVector{T}(Vector{T}(init, arg))
end

@inline function ArrayInterface.to_axis(::IndexStyle, axis::IdentityAxis, inds)
    return unsafe_reconstruct(axis, StaticInt(1):static_length(inds); keys=inds)
end

function print_axis(io::IO, axis::IdentityAxis)
    print(io, "idaxis(")
    print(io, Int(first(axis)))
    print(io, ")(")
    print(io, Int(last(axis)))
    print(io, " parent=$(parent(axis)))")
end


# script that outputs some statistics of the model fits

using DataFrames, Gadfly, Colors, MAT, CSV
import Cairo, Fontconfig

include("common.jl")
include("models.jl")
include("behaviorstats.jl")
include("simmodel.jl")
include("modelfitting.jl")

# settings (when manually evaluating the objective function)
const blockreps = 10
const simtrials = 10^5
const burnin = 1000
const nrats = 1 # do not divide by nrats = 4
const δt = 0.0001
const t = Theme(minor_label_font_size=6pt,
                major_label_font_size=6pt,
                grid_line_width=0pt, grid_color=RGB(1.0, 1.0, 1.0))
const fitpath = "fits"
const datapath = "../fitdata/$(fitpath)"
const figpath = "../fitdata/$(fitpath)"

# models
const modeldict = Dict("opt" => OptimalModel,
                       "optlapse" => OptimalLapseModel,
                       "optcollapse" => OptimalColLapseModel,
                       "optcolbiaslapse" => OptimalColBiasLapseModel,
                       "adflapse" => ADFLapseModel,
                       "adfcollapse" => ADFColLapseModel,
                       "adfbiaslapse" => ADFBiasLapseModel,
                       "adfcolbiaslapse" => ADFColBiasLapseModel,
                       "randomcolbiaslapse" => RandomColBiasLapseModel,
                       "randomcollapse" => RandomColLapseModel,
                       "gammalapse" => GammaLapseModel,
                       "gammacollapse" => GammaColLapseModel,
                       "gammabiaslapse" => GammaBiasLapseModel,
                       "gammacolbiaslapse" => GammaColBiasLapseModel,
                       "deltalapse" => DeltaLapseModel,
                       "deltacollapse" => DeltaColLapseModel,
                       "deltabiaslapse" => DeltaBiasLapseModel,
                       "deltacolbiaslapse" => DeltaColBiasLapseModel,
                       "deltaexpcolbiaslapse" => DeltaExpColBiasLapseModel,
                       "lcalapse" => LCALapseModel,
                       "lcacollapse" => LCAColLapseModel)
# separate list to preserve order that keys(modeldict) wouldn't have
# const modelnames = ["opt", "optlapse", "optcollapse",
#                     "adflapse", "adfcollapse", "adfbiaslapse", "adfcolbiaslapse",
#                     "gammalapse", "gammacollapse", "gammabiaslapse", "gammacolbiaslapse",
#                     "deltalapse", "deltacollapse", "deltabiaslapse", "deltacolbiaslapse",
#                     "deltaexpcolbiaslapse"]
const modelnames = ["opt", "optlapse", "optcollapse", "optcolbiaslapse",
                    "adflapse", "adfcollapse", "adfbiaslapse", "adfcolbiaslapse",
                    "randomcolbiaslapse", "randomcollapse",
                    "gammalapse", "gammacollapse", "gammabiaslapse", "gammacolbiaslapse",
                    "deltaexpcolbiaslapse", "lcalapse", "lcacollapse"]
# const modelnames = ["opt", "optlapse", "optcollapse",
#                     "adflapse", "adfcollapse", "adfbiaslapse", "adfcolbiaslapse",
#                     "gammalapse", "gammacollapse", "gammabiaslapse", "gammacolbiaslapse",
#                     "deltaexpcolbiaslapse"]
const basemodel = "adfcolbiaslapse"  # base model for model comparison
@assert issubset(Set(modelnames), Set(keys(modeldict))) # unordered comparison

# objective functions
# second tuple component is trials, 1 = ident, 2 = categ, 3 = itlvd
const objfndict = Dict("psychchron1" => (PsychChron1ObjFn, [1, 2]),
                       "psychchron1ident" => (PsychChron1IdentObjFn, [1]),
                       "psychchron1categ" => (PsychChron1CategObjFn, [2]),
                       "psychchron2" => (PsychChron2ObjFn, [1, 2]),
                       "psychchron3" => (PsychChron3ObjFn, [1, 2]),
                       "psychchronseq1" => (PsychChronSeq1ObjFn, [1, 2]),
                       "psychchronseq2" => (PsychChronSeq2ObjFn, [1, 2]),
                       "psychchronseq3" => (PsychChronSeq3ObjFn, [1, 2]),
                       "psychchron1itlvd" => (PsychChron1ItlvdObjFn, [3]))
#const objfns = ["psychchron1", "psychchron1ident", "psychchron2", "psychchron3",
#                "psychchronseq1", "psychchronseq2", "psychchronseq3"]
const objfns = ["psychchron1", "psychchronseq1"]
@assert issubset(Set(objfns), Set(keys(objfndict)))

# model fit criteria, converted to log-evidence in log(10)
BIC(avgllh, ϕn, trials) = (trials * avgllh - 0.5ϕn * log(trials)) / log(10)
AIC(avgllh, ϕn, trials) = (trials * avgllh - ϕn) / log(10)
AICc(avgllh, ϕn, trials) = AIC(avgllh, ϕn, trials) - (ϕn * (ϕn+1) / (trials-ϕn-1)) / log(10)

# returns number of trials in categ/ident dataset as vector
datatrials(ratid::Int=0) = [size(d, 1) / nrats
    for d in (load_data(:categ, ratid), load_data(:ident, ratid))]
datatrials(objfn::AbstractString, ratid::Int=0) = sum(datatrials(ratid)[objfndict[objfn][2]])

# returns dataframe with best-fitting parameters and avg. log-likelihood
function loadmodelfits(modelname, objfnname, ratid::Int=0)

    fitfilename = ratid == 0 ? "$(datapath)/fit_$(modelname)_$(objfnname).csv" :
                               "$(datapath)/fit_$(modelname)_$(objfnname)_rat$(ratid).csv"
    if !isfile(fitfilename)
        error("Could not find fit file $fitfilename")
    end
    d = CSV.read(fitfilename)
    # identify max-llh parameters
    maxllh, maxi = findmax(d[:llh])
    if maxllh != d[end,:llh]
        warning("Last element in $fitfilename is not best-fitting")
    end
    return d[maxi,:]
end

# prints the fit parameters for all objective functions for a given model and given rats
function printfitparams(modelname, ratids::Vector{Int}=[0])
    # use first objective function to get parameter names
    d = loadmodelfits(modelname, objfns[1], ratids[1])
    ϕnames = names(d)
    # print parameters
    objfnlen = maximum([length(s) for s in objfns])
    ratidlen = maximum([floor(Int, ratid > 0 ? log10(ratid) : 0.0)+1
                        for ratid in ratids]) + 3
    ratidstrs = [ratid > 0 ? "rat$(ratid)" : "all" for ratid in ratids]
    println(repeat("-", objfnlen + ratidlen + 1), "  ",
            join([@sprintf("%7s", string(ϕi)) for ϕi in ϕnames], "  "))
    # print fits
    for objfn in objfns
        for ratidx in 1:length(ratids)
            d = loadmodelfits(modelname, objfn, ratids[ratidx])
            if ratidx == 1
                print(repeat(" ", objfnlen - length(objfn)), objfn)
            else
                print(repeat(" ", objfnlen))
            end
            println(repeat(" ", ratidlen - length(ratidstrs[ratidx]) + 1),
                    ratidstrs[ratidx], "  ",
                    join([@sprintf("%7.2f", x) for x in Vector(d)], "  "))
        end
    end
end

# prints the fit quality for all models for the given objective function
function printfitquality(objfn, ratids::Vector{Int}=Int[])
    # remove 0's from ratids's vector, as they are default
    fitmeasures = ["BIC", "AIC", "AICc"]
    ratids = ratids[ratids .!= 0]
    ratnum = length(ratids)
    modelnum = length(modelnames)
    ratidlen = maximum([floor(Int, log10(ratid))+1 for ratid in ratids]) + 3
    ratidstrs = ["rat$(ratid)" for ratid in ratids]
    modellen = max(maximum([length(s) for s in modelnames]), 3) # include "all" here
    println(repeat("-", modellen + ratidlen + 1), "  ",
            join([@sprintf("%10s      ", s) for s in ["avg. llh", "BIC", "AIC", "AICc"]], "  "))
    fitstats = Array{Float64}(undef, ratnum, modelnum, length(fitmeasures)+1)
    fitstatsall = Array{Float64}(undef, modelnum, length(fitmeasures))
    for modelidx = 1:modelnum
        # fits across all rats
        modelname = modelnames[modelidx]
        d = loadmodelfits(modelname, objfn)
        avgllh = d[:llh]
        ϕn = length(d) - 1
        trials = datatrials(objfn)
        fitq = [avgllh; [f(avgllh, ϕn, trials) for f in [BIC, AIC, AICc]]]
        println(repeat(" ", modellen - length(modelname)), modelname, "  ",
                repeat(" ", ratidlen - 3), "all ",
                join([@sprintf("%10.2f      ", x) for x in fitq], "  "))
        fitstatsall[modelidx, :] = fitq[2:end]
        # fits for individual rats
        if length(ratids) == 0
            continue
        end
        for ratidx in 1:ratnum
            ratid = ratids[ratidx]
            d = loadmodelfits(modelname, objfn, ratid)
            avgllh = d[:llh]
            ϕn = length(d) - 1
            trials = datatrials(objfn, ratid)
            fitstats[ratidx,modelidx,:] = [avgllh; [f(avgllh, ϕn, trials) for f in [BIC, AIC, AICc]]]
            println(repeat(" ", modellen + 1), " ",
                    repeat(" ", ratidlen - length(ratidstrs[ratidx])), ratidstrs[ratidx], " ",
                    join([@sprintf("%10.2f      ", x) for x in fitstats[ratidx,modelidx,:]], "  "))
        end
        # averages across rats
        println(repeat(" ", modellen + 1), " ", repeat(" ", ratidlen - 3), "avg ",
                join([@sprintf("%10.2f±%6.2f", mean(fitstats[:,modelidx,i]), √(var(fitstats[:,modelidx,i]) / ratnum))
                      for i in 1:ratnum], " "))
    end
    return fitstats[:,:,2:end], fitstatsall
end

# returns the absolute and relative fit quality plots for all models for the
# given statistics
function genfitqualityplots(fitstats, fitstatsall, ratids::Vector{Int})
    fitmeasurenames = ["BIC", "AIC", "AICc"]
    ratnum = length(ratids)
    modelnum = length(modelnames)
    @assert(ratnum > 0)
    @assert(size(fitstats, 1) == ratnum)
    @assert(size(fitstats, 2) == modelnum)
    @assert(size(fitstats, 3) == length(fitmeasurenames))
    @assert(size(fitstatsall, 1) == modelnum)
    @assert(size(fitstatsall, 2) == length(fitmeasurenames))
    # separate plot per fit measure
    ps = Array{Any}(undef, length(fitmeasurenames))
    for fitmeasureidx = 1:length(fitmeasurenames)
        substats = fitstats[:,:,fitmeasureidx]
        fitmeasurename = fitmeasurenames[fitmeasureidx]
        # average statistics
        fitmean = mean(substats, dims=1)
        fitse = sqrt.(var(substats, dims=1) / ratnum)
        ps[fitmeasureidx] = plot(
            layer(x=1:modelnum, y=fitstatsall[:,fitmeasureidx] / ratnum,
                  Geom.point, Theme(t, default_color=RGB(0.8, 0.0, 0.0))),
            layer(x=1:modelnum, y=fitmean, ymin=fitmean.-fitse, ymax=fitmean.+fitse,
                  Geom.point, Geom.errorbar),
            layer(x=vec(repeat(collect(1:modelnum)', ratnum)), y=vec(substats),
                  Geom.point, Stat.x_jitter(range=0.4), Theme(t, default_color=RGB(0.5, 0.5, 0.5))),
            layer(yintercept=[0.0], Geom.hline(color="black")),
            Guide.xticks(ticks=collect(1:modelnum), orientation=:vertical),
            Scale.x_discrete(labels=(i->modelnames[i])),
            Guide.xlabel(""),
            Guide.ylabel("log model evidence (using $(fitmeasurename); log10)", orientation=:vertical),
            Coord.cartesian(ymax=0.0))
    end
    p_abs = vstack(ps...)
    # plot comparison to base model
    basemodelidx = findall(modelnames .== basemodel)
    fitdiff = fitstats[:, basemodelidx, :] .- fitstats[:, modelnames .!= basemodel, :]
    fitdiffall = fitstatsall[basemodelidx, :] .- fitstatsall[modelnames .!= basemodel, :]
    diffmodels = modelnames[modelnames .!= basemodel]
    ps = Array{Any}(undef, length(fitmeasurenames))
    for fitmeasureidx = 1:length(fitmeasurenames)
        substats = fitdiff[:,:,fitmeasureidx]
        fitmeasurename = fitmeasurenames[fitmeasureidx]
        # average statistics
        fitmean = mean(substats, dims=1)
        fitse = sqrt.(var(substats, dims=1) / ratnum)
        ps[fitmeasureidx] = plot(
            layer(x=1:(modelnum-1), y=fitdiffall[:,fitmeasureidx] / ratnum,
                  Geom.point, Theme(t, default_color=RGB(0.8, 0.0, 0.0))),
            layer(x=1:(modelnum-1), y=fitmean, ymin=fitmean.-fitse, ymax=fitmean.+fitse,
                  Geom.point, Geom.errorbar),
            layer(x=vec(repeat(collect(1:(modelnum-1))', ratnum)), y=vec(substats),
                  Geom.point, Stat.x_jitter(range=0.4), Theme(t, default_color=RGB(0.5, 0.5, 0.5))),
            layer(yintercept=[0.0], Geom.hline(color="black")),
            Guide.xticks(ticks=collect(1:(modelnum-1)), orientation=:vertical),
            Scale.x_discrete(labels=(i->diffmodels[i])),
            Guide.xlabel(""),
            Guide.ylabel("log model evidence (using $(fitmeasurename); log10)\n$(basemodel) vs. others",
                         orientation=:vertical))
    end
    p_rel = vstack(ps...)
    return p_abs, p_rel
end

# plots the fit quality for all models for the given objective function
function plotfitquality(objfn, ratids::Vector{Int})
    fitmeasures = [BIC, AIC, AICc]
    ratnum = length(ratids)
    @assert(ratnum > 0)
    # collect model fit statistics
    modelnum = length(modelnames)
    fitstats = Array{Float64}(undef, ratnum, modelnum, length(fitmeasures))
    fitstatsall = Array{Float64}(undef, modelnum, length(fitmeasures))
    for modelidx = 1:modelnum
        # first global, then for individual rats
        d = loadmodelfits(modelnames[modelidx], objfn)
        fitstatsall[modelidx, :] = [
            fitmeasure(d[:llh], length(d)-1, datatrials(objfn))
            for fitmeasure in fitmeasures]
        for ratidx = 1:ratnum
            d = loadmodelfits(modelnames[modelidx], objfn, ratids[ratidx])
            fitstats[ratidx, modelidx, :] = [
                fitmeasure(d[:llh], length(d)-1, datatrials(objfn, ratids[ratidx]))
                for fitmeasure in fitmeasures]
        end
    end
    # generate and save plots
    p_abs, p_rel = genfitqualityplots(fitstats, fitstatsall, ratids)
    draw(PDF("$figpath/fitqual_$(objfn).pdf", 8.5inch, 11inch), p_abs)
    println("Plot writting to $figpath/fitqual_$(objfn).pdf")
    draw(PDF("$figpath/fitqual_$(objfn)_relativ.pdf", 8.5inch, 11inch), p_rel)
    println("Plot written to $figpath/fitqual_$(objfn)_relativ.pdf")
end

# returns average log-likelihood for one objective function, using another
# objective function for the parameters. It returns this average
# log-likelihood and the number of trials, and the number of parameters
function avgcrossllh(paramobjfn, evalobjfn, modelname, ratid::Int=0)
    # trial sequence to fit data on
    d_ident = load_data(:ident, ratid)
    d_categ = load_data(:categ, ratid)
    # instantiate model and get model parameters
    m = modeldict[modelname]()
    ϕ = Vector(loadmodelfits(modelname, paramobjfn)[ϕnames(m)])
    ϕn = length(ϕ)
    # evaluate objective function
    evalobjfntype = objfndict[evalobjfn][1]
    evalobjfntrials = datatrials(evalobjfn, ratid)
    Random.seed!(0)
    evidence, ids = catevidence(ident_cs, categ_cs, 10, 10^5)
    o = evalobjfntype(m, evidence, ids, burnin, d_ident, d_categ, nrats)
    avgllh = llh(o, ϕ)
    return avgllh, evalobjfntrials, ϕn
end

# prints the cross fit quality, using parameters for the paramobjfn and
# evaluating it on the evalobjfn. It returns an array of fitmeasures per rat
# and model, and one per model for the meta rat (fits across all rats)
function printcrossfitquality(paramobjfn, evalobjfn, ratids::Vector{Int})
    fitmeasures = [(BIC, "BIC"), (AIC, "AIC"), (AICc, "AICc")]
    ratids = ratids[ratids .!= 0]
    ratnum = length(ratids)
    ratidlen = maximum([floor(Int, log10(ratid))+1 for ratid in ratids]) + 3
    ratidstrs = ["rat$(ratid)" for ratid in ratids]
    @assert(ratnum > 0)
    # collect model fit statistics
    modelnum = length(modelnames)
    modellen = max(maximum([length(s) for s in modelnames]), 3) # include "all" here
    fitstats = Array{Float64}(undef, ratnum, modelnum, length(fitmeasures)+1)
    fitstatsall = Array{Float64}(undef, modelnum, length(fitmeasures))
    for modelidx = 1:modelnum
        modelname = modelnames[modelidx]
        # fits across all rats
        avgllh, evalobjfntrials, ϕn = avgcrossllh(paramobjfn, evalobjfn, modelname)
        fitq = [avgllh; [f(avgllh, ϕn, evalobjfntrials) for f in [BIC, AIC, AICc]]]
        println(repeat(" ", modellen - length(modelname)), modelname, "  ",
                repeat(" ", ratidlen - 3), "all ",
                join([@sprintf("%10.2f      ", x) for x in fitq], "  "))
        fitstatsall[modelidx, :] = fitq[2:end]
        # fits for individual rats
        if length(ratids) == 0
            continue
        end
        avgllhs = Array{Float64}(undef, ratnum)
        for ratidx in 1:ratnum
            ratid = ratids[ratidx]
            avgllh, evalobjfntrials, ϕn = avgcrossllh(paramobjfn, evalobjfn, modelname, ratid)
            fitq = [avgllh; [f(avgllh, ϕn, evalobjfntrials) for f in [BIC, AIC, AICc]]]
            println(repeat(" ", modellen + 1), " ",
                    repeat(" ", ratidlen - length(ratidstrs[ratidx])), ratidstrs[ratidx], " ",
                    join([@sprintf("%10.2f      ", x) for x in fitq], "  "))
            fitstats[ratidx,modelidx,:] = fitq
        end
        # averages across rats
        println(repeat(" ", modellen + 1), " ", repeat(" ", ratidlen - 3), "avg ",
                join([@sprintf("%10.2f±%6.2f", mean(fitstats[:,modelidx,i]),
                               √(var(fitstats[:,modelidx,i]) / ratnum))
                      for i in 1:ratnum], " "))
    end
    # don't return the avgllh fit for individual rats
    return fitstats[:,:,2:end], fitstatsall
end

# writes the fit quality measures to a MAT file
function writefitquality(fitstats, fitstatsall, ratids::Vector{Int}, fname)
    fitmeasurenames = ["BIC", "AIC", "AICc"]
    ratnum = length(ratids)
    modelnum = length(modelnames)
    @assert(ratnum > 0)
    @assert(size(fitstats, 1) == ratnum)
    @assert(size(fitstats, 2) == modelnum)
    @assert(size(fitstats, 3) == length(fitmeasurenames))
    @assert(size(fitstatsall, 1) == modelnum)
    @assert(size(fitstatsall, 2) == length(fitmeasurenames))
    # write matrices to MAT file
    d = Dict(
        "fitstats" => fitstats,
        "fitstatsall" => fitstatsall,
        "fitmeasures" => fitmeasurenames,
        "modelnames" => modelnames,
        "ratids" => ratids
        )
    matwrite(fname, d)
end

# print model parameters
# for modelname in modelnames
#     println("Model $modelname")
#     printfitparams(modelname, [0, 1, 2, 3, 4])
#     println()
# end
# println()

# print model fit statistics
for objfn in objfns
    println("Objective function $objfn")
    fitstats, fitstatsall = printfitquality(objfn, [0, 1, 2, 3, 4])
    fname = "$datapath/fitqual_$(objfn).mat"
    writefitquality(fitstats, fitstatsall, [1, 2, 3, 4], fname)
    println("Data written to $fname")
    println()
end

# plot model fit stats
for objfn in objfns
    println("Generating plots for objective function $objfn")
    plotfitquality(objfn, [1, 2, 3, 4])
end

# # print model fit statistics when evaluated with psychchronseq1 objective
println("Fitted on psychchron1, evaluated on psychchronseq1")
fitstats, fitstatsall = printcrossfitquality("psychchron1", "psychchronseq1", [1, 2, 3, 4])
fname = "$datapath/fitqual_psychchron1_to_psychchronseq1.mat"
writefitquality(fitstats, fitstatsall, [1, 2, 3, 4], fname)
println("Data written to $fname")
println()
println("Generating plots")
p_abs, p_rel = genfitqualityplots(fitstats, fitstatsall, [1, 2, 3, 4])
draw(PDF("$figpath/fitqual_psychchron1_to_psychchronseq1.pdf", 8.5inch, 11inch), p_abs)
println("Plot written to $figpath/fitqual_psychchron1_to_psychchronseq1.pdf")
draw(PDF("$figpath/fitqual_psychchron1_to_psychchronseq1_relativ.pdf", 8.5inch, 11inch), p_rel)
println("Plot written to $figpath/fitqual_psychchron1_to_psychchronseq1_relativ.pdf")