#include "PyStage.h"

#include "PyScheduleMethods.h"

namespace Halide {
namespace PythonBindings {

void define_stage(py::module &m) {
    auto stage_class =
        py::class_<Stage>(m, "Stage")
            .def("dump_argument_list", &Stage::dump_argument_list)
            .def("name", &Stage::name)

            .def("rfactor", (Func(Stage::*)(std::vector<std::pair<RVar, Var>>)) & Stage::rfactor,
                 py::arg("preserved"))
            .def("rfactor", (Func(Stage::*)(const RVar &, const Var &)) & Stage::rfactor,
                 py::arg("r"), py::arg("v"))

            // These two variants of compute_with are specific to Stage
            .def("compute_with", (Stage & (Stage::*)(LoopLevel, const std::vector<std::pair<VarOrRVar, LoopAlignStrategy>> &)) & Stage::compute_with,
                 py::arg("loop_level"), py::arg("align"))
            .def("compute_with", (Stage & (Stage::*)(LoopLevel, LoopAlignStrategy)) & Stage::compute_with,
                 py::arg("loop_level"), py::arg("align") = LoopAlignStrategy::Auto);
    add_schedule_methods(stage_class);
}

}  // namespace PythonBindings
}  // namespace Halide
