#ifndef LIB_TRANSFORMS_LOWERPOLYNOMIALEVAL_PATTERNS_H_
#define LIB_TRANSFORMS_LOWERPOLYNOMIALEVAL_PATTERNS_H_

#include "lib/Dialect/Polynomial/IR/PolynomialOps.h"
#include "mlir/include/mlir/IR/MLIRContext.h"   // from @llvm-project
#include "mlir/include/mlir/IR/PatternMatch.h"  // from @llvm-project
#include "mlir/include/mlir/Support/LLVM.h"     // from @llvm-project

// Lowering patterns for polynomial.eval.

namespace mlir {
class DataFlowSolver;
namespace heir {

struct LoweringBase : public OpRewritePattern<polynomial::EvalOp> {
  LoweringBase(MLIRContext* context, bool force = false,
               PatternBenefit benefit = 1)
      : mlir::OpRewritePattern<polynomial::EvalOp>(context, benefit),
        force(force) {}

  bool shouldForce() const { return force; }

 private:
  // Force the use of this pattern, ignoring any heuristics on whether to apply
  // it.
  const bool force;
};

struct ChebyshevLoweringBase : public LoweringBase {
  using LoweringBase::LoweringBase;

  ChebyshevLoweringBase(MLIRContext* context, bool force = false,
                        double minCoefficientThreshold = 1e-12)
      : LoweringBase(context, force),
        minCoefficientThreshold(minCoefficientThreshold) {}

  double getMinCoefficientThreshold() const { return minCoefficientThreshold; }

 private:
  // Minimum threshold for coefficients to be included in the lowered polynomial
  const double minCoefficientThreshold;
};

// Lower polynomial.eval that uses a monomial float polynomial to a series of
// adds and muls via Horner's method. Supports scalar and tensor operands of
// floating point types.
struct LowerViaHorner : public LoweringBase {
  using LoweringBase::LoweringBase;

  LogicalResult matchAndRewrite(polynomial::EvalOp op,
                                PatternRewriter& rewriter) const override;
};

// Lower polynomial.eval that uses a monomial float polynomial to a series of
// adds and muls via the Paterson-Stockmeyer method. Supports scalar and tensor
// operands of floating point types.
struct LowerViaPatersonStockmeyerMonomial : public LoweringBase {
  using LoweringBase::LoweringBase;

  LogicalResult matchAndRewrite(polynomial::EvalOp op,
                                PatternRewriter& rewriter) const override;
};

// Lower polynomial.eval that uses a Chebyshev float polynomial to a series of
// adds and muls via the Paterson-Stockmeyer method. Supports scalar and tensor
// operands of floating point types.
struct LowerViaPatersonStockmeyerChebyshev : public ChebyshevLoweringBase {
  using ChebyshevLoweringBase::ChebyshevLoweringBase;

  LogicalResult matchAndRewrite(polynomial::EvalOp op,
                                PatternRewriter& rewriter) const override;
};

// Lower polynomial.eval that uses a Chebyshev float polynomial to
// kernel.eval_chebyshev.
struct LowerToKernelEvalChebyshev : public LoweringBase {
  LowerToKernelEvalChebyshev(MLIRContext* context, const DataFlowSolver& solver,
                             bool force = false)
      : LoweringBase(context, force, /*benefit=*/2), solver(solver) {}

  LogicalResult matchAndRewrite(polynomial::EvalOp op,
                                PatternRewriter& rewriter) const override;

 private:
  const DataFlowSolver& solver;
};

}  // namespace heir
}  // namespace mlir

#endif  // LIB_TRANSFORMS_LOWERPOLYNOMIALEVAL_PATTERNS_H_
