containers.jl

const MatVecType{T} = Union{Matrix{T}, Vector{T}, SRow{T}} where T
const ContainerTypes = Union{MatVecType, Set, Dict, Tuple, NamedTuple, Array}

# handle Vector and Matrix instances, i.e. Array instances with 1 or 2 dimensions
function type_params(obj::S) where {T, S <: MatVecType{T}}
  isempty(obj) && return TypeParams(S, TypeParams(T, nothing))
  
  params = type_params.(obj)
  @req allequal(params) "Not all params of the entries are the same, consider using a Tuple for serialization"
  return TypeParams(S, params[1])
end

# handle Array instances with >= 3 dimensions
function type_params(obj::S) where {N, T, S <: Array{T, N}}
  isempty(obj) && return TypeParams(S, :subtype_params => TypeParams(T, nothing), :dims => N)
  
  params = type_params.(obj)
  @req allequal(params) "Not all params of the entries are the same, consider using a Tuple for serialization"
  return TypeParams(S, :subtype_params => params[1], :dims => N)
end

has_empty_entries(obj) = false

has_empty_entries(obj::ContainerTypes) = isempty(obj) || any(has_empty_entries, obj)

# handle Vector and Matrix instances containing other containers
function type_params(obj::S) where {T <: ContainerTypes, S <: MatVecType{T}}
  isempty(obj) && return TypeParams(S, TypeParams(T, nothing))

  # empty entries can inherit params from the rest of the collection
  non_empty_entries = filter(!has_empty_entries, obj)
  params = type_params.(non_empty_entries)
  @req allequal(params) "Not all params of the entries are the same, consider using a Tuple for serialization"

  # need to check if any of the inner containers are non empty, and if there is at least one
  # then we can use that one to get the type for the entire nested container
  isempty(params) && return TypeParams(S, type_params(first(obj)))
  return TypeParams(S, params[1])
end

function save_type_params(s::SerializerState,
                          tp::TypeParams{T, <:Tuple{Vararg{Pair}}}) where {S, N,  T <: Array{S, N}}
  save_data_dict(s) do
    save_object(s, encode_type(T), :name)
    save_data_dict(s, :params) do
      for (k, v) in params(tp)
        if k == :dims
          save_object(s, v, k)
        else
          save_type_params(s, v, k)
        end
      end
    end
  end
end

# this function is forced to deal with all array types
function load_type_params(s::DeserializerState, T::Type{<: Union{Array, MatVecType}})
  !haskey(s, :params) && return T, nothing
  subtype, params = load_node(s, :params) do _
    if haskey(s, :dims)
      U = load_node(s, :subtype_params) do _
        return decode_type(s)
      end
       
      dims = load_object(s, Int, :dims)
      subtype, params = load_type_params(s, U, :subtype_params)
      return T{subtype, dims}, params
    else
      subtype, params = load_type_params(s, decode_type(s))
      return T{subtype}, params
    end
  end
end

################################################################################
# loads to handle

function load_object(s::DeserializerState, T::Type{Vector{S}},
                     key::Union{Symbol, Int}) where S
  load_node(s, key) do _
    load_object(s, T)
  end
end

################################################################################
# Saving and loading vectors
@register_serialization_type Vector

# see Array above for load_type_params for general Arrays which include Matrix
# and Vector

function save_object(s::SerializerState, x::AbstractVector{S}) where S
  save_data_array(s) do
    for elem in x
      if serialize_with_id(typeof(elem))
        ref = save_as_ref(s, elem)
        save_object(s, ref)
      else
        save_object(s, elem)
      end
    end
  end
end

function load_object(s::DeserializerState, T::Type{<: Vector{params}}) where params
  load_node(s) do v
    if serialize_with_id(params)
      loaded_v = load_array_node(s) do _
        load_ref(s)
      end::Vector{params}
    else
      isempty(s.obj) && return params[]
      loaded_v = load_array_node(s) do _
        load_object(s, params)
      end
    end
    return loaded_v
  end
