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

#include "onnx_extensions/patterns/collections/shape_pattern.h"

#include <algorithm>
#include <cstdint>
#include <string>
#include <vector>

#include "onnx_core/builder/graph_graph.h"
#include "onnx_extensions/patterns/collections/collections_utils.h"
#include "onnx_proto/onnx_helper.h"

namespace ONNX_LIGHT_NAMESPACE::onnx_patterns {

namespace {

using collections::FreeInitializerName;
using collections::IsDefaultOp;
using core::builder::BuilderError;

// Returns the integer ``name`` attribute of ``node`` in ``out`` when present.
bool TryGetAttr(const NodeProto &node, const char *name, int64_t &out) {
  const AttributeProto *attr = FindAttribute(node, name);
  if (attr == nullptr) {
    return false;
  }
  out = attr->i();
  return true;
}

// Normalises a possibly-negative dimension index and clamps it to ``[0, rank]``.
int64_t ClampDim(int64_t value, int64_t rank) {
  if (value < 0) {
    value += rank;
  }
  return std::max<int64_t>(0, std::min<int64_t>(value, rank));
}

bool HasRank(core::builder::GraphGraph &graph, const std::string &name) {
  return graph.HasShape(name);
}

int64_t GetRank(core::builder::GraphGraph &graph, const std::string &name) {
  return static_cast<int64_t>(graph.GetShape(name).Shape().Rank());
}

} // namespace

std::set<std::string> GatherShapePattern::FastOpType() const { return {"Gather"}; }

core::builder::MatchResult GatherShapePattern::Match(core::builder::GraphGraph &graph,
                                                     const NodeProto &candidate) const {
  if (!IsDefaultOp(candidate, "Gather")) {
    return NoMatch(candidate, "candidate is not a default-domain Gather");
  }
  if (GetAttributeOr<int64_t>(candidate, "axis", 0) != 0) {
    return NoMatch(candidate, "the Gather does not operate on axis 0");
  }
  const std::string &index_name = candidate.input()[1].value();
  if (!graph.IsConstant(index_name)) {
    return NoMatch(candidate, "the Gather index is not constant");
  }
  const TensorProto *indices = graph.GetComputedConstant(index_name);
  if (indices == nullptr ||
      static_cast<TensorProto::DataType>(indices->data_type()) != TensorProto::DataType::INT64) {
    return NoMatch(candidate, "the Gather index is not a materialised int64 tensor");
  }
  const int rank_index = indices->dims_size();
  if (rank_index > 1) {
    return NoMatch(candidate, "the Gather index is not a scalar or a vector");
  }
  std::vector<int64_t> values;
  if (!ReadIntegerValues(*indices, values) || values.empty()) {
    return NoMatch(candidate, "the Gather index could not be read");
  }
  if (rank_index == 1 && values.size() > 1) {
    for (std::size_t i = 1; i < values.size(); ++i) {
      if (values[i] != values[i - 1] + 1) {
        return NoMatch(candidate, "the Gather index is not a contiguous ascending range");
      }
    }
  }
  const NodeProto *shape = graph.NodeBefore(candidate.input()[0].value());
  if (shape == nullptr || !IsDefaultOp(*shape, "Shape")) {
    return NoMatch(candidate, "the gathered value is not produced by a Shape");
  }
  int64_t shape_start = 0;
  int64_t shape_end = 0;
  const bool has_start = TryGetAttr(*shape, "start", shape_start);
  const bool has_end = TryGetAttr(*shape, "end", shape_end);
  const bool negative_attr = (has_start && shape_start < 0) || (has_end && shape_end < 0);
  const bool negative_index = values.front() < 0;
  if ((negative_attr || negative_index) && !HasRank(graph, shape->input()[0].value())) {
    return NoMatch(candidate, "the shape origin rank is required to normalise negative bounds");
  }
  return core::builder::MatchResult{this, {shape, &candidate}, &candidate};
}

utils::RepeatedProtoField<NodeProto>
GatherShapePattern::Apply(core::builder::GraphGraph &graph,
                          const std::vector<const NodeProto *> &nodes) const {
  if (nodes.size() != 2 || nodes[0] == nullptr || nodes[1] == nullptr) {
    throw BuilderError("GatherShapePattern::Apply expects a Shape and a Gather node.");
  }
  const core::builder::MatchResult verified = Match(graph, *nodes[1]);
  if (verified.pattern == nullptr || verified.nodes != nodes) {
    throw BuilderError("GatherShapePattern::Apply received an unsafe or inconsistent match.");
  }
  const NodeProto &shape = *nodes[0];
  const NodeProto &gather = *nodes[1];

  const TensorProto *indices = graph.GetComputedConstant(gather.input()[1].value());
  std::vector<int64_t> values;
  if (indices == nullptr || !ReadIntegerValues(*indices, values) || values.empty()) {
    throw BuilderError("GatherShapePattern::Apply could not read the Gather index.");
  }
  const bool is_scalar = indices->dims_size() == 0;
  const std::string &x_name = shape.input()[0].value();

  int64_t shape_start = 0;
  int64_t shape_end = 0;
  const bool has_start = TryGetAttr(shape, "start", shape_start);
  const bool has_end = TryGetAttr(shape, "end", shape_end);

  int64_t s0 = has_start ? shape_start : 0;
  int64_t length = -1;
  const bool has_rank = HasRank(graph, x_name);
  if (has_rank) {
    const int64_t rank = GetRank(graph, x_name);
    const int64_t e0 = ClampDim(has_end ? shape_end : rank, rank);
    s0 = ClampDim(s0, rank);
    length = e0 - s0;
  } else if (has_end) {
    length = shape_end - s0;
  }

  int64_t gather_first = values.front();
  int64_t gather_last = values.back();
  if (gather_first < 0 && length >= 0) {
    gather_first += length;
  }
  if (gather_last < 0 && length >= 0) {
    gather_last += length;
  }
  const int64_t new_start = s0 + gather_first;
  const int64_t new_end = s0 + gather_last + 1;

  core::builder::GraphBuilder &builder = graph.Builder();
  const std::string name = "GatherShapePattern--" + gather.name().value();
  utils::RepeatedProtoField<NodeProto> replacements;
  if (graph.IsUsedMoreThanOnce(shape.output()[0].value())) {
    replacements.add() = shape;
  }

  if (is_scalar) {
    const std::string shape_out = builder.UniqueName(name + "_1d");
    NodeProto new_shape = MakeNode("Shape", {x_name}, {shape_out}, "", name.c_str());
    AddAttribute<int64_t>(new_shape, "start", new_start);
    AddAttribute<int64_t>(new_shape, "end", new_end);
    replacements.push_back(new_shape);

    const std::string axes_name = FreeInitializerName(builder, name + "_axes");
    builder.MakeInitializer(MakeInitializerShape(axes_name.c_str(), {0}));
    replacements.push_back(MakeNode("Squeeze", {shape_out, axes_name}, {gather.output()[0].value()},
                                    "", (name + "--squeeze").c_str()));
    return replacements;
  }

  NodeProto new_shape = MakeNode("Shape", {x_name}, {gather.output()[0].value()}, "", name.c_str());
  AddAttribute<int64_t>(new_shape, "start", new_start);
  AddAttribute<int64_t>(new_shape, "end", new_end);
  replacements.push_back(new_shape);
  return replacements;
}

} // namespace ONNX_LIGHT_NAMESPACE::onnx_patterns
