using ::tsl::monitoring::testing::Histogram;

static constexpr char kCompilationTimeStreamzName[] =
    "/tensorflow/core/tf2xla/api/v1/phase2_compilation_time";

static constexpr char kCompilationStatusStreamzName[] =
    "/tensorflow/core/tf2xla/api/v1/phase2_compilation_status";

static constexpr char kPlatformName[] = "Host";
constexpr char kEntryFuncName[] = "main";

static constexpr char kMlirModuleStr[] = R"(

using ::testing::UnorderedElementsAre;

using IsOpQuantizableStableHloTest = ::mlir::quant::QuantizationTestBase;

// Quantizable ops: constants
// Non-quantizable ops: normal StableHLO ops and terminators
constexpr absl::string_view kModuleConstantAdd = R"mlir(
  module {
    func.func @constant_add() -> (tensor<3x2xf32>) {
      %cst1 = stablehlo.constant dense<2.4> : tensor<3x2xf32>

#include "mlir/Interfaces/SideEffectInterfaces.h"  // from @llvm-project

namespace mlir {
namespace TFR {

constexpr char kAttrArgumentNameAttr[] = "tfr.name";
constexpr char kAttrArgumentDefaultAttr[] = "tfr.default";
constexpr char kAttrArgumentTypeAttr[] = "tfr.type";

class TFRDialect : public Dialect {
 public:
  explicit TFRDialect(MLIRContext *context);

// to represent calling a TF host callback function.
constexpr llvm::StringRef kTfTargetName = "tf.call_tf_function";

// `tf.backend_config` is a DictionaryAttr, JAX2TF sets the value of its
// string attribute `caller_name` to the TF host callback function's name.
constexpr llvm::StringRef kTfBackendConfigAttrName = "tf.backend_config";
constexpr llvm::StringRef kCalledFuncAttrName = "called_func";

}  // namespace

#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "mlir/IR/MLIRContext.h"  // from @llvm-project

namespace tensorflow {
namespace quantization {

inline constexpr absl::string_view kTfQuantConstantUnfreezingStepName =
    "tf_quant_constant_unfreezing";
inline constexpr absl::string_view kTfQuantInsertRestoreOpStepName =
    "tf_quant_insert_restore_op";

absl::Status UnfreezeConstantsAndSaveVariables(absl::string_view checkpoint_dir,

 public:
  using Base::Base;
  static constexpr mlir::StringLiteral name = "mlrt.compiler.future";
};

// The MLIR type represents a C++ mlrt::Promise.
class PromiseType
    : public mlir::Type::TypeBase<PromiseType, mlir::Type, mlir::TypeStorage> {
 public:
  using Base::Base;
  static constexpr mlir::StringLiteral name = "mlrt.compiler.promise";
};

// Presence of this attribute will mark the function as quantization target.
inline constexpr StringRef kFusedFunctionAttr = "tf_quant.composite_function";
// The keyword to detect if this is a `NullAttribute`.
inline constexpr StringRef kNullAttributeValue = "N/A";

// Prefixes attached to lifted functions.
constexpr StringRef kQuantizedFuncPrefix = "quantized_";
constexpr StringRef kCompositeFuncPrefix = "composite_";

using ::tsl::protobuf::io::ArrayInputStream;

// Empty (default) `QuantizationSpecs` proto.
constexpr absl::string_view kSpecsEmpty = R"pb(specs
                                               [])pb";

// Configure `QuantizationSpecs` to disable quantization for all dot_general
// quantizable units.
constexpr absl::string_view kSpecsDisableAllDotGeneral =
    R"pb(specs
         [ {

namespace mlir {
namespace TFL {
namespace tac {

// The copy can be non-consectutive copy. This is just fake data.
constexpr float kNNAPICopyUnitCost = 0.2;

// Default values.
constexpr float kNNAPIDefaultFixedValuedCost = 10000.0;

constexpr char NNAPIHardware::kId[];  // Define kId.

mlir::RewritePatternSet NNAPIHardware::GetTransformations(
    MLIRContext* context) const {

    EXPECT_TRUE(module_op_ref);
    return module_op_ref;
  }

  Env* env_{};
  mlir::MLIRContext ctx_{};
};

TEST_F(SaveVariablesToCheckpointTest, VariableSavedToCheckpoint) {
  constexpr absl::string_view kModuleCode = R"mlir(
    module attributes {tf_saved_model.semantics} {
      "tf_saved_model.session_initializer"() {initializers = [@init_func_restore_op]} : () -> ()