//===- SimpliyMath.cpp - Simplify Mathematical operations  ------------------ //
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This file implements a pass to lower custom ops generated by the Enzyme AD
// procedure to the MemRef dialect.
//===----------------------------------------------------------------------===//

#include "Dialect/Ops.h"
#include "PassDetails.h"
#include "Passes/Passes.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Complex/IR/Complex.h"

#include "mlir/IR/Matchers.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"

using namespace mlir;
using namespace enzyme;
using llvm::errs;
namespace {

struct AddSimplify : public OpRewritePattern<arith::AddFOp> {
  using OpRewritePattern<arith::AddFOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(arith::AddFOp op,
                                PatternRewriter &rewriter) const final {

    if (matchPattern(op.getLhs(), m_AnyZeroFloat())) {
      rewriter.replaceOp(op, op.getRhs());
      return success();
    }

    if (matchPattern(op.getRhs(), m_AnyZeroFloat())) {
      rewriter.replaceOp(op, op.getLhs());
      return success();
    }

    return failure();
  }
};

struct SubSimplify : public OpRewritePattern<arith::SubFOp> {
  using OpRewritePattern<arith::SubFOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(arith::SubFOp op,
                                PatternRewriter &rewriter) const final {

    if (matchPattern(op.getRhs(), m_AnyZeroFloat())) {
      rewriter.replaceOp(op, op.getLhs());
      return success();
    }

    if (matchPattern(op.getLhs(), m_AnyZeroFloat())) {
      rewriter.replaceOpWithNewOp<arith::NegFOp>(op, op.getRhs());
      return success();
    }

    return failure();
  }
};

bool isZero(mlir::Value v) {
  ArrayAttr lhs;
  matchPattern(v, m_Constant(&lhs));
  if (lhs) {
    for (auto e : lhs) {
      if (!e.cast<FloatAttr>().getValue().isZero())
        return false;
    }
    return true;
  }
  return false;
}

struct CAddSimplify : public OpRewritePattern<complex::AddOp> {
  using OpRewritePattern<complex::AddOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(complex::AddOp op,
                                PatternRewriter &rewriter) const final {

    if (isZero(op.getLhs())) {
      rewriter.replaceOp(op, op.getRhs());
      return success();
    }

    if (isZero(op.getRhs())) {
      rewriter.replaceOp(op, op.getLhs());
      return success();
    }

    return failure();
  }
};

struct CSubSimplify : public OpRewritePattern<complex::SubOp> {
  using OpRewritePattern<complex::SubOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(complex::SubOp op,
                                PatternRewriter &rewriter) const final {

    if (isZero(op.getRhs())) {
      rewriter.replaceOp(op, op.getLhs());
      return success();
    }

    if (isZero(op.getLhs())) {
      rewriter.replaceOpWithNewOp<complex::NegOp>(op, op.getRhs());
      return success();
    }

    return failure();
  }
};

struct MathematicSimplification
    : public enzyme::MathematicSimplificationPassBase<
          MathematicSimplification> {
  void runOnOperation() override {

    RewritePatternSet patterns(&getContext());
    patterns.insert<AddSimplify, SubSimplify, CAddSimplify, CSubSimplify>(
        &getContext());

    GreedyRewriteConfig config;
    (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns),
                                       config);
  };
};
} // end anonymous namespace

namespace mlir {
namespace enzyme {
std::unique_ptr<Pass> createMathematicSimplificationPass() {
  return std::make_unique<MathematicSimplification>();
}
} // namespace enzyme
} // namespace mlir
