#include "tensorflow/c/eager/abstract_operation.h"

namespace tensorflow {

// Abstract interface to a context.
//
// This serves as a factory for creating `AbstractOperation`s and for
// registering traced functions.
// Operations creation within a context can only be executed in that context
// (for now at least).
// Implementations of the context may contain some state e.g. an execution

} TP_OptimizerConfigs;

#define TP_OPTIMIZER_CONFIGS_STRUCT_SIZE \
  TF_OFFSET_OF_END(TP_OptimizerConfigs, scoped_allocator_optimization)

// Struct for Optimizer. Plugin authors must provide an optimize function.
// Creation and deletion functions are optional.
typedef struct TP_Optimizer {
  size_t struct_size;
  void* ext;  // reserved for future use

  // [Optional]
  // Create function for optimizer.
  void* (*create_func)();

  // followed by `TFE_ContextAsyncWait` to clear pending operations. It will be
  // used to cancel all other operations if any fails.
  //
  // Set step_id to configure the step id used for rendezvous creation. step id
  // of value -1 is reserved for global rendezvous and should not be set here.
  void StartExecute(TFE_Context* context,
                    const std::vector<ParallelTensor*>& inputs,

#ifdef __cplusplus
extern "C" {
#endif

// Resets `op_to_reset` with `op_or_function_name` and `raw_device_name`. This
// is for performance optimization by reusing an exiting unused op rather than
// creating a new op every time. If `raw_device_name` is `NULL` or empty, it
// does not set the device name. If it's not `NULL`, then it attempts to parse
// and set the device name. It's effectively `TFE_OpSetDevice`, but it is faster

//              The distinction between an empty array of operations and no
//              array of operations is necessary to distinguish the case of
//              creating a function with no body (e.g. identity or permutation)
//              and the case of creating a function whose body contains all
//              the nodes in the graph (except for the automatic skipping, see
//              below).

// The number of dimensions used to represent the tensor on device can be
// different from the number returned by TFE_TensorHandleNumDims.
// The return value was current at the time of TFE_TensorDebugInfo creation.
TF_CAPI_EXPORT extern int TFE_TensorDebugInfoOnDeviceNumDims(
    TFE_TensorDebugInfo* debug_info);

// Returns the number of elements in dimension `dim_index`.

// restricted stand-in for tf.Variable in Python.
class Variable {
 public:
  // Construct a Variable from a resource-dtype TFE_TensorHandle and an
  // indication of the dtype of the variable's value.
  //
  // Note that creating this resource-dtype handle can fail, so `Create` is a
  // separate static method which returns a status.
  Variable(TFE_TensorHandle* handle, TF_DataType type)
      : handle_(handle), type_(type) {}

  void* c_optimizer_;
};

// Registration function to register a CGraphOptimizer along with plugin configs
// and device type.
void CGraphOptimizerRegister(
    const PluginGraphOptimizerRegistry::Creator& creator,
    const TP_OptimizerConfigs tp_configs, const char* device_type);

}  // namespace grappler
}  // namespace tensorflow

};
// LINT.ThenChange(//tensorflow/c/eager/c_api.h)

// Abstract interface to a context.
//
// A context is responsible for creating key objects such as Tensors,
// TensorHandles & Operations.
class ImmediateExecutionContext : public AbstractContext {
 public:
  // Optimized scalar creation functions
  virtual AbstractTensorInterface* CreateInt64Scalar(int64_t value) = 0;

// allow us to trace the data dependencies between operations and hence compute
// gradients.
//
// `ZerosLike` is not expected to be called and returns a nullptr. The creation
// of default zeros grads is handled by the `DefaultGradientFunction` registered
// for each op.
// TODO(srbs): We need to define `ZerosLike` here to keep the compiler happy.
// Figure out a way to avoid this.