"""
"""
abstract type SparseMatrixAssembler <: Assembler end

"""
"""
function get_matrix_builder(a::SparseMatrixAssembler)
  @abstractmethod
end

"""
"""
function get_vector_builder(a::SparseMatrixAssembler)
  @abstractmethod
end

"""
"""
function symbolic_loop_matrix!(A,a::SparseMatrixAssembler,matdata)
  @abstractmethod
end

"""
"""
function symbolic_loop_vector!(b,a::SparseMatrixAssembler,vecdata)
  @abstractmethod
end

"""
"""
function symbolic_loop_matrix_and_vector!(A,b,a::SparseMatrixAssembler,data)
  @abstractmethod
end

"""
"""
function numeric_loop_matrix!(A,a::SparseMatrixAssembler,matdata)
  @abstractmethod
end

"""
"""
function numeric_loop_vector!(b,a::SparseMatrixAssembler,vecdata)
  @abstractmethod
end

"""
"""
function numeric_loop_matrix_and_vector!(A,b,a::SparseMatrixAssembler,data)
  @abstractmethod
end

get_matrix_type(a::SparseMatrixAssembler) = get_array_type(get_matrix_builder(a))
get_vector_type(a::SparseMatrixAssembler) = get_array_type(get_vector_builder(a))

function allocate_vector(a::SparseMatrixAssembler,vecdata)
  v1 = nz_counter(get_vector_builder(a),(get_rows(a),))
  symbolic_loop_vector!(v1,a,vecdata)
  v2 = nz_allocation(v1)
  symbolic_loop_vector!(v2,a,vecdata)
  v3 = create_from_nz(v2)
  v3
end

function assemble_vector!(b,a::SparseMatrixAssembler,vecdata)
  fill_entries!(b,zero(eltype(b)))
  assemble_vector_add!(b,a,vecdata)
end

function assemble_vector_add!(b,a::SparseMatrixAssembler,vecdata)
  numeric_loop_vector!(b,a,vecdata)
  create_from_nz(b)
end

function assemble_vector(a::SparseMatrixAssembler,vecdata)
  v1 = nz_counter(get_vector_builder(a),(get_rows(a),))
  symbolic_loop_vector!(v1,a,vecdata)
  v2 = nz_allocation(v1)
  numeric_loop_vector!(v2,a,vecdata)
  v3 = create_from_nz(v2)
  v3
end

function allocate_matrix(a::SparseMatrixAssembler,matdata)
  m1 = nz_counter(get_matrix_builder(a),(get_rows(a),get_cols(a)))
  symbolic_loop_matrix!(m1,a,matdata)
  m2 = nz_allocation(m1)
  symbolic_loop_matrix!(m2,a,matdata)
  m3 = create_from_nz(m2)
  m3
end

function assemble_matrix!(mat,a::SparseMatrixAssembler,matdata)
  fill_entries!(mat,zero(eltype(mat)))
  assemble_matrix_add!(mat,a,matdata)
end

function assemble_matrix_add!(mat,a::SparseMatrixAssembler,matdata)
  numeric_loop_matrix!(mat,a,matdata)
  create_from_nz(mat)
end

function assemble_matrix(a::SparseMatrixAssembler,matdata)
  m1 = nz_counter(get_matrix_builder(a),(get_rows(a),get_cols(a)))
  symbolic_loop_matrix!(m1,a,matdata)
  m2 = nz_allocation(m1)
  numeric_loop_matrix!(m2,a,matdata)
  m3 = create_from_nz(m2)
  m3
end

function allocate_matrix_and_vector(a::SparseMatrixAssembler,data)
  m1 = nz_counter(get_matrix_builder(a),(get_rows(a),get_cols(a)))
  v1 = nz_counter(get_vector_builder(a),(get_rows(a),))
  symbolic_loop_matrix_and_vector!(m1,v1,a,data)
  m2 = nz_allocation(m1)
  v2 = nz_allocation(v1)
  symbolic_loop_matrix_and_vector!(m2,v2,a,data)
  m3 = create_from_nz(m2)
  v3 = create_from_nz(v2)
  m3,v3
end

function assemble_matrix_and_vector!(A,b,a::SparseMatrixAssembler, data)
  fill_entries!(A,zero(eltype(A)))
  fill_entries!(b,zero(eltype(b)))
  assemble_matrix_and_vector_add!(A,b,a,data)
end

function assemble_matrix_and_vector_add!(A,b,a::SparseMatrixAssembler,data)
  numeric_loop_matrix_and_vector!(A,b,a,data)
  create_from_nz(A), create_from_nz(b)
end

function assemble_matrix_and_vector(a::SparseMatrixAssembler, data)
  m1 = nz_counter(get_matrix_builder(a),(get_rows(a),get_cols(a)))
  v1 = nz_counter(get_vector_builder(a),(get_rows(a),))
  symbolic_loop_matrix_and_vector!(m1,v1,a,data)
  m2 = nz_allocation(m1)
  v2 = nz_allocation(v1)
  numeric_loop_matrix_and_vector!(m2,v2,a,data)
  m3 = create_from_nz(m2)
  v3 = create_from_nz(v2)
  m3,v3
end

function test_sparse_matrix_assembler(a::SparseMatrixAssembler,matdata,vecdata,data)
  test_assembler(a,matdata,vecdata,data)
  _ = get_matrix_builder(a)
  _ = get_vector_builder(a)
end

struct GenericSparseMatrixAssembler <: SparseMatrixAssembler
  matrix_builder
  vector_builder
  rows::AbstractUnitRange
  cols::AbstractUnitRange
  strategy::AssemblyStrategy
end

function SparseMatrixAssembler(
  mat,
  vec,
  trial::FESpace,
  test::FESpace,
  strategy::AssemblyStrategy=DefaultAssemblyStrategy())

  rows = get_free_dof_ids(test)
  cols = get_free_dof_ids(trial)
  GenericSparseMatrixAssembler(
    SparseMatrixBuilder(mat),
    ArrayBuilder(vec),
    rows,
    cols,
    strategy)
end

function SparseMatrixAssembler(mat,trial::FESpace,test::FESpace)
  mat_builder = SparseMatrixBuilder(mat)
  T = eltype(get_array_type(mat_builder))
  SparseMatrixAssembler(mat_builder,Vector{T},trial,test)
end

"""
"""
function SparseMatrixAssembler(trial::FESpace,test::FESpace)
  T = get_dof_value_type(trial)
  matrix_type = SparseMatrixCSC{T,Int}
  vector_type = Vector{T}
  SparseMatrixAssembler(matrix_type,vector_type,trial,test)
end

get_rows(a::GenericSparseMatrixAssembler) = a.rows

get_cols(a::GenericSparseMatrixAssembler) = a.cols

get_matrix_builder(a::GenericSparseMatrixAssembler) = a.matrix_builder

get_vector_builder(a::GenericSparseMatrixAssembler) = a.vector_builder

get_assembly_strategy(a::GenericSparseMatrixAssembler) = a.strategy

function symbolic_loop_matrix!(A,a::GenericSparseMatrixAssembler,matdata)
  get_mat(a::Tuple) = a[1]
  get_mat(a) = a
  if LoopStyle(A) == DoNotLoop()
    return A
  end
  for (cellmat,_cellidsrows,_cellidscols) in zip(matdata...)
    cellidsrows = map_cell_rows(a.strategy,_cellidsrows)
    cellidscols = map_cell_cols(a.strategy,_cellidscols)
    rows_cache = array_cache(cellidsrows)
    cols_cache = array_cache(cellidscols)
    @assert length(cellidscols) == length(cellidsrows)
    if length(cellidscols) > 0
      mat1 = get_mat(first(cellmat))
      rows1 = getindex!(rows_cache,cellidsrows,1)
      cols1 = getindex!(cols_cache,cellidscols,1)
      touch! = TouchEntriesMap()
      touch_cache = return_cache(touch!,A,mat1,rows1,cols1)
      caches = touch_cache, rows_cache, cols_cache
      _symbolic_loop_matrix!(A,caches,cellidsrows,cellidscols,mat1)
    end
  end
  A
end

@noinline function _symbolic_loop_matrix!(A,caches,cell_rows,cell_cols,mat1)
  touch_cache, rows_cache, cols_cache = caches
  touch! = TouchEntriesMap()
  for cell in 1:length(cell_cols)
    rows = getindex!(rows_cache,cell_rows,cell)
    cols = getindex!(cols_cache,cell_cols,cell)
    evaluate!(touch_cache,touch!,A,mat1,rows,cols)
  end
end

function numeric_loop_matrix!(A,a::GenericSparseMatrixAssembler,matdata)
  for (cellmat,_cellidsrows,_cellidscols) in zip(matdata...)
    cellidsrows = map_cell_rows(a.strategy,_cellidsrows)
    cellidscols = map_cell_cols(a.strategy,_cellidscols)
    rows_cache = array_cache(cellidsrows)
    cols_cache = array_cache(cellidscols)
    vals_cache = array_cache(cellmat)
    @assert length(cellidscols) == length(cellidsrows)
    @assert length(cellmat) == length(cellidsrows)
    if length(cellmat) > 0
      mat1 = getindex!(vals_cache,cellmat,1)
      rows1 = getindex!(rows_cache,cellidsrows,1)
      cols1 = getindex!(cols_cache,cellidscols,1)
      add! = AddEntriesMap(+)
      add_cache = return_cache(add!,A,mat1,rows1,cols1)
      caches = add_cache, vals_cache, rows_cache, cols_cache
      _numeric_loop_matrix!(A,caches,cellmat,cellidsrows,cellidscols)
    end
  end
  A
end

@noinline function _numeric_loop_matrix!(mat,caches,cell_vals,cell_rows,cell_cols)
  add_cache, vals_cache, rows_cache, cols_cache = caches
  add! = AddEntriesMap(+)
  for cell in 1:length(cell_cols)
    rows = getindex!(rows_cache,cell_rows,cell)
    cols = getindex!(cols_cache,cell_cols,cell)
    vals = getindex!(vals_cache,cell_vals,cell)
    evaluate!(add_cache,add!,mat,vals,rows,cols)
  end
end

function symbolic_loop_vector!(b,a::GenericSparseMatrixAssembler,vecdata)
  @notimplementedif LoopStyle(b) == Loop()
  b
end

function numeric_loop_vector!(b,a::GenericSparseMatrixAssembler,vecdata)
  for (cellvec, _cellids) in zip(vecdata...)
    cellids = map_cell_rows(a.strategy,_cellids)
    rows_cache = array_cache(cellids)
    vals_cache = array_cache(cellvec)
    if length(cellvec) > 0
      vals1 = getindex!(vals_cache,cellvec,1)
      rows1 = getindex!(rows_cache,cellids,1)
      add! = AddEntriesMap(+)
      add_cache = return_cache(add!,b,vals1,rows1)
      caches = add_cache, vals_cache, rows_cache
      _numeric_loop_vector!(b,caches,cellvec,cellids)
    end
  end
  b
end

@noinline function _numeric_loop_vector!(vec,caches,cell_vals,cell_rows)
  add_cache, vals_cache, rows_cache = caches
  @assert length(cell_vals) == length(cell_rows)
  add! = AddEntriesMap(+)
  for cell in 1:length(cell_rows)
    rows = getindex!(rows_cache,cell_rows,cell)
    vals = getindex!(vals_cache,cell_vals,cell)
    evaluate!(add_cache,add!,vec,vals,rows)
  end
end

function symbolic_loop_matrix_and_vector!(A,b,a::GenericSparseMatrixAssembler,data)
  matvecdata, matdata, vecdata = data
  symbolic_loop_matrix!(A,a,matvecdata)
  symbolic_loop_matrix!(A,a,matdata)
  symbolic_loop_vector!(b,a,vecdata)
  A, b
end

function numeric_loop_matrix_and_vector!(A,b,a::GenericSparseMatrixAssembler,data)
  matvecdata, matdata, vecdata = data
  for (cellmatvec,_cellidsrows,_cellidscols) in zip(matvecdata...)
    cellidsrows = map_cell_rows(a.strategy,_cellidsrows)
    cellidscols = map_cell_cols(a.strategy,_cellidscols)
    rows_cache = array_cache(cellidsrows)
    cols_cache = array_cache(cellidscols)
    vals_cache = array_cache(cellmatvec)
    @assert length(cellidscols) == length(cellidsrows)
    @assert length(cellmatvec) == length(cellidsrows)
    if length(cellmatvec) > 0
      mat1, vec1 = getindex!(vals_cache,cellmatvec,1)
      rows1 = getindex!(rows_cache,cellidsrows,1)
      cols1 = getindex!(cols_cache,cellidscols,1)
      add! = AddEntriesMap(+)
      add_mat_cache = return_cache(add!,A,mat1,rows1,cols1)
      add_vec_cache = return_cache(add!,b,vec1,rows1)
      caches = add_mat_cache, add_vec_cache, vals_cache, rows_cache, cols_cache
      _numeric_loop_matvec!(A,b,caches,cellmatvec,cellidsrows,cellidscols)
    end
  end
  numeric_loop_matrix!(A,a,matdata)
  numeric_loop_vector!(b,a,vecdata)
  A, b
end

@noinline function _numeric_loop_matvec!(A,b,caches,cell_vals,cell_rows,cell_cols)
  add_mat_cache, add_vec_cache, vals_cache, rows_cache, cols_cache = caches
  add! = AddEntriesMap(+)
  for cell in 1:length(cell_cols)
    rows = getindex!(rows_cache,cell_rows,cell)
    cols = getindex!(cols_cache,cell_cols,cell)
    vals = getindex!(vals_cache,cell_vals,cell)
    matvals, vecvals = vals
    evaluate!(add_mat_cache,add!,A,matvals,rows,cols)
    evaluate!(add_vec_cache,add!,b,vecvals,rows)
  end
end


import TurbulenceConvection
const TC = TurbulenceConvection
import Thermodynamics
const TD = Thermodynamics

