//

// Created by serizba on 27/6/20.

//

#ifndef CPPFLOW2_TENSOR_H

#define CPPFLOW2_TENSOR_H

#include <memory>

#include <vector>

#include <cstring>

#include <string>

#include <tensorflow/c/tf_tensor.h>

#include <tensorflow/c/eager/c_api.h>

#include "context.h"

#include "datatype.h"

namespace cppflow {

/**

* @class tensor

* @brief A TensorFlow eager tensor wrapper

*

*/

class tensor {

public:

tensor()= default;

/**

* Creates a tensor with the given values and specified shape

* @tparam T A type that can be convertible into a tensor

* @param values The values to be converted (in a flattened version)

* @param shape The shape of the converted tensor

*/

template<typename T>

tensor(const std::vector<T>& values, const std::vector<int64_t>& shape);

/**

* Creates a flat tensor with the given values

* @tparam T A type that can be convertible into a tensor

* @param values The values to be converted

*/

template<typename T>

tensor(const std::initializer_list<T>& values);

/**

* Creates a tensor with the given value

* @tparam T A type that can be convertible into a tensor

* @param value The value to be converted

*/

template<typename T>

tensor(const T& value);

/**

* @return Shape of the tensor

*/

tensor shape() const;

/**

* @param on_memory If false, the function will return the name of the device that produced the tensor.

* If true, the function will return the name of the device in whose memory the tensor resides

* @return Returns the name of the device of the tensor

*/

std::string device(bool on_memory=false) const;

/**

* @return The tensor datatype

*/

datatype dtype() const;

/**

* Converts the tensor into a C++ vector

* @tparam T The c++ type (must be equivalent to the tensor type)

* @return A vector representing the flat tensor

*/

template<typename T>

std::vector<T> get_data() const;

~tensor() = default;

tensor(const tensor &tensor) = default;

tensor(tensor &&tensor) = default;

tensor &operator=(const tensor &other) = default;

tensor &operator=(tensor &&other) = default;

explicit tensor(TFE_TensorHandle* handle);

explicit tensor(TF_Tensor* t);

// NOTE: Usually, one should not call get_eager_handle() or get_tensor() below.

// They are designed for implementation details in cppflow.

// If you are calling them directly, it is likely that you are using some

// tenforflow APIs not supported in cppflow.

// Additional NOTE: TF_Tensor is an immutable tensor inside tensorflow.

// TFE_TensorHandle is a TF_Tensor and the associated device, plus some data cache

// TODO: Need to determine if we can mark the return value or *this as const

std::shared_ptr<TFE_TensorHandle> get_eager_handle() const { return tfe_handle; }

// Get the TF_Tensor data from the eager handle

// Call `get_data<T>()` instead if possible

// NOTE: Changes to the returned TF_Tensor may not be reflected in the actual device memory!

// Do *NOT* modify the returned TF_Tensor!

// See comments of `tf_tensor` for more details.

std::shared_ptr<TF_Tensor> get_tensor() const;

// DO NOT directly access this member, call get_eager_handle() instead

// TODO: This is kept as public to be compatible with existing code and should be mark as private

std::shared_ptr<TFE_TensorHandle> tfe_handle;

private:

// This member serves as a local cache of the data in tfe_handle.

// It refers to `local_mirrors_` if on device, or `data_` if on host CPU.

// Changes to this variable may not be reflected in the actual device memory,

// e.g. on GPUs or on remote nodes.

// Access it via get_tensor() if not in constructor

mutable std::shared_ptr<TF_Tensor> tf_tensor;

tensor(enum TF_DataType type, const void* data, size_t len, const std::vector<int64_t>& shape);

};

}

/******************************

* IMPLEMENTATION DETAILS *

******************************/

