import numpy as np

from tensorflow.compiler.mlir.quantization.stablehlo import quantization_config_pb2 as stablehlo_quant_config_pb2
from tensorflow.compiler.mlir.quantization.tensorflow.calibrator import calibration_algorithm
from tensorflow.compiler.mlir.quantization.tensorflow.calibrator import calibration_statistics_pb2 as calib_stats_pb2
from tensorflow.python.platform import test

_CalibrationMethod = (

#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/min_max_value.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/exported_model.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"

#include "tsl/platform/status.h"
#include "tsl/platform/status_matchers.h"

namespace tensorflow {
namespace {

using ::stablehlo::quantization::CalibrationOptions;
using ::tensorflow::calibrator::CalibrationStatistics;
using ::tensorflow::calibrator::CalibrationStatisticsMap;
using ::testing::Contains;
using ::testing::ElementsAre;
using ::testing::HasSubstr;
using ::testing::Key;
using ::testing::SizeIs;
using ::tsl::testing::StatusIs;

absl::StatusOr<mlir::quant::QuantizedType> GetQuantizedType(
    const tflite::TensorT& tensor, mlir::Builder builder,
    bool is_constant = false, mlir::Type storage_type = {});

// Imports float tensor with calibration value into calibrated quantized type.
absl::StatusOr<mlir::quant::QuantizedType> GetCalibratedQuantizedType(
    const tflite::TensorT& tensor, mlir::Builder builder);

absl::StatusOr<mlir::TensorType> GetTensorType(const tflite::TensorT& tensor,

#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/min_max_value.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/exported_model.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/python/py_function_lib.h"

# limitations under the License.
# ==============================================================================
from tensorflow.compiler.mlir.quantization.tensorflow.calibrator import calibration_statistics_pb2

# LINT.IfChange(clear_calibrator)
def clear_calibrator() -> None: ...

# LINT.ThenChange()

# LINT.IfChange(clear_data_from_calibrator)

syntax = "proto3";

package tensorflow.calibrator;

option cc_enable_arenas = true;

// Calibration algorithm's collecting statistics.
// NEXT_ID: 4
message CalibrationStatistics {
  message MinMaxStatistics {
    // global minimum of all sample datasets.
    float global_min = 1;

    // global maximum of all sample datasets.
    float global_max = 2;
  }

  message AverageMinMaxStatistics {

### Reviewing cheatsheet

Before merging the PR, comments starting with 
- ❌ ❓**must** be fixed
- 🤔 💅 **should** be fixed
- 💭 **may** be fixed

# limitations under the License.
# ==============================================================================
from typing import Any

from tensorflow.compiler.mlir.quantization.tensorflow.calibrator import calibration_statistics_pb2
from tensorflow.compiler.mlir.quantization.tensorflow.python import py_function_lib
from tensorflow.compiler.mlir.quantization.tensorflow.python import representative_dataset as rd

#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/cc/run_passes.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/exported_model.pb.h"