#pragma once

#include <cfloat>
#include <cmath>
#include "caffe2/core/context.h"
#include "caffe2/core/operator.h"
#include "caffe2/utils/math.h"

namespace caffe2 {

template <typename Context>
void lr_update(
    int n,
    const float* grad,
    const float* effgrad,
    const float* lr,
    float* nlr,
    float lr_alpha,
    bool normalized_lr_adaption,
    Context* /*context*/) {
  float x = 0;
  float y = 0, z = 0;
  const float kEps = 1e-12f;
  for (auto i = 0; i < n; i++) {
    x += grad[i] * effgrad[i];
    if (normalized_lr_adaption) {
      y += grad[i] * grad[i];
      z += effgrad[i] * effgrad[i];
    }
  }
  if (normalized_lr_adaption) {
    y = fmax(std::sqrt(y), kEps);
    z = fmax(std::sqrt(z), kEps);
    nlr[0] = lr[0] * (1 - lr_alpha * x / (y * z));
  } else {
    nlr[0] = lr[0] - lr_alpha * x;
  }
}

template <typename T, class Context>
class LearningRateAdaptionOp final : public Operator<Context> {
 public:
  LearningRateAdaptionOp(const OperatorDef& operator_def, Workspace* ws)
      : Operator<Context>(operator_def, ws),
        lr_alpha_(this->template GetSingleArgument<float>("lr_alpha", 0.01f)),
        normalized_lr_adaption_(this->template GetSingleArgument<bool>(
            "normalized_lr_adaption",
            true)) {}
  USE_OPERATOR_CONTEXT_FUNCTIONS;

  bool RunOnDevice() override {
    CAFFE_ENFORCE(Input(LR).numel() == 1);
    CAFFE_ENFORCE(Input(GRAD).numel() == Input(EFFGRAD).numel());
    Output(OUTPUT_LR)->ResizeLike(Input(LR));
    lr_update<Context>(
        Input(GRAD).numel(),
        Input(GRAD).template data<T>(),
        Input(EFFGRAD).template data<T>(),
        Input(LR).template data<T>(),
        Output(OUTPUT_LR)->template mutable_data<T>(),
        lr_alpha_,
        normalized_lr_adaption_,
        &context_);
    return true;
  }

 protected:
  T lr_alpha_{1e-2};
  bool normalized_lr_adaption_{true};
  INPUT_TAGS(LR, GRAD, EFFGRAD);
  OUTPUT_TAGS(OUTPUT_LR);
};

} // namespace caffe2