#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 {
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<SpatCompute>(op)) {
    compute.getProperties().setOperandSegmentSizes({weightCount, inputCount});
    return;
  }
  cast<SpatComputeBatch>(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;
}

} // namespace

std::optional<BlockArgument> SpatCompute::getWeightArgument(unsigned idx) { return getBlockArgument(getBody(), idx); }

std::optional<BlockArgument> SpatCompute::getInputArgument(unsigned idx) {
  return getBlockArgument(getBody(), getWeights().size() + idx);
}

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

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

std::optional<std::tuple<Value, BlockArgument>> SpatCompute::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(), weightCount + idx, input.getType(), loc);
  if (!blockArg)
    return std::nullopt;
  return std::make_tuple(getOperation()->getOperand(weightCount + idx), *blockArg);
}

CrossbarWeightSet SpatCompute::getCrossbarWeights() { return collectCrossbarWeights(getBody()); }

FailureOr<std::tuple<OpResult, SpatCompute>>
SpatCompute::insertOutput(RewriterBase& rewriter, unsigned idx, Type type, Location loc) {
  if (idx > getNumResults())
    return failure();

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

void SpatCompute::getAsmBlockArgumentNames(Region& region, OpAsmSetValueNameFn setNameFn) {
  if (region.empty())
    return;

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

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

std::optional<BlockArgument> SpatComputeBatch::getLaneArgument() { return getBlockArgument(getBody(), 0); }

std::optional<BlockArgument> SpatComputeBatch::getWeightArgument(unsigned idx) {
  return getBlockArgument(getBody(), 1 + idx);
}

std::optional<BlockArgument> SpatComputeBatch::getInputArgument(unsigned idx) {
  return getBlockArgument(getBody(), 1 + getWeights().size() + idx);
}

std::optional<BlockArgument> SpatComputeBatch::getOutputArgument(unsigned idx) {
  return getBlockArgument(getBody(), 1 + getWeights().size() + getInputs().size() + idx);
}

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

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

std::optional<std::tuple<Value, BlockArgument>> SpatComputeBatch::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 SpatComputeBatch::getCrossbarWeights() { return collectCrossbarWeights(getBody()); }

FailureOr<std::tuple<OpResult, BlockArgument, SpatComputeBatch>>
SpatComputeBatch::insertOutput(RewriterBase& rewriter, unsigned idx, Type type, Location loc) {
  if (idx > getNumResults())
    return failure();

  rewriter.setInsertionPoint(getOperation());
  SmallVector<Type> resultTypes(getResultTypes().begin(), getResultTypes().end());
  resultTypes.insert(resultTypes.begin() + idx, type);
  auto newBatch =
    SpatComputeBatch::create(rewriter, getLoc(), TypeRange(resultTypes), getLaneCountAttr(), getWeights(), getInputs());
  newBatch->setAttrs((*this)->getAttrs());
  setComputeOperandSegmentSizes(newBatch.getOperation(),
                                static_cast<int32_t>(newBatch.getWeights().size()),
                                static_cast<int32_t>(newBatch.getInputs().size()));
  rewriter.inlineRegionBefore(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 < getNumResults(); ++oldResultIdx)
    getResult(oldResultIdx)
      .replaceAllUsesWith(newBatch.getResult(oldResultIdx < idx ? oldResultIdx : oldResultIdx + 1));
  rewriter.eraseOp(getOperation());
  return std::make_tuple(cast<OpResult>(newBatch.getResult(idx)), blockArg, newBatch);
}

void SpatComputeBatch::getAsmBlockArgumentNames(Region& region, OpAsmSetValueNameFn setNameFn) {
  if (region.empty())
    return;

  if (auto laneArg = getLaneArgument())
    setNameFn(*laneArg, "lane");

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

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

  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) { return getOperation()->getParentOp()->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"