function initialize_edmf(edmf::TC.EDMF_PrognosticTKE, grid, state, Case, gm::TC.GridMeanVariables, TS::TC.TimeStepping)
    initialize_covariance(edmf, grid, state, gm, Case)
    up = edmf.UpdVar
    param_set = TC.parameter_set(gm)
    aux_tc = TC.center_aux_turbconv(state)
    prog_gm = TC.center_prog_grid_mean(state)
    p0_c = TC.center_ref_state(state).p0
    parent(aux_tc.prandtl_nvec) .= edmf.prandtl_number
    @inbounds for k in TC.real_center_indices(grid)
        ts = TC.thermo_state_pθq(param_set, p0_c[k], prog_gm.θ_liq_ice[k], prog_gm.q_tot[k])
        aux_tc.θ_virt[k] = TD.virtual_pottemp(ts)
    end
    TC.update_surface(Case, grid, state, gm, TS, param_set)
    TC.compute_updraft_surface_bc(edmf, grid, state, Case)
    if Case.casename == "DryBubble"
        initialize_updrafts_DryBubble(edmf, grid, state, up, gm)
    else
        initialize_updrafts(edmf, grid, state, up, gm)
    end
    TC.set_edmf_surface_bc(edmf, grid, state, up, Case.Sur)
    return
end

function initialize_covariance(edmf::TC.EDMF_PrognosticTKE, grid, state, gm, Case)

    kc_surf = TC.kc_surface(grid)
    en = edmf.EnvVar
    aux_gm = TC.center_aux_grid_mean(state)
    prog_en = TC.center_prog_environment(state)
    aux_en = TC.center_aux_environment(state)
    ρ0_c = TC.center_ref_state(state).ρ0
    aux_bulk = TC.center_aux_bulk(state)
    ae = 1 .- aux_bulk.area # area of environment

    aux_en.tke .= aux_gm.tke
    prog_en.ρatke .= aux_en.tke .* ρ0_c .* ae

    TC.get_GMV_CoVar(edmf, grid, state, :tke, :w)
    aux_gm.Hvar .= aux_gm.Hvar[kc_surf] .* aux_gm.tke
    aux_gm.QTvar .= aux_gm.QTvar[kc_surf] .* aux_gm.tke
    aux_gm.HQTcov .= aux_gm.HQTcov[kc_surf] .* aux_gm.tke

    prog_en.ρaHvar .= aux_gm.Hvar .* ρ0_c .* ae
    prog_en.ρaQTvar .= aux_gm.QTvar .* ρ0_c .* ae
    prog_en.ρaHQTcov .= aux_gm.HQTcov .* ρ0_c .* ae
    return
end

function initialize_updrafts(edmf, grid, state, up::TC.UpdraftVariables, gm::TC.GridMeanVariables)
    kc_surf = TC.kc_surface(grid)
    aux_up = TC.center_aux_updrafts(state)
    prog_gm = TC.center_prog_grid_mean(state)
    aux_up = TC.center_aux_updrafts(state)
    aux_up_f = TC.face_aux_updrafts(state)
    aux_gm = TC.center_aux_grid_mean(state)
    prog_up = TC.center_prog_updrafts(state)
    prog_up_f = TC.face_prog_updrafts(state)
    ρ0_c = TC.center_ref_state(state).ρ0
    @inbounds for i in 1:(up.n_updrafts)
        @inbounds for k in TC.real_face_indices(grid)
            aux_up_f[i].w[k] = 0
            prog_up_f[i].ρaw[k] = 0
        end

        @inbounds for k in TC.real_center_indices(grid)
            aux_up[i].buoy[k] = 0
            # Simple treatment for now, revise when multiple updraft closures
            # become more well defined
            aux_up[i].area[k] = 0
            aux_up[i].q_tot[k] = prog_gm.q_tot[k]
            aux_up[i].θ_liq_ice[k] = prog_gm.θ_liq_ice[k]
            aux_up[i].q_liq[k] = aux_gm.q_liq[k]
            aux_up[i].q_ice[k] = aux_gm.q_ice[k]
            aux_up[i].T[k] = aux_gm.T[k]
            prog_up[i].ρarea[k] = 0
            prog_up[i].ρaq_tot[k] = 0
            prog_up[i].ρaθ_liq_ice[k] = 0
        end

        aux_up[i].area[kc_surf] = edmf.area_surface_bc[i]
        prog_up[i].ρarea[kc_surf] = ρ0_c[kc_surf] * aux_up[i].area[kc_surf]
    end
    return
end

function initialize_updrafts_DryBubble(edmf, grid, state, up::TC.UpdraftVariables, gm::TC.GridMeanVariables)
    dz = grid.Δz

    # criterion 2: b>1e-4
    #! format: off
    z_in = [
          75.,  125.,  175.,  225.,  275.,  325.,  375.,  425.,  475.,
         525.,  575.,  625.,  675.,  725.,  775.,  825.,  875.,  925.,
         975., 1025., 1075., 1125., 1175., 1225., 1275., 1325., 1375.,
        1425., 1475., 1525., 1575., 1625., 1675., 1725., 1775., 1825.,
        1875., 1925., 1975., 2025., 2075., 2125., 2175., 2225., 2275.,
        2325., 2375., 2425., 2475., 2525., 2575., 2625., 2675., 2725.,
        2775., 2825., 2875., 2925., 2975., 3025., 3075., 3125., 3175.,
        3225., 3275., 3325., 3375., 3425., 3475., 3525., 3575., 3625.,
        3675., 3725., 3775., 3825., 3875., 3925.]

    θ_liq_in = [
        299.9882, 299.996 , 300.0063, 300.0205, 300.04  , 300.0594,
        300.0848, 300.1131, 300.1438, 300.1766, 300.2198, 300.2567,
        300.2946, 300.3452, 300.3849, 300.4245, 300.4791, 300.5182,
        300.574 , 300.6305, 300.6668, 300.7222, 300.7771, 300.8074,
        300.8591, 300.9092, 300.9574, 300.9758, 301.0182, 301.0579,
        301.0944, 301.1276, 301.1572, 301.1515, 301.1729, 301.1902,
        301.2033, 301.2122, 301.2167, 301.2169, 301.2127, 301.2041,
        301.1913, 301.1743, 301.1533, 301.1593, 301.1299, 301.097 ,
        301.0606, 301.0212, 300.9788, 300.9607, 300.9125, 300.8625,
        300.8108, 300.7806, 300.7256, 300.6701, 300.6338, 300.5772,
        300.5212, 300.482 , 300.4272, 300.3875, 300.3354, 300.2968,
        300.2587, 300.2216, 300.1782, 300.1452, 300.1143, 300.0859,
        300.0603, 300.0408, 300.0211, 300.0067, 299.9963, 299.9884    ]

    Area_in = [
        0.04 , 0.055, 0.07 , 0.08 , 0.085, 0.095, 0.1  , 0.105, 0.11 ,
        0.115, 0.115, 0.12 , 0.125, 0.125, 0.13 , 0.135, 0.135, 0.14 ,
        0.14 , 0.14 , 0.145, 0.145, 0.145, 0.15 , 0.15 , 0.15 , 0.15 ,
        0.155, 0.155, 0.155, 0.155, 0.155, 0.155, 0.16 , 0.16 , 0.16 ,
        0.16 , 0.16 , 0.16 , 0.16 , 0.16 , 0.16 , 0.16 , 0.16 , 0.16 ,
        0.155, 0.155, 0.155, 0.155, 0.155, 0.155, 0.15 , 0.15 , 0.15 ,
        0.15 , 0.145, 0.145, 0.145, 0.14 , 0.14 , 0.14 , 0.135, 0.135,
        0.13 , 0.13 , 0.125, 0.12 , 0.115, 0.115, 0.11 , 0.105, 0.1  ,
        0.095, 0.085, 0.08 , 0.07 , 0.055, 0.04    ]

    W_in = [
        0.017 , 0.0266, 0.0344, 0.0417, 0.0495, 0.0546, 0.061 , 0.0668,
        0.0721, 0.0768, 0.0849, 0.0887, 0.092 , 0.0996, 0.1019, 0.1037,
        0.1106, 0.1114, 0.1179, 0.1243, 0.1238, 0.1297, 0.1355, 0.1335,
        0.1387, 0.1437, 0.1485, 0.1448, 0.1489, 0.1527, 0.1564, 0.1597,
        0.1628, 0.1565, 0.1588, 0.1609, 0.1626, 0.1641, 0.1652, 0.166 ,
        0.1665, 0.1667, 0.1666, 0.1662, 0.1655, 0.1736, 0.1722, 0.1706,
        0.1686, 0.1664, 0.1639, 0.1698, 0.1667, 0.1634, 0.1599, 0.1641,
        0.1601, 0.1559, 0.1589, 0.1543, 0.1496, 0.1514, 0.1464, 0.1475,
        0.1422, 0.1425, 0.1424, 0.1419, 0.1361, 0.135 , 0.1335, 0.1316,
        0.1294, 0.1302, 0.1271, 0.1264, 0.1269, 0.1256    ]

    T_in = [
        299.2557, 298.775 , 298.2969, 297.8227, 297.3536, 296.8843,
        296.421 , 295.9603, 295.502 , 295.0456, 294.5994, 294.1468,
        293.6951, 293.2556, 292.8054, 292.3549, 291.9188, 291.4677,
        291.0325, 290.5978, 290.1434, 289.7073, 289.2706, 288.81  ,
        288.3698, 287.928 , 287.4842, 287.0118, 286.5622, 286.1099,
        285.6544, 285.1957, 284.7335, 284.2379, 283.7677, 283.2937,
        282.8157, 282.3337, 281.8476, 281.3574, 280.8631, 280.3649,
        279.8626, 279.3565, 278.8467, 278.362 , 277.8447, 277.3241,
        276.8006, 276.2742, 275.7454, 275.2388, 274.705 , 274.1694,
        273.6327, 273.1155, 272.576 , 272.0363, 271.514 , 270.9736,
        270.4339, 269.9094, 269.3711, 268.8465, 268.311 , 267.7877,
        267.2649, 266.7432, 266.2159, 265.698 , 265.1821, 264.6685,
        264.1574, 263.6518, 263.1461, 262.6451, 262.1476, 261.6524]
    #! format: on

    aux_up = TC.center_aux_updrafts(state)
    aux_up_f = TC.face_aux_updrafts(state)
    aux_gm = TC.center_aux_grid_mean(state)
    prog_gm = TC.center_prog_grid_mean(state)
    prog_up = TC.center_prog_updrafts(state)
    ρ_0_c = TC.center_ref_state(state).ρ0
    ρ_0_f = TC.face_ref_state(state).ρ0
    Area_in = TC.pyinterp(grid.zc, z_in, Area_in)
    θ_liq_in = TC.pyinterp(grid.zc, z_in, θ_liq_in)
    T_in = TC.pyinterp(grid.zc, z_in, T_in)
    @inbounds for i in 1:(up.n_updrafts)
        @inbounds for k in TC.real_face_indices(grid)
            if minimum(z_in) <= grid.zf[k] <= maximum(z_in)
                aux_up_f[i].w[k] = 0.0
            end
        end

        @inbounds for k in TC.real_center_indices(grid)
            if minimum(z_in) <= grid.zc[k] <= maximum(z_in)
                aux_up[i].area[k] = Area_in[k]
                aux_up[i].θ_liq_ice[k] = θ_liq_in[k]
                aux_up[i].q_tot[k] = 0.0
                aux_up[i].q_liq[k] = 0.0
                aux_up[i].q_ice[k] = 0.0

                # for now temperature is provided as diagnostics from LES
                aux_up[i].T[k] = T_in[k]
                prog_up[i].ρarea[k] = ρ_0_c[k] * aux_up[i].area[k]
                prog_up[i].ρaθ_liq_ice[k] = prog_up[i].ρarea[k] * aux_up[i].θ_liq_ice[k]
                prog_up[i].ρaq_tot[k] = prog_up[i].ρarea[k] * aux_up[i].q_tot[k]
            else
                aux_up[i].area[k] = 0.0
                aux_up[i].θ_liq_ice[k] = prog_gm.θ_liq_ice[k]
                aux_up[i].T[k] = aux_gm.T[k]
                prog_up[i].ρarea[k] = 0.0
                prog_up[i].ρaθ_liq_ice[k] = 0.0
                prog_up[i].ρaq_tot[k] = 0.0
            end
        end
    end
    return
end


module DataArrays
    using Base: promote_op
    using Base.Cartesian, Reexport
    @reexport using StatsBase
    @reexport using Missings
    using SpecialFunctions
    using Compat: AbstractRange, Nothing, Cvoid, uninitialized, invpermute!
    using Printf, Dates

    const DEFAULT_POOLED_REF_TYPE = UInt32

    import Base: ==, !=, >, <, >=, <=, +, -, *, !, &, |, ⊻, ^, /

    import StatsBase: autocor, inverse_rle, rle

    import Base.replace!

    export @data,
           @pdata,
           AbstractDataArray,
           AbstractDataMatrix,
           AbstractDataVector,
           array,
           autocor,
           compact,
           cut,
           data,
           DataArray,
           DataMatrix,
           DataVector,
           FastPerm,
           getpoolidx,
           gl,
           padmissing,
           pdata,
           PooledDataArray,
           PooledDataMatrix,
           PooledDataVecs,
           PooledDataVector,
           reorder,
           rep,
           replace!,
           setlevels!,
           setlevels

    include("utils.jl")
    include("abstractdataarray.jl")
    include("dataarray.jl")
    include("pooleddataarray.jl")
    include("datavector.jl")
    include("indexing.jl")
    include("datamatrix.jl")
    include("linalg.jl")
    include("operators.jl")
    include("reduce.jl")
    include("reducedim.jl")
    include("broadcast.jl")
    include("sort.jl")
    include("extras.jl")
    include("grouping.jl")
    include("statistics.jl")
    include("literals.jl")
    include("deprecated.jl")
end


using .Threads

function copyVerts!(b, sub, start)
    for i in 1:length(sub)
        b[start + i - 1] = sub[i]
    end
