// This file is part of OpenCV project.
// It is subject to the license terms in the LICENSE file found in the top-level directory
// of this distribution and at http://opencv.org/license.html.

#include "../precomp.hpp"
#include "layers_common.hpp"
#include "../net_impl.hpp"
//#include "../op_cuda.hpp"
//#include "../op_inf_engine.hpp"
//#include "../ie_ngraph.hpp"
//#include "../op_webnn.hpp"
//#include "../op_timvx.hpp"
//#include "../op_cann.hpp"

//#include <opencv2/dnn/shape_utils.hpp>

namespace cv
{
namespace dnn
{

/*
    Reshape2 layer, as defined in ONNX specification:
    https://onnx.ai/onnx/operators/onnx__Reshape.html

    Opset's 1 to 23 are covered.

    The layers Flatten, Reshape2, Squeeze and Unsqueeze all share the same
    implementation idea:
    1. calculate shape of the output tensor
    2. assuming that the input is continuous, just copy all the data to output tensor.
       reshapeAndCopyFirst() does that.
       The engine buffer allocator recognizes all these operations and tries to run
       them in-place. In such a case no copy operation is actually done,
       so the operations are really cheap.
*/

class Reshape2LayerImpl CV_FINAL : public Reshape2Layer
{
public:
    bool dynamicShapeSpec;

    static bool sameDims(const MatShape& a, const MatShape& b)
    {
        if (a.dims != b.dims)
            return false;
        for (int i = 0; i < a.dims; i++) {
            if (a[i] != b[i])
                return false;
        }
        return true;
    }

    DataLayout getOriginalLayout() const
    {
        Net::Impl* netimpl_ = getNetImpl(this);
        DataLayout layout = netimpl_ ? netimpl_->originalLayout : DATA_LAYOUT_NCHW;
        CV_Assert(layout == DATA_LAYOUT_NCHW || layout == DATA_LAYOUT_NHWC);
        return layout;
    }

    bool getConstShapeSpec(MatShape& shapeSpec) const
    {
        if (!dynamicShapeSpec) {
            shapeSpec = newShapeDesc;
            return true;
        }

        if (this->inputs.size() != 2)
            return false;

        Net::Impl* netimpl_ = getNetImpl(this);
        if (!netimpl_ || !netimpl_->isConstArg(this->inputs[1]))
            return false;

        Mat shapeTensor = netimpl_->argTensor(this->inputs[1]);
        shapeSpec = tensorToShape(shapeTensor);
        return true;
    }

    bool canKeepBlockLayout(const MatShape& inpShape, const MatShape& shapeSpec) const
    {
        if (inpShape.layout != DATA_LAYOUT_BLOCK || inpShape.dims != 5)
            return false;

        DataLayout origLayout = getOriginalLayout();
        MatShape semanticInp = inpShape.toLayout(origLayout);
        MatShape semanticOut = getOutShape(semanticInp, shapeSpec);
        if (sameDims(semanticOut, semanticInp))
            return true;
        if (semanticOut.dims != 4 || semanticInp.dims != 4)
            return false;
        if (origLayout == DATA_LAYOUT_NCHW)
            return semanticOut[0] == semanticInp[0] && semanticOut[1] == semanticInp[1];
        else
            return semanticOut[0] == semanticInp[0] && semanticOut[3] == semanticInp[3];
    }

    Reshape2LayerImpl(const LayerParams& params)
    {
        dynamicShapeSpec = true;
        setParamsFrom(params);
        if (params.has("shape"))
        {
            dynamicShapeSpec = false;

            const DictValue& shapeParam = params.get("shape");
            int i, ndims = shapeParam.size();
            newShapeDesc.resize(ndims);
            for (i = 0; i < ndims; i++) {
                newShapeDesc[i] = shapeParam.get<int>(i);
            }
        }
    }

    bool isDataShuffling() const CV_OVERRIDE { return true; }

