#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinTypes.h"

#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallBitVector.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/ADT/SmallVector.h"

#include <utility>

#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
#include "src/Dialect/ONNX/ONNXOps.hpp"

using namespace mlir;

namespace onnx_mlir {

namespace {

static bool hasStaticUnitStrides(tensor::ExtractSliceOp extractSliceOp) {
  return llvm::all_of(extractSliceOp.getStaticStrides(), [](int64_t stride) { return stride == 1; });
}

static bool hasConstantIndices(tensor::ExtractOp extractOp) {
  return llvm::all_of(extractOp.getIndices(), [](Value index) { return matchConstantIndexValue(index).has_value(); });
}

static bool isStaticTensorResult(Operation* op) {
  return llvm::all_of(op->getResultTypes(), [](Type type) {
    auto shapedType = dyn_cast<ShapedType>(type);
    return shapedType && shapedType.hasStaticShape();
  });
}

static FailureOr<DenseElementsAttr> transposeDenseElements(DenseElementsAttr denseAttr, ArrayRef<int64_t> perms) {
  auto tensorType = dyn_cast<RankedTensorType>(denseAttr.getType());
  if (!tensorType)
    return failure();

  int64_t rank = tensorType.getRank();
  if (static_cast<int64_t>(perms.size()) != rank)
    return failure();

  llvm::SmallBitVector seen(rank);
  SmallVector<int64_t> transposedShape;
  transposedShape.reserve(rank);
  for (int64_t perm : perms) {
    if (perm < 0 || perm >= rank || seen.test(perm))
      return failure();
    seen.set(perm);
    transposedShape.push_back(tensorType.getShape()[perm]);
  }

  auto transposedType = RankedTensorType::get(transposedShape, tensorType.getElementType(), tensorType.getEncoding());
  if (denseAttr.isSplat())
    return DenseElementsAttr::get(transposedType, denseAttr.getSplatValue<Attribute>());

  SmallVector<Attribute> originalValues(denseAttr.getValues<Attribute>());
  SmallVector<Attribute> transposedValues(originalValues.size());
  SmallVector<int64_t> originalStrides = computeRowMajorStrides(tensorType.getShape());
  SmallVector<int64_t> transposedStrides = computeRowMajorStrides(transposedShape);
  SmallVector<int64_t> originalIndices(rank);

  for (auto [linearIndex, value] : llvm::enumerate(originalValues)) {
    int64_t remaining = static_cast<int64_t>(linearIndex);
    for (int64_t dim = 0; dim < rank; ++dim) {
      originalIndices[dim] = remaining / originalStrides[dim];
      remaining %= originalStrides[dim];
    }

    int64_t transposedLinearIndex = 0;
    for (int64_t dim = 0; dim < rank; ++dim)
      transposedLinearIndex += originalIndices[perms[dim]] * transposedStrides[dim];

    transposedValues[transposedLinearIndex] = value;
  }

  return DenseElementsAttr::get(transposedType, transposedValues);
}

static FailureOr<DenseElementsAttr> reshapeDenseElements(DenseElementsAttr denseAttr, RankedTensorType resultType) {
  auto sourceType = dyn_cast<RankedTensorType>(denseAttr.getType());
  if (!sourceType || !resultType || sourceType.getNumElements() != resultType.getNumElements())
    return failure();

  if (denseAttr.isSplat())
    return DenseElementsAttr::get(resultType, denseAttr.getSplatValue<Attribute>());

  SmallVector<Attribute> values(denseAttr.getValues<Attribute>());
  return DenseElementsAttr::get(resultType, values);
}

static FailureOr<DenseElementsAttr> extractSliceDenseElements(DenseElementsAttr denseAttr,
                                                              tensor::ExtractSliceOp extractSliceOp) {
  auto sourceType = dyn_cast<RankedTensorType>(denseAttr.getType());
  auto resultType = dyn_cast<RankedTensorType>(extractSliceOp.getType());
  if (!sourceType || !resultType || !sourceType.hasStaticShape() || !resultType.hasStaticShape())
    return failure();

  ArrayRef<int64_t> offsets = extractSliceOp.getStaticOffsets();
  ArrayRef<int64_t> sizes = extractSliceOp.getStaticSizes();
  ArrayRef<int64_t> strides = extractSliceOp.getStaticStrides();
  if (llvm::any_of(offsets, [](int64_t value) { return ShapedType::isDynamic(value); })
      || llvm::any_of(sizes, [](int64_t value) { return ShapedType::isDynamic(value); })
      || llvm::any_of(strides, [](int64_t stride) { return ShapedType::isDynamic(stride) || stride != 1; }))
    return failure();

  if (denseAttr.isSplat())
    return DenseElementsAttr::get(resultType, denseAttr.getSplatValue<Attribute>());

  SmallVector<Attribute> sourceValues(denseAttr.getValues<Attribute>());
  SmallVector<int64_t> sourceStrides = computeRowMajorStrides(sourceType.getShape());
  SmallVector<int64_t> resultStrides = computeRowMajorStrides(resultType.getShape());
  SmallVector<Attribute> resultValues;
  resultValues.reserve(resultType.getNumElements());

  for (int64_t linearIndex = 0; linearIndex < resultType.getNumElements(); ++linearIndex) {
    int64_t remaining = linearIndex;
    int64_t sourceLinearIndex = 0;
    for (int64_t dim = 0; dim < resultType.getRank(); ++dim) {
      const int64_t resultIndex = resultStrides.empty() ? 0 : remaining / resultStrides[dim];
      remaining = resultStrides.empty() ? 0 : remaining % resultStrides[dim];
      sourceLinearIndex += (offsets[dim] + resultIndex) * sourceStrides[dim];
    }
    resultValues.push_back(sourceValues[sourceLinearIndex]);
  }

  return DenseElementsAttr::get(resultType, resultValues);
}

static DenseElementsAttr getDirectDenseConstantAttr(Value value) {
  if (auto constantOp = value.getDefiningOp<arith::ConstantOp>())
    return dyn_cast<DenseElementsAttr>(constantOp.getValue());
  if (auto constantOp = value.getDefiningOp<ONNXConstantOp>())
    return dyn_cast_or_null<DenseElementsAttr>(constantOp.getValueAttr());
  return nullptr;
}

static DenseElementsAttr getHostConstantDenseElementsAttrImpl(Value value, llvm::SmallPtrSetImpl<Operation*>& visited) {
  auto* definingOp = value.getDefiningOp();
  if (!definingOp || !visited.insert(definingOp).second)
    return nullptr;

  // Rebuild dense attributes through view-only host-foldable chains so later
  // lowering stages can still recognize grouped/sliced constants.
  if (auto denseAttr = getDirectDenseConstantAttr(value))
    return denseAttr;

  if (auto transposeOp = dyn_cast<ONNXTransposeOp>(definingOp)) {
    auto inputAttr = getHostConstantDenseElementsAttrImpl(transposeOp.getData(), visited);
    if (!inputAttr)
      return nullptr;

    SmallVector<int64_t> perm;
    perm.reserve(transposeOp.getPermAttr().size());
    for (IntegerAttr attr : transposeOp.getPermAttr().getAsRange<IntegerAttr>())
      perm.push_back(attr.getInt());
    auto transposedAttr = transposeDenseElements(inputAttr, perm);
    return succeeded(transposedAttr) ? *transposedAttr : nullptr;
  }

  if (auto transposeOp = dyn_cast<linalg::TransposeOp>(definingOp)) {
    auto inputAttr = getHostConstantDenseElementsAttrImpl(transposeOp.getInput(), visited);
    if (!inputAttr)
      return nullptr;

    SmallVector<int64_t> perm(transposeOp.getPermutation().begin(), transposeOp.getPermutation().end());
    auto transposedAttr = transposeDenseElements(inputAttr, perm);
    return succeeded(transposedAttr) ? *transposedAttr : nullptr;
  }

  if (auto collapseShapeOp = dyn_cast<tensor::CollapseShapeOp>(definingOp)) {
    auto inputAttr = getHostConstantDenseElementsAttrImpl(collapseShapeOp.getSrc(), visited);
    if (!inputAttr)
      return nullptr;
    auto reshapedAttr = reshapeDenseElements(inputAttr, cast<RankedTensorType>(collapseShapeOp.getType()));
    return succeeded(reshapedAttr) ? *reshapedAttr : nullptr;
  }

  if (auto expandShapeOp = dyn_cast<tensor::ExpandShapeOp>(definingOp)) {
    auto inputAttr = getHostConstantDenseElementsAttrImpl(expandShapeOp.getSrc(), visited);
    if (!inputAttr)
      return nullptr;
    auto reshapedAttr = reshapeDenseElements(inputAttr, cast<RankedTensorType>(expandShapeOp.getType()));
    return succeeded(reshapedAttr) ? *reshapedAttr : nullptr;
  }

  if (auto extractSliceOp = dyn_cast<tensor::ExtractSliceOp>(definingOp)) {
    auto inputAttr = getHostConstantDenseElementsAttrImpl(extractSliceOp.getSource(), visited);
    if (!inputAttr)
      return nullptr;
    auto slicedAttr = extractSliceDenseElements(inputAttr, extractSliceOp);
    return succeeded(slicedAttr) ? *slicedAttr : nullptr;
  }

  return nullptr;
}

static std::optional<CompileTimeSource>
getCompileTimeSourceImpl(Operation* op, llvm::SmallPtrSetImpl<Operation*>& visited, size_t chainLength = 0) {
  if (!op)
    return std::nullopt;

  if (!visited.insert(op).second)
    return {
      {op, chainLength}
    };

  if (isa<arith::ConstantOp, ONNXConstantOp, ONNXNoneOp>(op))
    return {
      {op, chainLength}
    };

  chainLength += 1;

  if (auto extractOp = dyn_cast<tensor::ExtractOp>(op))
    return hasConstantIndices(extractOp)
           ? getCompileTimeSourceImpl(extractOp.getTensor().getDefiningOp(), visited, chainLength)
           : std::nullopt;

  if (!isStaticTensorResult(op))
    return std::nullopt;

  if (auto transposeOp = dyn_cast<ONNXTransposeOp>(op))
    return getCompileTimeSourceImpl(transposeOp.getData().getDefiningOp(), visited, chainLength);

  if (auto transposeOp = dyn_cast<linalg::TransposeOp>(op))
    return getCompileTimeSourceImpl(transposeOp.getInput().getDefiningOp(), visited, chainLength);

  if (auto collapseShapeOp = dyn_cast<tensor::CollapseShapeOp>(op))
    return getCompileTimeSourceImpl(collapseShapeOp.getSrc().getDefiningOp(), visited, chainLength);

  if (auto expandShapeOp = dyn_cast<tensor::ExpandShapeOp>(op))
    return getCompileTimeSourceImpl(expandShapeOp.getSrc().getDefiningOp(), visited, chainLength);

  if (auto extractSliceOp = dyn_cast<tensor::ExtractSliceOp>(op))
    return hasStaticUnitStrides(extractSliceOp)
           ? getCompileTimeSourceImpl(extractSliceOp.getSource().getDefiningOp(), visited, chainLength)
           : std::nullopt;

  if (auto splatOp = dyn_cast<tensor::SplatOp>(op))
    return getCompileTimeSourceImpl(splatOp.getInput().getDefiningOp(), visited, chainLength);

  if (auto extractRowsOp = dyn_cast<spatial::SpatExtractRowsOp>(op))
    return getCompileTimeSourceImpl(extractRowsOp.getInput().getDefiningOp(), visited, chainLength);

  if (auto concatOp = dyn_cast<spatial::SpatConcatOp>(op)) {
    std::optional<CompileTimeSource> res = {};
    for (auto operandValue : concatOp.getOperands()) {
      auto partialRes = getCompileTimeSourceImpl(operandValue.getDefiningOp(), visited, chainLength);
      if (!partialRes)
        return std::nullopt;

      if (!res) {
        res = partialRes;
        continue;
      }
      if (res->chainLength < partialRes->chainLength)
        res = partialRes;
    }
    return res;
  }

  return std::nullopt;
}

} // namespace

std::optional<CompileTimeSource> getCompileTimeSource(Operation* op) {
  llvm::SmallPtrSet<Operation*, 8> visited;
  return getCompileTimeSourceImpl(op, visited);
}

bool isCompileTimeComputable(Value value) {
  auto* definingOp = value.getDefiningOp();
  if (!definingOp)
    return false;

  llvm::SmallPtrSet<Operation*, 8> visited;
  return getCompileTimeSourceImpl(definingOp, visited).has_value();
}

bool isCompileTimeOp(Operation* op) {
  llvm::SmallPtrSet<Operation*, 8> visited;
  return getCompileTimeSourceImpl(op, visited).has_value();
}

DenseElementsAttr getHostConstDenseElementsAttr(Value value) {
  llvm::SmallPtrSet<Operation*, 8> visited;
  return getHostConstantDenseElementsAttrImpl(value, visited);
}

} // namespace onnx_mlir