// Copyright (c) OpenMMLab. All rights reserved.
#include "shape.h"

#include "../ncnn_ops_definer.h"

namespace mmdeploy {
using namespace ncnn;
DEFINE_LAYER_CREATOR(Shape)
DEFINE_NCNN_OPS(Shape, Shape)
Shape::Shape() {
  one_blob_only = true;
  support_inplace = false;
}

int Shape::forward(const Mat &bottom_blob, Mat &top_blob, const Option &opt) const {
  int dims = bottom_blob.dims;
  int w = bottom_blob.w;
  size_t elemsize = sizeof(float);
  top_blob.create(dims + 1, elemsize, opt.blob_allocator);
  if (top_blob.empty()) {
    return -100;
  }
  float *outptr = top_blob;
  if (dims == 1) {
    outptr[0] = 1.0f;
    outptr[1] = w;
    return 0;
  }
  if (dims == 2) {
    int h = bottom_blob.h;
    outptr[0] = 1.0f;
    outptr[1] = h;
    outptr[2] = w;
    return 0;
  }
  if (dims == 3) {
    int h = bottom_blob.h;
    int channels = bottom_blob.c;
    outptr[0] = 1.0f;
    outptr[1] = channels;
    outptr[2] = h;
    outptr[3] = w;
    return 0;
  }
}

}  // namespace mmdeploy