#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SetVector.h"

#include <string>

#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"

using namespace mlir;

namespace onnx_mlir {
namespace spatial {

RankedTensorType getGraphBatchPhysicalResultType(int64_t laneCount, RankedTensorType fragmentType) {
  SmallVector<int64_t> shape {laneCount};
  llvm::append_range(shape, fragmentType.getShape());
  return RankedTensorType::get(shape, fragmentType.getElementType(), fragmentType.getEncoding());
}

FailureOr<RankedTensorType> getGraphBatchFragmentType(RankedTensorType physicalType, int64_t expectedLaneCount) {
  if (!physicalType || physicalType.getRank() < 1 || physicalType.getDimSize(0) != expectedLaneCount)
    return failure();
  return RankedTensorType::get(physicalType.getShape().drop_front(), physicalType.getElementType(), physicalType.getEncoding());
}
namespace {

std::optional<BlockArgument> getBlockArgument(Region& body, unsigned argIdx) {
  if (body.empty())
    return std::nullopt;

  Block& block = body.front();
  if (argIdx >= block.getNumArguments())
    return std::nullopt;
  return block.getArgument(argIdx);
}

std::optional<BlockArgument> insertBlockArgument(Region& body, unsigned argIdx, Type type, Location loc) {
  if (body.empty())
    return std::nullopt;
  return body.insertArgument(argIdx, type, loc);
}

void setComputeOperandSegmentSizes(Operation* op, int32_t weightCount, int32_t inputCount) {
  if (auto compute = dyn_cast<SpatGraphCompute>(op)) {
    compute.getProperties().setOperandSegmentSizes({weightCount, inputCount});
    return;
  }
  if (auto compute = dyn_cast<SpatScheduledCompute>(op)) {
    compute.getProperties().setOperandSegmentSizes({weightCount, inputCount});
    return;
  }
  if (auto batch = dyn_cast<SpatGraphComputeBatch>(op)) {
    batch.getProperties().setOperandSegmentSizes({weightCount, inputCount});
    return;
  }
  cast<SpatScheduledComputeBatch>(op).getProperties().setOperandSegmentSizes({weightCount, inputCount});
}

using CrossbarWeightSet = llvm::SetVector<Value, llvm::SmallVector<Value, 4>, llvm::SmallDenseSet<Value, 4>>;

CrossbarWeightSet collectCrossbarWeights(Region& body) {
  CrossbarWeightSet weights;
  body.walk([&](SpatVMMOp vmmOp) {
    Value weight = vmmOp.getWeight();
    weights.insert(weight);
  });
  return weights;
}

template <typename ComputeOpTy>
std::optional<BlockArgument> getComputeWeightArgument(ComputeOpTy compute, unsigned idx) {
  return getBlockArgument(compute.getBody(), idx);
}

template <typename ComputeOpTy>
std::optional<BlockArgument> getComputeInputArgument(ComputeOpTy compute, unsigned idx) {
  return getBlockArgument(compute.getBody(), compute.getWeights().size() + idx);
}

template <typename ComputeOpTy>
std::optional<std::tuple<Value, BlockArgument>>
insertComputeWeight(ComputeOpTy compute, unsigned idx, Value weight, Location loc) {
  if (auto existing = llvm::find(compute.getWeights(), weight); existing != compute.getWeights().end()) {
    auto index = std::distance(compute.getWeights().begin(), existing);
    return {{*existing, *getComputeWeightArgument(compute, index)}};
  }

  unsigned weightCount = compute.getWeights().size();
  unsigned inputCount = compute.getInputs().size();
  compute.getOperation()->insertOperands(idx, ValueRange {weight});
  setComputeOperandSegmentSizes(
    compute.getOperation(), static_cast<int32_t>(weightCount + 1), static_cast<int32_t>(inputCount));
  auto blockArg = insertBlockArgument(compute.getBody(), idx, weight.getType(), loc);
  if (!blockArg)
    return std::nullopt;
  return std::make_tuple(compute.getOperation()->getOperand(idx), *blockArg);
}

template <typename ComputeBatchOpTy>
std::optional<std::tuple<Value, BlockArgument>>
insertComputeBatchWeight(ComputeBatchOpTy batch, unsigned idx, Value weight, Location loc) {
  if (auto existing = llvm::find(batch.getWeights(), weight); existing != batch.getWeights().end()) {
    auto index = std::distance(batch.getWeights().begin(), existing);
    return {{*existing, *batch.getWeightArgument(index)}};
  }

  unsigned weightCount = batch.getWeights().size();
  unsigned inputCount = batch.getInputs().size();
  batch.getOperation()->insertOperands(idx, ValueRange {weight});
  setComputeOperandSegmentSizes(
    batch.getOperation(), static_cast<int32_t>(weightCount + 1), static_cast<int32_t>(inputCount));

  auto blockArg = insertBlockArgument(batch.getBody(), 1 + idx, weight.getType(), loc);
  if (!blockArg)
    return std::nullopt;
  return std::make_tuple(batch.getOperation()->getOperand(idx), *blockArg);
}

template <typename ComputeOpTy>
std::optional<std::tuple<Value, BlockArgument>>
insertComputeInput(ComputeOpTy compute, unsigned idx, Value input, Location loc) {
  unsigned weightCount = compute.getWeights().size();
  unsigned inputCount = compute.getInputs().size();
  compute.getOperation()->insertOperands(weightCount + idx, ValueRange {input});
  setComputeOperandSegmentSizes(
    compute.getOperation(), static_cast<int32_t>(weightCount), static_cast<int32_t>(inputCount + 1));
  auto blockArg = insertBlockArgument(compute.getBody(), weightCount + idx, input.getType(), loc);
  if (!blockArg)
    return std::nullopt;
  return std::make_tuple(compute.getOperation()->getOperand(weightCount + idx), *blockArg);
}

template <typename ComputeOpTy>
void setComputeAsmBlockArgumentNames(ComputeOpTy compute, Region& region, OpAsmSetValueNameFn setNameFn) {
  if (region.empty())
    return;

  for (unsigned index = 0; index < compute.getWeights().size(); ++index)
    if (auto weightArg = compute.getWeightArgument(index))
      setNameFn(*weightArg, ("w" + std::to_string(index)).c_str());

  for (unsigned index = 0; index < compute.getInputs().size(); ++index)
    if (auto inputArg = compute.getInputArgument(index))
      setNameFn(*inputArg, ("in" + std::to_string(index)).c_str());
}

template <typename ComputeOpTy>
FailureOr<std::tuple<OpResult, ComputeOpTy>>
insertComputeOutput(ComputeOpTy compute, RewriterBase& rewriter, unsigned idx, Type type, Location loc) {
  if (idx > compute.getNumResults())
    return failure();

  rewriter.setInsertionPoint(compute.getOperation());
  SmallVector<Type> resultTypes(compute.getResultTypes().begin(), compute.getResultTypes().end());
  resultTypes.insert(resultTypes.begin() + idx, type);
  auto newCompute =
    ComputeOpTy::create(rewriter, compute.getLoc(), TypeRange(resultTypes), compute.getWeights(), compute.getInputs());
  newCompute->setAttrs(compute->getAttrs());
  setComputeOperandSegmentSizes(newCompute.getOperation(),
                                static_cast<int32_t>(newCompute.getWeights().size()),
                                static_cast<int32_t>(newCompute.getInputs().size()));
  rewriter.inlineRegionBefore(compute.getBody(), newCompute.getBody(), newCompute.getBody().end());
  for (unsigned oldResultIdx = 0; oldResultIdx < compute.getNumResults(); ++oldResultIdx)
    compute.getResult(oldResultIdx)
      .replaceAllUsesWith(newCompute.getResult(oldResultIdx < idx ? oldResultIdx : oldResultIdx + 1));
  rewriter.eraseOp(compute.getOperation());
  return std::make_tuple(cast<OpResult>(newCompute.getResult(idx)), newCompute);
}

template <typename ComputeBatchOpTy>
FailureOr<std::tuple<OpResult, BlockArgument, ComputeBatchOpTy>>
insertComputeBatchOutput(ComputeBatchOpTy batch, RewriterBase& rewriter, unsigned idx, Type type, Location loc) {
  if (idx > batch.getNumResults())
    return failure();

  rewriter.setInsertionPoint(batch.getOperation());
  SmallVector<Type> resultTypes(batch.getResultTypes().begin(), batch.getResultTypes().end());
  resultTypes.insert(resultTypes.begin() + idx, type);
  auto newBatch =
    ComputeBatchOpTy::create(rewriter, batch.getLoc(), TypeRange(resultTypes), batch.getLaneCountAttr(), batch.getWeights(), batch.getInputs());
  newBatch->setAttrs(batch->getAttrs());
  setComputeOperandSegmentSizes(newBatch.getOperation(),
                                static_cast<int32_t>(newBatch.getWeights().size()),
                                static_cast<int32_t>(newBatch.getInputs().size()));
  rewriter.inlineRegionBefore(batch.getBody(), newBatch.getBody(), newBatch.getBody().end());
  if (newBatch.getBody().empty()) {
    rewriter.eraseOp(newBatch);
    return failure();
  }
  auto blockArg = newBatch.getBody().front().insertArgument(
    1 + newBatch.getWeights().size() + newBatch.getInputs().size() + idx, type, loc);
  for (unsigned oldResultIdx = 0; oldResultIdx < batch.getNumResults(); ++oldResultIdx)
    batch.getResult(oldResultIdx)
      .replaceAllUsesWith(newBatch.getResult(oldResultIdx < idx ? oldResultIdx : oldResultIdx + 1));
  rewriter.eraseOp(batch.getOperation());
  return std::make_tuple(cast<OpResult>(newBatch.getResult(idx)), blockArg, newBatch);
}

} // namespace

bool isGraphComputeLike(Operation* op) { return isa<SpatGraphCompute, SpatGraphComputeBatch>(op); }

bool isGraphBatchComputeLike(Operation* op) { return isa<SpatGraphComputeBatch>(op); }

bool isScheduledComputeLike(Operation* op) { return isa<SpatScheduledCompute, SpatScheduledComputeBatch>(op); }

bool isScheduledBatchComputeLike(Operation* op) { return isa<SpatScheduledComputeBatch>(op); }

bool isAnySpatialComputeLike(Operation* op) {
  return isa<SpatGraphCompute, SpatGraphComputeBatch, SpatScheduledCompute, SpatScheduledComputeBatch>(op);
}

bool isAnySpatialComputeBatchLike(Operation* op) { return isa<SpatGraphComputeBatch, SpatScheduledComputeBatch>(op); }

std::optional<BlockArgument> SpatGraphCompute::getWeightArgument(unsigned idx) { return getComputeWeightArgument(*this, idx); }
std::optional<BlockArgument> SpatGraphCompute::getInputArgument(unsigned idx) { return getComputeInputArgument(*this, idx); }
std::optional<std::tuple<Value, BlockArgument>> SpatGraphCompute::insertWeight(unsigned idx, Value weight, Location loc) {
  return insertComputeWeight(*this, idx, weight, loc);
}
std::optional<std::tuple<Value, BlockArgument>> SpatGraphCompute::insertInput(unsigned idx, Value input, Location loc) {
  return insertComputeInput(*this, idx, input, loc);
}
CrossbarWeightSet SpatGraphCompute::getCrossbarWeights() { return collectCrossbarWeights(getBody()); }
FailureOr<std::tuple<OpResult, SpatGraphCompute>>
SpatGraphCompute::insertOutput(RewriterBase& rewriter, unsigned idx, Type type, Location loc) {
  return insertComputeOutput(*this, rewriter, idx, type, loc);
}
void SpatGraphCompute::getAsmBlockArgumentNames(Region& region, OpAsmSetValueNameFn setNameFn) {
  setComputeAsmBlockArgumentNames(*this, region, setNameFn);
}

std::optional<BlockArgument> SpatScheduledCompute::getWeightArgument(unsigned idx) {
  return getComputeWeightArgument(*this, idx);
}
std::optional<BlockArgument> SpatScheduledCompute::getInputArgument(unsigned idx) { return getComputeInputArgument(*this, idx); }
std::optional<std::tuple<Value, BlockArgument>>
SpatScheduledCompute::insertWeight(unsigned idx, Value weight, Location loc) {
  return insertComputeWeight(*this, idx, weight, loc);
}
std::optional<std::tuple<Value, BlockArgument>>
SpatScheduledCompute::insertInput(unsigned idx, Value input, Location loc) {
  return insertComputeInput(*this, idx, input, loc);
}
CrossbarWeightSet SpatScheduledCompute::getCrossbarWeights() { return collectCrossbarWeights(getBody()); }
FailureOr<std::tuple<OpResult, SpatScheduledCompute>>
SpatScheduledCompute::insertOutput(RewriterBase& rewriter, unsigned idx, Type type, Location loc) {
  return insertComputeOutput(*this, rewriter, idx, type, loc);
}
void SpatScheduledCompute::getAsmBlockArgumentNames(Region& region, OpAsmSetValueNameFn setNameFn) {
  setComputeAsmBlockArgumentNames(*this, region, setNameFn);
}

SuccessorOperands SpatBlockYieldOp::getSuccessorOperands(unsigned index) {
  assert(index == 0 && "invalid successor index");
  return SuccessorOperands(getOutputsMutable());
}

Block* SpatBlockYieldOp::getSuccessorForOperands(ArrayRef<Attribute>) {
  return getOperation()->getNumSuccessors() == 0 ? nullptr : getOperation()->getSuccessor(0);
}

std::optional<BlockArgument> SpatGraphComputeBatch::getLaneArgument() { return getBlockArgument(getBody(), 0); }
std::optional<BlockArgument> SpatGraphComputeBatch::getWeightArgument(unsigned idx) {
  return getBlockArgument(getBody(), 1 + idx);
}
std::optional<BlockArgument> SpatGraphComputeBatch::getInputArgument(unsigned idx) {
  return getBlockArgument(getBody(), 1 + getWeights().size() + idx);
}
std::optional<BlockArgument> SpatGraphComputeBatch::getOutputArgument(unsigned idx) {
  return getBlockArgument(getBody(), 1 + getWeights().size() + getInputs().size() + idx);
}
std::optional<std::tuple<Value, BlockArgument>>
SpatGraphComputeBatch::insertWeight(unsigned idx, Value weight, Location loc) {
  return insertComputeBatchWeight(*this, idx, weight, loc);
}
std::optional<std::tuple<Value, BlockArgument>>
SpatGraphComputeBatch::insertInput(unsigned idx, Value input, Location loc) {
  unsigned weightCount = getWeights().size();
  unsigned inputCount = getInputs().size();
  getOperation()->insertOperands(weightCount + idx, ValueRange {input});
  setComputeOperandSegmentSizes(
    getOperation(), static_cast<int32_t>(weightCount), static_cast<int32_t>(inputCount + 1));
  auto blockArg = insertBlockArgument(getBody(), 1 + weightCount + idx, input.getType(), loc);
  if (!blockArg)
    return std::nullopt;
  return std::make_tuple(getOperation()->getOperand(weightCount + idx), *blockArg);
}
CrossbarWeightSet SpatGraphComputeBatch::getCrossbarWeights() { return collectCrossbarWeights(getBody()); }
FailureOr<std::tuple<OpResult, BlockArgument, SpatGraphComputeBatch>>
SpatGraphComputeBatch::insertOutput(RewriterBase& rewriter, unsigned idx, Type type, Location loc) {
  return insertComputeBatchOutput(*this, rewriter, idx, type, loc);
}
void SpatGraphComputeBatch::getAsmBlockArgumentNames(Region& region, OpAsmSetValueNameFn setNameFn) {
  if (region.empty())
    return;

  if (auto laneArg = getLaneArgument())
    setNameFn(*laneArg, "lane");
  setComputeAsmBlockArgumentNames(*this, region, setNameFn);
  for (unsigned index = 0; index < getNumResults(); ++index) {
    auto outputArg = getOutputArgument(index);
    if (!outputArg)
      continue;
    if (index == 0) {
      setNameFn(*outputArg, "out");
      continue;
    }
    setNameFn(*outputArg, ("out" + std::to_string(index)).c_str());
  }
}

std::optional<BlockArgument> SpatScheduledComputeBatch::getLaneArgument() { return getBlockArgument(getBody(), 0); }
std::optional<BlockArgument> SpatScheduledComputeBatch::getWeightArgument(unsigned idx) {
  return getBlockArgument(getBody(), 1 + idx);
}
std::optional<BlockArgument> SpatScheduledComputeBatch::getInputArgument(unsigned idx) {
  return getBlockArgument(getBody(), 1 + getWeights().size() + idx);
}
std::optional<BlockArgument> SpatScheduledComputeBatch::getOutputArgument(unsigned idx) {
  return getBlockArgument(getBody(), 1 + getWeights().size() + getInputs().size() + idx);
}
std::optional<std::tuple<Value, BlockArgument>>
SpatScheduledComputeBatch::insertWeight(unsigned idx, Value weight, Location loc) {
  return insertComputeBatchWeight(*this, idx, weight, loc);
}
std::optional<std::tuple<Value, BlockArgument>>
SpatScheduledComputeBatch::insertInput(unsigned idx, Value input, Location loc) {
  unsigned weightCount = getWeights().size();
  unsigned inputCount = getInputs().size();
  getOperation()->insertOperands(weightCount + idx, ValueRange {input});
  setComputeOperandSegmentSizes(
    getOperation(), static_cast<int32_t>(weightCount), static_cast<int32_t>(inputCount + 1));
  auto blockArg = insertBlockArgument(getBody(), 1 + weightCount + idx, input.getType(), loc);
  if (!blockArg)
    return std::nullopt;
  return std::make_tuple(getOperation()->getOperand(weightCount + idx), *blockArg);
}
CrossbarWeightSet SpatScheduledComputeBatch::getCrossbarWeights() { return collectCrossbarWeights(getBody()); }
FailureOr<std::tuple<OpResult, BlockArgument, SpatScheduledComputeBatch>>
SpatScheduledComputeBatch::insertOutput(RewriterBase& rewriter, unsigned idx, Type type, Location loc) {
  return insertComputeBatchOutput(*this, rewriter, idx, type, loc);
}
void SpatScheduledComputeBatch::getAsmBlockArgumentNames(Region& region, OpAsmSetValueNameFn setNameFn) {
  if (region.empty())
    return;

  if (auto laneArg = getLaneArgument())
    setNameFn(*laneArg, "lane");
  setComputeAsmBlockArgumentNames(*this, region, setNameFn);
  for (unsigned index = 0; index < getNumResults(); ++index) {
    auto outputArg = getOutputArgument(index);
    if (!outputArg)
      continue;
    if (index == 0) {
      setNameFn(*outputArg, "out");
      continue;
    }
    setNameFn(*outputArg, ("out" + std::to_string(index)).c_str());
  }
}

void SpatInParallelOp::build(OpBuilder& builder, OperationState& result) {
  OpBuilder::InsertionGuard guard(builder);
  Region* bodyRegion = result.addRegion();
  builder.createBlock(bodyRegion);
}

OpResult SpatInParallelOp::getParentResult(int64_t idx) {
  Operation* parent = getOperation()->getParentOp();
  assert(isAnySpatialComputeBatchLike(parent) && "expected Spatial compute batch parent");
  return parent->getResult(idx);
}

llvm::iterator_range<Block::iterator> SpatInParallelOp::getYieldingOps() { return getRegion().front().getOperations(); }

void SpatialDialect::initialize() {
  addTypes<
#define GET_TYPEDEF_LIST
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTypes.cpp.inc"

    >();
  addOperations<
#define GET_OP_LIST
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.cpp.inc"

    >();
}
} // namespace spatial
} // namespace onnx_mlir

//===----------------------------------------------------------------------===//
// TableGen'd op method definitions
//===----------------------------------------------------------------------===//

#define GET_OP_CLASSES
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.cpp.inc"

#define GET_TYPEDEF_CLASSES
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialDialect.cpp.inc"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTypes.cpp.inc"