    virtual bool dynamicOutputShapes() const CV_OVERRIDE
    {
        if (!dynamicShapeSpec)
            return false;
        // When the shape comes from a 2nd input, check if it's a constant tensor.
        // If so, the output shape can still be inferred.
        Net::Impl* netimpl_ = getNetImpl(this);
        if (netimpl_ && this->inputs.size() == 2)
            return !netimpl_->isConstArg(this->inputs[1]);
        return true;
    }

    virtual bool supportBackend(int backendId) CV_OVERRIDE
    {
        return backendId == DNN_BACKEND_OPENCV;
    }

    // Reshape just re-interprets the same contiguous buffer. Let the graph
    // buffer allocator alias input and output so the memcpy in forward() is a no-op.
    virtual bool alwaysSupportInplace() const CV_OVERRIDE { return true; }

    bool haveShapeSpec() const
    {
        return newShapeDesc.dims >= 0;
    }

    MatShape getOutShape(const MatShape& inpShape, const MatShape& shapeSpec) const
    {
        MatShape outShape = shapeSpec;
        int m1idx = -1;
        int i, ndims = outShape.dims;
        int64_t outTotal = 1;
        for (i = 0; i < ndims; i++) {
            if (outShape[i] < 0) {
                CV_Assert(outShape[i] == -1);
                if (m1idx >= 0) {
                    CV_Error(Error::StsBadArg, "invalid shape spec, there must be at most one '-1'");
                }
                m1idx = i;
            }
            else {
                if (outShape[i] == 0) {
                    if (i >= inpShape.dims) {
                        CV_Error(Error::StsBadArg, "cannot copy dimension from the input tensor");
                    }
                    outShape[i] = inpShape[i];
                }
                outTotal *= outShape[i];
            }
        }

        int64_t inpTotal = (int64_t)inpShape.total();
        if (m1idx >= 0) {
            int64_t autoSize = inpTotal/outTotal;
            CV_Assert(autoSize <= INT_MAX && autoSize*outTotal == inpTotal);
            outShape[m1idx] = (int)autoSize;
        }
        else if (outTotal != inpTotal && ndims == 2 && outShape[0] == 1 && outShape[1] > 0) {
            outShape[1] = (int)(inpTotal / outShape[0]);
        }
        else {
            CV_Assert(outTotal == inpTotal);
        }

        return outShape;
    }

    bool getMemoryShapes(const std::vector<MatShape> &inputs,
                         const int,
                         std::vector<MatShape> &outputs,
                         std::vector<MatShape> &internals) const CV_OVERRIDE
    {
        bool haveShapeSpec_ = haveShapeSpec();
        CV_Assert((inputs.size() == 1 && haveShapeSpec_) ||
                  (inputs.size() == 2 && !haveShapeSpec_));
        MatShape shapeSpec = newShapeDesc, outShape;

        if (inputs.size() == 2)
        {
            CV_Assert(this->inputs.size() == 2);
            Net::Impl* netimpl_ = getNetImpl(this);
            Mat shapeTensor = netimpl_->argTensor(this->inputs[1]);
            shapeSpec = tensorToShape(shapeTensor);
        } else {
            CV_Assert(shapeSpec.dims >= 0);
        }

        if (inputs[0].layout == DATA_LAYOUT_BLOCK) {
            if (canKeepBlockLayout(inputs[0], shapeSpec)) {
                DataLayout origLayout = getOriginalLayout();
                MatShape semanticInp = inputs[0].toLayout(origLayout);
                MatShape semanticOut = getOutShape(semanticInp, shapeSpec);
                if (sameDims(semanticOut, semanticInp)) {
                    outputs.assign(1, inputs[0]);
                } else {
                    outputs.assign(1, semanticOut.toLayout(DATA_LAYOUT_BLOCK, inputs[0][4]));
                }
            } else {
                DataLayout origLayout = getOriginalLayout();
                MatShape semanticInp = inputs[0].toLayout(origLayout);
                outputs.assign(1, getOutShape(semanticInp, shapeSpec));
            }
        } else {
            outputs.assign(1, getOutShape(inputs[0], shapeSpec));
        }
        internals.clear();
        return true;
    }