end

s1 = 1:nthreads()
s2 = s1 .+ 1
c1 = cumsum(s1)
c2 = cumsum(s2)
sa = [c1[i]:c2[i] for i in s1]
b = zeros(Int,maximum(c2))

@sync for i = 1:nthreads()
    Threads.@spawn copyVerts!(b, sa[i], c1[i])
end


module Unums

export
  AbstractUnum,

  FixedUnum64,
  BinaryUnum64,
  Unum64,
  mask64,

  FixedUnum16,
  BinaryUnum16,
  Unum16,
  mask16,

  isexact,
  isinexact,

  createmask

import Base: isequal, isfinite, isinf, isinteger,
             isless, isnan, isnull, isnumber, isreal, issubnormal,
             typemin, typemax, box, unbox

import Base: +, -, *, /, ==, <, >, <=, >=, !, &, |, <<, ~

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

include("types.jl")
include("utils.jl")
include("ops.jl")
include("convert/float.jl")
# include("convert/int.jl")

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

end # module


function inner_craft(its::Vector,lin=3,col=3)
	# 载入
	mat=Matrix{String}(undef,lin,col) # 辨识符
	filln=0x0 # 非空数
	for i in 1:lin
		for j in 1:col
			k=i*(col-1)+j
			mat[i,j]=its[k]
			if !isa(mat[i,j],EI)
				filln+=1
			end
		end
	end
# 左上角(mov_i,mov_j)
#=
	fl=false;mov_i=1;mov_j=1
	while mov_i<lin
		for j in 1:col
			if !isa(mat[mov_lin,j],EI)
				fl=true
				break
			end
		end
		if fl break end
		mov_i+=1
	end
	fl=false
	while mov_i<col
		for i in mov_j:lin
			if !isa(mat[i,mov_j],EI)
				fl=true
				break
			end
		end
		if fl break end
		mov_j+=1
	end
=#
	pattern=Vector{UInt8}()
	das=Vector{Item}(undef,filln)
	ids=Vector{String}(undef,filln)
	nas=Vector{String}(undef,filln)
	cou=1
	for i in 1:lin
		v=0
		for j in 1:col
			if !isa(mat[i,j],EI)
				das[cou]=mat[i,j]
				ids[cou]=id(mat[i,j])
				nas[cou]=mat[i,j].name.name
				v+=1<<(j-1)
				cou+=1
			end
		end
		if v!=0
			push!(pattern,v)
		end
	end
	res=Vector{Pair{Item,UInt8}}()
	# 检查普通合成表
	# [[3,1],["a1","a2","a3"],"a",2]
	for c in data[:普通合成表]
		if c[1]==pattern&&c[2]==ids
			push!(res,Pair(guess_struct(c[3]),length(c)==4 ? c[4] : 1))
		end
	end
	# 检查类合成表
	# [[3,1],[]]
	for c in data[:类合成表]
		if c[1]!=pattern
			continue
		end
		fl=false
		tarr=Vector{Union{UInt8,Nothing}}(nothing,4)
		for k in 1:length(c[2])>>1
			reqid::UInt8=c[2][k<<1-1]
			req=c[2][k<<1]
			if req!=nas[k]
				fl=true
				break
			elseif reqid!=0x0
				if tarr[reqid]===nothing
					tarr[reqid]=das[k].t
				elseif tarr[reqid]!=das[k].t
					fl=true
					break
				end
			end
		end
		if !fl
			str=c[3]
			s=replace(str,"\$"=>tarr[1])
			push!(res,Pair(guess_struct(s),length(c)==3 ? 0x1 : c[4]))
		end
	end
	return res
end


function t_interfaces()

N = 10	# variables
M = 25  # samples

vectordata = rand(M)
matrixdata = rand(N,M)
labels = round.(2*rand(M))
multilabels = round.(2*rand(2,M))

A01 = j4pr.datacell(vectordata)
A02 = j4pr.datacell(vectordata, labels)
A03 = j4pr.datacell(vectordata, multilabels)
A04 = j4pr.datacell(matrixdata)
A05 = j4pr.datacell(matrixdata, labels)
A06 = j4pr.datacell(matrixdata, multilabels)
AV = [A01, A02, A03, A04, A05, A06]

# Simple getindex type stability checks
#print("Checking getindex type stability for data cells... ")

# Index for ::Void and empty index
Test.@test nothing[1] isa Void
Test.@test nothing[:,1] isa Void
Test.@test nothing[:,[1,2]] isa Void
Test.@test nothing[:,1:2] isa Void

# Index with single integer (works for all data types e.g. vector and matrix)
for A in AV
	Test.@test A[5] isa typeof(A)
end

# Integer + Colon (works only for matrix data)
for A in AV[4:end]
	Test.@test A[5,:] isa typeof(A)
end

# Colon + Integer (works only for matrix data)
for A in AV[4:end]
	Test.@test A[:,5] isa typeof(A)
end

# Range/Vector (works only for vector data)
for A in AV[1:3]
	Test.@test A[[1,2,3]] isa typeof(A)
	Test.@test A[1:3] isa typeof(A)
end

# Range/Vector + Colon (works only for matrix data)
for A in AV[4:end]
	Test.@test A[[1,2,3],:] isa typeof(A)
	Test.@test A[1:3,:] isa typeof(A)
end

# Colon + Range/Vector (works only for matrix data)
for A in AV[4:end]
	Test.@test A[:,[1,2,3]] isa typeof(A)
	Test.@test A[:,1:3] isa typeof(A)
end
# Range/Vector + Range Vector (works only for matrix data)
for A in AV[4:end]
	Test.@test A[[1,2,3],[1,2,3]] isa typeof(A)
	Test.@test A[1:3,[1,2,3]] isa typeof(A)
	Test.@test A[1:3,[1,2,3]] isa typeof(A)
	Test.@test A[1:3,1:3] isa typeof(A)
end
#println("PASSED")

# Simple getindex type stability checks
#print("Checking setindex! for vector and matrix data... ")

Xvec = rand(3)
Xmat = rand(2,2)

for A in AV[1:3]
	A[1:length(Xvec)] = Xvec
	Test.@test A.x[1:length(Xvec)] == Xvec
end

for A in AV[4:end]
	A[1:size(Xmat,1),1:size(Xmat,2)] = Xmat
	Test.@test A.x[1:size(Xmat,1),1:size(Xmat,2)]  == Xmat
end
#println("PASSED")

# Simple pipe iteration checks
W1 = j4pr.ohenc("binary")
W2 = j4pr.scaler!(Dict(1=>"mean"))
W3 = j4pr.functioncell((x)->x)

PT = [W1;W2]
PP=[W1 W2]
PS = W1+W2

PC = [PT [PS;PP]]

#print("Checking pipe iteration... ")
for P in [PT, PP, PS]
	i = 1
	for W in P
		Test.@test W == P.x[i]
		i+=1
	end
end

for (i,P) in enumerate(PC)
	Test.@test P == PC.x[i]
end
#println("PASSED")

end


# Below, the symbols of GeometryPrimitives such as KDTree, bounds, surfpt_nearby, normal do
# not have to be re-defined here to be exported.  When GeometryPrimitives is imported to
# MaxwellBase by "using GeometryPrimitives", the symbols of GeometryPrimitives are already
# inside the namespace of MaxwellBase, so we simply need to export them to use outside
# MaxwellBase without qualification.
#
# One caveat, though.  In this file, bounds is defined for Interval, so if we simply export
# bounds, then only this is exported.  So, when defining bounds for Interval, we have to
# "extend" bounds for Shape by defining it as GeometryPrimitives.bounds(::Interval) = ....
# Then, exporting bounds exports this whole collection of bounds, both for Shape and Interval.
export Object
export shape, matparam, max∆l, sub_pind2matprm, add_obj!, periodize  #, surfpt_nearby, normal
# export lsf, bound_, L_, center_, dist2bound, bound_contains, ∆lmax, sphere, transform,
#     surfnormal, surfpoint  # functions
# import Base:size, getindex, contains, isless, union, intersect

mutable struct Object{K,S<:Shape{K}}  # K: dimension of space
    shp::S
    mat::Material
    ∆lmax::SVec{K,Float}
    Object{K,S}(shp, mat, ∆lmax) where {K,S} = new(shp, mat, ∆lmax)
end

Object(shp::S, mat::Material, ∆lmax::SVec{K}=@SVector(fill(Inf,K))) where {K,S<:Shape{K}} =
    Object{K,S}(shp, mat, ∆lmax)
Object(shp::Shape{K}, mat::Material, ∆lmax::AbsVec) where {K} = Object(shp, mat, SVec{K}(∆lmax))
Object(shp::Shape{K}, mat::Material, ∆lmax::Real) where {K} = Object(shp, mat, SVec(ntuple(k->∆lmax, Val(K))))

# Create the user interface such that maxwellsys.add(shape, material, ∆lmax) uses this function.
# setmat!(obj::Object, m::Material) = (obj.mat = m; obj)
# setmax∆l!(obj::Object{K}, ∆lmax::AbsVec) where {K} = (obj.∆lmax = SVec{K}(∆lmax); obj)
# setmax∆l!(obj::Object{K}, ∆lmax::Number) where {K} = setmax∆l!(obj, SVec(ntuple(k->∆lmax, Val(K))))

shape(obj::Object) = obj.shp
matparam(obj::Object, ft::FieldType) = matparam(obj.mat, ft)
max∆l(obj::Object) = obj.∆lmax

# Relationships between pind2matprm, oind2pind, oind2shp
#
# - pind2matprm is a vector of material parameter tensors.  pind2matprm[n] is a 3×3 material
# parameter tensor for the parameter index n.
#
# - oind2pind is a map from the object index to the parameter index.
#
# - oind2shp is a map from the object index to the shape.

# Returns a vector of submatrices of material parameter tensors.
# We return a vector of submatrices rather than a vector of submatrix views, because
# assign_param!() requires the entries of its argument pind2matprm to be SMatrix rather than
# views of SMatrix.
sub_pind2matprm(pind2matprm::AbsVec{S²ComplexF{Kf,Kf²}}, inds::SInteger) where {Kf,Kf²} =
    [mp[inds,inds] for mp = pind2matprm]  # mp[inds,inds]: S²Mat

# Consider using resize! on oind2obj.
function add_obj!(oind2shp::AbsVec{Shape{K,K²}}, oind2pind::Tuple2{AbsVec{ParamInd}},
                  pind2matprm::Tuple{AbsVec{S²ComplexF3},AbsVec{S²ComplexF3}},
                  objs::AbsVec{Object{K,<:Shape{K,K²}}}
                  ) where {K,K²}
    for obj = objs
        add_obj!(oind2obj, oind2pind, pind2matprm, obj)
    end
end

# Consider using resize! on oind2obj.
function add_obj!(oind2shp::AbsVec{Shape{K,K²}}, oind2pind::Tuple2{AbsVec{ParamInd}},
                  pind2matprm::Tuple{AbsVec{S²ComplexF3},AbsVec{S²ComplexF3}},
                  objs::Object{K,<:Shape{K,K²}}...
                  ) where {K,K²}
    for obj = objs
        add_obj!(oind2shp, oind2pind, pind2matprm, obj)
    end
end

# Populate `oind2shp`, `oind2pind`, `pind2matprm` with a given object.  See above for the
# meaning of `oind2shp`, `oind2pind`, `pind2matprm`.  In doing so, figure out the unique
# indices for the material parameters and use them as pind in `oind2pind` and `pind2matprm`.
#
# When I put an periodic array of objects, some object could be across a boundary.  In that
# case, the rule is to put two different objects at the negative-end and positive-end
# boundaries.  Make sure, though, the two objects have the same surface normal at the
# boundary points, so that the surface normal calculation is not ambiguous there.
#
# Note that pind2matprm stores the same material parameter "value" only once.  Therefore,
# if two objects are created with different Material instances with the same material
# parameter tensor, then the two objects' materials are assigned with the same pind.
function add_obj!(oind2shp::AbsVec{Shape{K,K²}},  # initially empty vector
                  oind2pind::Tuple2{AbsVec{ParamInd}},  # tuple of initially empty vectors
                  pind2matprm::Tuple2{AbsVec{S²ComplexF3}},  # tuple of initially empty vectors
                  obj::Object{K,<:Shape{K,K²}}) where {K,K²}
    push!(oind2shp, shape(obj))  # append obj (for potential use of oind2obj with KDTree, must use pushfirst! to prepend)

    # Assign the material parameter indices to obj.
    # If the material is aleady used in some object that was previously added, then use the
    # index already assigned to that material.
    for ft = EH
        nft = Int(ft)

        mp, on2pn, pn2mp = matparam(obj,ft), oind2pind[nft], pind2matprm[nft]
        pind = findlast(x -> x==mp, pn2mp)  # number if mp is already in pn2mp
        pind==nothing && (push!(pn2mp, mp); pind = length(pn2mp))  # update pind if mp is not in pn2mp
        push!(on2pn, pind)
    end

    return nothing
end

function GeometryPrimitives.periodize(obj::Object{K}, A::AbsMat, ∆range::Shape{K}) where {K}
    shp_array = periodize(obj.shp, A, ∆range)
    N = length(shp_array)
    obj_array = Vector{Object{K}}(N)
    for n = 1:N
        obj_array[n] = Object(shp_array[n], obj.mat, obj.∆lmax)
    end

    return obj_array
end

# The following includes operation on shapes.  These will need to be defined in GeometryPrimitives.
# Then, I will need to define the same functions for Object that delegate the operations to
# Shape's functions.
#
# See also test/object.jl

# # distance to the boundary.  If I want to get 0 for internal points, compare
# # dist(l,c) = (distance between l and center), and if it is less than L_(i) / 2,
# # return zero; otherwise return dist(l,c).
# @inline dist2bound(i::Interval1D, l::Real) = (a1 = abs(i.bound[nN]-l)) < (a2 = abs(i.bound[nP]-l)) ? a1 : a2

