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

#pragma once

#include "core/providers/dml/DmlExecutionProvider/inc/MLOperatorAuthor.h"
#include "core/common//common.h"

struct NoisyReluShapeInferrer : winrt::implements<NoisyReluShapeInferrer, IMLOperatorShapeInferrer> {
  STDMETHOD(InferOutputShapes)(IMLOperatorShapeInferenceContext* context) noexcept {
    try {
      uint32_t inputDimsSize;
      context->GetInputTensorDimensionCount(0, &inputDimsSize);

      uint32_t* inputDims = new uint32_t[inputDimsSize];
      context->GetInputTensorShape(0, inputDimsSize, inputDims);

      context->SetOutputTensorShape(0, inputDimsSize, inputDims);
      return S_OK;
    } catch (...) {
      return winrt::to_hresult();
    }
  }
};

struct NoisyReluOperator : winrt::implements<NoisyReluOperator, IMLOperatorKernel> {
  float m_mean;
  float m_variance;

  NoisyReluOperator(float mean, float variance) : m_mean(mean), m_variance(variance) {}

  // Computes the outputs of the kernel.  This may be called multiple times
  // simultaneously within the same instance of the class.  Implementations
  // of this method must be thread-safe.
  STDMETHOD(Compute)(IMLOperatorKernelContext* context) {
    try {
      // Get the input tensor
      winrt::com_ptr<IMLOperatorTensor> inputTensor;
      context->GetInputTensor(0, inputTensor.put());

      // Get the output tensor
      winrt::com_ptr<IMLOperatorTensor> outputTensor;
      context->GetOutputTensor(0, outputTensor.put());

      // Get the input and output shape sizes
      uint32_t inputDimsSize = inputTensor->GetDimensionCount();
      uint32_t outputDimsSize = outputTensor->GetDimensionCount();
      if (inputDimsSize != outputDimsSize) {
        return E_UNEXPECTED;
      }

      // Get the input shape
      std::vector<uint32_t> inputDims(inputDimsSize);
      outputTensor->GetShape(inputDimsSize, inputDims.data());

      // Get the output shape
      std::vector<uint32_t> outputDims(outputDimsSize);
      outputTensor->GetShape(outputDimsSize, outputDims.data());

      // For the number of total elements in the input and output shapes
      auto outputDataSize = std::accumulate(outputDims.begin(), outputDims.end(), 1, std::multiplies<uint32_t>());
      auto inputDataSize = std::accumulate(inputDims.begin(), inputDims.end(), 1, std::multiplies<uint32_t>());
      if (outputDataSize != inputDataSize) {
        return E_UNEXPECTED;
      }

      // If the tensor types are both float type
      if (outputTensor->GetTensorDataType() == MLOperatorTensorDataType::Float &&
          inputTensor->GetTensorDataType() == MLOperatorTensorDataType::Float) {
        // For cpu data
        if (outputTensor->IsCpuData() && inputTensor->IsCpuData()) {
          ComputeInternal<float>(inputTensor.get(), outputTensor.get(), inputDataSize);
        }
      } else if (outputTensor->GetTensorDataType() == MLOperatorTensorDataType::Double &&
                 inputTensor->GetTensorDataType() == MLOperatorTensorDataType::Double) {
        // For cpu data
        if (outputTensor->IsCpuData() && inputTensor->IsCpuData()) {
          ComputeInternal<double>(inputTensor.get(), outputTensor.get(), inputDataSize);
        }
      }

      return S_OK;
    } catch (...) {
      return winrt::to_hresult();
    }
  }

  template <typename T, typename U = T>
  void ComputeInternal(IMLOperatorTensor* pInputTensor, IMLOperatorTensor* pOutputTensor, uint32_t size) {
    // Create a normal distribution
    std::normal_distribution<> dist{m_mean, m_variance};
    std::random_device rd{};
    std::mt19937 gen{rd()};

    auto inputData = static_cast<T*>(pInputTensor->GetData());
    auto outputData = static_cast<U*>(pOutputTensor->GetData());

    for (uint32_t i = 0; i < size; i++) {
      outputData[i] = static_cast<U>(std::max<T>(0, static_cast<T>(inputData[i] + dist(gen))));
    }
  }
};

