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

#pragma once

#include <memory>
#include <string>
#include <functional>
#include <vector>

#include <ie_data.h>
#include <ie_blob.h>

#include <vpu/model/base.hpp>
#include <vpu/model/edges.hpp>
#include <vpu/model/data_desc.hpp>
#include <vpu/backend/blob_serializer.hpp>
#include <vpu/utils/enums.hpp>
#include <vpu/utils/func_ref.hpp>

namespace vpu {

namespace ie = InferenceEngine;

//
// DataUsage
//

//
// Describes how Data object is used in the Model:
//   * Input / Output : network input or output.
//   * Const : constant values (weights, biases, etc.).
//   * Intermediate : Data that are used for intermediate results.
//   * Temp : temporary buffer.
//   * Fake : fake Data object to fill stage input/output port.
//

VPU_DECLARE_ENUM(DataUsage,
    Input,
    Output,
    Const,
    Intermediate,
    Temp,
    Fake
)

//
// DataLocation
//

//
// Describes where Data object is located.
//

// Must be synchronized with MvTensor
VPU_DECLARE_ENUM(DataLocation,
    None = 0,
    Input = 1,
    Output = 2,
    Blob = 3,
    BSS = 4,
    CMX = 5
)

VPU_DECLARE_ENUM(MemoryType,
    DDR,
    CMX)

//
// DataContent
//

//
// Content of the Const Data object.
//

class DataContent {
public:
    using Ptr = std::shared_ptr<DataContent>;

    virtual ~DataContent();

    // TYPED pointer
    template <typename T>
    const T* get() const {
        return static_cast<const T*>(getRaw());
    }

    const DataDesc& desc() const {
        return _desc;
    }

private:
    // RAW pointer
    virtual const void* getRaw() const = 0;

private:
    DataDesc _desc;

    friend ModelObj;
};

//
// Data content that is calculated on the fly, using lazy calculation:
//
//   * It performs calculation on the first call and stores it in internal buffer.
//   * Next access will return the pointer to calculated buffer.
//
class CalculatedDataContent : public DataContent {
public:
    CalculatedDataContent() = default;
    explicit CalculatedDataContent(const SmallVector<DataContent::Ptr, 2>& baseContents) : _baseContents(baseContents) {}

private:
    const void* getRaw() const override;

