#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <pybind11/stl.h>
#include "qkmeans/core.hpp"
#include "qkmeans/distance.hpp"
#include "qkmeans/anns/faiss_backend.hpp"

namespace py = pybind11;

PYBIND11_MODULE(_qkmeans_cpp, m) {
    m.doc() = R"pbdoc(
        QKMeans: Fast k-means++ seeding under the manifold hypothesis

        This module provides the C++ backend for the QKMeans algorithm,
        which implements rejection-sampling-based k-means++ seeding
        with theoretical guarantees under the manifold hypothesis.
    )pbdoc";

    // ANNS type enum
    py::enum_<qkmeans::anns::ANNSType>(m, "ANNSType", "Approximate nearest neighbor search backend types")
        .value("LSH", qkmeans::anns::ANNSType::LSH, "Locality-Sensitive Hashing")
        .value("FAISS", qkmeans::anns::ANNSType::FAISS, "Facebook AI Similarity Search")
        .value("HNSW", qkmeans::anns::ANNSType::HNSW, "Hierarchical Navigable Small World graphs")
        .value("BRUTE_FORCE", qkmeans::anns::ANNSType::BRUTE_FORCE, "Exact nearest neighbor (brute force)")
        .export_values();

#ifdef QKMEANS_HAS_FAISS
    // FAISS index type enum
    py::enum_<qkmeans::anns::FAISSIndexType>(m, "FAISSIndexType", "FAISS index types")
        .value("FLAT", qkmeans::anns::FAISSIndexType::FLAT, "IndexFlatL2 - exact brute force search")
        .value("HNSW", qkmeans::anns::FAISSIndexType::HNSW, "IndexHNSWFlat - HNSW graph-based search")
        .value("LSH", qkmeans::anns::FAISSIndexType::LSH, "IndexLSH - locality-sensitive hashing")
        .value("IVF", qkmeans::anns::FAISSIndexType::IVF, "IndexIVFFlat - inverted file index")
        .export_values();