end


function load_object(s::DeserializerState, ::Type{Vector{T}}, params::S) where {T, S}
  isempty(s.obj) && return T[]
  v = load_array_node(s) do _
    if serialize_with_id(T)
      load_ref(s)
    else
      load_object(s, T, params)
    end
  end
  return v
end

function load_object(s::DeserializerState, T::Type{Vector{U}}, ::Nothing) where U
  isempty(s.obj) && return U[]
  return load_array_node(s) do _
    load_object(s, U)
  end
end

################################################################################
# Saving and loading matrices
@register_serialization_type Matrix

function save_object(s::SerializerState, mat::AbstractMatrix{S}) where {S}
  save_data_array(s) do
    for r in eachrow(mat)
      save_object(s, r)
    end
  end
end

# this needs to be Matrix to avoid ambiguities, this is also ok
# since types on load will be Matrix
function load_object(s::DeserializerState, T::Type{<:Matrix{S}}) where {S}
  load_node(s) do entries
    if isempty(entries)
      return T(undef, 0, 0)
    end
    len = length(entries)
    m = stack([
      load_object(s, Vector{S}, i) for i in 1:len
        ]; dims=1)
    return T(m)
  end
end

load_object(s::DeserializerState, T::Type{<:Matrix{S}}, ::Nothing) where {S} = load_object(s, T)

function load_object(s::DeserializerState, T::Type{<:Matrix{S}}, params) where {S}
  load_node(s) do entries
    if isempty(entries)
      return T(undef, 0, 0)
    end
    len = length(entries)
    m = stack([
      load_object(s, Vector{S}, params, i) for i in 1:len
        ]; dims=1)
    return T(m)
  end
end

function load_object(s::DeserializerState, T::Type{<:Matrix{S}}, key::Union{Int, Symbol}) where {S}
  return load_node(s, key) do _
    load_object(s, T)
  end
end

################################################################################
# Multidimensional Arrays
@register_serialization_type Array "MultiDimArray"

function save_object(s::SerializerState, arr::AbstractArray{T, N}) where {T, N}
  save_data_array(s) do
    for slice in eachslice(arr; dims=1)
      save_object(s, slice)
    end
  end
end

function load_object(s::DeserializerState, T::Type{Array{S, N}}) where {S, N}
  load_node(s) do entries
    if isempty(entries)
      return T(undef, [0 for _ in 1:N]...)
    end
    len = length(entries)
    m = [load_object(s, Array{S, N - 1}, i) for i in 1:len]
    return stack(m; dims=1)
  end
end

load_object(s::DeserializerState, T::Type{Array{S, N}}, ::Nothing) where {S, N} = load_object(s, T)

function load_object(s::DeserializerState, T::Type{Array{S, N}}, params::U) where {S, N, U}
  load_node(s) do entries
    if isempty(entries)
      return T(undef, [0 for _ in 1:N]...)
    end
    len = length(entries)
    m = [load_object(s, Array{S, N - 1}, params, i) for i in 1:len]
    return stack(m; dims=1)
  end
end

function load_object(s::DeserializerState, T::Type{Array{S, N}}, key::Union{Int, Symbol}) where {S, N}
  return load_node(s, key) do _
    load_object(s, T)
  end
end

################################################################################
# Saving and loading Tuple
@register_serialization_type Tuple

function type_params(obj::T) where T <: Tuple
  return TypeParams(T, type_params.(obj))
end

function save_type_params(s::SerializerState, tp::TypeParams{T}) where T <: Tuple
  save_data_dict(s) do
    save_object(s, encode_type(T), :name)
    save_data_array(s, :params) do
      for (i, param_tp) in enumerate(params(tp))
        save_type_params(s, param_tp)
      end
    end
  end
end