    virtual size_t getTempBufSize(const SmallVector<DataContent::Ptr, 2>& baseContents) const;
    virtual void fillTempBuf(const SmallVector<DataContent::Ptr, 2>& baseContents, void* tempBuf) const = 0;

private:
    mutable SmallVector<DataContent::Ptr, 2> _baseContents;
    mutable std::vector<uint8_t> _temp;
};

DataContent::Ptr ieBlobContent(
        const ie::Blob::Ptr& blob,
        int repeat = 1);

DataContent::Ptr replicateContent(float val, int count);
DataContent::Ptr replicateContent(const DataContent::Ptr& origContent, int count);

DataContent::Ptr scaleContent(const DataContent::Ptr& origContent, float scale);

// The function scales the major dimension of 4D origContent
DataContent::Ptr scaledChannelContent(
        const DataContent::Ptr& origContent,
        const DataContent::Ptr& scaleContent);

//
// DataNode
//

class DataNode final :
        public EnableHandle,
        public EnableCustomAttributes {
    //
    // Main attributes
    //

    VPU_MODEL_ATTRIBUTE(std::string, name, std::string())
    VPU_MODEL_ATTRIBUTE(DataUsage, usage, DataUsage::Fake)
    VPU_MODEL_ATTRIBUTE(DataDesc, desc, DataDesc())
    VPU_MODEL_ATTRIBUTE(StridesRequirement, requiredStrides, StridesRequirement::empty())

    //
    // Bindings with IE
    //

    VPU_MODEL_ATTRIBUTE(ie::DataPtr, origData, nullptr)

    //
    // Edges
    //

    VPU_MODEL_ATTRIBUTE(StageOutput, producerEdge, nullptr)
    VPU_MODEL_ATTRIBUTE_PTR_RANGE(StageInputList, consumerEdges)

    VPU_MODEL_ATTRIBUTE(StageTempBuffer, tempBufferEdge, nullptr)

    /**
     * Parent data edge actually allocates memory
     */
    VPU_MODEL_ATTRIBUTE(SharedAllocation, parentDataEdge, nullptr)

    /**
     * Children data edges uses parent's memory
     */
    VPU_MODEL_ATTRIBUTE_PTR_RANGE(SharedAllocationList, childDataEdges)

    //
    // Const data content
    //

    VPU_MODEL_ATTRIBUTE(DataContent::Ptr, content, nullptr)

    //
    // Allocation info
    //

    VPU_MODEL_ATTRIBUTE(MemoryType, memReqs, MemoryType::DDR)
    VPU_MODEL_ATTRIBUTE(DataLocation, location, DataLocation::None)
    VPU_MODEL_ATTRIBUTE(int, memoryOffset, 0)

    //
    // Edges wrappers
    //

    VPU_MODEL_ATTRIBUTE(Model, model, nullptr)

private:
    struct ConsumerAccess final {
        inline auto operator()(const StageInput& edge) const -> decltype(edge->consumer()) {
            return edge->consumer();
        }
    };

    struct ChildDataAccess final {
        inline auto operator()(const SharedAllocation& edge) const -> decltype(edge->child()) {
            return edge->child();
        }
    };

public:
    inline Stage producer() const {
        return _producerEdge == nullptr ? nullptr : _producerEdge->producer();
    }

    inline int numConsumers() const {
        return _consumerEdges.size();
    }
    inline auto consumers() const -> decltype(mapRange<ConsumerAccess>(consumerEdges())) {
        return mapRange<ConsumerAccess>(consumerEdges());
    }
    inline StageInput singleConsumerEdge() const {
        IE_ASSERT(_consumerEdges.size() == 1);
        return *_consumerEdges.begin();
    }
    inline Stage singleConsumer() const {
        return singleConsumerEdge()->consumer();
    }

    inline Data parentData() const {
        return _parentDataEdge == nullptr ? nullptr : _parentDataEdge->parent();
    }

    inline int numChildDatas() const {
        return _childDataEdges.size();
    }
    inline auto childDatas() const -> decltype(mapRange<ChildDataAccess>(childDataEdges())) {
        return mapRange<ChildDataAccess>(childDataEdges());
    }

    Data getTopParentData() const;

    //
    // DataDesc
    //

    DimValues strides() const;

    int totalByteSize() const;

    int elemOffset(const DimValues& coord) const;
    int lastElemOffset() const;

    bool canHaveAParent() const;

    //
    // Bindings with IE
    //

    inline void setOrigData(const ie::DataPtr& origData) { _origData = origData; }

    //
    // StridesRequirement
    //

    bool checkStrides(const StridesRequirement& reqs) const;

    inline void resetRequiredStrides() {
        _requiredStrides = StridesRequirement::empty();
    }

    void updateRequiredStrides(const StridesRequirement& newReqs);

    //
    // Allocation info
    //

    void clearAllocation();

    void setMemReqs(MemoryType mem);

    void setIOInfo(DataLocation location, int ioBufferOffset);

    void setAllocationInfo(DataLocation location, int memoryOffset);

    //
    // Backend utilities
    //

    // Serialize as-is for new MvTensor kernels that can work with ND data.
    // If `newOrder` is not empty, it will be used instead of original and missing dimensions will be set to 1.
    void serializeBuffer(
            BlobSerializer& serializer,
            DimsOrder newOrder = DimsOrder());

    void serializeIOInfo(BlobSerializer& serializer) const;

private:
    void serializeDescImpl(
            BlobSerializer& serializer,
            const DataDesc& storedDesc,
            const DimValues& storedStrides) const;

    void serializeBufferImpl(
            BlobSerializer& serializer,
            const DataDesc& storedDesc,
            const DimValues& storedStrides) const;

private:
    inline DataNode() :
        _consumerEdges(&StageInputEdge::_posInData),
        _childDataEdges(&SharedAllocationEdge::_posInData),
        _posInModel(this) {
    }

private:
    DataPtrList::iterator _ptrPosInModel;
    DataListNode _posInModel;

    friend ModelObj;
};

void printTo(std::ostream& os, const Data& data);

//
// loopOverData
//

VPU_DECLARE_ENUM(DataLoopStatus,
    NextChild,
    NextSibling,
    Stop)

void loopOverData(
        const Data& data,
        const FuncRef<DataLoopStatus(const Data&)>& op);

}  // namespace vpu
