//===- ComplexAutoDiffOpInterfaceImpl.cpp - Interface external model --===//
//
// 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 contains the external model implementation of the automatic
// differentiation op interfaces for the upstream MLIR complex dialect.
//
//===----------------------------------------------------------------------===//

#include "Implementations/CoreDialectsAutoDiffImplementations.h"
#include "Interfaces/AutoDiffOpInterface.h"
#include "Interfaces/AutoDiffTypeInterface.h"
#include "Interfaces/GradientUtils.h"
#include "Interfaces/GradientUtilsReverse.h"

#include "mlir/Dialect/Complex/IR/Complex.h"
#include "mlir/IR/DialectRegistry.h"
#include "mlir/Support/LogicalResult.h"

using namespace mlir;
using namespace mlir::enzyme;

namespace {
#include "Implementations/ComplexDerivatives.inc"

struct ComplexAddSimplifyMathInterface
    : public MathSimplifyInterface::ExternalModel<
          ComplexAddSimplifyMathInterface, complex::AddOp> {
  mlir::LogicalResult simplifyMath(Operation *src,
                                   PatternRewriter &rewriter) const {
    auto op = cast<complex::AddOp>(src);

    auto ATy = cast<AutoDiffTypeInterface>(op.getLhs().getType());

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

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

    return failure();
  }
};

struct ComplexSubSimplifyMathInterface
    : public MathSimplifyInterface::ExternalModel<
          ComplexSubSimplifyMathInterface, complex::SubOp> {
  mlir::LogicalResult simplifyMath(Operation *src,
                                   PatternRewriter &rewriter) const {
    auto op = cast<complex::SubOp>(src);

    auto ATy = cast<AutoDiffTypeInterface>(op.getLhs().getType());

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

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

    return failure();
  }
};

} // namespace

void mlir::enzyme::registerComplexDialectAutoDiffInterface(
    DialectRegistry &registry) {
  registry.addExtension(+[](MLIRContext *context, complex::ComplexDialect *) {
    complex::AddOp::attachInterface<ComplexAddSimplifyMathInterface>(*context);
    complex::SubOp::attachInterface<ComplexSubSimplifyMathInterface>(*context);
    registerInterfaces(context);
  });
}