function load_type_params(s::DeserializerState, T::Type{Tuple})
  !haskey(s, :params) && return T{}, nothing
  subtype, params = load_node(s, :params) do _
    tuple_params = load_array_node(s) do _
      load_type_params(s, decode_type(s))
    end
    return Tuple([x[1] for x in tuple_params]), Tuple(x[2] for x in tuple_params)
  end
  return T{subtype...}, params
end

function save_object(s::SerializerState, obj::Tuple)
  save_data_array(s) do
    for entry in obj
      if serialize_with_id(typeof(entry))
        ref = save_as_ref(s, entry)
        save_object(s, ref)
      else
        save_object(s, entry)
      end
    end
  end
end

function load_object(s::DeserializerState, T::Type{<:Tuple}, params::Tuple)
  entries = load_array_node(s) do (i, entry)
    S = fieldtype(T, i)
    if serialize_with_id(S)
      return load_ref(s)
    else
      return load_object(s, S, params[i])
    end
  end
  return Tuple(entries)
end

function load_object(s::DeserializerState, T::Type{<:Tuple})
  entries = load_array_node(s) do (i, entry)
    S = fieldtype(T, i)
    load_object(s, S)
  end
  return T(entries)
end

################################################################################
# Saving and loading NamedTuple
@register_serialization_type NamedTuple

function type_params(obj::T) where T <: NamedTuple
  return TypeParams(
    T, 
    NamedTuple(x.first => type_params(x.second) for x in pairs(obj))
  )
end

# Named Tuples need to preserve order so they are handled separate from Dict
function save_type_params(s::SerializerState, tp::TypeParams{<:NamedTuple})
  T = type(tp)
  save_data_dict(s) do
    save_object(s, encode_type(T), :name)
    save_data_dict(s, :params) do
      save_data_array(s, :names) do
        for name in keys(params(tp))
          save_object(s, name)
        end
      end
      save_data_array(s, :tuple_params) do
        for (i, param_tp) in enumerate(values(params(tp)))
          save_type_params(s, param_tp)
        end
      end
    end
  end
end

function load_type_params(s::DeserializerState, T::Type{NamedTuple})
  subtype, params = load_node(s, :params) do obj
    tuple_params = load_array_node(s, :tuple_params) do _
      load_type_params(s, decode_type(s))
    end
    tuple_types, named_tuple_params = collect(zip(tuple_params...))
    names = load_object(s, Vector{Symbol}, :names)
    return (tuple(names...), Tuple{tuple_types...}), named_tuple_params
  end
  return T{subtype...}, params
end

function save_object(s::SerializerState, obj::NamedTuple)
  save_object(s, values(obj))
end

function load_object(s::DeserializerState, T::Type{<: NamedTuple}, params::Tuple)
  return T(load_object(s, Tuple{Base.fieldtypes(T)...}, params))
end

################################################################################
# Saving and loading dicts
@register_serialization_type Dict

function type_params(obj::T) where {S <: Union{Symbol, Int, String},
                                    T <: Dict{S, Any}}
  return TypeParams(
    T,
    :key_params => TypeParams(S, nothing),
    (x.first => type_params(x.second) for x in pairs(obj))...
  )
end

function type_params(obj::T) where {U, S, T <: Dict{S, U}}
  if isempty(obj)
    return TypeParams(
      T,
      :key_params => TypeParams(S, nothing),
      :value_params => TypeParams(U, nothing)
    )
  end
  key_params = Oscar.type_params.(collect(keys(obj)))
  @req allequal(key_params) "Not all type parameters of keys in $obj are the same"

  value_params = type_params.(collect(values(obj)))
  @req allequal(value_params) "Not all type parameters of values in $obj are the same"

  return TypeParams(
    T, 
    :key_params => first(key_params),
    :value_params => first(value_params),
  )
end

