/*
SPDX-FileCopyrightText: Copyright 2022, Arm Limited and/or its affiliates.
SPDX-License-Identifier: Apache-2.0
*/
#include "tosa_checker.h"

#include <optional>
#include <sstream>
#include <string>

#include <pybind11/pybind11.h>
#include <pybind11/stl.h>

PYBIND11_MODULE(_tosa_checker_wrapper, m) {
  /**
   * tosa_checker::TOSAChecker
   */
  pybind11::class_<tosa_checker::TOSAChecker> tosa_checker_class(m,
                                                                 "TOSAChecker");
  tosa_checker_class.def(pybind11::init<const std::string&>(),
                         pybind11::arg("model_path"));

  tosa_checker_class.def(
      "is_tosa_compatible",
      [](tosa_checker::TOSAChecker& tc) { return tc.IsTOSACompatible(); },
      "Check if a model is compatible with the TOSA specification");

  tosa_checker_class.def(
      "_get_tosa_compatibility_for_ops",
      [](tosa_checker::TOSAChecker& tc, bool elide_large_elements_attrs) {
        return tc.GetTOSACompatibilityForOps(elide_large_elements_attrs);
      },
      pybind11::arg("elide_large_elements_attrs") = false,
      "Get all the operators of the models with a TOSA compatibility flag for "
      "each operator");

  tosa_checker_class.def(
      "_get_used_tosa_ops",
      [](tosa_checker::TOSAChecker& tc, bool elide_large_elements_attrs) {
        return tc.GetUsedTOSAOps(elide_large_elements_attrs);
      },
      pybind11::arg("elide_large_elements_attrs") = false,
      "Get the TOSA operators used by the model after its TOSA legalization");

  tosa_checker_class.def(
      "_get_mlir_model_representation",
      [](tosa_checker::TOSAChecker& tc, bool elide_large_elements_attrs) {
        return tc.GetMLIRModelRepresentation(elide_large_elements_attrs);
      },
      pybind11::arg("elide_large_elements_attrs") = false,
      "Get the MLIR representation of the model");

  tosa_checker_class.def(
      "_get_mlir_tosa_model_representation",
      [](tosa_checker::TOSAChecker& tc, bool elide_large_elements_attrs) {
        return tc.GetMLIRTOSAModelRepresentation(elide_large_elements_attrs);
      },
      pybind11::arg("elide_large_elements_attrs") = false,
      "Get the MLIR representation of the TOSA legalized model");

  /**
   * tosa_checker::TOSAChecker::Operator
   */
  pybind11::class_<tosa_checker::TOSAChecker::Operator>(tosa_checker_class,
                                                        "_Operator")
      .def_readonly("name", &tosa_checker::TOSAChecker::Operator::name)
      .def_readonly("location", &tosa_checker::TOSAChecker::Operator::location)
      .def_readonly("attributes",
                    &tosa_checker::TOSAChecker::Operator::attributes)
      .def_readonly("is_tosa_compatible",
                    &tosa_checker::TOSAChecker::Operator::is_tosa_compatible)
      .def_readonly("mlir_representation",
                    &tosa_checker::TOSAChecker::Operator::mlir_representation)
      .def("__repr__", [](const tosa_checker::TOSAChecker::Operator& o) {
        std::stringstream stream;
        stream << o;
        return stream.str();
      });
}