// Copyright (c) ONNX Project Contributors
//
// SPDX-License-Identifier: Apache-2.0

// Adapter for Add in default domain from version 6 to 5

#pragma once

#include <cstdint>
#include <memory>
#include <string>
#include <vector>

#include "onnx_lib/version_converter/adapters/adapter.h"

namespace ONNX_LIGHT_NAMESPACE::version_conversion {

class TypeRestriction : public Adapter {
public:
  explicit TypeRestriction(const std::string &op_name, const OpSetID &initial,
                           const OpSetID &target,
                           const std::vector<TensorProto_DataType> &unallowed_types)
      : Adapter(op_name, initial, target), unallowed_types_(unallowed_types) {}

  void adapt_type_restriction(const std::shared_ptr<Graph> & /*unused*/, const Node *node) const {
    // Since consumed_inputs is optional, no need to add it (as in batchnorm)
    // Iterate over all inputs and outputs
    for (const Value *input : node->inputs()) {
      isUnallowed(input);
    }
    for (const Value *output : node->outputs()) {
      isUnallowed(output);
    }
  }

  Node *adapt(std::shared_ptr<Graph> graph, Node *node) const override {
    adapt_type_restriction(graph, node);
    return node;
  }

private:
  std::vector<TensorProto_DataType> unallowed_types_;

  void isUnallowed(const Value *val) const {
    ONNX_ASSERTM(std::find(std::begin(unallowed_types_), std::end(unallowed_types_),
                           val->elemType()) == std::end(unallowed_types_),
                 "DataType (", val->elemType(), ") of Input or Output of operator '", name(),
                 "' is unallowed for Opset Version ",
                 static_cast<int64_t>(target_version().version()), ".")
  }
};

} // namespace ONNX_LIGHT_NAMESPACE::version_conversion