function save_type_params(
  s::SerializerState,
  tp::TypeParams{Dict{S, T}, <:Tuple{Vararg{Pair}}}) where {T, S <: Union{Symbol, Int, String}}
  save_data_dict(s) do
    save_object(s, encode_type(Dict), :name)
    save_data_dict(s, :params) do
      isempty(params(tp)) && save_object(s, encode_type(T), :value_params)
      for (k, param_tp) in params(tp)
        save_type_params(s, param_tp, Symbol(k))
      end
    end
  end
end

function load_type_params(s::DeserializerState, T::Type{Dict})
  subtype, params = load_node(s, :params) do obj
    if haskey(s, :value_params)
      S, key_params = load_node(s, :key_params) do params
        params isa String && return decode_type(s), nothing
        load_type_params(s, decode_type(s))
      end

      U, value_params = load_node(s, :value_params) do _
        load_type_params(s, decode_type(s))
      end

      return (S, U), Dict(:key_params => key_params, :value_params => value_params)
    else
      S, key_params = load_node(s, :key_params) do _
        decode_type(s), nothing
      end
      params_dict = Dict{S, Any}()
      value_types = Type[]
      for (k, _) in obj
        k == :key_params && continue
        key = S <: Integer ? parse(Int, string(k)) : S(k)
        params_dict[key] = load_node(s, k) do _
          return load_type_params(s, decode_type(s))
        end
        push!(value_types, params_dict[key][1])
      end
      params_dict = isempty(params_dict) ? nothing : params_dict
      return (S, Union{value_types...}), params_dict
    end
  end
  return Dict{subtype...}, params
end

function save_object(s::SerializerState, obj::Dict{S, T}) where {S <: Union{Symbol, String, Int}, T}
  save_data_dict(s) do
    for (k, v) in obj
      if !Base.issingletontype(typeof(v))
        save_object(s, v, Symbol(k))
      end
    end
  end
end

function save_object(s::SerializerState, obj::Dict{S, T}) where {S, T}
  save_data_array(s) do
    for (k, v) in obj
      save_object(s, (k, v))
    end
  end
end

function load_object(s::DeserializerState,
                     T::Type{<:Dict{S, Any}},
                     params::Dict{S, Any}) where {S <: Union{Symbol, String, Int}}
  dict = T()
  for k in keys(params)
    # has no data, hence no key was generated on the data side
    if Base.issingletontype(params[k][1])
      dict[k] = params[k][1]()
    else
      dict[k] = load_object(s, params[k]..., Symbol(k))
    end
  end
  return dict
end

function load_object(s::DeserializerState,
                     T::Type{<:Dict{S, U}}) where {S <: Union{Symbol, String}, U}
  dict = T()
  for k in keys(s.obj)
    dict[S(k)] = load_object(s, U, Symbol(k))
  end
  return dict
end

function load_object(s::DeserializerState,
                     T::Type{<:Dict{S, U}}, ::Nothing) where {S <: Union{Symbol, String}, U}
  dict = T()
  for k in keys(s.obj)
    dict[S(k)] = load_object(s, U, Symbol(k))
  end
  return dict
end

function load_object(s::DeserializerState,
                     T::Type{<:Dict{Int, Int}})
  dict = T()
  for k in keys(s.obj)
    dict[parse(Int, string(k))] = load_object(s, Int, k)
  end
  return dict
end

function load_object(s::DeserializerState,
                     T::Type{<:Dict{Int, S}}, params::Nothing) where S
  dict = T()
  for k in keys(s.obj)
    dict[parse(Int, string(k))] = load_object(s, S, k)
  end
  return dict
end

# here to handle ambiguities
function load_object(s::DeserializerState, T::Type{Dict{Int, Int}}, key::Union{Symbol, Int})
  load_node(s, key) do _
    load_object(s, T)
  end
end

# here to handle ambiguities
function load_object(s::DeserializerState, T::Type{Dict{String, Int}}, key::Union{Symbol, Int})
  load_node(s, key) do _
    load_object(s, T)
  end
end

# here to handle ambiguities
function load_object(s::DeserializerState, T::Type{Dict{Symbol, Int}}, key::Union{Symbol, Int})
  load_node(s, key) do _
    load_object(s, T)
  end
