// 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_MODULAR_ENCODING_ENC_MA_H_

#define LIB_JXL_MODULAR_ENCODING_ENC_MA_H_

#include <algorithm>

#include <array>

#include <cstddef>

#include <cstdint>

#include <vector>

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

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

#include "lib/jxl/enc_ans.h"

#include "lib/jxl/modular/encoding/dec_ma.h"

#include "lib/jxl/modular/modular_image.h"

#include "lib/jxl/modular/options.h"

namespace jxl {

struct ResidualToken {

  uint8_t tok;

  uint8_t nbits;

};

// Struct to collect all the data needed to build a tree.

struct TreeSamples {

  bool HasSamples() const {

    return !residuals.empty() && !residuals[0].empty();

  }

  size_t NumDistinctSamples() const { return sample_counts.size(); }

  size_t NumSamples() const { return num_samples; }

  // Set the predictor to use. Must be called before adding any samples.

  Status SetPredictor(Predictor predictor,

                      ModularOptions::TreeMode wp_tree_mode);

  // Set the properties to use. Must be called before adding any samples.

  Status SetProperties(const std::vector<uint32_t> &properties,

                       ModularOptions::TreeMode wp_tree_mode);

  const ResidualToken &RToken(size_t pred, size_t i) const {

    return residuals[pred][i];

  }

  size_t Token(size_t pred, size_t i) const { return residuals[pred][i].tok; }

  size_t Count(size_t i) const { return sample_counts[i]; }

  size_t PredictorIndex(Predictor predictor) const {

    const auto predictor_elem =

        std::find(predictors.begin(), predictors.end(), predictor);

    JXL_DASSERT(predictor_elem != predictors.end());

    return predictor_elem - predictors.begin();

  }

  size_t PropertyIndex(size_t property) const {

    const auto property_elem =

        std::find(props_to_use.begin(), props_to_use.end(), property);

    JXL_DASSERT(property_elem != props_to_use.end());

    return property_elem - props_to_use.begin();

  }

  size_t NumPropertyValues(size_t property_index) const {

    return compact_properties[property_index].size() + 1;

  }

  // Returns the *quantized* property value.

  size_t Property(size_t property_index, size_t i) const {

    return property_index < num_static_props

               ? static_props[property_index][i]

               : props[property_index - num_static_props][i];

  }

  int UnquantizeProperty(size_t property_index, uint32_t quant) const {

    JXL_DASSERT(quant < compact_properties[property_index].size());

    return compact_properties[property_index][quant];

  }

  Predictor PredictorFromIndex(size_t index) const {

    JXL_DASSERT(index < predictors.size());

    return predictors[index];

  }

  size_t PropertyFromIndex(size_t index) const {

    JXL_DASSERT(index < props_to_use.size());

    return props_to_use[index];

  }

  size_t NumPredictors() const { return predictors.size(); }

  size_t NumProperties() const { return props_to_use.size(); }

  // Preallocate data for a given number of samples. MUST be called before

  // adding any sample.

  void PrepareForSamples(size_t extra_num_samples);

  // Add a sample.

  void AddSample(pixel_type_w pixel, const Properties &properties,

                 const pixel_type_w *predictions);

  // Pre-cluster property values.

  void PreQuantizeProperties(

      const StaticPropRange &range,

      const std::vector<ModularMultiplierInfo> &multiplier_info,

      const std::vector<uint32_t> &group_pixel_count,

      const std::vector<uint32_t> &channel_pixel_count,

      std::vector<pixel_type> &pixel_samples,

      std::vector<pixel_type> &diff_samples, size_t max_property_values);

  void AllSamplesDone() { dedup_table_ = std::vector<uint32_t>(); }

  uint32_t QuantizeProperty(uint32_t prop, pixel_type v) const {

    JXL_DASSERT(prop >= num_static_props);

    v = jxl::Clamp1(v, -kPropertyRange, kPropertyRange) + kPropertyRange;

    return property_mapping[prop - num_static_props][v];

  }

  uint32_t QuantizeStaticProperty(uint32_t prop, pixel_type v) const {

    JXL_DASSERT(prop < num_static_props);

    v = jxl::Clamp1(v, -kPropertyRange, kPropertyRange) + kPropertyRange;

    return static_property_mapping[prop][v];

  }

  // Swaps samples in position a and b. Does nothing if a == b.

  void Swap(size_t a, size_t b);

 private:

  // TODO(veluca): as the total number of properties and predictors are known

  // before adding any samples, it might be better to interleave predictors,

  // properties and counts in a single vector to improve locality.

  // A first attempt at doing this actually results in much slower encoding,

  // possibly because of the more complex addressing.

  // Residual information: token and number of extra bits, per predictor.

  std::vector<std::vector<ResidualToken>> residuals;

  // Number of occurrences of each sample.

  std::vector<uint16_t> sample_counts;

  // Quantized static property values

  size_t num_static_props;

  std::array<std::vector<uint32_t>, kNumStaticProperties> static_props;

  // Property values, quantized to at most 256 distinct values.

  std::vector<std::vector<uint8_t>> props;

  // Decompactification info for `props`.

  std::vector<std::vector<int32_t>> compact_properties;

  // List of properties to use.

  std::vector<uint32_t> props_to_use;

  // List of predictors to use.

  std::vector<Predictor> predictors;

  // Mapping property value -> quantized property value.

  static constexpr int32_t kPropertyRange = 511;

  std::array<std::vector<uint16_t>, kNumStaticProperties>

      static_property_mapping;

  std::vector<std::vector<uint8_t>> property_mapping;

  // Number of samples seen.

  size_t num_samples = 0;

  // Table for deduplication.

  static constexpr uint32_t kDedupEntryUnused{static_cast<uint32_t>(-1)};

  std::vector<uint32_t> dedup_table_;

  // Functions for sample deduplication.

  bool IsSameSample(size_t a, size_t b) const;

  size_t Hash1(size_t a) const;

  size_t Hash2(size_t a) const;

  void InitTable(size_t log_size);

  // Returns true if `a` was already present in the table.

  bool AddToTableAndMerge(size_t a);

  void AddToTable(size_t a);

};

Status TokenizeTree(const Tree &tree, std::vector<Token> *tokens,

                    Tree *decoder_tree);

void CollectPixelSamples(const Image &image, const ModularOptions &options,

                         uint32_t group_id,

                         std::vector<uint32_t> &group_pixel_count,

                         std::vector<uint32_t> &channel_pixel_count,

                         std::vector<pixel_type> &pixel_samples,

                         std::vector<pixel_type> &diff_samples);

Status ComputeBestTree(TreeSamples &tree_samples, float threshold,

                       const std::vector<ModularMultiplierInfo> &mul_info,

                       StaticPropRange static_prop_range,

                       float fast_decode_multiplier, Tree *tree);

}  // namespace jxl

#endif  // LIB_JXL_MODULAR_ENCODING_ENC_MA_H_