Add bitrate accuracy estimation.

The function av1_q_mode_estimate_base_q accepts a GF_GROUP and a
bit budget, and outputs a base q index to use for encoding.

Edited the interface of av1_q_mode_compute_gop_q_indices to
accept an arf_q.

BUG=aomedia:3045

Change-Id: I072cf4e84d5d353d42dd73e3cc875010e72da208
diff --git a/av1/encoder/ratectrl.c b/av1/encoder/ratectrl.c
index 2c6d2a0..2218927 100644
--- a/av1/encoder/ratectrl.c
+++ b/av1/encoder/ratectrl.c
@@ -3135,12 +3135,7 @@
  * Intended to be used only with AOM_Q mode.
  */
 void av1_q_mode_compute_gop_q_indices(int gf_frame_index, int base_q_index,
-                                      int gfu_boost, int bit_depth,
-                                      double arf_boost_factor,
-                                      struct GF_GROUP *gf_group) {
-  int arf_q =
-      av1_get_arf_q_index(base_q_index, gfu_boost, bit_depth, arf_boost_factor);
-
+                                      int arf_q, struct GF_GROUP *gf_group) {
   for (int gf_index = gf_frame_index; gf_index < gf_group->size; ++gf_index) {
     int height = gf_group_pyramid_level(gf_group, gf_index);
     gf_group->q_val[gf_index] = av1_q_mode_get_q_index(
diff --git a/av1/encoder/ratectrl.h b/av1/encoder/ratectrl.h
index c88d5c0..a25d7f5 100644
--- a/av1/encoder/ratectrl.h
+++ b/av1/encoder/ratectrl.h
@@ -707,15 +707,12 @@
  *
  * \param[in]       gf_frame_index    Index of the current frame
  * \param[in]       base_q_index      Base q index
- * \param[in]       gfu_boost         GFU boost
  * \param[in]       bit_depth         Bit depth
- * \param[in]       arf_boost_factor  ARF boost factor
+ * \param[in]       arf_q             ARF q index
  * \param[out]      gf_group          Pointer to the GOP
  */
 void av1_q_mode_compute_gop_q_indices(int gf_frame_index, int base_q_index,
-                                      int gfu_boost, int bit_depth,
-                                      double arf_boost_factor,
-                                      struct GF_GROUP *gf_group);
+                                      int arf_q, struct GF_GROUP *gf_group);
 
 /*!\brief Compute the q_indices for a single frame.
  *
diff --git a/av1/encoder/tpl_model.c b/av1/encoder/tpl_model.c
index ac5812e..ab7423f 100644
--- a/av1/encoder/tpl_model.c
+++ b/av1/encoder/tpl_model.c
@@ -1958,3 +1958,52 @@
   fclose(fptr);
 }
 #endif  // CONFIG_RD_COMMAND
+
+/*
+ * Estimate the optimal base q index for a GOP.
+ */
+int av1_q_mode_estimate_base_q(GF_GROUP *gf_group,
+                               const TplTxfmStats *txfm_stats_list,
+                               double bit_budget, int gf_frame_index,
+                               int arf_q) {
+  int q_max = 255;  // Maximum q value.
+  int q_min = 0;    // Minimum q value.
+  int q = (q_max + q_min) / 2;
+
+  av1_q_mode_compute_gop_q_indices(gf_frame_index, q_max, arf_q, gf_group);
+  double q_max_estimate = av1_estimate_gop_bitrate(
+      gf_group->q_val, gf_group->size, txfm_stats_list);
+  av1_q_mode_compute_gop_q_indices(gf_frame_index, q_min, arf_q, gf_group);
+  double q_min_estimate = av1_estimate_gop_bitrate(
+      gf_group->q_val, gf_group->size, txfm_stats_list);
+
+  while (true) {
+    av1_q_mode_compute_gop_q_indices(gf_frame_index, q, arf_q, gf_group);
+
+    double estimate = av1_estimate_gop_bitrate(gf_group->q_val, gf_group->size,
+                                               txfm_stats_list);
+
+    // We want to find the lowest q that satisfies the bit budget constraint.
+    // A binary search narrows the result down to two values: q_min and q_max.
+    if (q_max <= q_min + 1 || estimate == bit_budget) {
+      // Pick the estimate that lands closest to the budget.
+      if (fabs(q_max_estimate - bit_budget) <
+          fabs(q_min_estimate - bit_budget)) {
+        q = q_max;
+      } else {
+        q = q_min;
+      }
+      break;
+    } else if (estimate > bit_budget) {
+      q_min = q;
+      q_min_estimate = estimate;
+      q = (q_max + q_min) / 2;
+    } else if (estimate < bit_budget) {
+      q_max = q;
+      q_max_estimate = estimate;
+      q = (q_max + q_min) / 2;
+    }
+  }
+
+  return q;
+}
diff --git a/av1/encoder/tpl_model.h b/av1/encoder/tpl_model.h
index c2fdb40..197bf4e 100644
--- a/av1/encoder/tpl_model.h
+++ b/av1/encoder/tpl_model.h
@@ -23,6 +23,7 @@
 struct AV1_SEQ_CODING_TOOLS;
 struct EncodeFrameParams;
 struct EncodeFrameInput;
+struct GF_GROUP;
 
 #include "config/aom_config.h"
 
@@ -441,6 +442,26 @@
 int av1_get_overlap_area(int row_a, int col_a, int row_b, int col_b, int width,
                          int height);
 
+/*!\brief Estimate the optimal base q index for a GOP.
+ *
+ * This function picks q based on a chosen bit rate. It
+ * estimates the bit rate using the starting base q, then uses
+ * a binary search to find q to achieve the specified bit rate.
+ *
+ * \param[in]       gf_group          GOP structure
+ * \param[in]       stats             Transform stats struct
+ * \param[in]       bit_budget        The specified bit budget to achieve
+ * \param[in]       gf_frame_index    current frame in the GOP
+ * \param[in]       bit_depth         Bit depth
+ * \param[in]       arf_q             ARF q index
+ *
+ * \return Returns the optimal base q index to use.
+ */
+int av1_q_mode_estimate_base_q(struct GF_GROUP *gf_group,
+                               const TplTxfmStats *txfm_stats_list,
+                               double bit_budget, int gf_frame_index,
+                               int arf_q);
+
 /*!\endcond */
 #ifdef __cplusplus
 }  // extern "C"
diff --git a/test/ratectrl_test.cc b/test/ratectrl_test.cc
index 7dfc947..735deee 100644
--- a/test/ratectrl_test.cc
+++ b/test/ratectrl_test.cc
@@ -41,9 +41,7 @@
 
 TEST(RatectrlTest, QModeComputeGOPQIndicesTest) {
   int base_q_index = 36;
-  int gfu_boost = 500;
-  int bit_depth = AOM_BITS_8;
-  double arf_boost_factor = 0.20;
+  int arf_q = 36;
 
   int gf_frame_index = 0;
   GF_GROUP gf_group = {};
@@ -57,11 +55,15 @@
     gf_group.update_type[i] = update_type[i];
   }
 
-  av1_q_mode_compute_gop_q_indices(gf_frame_index, base_q_index, gfu_boost,
-                                   bit_depth, arf_boost_factor, &gf_group);
+  av1_q_mode_compute_gop_q_indices(gf_frame_index, base_q_index, arf_q,
+                                   &gf_group);
 
   for (int i = 0; i < gf_group.size; i++) {
-    EXPECT_LE(gf_group.q_val[i], base_q_index);
+    if (arf_q == base_q_index) {
+      EXPECT_EQ(gf_group.q_val[i], base_q_index);
+    } else {
+      EXPECT_LE(gf_group.q_val[i], base_q_index);
+    }
   }
 }