    int getLayouts(const std::vector<DataLayout>& actualInputs,
                   std::vector<DataLayout>& desiredInputs,
                   const int requiredOutputs,
                   std::vector<DataLayout>& outputs) const CV_OVERRIDE
    {
        size_t ninputs = actualInputs.size();
        CV_Assert(ninputs == 1u || ninputs == 2u);

        desiredInputs.assign(ninputs, DATA_LAYOUT_UNKNOWN);
        desiredInputs[0] = actualInputs[0];
        outputs.assign(requiredOutputs, actualInputs[0]);

        if (actualInputs[0] == DATA_LAYOUT_BLOCK) {
            bool canUseBlock = false;
            MatShape shapeSpec;
            if (getConstShapeSpec(shapeSpec)) {
                DataLayout origLayout = getOriginalLayout();
                int cAxis = (origLayout == DATA_LAYOUT_NCHW) ? 1 : 3;
                if (shapeSpec.dims == 4 && cAxis < shapeSpec.dims && shapeSpec[cAxis] == 0)
                    canUseBlock = true;
            }
            if (!canUseBlock) {
                desiredInputs[0] = getOriginalLayout();
                outputs.assign(requiredOutputs, DATA_LAYOUT_UNKNOWN);
            }
        }

        if (ninputs > 1)
            desiredInputs[1] = DATA_LAYOUT_UNKNOWN;

        return outputs[0] == DATA_LAYOUT_BLOCK ? getNetImpl(this)->defaultC0 : 0;
    }

    void getTypes(const std::vector<MatType>& inputs,
        const int requiredOutputs,
        const int requiredInternals,
        std::vector<MatType>& outputs,
        std::vector<MatType>& internals) const CV_OVERRIDE
    {
        size_t ninputs = inputs.size();
        CV_Assert(ninputs == 1 || ninputs == 2);
        outputs.assign(requiredOutputs, inputs[0]);
        CV_Assert(requiredInternals == 0);
        internals.clear();
    }

    void finalize(InputArrayOfArrays, OutputArrayOfArrays outputs_arr) CV_OVERRIDE
    {
    }

    void forward(InputArrayOfArrays inputs_arr,
                 OutputArrayOfArrays outputs_arr,
                 OutputArrayOfArrays) CV_OVERRIDE
    {
        CV_TRACE_FUNCTION();
        CV_TRACE_ARG_VALUE(name, "name", name.c_str());

        Size size = inputs_arr.size();
        int ninputs = size.area();
        bool haveShapeSpec_ = haveShapeSpec();
        CV_Assert((ninputs == 1 && haveShapeSpec_) ||
                  (ninputs == 2 && !haveShapeSpec_));

        MatShape inpShape = inputs_arr.shape(0);
        CV_Assert(inpShape.layout != DATA_LAYOUT_BLOCK || inpShape.C > 0);
        MatShape shapeSpec = newShapeDesc;
        if (!haveShapeSpec_) {
            Mat shapeTensor = inputs_arr.getMat(1);
            shapeSpec = tensorToShape(shapeTensor);
        }
        MatShape outShape;
        if (inpShape.layout == DATA_LAYOUT_BLOCK) {
            if (canKeepBlockLayout(inpShape, shapeSpec)) {
                DataLayout origLayout = getOriginalLayout();
                MatShape semanticInp = inpShape.toLayout(origLayout);
                MatShape semanticOut = getOutShape(semanticInp, shapeSpec);
                if (sameDims(semanticOut, semanticInp)) {
                    outShape = inpShape;
                } else {
                    outShape = semanticOut.toLayout(DATA_LAYOUT_BLOCK, inpShape[4]);
                }
            } else {
                DataLayout origLayout = getOriginalLayout();
                MatShape semanticInp = inpShape.toLayout(origLayout);
                outShape = getOutShape(semanticInp, shapeSpec);
            }
        } else {
            outShape = getOutShape(inpShape, shapeSpec);
        }

        reshapeAndCopyFirst(inputs_arr, outputs_arr, outShape);
    }
};

Ptr<Reshape2Layer> Reshape2Layer::create(const LayerParams& params)
{
    return Ptr<Reshape2Layer>(new Reshape2LayerImpl(params));
}

}
}
