#pragma once

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

#include <disc/disc/Composition.hxx>

namespace sd::disc::pytools
{

namespace py = pybind11;

template <typename In, typename Out>
Out& pylist_to_itemset(const In& in, Out& out)
{
    out.clear();
    for (auto x : in)
    {
        out.insert(py::cast<size_t>(x));
    }
    return out;
}

template <typename Data>
void insert_pylists_into_data(const py::list& in, Data& out)
{
    sd::disc::itemset<typename Data::pattern_type> buffer;
    out.reserve(in.size());
    for (const auto& xs : in)
    {
        out.insert(pylist_to_itemset(xs, buffer));
    }
}

template <typename Data>
void insert_pylists_into_data(Data& data, const py::list& dataset, const py::list& labels)
{
    data.reserve(dataset.size());
    sd::disc::itemset<typename Data::pattern_type> buffer;

    if (dataset.size() != labels.size())
        throw std::domain_error("expect: dataset.size() == labels.size()");

    for (size_t i = 0; i < dataset.size(); ++i)
    {
        data.insert(pylist_to_itemset(dataset[i], buffer), py::cast<size_t>(labels[i]));
    }
}

template <typename Data>
void insert_data_into_pylist(const Data& in, const py::list& out)
{
    using namespace sd::disc;
    for (const auto& x : in)
    {
        py::list xx;
        sd::foreach(point(x), [&](size_t i) { xx.append(i); });
        out.append(std::move(xx));
    }
}

template <typename Trait>
py::dict translate_to_pydict(PatternsetResult<Trait> const& co)
{
    using float_type = typename Trait::float_type;
    py::list patternset;
    insert_data_into_pylist(co.summary, patternset);
    auto            frequencies = py::array_t<float_type>(co.summary.size());
    py::buffer_info info        = frequencies.request();
    std::copy_n(co.summary.template col<0>().data(),
                co.summary.size(),
                static_cast<float_type*>(info.ptr));

    py::dict r;

    r["pattern_set"]       = patternset;
    r["frequencies"]       = frequencies;
    r["initial_objective"] = co.initial_encoding.objective();
    r["objective"]         = co.encoding.objective();

    return r;
}

template <typename Trait>
py::dict translate_to_pydict(Composition<Trait> const& co)
{
    using float_type = typename Trait::float_type;
    py::list patternset;
    insert_data_into_pylist(co.summary, patternset);
    auto frequencies = py::array_t<float_type>();
    {
        frequencies.resize({co.frequency.extent(0), co.frequency.extent(1)});
        py::buffer_info info = frequencies.request();
        std::copy_n(
            co.frequency.data(), co.frequency.size(), static_cast<float_type*>(info.ptr));
    }

    py::list assignment_matrix;
    for (const auto& r : co.assignment)
    {
        py::list xx;
        for (auto i : r.container)
        {
            xx.append(i);
        };
        assignment_matrix.append(std::move(xx));
    }

    auto labels = py::array_t<int>(co.data.size());
    {
        py::buffer_info info = labels.request();
        std::copy_n(
            co.data.template col<0>().data(), co.data.size(), static_cast<int*>(info.ptr));
    }

    py::dict r;

    r["pattern_set"]       = patternset;
    r["frequencies"]       = frequencies;
    r["assignment_matrix"] = assignment_matrix;
    r["labels"]            = labels;
    r["initial_objective"] = co.initial_encoding.objective();
    r["objective"]         = co.encoding.objective();

    return r;
}

template <typename S>
constexpr const char* storage_type_to_str()
{
    if constexpr (std::is_same_v<S, tag_dense>)
    {
        return "dense";
    }
    else
    {
        return "sparse";
    }
}

template <typename T>
constexpr const char* float_storage_type_to_str()
{
    if constexpr (std::is_same_v<T, double>)
    {
        return "double";
    }
    else if constexpr (std::is_same_v<T, long double>)
    {
        return "ldouble";
    }
#if HAS_QUADMATH
    else if constexpr (std::is_same_v<T, boost::multiprecision::float128>)
    {
        return "float128";
    }
#endif
    else
    {
        return "unknown";
    }
}

} // namespace sd::disc::pytools