struct NoisyReluOperatorFactory : winrt::implements<NoisyReluOperatorFactory, IMLOperatorKernelFactory> {
  STDMETHOD(CreateKernel)(IMLOperatorKernelCreationContext* context, IMLOperatorKernel** kernel) {
    try {
      float mean;
      context->GetAttribute("mean", MLOperatorAttributeType::Float, 1, sizeof(float), reinterpret_cast<void*>(&mean));
      float variance;
      context->GetAttribute(
        "variance", MLOperatorAttributeType::Float, 1, sizeof(float), reinterpret_cast<void*>(&variance)
      );

      auto noisyReluOperator = winrt::make<NoisyReluOperator>(mean, variance);
      noisyReluOperator.copy_to(kernel);
      return S_OK;
    } catch (...) {
      return winrt::to_hresult();
    }
  }

  static MLOperatorEdgeDescription CreateEdgeDescriptor(MLOperatorEdgeType type, MLOperatorTensorDataType dataType) {
    ORT_UNUSED_PARAMETER(type);
    MLOperatorEdgeDescription desc;
    desc.edgeType = MLOperatorEdgeType::Tensor;
    desc.tensorDataType = dataType;
    return desc;
  }

  static void RegisterNoisyReluSchema(winrt::com_ptr<IMLOperatorRegistry> registry) {
    MLOperatorSetId operatorSetId;
    operatorSetId.domain = "";
    operatorSetId.version = 7;

    MLOperatorSchemaDescription noisyReluSchema;
    noisyReluSchema.name = "NoisyRelu";
    noisyReluSchema.operatorSetVersionAtLastChange = 1;

    MLOperatorSchemaEdgeDescription noisyReluXInput;
    noisyReluXInput.options = MLOperatorParameterOptions::Single;
    noisyReluXInput.typeFormat = MLOperatorSchemaEdgeTypeFormat::Label;
    noisyReluXInput.typeLabel = "T";

    std::vector<MLOperatorSchemaEdgeDescription> inputs{noisyReluXInput};
    noisyReluSchema.inputs = inputs.data();
    noisyReluSchema.inputCount = static_cast<uint32_t>(inputs.size());

    MLOperatorSchemaEdgeDescription noisyReluXOutput;
    noisyReluXOutput.options = MLOperatorParameterOptions::Single;
    noisyReluXOutput.typeFormat = MLOperatorSchemaEdgeTypeFormat::Label;
    noisyReluXOutput.typeLabel = "T";

    std::vector<MLOperatorSchemaEdgeDescription> outputs{noisyReluXOutput};
    noisyReluSchema.outputs = outputs.data();
    noisyReluSchema.outputCount = static_cast<uint32_t>(outputs.size());

    MLOperatorEdgeTypeConstrant typeConstraint;
    typeConstraint.typeLabel = "T";
    std::vector<MLOperatorEdgeDescription> allowedEdges{
      CreateEdgeDescriptor(MLOperatorEdgeType::Tensor, MLOperatorTensorDataType::Double),
      CreateEdgeDescriptor(MLOperatorEdgeType::Tensor, MLOperatorTensorDataType::Float),
      CreateEdgeDescriptor(MLOperatorEdgeType::Tensor, MLOperatorTensorDataType::Float16)
    };
    typeConstraint.allowedTypes = allowedEdges.data();
    typeConstraint.allowedTypeCount = static_cast<uint32_t>(allowedEdges.size());

    std::vector<MLOperatorEdgeTypeConstrant> typeConstraints{typeConstraint};
    noisyReluSchema.typeConstraints = typeConstraints.data();
    noisyReluSchema.typeConstraintCount = static_cast<uint32_t>(typeConstraints.size());

    MLOperatorAttribute noisyReluMeanAttribute;
    noisyReluMeanAttribute.name = "mean";
    noisyReluMeanAttribute.required = false;
    noisyReluMeanAttribute.type = MLOperatorAttributeType::Float;

    MLOperatorAttribute noisyReluVarianceAttribute;
    noisyReluVarianceAttribute.name = "variance";
    noisyReluVarianceAttribute.required = false;
    noisyReluVarianceAttribute.type = MLOperatorAttributeType::Float;

    std::vector<MLOperatorAttribute> attributes{noisyReluMeanAttribute, noisyReluVarianceAttribute};
    noisyReluSchema.attributes = attributes.data();
    noisyReluSchema.attributeCount = static_cast<uint32_t>(attributes.size());

    MLOperatorAttributeNameValue noisyReluMeanAttributeValue;
    noisyReluMeanAttributeValue.name = "mean";
    noisyReluMeanAttributeValue.type = MLOperatorAttributeType::Float;
    noisyReluMeanAttributeValue.valueCount = 1;
    static float defaultMeans[] = {0.f};
    noisyReluMeanAttributeValue.floats = defaultMeans;

    MLOperatorAttributeNameValue noisyReluVarianceAttributeValue;
    noisyReluVarianceAttributeValue.name = "variance";
    noisyReluVarianceAttributeValue.type = MLOperatorAttributeType::Float;
    noisyReluVarianceAttributeValue.valueCount = 1;
    static float defaultVariance[] = {1.f};
    noisyReluVarianceAttributeValue.floats = defaultVariance;

    std::vector<MLOperatorAttributeNameValue> attributeDefaultValues{
      noisyReluMeanAttributeValue, noisyReluVarianceAttributeValue
    };
    noisyReluSchema.defaultAttributes = attributeDefaultValues.data();
    noisyReluSchema.defaultAttributeCount = static_cast<uint32_t>(attributeDefaultValues.size());

    std::vector<const MLOperatorSchemaDescription*> schemas{&noisyReluSchema};
    registry->RegisterOperatorSetSchema(
      &operatorSetId, 6 /* baseline version */, schemas.data(), static_cast<uint32_t>(schemas.size()), nullptr, nullptr
    );
  }

