#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"

#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/ErrorHandling.h"
#include "llvm/Support/MathExtras.h"

#include <limits>

#include "AffineUtils.hpp"
#include "ConstantUtils.hpp"
#include "StaticIntSequence.hpp"

using namespace mlir;

namespace onnx_mlir {
namespace {

static bool getAffineValue(int64_t base, int64_t step, size_t index,
                           int64_t &value) {
  if (index > static_cast<size_t>(std::numeric_limits<int64_t>::max()))
    return false;
  int64_t scaled;
  return !llvm::MulOverflow(static_cast<int64_t>(index), step, scaled)
      && !llvm::AddOverflow(base, scaled, value);
}

static FailureOr<SmallVector<int64_t>> getI64Values(Operation *op,
                                                     StringRef name) {
  Attribute attr = op->getAttr(name);
  if (!attr)
    return op->emitOpError() << "is missing " << name << " metadata",
           failure();
  if (auto scalar = dyn_cast<IntegerAttr>(attr))
    return SmallVector<int64_t> {scalar.getInt()};
  if (auto array = dyn_cast<DenseI64ArrayAttr>(attr))
    return SmallVector<int64_t>(array.asArrayRef());
  auto elements = dyn_cast<DenseIntElementsAttr>(attr);
  auto type = elements ? dyn_cast<RankedTensorType>(elements.getType())
                       : RankedTensorType();
  if (!elements || !type || type.getRank() != 1
      || !type.getElementType().isInteger(64))
    return op->emitOpError() << "has invalid " << name << " metadata",
           failure();
  SmallVector<int64_t> values;
  values.reserve(elements.getNumElements());
  for (APInt value : elements.getValues<APInt>())
    values.push_back(value.getSExtValue());
  return values;
}

} // namespace

StaticIntSequence StaticIntSequence::uniform(int64_t value, size_t count) {
  assert(count != 0 && "empty static integer sequence");
  StaticIntSequence result;
  result.kind = StaticIntSequenceKind::Uniform;
  result.count = count;
  result.base = value;
  return result;
}

StaticIntSequence StaticIntSequence::affine(int64_t base, int64_t step,
                                             size_t count) {
  assert(count != 0 && "empty static integer sequence");
  int64_t last;
  assert(getAffineValue(base, step, count - 1, last)
         && "overflowing static affine sequence");
  if (count == 1 || step == 0)
    return uniform(base, count);
  StaticIntSequence result;
  result.kind = StaticIntSequenceKind::Affine;
  result.count = count;
  result.base = base;
  result.step = step;
  return result;
}

StaticIntSequence StaticIntSequence::runLengthEncoded(
  ArrayRef<int64_t> runs, size_t count) {
  assert(count != 0 && runs.size() % 2 == 0
         && "invalid run-length encoded sequence");
  StaticIntSequence result;
  result.kind = StaticIntSequenceKind::RunLengthEncoded;
  result.count = count;
  result.data.assign(runs);
  return result;
}

StaticIntSequence StaticIntSequence::fromValues(ArrayRef<int64_t> values) {
  assert(!values.empty() && "empty static integer sequence");
  if (llvm::all_equal(values))
    return uniform(values.front(), values.size());
  int64_t step;
  bool isAffine = !llvm::SubOverflow(values[1], values[0], step);
  for (size_t index = 1; isAffine && index < values.size(); ++index) {
    int64_t difference;
    isAffine = !llvm::SubOverflow(values[index], values[index - 1],
                                  difference)
        && difference == step;
  }
  if (isAffine)
    return affine(values.front(), step, values.size());

  SmallVector<int64_t> runs;
  for (int64_t value : values) {
    if (!runs.empty() && runs[runs.size() - 2] == value) {
      ++runs.back();
      continue;
    }
    runs.push_back(value);
    runs.push_back(1);
  }
  if (runs.size() < values.size())
    return runLengthEncoded(runs, values.size());
  StaticIntSequence result;
  result.kind = StaticIntSequenceKind::Dense;
  result.count = values.size();
  result.data.assign(values);
  return result;
}

int64_t StaticIntSequence::valueAt(size_t index) const {
  assert(index < count && "static integer sequence index out of range");
  if (kind == StaticIntSequenceKind::Uniform)
    return base;
  if (kind == StaticIntSequenceKind::Affine) {
    int64_t value;
    bool valid = getAffineValue(base, step, index, value);
    assert(valid && "overflowing static affine sequence");
    return value;
  }
  if (kind == StaticIntSequenceKind::Dense)
    return data[index];
  for (size_t run = 0; run < data.size(); run += 2) {
    size_t length = static_cast<size_t>(data[run + 1]);
    if (index < length)
      return data[run];
    index -= length;
  }
  llvm_unreachable("malformed run-length encoded sequence");
}

std::optional<size_t> StaticIntSequence::find(int64_t value, size_t begin,
                                              size_t length) const {
  assert(begin <= count && length <= count - begin
         && "invalid static integer sequence search");
  if (length == 0)
    return std::nullopt;
  size_t end = begin + length;
  if (kind == StaticIntSequenceKind::Uniform)
    return value == base ? std::optional<size_t>(begin) : std::nullopt;
  if (kind == StaticIntSequenceKind::Affine) {
    int64_t delta;
    if (llvm::SubOverflow(value, base, delta) || delta % step != 0)
      return std::nullopt;
    int64_t index = delta / step;
    return index >= static_cast<int64_t>(begin)
            && index < static_cast<int64_t>(end)
        ? std::optional<size_t>(index)
        : std::nullopt;
  }
  if (kind == StaticIntSequenceKind::Dense) {
    ArrayRef<int64_t> selected = ArrayRef(data).slice(begin, length);
    auto found = llvm::find(selected, value);
    return found == selected.end()
        ? std::nullopt
        : std::optional<size_t>(begin + (found - selected.begin()));
  }
  size_t runBegin = 0;
  for (size_t run = 0; run < data.size(); run += 2) {
    size_t runEnd = runBegin + static_cast<size_t>(data[run + 1]);
    if (runEnd > begin && runBegin < end && data[run] == value)
      return std::max(begin, runBegin);
    if (runBegin >= end)
      break;
    runBegin = runEnd;
  }
  return std::nullopt;
}

StaticIntSequence StaticIntSequence::slice(size_t begin, size_t length) const {
  assert(length != 0 && begin <= count - length && "invalid sequence slice");
  if (kind == StaticIntSequenceKind::Uniform)
    return uniform(base, length);
  if (kind == StaticIntSequenceKind::Affine)
    return affine(valueAt(begin), step, length);
  if (kind == StaticIntSequenceKind::Dense)
    return fromValues(ArrayRef(data).slice(begin, length));
  SmallVector<int64_t> runs;
  size_t end = begin + length;
  forEachEqualRun([&](int64_t value, size_t runBegin, size_t runCount) {
    size_t selectedBegin = std::max(begin, runBegin);
    size_t selectedEnd = std::min(end, runBegin + runCount);
    if (selectedBegin >= selectedEnd)
      return;
    if (!runs.empty() && runs[runs.size() - 2] == value)
      runs.back() += selectedEnd - selectedBegin;
    else {
      runs.push_back(value);
      runs.push_back(selectedEnd - selectedBegin);
    }
  });
  if (runs.size() == 2)
    return uniform(runs.front(), length);
  if (runs.size() < length)
    return runLengthEncoded(runs, length);
  SmallVector<int64_t> values;
  for (size_t run = 0; run < runs.size(); run += 2)
    values.append(runs[run + 1], runs[run]);
  return fromValues(values);
}

StaticIntSequence StaticIntSequence::remap(ArrayRef<unsigned> indices) const {
  assert(!indices.empty() && "empty static integer sequence remap");
  SmallVector<int64_t> values;
  values.reserve(indices.size());
  for (unsigned index : indices)
    values.push_back(valueAt(index));
  return fromValues(values);
}

bool StaticIntSequence::operator==(const StaticIntSequence& other) const {
  return kind == other.kind && count == other.count && base == other.base
      && step == other.step && data == other.data;
}

llvm::hash_code StaticIntSequence::hash() const {
  return llvm::hash_combine(kind, count, base, step,
                            llvm::hash_combine_range(data.begin(), data.end()));
}

void StaticIntSequence::forEachEqualRun(
  llvm::function_ref<void(int64_t, size_t, size_t)> callback) const {
  if (kind == StaticIntSequenceKind::Uniform) {
    callback(base, 0, count);
    return;
  }
  if (kind == StaticIntSequenceKind::RunLengthEncoded) {
    size_t begin = 0;
    for (size_t run = 0; run < data.size(); run += 2) {
      size_t runCount = static_cast<size_t>(data[run + 1]);
      callback(data[run], begin, runCount);
      begin += runCount;
    }
    return;
  }
  size_t begin = 0;
  while (begin < count) {
    int64_t value = valueAt(begin);
    size_t end = begin + 1;
    while (end < count && valueAt(end) == value)
      ++end;
    callback(value, begin, end - begin);
    begin = end;
  }
}

void StaticIntSequenceChain::append(const StaticIntSequence &sequence,
                                    size_t begin, size_t length) {
  assert(length != 0 && begin <= sequence.size() - length
         && "invalid static integer sequence chain slice");
  if (!slices.empty()) {
    StaticIntSequenceSlice &last = slices.back();
    if (last.sequence == &sequence && last.begin + last.count == begin) {
      last.count += length;
      count += length;
      return;
    }
    auto affinePart = [](const StaticIntSequenceSlice &slice,
                         int64_t &base, int64_t &step) {
      base = slice.sequence->valueAt(slice.begin);
      if (slice.count == 1) {
        step = 0;
        return true;
      }
      return !llvm::SubOverflow(slice.sequence->valueAt(slice.begin + 1),
                                base, step)
          && (slice.sequence->kind == StaticIntSequenceKind::Uniform
              || slice.sequence->kind == StaticIntSequenceKind::Affine);
    };
    StaticIntSequenceSlice next {&sequence, begin, length};
    int64_t leftBase, leftStep, rightBase, rightStep, expected;
    if (affinePart(last, leftBase, leftStep)
        && affinePart(next, rightBase, rightStep)
        && (last.count == 1 || length == 1 || leftStep == rightStep)) {
      int64_t step = last.count == 1 ? rightStep : leftStep;
      if (getAffineValue(leftBase, step, last.count, expected)
          && expected == rightBase) {
        owned.push_back(std::make_unique<StaticIntSequence>(
          StaticIntSequence::affine(leftBase, step, last.count + length)));
        last = {owned.back().get(), 0, last.count + length};
        count += length;
        return;
      }
    }
  }
  slices.push_back({&sequence, begin, length});
  count += length;
}

void StaticIntSequenceChain::append(StaticIntSequence sequence) {
  size_t length = sequence.size();
  owned.push_back(std::make_unique<StaticIntSequence>(std::move(sequence)));
  append(*owned.back(), 0, length);
}

int64_t StaticIntSequenceChain::valueAt(size_t index) const {
  assert(index < count && "static integer sequence chain index out of range");
  for (const StaticIntSequenceSlice &slice : slices) {
    if (index < slice.count)
      return slice.sequence->valueAt(slice.begin + index);
    index -= slice.count;
  }
  llvm_unreachable("malformed static integer sequence chain");
}

void StaticIntSequenceChain::forEachSegment(llvm::function_ref<void(
  const StaticIntSequence &, size_t, size_t)> callback) const {
  for (const StaticIntSequenceSlice &slice : slices)
    callback(*slice.sequence, slice.begin, slice.count);
}

void StaticIntSequenceChain::forEachEqualRun(
  llvm::function_ref<void(int64_t, size_t, size_t)> callback) const {
  std::optional<int64_t> pendingValue;
  size_t pendingBegin = 0, pendingCount = 0, chainBegin = 0;
  auto flush = [&] {
    if (pendingValue)
      callback(*pendingValue, pendingBegin, pendingCount);
  };
  for (const StaticIntSequenceSlice &slice : slices) {
    size_t sliceEnd = slice.begin + slice.count;
    slice.sequence->forEachEqualRun(
      [&](int64_t value, size_t runBegin, size_t runCount) {
        size_t begin = std::max(slice.begin, runBegin);
        size_t end = std::min(sliceEnd, runBegin + runCount);
        if (begin >= end)
          return;
        size_t selectedCount = end - begin;
        size_t globalBegin = chainBegin + begin - slice.begin;
        if (pendingValue && *pendingValue == value
            && pendingBegin + pendingCount == globalBegin) {
          pendingCount += selectedCount;
          return;
        }
        flush();
        pendingValue = value;
        pendingBegin = globalBegin;
        pendingCount = selectedCount;
      });
    chainBegin += slice.count;
  }
  flush();
}

StaticIntSequence StaticIntSequenceChain::canonicalize() const {
  assert(count != 0 && "empty static integer sequence chain");
  int64_t first = valueAt(0);
  bool uniform = true;
  forEachEqualRun([&](int64_t value, size_t, size_t) {
    uniform &= value == first;
  });
  if (uniform)
    return StaticIntSequence::uniform(first, count);

  int64_t step = 0, previous = first;
  bool affine = true, haveStep = false;
  size_t position = 0;
  forEachSegment([&](const StaticIntSequence &sequence, size_t begin,
                     size_t length) {
    if (!affine)
      return;
    for (size_t index = 0; index < length; ++index) {
      int64_t value = sequence.valueAt(begin + index);
      if (position++ == 0) {
        previous = value;
        continue;
      }
      if (!haveStep) {
        affine = !llvm::SubOverflow(value, previous, step);
        haveStep = true;
      } else if (haveStep) {
        int64_t difference;
        affine = !llvm::SubOverflow(value, previous, difference)
            && difference == step;
      }
      previous = value;
      if (!affine)
        break;
    }
  });
  if (affine && haveStep)
    return StaticIntSequence::affine(first, step, count);

  SmallVector<int64_t> runs;
  forEachEqualRun([&](int64_t value, size_t, size_t runCount) {
    runs.push_back(value);
    runs.push_back(runCount);
  });
  if (runs.size() < count)
    return StaticIntSequence::runLengthEncoded(runs, count);
  SmallVector<int64_t> values;
  values.reserve(count);
  for (size_t run = 0; run < runs.size(); run += 2)
    values.append(runs[run + 1], runs[run]);
  return StaticIntSequence::fromValues(values);
}

int64_t StaticIntSequenceChainCursor::value() const {
  assert(!done() && "static integer sequence chain cursor is done");
  const StaticIntSequenceSlice &current = chain.slices[slice];
  return current.sequence->valueAt(current.begin + offset);
}

void StaticIntSequenceChainCursor::advance() {
  assert(!done() && "static integer sequence chain cursor is done");
  if (++offset != chain.slices[slice].count)
    return;
  offset = 0;
  ++slice;
}

void setStaticIntSequenceAttr(Operation *op, StringRef name,
                              const StaticIntSequence &sequence,
                              size_t logicalCount) {
  assert(sequence.size() == logicalCount && logicalCount != 0
         && "invalid static integer metadata count");
  SmallVector<int64_t> values;
  StringRef encoding;
  switch (sequence.kind) {
  case StaticIntSequenceKind::Uniform:
    encoding = "uniform";
    values.push_back(sequence.base);
    break;
  case StaticIntSequenceKind::Affine:
    encoding = "affine";
    values = {sequence.base, sequence.step};
    break;
  case StaticIntSequenceKind::RunLengthEncoded:
    encoding = "rle";
    values = sequence.data;
    break;
  case StaticIntSequenceKind::Dense:
    encoding = "dense";
    values = sequence.data;
    break;
  }
  OpBuilder builder(op);
  auto type = RankedTensorType::get(
    {static_cast<int64_t>(values.size())}, builder.getI64Type());
  op->setAttr(name, DenseIntElementsAttr::get(type, values));
  if (sequence.kind != StaticIntSequenceKind::Dense)
    op->setAttr((name + "_encoding").str(), builder.getStringAttr(encoding));
}

FailureOr<StaticIntSequence> getStaticIntSequenceAttr(
  Operation *op, StringRef name, size_t logicalCount) {
  if (logicalCount == 0)
    return op->emitOpError() << "has zero logical count for " << name,
           failure();
  auto values = getI64Values(op, name);
  if (failed(values))
    return failure();
  auto encoding = op->getAttrOfType<StringAttr>((name + "_encoding").str());
  if (!encoding) {
    if (values->size() != logicalCount)
      return op->emitOpError() << "has invalid dense " << name << " count",
             failure();
    return StaticIntSequence::fromValues(*values);
  }
  if (encoding.getValue() == "uniform") {
    if (values->size() != 1)
      return op->emitOpError() << "has invalid uniform " << name,
             failure();
    return StaticIntSequence::uniform(values->front(), logicalCount);
  }
  if (encoding.getValue() == "affine") {
    int64_t last;
    if (values->size() != 2
        || !getAffineValue((*values)[0], (*values)[1], logicalCount - 1, last))
      return op->emitOpError() << "has invalid affine " << name,
             failure();
    return StaticIntSequence::affine((*values)[0], (*values)[1], logicalCount);
  }
  if (encoding.getValue() == "rle") {
    size_t count = 0;
    if (values->empty() || values->size() % 2 != 0)
      return op->emitOpError() << "has invalid RLE " << name, failure();
    for (size_t index = 1; index < values->size(); index += 2) {
      if ((*values)[index] <= 0
          || static_cast<uint64_t>((*values)[index]) > logicalCount - count)
        return op->emitOpError() << "has invalid RLE " << name, failure();
      count += (*values)[index];
    }
    if (count != logicalCount)
      return op->emitOpError() << "has mismatched RLE " << name << " count",
             failure();
    return StaticIntSequence::runLengthEncoded(*values, count);
  }
  if (encoding.getValue() == "dense") {
    if (values->size() != logicalCount)
      return op->emitOpError() << "has invalid dense " << name << " count",
             failure();
    return StaticIntSequence::fromValues(*values);
  }
  return op->emitOpError() << "has unknown " << name << " encoding",
         failure();
}

Value emitStaticIntLookup(const StaticIntSequence& sequence, Value position,
                          Operation* constantAnchor,
                          ConstantPool& constants, OpBuilder& builder,
                          Location loc) {
  if (sequence.getKind() == StaticIntSequenceKind::Uniform)
    return constants.getIndex(sequence.valueAt(0));
  if (sequence.getKind() == StaticIntSequenceKind::Affine) {
    Value scaled = affineMulConst(builder, loc, position,
                                  sequence.valueAt(1) - sequence.valueAt(0),
                                  constantAnchor);
    return affineAddConst(builder, loc, scaled, sequence.valueAt(0),
                          constantAnchor);
  }
  SmallVector<int64_t> values;
  values.reserve(sequence.size());
  sequence.forEachEqualRun([&](int64_t value, size_t, size_t count) {
    values.append(count, value);
  });
  auto type = RankedTensorType::get(
    {static_cast<int64_t>(values.size())}, builder.getI64Type());
  Value table = constants.get(type,
    DenseElementsAttr::get(type, ArrayRef<int64_t>(values)));
  Value selected = tensor::ExtractOp::create(
    builder, loc, table, ValueRange {position});
  return arith::IndexCastOp::create(
    builder, loc, builder.getIndexType(), selected);
}

} // namespace onnx_mlir