#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/DialectConversion.h"

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

#include "Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
#include "src/Accelerators/PIM/Common/PimCommon.hpp"
#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
#include "src/Accelerators/PIM/Pass/PIMPasses.h"
#include "src/Dialect/ONNX/ONNXOps.hpp"

using namespace mlir;

namespace onnx_mlir {
namespace {

static constexpr StringLiteral kDenseLayout = "dense_nchw";
static constexpr StringLiteral kRowStripLayout = "nchw_row_strip";

static FailureOr<RowStripPhysicalValue> getRowStripValue(llvm::DenseMap<Value, RowStripPhysicalValue>& rowStripValues,
                                                         Value value) {
  auto it = rowStripValues.find(value);
  if (it == rowStripValues.end())
    return failure();
  return it->second;
}

static FailureOr<RowStripPhysicalValue> buildRowStripValue(spatial::SpatBlueprintOp blueprint,
                                                           Value storage) {
  auto logicalType = dyn_cast<RankedTensorType>(blueprint.getOutput().getType());
  if (!logicalType)
    return blueprint.emitOpError("requires ranked logical output type"), failure();
  RowStripPhysicalValue value;
  value.storage = storage;
  value.logicalType = logicalType;
  value.fragmentOffsets.append(blueprint.getFragmentOffsets().begin(), blueprint.getFragmentOffsets().end());
  value.fragmentSizes.append(blueprint.getFragmentSizes().begin(), blueprint.getFragmentSizes().end());
  if (blueprint.getIndexMap() != kRowStripIndexMap)
    return blueprint.emitOpError("requires the canonical row-strip index map"), failure();
  auto storageType = dyn_cast<RankedTensorType>(storage.getType());
  if (!storageType || storageType != getRowStripStorageType(logicalType))
    return blueprint.emitOpError("requires physical row-strip fragment storage"), failure();
  return value;
}

static FailureOr<Value>
lowerRowStripRelu(const RowStripPhysicalValue& input, spatial::SpatReluPlanOp planOp, PatternRewriter& rewriter) {
  return applyRowStripRelu(input.storage, input.logicalType, rewriter, planOp.getLoc());
}

static FailureOr<Value> lowerRowStripBiasAdd(const RowStripPhysicalValue& input,
                                             spatial::SpatBiasAddPlanOp planOp,
                                             PatternRewriter& rewriter) {
  return applyRowStripBiasAdd(input.storage, input.logicalType, planOp.getBias(), rewriter, planOp.getLoc());
}

static FailureOr<Value>
materializeRowStripToDense(const RowStripPhysicalValue& rowStripValue, Location loc, PatternRewriter& rewriter) {
  if (rowStripValue.logicalType.getRank() != 4 || !rowStripValue.logicalType.hasStaticShape())
    return failure();
  auto [expectedOffsets, expectedSizes] = buildRowStripMetadata(rowStripValue.logicalType);
  if (!llvm::equal(rowStripValue.fragmentOffsets, expectedOffsets) || !llvm::equal(rowStripValue.fragmentSizes, expectedSizes))
    return failure();
  return materializeRowStripStorageToDense(rowStripValue.storage, rowStripValue.logicalType, rewriter, loc);
}

struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, OperationPass<ModuleOp>> {
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LowerSpatialPlansPass)

  StringRef getArgument() const override { return "lower-spatial-plans"; }
  StringRef getDescription() const override { return "Lower selected Spatial planning ops to low-level Spatial IR."; }

