//===- QUIRCircuitsAnalsysis.cpp - Cache values for circuits ----*- C++ -*-===//
//
// (C) Copyright IBM 2023.
//
// Any modifications or derivative works of this code must retain this
// copyright notice, and modified files need to carry a notice indicating
// that they have been altered from the originals.
//
//===----------------------------------------------------------------------===//
///
/// This file implements a analysis for caching argument values with default
/// values for angle and duration arguments.
///
//===----------------------------------------------------------------------===//

#include "Dialect/QUIR/Transforms/QUIRCircuitAnalysis.h"

#include "Dialect/OQ3/IR/OQ3Ops.h"
#include "Dialect/QCS/IR/QCSOps.h"
#include "Dialect/QUIR/IR/QUIRAttributes.h"
#include "Dialect/QUIR/IR/QUIROps.h"
#include "Dialect/QUIR/IR/QUIRTypes.h"
#include "Dialect/QUIR/Utils/Utils.h"

#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/Value.h"
#include "mlir/Pass/AnalysisManager.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Support/LLVM.h"

#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Error.h"

#include <cassert>
#include <map>
#include <sys/types.h>
#include <unordered_map>
#include <utility>

using namespace mlir;

namespace mlir::quir {

double parameterValToDouble(mlir::qcs::ParameterLoadOp defOp) {
  return std::get<double>(defOp.getInitialValue());
}

llvm::Expected<double>
angleValToDouble(mlir::Value inVal,
                 mlir::quir::QUIRCircuitAnalysis *circuitAnalysis) {

  llvm::StringRef errorStr;

  if (auto defOp = inVal.getDefiningOp<mlir::quir::ConstantOp>())
    return defOp.getAngleValueFromConstant().convertToDouble();

  if (auto defOp = inVal.getDefiningOp<mlir::qcs::ParameterLoadOp>())
    return parameterValToDouble(defOp);

  if (auto blockArg = inVal.dyn_cast<mlir::BlockArgument>()) {
    auto circuitOp = mlir::dyn_cast<mlir::quir::CircuitOp>(
        inVal.getParentBlock()->getParentOp());
    assert(circuitOp && "can only handle circuit arguments");

    auto argNum = blockArg.getArgNumber();
    if (circuitAnalysis == nullptr) {

      auto argAttr = circuitOp.getArgAttrOfType<mlir::quir::AngleAttr>(
          argNum, mlir::quir::getAngleAttrName());
      return argAttr.getValue().convertToDouble();
    }
    auto parentModuleOp = circuitOp->getParentOfType<mlir::ModuleOp>();
    return std::get<QUIRCircuitAnalysisEntry::ANGLE>(
        circuitAnalysis->getAnalysisMap()[parentModuleOp][circuitOp][argNum]);
  }

  if (auto castOp = inVal.getDefiningOp<mlir::oq3::CastOp>()) {
    auto defOp = castOp.getArg().getDefiningOp<mlir::qcs::ParameterLoadOp>();
    if (defOp)
      return parameterValToDouble(defOp);
    if (auto constOp =
            castOp.getArg().getDefiningOp<mlir::arith::ConstantOp>()) {
      if (auto angleAttr = constOp.getValue().dyn_cast<mlir::quir::AngleAttr>())
        return angleAttr.getValue().convertToDouble();
      if (auto floatAttr = constOp.getValue().dyn_cast<mlir::FloatAttr>())
        return floatAttr.getValue().convertToDouble();
      errorStr = "unable to cast Angle from constant op";
    } else {
      errorStr = "unable to cast Angle from defining op";
    }
  } else {
    errorStr = "Non-constant angles are not supported!";
  }
  return llvm::createStringError(llvm::inconvertibleErrorCode(), errorStr);
} // angleValToDouble

double QUIRCircuitAnalysis::getAngleValue(mlir::Value operand) {
  auto valueOrError = angleValToDouble(operand);
  if (auto err = valueOrError.takeError()) {
    operand.getDefiningOp()->emitOpError() << toString(std::move(err)) + "\n";
    assert(false && "unhandled value in angleValToDouble");
  }
  return *valueOrError;
}

llvm::StringRef QUIRCircuitAnalysis::getParameterName(mlir::Value operand) {
  llvm::StringRef parameterName = {};
  qcs::ParameterLoadOp parameterLoad;
  parameterLoad = dyn_cast<qcs::ParameterLoadOp>(operand.getDefiningOp());

  if (!parameterLoad) {
    auto castOp = dyn_cast<mlir::oq3::CastOp>(operand.getDefiningOp());
    if (castOp)
      parameterLoad =
          dyn_cast<qcs::ParameterLoadOp>(castOp.getArg().getDefiningOp());
  }

  if (parameterLoad)
    parameterName = parameterLoad.getParameterName();
  return parameterName;
}

quir::DurationAttr QUIRCircuitAnalysis::getDuration(mlir::Value operand) {
  quir::DurationAttr duration;
  auto constantOp = dyn_cast<quir::ConstantOp>(operand.getDefiningOp());

  if (constantOp)
    return constantOp.getValue().dyn_cast<DurationAttr>();
  return duration;
}

QUIRCircuitAnalysis::QUIRCircuitAnalysis(mlir::Operation *moduleOp,
                                         AnalysisManager &am) {

  if (not invalid_)
    return;

  std::unordered_map<mlir::Operation *, std::map<llvm::StringRef, Operation *>>
      circuitOps;

  moduleOp->walk([&](CircuitOp circuitOp) {
    circuitOps[circuitOp->getParentOfType<ModuleOp>()][circuitOp.getSymName()] =
        circuitOp.getOperation();
  });

  moduleOp->walk([&](CallCircuitOp callCircuitOp) {
    auto search = circuitOps[callCircuitOp->getParentOfType<ModuleOp>()].find(
        callCircuitOp.getCalleeAttr().getValue());

    if (search ==
        circuitOps[callCircuitOp->getParentOfType<ModuleOp>()].end()) {
      callCircuitOp->emitOpError("Could not find circuit.");
      return;
    }

    auto circuitOp = dyn_cast<CircuitOp>(search->second);
    auto parentModuleOp = circuitOp->getParentOfType<ModuleOp>();

    for (uint ii = 0; ii < callCircuitOp.getOperands().size(); ++ii) {

      double value = 0;
      llvm::StringRef parameterName = {};
      quir::DurationAttr duration;

      auto operand = callCircuitOp.getOperands()[ii];

      // cache angle values and parameter names
      if (auto angType = operand.getType().dyn_cast<quir::AngleType>()) {

        value = getAngleValue(operand);
        parameterName = getParameterName(operand);
        circuitOperands[parentModuleOp][circuitOp][ii] = {value, parameterName,
                                                          duration};
      }

      // cache durations
      if (auto durType = operand.getType().dyn_cast<quir::DurationType>()) {

        duration = getDuration(operand);
        circuitOperands[parentModuleOp][circuitOp][ii] = {value, parameterName,
                                                          duration};
      }
    }
  });
  invalid_ = false;
}

void QUIRCircuitAnalysisPass::runOnOperation() {
  mlir::Pass::getAnalysis<QUIRCircuitAnalysis>();
}

llvm::StringRef QUIRCircuitAnalysisPass::getArgument() const {
  return "quir-circuit-analysis";
}

llvm::StringRef QUIRCircuitAnalysisPass::getDescription() const {
  return "Analyze Circuit Inputs";
}

llvm::StringRef QUIRCircuitAnalysisPass::getName() const {
  return "QUIR Circuit Analysis Pass";
}

} // namespace mlir::quir
