// Copyright (c) the JPEG XL Project Authors. All rights reserved.

//

// Use of this source code is governed by a BSD-style

// license that can be found in the LICENSE file.

#ifndef LIB_JXL_ENC_ANS_H_

#define LIB_JXL_ENC_ANS_H_

// Library to encode the ANS population counts to the bit-stream and encode

// symbols based on the respective distributions.

#include <jxl/memory_manager.h>

#include <cstddef>

#include <cstdint>

#include <vector>

#include "lib/jxl/ans_params.h"

#include "lib/jxl/base/status.h"

#include "lib/jxl/dec_ans.h"

#include "lib/jxl/enc_ans_params.h"

#include "lib/jxl/enc_bit_writer.h"

namespace jxl {

struct AuxOut;

enum class LayerType : uint8_t;

#define USE_MULT_BY_RECIPROCAL

// precision must be equal to:  #bits(state_) + #bits(freq)

#define RECIPROCAL_PRECISION (32 + ANS_LOG_TAB_SIZE)

// Data structure representing one element of the encoding table built

// from a distribution.

// TODO(veluca): split this up, or use an union.

struct ANSEncSymbolInfo {

  // ANS

  uint16_t freq_;

  std::vector<uint16_t> reverse_map_;

#ifdef USE_MULT_BY_RECIPROCAL

  uint64_t ifreq_;

#endif

  // Prefix coding.

  uint8_t depth;

  uint16_t bits;

};

class ANSCoder {

 public:

  ANSCoder() : state_(ANS_SIGNATURE << 16) {}

  uint32_t PutSymbol(const ANSEncSymbolInfo& t, uint8_t* nbits) {

    uint32_t bits = 0;

    *nbits = 0;

    if ((state_ >> (32 - ANS_LOG_TAB_SIZE)) >= t.freq_) {

      bits = state_ & 0xffff;

      state_ >>= 16;

      *nbits = 16;

    }

#ifdef USE_MULT_BY_RECIPROCAL

    // We use mult-by-reciprocal trick, but that requires 64b calc.

    const uint32_t v = (state_ * t.ifreq_) >> RECIPROCAL_PRECISION;

    const uint32_t offset = t.reverse_map_[state_ - v * t.freq_];

    state_ = (v << ANS_LOG_TAB_SIZE) + offset;

#else

    state_ = ((state_ / t.freq_) << ANS_LOG_TAB_SIZE) +

             t.reverse_map_[state_ % t.freq_];

#endif

    return bits;

  }

  uint32_t GetState() const { return state_; }

 private:

  uint32_t state_;

};

static const int kNumFixedHistograms = 1;

// Integer to be encoded by an entropy coder, either ANS or Huffman.

struct Token {

  Token() = default;

  Token(uint32_t c, uint32_t value)

      : is_lz77_length(false), context(c), value(value) {}

  uint32_t is_lz77_length : 1;

  uint32_t context : 31;

  uint32_t value;

};

struct SizeWriter {

  size_t size = 0;

  void Write(size_t num, size_t bits) { size += num; }

};

struct EntropyEncodingData {

  std::vector<std::vector<ANSEncSymbolInfo>> encoding_info;

  bool use_prefix_code;

  std::vector<HybridUintConfig> uint_config;

  size_t log_alpha_size;

  LZ77Params lz77;

  std::vector<BitWriter> encoded_histograms;

  std::vector<uint8_t> context_map;

  StatusOr<size_t> BuildAndStoreEntropyCodes(

      JxlMemoryManager* memory_manager, const HistogramParams& params,

      const std::vector<std::vector<Token>>& tokens,

      const std::vector<Histogram>& builder, BitWriter* writer, LayerType layer,

      AuxOut* aux_out);

  StatusOr<size_t> BuildAndStoreANSEncodingData(

      JxlMemoryManager* memory_manager,

      HistogramParams::ANSHistogramStrategy ans_histogram_strategy,

      const Histogram& histogram, BitWriter* writer);

 private:

  Status ChooseUintConfigs(JxlMemoryManager* memory_manager,

                           const HistogramParams& params,

                           const std::vector<std::vector<Token>>& tokens,

                           std::vector<Histogram>& clustered_histograms);

};

// Writes the context map to the bitstream and concatenates the individual

// histogram bitstreams in codes.encoded_histograms. Used in streaming mode.

Status EncodeHistograms(const EntropyEncodingData& codes, BitWriter* writer,

                        LayerType layer, AuxOut* aux_out);

// Apply context clustering, compute histograms and encode them. Returns an

// estimate of the total bits used for encoding the stream. If `writer` ==

// nullptr, the bit estimate will not take into account the context map (which

// does not get written if `num_contexts` == 1).

// Returns cost

StatusOr<size_t> BuildAndEncodeHistograms(

    JxlMemoryManager* memory_manager, const HistogramParams& params,

    size_t num_contexts, std::vector<std::vector<Token>>& tokens,

    EntropyEncodingData* codes, BitWriter* writer, LayerType layer,

    AuxOut* aux_out);

// Write the tokens to a string.

Status WriteTokens(const std::vector<Token>& tokens,

                   const EntropyEncodingData& codes, size_t context_offset,

                   BitWriter* writer, LayerType layer, AuxOut* aux_out);

// Same as above, but assumes allotment created by caller.

size_t WriteTokens(const std::vector<Token>& tokens,

                   const EntropyEncodingData& codes, size_t context_offset,

                   BitWriter* writer);

// Exposed for tests; to be used with Writer=BitWriter only.

template <typename Writer>

void EncodeUintConfigs(const std::vector<HybridUintConfig>& uint_config,

                       Writer* writer, size_t log_alpha_size);

extern template void EncodeUintConfigs(const std::vector<HybridUintConfig>&,

                                       BitWriter*, size_t);

// Globally set the option to create fuzzer-friendly ANS streams. Negatively

// impacts compression. Not thread-safe.

void SetANSFuzzerFriendly(bool ans_fuzzer_friendly);

}  // namespace jxl

#endif  // LIB_JXL_ENC_ANS_H_