using ::testing::Truly;

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

  const QuantizationConfig new_config = PopulateDefaults(config);
  EXPECT_TRUE(new_config.pipeline_config().unpack_quantized_types());
}

TEST(PopulateDefaultsTest, PopulateDefaultsForConfigWithUnpackQuantizedTypes) {
  QuantizationConfig config{};
  config.mutable_pipeline_config()->set_unpack_quantized_types(false);

// `::stablehlo::quantization::PopulateDefaults`.
QuantizationConfig PywrapPopulateDefaults(
    const QuantizationConfig& user_provided_config);

// Function used by the pywrap_quantization module to mirror
// `::stablehlo::quantization::ExpandPresets`.
QuantizationConfig PywrapExpandPresets(const QuantizationConfig& config);

}  // namespace stablehlo::quantization::pywrap

// Returns a copy of `user_provided_config` with default values populated where
// the user did not explicitly specify.
QuantizationConfig PopulateDefaults(
    const QuantizationConfig& user_provided_config);

// Returns a copy of `QuantizationConfig` where presets are expanded and
// transformed into other fields in `QuantizationConfig`.
//
// The expansion rules are as follows:
// * StaticRangePtqPreset

  config.mutable_specs()->Swap(&new_specs);
}

}  // namespace

QuantizationConfig ExpandPresets(const QuantizationConfig& config) {
  QuantizationConfig new_config = config;

  // Update the `new_config` with each preset's expansions.
  switch (config.preset_case()) {
    case QuantizationConfig::kStaticRangePtqPreset:
      ExpandStaticRangePtqPreset(config.static_range_ptq_preset(), new_config);

    single representative dataset.
  """
  # Updates user-provided `QuantizationConfig`s for the internal quantization
  # pipeline to work with.
  print('=== User-provided QuantizationConfig ===')
  print(config)
  config = qc.QuantizationConfig.FromString(
      pywrap_quantization.populate_default_configs(config.SerializeToString())
  )
  config = qc.QuantizationConfig.FromString(

  // LINT.IfChange(expand_preset_configs)
  m.def("expand_preset_configs", &PywrapExpandPresets, R"pbdoc(
        Expands presets to other fields in `QuantizationConfig`.

        Each preset is expressible by other fields in `QuantizationConfig`.
        Returns a copy of `QuantizationConfig` (serialized) where the fields are
        expanded from presets. If no preset has been set, it is a no-op and
        returns the same copy of the input.

def populate_default_configs(
    user_provided_quantization_config_serialized: bytes,
) -> bytes: ...  # QuantizationConfig

# LINT.ThenChange()

# LINT.IfChange(expand_preset_configs)
def expand_preset_configs(
    quantization_config_serialized: bytes,
) -> bytes: ...  # QuantizationConfig

                               py_function_library);
}

QuantizationConfig PywrapPopulateDefaults(
    const QuantizationConfig& user_provided_config) {
  return PopulateDefaults(user_provided_config);
}

QuantizationConfig PywrapExpandPresets(const QuantizationConfig& config) {
  return ExpandPresets(config);
}

#include "tsl/platform/errors.h"
#include "tsl/platform/statusor.h"

namespace mlir::quant::stablehlo {

using ::stablehlo::quantization::GetReportFilePath;
using ::stablehlo::quantization::QuantizationConfig;
using ::tensorflow::SignatureDef;
using ::tensorflow::quantization::ExportedModel;
using ::tensorflow::quantization::PyFunctionLibrary;
using ::tensorflow::quantization::RunPasses;

    path_map = {'serving_default': dataset_path}
    repr_dataset.TfRecordRepresentativeDatasetSaver(path_map).save(
        {'serving_default': data_gen()}
    )

    config = qc.QuantizationConfig(
        static_range_ptq_preset=qc.StaticRangePtqPreset(
            representative_datasets=[
                qc.RepresentativeDatasetConfig(
                    tf_record=qc.TfRecordFile(path=dataset_path)