// Copyright (C) 2018-2022 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//

/**
 * @brief This is a header file for the OpenVINO Runtime tensor API
 *
 * @file openvino/runtime/tensor.hpp
 */
#pragma once

#include <type_traits>

#include "openvino/core/coordinate.hpp"
#include "openvino/core/rtti.hpp"
#include "openvino/core/shape.hpp"
#include "openvino/core/type/element_type.hpp"
#include "openvino/runtime/allocator.hpp"

namespace InferenceEngine {
class Blob;
}  // namespace InferenceEngine

namespace ov {

class Core;
class InferRequest;
class RemoteContext;
class VariableState;

/**
 * @brief Tensor API holding host memory
 *
 * It can throw exceptions safely for the application, where it is properly handled.
 */
class OPENVINO_API Tensor {
protected:
    std::shared_ptr<InferenceEngine::Blob> _impl;  //!< Shared pointer to internal tensor representation
    std::shared_ptr<void> _so;                     //!< Reference to dynamically loaded library

    /**
     * @brief Constructs Tensor from the initialized std::shared_ptr
     * @param impl Initialized shared pointer
     * @param so Plugin to use. This is required to ensure that Tensor can work properly even if plugin object is
     * destroyed.
     */
    Tensor(const std::shared_ptr<InferenceEngine::Blob>& impl, const std::shared_ptr<void>& so);

    friend class ov::Core;
    friend class ov::InferRequest;
    friend class ov::RemoteContext;
    friend class ov::VariableState;

public:
    /// @brief Default constructor
    Tensor() = default;

    /// @brief Default copy constructor
    /// @param other other Tensor object
    Tensor(const Tensor& other) = default;

    /// @brief Default copy assignment operator
    /// @param other other Tensor object
    /// @return reference to the current object
    Tensor& operator=(const Tensor& other) = default;

    /// @brief Default move constructor
    /// @param other other Tensor object
    Tensor(Tensor&& other) = default;

    /// @brief Default move assignment operator
    /// @param other other Tensor object
    /// @return reference to the current object
    Tensor& operator=(Tensor&& other) = default;

    /**
     * @brief Destructor preserves unloading order of implementation object and reference to library
     */
    ~Tensor();

    /**
     * @brief Checks openvino tensor type
     * @param tensor a tensor which type will be checked
     * @throw Exception if type check with specified tensor is not pass
     */
    static void type_check(const Tensor& tensor);

    /**
     * @brief Constructs Tensor using element type and shape. Allocate internal host storage using default allocator
     * @param type Tensor element type
     * @param shape Tensor shape
     * @param allocator allocates memory for internal tensor storage
     */
    Tensor(const element::Type type, const Shape& shape, const Allocator& allocator = {});

    /**
     * @brief Constructs Tensor using element type and shape. Wraps allocated host memory.
     * @note Does not perform memory allocation internally
     * @param type Tensor element type
     * @param shape Tensor shape
     * @param host_ptr Pointer to pre-allocated host memory
     * @param strides Optional strides parameters in bytes. Strides are supposed to be computed automatically based
     * on shape and element size
     */
    Tensor(const element::Type type, const Shape& shape, void* host_ptr, const Strides& strides = {});

    /**
     * @brief Constructs region of interest (ROI) tensor form another tensor.
     * @note Does not perform memory allocation internally
     * @param other original tensor
     * @param begin start coordinate of ROI object inside of the original object.
     * @param end end coordinate of ROI object inside of the original object.
     * @note A Number of dimensions in `begin` and `end` must match number of dimensions in `other.get_shape()`
     */
    Tensor(const Tensor& other, const Coordinate& begin, const Coordinate& end);

    /**
     * @brief Set new shape for tensor, deallocate/allocate if new total size is bigger than previous one.
     * @note Memory allocation may happen
     * @param shape A new shape
     */
    void set_shape(const ov::Shape& shape);

    /**
     * @return A tensor element type
     */
    element::Type get_element_type() const;

    /**
     * @return A tensor shape
     */
    Shape get_shape() const;

    /**
     * @brief Returns the total number of elements (a product of all the dims or 1 for scalar)
     * @return The total number of elements
     */
    size_t get_size() const;

    /**
     * @brief Returns the size of the current Tensor in bytes.
     * @return Tensor's size in bytes
     */
    size_t get_byte_size() const;

    /**
     * @return Tensor's strides in bytes
     */
    Strides get_strides() const;

    /**
     * @brief Provides an access to the underlaying host memory
     * @param type Optional type parameter.
     * @note If type parameter is specified, the method throws an exception
     * if specified type's fundamental type does not match with tensor element type's fundamental type
     * @return A host pointer to tensor memory
     */
    void* data(const element::Type type = {}) const;

    /**
     * @brief Provides an access to the underlaying host memory casted to type `T`
     * @return A host pointer to tensor memory casted to specified type `T`.
     * @note Throws exception if specified type does not match with tensor element type
     */
    template <typename T, typename datatype = typename std::decay<T>::type>
    T* data() const {
        return static_cast<T*>(data(element::from<datatype>()));
    }

    /**
     * @brief Checks if current Tensor object is not initialized
     * @return `true` if current Tensor object is not initialized, `false` - otherwise
     */
    bool operator!() const noexcept;

    /**
     * @brief Checks if current Tensor object is initialized
     * @return `true` if current Tensor object is initialized, `false` - otherwise
     */
    explicit operator bool() const noexcept;

    /**
     * @brief Checks if the Tensor object can be cast to the type T
     *
     * @tparam T Type to be checked. Must represent a class derived from the Tensor
     * @return true if this object can be dynamically cast to the type const T*. Otherwise, false
     */
    template <typename T>
    typename std::enable_if<std::is_base_of<Tensor, T>::value, bool>::type is() const noexcept {
        try {
            T::type_check(*this);
        } catch (...) {
            return false;
        }
        return true;
    }

    /**
     * @brief Casts this Tensor object to the type T.
     *
     * @tparam T Type to cast to. Must represent a class derived from the Tensor
     * @return T object
     */
    template <typename T>
    const typename std::enable_if<std::is_base_of<Tensor, T>::value, T>::type as() const {
        T::type_check(*this);
        return *static_cast<const T*>(this);
    }
};

using TensorVector = std::vector<Tensor>;

namespace runtime {
using ov::Tensor;
using ov::TensorVector;
}  // namespace runtime

}  // namespace ov