#endif

    // Configuration struct
    py::class_<qkmeans::QKMeansConfig>(m, "QKMeansConfig", "Configuration for QKMeans algorithm")
        .def(py::init<>())
        .def_readwrite("k", &qkmeans::QKMeansConfig::k, "Number of clusters")
        .def_readwrite("max_rejections", &qkmeans::QKMeansConfig::max_rejections,
            "Maximum rejection iterations per center (0 = auto)")
        .def_readwrite("anns_type", &qkmeans::QKMeansConfig::anns_type,
            "ANNS backend type")
        .def_readwrite("random_seed", &qkmeans::QKMeansConfig::random_seed,
            "Random seed (0 = random)")
        .def_readwrite("n_threads", &qkmeans::QKMeansConfig::n_threads,
            "Number of threads (-1 = all available)")
        // LSH parameters
        .def_readwrite("lsh_num_tables", &qkmeans::QKMeansConfig::lsh_num_tables,
            "LSH: number of hash tables")
        .def_readwrite("lsh_num_hashes", &qkmeans::QKMeansConfig::lsh_num_hashes,
            "LSH: number of hash functions per table")
        .def_readwrite("lsh_rho", &qkmeans::QKMeansConfig::lsh_rho,
            "LSH: approximation parameter rho")
        // HNSW parameters
        .def_readwrite("hnsw_M", &qkmeans::QKMeansConfig::hnsw_M,
            "HNSW: M parameter (number of bi-directional links)")
        .def_readwrite("hnsw_ef_construction", &qkmeans::QKMeansConfig::hnsw_ef_construction,
            "HNSW: ef_construction parameter")
        .def_readwrite("hnsw_ef_search", &qkmeans::QKMeansConfig::hnsw_ef_search,
            "HNSW: ef_search parameter")
        // FAISS parameters
        .def_readwrite("faiss_index_type", &qkmeans::QKMeansConfig::faiss_index_type,
            "FAISS index type: 0=FLAT, 1=HNSW, 2=LSH, 3=IVF")
        .def_readwrite("faiss_hnsw_M", &qkmeans::QKMeansConfig::faiss_hnsw_M,
            "FAISS HNSW: M parameter (number of links)")
        .def_readwrite("faiss_hnsw_efConstruction", &qkmeans::QKMeansConfig::faiss_hnsw_efConstruction,
            "FAISS HNSW: efConstruction parameter")
        .def_readwrite("faiss_hnsw_efSearch", &qkmeans::QKMeansConfig::faiss_hnsw_efSearch,
            "FAISS HNSW: efSearch parameter")
        .def_readwrite("faiss_lsh_nbits", &qkmeans::QKMeansConfig::faiss_lsh_nbits,
            "FAISS LSH: number of bits (0 = auto: 2*dim)")
        .def_readwrite("faiss_nlist", &qkmeans::QKMeansConfig::faiss_nlist,
            "FAISS IVF: number of Voronoi cells/inverted lists")
        .def_readwrite("faiss_nprobe", &qkmeans::QKMeansConfig::faiss_nprobe,
            "FAISS IVF: number of cells to probe during search");

    // Result struct
    py::class_<qkmeans::QKMeansResult>(m, "QKMeansResult", "Result of QKMeans seeding")
        .def_readonly("centers", &qkmeans::QKMeansResult::centers,
            "Cluster centers as flattened array (k x d)")
        .def_readonly("center_indices", &qkmeans::QKMeansResult::center_indices,
            "Indices of data points chosen as centers")
        .def_readonly("total_cost", &qkmeans::QKMeansResult::total_cost,
            "Total k-means cost (sum of squared distances)")
        .def_readonly("total_rejections", &qkmeans::QKMeansResult::total_rejections,
            "Number of rejected samples")
        .def_readonly("total_accepted", &qkmeans::QKMeansResult::total_accepted,
            "Number of accepted samples")
        .def_readonly("fallback_samples", &qkmeans::QKMeansResult::fallback_samples,
            "Number of uniform fallback samples")
        .def_readonly("seeding_time_ms", &qkmeans::QKMeansResult::seeding_time_ms,
            "Time for seeding in milliseconds")
        .def_readonly("preprocess_time_ms", &qkmeans::QKMeansResult::preprocess_time_ms,
            "Time for preprocessing in milliseconds");

    // Main seeding function
    m.def("qkmeans_seed",
        [](py::array_t<double, py::array::c_style | py::array::forcecast> X,
           const qkmeans::QKMeansConfig& config) {
            py::buffer_info buf = X.request();

            if (buf.ndim != 2) {
                throw std::runtime_error("X must be a 2D array");
            }

            size_t n = buf.shape[0];
            size_t d = buf.shape[1];

            qkmeans::QKMeans algo(config);
            return algo.fit(static_cast<const double*>(buf.ptr), n, d);
        },
        py::arg("X"),
        py::arg("config"),
        R"pbdoc(
            Run QKMeans seeding algorithm.

            Parameters
            ----------
            X : numpy.ndarray
                Input data matrix of shape (n_samples, n_features), dtype float64.
            config : QKMeansConfig
                Algorithm configuration.

            Returns
            -------
            QKMeansResult
                Result containing centers and statistics.
        )pbdoc"
    );

    // Utility functions
    m.def("get_simd_level",
        []() {
            return qkmeans::distance::simd_level_name(
                qkmeans::distance::detect_simd_level()
            );
        },
        "Get the SIMD instruction set being used (SCALAR, SSE4, or AVX2)"
    );

    m.def("is_anns_available",
        [](qkmeans::anns::ANNSType type) {
            return qkmeans::anns::ANNSBase::is_available(type);
        },
        py::arg("type"),
        "Check if an ANNS backend is available"
    );

    // Version info
    m.attr("__version__") = "0.1.0";
}
