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

#include "lib/jxl/enc_modular.h"

#include <jxl/cms_interface.h>

#include <jxl/memory_manager.h>

#include <jxl/types.h>

#include <algorithm>

#include <array>

#include <cmath>

#include <cstddef>

#include <cstdint>

#include <cstdlib>

#include <cstring>

#include <limits>

#include <memory>

#include <tuple>

#include <utility>

#include <vector>

#include "lib/jxl/ac_strategy.h"

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

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

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

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

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

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

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

#include "lib/jxl/chroma_from_luma.h"

#include "lib/jxl/common.h"

#include "lib/jxl/compressed_dc.h"

#include "lib/jxl/dec_ans.h"

#include "lib/jxl/dec_modular.h"

#include "lib/jxl/enc_ans.h"

#include "lib/jxl/enc_ans_params.h"

#include "lib/jxl/enc_aux_out.h"

#include "lib/jxl/enc_bit_writer.h"

#include "lib/jxl/enc_cache.h"

#include "lib/jxl/enc_fields.h"

#include "lib/jxl/enc_gaborish.h"

#include "lib/jxl/enc_modular_simd.h"

#include "lib/jxl/enc_params.h"

#include "lib/jxl/enc_patch_dictionary.h"

#include "lib/jxl/enc_quant_weights.h"

#include "lib/jxl/fields.h"

#include "lib/jxl/frame_dimensions.h"

#include "lib/jxl/frame_header.h"

#include "lib/jxl/image.h"

#include "lib/jxl/image_metadata.h"

#include "lib/jxl/image_ops.h"

#include "lib/jxl/memory_manager_internal.h"

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

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

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

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

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

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

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

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

#include "lib/jxl/modular/transform/enc_rct.h"

#include "lib/jxl/modular/transform/enc_transform.h"

#include "lib/jxl/modular/transform/squeeze.h"

#include "lib/jxl/modular/transform/squeeze_params.h"

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

#include "lib/jxl/pack_signed.h"

#include "lib/jxl/passes_state.h"

#include "lib/jxl/quant_weights.h"

#include "modular/options.h"

namespace jxl {

namespace {

// constexpr bool kPrintTree = false;

// Squeeze default quantization factors

// these quantization factors are for -Q 50  (other qualities simply scale the

// factors; things are rounded down and obviously cannot get below 1)

const float squeeze_quality_factor =

    0.35;  // for easy tweaking of the quality range (decrease this number for

           // higher quality)

const float squeeze_luma_factor =

    1.1;  // for easy tweaking of the balance between luma (or anything

          // non-chroma) and chroma (decrease this number for higher quality

