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

// Adapter for Cast in default domain from version 9 to 8

#pragma once

#include <cinttypes>
#include <cstdint>
#include <memory>
#include <vector>

#include "onnx/common/assertions.h"
#include "onnx/version_converter/adapters/type_restriction.h"

namespace ONNX_NAMESPACE::version_conversion {

static const std::vector<TensorProto_DataType> q_dq_20_unallowed_types = {
    TensorProto_DataType_UINT16,
    TensorProto_DataType_INT16,
    TensorProto_DataType_UINT4,
    TensorProto_DataType_INT4};

class QuantizeLinear_21_20 final : public TypeRestriction {
 public:
  explicit QuantizeLinear_21_20()
      : TypeRestriction("QuantizeLinear", OpSetID(21), OpSetID(20), q_dq_20_unallowed_types) {}

  void adapt_quantize_linear_21_20(const std::shared_ptr<Graph>& /*unused*/, Node* node) const {
    if (node->hasAttribute(kblock_size)) {
      if ((node->i(kblock_size) != 0)) {
        ONNX_ASSERTM(
            false,
            "Blocked quantization is not supported for Opset Version ",
            static_cast<int64_t>(target_version().version()),
            ".")
      }
      node->removeAttribute(kblock_size);
    }
    if (node->hasAttribute(koutput_dtype)) {
      if (node->i(koutput_dtype) != TensorProto_DataType_UINT8 && node->inputs().size() < 3) {
        ONNX_ASSERTM(
            false,
            "Attribute output_dtype is not supported for Opset Version ",
            static_cast<int64_t>(target_version().version()),
            ", supply a zero-point tensor instead")
      }
      node->removeAttribute(koutput_dtype);
    }
  }

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

class DequantizeLinear_21_20 final : public TypeRestriction {
 public:
  explicit DequantizeLinear_21_20()
      : TypeRestriction("DequantizeLinear", OpSetID(21), OpSetID(20), q_dq_20_unallowed_types) {}

  void adapt_dequantize_linear_21_20(const std::shared_ptr<Graph>& /*unused*/, Node* node) const {
    if (node->hasAttribute(kblock_size)) {
      if ((node->i(kblock_size) != 0)) {
        ONNX_ASSERTM(
            false,
            "Blocked quantization is not supported for Opset Version ",
            static_cast<int64_t>(target_version().version()),
            ".")
      }
      node->removeAttribute(kblock_size);
    }
  }

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

} // namespace ONNX_NAMESPACE::version_conversion