  static void RegisterNoisyReluKernel(winrt::com_ptr<IMLOperatorRegistry> registry) {
    MLOperatorKernelDescription kernelDescription;
    kernelDescription.domain = "";
    kernelDescription.name = "NoisyRelu";
    kernelDescription.minimumOperatorSetVersion = 1;
    kernelDescription.executionType = MLOperatorExecutionType::Cpu;

    MLOperatorEdgeTypeConstrant typeConstraint;
    typeConstraint.typeLabel = "T";
    std::vector<MLOperatorEdgeDescription> allowedEdges{
      CreateEdgeDescriptor(MLOperatorEdgeType::Tensor, MLOperatorTensorDataType::Double),
      CreateEdgeDescriptor(MLOperatorEdgeType::Tensor, MLOperatorTensorDataType::Float),
      CreateEdgeDescriptor(MLOperatorEdgeType::Tensor, MLOperatorTensorDataType::Float16)
    };
    typeConstraint.allowedTypes = allowedEdges.data();
    typeConstraint.allowedTypeCount = static_cast<uint32_t>(allowedEdges.size());

    std::vector<MLOperatorEdgeTypeConstrant> typeConstraints{typeConstraint};
    kernelDescription.typeConstraints = typeConstraints.data();
    kernelDescription.typeConstraintCount = static_cast<uint32_t>(typeConstraints.size());

    MLOperatorAttributeNameValue noisyReluMeanAttributeValue;
    noisyReluMeanAttributeValue.name = "mean";
    noisyReluMeanAttributeValue.type = MLOperatorAttributeType::Float;
    noisyReluMeanAttributeValue.valueCount = 1;
    static float defaultMeans[] = {0.f};
    noisyReluMeanAttributeValue.floats = defaultMeans;

    MLOperatorAttributeNameValue noisyReluVarianceAttributeValue;
    noisyReluVarianceAttributeValue.name = "variance";
    noisyReluVarianceAttributeValue.type = MLOperatorAttributeType::Float;
    noisyReluVarianceAttributeValue.valueCount = 1;
    static float defaultVariance[] = {1.f};
    noisyReluVarianceAttributeValue.floats = defaultVariance;

    std::vector<MLOperatorAttributeNameValue> attributeDefaultValues{
      noisyReluMeanAttributeValue, noisyReluVarianceAttributeValue
    };
    kernelDescription.defaultAttributes = attributeDefaultValues.data();
    kernelDescription.defaultAttributeCount = static_cast<uint32_t>(attributeDefaultValues.size());
    kernelDescription.options = MLOperatorKernelOptions::None;
    kernelDescription.executionOptions = 0;

    auto factory = winrt::make<NoisyReluOperatorFactory>();
    auto shareInferrer = winrt::make<NoisyReluShapeInferrer>();

    registry->RegisterOperatorKernel(&kernelDescription, factory.get(), shareInferrer.get());
  }
};