  void runOnOperation() override {
    ModuleOp moduleOp = getOperation();
    MLIRContext* ctx = moduleOp.getContext();
    auto entryFunc = getPimEntryFunc(moduleOp);
    if (failed(entryFunc)) {
      moduleOp.emitError("failed to locate the PIM entry function during LowerSpatialPlans");
      signalPassFailure();
      return;
    }
    func::FuncOp funcOp = *entryFunc;
    PatternRewriter rewriter(ctx);
    llvm::DenseMap<Value, RowStripPhysicalValue> rowStripValues;
    llvm::SmallPtrSet<Operation*, 16> eraseAfterLowering;
    auto verifyLogicalPhase = [&](StringRef stage) -> bool {
      if (succeeded(verifyLogicalSpatialGraphInvariants(*entryFunc)))
        return true;
      moduleOp.emitError() << "logical Spatial graph verification failed " << stage;
      signalPassFailure();
      return false;
    };

    if (!verifyLogicalPhase("at the start of LowerSpatialPlans"))
      return;
    for (Operation& op : llvm::make_early_inc_range(funcOp.getBody().front())) {
      if (auto planOp = dyn_cast<spatial::SpatConv2DPlanOp>(&op)) {
        FailureOr<RowStripPhysicalValue> rowStripInput = getRowStripValue(rowStripValues, planOp.getInput());
        auto rowStripBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
          auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(user);
          return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout;
        });
        if (rowStripBlueprint != planOp.getResult().getUsers().end()) {
          rewriter.setInsertionPoint(planOp);
          FailureOr<Value> lowered = lowerSelectedConv2DPlan(
            planOp,
            succeeded(rowStripInput) ? std::optional<Value> {rowStripInput->storage} : std::nullopt,
            /*emitRowStripLayout=*/true,
            rewriter);
          if (failed(lowered)) {
            planOp.emitOpError("failed to lower selected row-strip Spatial Conv plan");
            signalPassFailure();
            return;
          }
          auto blueprint = cast<spatial::SpatBlueprintOp>(*rowStripBlueprint);
          FailureOr<RowStripPhysicalValue> rowStripValue = buildRowStripValue(blueprint, *lowered);
          if (failed(rowStripValue)) {
            signalPassFailure();
            return;
          }
          rowStripValues[blueprint.getResult()] = *rowStripValue;
          eraseAfterLowering.insert(planOp);
          eraseAfterLowering.insert(blueprint);
          continue;
        }
        rewriter.setInsertionPoint(planOp);
        FailureOr<Value> lowered =
          lowerSelectedConv2DPlan(planOp, std::nullopt, /*emitRowStripLayout=*/false, rewriter);
        if (failed(lowered)) {
          planOp.emitOpError("failed to lower selected Spatial Conv plan");
          signalPassFailure();
          return;
        }
        rewriter.replaceOp(planOp, *lowered);
        continue;
      }

      if (auto planOp = dyn_cast<spatial::SpatReluPlanOp>(&op)) {
        if (succeeded(getRowStripValue(rowStripValues, planOp.getInput()))) {
          auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
            auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(user);
            return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout;
          });
          if (outputBlueprint == planOp.getResult().getUsers().end()) {
            planOp.emitOpError("row-strip Relu plan requires a row-strip blueprint result");
            signalPassFailure();
            return;
          }

          FailureOr<RowStripPhysicalValue> input = getRowStripValue(rowStripValues, planOp.getInput());
          rewriter.setInsertionPoint(planOp);
          FailureOr<Value> lowered = lowerRowStripRelu(*input, planOp, rewriter);
          if (failed(lowered)) {
            planOp.emitOpError("failed to lower selected row-strip Spatial Relu plan");
            signalPassFailure();
            return;
          }
          auto blueprint = cast<spatial::SpatBlueprintOp>(*outputBlueprint);
          FailureOr<RowStripPhysicalValue> output = buildRowStripValue(blueprint, *lowered);
          if (failed(output)) {
            signalPassFailure();
            return;
          }
          rowStripValues[blueprint.getResult()] = *output;
          eraseAfterLowering.insert(planOp);
          eraseAfterLowering.insert(blueprint);
          continue;
        }