# # Functions common to all Shape's
#

# # contains
# # Sometimes, we only have level-set functions rather than Shape instances.
# @inline contains(lsf::Function, l::Tuple3{Real}, isinclusive::Bool=true) = isinclusive ? lsf(l) ≥ 0 : lsf(l) > 0
# @inline contains(s::Shape, l::Tuple3{Real}, isinclusive::Bool=true) = contains(lsf(s), l, isinclusive)
#
# # Other functions
# @inline bound_(s::Shape) = bound_(s.cbox)
# @inline bound_(s::Shape, w::Axis) = bound_(s.cbox, w)
# @inline bound_(s::Shape, w::Axis, sn::Sign) = bound_(s.cbox, w, sn)
# @inline bound_contains(s::Shape, l::Tuple3{Real}, isinclusive::Bool=true) = contains(s.cbox, l, isinclusive)
# @inline L_(s::Shape) = L_(s.cbox)  # scalar
# @inline center_(s::Shape) = center_(s.cbox)  # scalar
#
# # Level-set function manipulation
# @inline transform(lsf::Function, σ::Tuple3{Real}, ∆::Tuple3{Real}) = l -> lsf(σ .* (l .- ∆))  # σ[w] = 1 for shift, σ[w] = -1 for flip
#
# @inline union(lsf::Function...) = l -> maximum(map(f->f(l), lsf))
# @inline union(lsf1::Function, lsf2::Function) = lsf1==lsf2 ? lsf1 : l->max(lsf1(l), lsf2(l))
#
# @inline intersect(lsf::Function...) = l -> minimum(map(f->f(l), lsf))
#
# # Numerical gradient, and closest boundary point to a given point (needed for subpixel
# # smoothing of material parameters)
#
# # Calculates the direction normal to the curved surface crossing a voxel by numerical gradient
# # of the level-set function describing the surface.  The gradient is evaluated at a given point
# # inside the voxel.
# function surfnormal(lsf::Function, r::Tuple3{Float}, vxl::Tuple32{Float})
#     # println("lsf((0,0,0)) = $(lsf((0,0,0))),
#     #     lsf((-1,0,0)) = $(lsf((-1,0,0))), lsf((1,0,0)) = $(lsf((1,0,0))),
#     #     lsf((0,-1,0)) = $(lsf((0,-1,0))), lsf((0,1,0)) = $(lsf((0,1,0))),
#     #     lsf((0,0,-1)) = $(lsf((0,0,-1))), lsf((0,0,1)) = $(lsf((0,0,1))),
#     #     lsf((-1,-1,-1)) = $(lsf((-1,-1,-1))), lsf((1,-1,-1)) = $(lsf((1,-1,-1))),
#     #     lsf((-1,1,-1)) = $(lsf((-1,1,-1))), lsf((1,1,-1)) = $(lsf((1,1,-1))),
#     #     lsf((-1,-1,1)) = $(lsf((-1,-1,1))), lsf((1,-1,1)) = $(lsf((1,-1,1))),
#     #     lsf((-1,1,1)) = $(lsf((-1,1,1))), lsf((1,1,1)) = $(lsf((1,1,1))),
#     #     ")
#     lsf2(s) = 1-hypot(s[1], s[2], s[3])
#     # println("lsf($r) = $(lsf(r))")
#     # println("lsf2($r) = $(lsf2(r))")
#     # println("r = $r, vxl = $vxl")
#     α = 5e-4  # 1e-3 / 2
#
#     ∆x, ∆y, ∆z = α * (vxl[nX][nP] - vxl[nX][nN]), α * (vxl[nY][nP] - vxl[nY][nN]), α * (vxl[nZ][nP] - vxl[nZ][nN])
#     # ∆x, ∆y, ∆z = α .* (t_ind(vxl, nP, nP, nP) .- t_ind(vxl, nN, nN, nN))
#     rx, ry, rz = r
#
#     # println("(∆x, ∆y, ∆z) = $((∆x, ∆y, ∆z))")
#     # println("(rx, ry, rz) = $((rx, ry, rz))")
#     # # println("lsf(($rx, $ry, $rz)) = $(lsf((rx,ry,rz)))")
#     fxn::Float, fxp::Float = lsf((rx-∆x, ry, rz)), lsf((rx+∆x, ry, rz))
#     fyn::Float, fyp::Float = lsf((rx, ry-∆y, rz)), lsf((rx, ry+∆y, rz))
#     fzn::Float, fzp::Float = lsf((rx, ry, rz-∆z)), lsf((rx, ry, rz+∆z))
#     # println("((fxn, fxp), (fyn, fyp), (fzn, fzp)) = $(((fxn, fxp), (fyn, fyp), (fzn, fzp)))")
#
#     fxn2::Float, fxp2::Float = lsf2((rx-∆x, ry, rz)), lsf2((rx+∆x, ry, rz))
#     fyn2::Float, fyp2::Float = lsf2((rx, ry-∆y, rz)), lsf2((rx, ry+∆y, rz))
#     fzn2::Float, fzp2::Float = lsf2((rx, ry, rz-∆z)), lsf2((rx, ry, rz+∆z))
#     # println("((fxn2, fxp2), (fyn2, fyp2), (fzn2, fzp2)) = $(((fxn2, fxp2), (fyn2, fyp2), (fzn2, fzp2)))")
#
#     neg∂f∂x, neg∂f∂y, neg∂f∂z = (fxn - fxp) / ∆x, (fyn - fyp) / ∆y, (fzn - fzp) / ∆z
#
#     return neg∂f∂x, neg∂f∂y, neg∂f∂z  # outward normal is negative gradient
# end
#
# # Find the point on a curved surface by moving from a point along a direction.
# function surfpoint(lsf::Function, r::Tuple3{Float}, v::Tuple3{Float})
#     # println("r = $r, v = $v")
#     f(t) = lsf(r .+ t.*v)
#     tsol, isconverged = newtsol(0., f)
#
#     p = (r[nX] + tsol*v[nX], r[nY] + tsol*v[nY], r[nZ] + tsol*v[nZ])
#
#     return p, isconverged
#     # return r .+ tsol.*v
# end


"""
The submodule `MLLabelUtils.LabelEnc` contains a selection of
popular label encodings:

Strictly binary label encodings:

- [`LabelEnc.FuzzyBinary`](@ref)
- [`LabelEnc.TrueFalse`](@ref)
- [`LabelEnc.ZeroOne`](@ref)
- [`LabelEnc.MarginBased`](@ref)

Multiclass label encodings:

- [`LabelEnc.Indices`](@ref)
- [`LabelEnc.NativeLabels`](@ref)
- [`LabelEnc.OneOfK`](@ref)

Multiclass to binary:

- [`LabelEnc.OneVsRest`](@ref)
"""
module LabelEnc
    import ..MLLabelUtils.LabelEncoding
    import ..MLLabelUtils.BinaryLabelEncoding

    """
        LabelEnc.FuzzyBinary <: LabelEncoding{Any,2,1}

    A vector-based binary label encoding without a specific labeltype.

    Primarily used for fuzzy comparision of binary true targets
    and predicted targets. It basically assumes that the encoding
    is either `TrueFalse`, `ZeroOne`, or `MarginBased` by treating
    all non-negative values as positive outputs.
    """
    struct FuzzyBinary <: BinaryLabelEncoding{Any,1} end

    """
        LabelEnc.TrueFalse <: LabelEncoding{Bool,2,1}

    A vector-based binary label encoding with labeltype `Bool`.

    Represents binary classification-labels as boolean, in which `true`
    represents the positive class, and `false` the negative class.
    """
    struct TrueFalse <: BinaryLabelEncoding{Bool,1} end

    """
        LabelEnc.ZeroOne{T<:Number} <: LabelEncoding{T,2,1}

    A vector-based binary label encoding with a numeric labeltype.

    Represents binary classification-labels as numbers, in which
    `one(T)` represents the positive class, and `zero(T)` the
    negative class.
    """
    struct ZeroOne{T<:Number,R<:Number} <: BinaryLabelEncoding{T,1}
        cutoff::R
        function ZeroOne{T,R}(cutoff::R = R(0.5)) where {T,R}
            @assert 0 <= cutoff <= 1
            new{T,R}(cutoff)
        end
    end
    ZeroOne(t::Type{T} = Float64, cutoff::R = 0.5) where {T<:Number,R<:Number} = ZeroOne{T,R}(cutoff)
    ZeroOne(cutoff::R) where {R<:Number} = ZeroOne(R, cutoff)

    """
        LabelEnc.MarginBased{T<:Number} <: LabelEncoding{T,2,1}

    A vector-based binary label encoding with a numeric labeltype.

    Represents binary classification-labels as numbers, in which
    `one(T)` represents the positive class, and `-one(T)` the
    negative class.
    """
    struct MarginBased{T<:Number} <: BinaryLabelEncoding{T,1} end
    MarginBased(::Type{T} = Float64) where {T<:Number} = MarginBased{T}()

    """
        LabelEnc.OneVsRest{T} <: LabelEncoding{T,2,1}

    A vector-based binary label encoding with a arbitrary labeltype.

    A special label-encoding in that it only uses a positive
    label to determine which class some element belongs to.
    If some value matches the specified positive label, then it is
    considered positive, otherwise it is considered negative.

    That said `OneVsRest` requires a negative label to be specified
    in order to be able to denote it somehow. Note, however,
    that the specified negative label is purely for asthetic
    reasons and is not used to determine class membership
    """
    struct OneVsRest{T} <: BinaryLabelEncoding{T,1}
        poslabel::T
        neglabel::T
    end
    OneVsRest(poslabel::Bool) = OneVsRest{Bool}(poslabel, !poslabel)
    OneVsRest(poslabel::String) = OneVsRest{String}(poslabel, "not_$(poslabel)")
    OneVsRest(poslabel::Symbol) = OneVsRest{Symbol}(poslabel, Symbol(:not_, poslabel))
    OneVsRest(poslabel::T) where {T<:Number} = OneVsRest{T}(poslabel, ifelse(poslabel == 0, poslabel+T(1), T(0)))

    """
        LabelEnc.Indices{T,K<:Number} <: LabelEncoding{T,K,1}

    A vector-based multi-label encoding with a numeric labeltype.

    Represents the class-labels as indices starting from `T(1)` up to
    (including) `T(K)`.
    In a binary setting `T(1)` corresponds to te positive class and
    `T(2)` corresponds to the negative class.
    """
    struct Indices{T<:Number,K} <: LabelEncoding{T,K,1}
        function Indices{T,K}() where {T,K}
            typeof(K) <: Int || throw(TypeError(:Indices,"constructor when checking typeof(K)",Type{Int},typeof(K)))
            new{T,K}()
        end
    end
    Indices(::Type{T}, ::Type{Val{K}}) where {T,K} = Indices{T,K}()
    Indices(::Type{Val{K}}) where {K} = Indices(Int,Val{K})
    Indices(::Type{T}, K::Number) where {T} = Indices(T,Val{Int(K)})
    Indices(K::Number) = Indices(Int,Val{Int(K)})

    """
        LabelEnc.OneOfK{T,K<:Number} <: LabelEncoding{T,K,2}

    A matrix-based multi-label encoding with a numeric labeltype.

    Represents the class labels in a one-hot encoding scheme.
    That means a matrix in which for every observation one of
    `K` elements is set to 1.
    """
    struct OneOfK{T<:Number,K} <: LabelEncoding{T,K,2}
        function OneOfK{T,K}() where {T,K}
            typeof(K) <: Int || throw(TypeError(:OneOfK,"constructor when checking typeof(K)",Type{Int},typeof(K)))
            new{T,K}()
        end
    end
    OneOfK(::Type{T}, ::Type{Val{K}}) where {T,K} = OneOfK{T,K}()
    OneOfK(::Type{Val{K}}) where {K} = OneOfK(Int,Val{K})
    OneOfK(::Type{T}, K::Number) where {T} = OneOfK(T,Val{Int(K)})
    OneOfK(K::Number) = OneOfK(Int,Val{Int(K)})

    """
        LabelEnc.NativeLabels{T,K,F} <: LabelEncoding{T,K,1}

    A vector-based multi-label encoding with a arbitrary labeltype.

    Represents arbitrary class labels by storing the possible values
    as a member variable.

    Constructor `LabelEnc.NativeLabels([fallbacklabel], classlabels)`
     - `classlabels` is a collection of labels that are allowed by the
       encoding. Each one representing a class.
     - `fallbacklabel` is an optional first argument.
         - It is either:
            - a value from `classlabels` which is used to encode any observed
              label that does not occur in `classlabels`
            - or; a function which takes as its input a label which does not
              occur in `classlabels` and returns one that does.
         - This can be used for example to handle out of vocabulary words in NLP
           as in `LabelEnc.NativeLabels("<OOV>", [red", "green", blue", ..., "<OOV>"])`
         - Or to standardize inputs: as in
           ```
            LabelEnc.NativeLabels([..., v"0.5.2", v"0.6.0", v"0.6.1", ....]) do oov
                # Discard build and prerelease fields, `v"0.7.0-beta.18"`->`v"0.7.0"`
                VersionNumber(oov.major, oov.minor, oov.patch)
            end
           ```

    In a binary setting the first element of the stored class labels
    represents the positive label and the second element the negative
    label.
    In a multiclass setting, the labels will `convertenc` to
    `LabelEnc.Indices`  using their position in the `label` argument
    to the constructor.
    """
    struct NativeLabels{T,K,F} <: LabelEncoding{T,K,1}
        getfallbacklabel::F
        label::Vector{T}
        invlabel::Dict{T,Int}
        function NativeLabels{T,K,F}(getfallbacklabel::F, label::Vector{T}) where {T,K,F<:Function}
            typeof(K) <: Int || throw(TypeError(:NativeLabels,"constructor when checking typeof(K)",Type{Int},typeof(K)))
            @assert length(label) == length(unique(label)) == K
            new{T,K,F}(getfallbacklabel, label, Dict(zip(label,1:K)))
        end
    end
    const default_getfallbacklabel = identity # by default leave label as is, so parent can change

    NativeLabels{T,K}(getfallbacklabel::F, label::Vector{T}) where {T,K,F<:Function} = NativeLabels{T,K,F}(getfallbacklabel, label)
    NativeLabels{T,K}(fallbacklabel, label::Vector{T}) where {T,K} = NativeLabels{T,K}(oov->fallbacklabel, label)
    NativeLabels{T,K}(label::Vector{T}) where {T,K} = NativeLabels{T,K}(default_getfallbacklabel, label)

    NativeLabels{T,K}(getfallbacklabel::F, label::AbstractVector{T}) where {F,T,K} = NativeLabels{T,K}(getfallbacklabel, collect(label))
    NativeLabels{T,K}(fallbacklabel::T, label::AbstractVector{T}) where {T,K} = NativeLabels{T,K}(oov->fallbacklabel, label)
    NativeLabels{T,K}(label::AbstractVector{T}) where {T,K} = NativeLabels{T,K}(default_getfallbacklabel, label)

    NativeLabels(getfallbacklabel::Function, label::AbstractVector{T}, ::Type{Val{K}})  where {T,K} = NativeLabels{T,K}(getfallbacklabel, label)
    NativeLabels(fallbacklabel, label::AbstractVector{T}, ::Type{Val{K}})  where {T,K} = NativeLabels(oov->fallbacklabel, label, Val{K})
    NativeLabels(label::AbstractVector, ::Type{Val{K}})  where {K} = NativeLabels(default_getfallbacklabel, label, Val{K})

    NativeLabels(getfallbacklabel, label) = NativeLabels(getfallbacklabel, label, Val{length(label)})
    NativeLabels(label) = NativeLabels(default_getfallbacklabel, label)

    Base.hash(a::NativeLabels, h::UInt) = hash(a.getfallbacklabel, hash(a.label, hash(:NativeLabels, h)))
    Base.:(==)(a::NativeLabels, b::NativeLabels) = isequal(a.label, b.label) && isequal(a.getfallbacklabel, b.getfallbacklabel)

