#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <random>
#include <map>
#include <set>
#include <string>
#include <vector>

class Dice {
    std::vector<std::string> sides;
    std::vector<double> probs;
    std::mt19937 gen;
    std::discrete_distribution<> dist;

public:
    size_t accepted = 0;
    size_t rejected = 0;

    Dice(const std::vector<std::string>& sides, const std::vector<double>& probs)
        : sides(sides), probs(probs), gen(std::random_device{}()), dist(probs.begin(), probs.end()) {}

    std::string fair_dice() {
        while (true) {
            std::vector<std::string> res;
            for (size_t i = 0; i < sides.size(); ++i)
                res.push_back(sides[dist(gen)]);
            std::set<std::string> uniq(res.begin(), res.end());
            if (uniq.size() == res.size()) {
                ++accepted;
                return res[0];
            }
            ++rejected;
        }
    }

    std::map<std::string,int> throw_dice(int n) {
        std::map<std::string,int> res;
        for (auto& s : sides) res[s] = 0;
        for (int i = 0; i < n; ++i)
            res[fair_dice()]++;
        return res;
    }

    std::pair<size_t,size_t> stats() const { return {accepted, rejected}; }

    void reset_stats() { accepted = rejected = 0; }
};

namespace py = pybind11;

PYBIND11_MODULE(dice, m) {
    py::class_<Dice>(m, "Dice")
        .def(py::init<const std::vector<std::string>&, const std::vector<double>&>())
        .def("fair_dice", &Dice::fair_dice)
        .def("throw_dice", &Dice::throw_dice)
        .def("stats", &Dice::stats)
        .def("reset_stats", &Dice::reset_stats)
        .def_readwrite("accepted", &Dice::accepted)
        .def_readwrite("rejected", &Dice::rejected);
}