// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#pragma once

#include "core/common/common.h"
#include "core/framework/op_kernel.h"
#include "core/util/math_cpuonly.h"

namespace onnxruntime {

template <typename T>
class ImageScaler final : public OpKernel {
 public:
  ImageScaler(const OpKernelInfo& info) : OpKernel(info) {
    ORT_ENFORCE(info.GetAttr<float>("scale", &scale_).IsOK());
    ORT_ENFORCE(info.GetAttrs<float>("bias", bias_).IsOK());
  }

  Status Compute(OpKernelContext* context) const override {
    const Tensor* X = context->Input<Tensor>(0);
    if (X == nullptr) return Status(common::ONNXRUNTIME, common::FAIL, "input count mismatch");

    const auto dims = X->Shape().GetDims();

    if (dims.size() < 4) {
      return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
                             "Input is expected to have four dimensions corresponding to [N,C,H,W], got ", dims.size());
    }

    const int64_t N = dims[0];
    const int64_t C = dims[1];
    const int64_t H = dims[2];
    const int64_t W = dims[3];

    if (!bias_.empty() && bias_.size() != static_cast<size_t>(C)) {
      return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Bias size (", bias_.size(), ") does not match the number of channels (", C, ")");
    }

    Tensor* Y = context->Output(0, TensorShape({N, C, H, W}));
    ConstEigenArrayMap<T> X_arr(X->template Data<T>(), H * W, N * C);
    EigenArrayMap<T> Y_arr(Y->template MutableData<T>(), H * W, N * C);

    for (int64_t nc = 0; nc < N * C; ++nc) {
      Y_arr.col(nc) = scale_ * X_arr.col(nc) + bias_[nc % C];
    }
    return Status::OK();
  }

 protected:
  float scale_;
  std::vector<float> bias_;
};

}  //namespace onnxruntime