end # submodule

standardize_label(lbl, lm::LabelEnc.NativeLabels) = haskey(lm.invlabel, lbl) ? lbl : lm.getfallbacklabel(lbl)

_ambiguous() = throw(ArgumentError("Can't infer the label meaning because argument types or values are ambiguous. Please specify the desired LabelEncoding manually."))

# Query the index
label2ind(lbl, lm::BinaryLabelEncoding) = ifelse(isposlabel(lbl, lm), 1, 2)
label2ind(lbl::Union{Number,T}, lm::LabelEnc.Indices{T}) where {T} = Int(lbl)
label2ind(lbl::Union{Number,T}, lm::LabelEnc.OneOfK{T}) where {T} = Int(lbl)
label2ind(lbl::AbstractVector, lm::LabelEnc.OneOfK) = argmax(lbl)
function label2ind(lbl, lm::LabelEnc.NativeLabels)
    std_lbl = standardize_label(lbl, lm)
    Int(lm.invlabel[std_lbl])
end

# Query the label
ind2label(i::Integer, lm::BinaryLabelEncoding) = ifelse(i == 1, ind2label(Val{1},lm), ind2label(Val{2},lm))
ind2label(i::Integer, ::LabelEnc.Indices{T}) where {T} = T(i)
ind2label(i::Integer, ::LabelEnc.OneOfK{T,K}) where {T,K} = (x = zeros(T,K); x[i] = T(1); x)
ind2label(i::Integer, lm::LabelEnc.NativeLabels) = lm.label[Int(i)]

ind2label(::Type{Val{1}}, ::LabelEnc.TrueFalse) = true
ind2label(::Type{Val{2}}, ::LabelEnc.TrueFalse) = false
ind2label(::Type{Val{1}}, ::LabelEnc.ZeroOne{T}) where {T} = T(1)
ind2label(::Type{Val{2}}, ::LabelEnc.ZeroOne{T}) where {T} = T(0)
ind2label(::Type{Val{1}}, ::LabelEnc.MarginBased{T}) where {T} = T(1)
ind2label(::Type{Val{2}}, ::LabelEnc.MarginBased{T}) where {T} = -T(1)
ind2label(::Type{Val{1}}, ovr::LabelEnc.OneVsRest) = ovr.poslabel
ind2label(::Type{Val{2}}, ovr::LabelEnc.OneVsRest) = ovr.neglabel
ind2label(::Type{Val{K}}, ::LabelEnc.Indices{T}) where {T,K} = T(K)
ind2label(::Type{Val{I}}, ::LabelEnc.OneOfK{T,K}) where {T,I,K} = (x = zeros(T,K); x[I] = T(1); x)
ind2label(::Type{Val{K}}, lm::LabelEnc.NativeLabels) where {K} = lm.label[K]

poslabel(lm::LabelEnc.BinaryLabelEncoding) = ind2label(Val{1}, lm)
neglabel(lm::LabelEnc.BinaryLabelEncoding) = ind2label(Val{2}, lm)

label(lm::LabelEnc.BinaryLabelEncoding{T,1}) where {T} = [poslabel(lm), neglabel(lm)]
label(::LabelEnc.Indices{T,K}) where {T,K} = collect(T(1):T(K))
label(::LabelEnc.OneOfK{T,K}) where {T,K} = collect(1:K)
label(lm::LabelEnc.NativeLabels) = lm.label

labeltype(::Type{LabelEnc.ZeroOne}) = Number
labeltype(::Type{LabelEnc.MarginBased}) = Number
labeltype(::Type{LabelEnc.Indices}) = Number
labeltype(::Type{LabelEnc.OneOfK}) = Number

# What it means to be a positive label
isposlabel(value, ::LabelEnc.FuzzyBinary) = _ambiguous()
isposlabel(value, ovr::LabelEnc.OneVsRest) = (value == ovr.poslabel)
isposlabel(value::Bool, ::LabelEnc.FuzzyBinary) = value
isposlabel(value::Bool, ::LabelEnc.TrueFalse)   = value
isposlabel(value::Bool, ::LabelEnc.MarginBased) = throw(MethodError(isposlabel,(value,)))
isposlabel(value::Bool, ::LabelEnc.Indices)     = throw(MethodError(isposlabel,(value,)))
isposlabel(value::Bool, ::LabelEnc.OneOfK)      = throw(MethodError(isposlabel,(value,)))
isposlabel(value::T, ::LabelEnc.FuzzyBinary) where {T<:Number} = (value > T(0))
isposlabel(value::T, zo::LabelEnc.ZeroOne)   where {T<:Number} = (value >= zo.cutoff)
isposlabel(value::T, ::LabelEnc.MarginBased) where {T<:Number} = (sign(value) >= T(0))
isposlabel(value::Number, lm::LabelEnc.Indices{T,2}) where {T} = value == poslabel(lm)
isposlabel(value::Number, lm::LabelEnc.OneOfK{T,2})  where {T} = value == T(1)
isposlabel(value::AbstractVector{<:Number}, lm::LabelEnc.OneOfK{T,2}) where {T} = argmax(value) == 1
isposlabel(value, lm::LabelEnc.NativeLabels{T,2}) where {T} = standardize_label(value, lm) == poslabel(lm)

# What it means to be a negative label
isneglabel(value, ::LabelEnc.FuzzyBinary) = _ambiguous()
isneglabel(value, ovr::LabelEnc.OneVsRest) = (value != ovr.poslabel)
isneglabel(value::Bool, ::LabelEnc.FuzzyBinary) = !value
isneglabel(value::Bool, ::LabelEnc.TrueFalse)   = !value
isneglabel(value::Bool, ::LabelEnc.MarginBased) = throw(MethodError(isneglabel,(value,)))
isneglabel(value::Bool, ::LabelEnc.Indices)     = throw(MethodError(isneglabel,(value,)))
isneglabel(value::Bool, ::LabelEnc.OneOfK)      = throw(MethodError(isneglabel,(value,)))
isneglabel(value::T, ::LabelEnc.FuzzyBinary) where {T<:Number} = (value <= T(0))
isneglabel(value::T, zo::LabelEnc.ZeroOne)   where {T<:Number} = (value < zo.cutoff)
isneglabel(value::T, ::LabelEnc.MarginBased) where {T<:Number} = (sign(value) == -1)
isneglabel(value::Number, lm::LabelEnc.Indices{T,2}) where {T} = value == neglabel(lm)
isneglabel(value::Number, lm::LabelEnc.OneOfK{T,2})  where {T} = value == T(2)
isneglabel(value::AbstractVector{<:Number}, lm::LabelEnc.OneOfK{T,2}) where {T} = argmax(value) == 2
isneglabel(value, lm::LabelEnc.NativeLabels{T,2}) where {T} = standardize_label(value, lm) == neglabel(lm)

# Check if the encoding is appropriate
islabelenc(targets::AbstractArray, args...) = false
islabelenc(targets::AbstractVector{T}, ::LabelEnc.FuzzyBinary)       where {T<:Union{Number,Bool}} = true
islabelenc(targets::AbstractVector{T}, ::Type{LabelEnc.FuzzyBinary}) where {T<:Union{Number,Bool}} = true
islabelenc(targets::AbstractVector{T}, ::LabelEnc.ZeroOne{T})        where {T<:Union{Number,Bool}} = all(x == 0 || x == 1 for x in targets)
islabelenc(targets::AbstractVector{T}, ::Type{LabelEnc.ZeroOne{T}})  where {T<:Union{Number,Bool}} = all(x == 0 || x == 1 for x in targets)
islabelenc(targets::AbstractVector{T}, ::Type{LabelEnc.ZeroOne})     where {T<:Union{Number,Bool}} = all(x == 0 || x == 1 for x in targets)
islabelenc(targets::AbstractVector{T}, ::LabelEnc.MarginBased{T})       where {T<:Number} = all(x == -1 || x == 1 for x in targets)
islabelenc(targets::AbstractVector{T}, ::Type{LabelEnc.MarginBased{T}}) where {T<:Number} = all(x == -1 || x == 1 for x in targets)
islabelenc(targets::AbstractVector{T}, ::Type{LabelEnc.MarginBased})    where {T<:Number} = all(x == -1 || x == 1 for x in targets)
islabelenc(targets::AbstractVector{Bool}, ::LabelEnc.TrueFalse) = true
islabelenc(targets::AbstractVector{Bool}, ::Type{LabelEnc.TrueFalse}) = true
islabelenc(targets::AbstractVector{T}, lm::LabelEnc.OneVsRest{T})   where {T}           = any(x == lm.poslabel for x in targets)
islabelenc(targets::AbstractVector{T}, ::LabelEnc.Indices{T,K})     where {T<:Number,K} = all(0 < x <= T(K) && isinteger(x) for x in targets)
islabelenc(targets::AbstractVector{T}, ::Type{LabelEnc.Indices{T}}) where {T<:Number}   = all(0 < x && isinteger(x) for x in targets)
islabelenc(targets::AbstractVector{T}, ::Type{LabelEnc.Indices})    where {T<:Number}   = all(0 < x && isinteger(x) for x in targets)

function islabelenc(targets::AbstractVector, lm::LabelEnc.NativeLabels; strict=true)
    if strict
        all(haskey(lm.invlabel, x) for x in targets)
    else
        all(haskey(lm.invlabel, standardize_label(x, lm)) for x in targets)
    end
end

function islabelenc(targets::AbstractMatrix{<:Union{Bool,Number}}, lm; obsdim = LearnBase.default_obsdim(targets))
    islabelenc(targets, lm, convert(LearnBase.ObsDimension,obsdim))
end

function islabelenc(targets::AbstractMatrix{<:Union{Bool,Number}}, lm::LabelEnc.OneOfK, obsdim)
    islabelenc(targets, typeof(lm), obsdim)
end

function islabelenc(targets::AbstractMatrix{T}, ::Type{LabelEnc.OneOfK{T}}, obsdim) where {T<:Union{Bool,Number}}
    islabelenc(targets, LabelEnc.OneOfK, obsdim)
end

function islabelenc(targets::AbstractMatrix{T}, ::Type{LabelEnc.OneOfK{T,K}}, ::Union{ObsDim.Last,ObsDim.Constant{2}}) where {T<:Union{Bool,Number},K}
    k, n = size(targets)
    ifelse(k != K, false, islabelenc(targets, LabelEnc.OneOfK, ObsDim.Last()))
end

function islabelenc(targets::AbstractMatrix{T}, ::Type{LabelEnc.OneOfK{T,K}}, ::ObsDim.First) where {T<:Union{Bool,Number},K}
    n, k = size(targets)
    ifelse(k != K, false, islabelenc(targets, LabelEnc.OneOfK, ObsDim.First()))
end

function islabelenc(targets::AbstractMatrix{T}, ::Type{LabelEnc.OneOfK}, ::Union{ObsDim.Last,ObsDim.Constant{2}}) where {T<:Union{Bool,Number}}
    k, n = size(targets)
    @inbounds for i in 1:n
        found = false
        for j in 1:k
            tcur = targets[j,i]
            if tcur == T(1)
                if found
                    return false
                end
                found = true
            elseif tcur == T(0)
                # this is fine
            else
                return false
            end
        end
        if !found
            return false
        end
    end
    return true
end

function islabelenc(targets::AbstractMatrix{T}, ::Type{LabelEnc.OneOfK}, ::ObsDim.First) where {T<:Union{Bool,Number}}
    n, k = size(targets)
    @inbounds for i in 1:n
        found = false
        for j in 1:k
            tcur = targets[i,j]
            if tcur == T(1)
                if found
                    return false
                end
                found = true
            elseif tcur == T(0)
                # this is fine
            else
                return false
            end
        end
        if !found
            return false
        end
    end
    return true
end

# Automatic determination of label mode
labelenc(target) = _ambiguous()
labelenc(targets::AbstractVector{Bool}) = LabelEnc.TrueFalse()

function labelenc(targets::AbstractVector)
    lbls = label(targets)
    LabelEnc.NativeLabels(lbls)
