#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Interfaces/DestinationStyleOpInterface.h"

#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/ADT/SmallVector.h"

#include "src/Accelerators/PIM/Common/PimCommon.hpp"
#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
#include "src/Accelerators/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.hpp"
#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"

using namespace mlir;

namespace onnx_mlir::pim {
namespace {

struct OperationOrdering {
  DenseMap<Operation*, uint64_t> position;
  DenseMap<Operation*, uint64_t> subtreeEnd;
  uint64_t next = 0;
};

static void orderOperation(Operation* op, OperationOrdering& ordering) {
  uint64_t end = ordering.next;
  ordering.position[op] = ordering.next++;
  for (Region& region : op->getRegions())
    for (Block& block : region)
      for (Operation& nested : block) {
        orderOperation(&nested, ordering);
        end = std::max(end, ordering.subtreeEnd.lookup(&nested));
      }
  ordering.subtreeEnd[op] = end;
}

static OperationOrdering buildOrdering(Operation* coreLikeOp) {
  OperationOrdering ordering;
  if (!coreLikeOp || coreLikeOp->getNumRegions() != 1 || coreLikeOp->getRegion(0).empty())
    return ordering;
  for (Operation& op : coreLikeOp->getRegion(0).front())
    orderOperation(&op, ordering);
  return ordering;
}

static bool isRuntimeMemoryTouch(Operation* op) {
  return isa<PimMemCopyHostToDevOp,
             PimMemCopyDevToHostOp,
             PimMemCopyOp,
             PimReceiveOp,
             PimSendOp,
             PimConcatOp,
             PimVMMOp,
             PimTransposeOp,
             PimVVAddOp,
             PimVVSubOp,
             PimVVMulOp,
             PimVVMaxOp,
             PimVVDMulOp,
             PimVAvgOp,
             PimVReluOp,
             PimVTanhOp,
             PimVSigmOp,
             PimVSoftmaxOp>(op);
}

static bool isIgnoredUser(Operation* op) {
  return isLocalMemoryAliasOp(op) || isa<scf::ForOp, scf::YieldOp, memref::DeallocOp>(op)
      || isCoreStaticAddressOp(op);
}

static bool isWithin(Value value, Region* region) {
  if (auto argument = dyn_cast<BlockArgument>(value))
    return argument.getOwner()->getParent() == region;
  if (Operation* definingOp = value.getDefiningOp())
    return definingOp->getParentRegion() == region || region->isAncestor(definingOp->getParentRegion());
  return false;
}

static std::pair<uint64_t, uint64_t> getTouchRange(Value definingValue, Operation* user, const OperationOrdering& ordering) {
  uint64_t start = ordering.position.lookup(user);
  uint64_t end = start;
  for (Operation* current = user; current; current = current->getParentOp())
    if (auto forOp = dyn_cast<scf::ForOp>(current); forOp && !isWithin(definingValue, &forOp.getRegion())) {
      start = std::min(start, ordering.position.lookup(forOp));
      end = std::max(end, ordering.subtreeEnd.lookup(forOp));
    }
  return {start, end};
}

static FailureOr<size_t> getAllocationSize(memref::AllocOp allocation) {
  auto type = dyn_cast<ShapedType>(allocation.getType());
  if (!type)
    return failure();
  auto bytes = getCheckedShapedTypeSizeInBytes(type, allocation, "memory allocation byte size");
  if (failed(bytes))
    return failure();
  return checkedSize(*bytes, allocation, "memory allocation byte size");
}

} // namespace

bool isLocalMemoryAliasOp(Operation* op) {
  return isa<memref::SubViewOp, memref::CastOp, memref::CollapseShapeOp, memref::ExpandShapeOp>(op);
}

LogicalResult walkLocalMemoryUses(Value root, LocalMemoryUseCallback callback) {
  llvm::SmallPtrSet<Value, 16> visitedValues;
  llvm::SmallPtrSet<OpOperand*, 32> visitedUses;
  SmallVector<Value> pending = {root};
  while (!pending.empty()) {
    Value value = pending.pop_back_val();
    if (!visitedValues.insert(value).second)
      continue;
    for (OpOperand& use : value.getUses()) {
      if (!visitedUses.insert(&use).second || failed(callback(use)))
        return failure();
      Operation* user = use.getOwner();

      if (isLocalMemoryAliasOp(user) && use.getOperandNumber() == 0) {
        if (user->getNumResults() != 1)
          return failure();
        pending.push_back(user->getResult(0));
      }

      if (auto dpsOp = dyn_cast<DestinationStyleOpInterface>(user); dpsOp && dpsOp.isDpsInit(&use))
        pending.push_back(dpsOp.getTiedOpResult(&use));

      if (auto forOp = dyn_cast<scf::ForOp>(user))
        for (auto [index, init] : llvm::enumerate(forOp.getInitArgsMutable()))
          if (&init == &use) {
            pending.push_back(forOp.getRegionIterArgs()[index]);
            pending.push_back(forOp.getResult(index));
          }

      if (auto yieldOp = dyn_cast<scf::YieldOp>(user)) {
        unsigned index = use.getOperandNumber();
        Operation* parent = yieldOp->getParentOp();
        if (auto forOp = dyn_cast<scf::ForOp>(parent))
          pending.push_back(forOp.getResult(index));
        else if (auto ifOp = dyn_cast<scf::IfOp>(parent))
          pending.push_back(ifOp.getResult(index));
        else if (auto switchOp = dyn_cast<scf::IndexSwitchOp>(parent))
          pending.push_back(switchOp.getResult(index));
        else
          return failure();
      }
    }
  }
  return success();
}

FailureOr<SmallVector<LocalMemoryInterval, 0>> analyzeLocalMemoryLifetimes(Operation* coreLikeOp) {
  SmallVector<LocalMemoryInterval, 0> intervals;
  OperationOrdering ordering = buildOrdering(coreLikeOp);
  if (ordering.position.empty())
    return intervals;
  uint64_t fallbackEnd = ordering.next - 1;
  bool failedSize = false;
  coreLikeOp->walk([&](memref::AllocOp allocation) {
    auto size = getAllocationSize(allocation);
    if (failed(size)) {
      failedSize = true;
      return;
    }

    LocalMemoryInterval interval {allocation, ordering.position.lookup(allocation), ordering.position.lookup(allocation), *size};
    bool hasRuntimeTouch = false;
    bool needsFallback = false;
    if (failed(walkLocalMemoryUses(allocation.getResult(), [&](OpOperand& use) {
      Operation* user = use.getOwner();
      if (isRuntimeMemoryTouch(user)) {
        auto [start, end] = getTouchRange(allocation.getResult(), user, ordering);
        if (!hasRuntimeTouch) {
          interval.start = start;
          interval.end = end;
          hasRuntimeTouch = true;
        }
        else {
          interval.start = std::min(interval.start, start);
          interval.end = std::max(interval.end, end);
        }
      }
      else if (!isIgnoredUser(user)) {
        needsFallback = true;
      }
      return success();
    })))
      needsFallback = true;
    if (!hasRuntimeTouch || needsFallback) {
      interval.start = std::min(interval.start, ordering.position.lookup(allocation));
      interval.end = fallbackEnd;
    }
    intervals.push_back(interval);
  });
  if (failedSize)
    return failure();
  return intervals;
}

bool localMemoryLifetimesOverlap(const LocalMemoryInterval& lhs, const LocalMemoryInterval& rhs) {
  return !(lhs.end < rhs.start || rhs.end < lhs.start);
}

} // namespace onnx_mlir::pim