namespace cppflow {

inline tensor::tensor(enum TF_DataType type, const void *data, size_t len, const std::vector<int64_t> &shape) {

this->tf_tensor = {TF_AllocateTensor(type, shape.data(), shape.size(), len), TF_DeleteTensor};

memcpy(TF_TensorData(this->tf_tensor.get()), data, TF_TensorByteSize(this->tf_tensor.get()));

this->tfe_handle = {TFE_NewTensorHandle(this->tf_tensor.get(), context::get_status()), TFE_DeleteTensorHandle};

status_check(context::get_status());

}

template<typename T>

tensor::tensor(const std::vector<T>& values, const std::vector<int64_t>& shape) :

tensor(deduce_tf_type<T>(), values.data(), values.size() * sizeof(T), shape) {}

template<typename T>

tensor::tensor(const std::initializer_list<T>& values) :

tensor(std::vector<T>(values), {(int64_t) values.size()}) {}

template<typename T>

tensor::tensor(const T& value) :

tensor(std::vector<T>({value}), {}) {}

#ifdef TENSORFLOW_C_TF_TSTRING_H_

// For future version TensorFlow 2.4

template<>

inline tensor::tensor(const std::string& value) {

TF_TString tstr[1];

TF_TString_Init(&tstr[0]);

TF_TString_Copy(&tstr[0], value.c_str(), value.size());

*this = tensor(static_cast<enum TF_DataType>(TF_STRING), (void *) tstr, sizeof(tstr), {});

}

#else

template<>

inline tensor::tensor(const std::string& value) {

size_t size = 8 + TF_StringEncodedSize(value.length());

char* data = new char[value.size() + 8];

for (int i=0; i<8; i++) {data[i]=0;}

TF_StringEncode(value.c_str(), value.size(), data + 8, size - 8, context::get_status());

status_check(context::get_status());

*this = tensor(static_cast<enum TF_DataType>(TF_STRING), (void *) data, size, {});

delete [] data;

}

#endif // TENSORFLOW_C_TF_TSTRING_H_

inline tensor::tensor(TFE_TensorHandle* handle) {

this->tfe_handle = {handle, TFE_DeleteTensorHandle};

}

inline tensor::tensor(TF_Tensor* t) {

this->tf_tensor = {t, TF_DeleteTensor};

this->tfe_handle = {TFE_NewTensorHandle(this->tf_tensor.get(), context::get_status()), TFE_DeleteTensorHandle};

status_check(context::get_status());

}

inline tensor tensor::shape() const {

auto op = TFE_NewOp(context::get_context(), "Shape", context::get_status());

status_check(context::get_status());

TFE_OpAddInput(op, this->tfe_handle.get(), context::get_status());

status_check(context::get_status());

// Output type should be int64_t

TFE_OpSetAttrType(op, "out_type", cppflow::datatype::TF_INT64);

// EXECUTE

int n = 1;

TFE_TensorHandle* res[1] = { nullptr };

TFE_Execute(op, res, &n, context::get_status());

status_check(context::get_status());

TFE_DeleteOp(op);

return tensor(res[0]);

}

inline std::string tensor::device(bool on_memory) const {

std::string res;

if (on_memory)

res = TFE_TensorHandleBackingDeviceName(this->tfe_handle.get(), context::get_status());

else

res = std::string(TFE_TensorHandleDeviceName(this->tfe_handle.get(), context::get_status()));

status_check(context::get_status());

return res;

}

template<typename T>

std::vector<T> tensor::get_data() const {

// Check if asked datatype and tensor datatype match

if (this->dtype() != deduce_tf_type<T>()) {

auto type1 = cppflow::to_string(deduce_tf_type<T>());

auto type2 = cppflow::to_string(this->dtype());

auto error = "Datatype in function get_data (" + type1 + ") does not match tensor datatype (" + type2 + ")";

throw std::runtime_error(error);

}

auto res_tensor = get_tensor();

// Check tensor data is not empty

auto raw_data = TF_TensorData(res_tensor.get());

//this->error_check(raw_data != nullptr, "Tensor data is empty");

size_t size = TF_TensorByteSize(res_tensor.get()) / TF_DataTypeSize(TF_TensorType(res_tensor.get()));

// Convert to correct type

const auto T_data = static_cast<T*>(raw_data);

std::vector<T> r(T_data, T_data + size);

return r;

}

inline datatype tensor::dtype() const {

return TFE_TensorHandleDataType(this->tfe_handle.get());

}

// NOTE: Changes to the returned TF_Tensor are not reflected in the actual device memory!

inline std::shared_ptr<TF_Tensor> tensor::get_tensor() const {

if(!tf_tensor) {

tf_tensor = {TFE_TensorHandleResolve(tfe_handle.get(), context::get_status()), TF_DeleteTensor};

status_check(context::get_status());

}

return tf_tensor;

}

}

#endif //CPPFLOW2_TENSOR_H