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

#include "ioh.hpp"

namespace py = pybind11;
using namespace ioh::suite;

template <typename SuiteType>
void define_base_class(py::module &m, const std::string &name)
{
    py::class_<SuiteType, std::shared_ptr<SuiteType>>(m, name.c_str(), py::buffer_protocol(),
                                                      fmt::format("{} suite", name).c_str())
        .def("reset", &SuiteType::reset,
             R"pbdoc(
                Reset the state variables of the current problem.
            )pbdoc")
        .def("attach_logger", &SuiteType::attach_logger,
             R"pbdoc(
                Attach a logger to the problem to track the performance and internal variables.
            )pbdoc")
        .def("detach_logger", &SuiteType::detach_logger,
             R"pbdoc(
                Remove the specified logger from the problem
            )pbdoc")
        .def_property_readonly(
            "problem_ids",
            [](const SuiteType &c) {
                const auto problem_ids = c.problem_ids();
                return py::array(problem_ids.size(), problem_ids.data());
            },
            R"pbdoc(
                The list of all problems ids contained in the current suite.
            )pbdoc")
        .def_property_readonly(
            "dimensions",
            [](const SuiteType &c) {
                const auto dimensions = c.dimensions();
                return py::array(dimensions.size(), dimensions.data());
            },
            R"pbdoc(
                The list of all problems ids contained in the current suite.
            )pbdoc")
        .def_property_readonly(
            "instances",
            [](const SuiteType &c) {
                const auto instances = c.instances();
                return py::array(instances.size(), instances.data());
            },
            R"pbdoc(
                The list of all instance ids contained in the current suite.
            )pbdoc")
        .def_property_readonly("name", &SuiteType::name,
                               R"pbdoc(
                The name of the suite.
            )pbdoc")
        .def("__len__", &SuiteType::size)
        .def(
            "__iter__", [](SuiteType &s) { return py::make_iterator(s.begin(), s.end()); }, py::keep_alive<0, 1>());
}


template <typename S, typename P>
void define_suite(py::module &m, const std::string &name, const int default_dim = 5)
{
    py::class_<S, P, std::shared_ptr<S>>(m, name.c_str(),
                                         fmt::format(R"pbdoc(
        Suite with {} functions included in IOH. 

        This class is an iterator for problems, you can define a set of problems to iterate over.

        Parameters
        ----------
        problem_ids: list[int]
            A list of problems ids to include in this instantiation of the suite
        instances: list[int]
            A list of problem instances to include in this instantiation of the suite
        problem_ids: list[int]
            A list of problem dimensions to include in this instantiation of the suite          
    )pbdoc",
                                                     name)
                                             .c_str())
        .def(py::init<std::vector<int>, std::vector<int>, std::vector<int>>(), py::arg("problem_ids"),
             py::arg("instances") = std::vector<int>{1}, py::arg("dimensions") = std::vector<int>{default_dim});
}

void define_suites(py::module &m)
{
    using RealBase = Suite<ioh::problem::RealSingleObjective>;
    define_base_class<RealBase>(m, "RealBase");
    define_suite<Real, RealBase>(m, "Real");
    define_suite<BBOB, RealBase>(m, "BBOB");
    define_suite<SBOX, RealBase>(m, "SBOX");
    define_suite<RealStarDiscrepancy, RealBase>(m, "RealStarDiscrepancy");

    using IntegerBase = Suite<ioh::problem::IntegerSingleObjective>;
    define_base_class<IntegerBase>(m, "IntegerBase");
    define_suite<Integer, IntegerBase>(m, "Integer");
    define_suite<PBO, IntegerBase>(m, "PBO", 16);
    define_suite<Submodular, IntegerBase>(m, "Submodular", 1);
    define_suite<IntegerStarDiscrepancy, IntegerBase>(m, "IntegerStarDiscrepancy");
}
