#include "xla/mlir_hlo/mhlo/transforms/passes.h"
#include "tsl/platform/errors.h"

namespace mlir::quant::stablehlo {

using ::stablehlo::quantization::GetReportFilePath;
using ::stablehlo::quantization::PipelineConfig;
using ::stablehlo::quantization::QuantizationConfig;
using ::stablehlo::quantization::QuantizationSpecs;
using ::tensorflow::quantization::RunPasses;

PostCalibrationComponent::PostCalibrationComponent(

#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/testing/passes.h.inc"

namespace {

using ::stablehlo::quantization::ExpandPresets;
using ::stablehlo::quantization::PipelineConfig;
using ::stablehlo::quantization::QuantizationConfig;

class TestPostCalibrationComponentPass
    : public impl::TestPostCalibrationComponentPassBase<
          TestPostCalibrationComponentPass> {
 public:

    OpPassManager& pm,
    const ::stablehlo::quantization::PipelineConfig& pipeline_config,
    const ::stablehlo::quantization::QuantizationSpecs& specs);

// Adds passes for weight-only quantization.
void AddWeightOnlyQuantizationPasses(
    OpPassManager& pm,
    const ::stablehlo::quantization::QuantizationSpecs& quantization_specs,
    const ::stablehlo::quantization::PipelineConfig& pipeline_config,

#include "xla/mlir_hlo/mhlo/transforms/passes.h"

namespace mlir::quant::stablehlo {

using ::stablehlo::quantization::CalibrationOptions;
using ::stablehlo::quantization::DebuggerConfig;
using ::stablehlo::quantization::PipelineConfig;
using ::stablehlo::quantization::QuantizationSpecs;

void AddPreCalibrationPasses(OpPassManager& pm,
                             const CalibrationOptions& calibration_options,

  void AddPasses(
      OpPassManager& pm,
      const ::stablehlo::quantization::QuantizationSpecs& specs,
      const ::stablehlo::quantization::PipelineConfig& pipeline_config) const;

 private:
  absl::Nonnull<MLIRContext*> ctx_;
};

}  // namespace mlir::quant::stablehlo

}

QuantizationConfig PopulateDefaults(
    const QuantizationConfig& user_provided_config) {
  QuantizationConfig config = user_provided_config;

  PopulateDefaultCalibrationOptions(config);

  PipelineConfig& pipeline_config = *config.mutable_pipeline_config();
  if (!pipeline_config.has_unpack_quantized_types()) {
    pipeline_config.set_unpack_quantized_types(true);
  }

  return config;
}

  // Set of tags that uniquely identify the `MetaGraphDef` existing in the
  // input SavedModel.
  repeated string tags = 1;
}

// Configures the graph transformation pipeline for quantization.
message PipelineConfig {
  // When set to True, unpacks ops with uniform quantized types into operations
  // without uniform quantized types (mostly i8 or i32). Useful when the target
  // hardware performs better with integer ops.

                    tf_record=qc.TfRecordFile(path=dataset_path)
                )
            ]
        ),
        tf_saved_model=qc.TfSavedModelConfig(tags=[tag_constants.SERVING]),
        pipeline_config=qc.PipelineConfig(
            merge_fusion_with_dequantize=merge_fusion_with_dequantize
        ),
    )
    quantization.quantize_saved_model(
        self._input_saved_model_path,
        self._output_saved_model_path,