end

function labelenc(targets::AbstractVector{T}) where {T<:Number}
    lbls = label(targets)
    if length(lbls) == 1
        if lbls[1] == 0
            LabelEnc.ZeroOne(T)
        elseif lbls[1] == -1
            LabelEnc.MarginBased(T)
        else
            _ambiguous() # could be ZeroOne or MarginBased
        end
    elseif length(lbls) == 2
        mi, ma = extrema(lbls)
        if mi == 0 && ma == 1
            LabelEnc.ZeroOne(T)
        elseif mi == -1 && ma == 1
            LabelEnc.MarginBased(T)
        elseif mi == 1 && ma == 2
            LabelEnc.Indices(T, Val{2})
        else
            LabelEnc.NativeLabels([ma, mi])
        end
    elseif all(x > 0 && isinteger(x) for x in lbls)
        LabelEnc.Indices(T, maximum(lbls))
    elseif length(lbls) > 10 && any(!isinteger(x) for x in lbls)
        @warn("The number of distinct floating point numbers (including at least one that is non-integer!) in the label vector is quite large. Are you sure you want to perform classification and not regression?")
        LabelEnc.NativeLabels(lbls)
    else
        LabelEnc.NativeLabels(lbls)
    end
end

# TODO: Multilabel (Matrix as targets)
function labelenc(targets::AbstractMatrix{<:Number}; obsdim = LearnBase.default_obsdim(targets))
    labelenc(targets, convert(LearnBase.ObsDimension,obsdim))
end

function labelenc(targets::AbstractMatrix{T}, ::Union{ObsDim.Last,ObsDim.Constant{2}}) where {T<:Number}
    LabelEnc.OneOfK(T,size(targets,1))
end

function labelenc(targets::AbstractMatrix{T}, ::ObsDim.First) where {T<:Number}
    LabelEnc.OneOfK(T,size(targets,2))
end


abstract type AbstractCase end
export AbstractCase, EnumCase, SwitchCase,
       Leaf, Shaped

Continuation = Symbol
Branch = Pair{PatternInfo, Tuple{LineNumberNode, Continuation}}

"""for generating patterns with one-by-one checks
"""
struct EnumCase <: AbstractCase
    cases :: Vector{AbstractCase}
end

"""for generating patterns with a optimized switch
"""
struct SwitchCase <: AbstractCase
    cases :: Dict{TypeObject, AbstractCase}
end

"""specifying the code body of a case
"""
struct Leaf <: AbstractCase
    cont :: Continuation
end

"""checking the shape of the datum with the defined patterns
"""
struct Shaped <: AbstractCase
    pattern :: PatternInfo
    ln :: LineNumberNode
    case :: AbstractCase
end

@nospecialize

function case_split!(result::Vector{Branch}, branches :: Vector{Branch})
    for (p, (ln, branch)) in branches
        if p.pattern isa Or
            case_split!(result, Branch[info => (ln, branch) for info in p.pattern.ps])
        else
            push!(result, p => (ln, branch))
        end
    end
end

function build_dyn(top::TypeObject, branches::Vector{Branch})::AbstractCase
    length(branches) === 1 && return begin
        br = branches[1]
        Shaped(br.first, br.second[1], Leaf(br.second[2]))
    end
    @assert !isempty(branches)
    enum_cases = AbstractCase[]
    labels = Vector{TypeObject}(undef, length(branches))
    groups = [1]
    for i in eachindex(branches)
        current_type = branches[i].first.typetag
        @assert current_type <: top
        if current_type === top
            # Detected forcely specified matching order.
            # e.g:
            #    match val with
            #    | (_: Int) -> ...
            #    | _ -> ...
            #    | (_ : String) -> ..
            # Although Int and String are orthogonal,
            # the second pattern `_` breaks their merging
           push!(groups, i)
           push!(groups, i + 1)
           labels[i] = top
           continue
        end

        last_group_start = groups[end]
        non_orthogonal_types = TypeObject[]
        non_orthogonal_indices = Int[]
        for j in last_group_start : i - 1
            if typeintersect(current_type, labels[j]) !== Base.Bottom
            # Merging branches
            # e.g:
            #    match val with
            #   | (_ : Int) ->
            #   | (_ : String) ->
            #   | (_ : Float64) ->
            # We gives the 3 patterns such a label array:
            #   [Real, String, Real]
            # where, the 1st and 3rd branches merged
                push!(non_orthogonal_indices, j)
                push!(non_orthogonal_types, labels[j])
            end
        end
        merged_type = reduce(typejoin, non_orthogonal_types, init=current_type)
        for j in non_orthogonal_indices
            labels[j] = merged_type
        end
        labels[i] = merged_type
    end

    push!(groups, length(branches) + 1)
    n_groups = length(groups)
    for i_group in 1:n_groups-1
        start = groups[i_group]
        final = groups[i_group+1] - 1

        if start === final
            br = branches[start]
            push!(enum_cases, Shaped(br.first, br.second[1], Leaf(br.second[2])))
            continue
        elseif start > final
            continue
        end

        switch_map = Dict{TypeObject, Vector{Branch}}()
        for i in start:final
            vec = get!(switch_map, labels[i]) do
                Branch[]
            end
            push!(vec, branches[i])
        end
        switch = Pair{TypeObject, AbstractCase}[
            top => build_dyn(top, branches) for (top, branches) in switch_map
        ]
        push!(
            enum_cases,
            SwitchCase(Dict(switch))
        )
    end

    if length(enum_cases) === 1
        enum_cases[1]
    else
        EnumCase(enum_cases)
    end
end

function case_merge(branches::Vector{Branch})
    top = reduce(typejoin, TypeObject[pattern.typetag for (pattern, _) in branches])
    build_dyn(top, branches)
end

@specialize


# Testing a pretrained normalizing flow for seismic denosing
# Authors: Ali Siahkoohi, alisk@gatech.edu
#          Gabrio Rizzuti, rizzuti.gabrio@gatech.edu
#          Philipp Witte, pwitte3@gatech.edu
# Date: September 2020
"""
Example:
julia test_hint_denoising.jl --batchsize 256 --max_epoch 5
"""

using DrWatson
@quickactivate :FastApproximateInference

using InvertibleNetworks
using JLD
using JLD2
using Random
using Statistics
using ArgParse
using BSON
using Printf
using PyPlot
using Seaborn
using Flux: gpu, cpu
using Flux.Data: DataLoader

set_style("whitegrid")
rc("font", family="serif", size=13)
font_prop = matplotlib.font_manager.FontProperties(
    family="serif",
    style="normal",
    size=13
)
sfmt=matplotlib.ticker.ScalarFormatter(useMathText=true)
sfmt.set_powerlimits((0, 0))

# Random seed
# Random.seed!(19)

s = ArgParseSettings()
@add_arg_table! s begin
    "--max_epoch"
        help = "Maximum # of epochs"
        arg_type = Int
        default = 50
    "--epoch"
        help = "Epoch to load net params"
        arg_type = Int
        default = -1
    "--lr"
        help = "starting learning rate"
        arg_type = Float32
        default = 1f-3
    "--lr_step"
        help = "lr scheduler applied at evey lr_step epoch"
        arg_type = Int
        default = 1
    "--batchsize"
        help = "batch size"
        arg_type = Int
        default = 4
    "--n_hidden"
        help = "# hidden channels"
        arg_type = Int
        default = 64
    "--depth"
        help = "depth of the network"
        arg_type = Int
        default = 12
    "--sigma"
        help = "noise std"
        arg_type = Float32
        default = 2f-2
    "--sim_name"
        help = "simulation name"
        arg_type = String
        default = "seismic-denoising-256x256-0.3CS"
end
parsed_args = parse_args(s)

max_epoch = parsed_args["max_epoch"]
lr = parsed_args["lr"]
lr_step = parsed_args["lr_step"]
batchsize = parsed_args["batchsize"]
n_hidden = parsed_args["n_hidden"]
depth = parsed_args["depth"]
sigma = parsed_args["sigma"]
sim_name = parsed_args["sim_name"]

if parsed_args["epoch"] == -1
    parsed_args["epoch"] = parsed_args["max_epoch"]
end
epoch = parsed_args["epoch"]

# Define raw data directory
mkpath(datadir("training-data"))
data_path = datadir("training-data", "seismic_samples_256_by_256_num_10k.jld")
label_path = datadir("training-data", "seismic_samples_256_by_256_num_10k_labels.jld")

# Download the dataset into the data directory if it does not exist
if isfile(data_path) == false
    run(`wget https://www.dropbox.com/s/vapl62yhh8fxgwy/'
        'seismic_samples_256_by_256_num_10k.jld -q -O $data_path`)
end
if isfile(label_path) == false
    run(`wget https://www.dropbox.com/s/blxkh6tdszudlcq/'
        'seismic_samples_256_by_256_num_10k_labels.jld -q -O $label_path`)
end

# Load seismic images and create training and testing data
JLD.jldopen(data_path, "r") do file
    X_orig = read(file, "X")

    labels = JLD.jldopen(label_path, "r")["labels"][:]
    idx = findall(x -> x == 1, labels)

    X_orig = X_orig[:, :, :, idx]
    global nx, ny, nc, nsamples = size(X_orig)

    # Whiten data
    X_orig = whiten_input(X_orig)

    # Load observed data
    Y_orig = wload(
        datadir("testing-data", "noisy_seismic_samples_256_by_256_num_10k.jld2")
    )["Y_orig"]

    # Load split in training/testing
    train_idx = wload(
        datadir("testing-data", "noisy_seismic_samples_256_by_256_num_10k.jld2")
    )["train_idx"]

    # Split in training - testing
    global test_size = 5
    global test_idx = shuffle(setdiff(1:nsamples, train_idx))[1:5]

    # Dimensions after wavelet squeeze to increase no. of channels
    global nx = Int(nx/2)
    global ny = Int(ny/2)
    global n_in = Int(nc*4)

    # Apply wavelet squeeze (change dimensions to -> n1/2 x n2/2 x nc*4 x ntrain)
    global X_test = zeros(Float32, nx, ny, n_in, test_size)
    global Y_test = zeros(Float32, nx, ny, n_in, test_size)
    for (k, j) in enumerate(test_idx)
        X_test[:, :, :, k:k] = wavelet_squeeze(X_orig[:, :, :, j:j])
        Y_test[:, :, :, k:k] = wavelet_squeeze(Y_orig[:, :, :, j:j])
    end
end

# Create network
CH = NetworkConditionalHINT(
    n_in, n_hidden, depth, k1=3, k2=3, p1=1, p2=1
)

# Loading the experiment—only network weights and training loss
Params, fval = load_experiment(parsed_args)
put_params!(CH, convert(Array{Any,1}, Params))

Zx = randn(Float32, nx, ny, n_in, test_size)
Zy = randn(Float32, nx, ny, n_in, test_size)

# Precited model and data
X_, Y_ = CH.inverse(Zx, Zy)

# Now select single fixed sample from all Ys
X_fixed = X_test[:, :, :, 1:1]
Y_fixed = Y_test[:, :, :, 1:1]
Zy_fixed = CH.forward_Y(Y_fixed)

# Draw new Zx, while keeping Zy fixed
n_samples = 1000
X_post = zeros(Float32, nx, ny, n_in, n_samples)
CH = CH |> gpu

test_loader = DataLoader(
    (randn(Float32, nx, ny, n_in, n_samples), repeat(Zy_fixed, 1, 1, 1, n_samples)),
    batchsize=8, shuffle=false
)

for (itr, (X, Y)) in enumerate(test_loader)
    counter = (itr - 1)*size(X)[4] + 1

    X = X |> gpu
    Y = Y |> gpu

    X_post[:, :, :, counter:(counter+size(X)[4]-1)] = (CH.inverse(X, Y)[1] |> cpu)
    @printf("Sampling: [%d/%d]\r", counter, n_samples)
end

X_post = wavelet_unsqueeze(X_post)

# Some stats
X_post_mean = mean(X_post; dims=4)
X_post_std = std(X_post; dims=4)

# Unsqueeze all tensors
X_test = wavelet_unsqueeze(X_test)
Y_test = wavelet_unsqueeze(Y_test)

X_ = wavelet_unsqueeze(X_)
Y_ = wavelet_unsqueeze(Y_)
Zx = wavelet_unsqueeze(Zx)
Zy = wavelet_unsqueeze(Zy)

X_fixed = wavelet_unsqueeze(X_fixed)
Y_fixed = wavelet_unsqueeze(Y_fixed)
Zy_fixed = wavelet_unsqueeze(Zy_fixed)

save_dict = @strdict max_epoch epoch lr lr_step batchsize n_hidden depth sim_name
save_path = plotsdir(sim_name, savename(save_dict; digits=6))

# Training loss
fig = figure("training logs", figsize=(7, 2.5))
plot(fval, color="#d48955")
title("Negative log-likelihood")
ylabel("Training objective")
xlabel("Iterations")
safesave(joinpath(save_path, "log.png"), fig)
close(fig)

# Plot the true model
fig = figure("x", figsize=(5, 5))
imshow(
    X_fixed[:, :, 1, 1], vmin=-4.0, vmax=4.0, aspect=1, cmap="Greys", resample=true,
    interpolation="lanczos", filterrad=1
)
title(L"True model: $\mathbf{x} \sim \widehat{\pi}_x (\mathbf{x})$")
colorbar(fraction=0.047, pad=0.01, format=sfmt)
grid(false)
ax = plt.gca()
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
safesave(joinpath(save_path, "true_model" * string(test_idx[1]) * ".png"), fig)
close(fig)

# Plot the observed data
fig = figure("y", figsize=(5, 5))
imshow(
    Y_fixed[:, :, 1, 1], vmin=-4.0, vmax=4.0, aspect=1, cmap="Greys", resample=true,
    interpolation="lanczos", filterrad=1
)
title(L"Observed data: $\mathbf{y} \sim \widehat{\pi}_y (\mathbf{y})$")
colorbar(fraction=0.047, pad=0.01, format=sfmt)
grid(false)
ax = plt.gca()
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
safesave(joinpath(save_path, "observed_data" * string(test_idx[1]) * ".png"), fig)
close(fig)

