// Usage: go test -run='^TestCalibrate$' -v -calibrate

package big

import (
	"flag"
	"fmt"
	"testing"
	"time"
)

var calibrate = flag.Bool("calibrate", false, "run calibration test")

const (
	sqrModeMul       = "mul(x, x)"
	sqrModeBasic     = "basicSqr(x)"
	sqrModeKaratsuba = "karatsubaSqr(x)"
)

func TestCalibrate(t *testing.T) {
	if !*calibrate {
		return
	}

#include "tsl/platform/statusor.h"

namespace stablehlo::quantization {
namespace {

using ::stablehlo::quantization::CalibrationOptions;
using ::tensorflow::calibrator::CalibrationStatistics;
using ::tensorflow::calibrator::CalibrationStatisticsMap;
using ::tensorflow::quantization::PyFunctionLibrary;
using CalibrationStatisticsFlatMap =
    absl::flat_hash_map<std::string, CalibrationStatistics>;

}  // namespace

#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/python/py_function_lib.h"

namespace stablehlo::quantization {

// Reads the calibration statistics from the given directory.
absl::StatusOr<absl::flat_hash_map<
    std::string, tensorflow::calibrator::CalibrationStatistics>>
ReadStatistics(absl::string_view calibration_data_dir);

// CHECK-SAME: effective_hidden_scale_intermediate = tensor<*x!quant.calibrated<f32<-5.000000e-01:5.000000e-01>>>
// CHECK-SAME: input_to_cell_intermediate = tensor<*x!quant.calibrated<f32<-4.000000e+00:4.000000e+00>>>
// CHECK-SAME: input_to_forget_intermediate = tensor<*x!quant.calibrated<f32<-1.600000e+01:1.600000e+01>>>
// CHECK-SAME: input_to_input_intermediate = tensor<*x!quant.calibrated<f32<-3.200000e+01:3.200000e+01>>>

#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics_collector_min_max.h"

#include <algorithm>
#include <cstdint>
#include <limits>
#include <optional>

#include "absl/types/span.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics.pb.h"

namespace tensorflow {
namespace calibrator {

void CalibrationStatisticsCollectorMinMax::ClearData() {

  // Runs calibration on `module_op` and returns a calibrated ModuleOp with
  // calibrated statistics embedded.
  absl::StatusOr<ModuleOp> Run(
      ModuleOp module_op,
      const ::stablehlo::quantization::QuantizationConfig& config) override;

 private:
  // Exports `module_op` to SavedModel at `dst_saved_model_path`. This is used
  // to export the pre-calibrated `module_op` to SavedModel so that the

using ::tensorflow::calibrator::CalibrationStatistics;
using ::tensorflow::calibrator::CalibrationStatisticsCollectorAverageMinMax;
using ::tensorflow::calibrator::CalibrationStatisticsCollectorBase;
using ::tensorflow::calibrator::CalibrationStatisticsCollectorHistogram;
using ::tensorflow::calibrator::CalibrationStatisticsCollectorMinMax;
using ::tensorflow::calibrator::CalibrationStatisticsMap;

}  // namespace

#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics_collector_average_min_max.h"

#include <cstdint>
#include <optional>

#include "absl/types/span.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics.pb.h"

namespace tensorflow {
namespace calibrator {

void CalibrationStatisticsCollectorAverageMinMax::ClearData() {

#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics_collector_base.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"

namespace tensorflow {
namespace calibrator {

using ::stablehlo::quantization::CalibrationOptions;

#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics_collector_base.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"

namespace tensorflow {
namespace calibrator {

using ::stablehlo::quantization::CalibrationOptions;