        rewriter.setInsertionPoint(planOp);
        auto computeOp = createSpatCompute<1>(
          rewriter, planOp.getLoc(), planOp.getOutput().getType(), {}, planOp.getInput(), [&](Value x) {
            auto relu = spatial::SpatReluOp::create(rewriter, planOp.getLoc(), planOp.getOutput().getType(), x);
            spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), relu.getResult());
          });
        rewriter.replaceOp(planOp, computeOp.getResults());
        continue;
      }
      if (auto planOp = dyn_cast<spatial::SpatBiasAddPlanOp>(&op)) {
        if (succeeded(getRowStripValue(rowStripValues, planOp.getInput()))) {
          auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
            auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(user);
            return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout;
          });
          if (outputBlueprint == planOp.getResult().getUsers().end()) {
            planOp.emitOpError("row-strip bias_add plan requires a row-strip blueprint result");
            signalPassFailure();
            return;
          }

          FailureOr<RowStripPhysicalValue> input = getRowStripValue(rowStripValues, planOp.getInput());
          rewriter.setInsertionPoint(planOp);
          FailureOr<Value> lowered = lowerRowStripBiasAdd(*input, planOp, rewriter);
          if (failed(lowered)) {
            planOp.emitOpError("failed to lower selected row-strip Spatial bias_add plan");
            signalPassFailure();
            return;
          }
          auto blueprint = cast<spatial::SpatBlueprintOp>(*outputBlueprint);
          FailureOr<RowStripPhysicalValue> output = buildRowStripValue(blueprint, *lowered);
          if (failed(output)) {
            signalPassFailure();
            return;
          }
          rowStripValues[blueprint.getResult()] = *output;
          eraseAfterLowering.insert(planOp);
          eraseAfterLowering.insert(blueprint);
          continue;
        }

        auto resultType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
        if (!resultType) {
          planOp.emitOpError("requires ranked output type");
          signalPassFailure();
          return;
        }
        rewriter.setInsertionPoint(planOp);
        FailureOr<Value> denseBias = materializeDenseBiasAddTensor(planOp.getBias(), resultType, rewriter, planOp.getLoc());
        if (failed(denseBias)) {
          planOp.emitOpError("failed to materialize dense Conv-style bias");
          signalPassFailure();
          return;
        }
        auto computeOp = createSpatCompute<2>(rewriter,
                                              planOp.getLoc(),
                                              planOp.getOutput().getType(),
                                              {},
                                              ValueRange {planOp.getInput(), *denseBias},
                                              [&](Value x, Value y) {
                                                auto added = spatial::SpatVAddOp::create(
                                                  rewriter, planOp.getLoc(), planOp.getOutput().getType(), x, y);
                                                spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), added.getResult());
                                              });
        rewriter.replaceOp(planOp, computeOp.getResults());
        continue;
      }
      if (auto materializeOp = dyn_cast<spatial::SpatMaterializeLayoutOp>(&op)) {
        if (materializeOp.getSourcePhysicalLayout() == kDenseLayout
            && materializeOp.getTargetPhysicalLayout() == kDenseLayout) {
          rewriter.replaceOp(materializeOp, materializeOp.getInput());
          continue;
        }
        if (materializeOp.getSourcePhysicalLayout() != kRowStripLayout
            || materializeOp.getTargetPhysicalLayout() != kDenseLayout) {
          materializeOp.emitOpError("non-dense materialize_layout lowering is not supported yet");
          signalPassFailure();
          return;
        }
        FailureOr<RowStripPhysicalValue> rowStripValue = getRowStripValue(rowStripValues, materializeOp.getInput());
        if (failed(rowStripValue)) {
          materializeOp.emitOpError("expected a row-strip blueprint input during row-strip materialization");
          signalPassFailure();
          return;
        }
        rewriter.setInsertionPoint(materializeOp);
        FailureOr<Value> dense = materializeRowStripToDense(*rowStripValue, materializeOp.getLoc(), rewriter);
        if (failed(dense)) {
          materializeOp.emitOpError("failed to materialize selected row-strip layout back to dense NCHW");
          signalPassFailure();
          return;
        }
        rewriter.replaceOp(materializeOp, *dense);
        continue;
      }
      if (auto blueprintOp = dyn_cast<spatial::SpatBlueprintOp>(&op)) {
        if (blueprintOp.getPhysicalLayout() == kDenseLayout) {
          rewriter.replaceOp(blueprintOp, blueprintOp.getInput());
          continue;
        }
        if (blueprintOp.getPhysicalLayout() != kRowStripLayout) {
          blueprintOp.emitOpError("non-dense blueprint lowering is not supported yet");
          signalPassFailure();
          return;
        }
        if (!eraseAfterLowering.contains(blueprintOp)) {
          blueprintOp.emitOpError("unhandled row-strip blueprint remained during LowerSpatialPlans");
          signalPassFailure();
          return;
        }
      }
    }
    bool erasedAny = true;
    while (erasedAny) {
      erasedAny = false;
      for (Operation& op : llvm::make_early_inc_range(funcOp.getBody().front())) {
        if (!eraseAfterLowering.contains(&op))
          continue;
        if (!op.use_empty())
          continue;
        eraseAfterLowering.erase(&op);
        rewriter.eraseOp(&op);
        erasedAny = true;
      }
    }
    if (!eraseAfterLowering.empty()) {
      for (Operation& op : funcOp.getBody().front())
        if (eraseAfterLowering.contains(&op))
          op.emitOpError("selected row-strip planning op could not be fully eliminated during LowerSpatialPlans");
      signalPassFailure();
      return;
    }
    ConversionTarget helperTarget(*ctx);
    helperTarget.addLegalDialect<spatial::SpatialDialect,
                                 tensor::TensorDialect,
                                 linalg::LinalgDialect,
                                 affine::AffineDialect,
                                 arith::ArithDialect,
                                 scf::SCFDialect,
                                 func::FuncDialect>();
    helperTarget.addLegalOp<spatial::SpatGraphCompute, spatial::SpatGraphComputeBatch>();
    helperTarget.addIllegalOp<ONNXGemmOp, ONNXTransposeOp>();
    helperTarget.markOpRecursivelyLegal<spatial::SpatGraphCompute, spatial::SpatGraphComputeBatch>();

    RewritePatternSet helperPatterns(ctx);
    populateGemmPatterns(helperPatterns, ctx);
    populateTransposePatterns(helperPatterns, ctx);
    if (failed(applyPartialConversion(moduleOp, helperTarget, std::move(helperPatterns)))) {
      moduleOp.emitError("failed to lower helper ONNX ops emitted by selected Spatial plan lowering");
      signalPassFailure();
      return;
    }
    FrozenRewritePatternSet nestedHelperPatterns([&] {
      RewritePatternSet patterns(ctx);
      populateGemmPatterns(patterns, ctx);
      populateTransposePatterns(patterns, ctx);
      return patterns;
    }());
    ConversionTarget nestedHelperTarget(*ctx);
    nestedHelperTarget.addLegalDialect<spatial::SpatialDialect,
                                       tensor::TensorDialect,
                                       linalg::LinalgDialect,
                                       affine::AffineDialect,
                                       arith::ArithDialect,
                                       scf::SCFDialect,
                                       func::FuncDialect>();
    nestedHelperTarget.addIllegalOp<ONNXGemmOp, ONNXTransposeOp>();
    SmallVector<Operation*> computeLikeOps;
    funcOp.walk([&](Operation* op) {
      if (isa<spatial::SpatGraphCompute, spatial::SpatGraphComputeBatch>(op))
        computeLikeOps.push_back(op);
    });
    for (Operation* op : computeLikeOps) {
      if (failed(applyFullConversion(op, nestedHelperTarget, nestedHelperPatterns))) {
        op->emitOpError("failed to lower nested helper ONNX ops emitted by selected Spatial plan lowering");
        signalPassFailure();
        return;
      }
    }
    if (!verifyLogicalPhase("after nested helper conversions"))
      return;
    bool hasIllegalOps = false;
    moduleOp.walk([&](Operation* op) {
      if (isa<ONNXEntryPointOp>(op))
        return;
      if (isa<spatial::SpatConv2DPlanOp,
              spatial::SpatBiasAddPlanOp,
              spatial::SpatReluPlanOp,
              spatial::SpatBlueprintOp,
              spatial::SpatMaterializeLayoutOp>(op)
          || op->getDialect()->getNamespace() == "onnx") {
        op->emitOpError("operation must not remain after LowerSpatialPlans");
        hasIllegalOps = true;
      }
    });
    if (hasIllegalOps) {
      signalPassFailure();
    } else {
      dumpModule(moduleOp, "spatial1_graph");
      spatial::SpatialDataflowExportStage exportMode = spatial::getSpatialDataflowExportStage();
      if (spatial::shouldExportSpatialDataflowStage(exportMode, spatial::SpatialDataflowExportStage::Pre)
          && failed(spatial::exportSpatialDataflowCsvPre(funcOp))) {
        signalPassFailure();
        return;
      }
    }

    if (!verifyLogicalPhase("at the end of LowerSpatialPlans"))
      return;
  }
};

} // namespace

std::unique_ptr<Pass> createLowerSpatialPlansPass() { return std::make_unique<LowerSpatialPlansPass>(); }

} // namespace onnx_mlir