          // luma)

const float squeeze_quality_factor_xyb = 4.0f;

const float squeeze_quality_factor_y = 1.5f;

const float squeeze_xyb_qtable[3][16] = {

    {163.84, 81.92, 40.96, 20.48, 10.24, 5.12, 2.56, 1.28, 0.64, 0.32, 0.16,

     0.08, 0.04, 0.02, 0.01, 0.005},  // Y

    {1024, 512, 256, 128, 64, 32, 16, 8, 4, 2, 1, 0.5, 0.5, 0.5, 0.5,

     0.5},  // X

    {2048, 1024, 512, 256, 128, 64, 32, 16, 8, 4, 2, 1, 0.5, 0.5, 0.5,

     0.5},  // B-Y

};

const float squeeze_luma_qtable[16] = {163.84, 81.92, 40.96, 20.48, 10.24, 5.12,

                                       2.56,   1.28,  0.64,  0.32,  0.16,  0.08,

                                       0.04,   0.02,  0.01,  0.005};

// for 8-bit input, the range of YCoCg chroma is -255..255 so basically this

// does 4:2:0 subsampling (two most fine grained layers get quantized away)

const float squeeze_chroma_qtable[16] = {

    1024, 512, 256, 128, 64, 32, 16, 8, 4, 2, 1, 0.5, 0.5, 0.5, 0.5, 0.5};

// Merges the trees in `trees` using nodes that decide on stream_id, as defined

// by `tree_splits`.

Status MergeTrees(const std::vector<Tree>& trees,

                  const std::vector<size_t>& tree_splits, size_t begin,

                  size_t end, Tree* tree) {

  JXL_ENSURE(trees.size() + 1 == tree_splits.size());

  JXL_ENSURE(end > begin);

  JXL_ENSURE(end <= trees.size());

  if (end == begin + 1) {

    // Insert the tree, adding the opportune offset to all child nodes.

    // This will make the leaf IDs wrong, but subsequent roundtripping will fix

    // them.

    size_t sz = tree->size();

    tree->insert(tree->end(), trees[begin].begin(), trees[begin].end());

    for (size_t i = sz; i < tree->size(); i++) {

      (*tree)[i].lchild += sz;

      (*tree)[i].rchild += sz;

    }

    return true;

  }

  size_t mid = (begin + end) / 2;

  size_t splitval = tree_splits[mid] - 1;

  size_t cur = tree->size();

  tree->emplace_back(1 /*stream_id*/, static_cast<int>(splitval), 0, 0,

                     Predictor::Zero, 0, 1);

  (*tree)[cur].lchild = tree->size();

  JXL_RETURN_IF_ERROR(MergeTrees(trees, tree_splits, mid, end, tree));

  (*tree)[cur].rchild = tree->size();

  JXL_RETURN_IF_ERROR(MergeTrees(trees, tree_splits, begin, mid, tree));

  return true;

}

void QuantizeChannel(Channel& ch, const int q) {

  if (q == 1) return;

  for (size_t y = 0; y < ch.plane.ysize(); y++) {

    pixel_type* row = ch.plane.Row(y);

    for (size_t x = 0; x < ch.plane.xsize(); x++) {

      if (row[x] < 0) {

        row[x] = -((-row[x] + q / 2) / q) * q;

      } else {

        row[x] = ((row[x] + q / 2) / q) * q;

      }

    }

  }

}

// convert binary32 float that corresponds to custom [bits]-bit float (with

// [exp_bits] exponent bits) to a [bits]-bit integer representation that should

// fit in pixel_type

Status float_to_int(const float* const row_in, pixel_type* const row_out,

                    size_t xsize, unsigned int bits, unsigned int exp_bits,

                    bool fp, double dfactor) {

  JXL_ENSURE(sizeof(pixel_type) * 8 >= bits);

  if (!fp) {

    if (bits > 22) {

      for (size_t x = 0; x < xsize; ++x) {

        row_out[x] = row_in[x] * dfactor + (row_in[x] < 0 ? -0.5 : 0.5);

      }

    } else {

      float factor = dfactor;

      for (size_t x = 0; x < xsize; ++x) {

        row_out[x] = row_in[x] * factor + (row_in[x] < 0 ? -0.5f : 0.5f);

      }

    }

    return true;

  }

  if (bits == 32 && fp) {

    JXL_ENSURE(exp_bits == 8);

    memcpy(static_cast<void*>(row_out), static_cast<const void*>(row_in),

           4 * xsize);

    return true;

  }

  JXL_ENSURE(bits > 0);

  int exp_bias = (1 << (exp_bits - 1)) - 1;

  int max_exp = (1 << exp_bits) - 1;

  uint32_t sign = (1u << (bits - 1));

  int mant_bits = bits - exp_bits - 1;

  int mant_shift = 23 - mant_bits;

  for (size_t x = 0; x < xsize; ++x) {

    uint32_t f;

    memcpy(&f, &row_in[x], 4);

    int signbit = (f >> 31);

    f &= 0x7fffffff;

    if (f == 0) {

      row_out[x] = (signbit ? sign : 0);

      continue;

    }

    int exp = (f >> 23) - 127;

    int mantissa = (f & 0x007fffff);

    // broke up the binary32 into its parts, now reassemble into

    // arbitrary float

    if (exp == 128) {

      // NaN or infinity

      f = (signbit ? sign : 0);

      f |= ((1 << exp_bits) - 1) << mant_bits;

      f |= mantissa >> mant_shift;

      row_out[x] = static_cast<pixel_type>(f);

      continue;

    }

    exp += exp_bias;

    if (exp <= 0) {  // will become a subnormal number

      // add implicit leading 1 to mantissa

      mantissa |= 0x00800000;

      if (exp < -mant_bits) {

        return JXL_FAILURE(

            "Invalid float number: %g cannot be represented with %i "

            "exp_bits and %i mant_bits (exp %i)",

            row_in[x], exp_bits, mant_bits, exp);

      }

      mantissa >>= 1 - exp;

      exp = 0;

    }

    // exp should be representable in exp_bits, otherwise input was

    // invalid; max_exp is NaN or infinity

    if (exp >= max_exp) return JXL_FAILURE("Invalid float exponent");

    if (mantissa & ((1 << mant_shift) - 1)) {

      return JXL_FAILURE("%g is losing precision (mant: %x)", row_in[x],

                         mantissa);

    }

    mantissa >>= mant_shift;

    f = (signbit ? sign : 0);

    f |= (exp << mant_bits);

    f |= mantissa;

    row_out[x] = static_cast<pixel_type>(f);

  }

  return true;

}

float EstimateWPCost(const Image& img, size_t i) {

  size_t extra_bits = 0;

  float histo_cost = 0;

  HybridUintConfig config;

  int32_t cutoffs[] = {-500, -392, -255, -191, -127, -95, -63, -47, -31,

                       -23,  -15,  -11,  -7,   -4,   -3,  -1,  0,   1,

                       3,    5,    7,    11,   15,   23,  31,  47,  63,

                       95,   127,  191,  255,  392,  500};

  constexpr size_t nc = sizeof(cutoffs) / sizeof(*cutoffs) + 1;

  Histogram histo[nc] = {};

  weighted::Header wp_header;

  PredictorMode(i, &wp_header);

  for (const Channel& ch : img.channel) {

    const ptrdiff_t onerow = ch.plane.PixelsPerRow();

    weighted::State wp_state(wp_header, ch.w, ch.h);

    Properties properties(1);

    for (size_t y = 0; y < ch.h; y++) {

      const pixel_type* JXL_RESTRICT r = ch.Row(y);

      for (size_t x = 0; x < ch.w; x++) {

        size_t offset = 0;

        pixel_type_w left = (x ? r[x - 1] : y ? *(r + x - onerow) : 0);

        pixel_type_w top = (y ? *(r + x - onerow) : left);

        pixel_type_w topleft = (x && y ? *(r + x - 1 - onerow) : left);

        pixel_type_w topright =

            (x + 1 < ch.w && y ? *(r + x + 1 - onerow) : top);

        pixel_type_w toptop = (y > 1 ? *(r + x - onerow - onerow) : top);

        pixel_type guess = wp_state.Predict</*compute_properties=*/true>(

            x, y, ch.w, top, left, topright, topleft, toptop, &properties,

            offset);

        size_t ctx = 0;

        for (int c : cutoffs) {

          ctx += (c >= properties[0]) ? 1 : 0;

        }

        pixel_type res = r[x] - guess;

        uint32_t token;

        uint32_t nbits;

        uint32_t bits;

        config.Encode(PackSigned(res), &token, &nbits, &bits);

        histo[ctx].Add(token);

        extra_bits += nbits;

        wp_state.UpdateErrors(r[x], x, y, ch.w);

      }

    }

    for (auto& h : histo) {

      histo_cost += h.ShannonEntropy();

      h.Clear();

    }

  }

  return histo_cost + extra_bits;

}

bool do_transform(Image& image, const Transform& tr,

                  const weighted::Header& wp_header,

                  jxl::ThreadPool* pool = nullptr, bool force_jxlart = false) {

  Transform t = tr;

  bool did_it = true;

  if (force_jxlart) {

    if (!t.MetaApply(image)) return false;

  } else {

    did_it = TransformForward(t, image, wp_header, pool);

  }

  if (did_it) image.transform.push_back(t);

  return did_it;

}

StatusOr<bool> maybe_do_transform(Image& image, const Transform& tr,

                                  const CompressParams& cparams,

                                  const weighted::Header& wp_header,

                                  float cost_before,

                                  jxl::ThreadPool* pool = nullptr,

                                  bool force_jxlart = false) {

  if (force_jxlart || cparams.speed_tier >= SpeedTier::kSquirrel) {

    return do_transform(image, tr, wp_header, pool, force_jxlart);

  }

  bool did_it = do_transform(image, tr, wp_header, pool);

  if (did_it) {

    JXL_ASSIGN_OR_RETURN(float cost_after, EstimateCost(image));

    JXL_DEBUG_V(7, "Cost before: %f  cost after: %f", cost_before, cost_after);

    if (cost_after > cost_before) {

      Transform t = image.transform.back();

      if (!t.Inverse(image, wp_header, pool)) {

        return false;

      }

      image.transform.pop_back();

      did_it = false;

    }

  }

  return did_it;

}

Status try_palettes(Image& gi, int& max_bitdepth, int& maxval,

                    const CompressParams& cparams_,

                    float channel_colors_percent,

                    jxl::ThreadPool* pool = nullptr) {

  float cost_before = 0.f;

  size_t did_palette = 0;

  float nb_pixels = gi.channel[0].w * gi.channel[0].h;

  int nb_chans = gi.channel.size() - gi.nb_meta_channels;

  // arbitrary estimate: 4.8 bpp for 8-bit RGB

  float arbitrary_bpp_estimate = 0.2f * gi.bitdepth * nb_chans;

  if (cparams_.palette_colors != 0 || cparams_.lossy_palette) {

    // when not estimating, assume some arbitrary bpp

    if (cparams_.speed_tier <= SpeedTier::kSquirrel) {

      JXL_ASSIGN_OR_RETURN(cost_before, EstimateCost(gi));

    } else {

      cost_before = nb_pixels * arbitrary_bpp_estimate;

    }

    // all-channel palette (e.g. RGBA)

    if (nb_chans > 1) {

      Transform maybe_palette(TransformId::kPalette);

      maybe_palette.begin_c = gi.nb_meta_channels;

      maybe_palette.num_c = nb_chans;

      // Heuristic choice of max colors for a palette:

      // max_colors = nb_pixels * estimated_bpp_without_palette * 0.0005 +

      //              + nb_pixels / 128 + 128

      //       (estimated_bpp_without_palette = cost_before / nb_pixels)

      // Rationale: small image with large palette is not effective;

      // also if the entropy (estimated bpp) is low (e.g. mostly solid/gradient

      // areas), palette is less useful and may even be counterproductive.

      maybe_palette.nb_colors = std::min(

          static_cast<int>(cost_before * 0.0005f + nb_pixels / 128 + 128),

          std::abs(cparams_.palette_colors));

      maybe_palette.ordered_palette = cparams_.palette_colors >= 0;

      maybe_palette.lossy_palette =

          (cparams_.lossy_palette && maybe_palette.num_c == 3);

      if (maybe_palette.lossy_palette) {

        maybe_palette.predictor = Predictor::Average4;

      }

      // TODO(veluca): use a custom weighted header if using the weighted

      // predictor.

      JXL_ASSIGN_OR_RETURN(

          did_palette,

          maybe_do_transform(gi, maybe_palette, cparams_, weighted::Header(),

                             cost_before, pool, cparams_.options.zero_tokens));

    }

    // all-minus-one-channel palette (RGB with separate alpha, or CMY with

    // separate K)

    if (!did_palette && nb_chans > 3) {

      Transform maybe_palette_3(TransformId::kPalette);

      maybe_palette_3.begin_c = gi.nb_meta_channels;

      maybe_palette_3.num_c = nb_chans - 1;

      maybe_palette_3.nb_colors = std::min(

          static_cast<int>(cost_before * 0.0005f + nb_pixels / 128 + 128),

          std::abs(cparams_.palette_colors));

      maybe_palette_3.ordered_palette = cparams_.palette_colors >= 0;

      maybe_palette_3.lossy_palette = cparams_.lossy_palette;

      if (maybe_palette_3.lossy_palette) {

        maybe_palette_3.predictor = Predictor::Average4;

      }

      JXL_ASSIGN_OR_RETURN(

          did_palette,

          maybe_do_transform(gi, maybe_palette_3, cparams_, weighted::Header(),

                             cost_before, pool, cparams_.options.zero_tokens));

    }

  }

  if (channel_colors_percent > 0) {

    // single channel palette (like FLIF's ChannelCompact)

    size_t nb_channels = gi.channel.size() - gi.nb_meta_channels - did_palette;

    int orig_bitdepth = max_bitdepth;

    max_bitdepth = 0;

    if (nb_channels > 0 && (did_palette || cost_before == 0)) {

      if (cparams_.speed_tier < SpeedTier::kSquirrel) {

        JXL_ASSIGN_OR_RETURN(cost_before, EstimateCost(gi));

      } else {

        cost_before = 0;

      }

    }

    for (size_t i = did_palette; i < nb_channels + did_palette; i++) {

      int32_t min;

      int32_t max;

      compute_minmax(gi.channel[gi.nb_meta_channels + i], &min, &max);

      int64_t colors = static_cast<int64_t>(max) - min + 1;

      JXL_DEBUG_V(10, "Channel %" PRIuS ": range=%i..%i", i, min, max);

      Transform maybe_palette_1(TransformId::kPalette);

      maybe_palette_1.begin_c = i + gi.nb_meta_channels;

      maybe_palette_1.num_c = 1;

      // simple heuristic: if less than X percent of the values in the range

      // actually occur, it is probably worth it to do a compaction

      // (but only if the channel palette is less than 6% the size of the

      // image itself)

      maybe_palette_1.nb_colors =

          std::min(static_cast<int>(nb_pixels / 16),

                   static_cast<int>(channel_colors_percent / 100. * colors));

      JXL_ASSIGN_OR_RETURN(

          bool did_ch_palette,

          maybe_do_transform(gi, maybe_palette_1, cparams_, weighted::Header(),

                             cost_before, pool));

      if (did_ch_palette) {

        // effective bit depth is lower, adjust quantization accordingly

        compute_minmax(gi.channel[gi.nb_meta_channels + i], &min, &max);

        if (max < maxval) maxval = max;

        int ch_bitdepth =

            (max > 0 ? CeilLog2Nonzero(static_cast<uint32_t>(max)) : 0);

        if (ch_bitdepth > max_bitdepth) max_bitdepth = ch_bitdepth;

      } else {

        max_bitdepth = orig_bitdepth;

      }

    }

  }

  return true;

}

}  // namespace

StatusOr<std::unique_ptr<ModularFrameEncoder>> ModularFrameEncoder::Create(

    JxlMemoryManager* memory_manager, const FrameHeader& frame_header,

    const CompressParams& cparams_orig, bool streaming_mode) {

  auto self = std::unique_ptr<ModularFrameEncoder>(

      new ModularFrameEncoder(memory_manager));

  JXL_RETURN_IF_ERROR(self->Init(frame_header, cparams_orig, streaming_mode));

  return self;

}

ModularFrameEncoder::ModularFrameEncoder(JxlMemoryManager* memory_manager)

