#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/threadpool_options.h"
#include "tensorflow/core/public/session.h"

namespace mlir {
namespace tf_saved_model {

using ::tensorflow::DeviceMgr;
using ::tensorflow::Session;
using ::tensorflow::Status;
using ::tensorflow::Tensor;

// FakeSession is for testing only.
class FakeSession : public tensorflow::Session {

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/convert_tensor.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/session_utils.h"
#include "tensorflow/core/framework/resource_var.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/public/session.h"

namespace mlir {
namespace tf_saved_model {
namespace {

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"

namespace mlir {
namespace quant {
namespace {

using ::mlir::tf_saved_model::GetInitializerFunction;
using ::mlir::tf_saved_model::kTfSavedModelInitializerRestoreType;

constexpr StringRef kTfQuantSaveV2OpName = "tf_quant__save_save_v2";

struct FreezeVariableTestPass
    : public impl::FreezeVariablesTestPassBase<FreezeVariableTestPass> {
  void runOnOperation() override {
    TF::test_util::FakeSession session;
    if (tf_saved_model::FreezeVariables(getOperation(), &session).failed())
      signalPassFailure();
  }
};
}  // namespace

std::unique_ptr<OperationPass<ModuleOp>> CreateFreezeVariableTestPass() {

using ::mlir::TensorType;
using ::mlir::tf_saved_model::AssetOp;
using ::mlir::tf_saved_model::GlobalTensorOp;
using ::mlir::tf_saved_model::kTfSavedModelExportedNamesAttr;
using ::mlir::tf_saved_model::kTfSavedModelIndexPathAttr;
using ::mlir::tf_saved_model::kTfSavedModelInitializerInitType;
using ::mlir::tf_saved_model::kTfSavedModelInitializerRestoreType;
using ::mlir::tf_saved_model::kTfSavedModelInitializerTypeAttr;

# RUN: (! %p/exported_python_args 2>&1) | FileCheck %s

# pylint: disable=missing-docstring,line-too-long,dangerous-default-value
import tensorflow.compat.v2 as tf
from tensorflow.compiler.mlir.tensorflow.tests.tf_saved_model import common


class TestModule(tf.Module):

  @tf.function(input_signature=[tf.TensorSpec([], tf.float32)])
  def some_function(self, x):
    return self.callee(x)

          InitializeVariablesInSessionInitializerPass> {
 public:
  void runOnOperation() final {
    static tensorflow::Session* session = new TF::test_util::FakeSession();
    if (failed(tf_saved_model::InitializeVariablesInSessionInitializer(
            getOperation(), session)))
      signalPassFailure();
  }

 private:
};
}  // anonymous namespace

std::unique_ptr<OperationPass<ModuleOp>>

#include "tensorflow/compiler/mlir/tensorflow/analysis/resource_dataflow.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"

namespace mlir {
namespace TF {

namespace {

#define GEN_PASS_DEF_LOCALIZEVARHANDLESPASS

# RUN: %p/control_flow_upgrade_legacy_v1 | FileCheck %s

# pylint: disable=missing-docstring,line-too-long
import tensorflow.compat.v1 as tf
from tensorflow.compiler.mlir.tensorflow.tests.tf_saved_model import common_v1
from tensorflow.python.ops import control_flow_ops

# Tests V1 control flow is functionalized.

# CHECK-NOT: tf_executor.Switch
# CHECK-NOT: tf_executor.Merge
# CHECK: "tf.If"

    graph = ops.get_default_graph()
    for op in graph.get_operations():
      if op.type == '_Retval':
        expected_node_name = op.name
        if op.get_attr('tf_saved_model.index_path') is not None:
          index_path_name = op.get_attr('tf_saved_model.index_path')[0]
          index_path_name = index_path_name.decode('utf-8').split(':')[0]
          try:
            # Only use the index_path name if it points to a Retval node.