# Plot the conditional mean estimate
fig = figure("x_cm", figsize=(5, 5))
imshow(
    X_post_mean[:, :, 1, 1], vmin=-4.0, vmax=4.0, aspect=1, cmap="Greys", resample=true,
    interpolation="lanczos", filterrad=1
)
title("Conditional mean")
colorbar(fraction=0.047, pad=0.01, format=sfmt)
grid(false)
ax = plt.gca()
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
safesave(joinpath(save_path, "conditional_mean" * string(test_idx[1]) * ".png"), fig)
close(fig)

# Plot the pointwise standard deviation
fig = figure("x_std", figsize=(5, 5))
imshow(
    X_post_std[:, :, 1, 1], vmin=.30,
    vmax=.42, aspect=1, cmap="OrRd", resample=true,
    interpolation="lanczos", filterrad=1
)
title("Pointwise STD")
colorbar(fraction=0.047, pad=0.01, format=sfmt)
grid(false)
ax = plt.gca()
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
safesave(joinpath(save_path, "pointwise_std" * string(test_idx[1]) * ".png"), fig)
close(fig)

fig = figure("err", figsize=(5, 5))
imshow(
    2f0.*(X_fixed[:, :, 1, 1] - X_post_mean[:, :, 1, 1])./(abs.(X_fixed[:, :, 1, 1])  + abs.(X_post_mean[:, :, 1, 1])), vmin=-2, vmax=2, aspect=1, cmap="twilight_shifted", resample=true,
    interpolation="lanczos", filterrad=1
)
title("Estimation error — accelerated")
colorbar(fraction=0.047, pad=0.01, format=sfmt)
grid(false)
ax = plt.gca()
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
safesave(joinpath(save_path, "conditional_mean_err.png"), fig)
close(fig)


using jcontrol.visuals
using jcontrol
import jcontrol.comparisons: track_segments

using Plots
import MyterialColors: teal, indigo_dark

"""
Visualize all the track's segments used for the multiple
horizons analysis
"""

track = FULLTRACK

p1 = draw(:arena)
draw!(track; alpha=.1)
draw!(track_segments[1])

p2 = draw(:arena)
draw!(track; alpha=.1)
draw!(track_segments[2])

p3 = draw(:arena)
draw!(track; alpha=.1)
draw!(track_segments[3])

p4 = draw(:arena)
draw!(track; alpha=.1)
draw!(track_segments[4])

p5 = draw(:arena)
draw!(track; alpha=.1)
draw!(track_segments[5])

p6 = draw(:arena)
draw!(track; alpha=.1)
draw!(track_segments[6])

display(
    plot(
        p1, p2, p3, p4, p5, p6,
        layout=grid(2, 3), size=(1200, 1200))
)


export PC, petscview

# preconditioner context
mutable struct PC{T}
  p::C.PC{T}
  function PC(p::C.PC{T}) where {T}
    o = new{T}(p)
    if (!deactivate_finalizers)
      finalizer(PetscDestroy,o)
    end
    return o
  end
end

comm(a::PC{T}) where {T} = C.PetscObjectComm(T, a.p.pobj)

function PetscDestroy(o::PC{T}) where {T}
  PetscFinalized(T) || C.PCDestroy(Ref(o.p))
end

function petscview(o::PC{T}) where {T}
  viewer = C.PetscViewer{T}(C_NULL)
  chk(C.PCView(o.p, viewer))
end

"""
    PC(A::Mat, PA=A, kws...)

Create a preconditioner (PC) context, given the matrix `A` of the
linear system to be solved, and optionally a different
matrix `PA` from which the preconditioner is constructed.

The remaining keywords specify zero or more additional options:
* `pc_type="a"`: use preconditioning algorithm `a`
* `pc_use_amat=true`: use Amat (instead of Pmat) to define preconditioner in nested inner solves
* ... additional options that depend on the preconditioner type ...
"""
function PC(::Type{T}; comm::MPI.Comm=MPI.COMM_SELF, kws...) where {T}
  pc_c = Ref{C.PC{T}}()
  chk(C.PCCreate(comm, pc_c))
  withoptions(T, kws) do
    chk(C.PCSetFromOptions(pc_c[]))
  end
  return PC(pc_c[])
end

function PC(A::Mat{T}, PA::Mat{T}=A; kws...) where {T}
  pc_c = Ref{C.PC{T}}()
  chk(C.PCCreate(comm(A), pc_c))
  pc = pc_c[]
  chk(C.PCSetOperators(pc, A.p, PA.p))
  withoptions(T, kws) do
    chk(C.PCSetFromOptions(pc))
  end
  return PC(pc)
end


if !isdefined(Makie, :FigurePosition)
    const FigurePosition = Makie.GridPosition
    get_parent(x::Makie.GridPosition) = x.layout.parent
else
    using Makie: FigurePosition
    get_parent(x::FigurePosition) = x.fig
end

# We can't rely on inference to always give us fully typed
# Vector{<: Number} so we add `{T} where T` to the the mix
# This makes the number like type a bit absurd, but is still nice for
# documentation purposes!
const NumberLike = Union{Number,Missing,Nothing,T} where {T}
const NumberVector = AbstractVector{<:NumberLike}
const NumberMatrix = AbstractMatrix{<:NumberLike}

#####
##### Helpers for theming and color generation...May want to move them to Colors.jl / Makie.jl
#####
using Makie.Colors: LCHab, distinguishable_colors, RGB, Colorant

function get_theme(scene, key1::Symbol, key2::Symbol; defaults...)
    return get_theme(Makie.get_scene(scene), key1, key2; defaults...)
end

function get_theme(fig::FigurePosition, key1::Symbol, key2::Symbol; defaults...)
    return get_theme(get_parent(fig), key1, key2; defaults...)
end

# This function helps us to get the theme from a scene, that we can apply to our plotting functions
function get_theme(scene::Scene, key1::Symbol, key2::Symbol; defaults...)
    scene_theme = theme(scene)
    sub_theme = get(scene_theme, key1, Theme())
    # The priority is key1.key2 > defaults > key2
    # this way defaults are overwritten by theme options specifically set for our theme.
    # Consider Kappas.Axis for key1/2, what we want is, that if there are defaults in Kappas.Axis
    # they should overwrite our generic lighthouse defaults. But anything not specified in Kappas.Axis/defaults,
    # should fall back to scene_theme.Axis
    return merge(get(sub_theme, key2, Theme()), Theme(; defaults...),
                 get(scene_theme, key2, Theme()))
end

function high_contrast(background_color::Colorant, target_color::Colorant;
                       # chose from whole lightness spectrum
                       lchoices=range(0; stop=100, length=15))
    target = LCHab(target_color)
    color = distinguishable_colors(1, [RGB(background_color)]; dropseed=true, lchoices=lchoices,
                                   cchoices=[target.c], hchoices=[target.h])
    return RGBAf0(color[1], Makie.Colors.alpha(target_color))
end

"""
    Tuple{<:NumberVector, <: NumberVector}

Tuple of X, Y coordinates
"""
const XYVector = Tuple{<:NumberVector,<:NumberVector}

"""
    Union{XYVector, AbstractVector{<: XYVector}}

A series of XYVectors, or a single xyvector.
"""
const SeriesCurves = Union{XYVector,AbstractVector{<:XYVector}}

function series_plot!(subfig::FigurePosition, per_class_pr_curves::SeriesCurves,
                      class_labels::Union{Nothing,AbstractVector{String}}; legend=:lt, title="No title",
                      xlabel="x label", ylabel="y label", solid_color=nothing, color=nothing,
                      linewidth=nothing, scatter=NamedTuple())

    axis_theme = get_theme(subfig, :SeriesPlot, :Axis; title=title, titlealign=:left, xlabel=xlabel,
                           ylabel=ylabel, aspect=AxisAspect(1),
                           xticks=0:0.2:1, yticks=0.2:0.2:1)

    ax = Axis(subfig; axis_theme...)
    # Not the most elegant, but this way we can let the theming to the Series theme, or
    # pass it through explicitely
    series_theme = get_theme(subfig, :SeriesPlot, :Series; scatter...)
    isnothing(solid_color) || (series_theme[:solid_color] = solid_color)
    isnothing(color) || (series_theme[:color] = color)
    isnothing(linewidth) || (series_theme[:linewidth] = linewidth)
    series_theme = merge(series_theme, Attributes(;scatter...))
    hidedecorations!(ax; label=false, ticklabels=false, grid=false)
    limits!(ax, 0, 1, 0, 1)
    Makie.series!(ax, per_class_pr_curves; labels=class_labels, series_theme...)
    if !isnothing(legend)
        axislegend(ax; position=legend)
    end
    return ax
end

function plot_pr_curves!(subfig::FigurePosition, per_class_pr_curves::SeriesCurves,
                         class_labels::Union{Nothing,AbstractVector{String}}; legend=:lt, title="PR curves",
                         xlabel="True positive rate", ylabel="Precision", scatter=NamedTuple(),
                         solid_color=nothing)
    return series_plot!(subfig, per_class_pr_curves, class_labels; legend=legend, title=title, xlabel=xlabel,
                        ylabel=ylabel, scatter=scatter, solid_color=solid_color)
end

function plot_roc_curves!(subfig::FigurePosition, per_class_roc_curves::SeriesCurves,
                          per_class_roc_aucs::NumberVector, class_labels::AbstractVector{<:String};
                          legend=:rb, title="ROC curves", xlabel="False positive rate",
                          ylabel="True positive rate")
    auc_labels = [@sprintf("%s (AUC: %.3f)", class, per_class_roc_aucs[i])
                  for (i, class) in enumerate(class_labels)]

    return series_plot!(subfig, per_class_roc_curves, auc_labels; legend=legend, title=title, xlabel=xlabel,
                        ylabel=ylabel)
end

function plot_reliability_calibration_curves!(subfig::FigurePosition,
                                              per_class_reliability_calibration_curves::SeriesCurves,
                                              per_class_reliability_calibration_scores::NumberVector,
                                              class_labels::AbstractVector{String}; legend=:rb)
    calibration_score_labels = map(enumerate(class_labels)) do (i, class)
        @sprintf("%s (MSE: %.3f)", class, per_class_reliability_calibration_scores[i])
    end

    scatter_theme = get_theme(subfig, :ReliabilityCalibrationCurves, :Scatter; marker=Circle,
                              markersize=5, strokewidth=0)
    ideal_theme = get_theme(subfig, :ReliabilityCalibrationCurves, :Ideal; color=(:black, 0.5),
                            linestyle=:dash, linewidth=2)

    ax = series_plot!(subfig, per_class_reliability_calibration_curves, calibration_score_labels;
                      legend=legend, title="Prediction reliability calibration",
                      xlabel="Predicted probability bin", ylabel="Fraction of positives",
                      scatter=scatter_theme)
    #TODO: mean predicted value histogram underneath?? Maybe important...
    # https://scikit-learn.org/stable/modules/calibration.html
    linesegments!(ax, [0, 1], [0, 1]; ideal_theme...)
    return ax
end

function set_from_kw!(theme, key, kw, default)
    if haskey(kw, key)
        theme[key] = getproperty(kw, key)
    elseif !haskey(theme, key)
        theme[key] = default
    end
    return
end

function plot_binary_discrimination_calibration_curves!(subfig::FigurePosition,
                                                        calibration_curve::SeriesCurves, calibration_score,
                                                        per_expert_calibration_curves::SeriesCurves,
                                                        per_expert_calibration_scores, optimal_threshold,
                                                        discrimination_class::AbstractString; kw...)
    kw = values(kw)
    scatter_theme = get_theme(subfig, :BinaryDiscriminationCalibrationCurves, :Scatter; strokewidth=0)
    # Hayaah, this theme merging is getting out of hand
    # but we want kw > BinaryDiscriminationCalibrationCurves > Scatter, so we need to somehow set things
    # after the theme merging above, especially, since we also pass those to series!,
    # which then again tries to merge the kw args with the theme.
    set_from_kw!(scatter_theme, :makersize, kw, 5)
    set_from_kw!(scatter_theme, :marker, kw, :rect)

    per_expert = get_theme(subfig, :BinaryDiscriminationCalibrationCurves, :PerExpert; solid_color=:darkgrey,
                           color=nothing)
    set_from_kw!(per_expert, :linewidth, kw, 2)
    ax = series_plot!(subfig, per_expert_calibration_curves, nothing; legend=nothing,
                      title="Detection calibration", xlabel="Expert agreement rate",
                      ylabel="Predicted positive probability", scatter=scatter_theme, per_expert...)

    calibration = get_theme(subfig, :BinaryDiscriminationCalibrationCurves, :CalibrationCurve;
                            solid_color=:navyblue, markerstrokewidth=0)

    set_from_kw!(calibration, :markersize, kw, 5)
    set_from_kw!(calibration, :marker, kw, :rect)
    set_from_kw!(calibration, :linewidth, kw, 2)

    Makie.series!(ax, calibration_curve; calibration...)

    ideal_theme = get_theme(subfig, :BinaryDiscriminationCalibrationCurves, :Ideal; color=(:black, 0.5),
                            linestyle=:dash)
    set_from_kw!(ideal_theme, :linewidth, kw, 2)
    linesegments!(ax, [0, 1], [0, 1]; label="Ideal", ideal_theme...)
    #TODO: expert agreement histogram underneath?? Maybe important...
    # https://scikit-learn.org/stable/modules/calibration.html
    return ax
end

