#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;
}