end

function load_object(s::DeserializerState, T::Type{<:Dict{S, U}}, ::Nothing) where {S, U}
  pairs = load_array_node(s) do _
    load_object(s, Tuple{S, U}, (nothing, nothing))
  end::Vector{Tuple{S, U}}
  return T(k => v for (k, v) in pairs)
end

function load_object(s::DeserializerState,
                     T::Type{<:Dict{S, U}},
                     params::Dict) where {S <: Union{Int, String, Symbol}, U}
  if haskey(params, :value_params)
    pairs = Tuple{S, U}[]
    for k in keys(s.obj)
      key = S <: Integer ? parse(S, string(k)) : S(k)
      push!(pairs, (key, load_object(s, U, params[:value_params], k)))
    end
    isempty(pairs) && return T()
    _, v = first(pairs)

    return Dict{S, typeof(v)}(k => v for (k, v) in pairs)
  else
    dict = Dict{S, Any}()
    value_types = Type[]
    for k in keys(s.obj)
      key = S <: Integer ? parse(S, string(k)) : S(k)
      value_type, param = params[key]
      v = load_object(s, value_type, param, k)
      dict[key] = load_object(s, value_type, param, k)
      push!(value_types, typeof(v))
    end
    isempty(value_types) && return T()
    value_params = type_params.(collect(values(dict)))
    value_type = allequal(value_params) ? typejoin(unique(value_types)...) : Any
    return Dict{S, value_type}(dict)
  end
end

function load_object(s::DeserializerState,
                     T::Type{<:Dict{S, U}},
                     params::Dict) where {S, U}
  pairs = load_array_node(s) do _
    load_object(s, Tuple{S, U}, (params[:key_params], params[:value_params]))
  end
  isempty(pairs) && return T()
  k1, v1 = first(pairs)

  return Dict{typeof(k1), typeof(v1)}(k => v for (k, v) in pairs)
end

################################################################################
# Saving and loading sets
@register_serialization_type Set

function type_params(obj::T) where {S, T <: Set{S}}
  isempty(obj) && return TypeParams(T, TypeParams(S, nothing))
  return TypeParams(T, type_params(first(obj)))
end

function load_type_params(s::DeserializerState, T::Type{<: Set})
  !haskey(s, :params) && return T, nothing
  subtype, params = load_node(s, :params) do _
    subtype, params = load_type_params(s, decode_type(s))
  end
  return T{subtype}, params
end

function save_object(s::SerializerState, x::Set)
  save_data_array(s) do
    for elem in x
      if serialize_with_id(typeof(elem))
        ref = save_as_ref(s, elem)
        save_object(s, ref)
      else
        save_object(s, elem)
      end
    end
  end
end

function load_object(s::DeserializerState, S::Type{<:Set{T}}, params::Any) where T
  elems = load_array_node(s) do _
    load_object(s, T, params)
  end
  return S(elems)
end

function load_object(s::DeserializerState, S::Type{<:Set{T}}, ::Nothing) where T
  elems = load_array_node(s) do _
    load_object(s, T, nothing)
  end
  return S(elems)
end

function load_object(s::DeserializerState, S::Type{<:Set{T}}) where T
  elems = load_array_node(s) do _
    load_object(s, T)
  end
  return S(elems)
end

################################################################################
# Sparse rows

@register_serialization_type SRow 

function save_object(s::SerializerState, obj::SRow)
  save_data_array(s) do
    for (i, v) in obj
      save_object(s, (i, v))
    end
  end
end

function load_object(s::DeserializerState, ::Type{<:SRow}, params::NCRing)
  pos = Int[]
  entry_type = elem_type(params)
  values = entry_type[]
  load_array_node(s) do _
    push!(pos, load_object(s, Int, 1))
    push!(values, load_object(s, entry_type, params, 2))
  end
  return sparse_row(params, pos, values)
end