    : memory_manager_(memory_manager) {}

Status ModularFrameEncoder::Init(const FrameHeader& frame_header,

                                 const CompressParams& cparams_orig,

                                 bool streaming_mode) {

  frame_dim_ = frame_header.ToFrameDimensions();

  cparams_ = cparams_orig;

  size_t num_streams =

      ModularStreamId::Num(frame_dim_, frame_header.passes.num_passes);

  // Progressive lossless only benefits from levels 2 and higher

  // Lower levels of faster decoding can outperform higher tiers

  // depending on the PC

  if (cparams_.responsive == 1 && cparams_.IsLossless() &&

      cparams_.decoding_speed_tier == 1) {

    cparams_.decoding_speed_tier = 2;

  }

  if (cparams_.responsive == 1 && cparams_.IsLossless()) {

    // RCT selection seems bugged with Squeeze, YCoCg works well.

    if (cparams_.colorspace < 0) {

      cparams_.colorspace = 6;

    }

  }

  if (cparams_.ModularPartIsLossless()) {

    const auto disable_wp = [this] () {

        cparams_.options.wp_tree_mode = ModularOptions::TreeMode::kNoWP;

        if (cparams_.options.predictor == Predictor::Weighted) {

          // Predictor::Best turns to Predictor::Gradient anyways.

          cparams_.options.predictor = Predictor::Gradient;

        }

    };

    switch (cparams_.decoding_speed_tier) {

      case 0:

        cparams_.options.fast_decode_multiplier = 1.001f;

        break;

      case 1:  // No Weighted predictor

        cparams_.options.fast_decode_multiplier = 1.005f;

        disable_wp();

        break;

      case 2: {  // No Weighted predictor and Group size 0 defined in

                 // enc_frame.cc

        cparams_.options.fast_decode_multiplier = 1.015f;

        disable_wp();

        break;

      }

      case 3: {  // Gradient only, Group size 0, and Fast MA tree

        cparams_.options.wp_tree_mode = ModularOptions::TreeMode::kGradientOnly;

        cparams_.options.predictor = Predictor::Gradient;

        break;

      }

      default: {  // Gradient only, Group size 0, and No MA tree

        cparams_.options.wp_tree_mode = ModularOptions::TreeMode::kGradientOnly;

        cparams_.options.predictor = Predictor::Gradient;

        cparams_.options.nb_repeats = 0;

        // Disabling MA Trees sometimes doesn't increase decode speed

        // depending on PC

        break;

      }

    }

  }

  for (size_t i = 0; i < num_streams; ++i) {

    stream_images_.emplace_back(memory_manager_);

  }

  // use a sensible default if nothing explicit is specified:

  // Squeeze for lossy, no squeeze for lossless

  if (cparams_.responsive < 0) {

    if (cparams_.ModularPartIsLossless()) {

      cparams_.responsive = 0;

    } else {

      cparams_.responsive = 1;

    }

  }

  cparams_.options.splitting_heuristics_node_threshold =

      75 + 14 * static_cast<int>(cparams_.speed_tier) +

      10 * cparams_.decoding_speed_tier;

  {

    // Set properties.

    std::vector<uint32_t> prop_order;

    if (cparams_.responsive) {

      // Properties in order of their likelihood of being useful for Squeeze

      // residuals.

      prop_order = {0, 1, 4, 5, 6, 7, 8, 15, 9, 10, 11, 12, 13, 14, 2, 3};

    } else {

      // Same, but for the non-Squeeze case.

      prop_order = {0, 1, 15, 9, 10, 11, 12, 13, 14, 2, 3, 4, 5, 6, 7, 8};

      // if few groups, don't use group as a property

      if (num_streams < 30 && cparams_.speed_tier > SpeedTier::kTortoise &&

          cparams_orig.ModularPartIsLossless()) {

        prop_order.erase(prop_order.begin() + 1);

      }

    }

    int max_properties = std::min<int>(

        cparams_.options.max_properties,

        static_cast<int>(

            frame_header.nonserialized_metadata->m.num_extra_channels) +

            (frame_header.encoding == FrameEncoding::kModular ? 2 : -1));

    switch (cparams_.speed_tier) {

      case SpeedTier::kHare:

        cparams_.options.splitting_heuristics_properties.assign(

            prop_order.begin(), prop_order.begin() + 4);

        cparams_.options.max_property_values = 48;

        cparams_.options.nb_repeats *= 0.5f;

        break;

      case SpeedTier::kWombat:

        cparams_.options.splitting_heuristics_properties.assign(

            prop_order.begin(), prop_order.begin() + 5);

        cparams_.options.max_property_values = 64;

        cparams_.options.nb_repeats *= 0.7f;

        break;

      case SpeedTier::kSquirrel:

        cparams_.options.splitting_heuristics_properties.assign(

            prop_order.begin(), prop_order.begin() + 7);

        cparams_.options.max_property_values = 96;

        break;

      case SpeedTier::kKitten:

        cparams_.options.splitting_heuristics_properties.assign(

            prop_order.begin(), prop_order.begin() + 10);

        cparams_.options.max_property_values = 128;

        cparams_.options.nb_repeats *= 1.1f;

        break;

      case SpeedTier::kGlacier:

      case SpeedTier::kTortoise:

        cparams_.options.splitting_heuristics_properties = prop_order;

        cparams_.options.max_property_values = 256;

        cparams_.options.nb_repeats *= 1.3f;

        break;

      default:

        cparams_.options.splitting_heuristics_properties.assign(

            prop_order.begin(), prop_order.begin() + 3);

        cparams_.options.max_property_values = 32;

        cparams_.options.nb_repeats *= 0.3f;

        break;

    }

    if (cparams_.speed_tier > SpeedTier::kTortoise) {

      // Gradient in previous channels.

      for (int i = 0; i < max_properties; i++) {

        cparams_.options.splitting_heuristics_properties.push_back(

            kNumNonrefProperties + i * 4 + 3);

      }

    } else {

      // All the extra properties in Tortoise mode.

      for (int i = 0; i < max_properties * 4; i++) {

        cparams_.options.splitting_heuristics_properties.push_back(

            kNumNonrefProperties + i);

      }

    }

  }

  cparams_.options.nb_repeats = std::min(1.0f, cparams_.options.nb_repeats);

  if ((cparams_.options.predictor == Predictor::Average0 ||

       cparams_.options.predictor == Predictor::Average1 ||

       cparams_.options.predictor == Predictor::Average2 ||

       cparams_.options.predictor == Predictor::Average3 ||

       cparams_.options.predictor == Predictor::Average4 ||

       cparams_.options.predictor == Predictor::Weighted) &&

      !cparams_.ModularPartIsLossless()) {

    // Lossy + Average/Weighted predictors does not work, so switch to default

    // predictors.

    cparams_.options.predictor = kUndefinedPredictor;

  }

  if (cparams_.options.predictor == kUndefinedPredictor) {

    // no explicit predictor(s) given, set a good default

    if ((cparams_.speed_tier <= SpeedTier::kGlacier ||

         cparams_.modular_mode == false) &&

        cparams_.IsLossless() && cparams_.responsive == JXL_FALSE) {

      // TODO(veluca): allow all predictors that don't break residual

      // multipliers in lossy mode.

      cparams_.options.predictor = Predictor::Variable;

    } else if (cparams_.responsive || cparams_.lossy_palette) {

      // zero predictor for Squeeze residues and lossy palette indices

      // TODO: Try adding 'Squeezed' predictor set, with the most

      // common predictors used by Variable in squeezed images, including none.

      cparams_.options.predictor = Predictor::Zero;

    } else if (!cparams_.IsLossless()) {

      // If not responsive and lossy. TODO(veluca): use near_lossless instead?

      cparams_.options.predictor = Predictor::Gradient;

    } else if (cparams_.speed_tier < SpeedTier::kFalcon) {

      // try median and weighted predictor for anything else

      cparams_.options.predictor = Predictor::Best;

    } else if (cparams_.speed_tier == SpeedTier::kFalcon) {

      // just weighted predictor in falcon mode

      cparams_.options.predictor = Predictor::Weighted;

    } else if (cparams_.speed_tier > SpeedTier::kFalcon) {

      // just gradient predictor in thunder mode

      cparams_.options.predictor = Predictor::Gradient;

    }

  } else {

    if (cparams_.lossy_palette) cparams_.options.predictor = Predictor::Zero;

  }

  if (!cparams_.ModularPartIsLossless()) {

    if (cparams_.options.predictor == Predictor::Weighted ||

        cparams_.options.predictor == Predictor::Variable ||

        cparams_.options.predictor == Predictor::Best)

      cparams_.options.predictor = Predictor::Zero;

  }

  tree_splits_.push_back(0);

  if (cparams_.modular_mode == false) {

    JXL_ASSIGN_OR_RETURN(ModularStreamId qt0, ModularStreamId::QuantTable(0));

    cparams_.options.fast_decode_multiplier = 1.0f;

    tree_splits_.push_back(ModularStreamId::VarDCTDC(0).ID(frame_dim_));

    tree_splits_.push_back(ModularStreamId::ModularDC(0).ID(frame_dim_));

    tree_splits_.push_back(ModularStreamId::ACMetadata(0).ID(frame_dim_));

    tree_splits_.push_back(qt0.ID(frame_dim_));

    tree_splits_.push_back(ModularStreamId::ModularAC(0, 0).ID(frame_dim_));

    ac_metadata_size.resize(frame_dim_.num_dc_groups);

    extra_dc_precision.resize(frame_dim_.num_dc_groups);

  }

  tree_splits_.push_back(num_streams);

  cparams_.options.max_chan_size = frame_dim_.group_dim;

  cparams_.options.group_dim = frame_dim_.group_dim;

  // TODO(veluca): figure out how to use different predictor sets per channel.

  stream_options_.resize(num_streams, cparams_.options);

  stream_options_[0] = cparams_.options;

  if (cparams_.speed_tier == SpeedTier::kFalcon) {

    stream_options_[0].tree_kind = ModularOptions::TreeKind::kWPFixedDC;

  } else if (cparams_.speed_tier == SpeedTier::kThunder) {

    stream_options_[0].tree_kind = ModularOptions::TreeKind::kGradientFixedDC;

  }

  stream_options_[0].histogram_params =

      HistogramParams::ForModular(cparams_, {}, streaming_mode);

  return true;

}

Status ModularFrameEncoder::ComputeEncodingData(

    const FrameHeader& frame_header, const ImageMetadata& metadata,

    Image3F* JXL_RESTRICT color, const std::vector<ImageF>& extra_channels,

    const Rect& group_rect, const FrameDimensions& patch_dim,

    const Rect& frame_area_rect, PassesEncoderState* JXL_RESTRICT enc_state,

    const JxlCmsInterface& cms, ThreadPool* pool, AuxOut* aux_out,

    bool do_color) {

  JxlMemoryManager* memory_manager = enc_state->memory_manager();

  JXL_DEBUG_V(6, "Computing modular encoding data for frame %s",

              frame_header.DebugString().c_str());

  bool groupwise = enc_state->streaming_mode;

  if (do_color && frame_header.loop_filter.gab && !groupwise) {

    float w = 0.9908511000000001f;

    float weights[3] = {w, w, w};

    JXL_RETURN_IF_ERROR(GaborishInverse(color, Rect(*color), weights, pool));

  }

  if (do_color && metadata.bit_depth.bits_per_sample <= 16 &&

      cparams_.speed_tier < SpeedTier::kCheetah &&

      cparams_.decoding_speed_tier < 2 && !groupwise) {

    JXL_RETURN_IF_ERROR(FindBestPatchDictionary(

        *color, enc_state, cms, nullptr, aux_out,

        cparams_.color_transform == ColorTransform::kXYB));

    JXL_RETURN_IF_ERROR(PatchDictionaryEncoder::SubtractFrom(

        enc_state->shared.image_features.patches, color));

  }

  if (cparams_.custom_splines.HasAny()) {

    PassesSharedState& shared = enc_state->shared;

    ImageFeatures& image_features = shared.image_features;

    image_features.splines = cparams_.custom_splines;

  }

  // Convert ImageBundle to modular Image object

  const size_t xsize = patch_dim.xsize;

  const size_t ysize = patch_dim.ysize;

  int nb_chans = 3;

  if (metadata.color_encoding.IsGray() &&

      cparams_.color_transform == ColorTransform::kNone) {

    nb_chans = 1;

  }

  if (!do_color) nb_chans = 0;

  nb_chans += extra_channels.size();

  bool fp = metadata.bit_depth.floating_point_sample &&

            cparams_.color_transform != ColorTransform::kXYB;

  // bits_per_sample is just metadata for XYB images.

  if (metadata.bit_depth.bits_per_sample >= 32 && do_color &&

      cparams_.color_transform != ColorTransform::kXYB) {

    if (metadata.bit_depth.bits_per_sample == 32 && fp == false) {

      return JXL_FAILURE("uint32_t not supported in enc_modular");

    } else if (metadata.bit_depth.bits_per_sample > 32) {

      return JXL_FAILURE("bits_per_sample > 32 not supported");

    }

  }

  // in the non-float case, there is an implicit 0 sign bit

  int max_bitdepth =

      do_color ? metadata.bit_depth.bits_per_sample + (fp ? 0 : 1) : 0;

  Image& gi = stream_images_[0];

  JXL_ASSIGN_OR_RETURN(

      gi, Image::Create(memory_manager, xsize, ysize,

                        metadata.bit_depth.bits_per_sample, nb_chans));

  int c = 0;

  if (cparams_.color_transform == ColorTransform::kXYB &&

      cparams_.modular_mode == true) {

    float enc_factors[3] = {65536.0f, 4096.0f, 4096.0f};

    if (cparams_.butteraugli_distance > 0 && !cparams_.responsive) {

      // quantize XYB here and then treat it as a lossless image

      enc_factors[0] *= 1.f / (1.f + 23.f * cparams_.butteraugli_distance);

      enc_factors[1] *= 1.f / (1.f + 14.f * cparams_.butteraugli_distance);

      enc_factors[2] *= 1.f / (1.f + 14.f * cparams_.butteraugli_distance);

      cparams_.butteraugli_distance = 0;

    }

    if (cparams_.manual_xyb_factors.size() == 3) {

      JXL_RETURN_IF_ERROR(DequantMatricesSetCustomDC(

          memory_manager, &enc_state->shared.matrices,

          cparams_.manual_xyb_factors.data()));

      // TODO(jon): update max_bitdepth in this case

    } else {

      JXL_RETURN_IF_ERROR(DequantMatricesSetCustomDC(

          memory_manager, &enc_state->shared.matrices, enc_factors));

      max_bitdepth = 12;

    }

  }

  pixel_type maxval = gi.bitdepth < 32 ? (1u << gi.bitdepth) - 1 : 0;

  if (do_color) {

    for (; c < 3; c++) {

      if (metadata.color_encoding.IsGray() &&

          cparams_.color_transform == ColorTransform::kNone &&

          c != (cparams_.color_transform == ColorTransform::kXYB ? 1 : 0))

        continue;

      int c_out = c;

      // XYB is encoded as YX(B-Y)

      if (cparams_.color_transform == ColorTransform::kXYB && c < 2)

        c_out = 1 - c_out;

      double factor = maxval;

      if (cparams_.color_transform == ColorTransform::kXYB)

        factor = enc_state->shared.matrices.InvDCQuant(c);

      if (c == 2 && cparams_.color_transform == ColorTransform::kXYB) {

        JXL_ENSURE(!fp);

        for (size_t y = 0; y < ysize; ++y) {

          const float* const JXL_RESTRICT row_in = color->PlaneRow(c, y);

          pixel_type* const JXL_RESTRICT row_out = gi.channel[c_out].Row(y);

          pixel_type* const JXL_RESTRICT row_Y = gi.channel[0].Row(y);

          for (size_t x = 0; x < xsize; ++x) {

            // TODO(eustas): check if std::roundf is appropriate

            row_out[x] = row_in[x] * factor + 0.5f;

            row_out[x] -= row_Y[x];

          }

        }

      } else {

        int bits = metadata.bit_depth.bits_per_sample;

        int exp_bits = metadata.bit_depth.exponent_bits_per_sample;

        gi.channel[c_out].hshift = frame_header.chroma_subsampling.HShift(c);

        gi.channel[c_out].vshift = frame_header.chroma_subsampling.VShift(c);

        size_t xsize_shifted = DivCeil(xsize, 1 << gi.channel[c_out].hshift);

        size_t ysize_shifted = DivCeil(ysize, 1 << gi.channel[c_out].vshift);

        JXL_RETURN_IF_ERROR(

            gi.channel[c_out].shrink(xsize_shifted, ysize_shifted));

        const auto process_row = [&](const int task,

                                     const int thread) -> Status {

          const size_t y = task;

          const float* const JXL_RESTRICT row_in =

              color->PlaneRow(c, y + group_rect.y0()) + group_rect.x0();

          pixel_type* const JXL_RESTRICT row_out = gi.channel[c_out].Row(y);

          JXL_RETURN_IF_ERROR(float_to_int(row_in, row_out, xsize_shifted, bits,

                                           exp_bits, fp, factor));

          return true;

        };

        JXL_RETURN_IF_ERROR(RunOnPool(pool, 0, ysize_shifted,

                                      ThreadPool::NoInit, process_row,

                                      "float2int"));

      }

    }

    if (metadata.color_encoding.IsGray() &&

        cparams_.color_transform == ColorTransform::kNone)

      c = 1;

  }

  for (size_t ec = 0; ec < extra_channels.size(); ec++, c++) {

    const ExtraChannelInfo& eci = metadata.extra_channel_info[ec];

    size_t ecups = frame_header.extra_channel_upsampling[ec];

    JXL_RETURN_IF_ERROR(

        gi.channel[c].shrink(DivCeil(patch_dim.xsize_upsampled, ecups),

                             DivCeil(patch_dim.ysize_upsampled, ecups)));

    gi.channel[c].hshift = gi.channel[c].vshift =

        CeilLog2Nonzero(ecups) - CeilLog2Nonzero(frame_header.upsampling);

    int bits = eci.bit_depth.bits_per_sample;

    int exp_bits = eci.bit_depth.exponent_bits_per_sample;

    bool ec_fp = eci.bit_depth.floating_point_sample;

    double factor = (ec_fp ? 1 : ((1u << eci.bit_depth.bits_per_sample) - 1));

    if (bits + (ec_fp ? 0 : 1) > max_bitdepth) {

      max_bitdepth = bits + (ec_fp ? 0 : 1);

    }

    const auto process_row = [&](const int task, const int thread) -> Status {

      const size_t y = task;

      const float* const JXL_RESTRICT row_in =

          extra_channels[ec].Row(y + group_rect.y0()) + group_rect.x0();

      pixel_type* const JXL_RESTRICT row_out = gi.channel[c].Row(y);

      JXL_RETURN_IF_ERROR(float_to_int(row_in, row_out,

                                       gi.channel[c].plane.xsize(), bits,

                                       exp_bits, ec_fp, factor));

      return true;

    };

    JXL_RETURN_IF_ERROR(RunOnPool(pool, 0, gi.channel[c].plane.ysize(),

                                  ThreadPool::NoInit, process_row,

                                  "float2int"));

  }

  JXL_ENSURE(c == nb_chans);

  int level_max_bitdepth = (cparams_.level == 5 ? 16 : 32);

  if (max_bitdepth > level_max_bitdepth) {

    return JXL_FAILURE(

        "Bitdepth too high for level %i (need %i bits, have only %i in this "

        "level)",

        cparams_.level, max_bitdepth, level_max_bitdepth);

  }

  // Set options and apply transformations

  if (!cparams_.ModularPartIsLossless()) {

    if (cparams_.palette_colors != 0) {

      JXL_DEBUG_V(3, "Lossy encode, not doing palette transforms");

    }

    if (cparams_.color_transform == ColorTransform::kXYB) {

      cparams_.channel_colors_pre_transform_percent = 0;

    }

    cparams_.channel_colors_percent = 0;

    cparams_.palette_colors = 0;

    cparams_.lossy_palette = false;

  }

  // Global palette transforms

  float channel_colors_percent = 0;

  if (!cparams_.lossy_palette &&

      (cparams_.speed_tier <= SpeedTier::kThunder ||

       (do_color && metadata.bit_depth.bits_per_sample > 8))) {

    channel_colors_percent = cparams_.channel_colors_pre_transform_percent;

  }

  if (!groupwise) {

    JXL_RETURN_IF_ERROR(try_palettes(gi, max_bitdepth, maxval, cparams_,

                                     channel_colors_percent, pool));

  }

  // don't do an RCT if we're short on bits

  if (cparams_.color_transform == ColorTransform::kNone && do_color &&

      gi.channel.size() - gi.nb_meta_channels >= 3 &&

      max_bitdepth + 1 < level_max_bitdepth) {

    if (cparams_.colorspace < 0 && (!cparams_.ModularPartIsLossless() ||

                                    cparams_.speed_tier > SpeedTier::kHare)) {

      Transform ycocg{TransformId::kRCT};

      ycocg.rct_type = 6;

      ycocg.begin_c = gi.nb_meta_channels;

      do_transform(gi, ycocg, weighted::Header(), pool);

      max_bitdepth++;

    } else if (cparams_.colorspace > 0) {

      Transform sg(TransformId::kRCT);

      sg.begin_c = gi.nb_meta_channels;

      sg.rct_type = cparams_.colorspace;

      do_transform(gi, sg, weighted::Header(), pool);

      max_bitdepth++;

    }

  }

  if (cparams_.move_to_front_from_channel > 0) {

    for (size_t tgt = 0;

         tgt + cparams_.move_to_front_from_channel < gi.channel.size(); tgt++) {

      size_t pos = cparams_.move_to_front_from_channel;

      while (pos > 0) {

        Transform move(TransformId::kRCT);

        if (pos == 1) {

          move.begin_c = tgt;

          move.rct_type = 28;  // RGB -> GRB

          pos -= 1;

        } else {

          move.begin_c = tgt + pos - 2;

          move.rct_type = 14;  // RGB -> BRG

          pos -= 2;

        }

        do_transform(gi, move, weighted::Header(), pool);

      }

    }

  }

  // don't do squeeze if we don't have some spare bits

  if (!groupwise && cparams_.responsive && !gi.channel.empty() &&

      max_bitdepth + 2 < level_max_bitdepth) {

    Transform t(TransformId::kSqueeze);

    // Check if default squeeze parameters are ok.

    std::vector<SqueezeParams> params;

    DefaultSqueezeParameters(&params, gi);

    // If image is smaller than group_dim, then default squeeze parameters

    // are not going too far. Else, channel size don't turn zero. Thus we only

    // check if tile does not go to zero-dim.

    size_t shift_cap = 7 + frame_header.group_size_shift;

    size_t hshift = 0;

    size_t vshift = 0;

    for (size_t i = 0; i < params.size(); ++i) {

      if (params[i].horizontal) {

        hshift++;

      } else {

        vshift++;

      }

      size_t dc_boost = (std::min(hshift, vshift) >= 3) ? 3 : 0;

      // In case we squeeze too much, truncate squeeze script.

      if (std::max(hshift, vshift) > shift_cap + dc_boost) {

        params.resize(i - 1);

        t.squeezes = params;

        break;

      }

    }

    do_transform(gi, t, weighted::Header(), pool);

    max_bitdepth += 2;

  }

  if (max_bitdepth + 1 > level_max_bitdepth) {

    // force no group RCTs if we don't have a spare bit

    cparams_.colorspace = 0;

  }

  JXL_ENSURE(max_bitdepth <= level_max_bitdepth);

  if (!cparams_.ModularPartIsLossless()) {

    quants_.resize(gi.channel.size(), 1);

    float quantizer = 0.25f;

    if (!cparams_.responsive) {

      JXL_DEBUG_V(1,

                  "Warning: lossy compression without Squeeze "

                  "transform is just color quantization.");

      quantizer *= 0.1f;

    }

    float bitdepth_correction = 1.f;

    if (cparams_.color_transform != ColorTransform::kXYB) {

      bitdepth_correction = maxval / 255.f;

    }

    std::vector<float> quantizers;

    for (size_t i = 0; i < 3; i++) {

      float dist = cparams_.butteraugli_distance;

      quantizers.push_back(quantizer * powf(dist, 1.2) * bitdepth_correction);

    }

    for (size_t i = 0; i < extra_channels.size(); i++) {

      int ec_bitdepth =

          metadata.extra_channel_info[i].bit_depth.bits_per_sample;

      pixel_type ec_maxval = ec_bitdepth < 32 ? (1u << ec_bitdepth) - 1 : 0;

      bitdepth_correction = ec_maxval / 255.f;

      float dist = 0;

      if (i < cparams_.ec_distance.size()) dist = cparams_.ec_distance[i];

      if (dist < 0) dist = cparams_.butteraugli_distance;

      quantizers.push_back(quantizer * dist * bitdepth_correction);

    }