#include "lib/Transforms/ConvertToCiphertextSemantics/AssignLayout.h"

#include <cassert>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <functional>
#include <string>
#include <utility>
#include <vector>

#include "lib/Dialect/TensorExt/IR/TensorExtAttributes.h"
#include "lib/Dialect/TensorExt/IR/TensorExtOps.h"
#include "lib/Transforms/ConvertToCiphertextSemantics/TypeConversion.h"
#include "lib/Utils/Layout/Codegen.h"
#include "lib/Utils/Layout/IslConversion.h"
#include "lib/Utils/Layout/Utils.h"
#include "llvm/include/llvm/ADT/BitVector.h"          // from @llvm-project
#include "llvm/include/llvm/ADT/DynamicAPInt.h"       // from @llvm-project
#include "llvm/include/llvm/ADT/STLExtras.h"          // from @llvm-project
#include "llvm/include/llvm/ADT/SmallVector.h"        // from @llvm-project
#include "llvm/include/llvm/ADT/SmallVectorExtras.h"  // from @llvm-project
#include "llvm/include/llvm/Support/Debug.h"          // from @llvm-project
#include "llvm/include/llvm/Support/raw_ostream.h"    // from @llvm-project
#include "mlir/include/mlir/Analysis/FlatLinearValueConstraints.h"  // from @llvm-project
#include "mlir/include/mlir/Analysis/Presburger/IntegerRelation.h"  // from @llvm-project
#include "mlir/include/mlir/Analysis/Presburger/PresburgerSpace.h"  // from @llvm-project
#include "mlir/include/mlir/Dialect/Arith/IR/Arith.h"    // from @llvm-project
#include "mlir/include/mlir/Dialect/SCF/IR/SCF.h"        // from @llvm-project
#include "mlir/include/mlir/Dialect/Tensor/IR/Tensor.h"  // from @llvm-project
#include "mlir/include/mlir/Dialect/Utils/StructuredOpsUtils.h"  // from @llvm-project
#include "mlir/include/mlir/IR/AsmState.h"  // from @llvm-project
#include "mlir/include/mlir/IR/BuiltinAttributeInterfaces.h"  // from @llvm-project
#include "mlir/include/mlir/IR/BuiltinAttributes.h"  // from @llvm-project
#include "mlir/include/mlir/IR/BuiltinTypes.h"       // from @llvm-project
#include "mlir/include/mlir/IR/DialectResourceBlobManager.h"  // from @llvm-project
#include "mlir/include/mlir/IR/Location.h"       // from @llvm-project
#include "mlir/include/mlir/IR/Matchers.h"       // from @llvm-project
#include "mlir/include/mlir/IR/OpDefinition.h"   // from @llvm-project
#include "mlir/include/mlir/IR/Operation.h"      // from @llvm-project
#include "mlir/include/mlir/IR/TypeUtilities.h"  // from @llvm-project
#include "mlir/include/mlir/IR/Value.h"          // from @llvm-project
#include "mlir/include/mlir/IR/ValueRange.h"     // from @llvm-project
#include "mlir/include/mlir/Support/LLVM.h"      // from @llvm-project

#define DEBUG_TYPE "convert-to-ciphertext-semantics"

