#include "Common.hpp"
#include "src/Accelerators/PIM/Common/PimCommon.hpp"

using namespace mlir;

namespace onnx_mlir {

Value stripMemRefCasts(Value value) {
  while (auto castOp = value.getDefiningOp<memref::CastOp>())
    value = castOp.getSource();
  return value;
}

Value stripMemRefViewOps(Value value) {
  while (true) {
    if (auto castOp = value.getDefiningOp<memref::CastOp>()) {
      value = castOp.getSource();
      continue;
    }
    if (auto collapseOp = value.getDefiningOp<memref::CollapseShapeOp>()) {
      value = collapseOp.getSrc();
      continue;
    }
    if (auto expandOp = value.getDefiningOp<memref::ExpandShapeOp>()) {
      value = expandOp.getSrc();
      continue;
    }
    return value;
  }
}

memref::GlobalOp createFoldedGlobal(ModuleOp moduleOp,
                                    Location loc,
                                    MemRefType globalType,
                                    DenseElementsAttr denseAttr,
                                    StringRef nameStem,
                                    IntegerAttr alignment) {
  auto globalName = nameStem.str();
  unsigned suffix = 0;
  while (moduleOp.lookupSymbol(globalName))
    globalName = (nameStem + "_" + std::to_string(++suffix)).str();

  auto visibility = StringAttr::get(moduleOp.getContext(), "private");
  OpBuilder moduleBuilder(moduleOp.getBodyRegion());
  moduleBuilder.setInsertionPointToStart(moduleOp.getBody());
  return memref::GlobalOp::create(moduleBuilder,
                                  loc,
                                  globalName,
                                  visibility,
                                  globalType,
                                  denseAttr,
                                  /*constant=*/true,
                                  alignment);
}

FailureOr<DenseElementsAttr> getDenseGlobalValue(ModuleOp moduleOp, Value value) {
  value = stripMemRefCasts(value);

  auto getGlobalOp = value.getDefiningOp<memref::GetGlobalOp>();
  if (!getGlobalOp)
    return failure();

  auto globalOp = lookupGlobalForGetGlobal(moduleOp, getGlobalOp);
  if (!globalOp || !globalOp.getConstant() || !globalOp.getInitialValue())
    return failure();

  auto denseAttr = dyn_cast<DenseElementsAttr>(*globalOp.getInitialValue());
  if (!denseAttr)
    return failure();
  return denseAttr;
}

FailureOr<StaticSubviewInfo> getStaticSubviewInfo(Value value) {
  value = stripMemRefViewOps(value);
  auto subviewOp = value.getDefiningOp<memref::SubViewOp>();
  if (!subviewOp)
    return failure();

  auto source = stripMemRefCasts(subviewOp.getSource());
  auto sourceType = dyn_cast<MemRefType>(source.getType());
  auto subviewType = dyn_cast<MemRefType>(subviewOp.getType());
  if (!sourceType || !subviewType || !sourceType.hasStaticShape() || !subviewType.hasStaticShape())
    return failure();

  StaticSubviewInfo info;
  info.source = source;
  info.sourceShape.assign(sourceType.getShape().begin(), sourceType.getShape().end());
  for (OpFoldResult offset : subviewOp.getMixedOffsets()) {
    auto staticOffset = getConstantIntValue(offset);
    if (!staticOffset)
      return failure();
    info.offsets.push_back(*staticOffset);
  }
  for (OpFoldResult size : subviewOp.getMixedSizes()) {
    auto staticSize = getConstantIntValue(size);
    if (!staticSize)
      return failure();
    info.sizes.push_back(*staticSize);
  }
  for (OpFoldResult stride : subviewOp.getMixedStrides()) {
    auto staticStride = getConstantIntValue(stride);
    if (!staticStride)
      return failure();
    info.strides.push_back(*staticStride);
  }
  return info;
}

int64_t
getSubviewChunkOffsetBytes(const StaticSubviewInfo& info, ArrayRef<int64_t> outerIndices, int64_t elementByteWidth) {
  SmallVector<int64_t> sourceIndices;
  sourceIndices.reserve(info.sourceShape.size());
  for (size_t dim = 0; dim + 1 < info.sourceShape.size(); ++dim)
    sourceIndices.push_back(info.offsets[dim] + outerIndices[dim] * info.strides[dim]);
  sourceIndices.push_back(info.offsets.back());
  return linearizeIndex(sourceIndices, computeRowMajorStrides(info.sourceShape)) * elementByteWidth;
}

} // namespace onnx_mlir
