/* -----------------------------------------------------------------------------
Copyright 2021 Kevin P. Barry

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
----------------------------------------------------------------------------- */

// Author: Kevin P. Barry [ta0kira@gmail.com]

// A read-only categorical distribution.
//
// Params:
// - #c: The type of category (i.e., object) used in the distribution.
@value interface CategoricalReader<#c> {
  // Get the sum of all weights in the distribution.
  getTotal () -> (Int)

  // Get the relative weight of a value.
  //
  // Notes:
  // - If the category isn't present, its weight is 0.
  getWeight (#c) -> (Int)

  // Return the value at the given offset.
  //
  // Notes:
  // - The offset must be within [0, getTotal()). A uniform selection in that
  //   range will provide samples that follow the categorical distribution
  //   corresponding to the relative weights of the respective #c.
  // - The return value is deterministic. If you were to iterate over
  //   [0, getTotal()), you'd get an increasing sequence of all #c in the
  //   CategoricalReader, each repeated the number of times indicated by its
  //   respective weight.
  // - Also see RandomCategorical in random.0rp.
  locate (Int) -> (#c)
}

// A categorical distribution represented as a tree.
//
// Params:
// - #c: The type of category (i.e., object) used in the distribution.
//
// Notes:
// - CategoricalTree is intended for use in random sampling of arbitrary objects
//   based on relative weights that can be dynamically set.
// - The distribution automatically updates every time a weight is changed. This
//   can be used for updating Bayesian observations while simultaneously being
//   able to sample values from the distribution.
// - The complexity of most operations is O(log n) with n distinct #c values.
// - The required storage space is independent of the weights; it only depends
//   on the number of distinct #c values with non-zero weights.
concrete CategoricalTree<#c> {
  defines Default
  refines CategoricalReader<#c>
  refines Duplicate
  refines ReadAt<#c>
  #c immutable
  #c defines LessThan<#c>

  // Create a new distribution.
  @type new () -> (#self)

  // Set the relative weight of a value.
  //
  // Notes:
  // - The weight must not be negative.
  // - The sum of all weights (see getTotal()) must not exceed Int.maxBound().
  @value setWeight (#c, Int) -> (#self)

  // Increments the weight of the value by 1.
  @value incrWeight (#c) -> (#self)

  // Decrements the weight of the value by 1.
  //
  // Notes:
  // - Do not call this for values that already have 0 weight. This is primarily
  //   meant for use when locate is used to choose a value to decrement.
  @value decrWeight (#c) -> (#self)

  // Identical to locate(). (From ReadAt.)
  @value readAt (Int) -> (#c)

  // Identical to getTotal(). (From ReadAt/Container.)
  @value size () -> (Int)
}