from tensorflow.python.platform import test

_CalibrationMethod = (
    stablehlo_quant_config_pb2.CalibrationOptions.CalibrationMethod
)


class CalibrationAlgorithmTest(test.TestCase, parameterized.TestCase):

  def test_min_max_max(self):
    calib_opts = stablehlo_quant_config_pb2.CalibrationOptions(
        calibration_method=_CalibrationMethod.CALIBRATION_METHOD_MIN_MAX
    )

          [1.0, 2.0, 3.0, 4.0, 5.0], dtypes.float32
      )

      aggregator = custom_aggregator_op_wrapper.custom_aggregator(
          input_tensor,
          id='1',
          calibration_method=_CalibrationMethod.CALIBRATION_METHOD_MIN_MAX,
      )
      aggregator_output = self.evaluate(aggregator)
      self.assertAllEqual(aggregator_output.output, [1.0, 2.0, 3.0, 4.0, 5.0])

// CKECK-SAME: <{calibration_method = 5 : i32, id = "skipping_id", num_bins = 32 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32}> : (tensor<1x3x4x3xf32>) -> (tensor<1x3x4x3xf32>, tensor<f32>, tensor<f32>, tensor<512xi64>)
// CHECK: %[[CUSTOM_AGGREGATOR_1:.*]], %[[MIN_1:.*]], %[[MAX_1:.*]], %[[HISTOGRAM_1:.*]] = "tf.CustomAggregator"

    auto calibration_method =
        static_cast<CalibrationMethod>(calibration_method_value);
    OP_REQUIRES(
        context,
        calibration_method !=
            CalibrationOptions::CALIBRATION_METHOD_UNSPECIFIED,
        absl::AbortedError("The calibration method must be specified."));

    calib_opts_.set_calibration_method(calibration_method);

  %0:4 = "tf.CustomAggregator"(%arg0) {min = -0.1 : f32, max = 0.2 : f32, id = "0", calibration_method = 1 : i32, num_bins = 0 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<8x8x8x8xf32>) -> (tensor<8x8x8x8xf32>, tensor<f32>, tensor<f32>, tensor<*xi64>)

// CKECK-SAME: <{calibration_method = 5 : i32, id = "0", num_bins = 32 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32}> : (tensor<1x3x4x3xf32>) -> (tensor<1x3x4x3xf32>, tensor<f32>, tensor<f32>, tensor<512xi64>)
// CHECK: %[[CUSTOM_AGGREGATOR_1:.*]], %[[MIN_1:.*]], %[[MAX_1:.*]], %[[HISTOGRAM_1:.*]] = "tf.CustomAggregator"

    OP_REQUIRES_OK(context, context->GetAttr("ids", &ids_));
    OP_REQUIRES_OK(context, context->GetAttr("calibration_methods",
                                             &calibration_methods_));
    OP_REQUIRES(
        context, ids_.size() == calibration_methods_.size(),
        absl::AbortedError(
            "The `ids` and `calibration_methods` must have the same size."));

    // Check the number and type of inputs.

    %1 = "tf.Relu6"(%0) : (tensor<*xf32>) -> tensor<*xf32>
    func.return %1 : tensor<*xf32>
  }
}

// CalibrationOptions(calibration_method=CALIBRATION_METHOD_MIN_MAX)
// MIN-MAX-CHECK: func @wrap_composite_func

    %0 = stablehlo.constant dense<1.000000e+03> : tensor<1024x3xf32>
    %1 = stablehlo.constant dense<1.000000e+03> : tensor<1x3xf32>
    %2:4 = "tf.CustomAggregator"(%arg0) {calibration_method = 1 : i32, id = "0", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32} : (tensor<1x1024xf32>) -> (tensor<1x1024xf32>, tensor<f32>, tensor<f32>, tensor<*xi64>)

}

TEST(PopulateDefaultsTest, DefaultCalibrationOptionsPopulated) {
  QuantizationConfig config{};

  const QuantizationConfig new_config = PopulateDefaults(config);
  EXPECT_THAT(new_config.calibration_options().calibration_method(),
              Eq(CalibrationOptions::CALIBRATION_METHOD_MIN_MAX));
}

TEST(PopulateDefaultsTest,
     DefaultCalibrationOptionsPopulatedForUnspecifiedMethod) {
  QuantizationConfig config{};