/*
 * Copyright (c) 2023 - 2025 Chair for Design Automation, TUM
 * Copyright (c) 2025 Munich Quantum Software Company GmbH
 * All rights reserved.
 *
 * SPDX-License-Identifier: MIT
 *
 * Licensed under the MIT License
 */

#pragma once

#include "mlir/Dialect/MQTOpt/Transforms/Transpilation/Common.h"

#include <cstddef>
#include <cstdint>
#include <llvm/ADT/DenseSet.h>
#include <llvm/ADT/SmallVector.h>
#include <llvm/ADT/StringRef.h>
#include <memory>
#include <string>
#include <string_view>

namespace mqt::ir::opt {

/**
 * @brief A quantum accelerator's architecture.
 * @details Computes all-shortest paths at construction.
 */
class Architecture {
public:
  using CouplingSet = mlir::DenseSet<std::pair<QubitIndex, QubitIndex>>;
  using NeighbourVector = mlir::SmallVector<mlir::SmallVector<QubitIndex, 4>>;

  explicit Architecture(std::string name, std::size_t nqubits,
                        CouplingSet couplingSet)
      : name_(std::move(name)), nqubits_(nqubits),
        couplingSet_(std::move(couplingSet)), neighbours_(nqubits),
        dist_(nqubits, llvm::SmallVector<std::size_t>(nqubits, UINT64_MAX)),
        prev_(nqubits, llvm::SmallVector<std::size_t>(nqubits, UINT64_MAX)) {
    floydWarshallWithPathReconstruction();
    collectNeighbours();
  }

  /**
   * @brief Return the architecture's name.
   */
  [[nodiscard]] constexpr std::string_view name() const { return name_; }

  /**
   * @brief Return the architecture's number of qubits.
   */
  [[nodiscard]] constexpr std::size_t nqubits() const { return nqubits_; }

  /**
   * @brief Return true if @p u and @p v are adjacent.
   */
  [[nodiscard]] bool areAdjacent(QubitIndex u, QubitIndex v) const {
    return couplingSet_.contains({u, v});
  }

  /**
   * @brief Collect the shortest path between @p u and @p v.
   * @returns The path from the destination (v) to source (u) qubit.
   */
  [[nodiscard]] llvm::SmallVector<std::size_t>
  shortestPathBetween(QubitIndex u, QubitIndex v) const;

  /**
   * @brief Return the length of the shortest path between @p u and @p v.
   */
  [[nodiscard]] std::size_t distanceBetween(QubitIndex u, QubitIndex v) const;

  /**
   * @brief Collect all neighbours of @p u.
   */
  [[nodiscard]] llvm::SmallVector<QubitIndex, 4>
  neighboursOf(QubitIndex u) const;

private:
  using Matrix = llvm::SmallVector<llvm::SmallVector<std::size_t>>;

  /**
   * @brief Find all shortest paths in the coupling map between two qubits.
   * @details Vertices are the qubits. Edges connected two qubits. Has a time
   * and memory complexity of O(nqubits^3) and O(nqubits^2), respectively.
   * @link Adapted from https://en.wikipedia.org/wiki/Floyd–Warshall_algorithm
   */
  void floydWarshallWithPathReconstruction();

  /**
   * @brief Collect the neighbours of all qubits.
   * @details Has a time complexity of O(nqubits)
   */
  void collectNeighbours();

  std::string name_;
  std::size_t nqubits_;
  CouplingSet couplingSet_;
  NeighbourVector neighbours_;

  Matrix dist_;
  Matrix prev_;
};

/**
 * @brief Get architecture by its name.
 */
std::unique_ptr<Architecture> getArchitecture(llvm::StringRef name);

}; // namespace mqt::ir::opt