namespace mlir {
namespace heir {

using ::mlir::presburger::IntegerRelation;
using tensor_ext::LayoutAttr;

static std::string printRelation(const IntegerRelation& rel) {
  std::string str;
  llvm::raw_string_ostream os(str);
  rel.print(os);
  return str;
}

static FailureOr<Value> implementUnpackOpStep(
    Value input, LayoutAttr layout, Type targetType,
    ImplicitLocOpBuilder& builder,
    const std::function<void(Operation*)>& createdOpCallback) {
  IntegerRelation rel = layout.getIntegerRelation();
  LLVM_DEBUG(llvm::dbgs() << "implementUnpackOpStep with layout: " << layout
                          << " targetType: " << targetType << "\n");

  RankedTensorType unpackedTensorType = dyn_cast<RankedTensorType>(targetType);

  // Restrict the layout relation's domain bounds to the valid elements
  // [0, dimSize - 1] of targetType. This has to happen before any use of `rel`
  // below, including the single-element shortcut: that shortcut samples an
  // arbitrary point from the relation's range, and with an unrestricted domain
  // the sample can be the slot belonging to a different element.
  if (unpackedTensorType) {
    for (unsigned i = 0; i < unpackedTensorType.getRank(); ++i) {
      if (unpackedTensorType.isDynamicDim(i)) continue;
      rel.addBound(presburger::BoundType::UB,
                   rel.getVarKindOffset(presburger::VarKind::Domain) + i,
                   unpackedTensorType.getDimSize(i) - 1);
    }
  }

  bool isSingleElementTensor = unpackedTensorType &&
                               unpackedTensorType.hasStaticShape() &&
                               unpackedTensorType.getNumElements() == 1;
  bool isScalar = !unpackedTensorType;

  if (isScalar || isSingleElementTensor) {
    // Extract the lone element from any slot in the mapping.
    std::vector<int64_t> point = anyRangePoint(rel);
    if (point.empty()) {
      return builder.emitError()
             << "Failed to find any point in the range of relation "
             << printRelation(rel);
    }

    SmallVector<Value> indices;
    for (int64_t idx : point) {
      indices.push_back(arith::ConstantIndexOp::create(builder, idx));
    }
    auto extractOp = tensor::ExtractOp::create(builder, input, indices);
    createdOpCallback(extractOp);

    if (isScalar) {
      return extractOp.getResult();
    }

    // isSingleElementTensor -> wrap back into single-element tensor.
    auto fromElementsOp = tensor::FromElementsOp::create(
        builder, unpackedTensorType, ValueRange{extractOp.getResult()});
    createdOpCallback(fromElementsOp);
    return fromElementsOp.getResult();
  }

  SmallVector<int> domainSchedule;
  for (unsigned i = 0; i < rel.getNumDomainVars(); ++i) {
    domainSchedule.push_back(i);
  }

  MLIRLoopNestGenerator generator(builder);
  auto loopNestCstr = generateLoopNestAsCStr(rel, domainSchedule);
  if (failed(loopNestCstr)) {
    return builder.emitError() << "Failed to generate loop nest for relation "
                               << printRelation(rel);
  }
  LLVM_DEBUG(llvm::dbgs() << "Generating loop nest assignment for relation "
                          << loopNestCstr.value() << "\n");

  TypedValue<RankedTensorType> ciphertextTensor =
      cast<TypedValue<RankedTensorType>>(input);
  RankedTensorType dataSemanticType = unpackedTensorType;

  auto dataTensor = cast<TypedValue<RankedTensorType>>(
      mlir::arith::ConstantOp::create(builder,
                                      builder.getZeroAttr(dataSemanticType))
          .getResult());
  auto loop = generator.generateForLoop(
      rel, {dataTensor},
      [&](mlir::OpBuilder& builder, Location loc, ValueRange exprs,
          ValueRange iterArgs) {
        SmallVector<Value> extractIndices;
        for (Value idx : exprs.drop_front(dataSemanticType.getRank())) {
          extractIndices.push_back(arith::IndexCastOp::create(
              builder, loc, builder.getIndexType(), idx));
        }
        // Extract from ciphertext and insert into dataTensor
        auto extracted = tensor::ExtractOp::create(
            builder, loc, ciphertextTensor, extractIndices);
        SmallVector<Value> insertIndices;
        for (Value idx :
             exprs.drop_back(ciphertextTensor.getType().getRank())) {
          insertIndices.push_back(arith::IndexCastOp::create(
              builder, loc, builder.getIndexType(), idx));
        }
        auto inserted = tensor::InsertOp::create(builder, loc, extracted,
                                                 iterArgs[0], insertIndices);
        return scf::ValueVector({inserted});
      },
      domainSchedule,
      /*reverse=*/true);
  if (failed(loop)) {
    return builder.emitError() << "Failed to generate loop nest for relation "
                               << printRelation(rel);
  }
  createdOpCallback(loop.value());
  return loop.value().getResults()[0];
}

static FailureOr<Value> implementAssignLayoutPermutation(
    Value input, DenseIntElementsAttr permutation, Type targetTypeTy,
    ImplicitLocOpBuilder& builder,
    const std::function<void(Operation*)>& createdOpCallback) {
  auto tensorType = dyn_cast<RankedTensorType>(input.getType());
  if (!tensorType)
    return builder.emitError() << "Permutation layout requires tensor input";

  if (tensorType.getRank() > 2)
    return builder.emitError() << "Arbitrary permutation layout only "
                                  "supports up to 2D tensors for now";

  RankedTensorType targetType = cast<RankedTensorType>(targetTypeTy);
  auto zeroCtxt = arith::ConstantOp::create(builder, targetType,
                                            builder.getZeroAttr(targetType));
  createdOpCallback(zeroCtxt);
  Value result = zeroCtxt.getResult();

  int64_t ctBound = (tensorType.getRank() == 2) ? tensorType.getDimSize(0) : 1;
  int64_t srcSlotBound = tensorType.getDimSize(tensorType.getRank() - 1);
  int64_t dstSlotBound = targetType.getDimSize(targetType.getRank() - 1);

  for (auto it = permutation.value_begin<APInt>();
       it != permutation.value_end<APInt>();) {
    int64_t srcCt = (*it++).getSExtValue();
    int64_t srcSlot = (*it++).getSExtValue();
    int64_t dstCt = (*it++).getSExtValue();
    int64_t dstSlot = (*it++).getSExtValue();

    if (srcCt >= ctBound || srcSlot >= srcSlotBound || dstCt >= ctBound ||
        dstSlot >= dstSlotBound) {
      return builder.emitError()
             << "Permutation index out of bounds: " << "src_ct=" << srcCt
             << ", src_slot=" << srcSlot << " (input bounds: ct < " << ctBound
             << ", slot < " << srcSlotBound << "); " << "dst_ct=" << dstCt
             << ", dst_slot=" << dstSlot << " (target bounds: ct < " << ctBound
             << ", slot < " << dstSlotBound << ")";
    }

    SmallVector<Value> extractIndices;

    if (tensorType.getRank() == 2) {
      auto srcCtIdxOp = arith::ConstantIndexOp::create(builder, srcCt);
      createdOpCallback(srcCtIdxOp);
      Value srcCtIdx = srcCtIdxOp.getResult();
      extractIndices.push_back(srcCtIdx);
    }
    auto srcSlotIdxOp = arith::ConstantIndexOp::create(builder, srcSlot);
    createdOpCallback(srcSlotIdxOp);
    extractIndices.push_back(srcSlotIdxOp.getResult());
    auto extracted = tensor::ExtractOp::create(builder, input, extractIndices);
    createdOpCallback(extracted);

    auto dstCtIdxOp = arith::ConstantIndexOp::create(builder, dstCt);
    createdOpCallback(dstCtIdxOp);
    auto dstSlotIdxOp = arith::ConstantIndexOp::create(builder, dstSlot);
    createdOpCallback(dstSlotIdxOp);
    auto insertOp = tensor::InsertOp::create(
        builder, extracted.getResult(), result,
        ValueRange{dstCtIdxOp.getResult(), dstSlotIdxOp.getResult()});
    createdOpCallback(insertOp);

    result = insertOp.getResult();
  }

  return result;
}

static FailureOr<Type> getIntermediateTargetType(
    Type elementType, LayoutAttr layout, ImplicitLocOpBuilder& builder) {
  presburger::IntegerRelation rel = layout.getIntegerRelation();
  unsigned numRangeVars = rel.getNumRangeVars();
  SmallVector<int64_t> targetShape;
  for (unsigned j = 0; j < numRangeVars; ++j) {
    auto ub = rel.getConstantBound64(
        presburger::BoundType::UB,
        rel.getVarKindOffset(presburger::VarKind::Range) + j);
    if (!ub.has_value()) {
      return builder.emitError() << "Unbounded range variable in relation";
    }
    targetShape.push_back(ub.value() + 1);
  }
  return RankedTensorType::get(targetShape, elementType);
}

static FailureOr<Value> implementCrtAssignLayoutStep(
    Value input, RankedTensorType inputType, RankedTensorType targetType,
    ImplicitLocOpBuilder& builder,
    const std::function<void(Operation*)>& createdOpCallback) {
  auto zeroOp = arith::ConstantOp::create(builder, targetType,
                                          builder.getZeroAttr(targetType));
  createdOpCallback(zeroOp);
  Value targetTensor = zeroOp.getResult();

  int64_t numSlots = targetType.getShape().back();
  int64_t numCts = (targetType.getRank() == 2) ? targetType.getDimSize(0) : 1;
  int64_t totalSlots = numCts * numSlots;

  Value c0 = arith::ConstantIndexOp::create(builder, 0);
  Value c1 = arith::ConstantIndexOp::create(builder, 1);
  Value cTotalSlots = arith::ConstantIndexOp::create(builder, totalSlots);
  Value cNumSlots = arith::ConstantIndexOp::create(builder, numSlots);

  SmallVector<Value> dimSizes =
      llvm::map_to_vector(inputType.getShape(), [&](int64_t dim) -> Value {
        return arith::ConstantIndexOp::create(builder, dim);
      });

  auto loop = scf::ForOp::create(
      builder, c0, cTotalSlots, c1, ValueRange{targetTensor},
      [&](OpBuilder& b, Location loc, Value k, ValueRange iterArgs) {
        ImplicitLocOpBuilder ib(loc, b);

        SmallVector<Value> extractIndices =
            llvm::map_to_vector(dimSizes, [&](Value dimSize) -> Value {
              return arith::RemSIOp::create(ib, k, dimSize);
            });

        Value extracted = tensor::ExtractOp::create(ib, input, extractIndices);

        SmallVector<Value> insertIndices;
        if (targetType.getRank() == 2) {
          if (numCts == 1) {
            insertIndices.push_back(c0);
            insertIndices.push_back(k);
          } else {
            insertIndices.push_back(arith::DivSIOp::create(ib, k, cNumSlots));
            insertIndices.push_back(arith::RemSIOp::create(ib, k, cNumSlots));
          }
        } else {
          insertIndices.push_back(k);
        }

        Value inserted =
            tensor::InsertOp::create(ib, extracted, iterArgs[0], insertIndices);
        scf::YieldOp::create(ib, ValueRange{inserted});
      });
  createdOpCallback(loop);
  return loop.getResult(0);
}

static FailureOr<Value> implementAssignLayoutStep(
    Value input, LayoutAttr layout, Type targetTypeTy,
    ImplicitLocOpBuilder& builder,
    const std::function<void(Operation*)>& createdOpCallback, bool isLast,
    ArrayRef<int64_t> domainSchedule = {},
    CodegenStrategy strategy = CodegenStrategy::AUTO) {
  presburger::IntegerRelation rel = layout.getIntegerRelation();
  RankedTensorType targetType = cast<RankedTensorType>(targetTypeTy);
  auto elementType = getElementTypeOrSelf(input.getType());
  unsigned numRangeVars = rel.getNumRangeVars();

  auto dataSemanticType = dyn_cast<RankedTensorType>(input.getType());
  if (dataSemanticType) {
    for (unsigned i = 0; i < dataSemanticType.getRank(); ++i) {
      rel.addBound(presburger::BoundType::LB,
                   rel.getVarKindOffset(presburger::VarKind::Domain) + i, 0);
      rel.addBound(presburger::BoundType::UB,
                   rel.getVarKindOffset(presburger::VarKind::Domain) + i,
                   dataSemanticType.getDimSize(i) - 1);
    }
  }

  LLVM_DEBUG(llvm::dbgs() << "implementAssignLayoutStep on input: "
                          << input.getType() << " with layout: " << layout
                          << " targetType: " << targetType << "\n");

  // If the input is a constant zero, then the output will be a constant zero
  // since we zero-fill slots.
  if (matchPattern(input, m_AnyZeroFloat()) || matchPattern(input, m_Zero())) {
    auto zeroOp = arith::ConstantOp::create(builder, targetType,
                                            builder.getZeroAttr(targetType));
    createdOpCallback(zeroOp);
    return zeroOp.getResult();
  }

  // The input came from an empty tensor, so we can just create an empty
  // ciphertext semantic tensor type.
  if (auto emptyOp = dyn_cast_or_null<tensor::EmptyOp>(input.getDefiningOp())) {
    auto emptyCiphertextOp = tensor::EmptyOp::create(
        builder, builder.getLoc(), targetType.getShape(),
        targetType.getElementType());
    createdOpCallback(emptyCiphertextOp);
    return emptyCiphertextOp.getResult();
  }

  // If the input has a bicyclic or tricyclic CRT layout, we directly compute
  // logical coordinates via modular remainder operations (arith.remsi) to
  // avoid ISL codegen overhead.
  int64_t numSlots = targetType.getShape().back();
  if (dataSemanticType &&
      (isRelationBicyclic(dataSemanticType, numSlots, rel) ||
       isRelationTricyclic(dataSemanticType, numSlots, rel))) {
    return implementCrtAssignLayoutStep(input, dataSemanticType, targetType,
                                        builder, createdOpCallback);
  }

  // The result can be simplified if the layout is dense in the ciphertext type,
  // and the input is a scalar, a constant splat, or a single-element tensor
  // (whose dense packing is a broadcast of its lone element).
  SplatElementsAttr splatAttr;
  bool inputIsScalar = !dataSemanticType;
  bool inputIsSplatConstant = matchPattern(input, m_Constant(&splatAttr));
  bool inputIsSingleElementTensor = dataSemanticType &&
                                    dataSemanticType.hasStaticShape() &&
                                    dataSemanticType.getNumElements() == 1;
  if ((inputIsScalar || inputIsSplatConstant || inputIsSingleElementTensor) &&
      isDenseLayout(rel, targetType)) {
    // Regardless of being constant or not, a scalar can be splat into the
    // ciphertext tensor.
    if (inputIsScalar) {
      auto splatOp = tensor::SplatOp::create(builder, targetType, input);
      createdOpCallback(splatOp);
      return splatOp.getResult();
    }
    if (inputIsSplatConstant) {
      auto constantOp = arith::ConstantOp::create(
          builder, targetType,
          SplatElementsAttr::get(targetType,
                                 splatAttr.getSplatValue<TypedAttr>()));
      createdOpCallback(constantOp);
      return constantOp.getResult();
    }
    // A non-constant single-element tensor: the loop-generator path below
    // handles it correctly, but emits one iteration per slot whose body is
    // loop-invariant; extract the element once and splat it instead.
    auto zero = arith::ConstantIndexOp::create(builder, 0);
    createdOpCallback(zero);
    SmallVector<Value> zeroIndices(dataSemanticType.getRank(),
                                   zero.getResult());
    auto extractOp = tensor::ExtractOp::create(builder, input, zeroIndices);
    createdOpCallback(extractOp);
    auto splatOp =
        tensor::SplatOp::create(builder, targetType, extractOp.getResult());
    createdOpCallback(splatOp);
    return splatOp.getResult();
  }

  DenseElementsAttr constantAttr;
  DenseResourceElementsAttr resourceAttr;
  ArrayRef<char> resourceRaw;
  bool isDenseConstant = matchPattern(input, m_Constant(&constantAttr));
  if (!isDenseConstant) {
    if (auto cstOp =
            dyn_cast_or_null<arith::ConstantOp>(input.getDefiningOp())) {
      if ((resourceAttr =
               dyn_cast<DenseResourceElementsAttr>(cstOp.getValue()))) {
        resourceRaw = resourceAttr.getData();
      }
    }
  }
  bool isResourceConstant = !resourceRaw.empty() &&
                            elementType.isIntOrFloat() &&
                            elementType.getIntOrFloatBitWidth() % 8 == 0;
  auto tryFolding = [&]() -> FailureOr<Value> {
    if (!dataSemanticType) return failure();
    LLVM_DEBUG(llvm::dbgs() << "Detected constant input, evaluating layout\n");
    int64_t numTargetElements = targetType.getNumElements();

    PointPairCollector collector(dataSemanticType.getRank(),
                                 /*rangeDims=*/targetType.getRank());
    enumeratePoints(rel, collector);

    // Row-major strides let us flatten a multi-dimensional point with plain
    // integer math, avoiding per-element getFlattenedIndex() interface lookups
    auto rowMajorStrides = [](ArrayRef<int64_t> shape) {
      SmallVector<int64_t> strides(shape.size(), 1);
      for (int i = static_cast<int>(shape.size()) - 2; i >= 0; --i)
        strides[i] = strides[i + 1] * shape[i + 1];
      return strides;
    };
    SmallVector<int64_t> srcStrides =
        rowMajorStrides(dataSemanticType.getShape());
    SmallVector<int64_t> dstStrides = rowMajorStrides(targetType.getShape());
    auto flatten = [](ArrayRef<int64_t> point, ArrayRef<int64_t> strides) {
      int64_t flat = 0;
      for (size_t i = 0; i < point.size(); ++i) flat += point[i] * strides[i];
      return flat;
    };
    bool srcIsSplat = isDenseConstant && constantAttr.isSplat();

    // Fast path: for byte-aligned int/float element types, pack directly into a
    // raw byte buffer rather than building an N-element SmallVector<Attribute>
    // and re-encoding it via DenseElementsAttr::get().
    if (elementType.isIntOrFloat() &&
        elementType.getIntOrFloatBitWidth() % 8 == 0) {
      unsigned byteWidth = elementType.getIntOrFloatBitWidth() / 8;
      ArrayRef<char> srcRaw =
          isDenseConstant ? constantAttr.getRawData() : resourceRaw;
      std::vector<char> rawBuffer(
          static_cast<size_t>(numTargetElements) * byteWidth, 0);

      for (const auto& [domainPoint, rangePoint] : collector.points) {
        int64_t dstFlat = flatten(rangePoint, dstStrides);
        if (dstFlat < 0 || dstFlat >= numTargetElements) continue;
        int64_t srcFlat = srcIsSplat ? 0 : flatten(domainPoint, srcStrides);
        std::memcpy(rawBuffer.data() + static_cast<size_t>(dstFlat) * byteWidth,
                    srcRaw.data() + static_cast<size_t>(srcFlat) * byteWidth,
                    byteWidth);
      }

      TypedAttr packedConstantAttr;
      ArrayRef<char> packedRaw(rawBuffer.data(), rawBuffer.size());
      if (isResourceConstant) {
        std::string resourceName = resourceAttr.getRawHandle().getKey().str();
        resourceName += "_packed";
        auto packedBlob = HeapAsmResourceBlob::allocateAndCopyWithAlign(
            packedRaw,
            resourceAttr.getRawHandle().getBlob()->getDataAlignment(),
            /*dataIsMutable=*/false);
        packedConstantAttr = DenseResourceElementsAttr::get(
            targetType, resourceName, std::move(packedBlob));
      } else {
        packedConstantAttr =
            DenseElementsAttr::getFromRawBuffer(targetType, packedRaw);
      }
      auto constantOp = arith::ConstantOp::create(builder, builder.getLoc(),
                                                  packedConstantAttr);
      createdOpCallback(constantOp);
      return constantOp.getResult();
    }

    // Fallback for sub-byte element types (e.g. i1): build the packed constant
    // from boxed Attributes.
    Attribute zeroAttr = builder.getZeroAttr(elementType);
    SmallVector<Attribute> packedValues(numTargetElements, zeroAttr);
    Attribute splatValue =
        srcIsSplat ? constantAttr.getSplatValue<Attribute>() : Attribute();
    auto srcValues = constantAttr.getValues<Attribute>();
    for (const auto& [domainPoint, rangePoint] : collector.points) {
      int64_t dstFlat = flatten(rangePoint, dstStrides);
      if (dstFlat < 0 || dstFlat >= numTargetElements) continue;
      packedValues[dstFlat] = srcIsSplat
                                  ? splatValue
                                  : srcValues[static_cast<size_t>(
                                        flatten(domainPoint, srcStrides))];
    }

    auto packedConstantAttr =
        DenseElementsAttr::get(targetType, ArrayRef<Attribute>(packedValues));
    auto constantOp = arith::ConstantOp::create(builder, builder.getLoc(),
                                                packedConstantAttr);
    createdOpCallback(constantOp);
    return constantOp.getResult();
  };

  bool shouldFold = false;
  if ((isDenseConstant || isResourceConstant) && dataSemanticType) {
    if (strategy == CodegenStrategy::FOLD_WHEN_POSSIBLE ||
        (strategy == CodegenStrategy::AUTO && isResourceConstant)) {
      shouldFold = true;
    } else if (strategy == CodegenStrategy::AUTO) {
      if (rel.getNumLocalVars() > 5) {
        LLVM_DEBUG(
            llvm::dbgs()
            << "Relation has many existentials, folding constant to avoid "
               "loop generation hangs\n");
        shouldFold = true;
      } else {
        int64_t relSize = relationSize(rel);
        if (relSize >= 0 && relSize <= 16384) {
          shouldFold = true;
        } else {
          LLVM_DEBUG(
              llvm::dbgs()
              << "Relation size " << relSize
              << " exceeds threshold 16384, skipping constant folding\n");
        }
      }
    }
  }

  if (shouldFold) {
    auto folded = tryFolding();
    if (succeeded(folded)) {
      return folded.value();
    }
  }

  auto zeroOp = arith::ConstantOp::create(builder, targetType,
                                          builder.getZeroAttr(targetType));
  createdOpCallback(zeroOp);
  Value targetTensor = zeroOp.getResult();

  MLIRLoopNestGenerator generator(builder, createdOpCallback);

  SmallVector<int> domainIndices = llvm::to_vector(llvm::map_range(
      domainSchedule, [](int64_t idx) { return static_cast<int>(idx); }));

  int numDomainVars = rel.getNumDomainVars();
  llvm::BitVector present(numDomainVars);
  for (int var : domainIndices) {
    if (var >= 0 && var < numDomainVars) {
      present.set(var);
    }
  }
  for (int i = 0; i < numDomainVars; ++i) {
    if (!present.test(i)) {
      domainIndices.push_back(i);
    }
  }
  auto loopNestCstr = generateLoopNestAsCStr(rel, domainIndices);
  if (failed(loopNestCstr)) {
    return builder.emitError() << "Failed to generate loop nest for relation "
                               << printRelation(rel);
  }
  LLVM_DEBUG(llvm::dbgs() << "Generating loop nest assignment for relation "
                          << loopNestCstr.value() << "\n");

  auto loop = generator.generateForLoop(
      rel, {targetTensor},
      [&](OpBuilder& builder, Location loc, ValueRange exprs,
          ValueRange iterArgs) {
        Value extracted = input;
        auto dataSemanticType = dyn_cast<RankedTensorType>(input.getType());
        if (dataSemanticType) {
          SmallVector<Value> extractIndices;
          for (Value idx : exprs.take_front(dataSemanticType.getRank())) {
            extractIndices.push_back(arith::IndexCastOp::create(
                builder, loc, builder.getIndexType(), idx));
          }
          extracted =
              tensor::ExtractOp::create(builder, loc, input, extractIndices);
        }

        SmallVector<Value> insertIndices;
        unsigned relationDomainSize = rel.getNumDomainVars();
        for (Value idx : exprs.slice(relationDomainSize, numRangeVars)) {
          insertIndices.push_back(arith::IndexCastOp::create(
              builder, loc, builder.getIndexType(), idx));
        }
        auto inserted = tensor::InsertOp::create(builder, loc, extracted,
                                                 iterArgs[0], insertIndices);
        return scf::ValueVector({inserted});
      },
      domainIndices);

  if (failed(loop)) {
    return builder.emitError() << "Failed to generate loop nest";
  }

  createdOpCallback(loop.value());
  return loop.value().getResults()[0];
}

FailureOr<Value> implementAssignLayout(
    Value input, Attribute layout, int64_t minSlotCount,
    ImplicitLocOpBuilder& builder,
    const std::function<void(Operation*)>& createdOpCallback,
    ArrayRef<int64_t> domainSchedule, CodegenStrategy strategy) {
  OpBuilder::InsertionGuard guard(builder);

  if (auto arrayAttr = dyn_cast<ArrayAttr>(layout)) {
    Value currentInput = input;
    for (size_t i = 0; i < arrayAttr.size(); ++i) {
      auto layoutAttr = dyn_cast<LayoutAttr>(arrayAttr[i]);
      if (!layoutAttr) {
        return builder.emitError()
               << "All layouts in the array must be LayoutAttr";
      }
      bool isLast = (i == arrayAttr.size() - 1);
      Type targetType;
      auto elementType = getElementTypeOrSelf(currentInput.getType());
      if (isLast) {
        targetType = materializeLayout(elementType, layoutAttr, minSlotCount);
      } else {
        auto intermediateType =
            getIntermediateTargetType(elementType, layoutAttr, builder);
        if (failed(intermediateType)) return failure();
        targetType = intermediateType.value();
      }
      auto result =
          implementAssignLayoutStep(currentInput, layoutAttr, targetType,
                                    builder, createdOpCallback, isLast,
                                    /*domainSchedule=*/{}, strategy);
      if (failed(result)) {
        return failure();
      }
      currentInput = result.value();
    }
    return currentInput;
  }

  if (LayoutAttr layoutAttr = dyn_cast<LayoutAttr>(layout)) {
    auto elementType = getElementTypeOrSelf(input.getType());
    Type targetType = materializeLayout(elementType, layoutAttr, minSlotCount);
    return implementAssignLayoutStep(input, layoutAttr, targetType, builder,
                                     createdOpCallback, /*isLast=*/true,
                                     domainSchedule, strategy);
  } else if (DenseIntElementsAttr elementAttr =
                 dyn_cast<DenseIntElementsAttr>(layout)) {
    Type targetType = materializePermutationLayout(input.getType(), elementAttr,
                                                   minSlotCount);
    return implementAssignLayoutPermutation(input, elementAttr, targetType,
                                            builder, createdOpCallback);
  }
  return builder.emitError() << "Unsupported layout attribute type: " << layout;
}

FailureOr<Value> implementUnpackOp(
    tensor_ext::UnpackOp op, ImplicitLocOpBuilder& builder,
    const std::function<void(Operation*)>& createdOpCallback) {
  OpBuilder::InsertionGuard guard(builder);
  if (auto layoutAttr = dyn_cast<LayoutAttr>(op.getLayout())) {
    return implementUnpackOpStep(op.getValue(), layoutAttr,
                                 op.getResult().getType(), builder,
                                 createdOpCallback);
  }

  if (auto arrayAttr = dyn_cast<ArrayAttr>(op.getLayout())) {
    Value currentInput = op.getValue();
    Type finalTargetType = op.getResult().getType();

    for (int i = arrayAttr.size() - 1; i >= 0; --i) {
      auto layoutAttr = cast<LayoutAttr>(arrayAttr[i]);
      Type stepTargetType;
      if (i == 0) {
        stepTargetType = finalTargetType;
      } else {
        presburger::IntegerRelation rel = layoutAttr.getIntegerRelation();
        unsigned numDomainVars = rel.getNumDomainVars();
        SmallVector<int64_t> domainShape;
        for (unsigned d = 0; d < numDomainVars; ++d) {
          auto ub = rel.getConstantBound64(
              presburger::BoundType::UB,
              rel.getVarKindOffset(presburger::VarKind::Domain) + d);
          if (!ub.has_value()) {
            return builder.emitError()
                   << "Unbounded domain variable in relation";
          }
          domainShape.push_back(ub.value() + 1);
        }
        Type elementType = getElementTypeOrSelf(finalTargetType);
        stepTargetType = RankedTensorType::get(domainShape, elementType);
      }

      auto res = implementUnpackOpStep(currentInput, layoutAttr, stepTargetType,
                                       builder, createdOpCallback);
      if (failed(res)) {
        return failure();
      }
      currentInput = res.value();
    }
    return currentInput;
  }

  return builder.emitError()
         << "Unsupported layout attribute type: " << op.getLayout();
}

}  // namespace heir
}  // namespace mlir