function plot_confusion_matrix!(subfig::FigurePosition, confusion::NumberMatrix,
                                class_labels::AbstractVector{String}, normalize_by::Symbol;
                                annotation_text_size=20, colormap=:Blues)
    normdim = get((Row=2, Column=1), normalize_by) do
        return error("normalize_by must be either :Row or :Column, found: $(normalize_by)")
    end

    nclasses = length(class_labels)
    if size(confusion) != (nclasses, nclasses)
        error("Labels must match size of square confusion matrix. Found $(nclasses) labels for an $(size(confusion)) matrix")
    end
    confusion = round.(confusion ./ sum(confusion; dims=normdim); digits=3)
    class_indices = 1:nclasses
    text_theme = get_theme(subfig, :ConfusionMatrix, :Text; textsize=annotation_text_size)
    heatmap_theme = get_theme(subfig, :ConfusionMatrix, :Heatmap; nan_color=(:black, 0.0))
    axis_theme = get_theme(subfig, :ConfusionMatrix, :Axis; xticklabelrotation=pi / 4,
                           titlealign=:left, title="$(string(normalize_by))-Normalized Confusion",
                           xlabel="Elected Class", ylabel="Predicted Class",
                           xticks=(class_indices, class_labels),
                           yticks=(class_indices, class_labels),
                           aspect=AxisAspect(1))

    ax = Axis(subfig; axis_theme...)

    hidedecorations!(ax; label=false, ticklabels=false, grid=false)
    ylims!(ax, nclasses, 0)
    tightlimits!(ax)
    plot_bg_color = to_color(ax.backgroundcolor[])
    crange = (0.0, 1.0)
    nan_color = to_color(heatmap_theme.nan_color[])
    cmap = to_colormap(to_value(pop!(heatmap_theme, :colormap, colormap)))
    heatmap!(ax, confusion'; colorrange=crange, colormap=cmap, nan_color=nan_color, heatmap_theme...)
    text_color = to_color(to_value(pop!(text_theme, :color, :black)))
    function label_color(i, j)
        c = confusion[i, j]
        bg_color = if isnan(c) || ismissing(c)
            Makie.Colors.alpha(nan_color) <= 0.0 ? plot_bg_color : nan_color
        else
            Makie.interpolated_getindex(cmap, c, crange)
        end
        return high_contrast(bg_color, text_color)
    end
    annos = vec([(string(confusion[i, j]), Point2f0(j, i)) for i in class_indices, j in class_indices])
    colors = vec([label_color(i, j) for i in class_indices, j in class_indices])
    text!(ax, annos; align=(:center, :center), color=colors, textsize=annotation_text_size, text_theme...)
    return ax
end

function text_attributes(values, groups, bar_colors, bg_color, text_color)
    aligns = NTuple{2,Symbol}[]
    offsets = NTuple{2,Int}[]
    text_colors = RGBAf0[]
    for (i, k) in enumerate(values)
        group = groups isa AbstractVector ? groups[i] : groups
        bar_color = bar_colors[group]
        # Plot text inside bar
        if k > 0.85
            push!(aligns, (:right, :center))
            push!(offsets, (-10, 0))
            push!(text_colors, high_contrast(bar_color, text_color))
        else
            # plot text next to bar
            push!(aligns, (:left, :center))
            push!(offsets, (10, 0))
            push!(text_colors, high_contrast(bg_color, text_color))
        end
    end
    return aligns, offsets, text_colors
end

function plot_kappas!(subfig::FigurePosition, per_class_kappas::NumberVector,
                      class_labels::AbstractVector{String}, per_class_IRA_kappas=nothing;
                      color=[:lightgrey, :lightblue], annotation_text_size=20)
    nclasses = length(class_labels)
    axis_theme = get_theme(subfig, :Kappas, :Axis; aspect=AxisAspect(1), titlealign=:left,
                           xlabel="Cohen's kappa", xticks=[0, 1], yreversed=true,
                           yticks=(1:nclasses, class_labels))

    text_theme = get_theme(subfig, :Kappas, :Text; textsize=annotation_text_size)
    text_color = to_color(to_value(pop!(text_theme, :color, to_color(:black))))
    bars_theme = get_theme(subfig, :Kappas, :BarPlot; color=color)
    bar_colors = to_color.(bars_theme.color[])

    ax = Axis(subfig[1, 1]; axis_theme...)
    bg_color = to_color(ax.backgroundcolor[])
    xlims!(ax, 0, 1)
    if isnothing(per_class_IRA_kappas)
        ax.title = "Algorithm-expert agreement"
        annotations = map(enumerate(per_class_kappas)) do (i, k)
            return (string(round(k; digits=3)), Point2f0(max(0, k), i))
        end
        aligns, offsets, text_colors = text_attributes(per_class_kappas, 2, bar_colors, bg_color, text_color)
        barplot!(ax, per_class_kappas; direction=:x, color=bar_colors[2])
        text!(ax, annotations; align=aligns, offset=offsets, color=text_colors, text_theme...)
    else
        ax.title = "Inter-rater reliability"
        values = vcat(per_class_kappas, per_class_IRA_kappas)
        groups = vcat(fill(2, nclasses), fill(1, nclasses))
        xvals = vcat(1:nclasses, 1:nclasses)
        cmap = bar_colors
        bars = barplot!(ax, xvals, max.(0, values); dodge=groups, color=groups, direction=:x, colormap=cmap)
        # This is a bit hacky, but for now the easiest way to figure out the exact, dodged positions
        rectangles = bars.plots[][1][]
        dodged_y = last.(minimum.(rectangles)) .+ (last.(widths.(rectangles)) ./ 2)
        textpos = Point2f0.(max.(0, values), dodged_y)

        labels = string.(round.(values; digits=3))
        aligns, offsets, text_colors = text_attributes(values, groups, bar_colors, bg_color, text_color)
        text!(ax, labels; position=textpos, align=aligns, offset=offsets, color=text_colors, text_theme...)
        labels = ["Expert-vs-expert IRA", "Algorithm-vs-expert"]
        entries = map(c -> PolyElement(; color=c, strokewidth=0, strokecolor=:white), cmap)
        legend = Legend(subfig[1, 1, Bottom()], entries, labels; tellwidth=false, tellheight=true,
                        labelsize=12, padding=(0, 0, 0, 0), framevisible=false, patchsize=(10, 10),
                        patchlabelgap=6)
        legend.margin = (0, 0, 0, 60)
    end
    return ax
end

"""
    evaluation_metrics_plot(data::Dict; resolution=(1000, 1000), textsize=12)

Plots all evaluation metrics Lighthouse has to offer.
"""
function evaluation_metrics_plot(data::Dict; resolution=(1000, 1000), textsize=12)
    fig = Figure(; resolution=resolution, Axis=(titlesize=17,))

    # Confusion
    plot_confusion_matrix!(fig[1, 1], data["confusion_matrix"], data["class_labels"], :Column;
                           annotation_text_size=textsize)
    plot_confusion_matrix!(fig[1, 2], data["confusion_matrix"], data["class_labels"], :Row;
                           annotation_text_size=textsize)
    # Kappas
    IRA_kappa_data = nothing
    multiclass = length(data["class_labels"]) > 2
    labels = multiclass ? vcat("Multiclass", data["class_labels"]) : data["class_labels"]
    kappa_data = multiclass ? vcat(data["multiclass_kappa"], data["per_class_kappas"]) :
                 data["per_class_kappas"]

    if issubset(["multiclass_IRA_kappas", "per_class_IRA_kappas"], keys(data))
        IRA_kappa_data = multiclass ? vcat(data["multiclass_IRA_kappas"], data["per_class_IRA_kappas"]) :
                         data["per_class_IRA_kappas"]
    end

    plot_kappas!(fig[1, 3], kappa_data, labels, IRA_kappa_data; annotation_text_size=textsize)

    # Curves
    ax = plot_roc_curves!(fig[2, 1], data["per_class_roc_curves"], data["per_class_roc_aucs"],
                          data["class_labels"]; legend=nothing)

    plot_pr_curves!(fig[2, 2], data["per_class_pr_curves"], data["class_labels"]; legend=nothing)

    plot_reliability_calibration_curves!(fig[3, 1], data["per_class_reliability_calibration_curves"],
                                         data["per_class_reliability_calibration_scores"],
                                         data["class_labels"]; legend=nothing)

    legend_pos = 2:3
    if haskey(data, "discrimination_calibration_curve")
        legend_pos = 3
        plot_binary_discrimination_calibration_curves!(fig[3, 2], data["discrimination_calibration_curve"],
                                                       data["discrimination_calibration_score"],
                                                       data["per_expert_discrimination_calibration_curves"],
                                                       data["per_expert_discrimination_calibration_scores"],
                                                       data["optimal_threshold"],
                                                       data["class_labels"][data["optimal_threshold_class"]])
    end
    legend_plots = filter(Makie.MakieLayout.get_plots(ax)) do plot
        return haskey(plot, :label)
    end
    elements = map(legend_plots) do elem
        return [PolyElement(; color=elem.color, strokecolor=:transparent)]
    end

    function label_str(i)
        auc = round(data["per_class_roc_aucs"][i]; digits=2)
        mse = round(data["per_class_reliability_calibration_scores"][i]; digits=2)
        return ["""ROC AUC  $auc
                   Cal. MSE    $mse
                   """]
    end
    classes = data["class_labels"]
    nclasses = length(classes)
    class_labels = label_str.(1:nclasses)
    Legend(fig[3, legend_pos], elements, class_labels, classes; nbanks=2, tellwidth=false, tellheight=false,
           labelsize=11, titlesize=14, titlegap=5, groupgap=6, labelhalign=:left, labelvalign=:center)
    colgap!(fig.layout, 2)
    rowgap!(fig.layout, 4)
    return fig
end

# Helper to more easily define the non mutating versions
function axisplot(func, args; resolution=(800, 600), plot_kw...)
    fig = Figure(; resolution=resolution)
    ax = func(fig[1, 1], args...; plot_kw...)
    # ax.plots[1] is not really that great, but there isn't a FigureAxis object right now
    # this will need to wait for when we figure out a better recipe integration
    return Makie.FigureAxisPlot(fig, ax, ax.scene.plots[1])
end

"""
    plot_confusion_matrix!(subfig::FigurePosition, args...; kw...)

    plot_confusion_matrix(confusion::AbstractMatrix{<: Number}, class_labels::AbstractVector{String}, normalize_by::Symbol;
                          resolution=(800,600),
                          annotation_text_size=20)

Lighthouse plots confusion matrices, which are simple tables
showing the empirical distribution of predicted class (the rows)
versus the elected class (the columns). These come in two variants:

* row-normalized: this means each row has been normalized to sum to 1. Thus, the row-normalized confusion matrix shows the empirical distribution of elected classes for a given predicted class. E.g. the first row of the row-normalized confusion matrix shows the empirical probabilities of the elected classes for a sample which was predicted to be in the first class.
* column-normalized: this means each column has been normalized to sum to 1. Thus, the column-normalized confusion matrix shows the empirical distribution of predicted classes for a given elected class. E.g. the first column of the column-normalized confusion matrix shows the empirical probabilities of the predicted classes for a sample which was elected to be in the first class.

```
fig, ax, p = plot_confusion_matrix(rand(2, 2), ["1", "2"], :Row)
fig = Figure()
ax = plot_confusion_matrix!(fig[1, 1], rand(2, 2), ["1", "2"], :Column)
```
"""
plot_confusion_matrix(args...; kw...) = axisplot(plot_confusion_matrix!, args; kw...)

"""
    plot_reliability_calibration_curves!(fig::SubFigure, args...; kw...)

    plot_reliability_calibration_curves(per_class_reliability_calibration_curves::SeriesCurves,
                                        per_class_reliability_calibration_scores::NumberVector,
                                        class_labels::AbstractVector{String};
                                        legend=:rb, resolution=(800, 600))

"""
function plot_reliability_calibration_curves(args...; kw...)
    return axisplot(plot_reliability_calibration_curves!, args; kw...)
end

"""
    plot_binary_discrimination_calibration_curves!(fig::SubFigure, args...; kw...)

    plot_binary_discrimination_calibration_curves!(calibration_curve::SeriesCurves, calibration_score,
                                                   per_expert_calibration_curves::SeriesCurves,
                                                   per_expert_calibration_scores, optimal_threshold,
                                                   discrimination_class::AbstractString;
                                                   marker=:rect, markersize=5, linewidth=2)

"""
function plot_binary_discrimination_calibration_curves(args...; kw...)
    return axisplot(plot_binary_discrimination_calibration_curves!, args; kw...)
end

"""
    plot_pr_curves!(subfig::FigurePosition, args...; kw...)

    plot_pr_curves(per_class_pr_curves::SeriesCurves,
                class_labels::AbstractVector{<: String};
                resolution=(800, 600),
                legend=:lt, title="PR curves",
                xlabel="True positive rate", ylabel="Precision",
                linewidth=2, scatter=NamedTuple(), color=:darktest)

- `scatter::Union{Nothing, NamedTuple}`: can be set to a named tuples of attributes that are forwarded to the scatter call (e.g. markersize). If nothing, no scatter is added.

"""
plot_pr_curves(args...; kw...) = axisplot(plot_pr_curves!, args; kw...)

"""
    plot_roc_curves!(subfig::FigurePosition, args...; kw...)

    plot_roc_curves(per_class_roc_curves::SeriesCurves,
                    per_class_roc_aucs::NumberVector,
                    class_labels::AbstractVector{<: String};
                    resolution=(800, 600),
                    legend=:lt,
                    title="ROC curves",
                    xlabel="False positive rate",
                    ylabel="True positive rate",
                    linewidth=2, scatter=NamedTuple(), color=:darktest)

- `scatter::Union{Nothing, NamedTuple}`: can be set to a named tuples of attributes that are forwarded to the scatter call (e.g. markersize). If nothing, no scatter is added.

"""
plot_roc_curves(args...; kw...) = axisplot(plot_roc_curves!, args; kw...)

"""
    plot_kappas!(subfig::FigurePosition, args...; kw...)

    plot_kappas(per_class_kappas::NumberVector,
                class_labels::AbstractVector{String},
                per_class_IRA_kappas=nothing;
                resolution=(800, 600),
                annotation_text_size=20)
"""
plot_kappas(args...; kw...) = axisplot(plot_kappas!, args; kw...)